Plugins: Allow exchanging analysis providers and inference providers alike

docs/PLUGIN.md

@@ -24,7 +24,8 @@ Everything a plugin can do, at a glance. The last column tells you where to find
 | 12 | Named worker tasks | `add_task` | No - see "API reference" |
 | 13 | Startup hooks | `on_flask_start` / `on_worker_start` | No - see "The plugin lifecycle" |
 | 14 | Create media-server playlists | `tasks.mediaserver` | No - see "Create a playlist on your media server" |
-| 15 | Extra ONNX execution provider | `register_onnx_provider` | No - see "API reference" (advanced) |
+| 15 | Extra ONNX execution provider (optionally scoped per model) | `register_onnx_provider` | No - see "API reference" (advanced) |
+| 16 | Replace an analysis component (e.g. the ASR/Whisper backend) | `register_analysis_provider` | No - see "API reference" (advanced) |
 
 ## Installing and managing plugins
 
@@ -468,7 +469,66 @@ And the methods on the `ctx` object in `register(ctx)`:
 | `add_task(name, func, queue='default')` | Register a named worker task. It appears on the Scheduled Tasks page like a cron task, so the admin can schedule it or run it now. |
 | `on_install(func)` | Run once at install and on every update. Gets the database connection. |
 | `on_flask_start(func)` / `on_worker_start(func)` | Run at every start of that container, after the plugin loads. |
-| `register_onnx_provider(name, options, position)` | Advanced: offer an extra ONNX Runtime execution provider (for example a GPU) for analysis. |
+| `register_onnx_provider(name, options, position, only_models=None, exclude_models=None, needs_static_shapes=False)` | Advanced: offer an extra ONNX Runtime execution provider (for example a GPU) for analysis. See "Offer an ONNX execution provider" below. |
+| `register_analysis_provider(component, factory, cache=True)` | Advanced: replace a whole analysis component with a plugin-supplied implementation. `component` is currently `asr` (the Whisper backend); `factory` is the replacement module/object (or a zero-arg callable returning one) matching the built-in surface `load_whisper_model`/`transcribe`/`is_loaded`/`unload`. See "Replace an analysis component" below for what each method must return. Consulted before the built-in; an unknown `component` or a backend missing part of that surface is ignored with a warning in the logs, and the built-in is used. `cache=True` (the default) resolves a callable `factory` once and reuses the result, like the built-in modules, which stay loaded for a whole album; pass `cache=False` to be called for every use, and then unload what you hand out yourself. |
+
+### Offer an ONNX execution provider
+
+`register_onnx_provider` adds your provider to the chain AudioMuse-AI builds for every analysis model. `position` is `before_cpu` (the default) or `before_cuda` when your provider should win over CUDA. Set `needs_static_shapes=True` if your graph compiler cannot handle symbolic dimensions, so core pins them before it builds the session; today core pins them only for the CLAP audio model. Other models (the GTE embedder and the Whisper decoder, for example) also carry symbolic axes but are never pinned, so a provider that cannot compile them must be scoped away from those models with `only_models`/`exclude_models`.
+
+Your provider must already exist in the `onnxruntime` build of the image you run, because core only offers providers that `onnxruntime.get_available_providers()` reports. A plugin cannot add one by pip-installing into its own `_lib`: execution providers are compiled into `onnxruntime` itself, so a provider that needs a different build needs a different image. A provider core has already put in the chain for that session (`CUDAExecutionProvider`, or `CoreMLExecutionProvider` on macOS) wins: your entry for that same name is dropped and core's own options are kept, so use `position` to order yourself around them instead of re-registering them.
+
+Scope the provider to the models it can actually compile with `only_models`/`exclude_models`. Both take a session label or a list of them; an unknown label is logged as a warning. The labels are:
+
+| Label | Model | Default without a plugin |
+|---|---|---|
+| `musicnn` | MusiCNN embedding + prediction | CUDA, else CPU |
+| `clap` | CLAP audio model | CUDA / CoreML, else CPU |
+| `clap_text` | CLAP text model, on the worker | CUDA, else CPU |
+| `whisper_encoder` | Whisper encoder (lyrics) | CUDA, else CPU |
+| `whisper_decoder` | Whisper decoder (lyrics) | CUDA, else CPU |
+| `gte` | GTE lyrics text embedding | CPU only |
+| `silero_vad` | Silero voice activity detection | CPU only |
+
+The last two are CPU by default and stay that way unless a plugin names them in `only_models`, so nothing changes for a normal user. The CLAP text model runs on CPU in the Flask process for thread safety and is not offered to plugins there; `clap_text` applies on the worker only.
+
+```python
+def register(ctx):
+    # A provider that handles MusiCNN and the Whisper encoder, but not the
+    # decoder and not CLAP.
+    ctx.register_onnx_provider(
+        'MyExecutionProvider',
+        {'device_id': 0},
+        only_models=['musicnn', 'whisper_encoder'],
+    )
+```
+
+### Replace an analysis component
+
+Sometimes a different execution provider is not enough and the step needs a different library altogether: some execution providers cannot run the ONNX Whisper decoder at all, so a plugin can swap in faster-whisper instead. `register_analysis_provider(component, factory)` hands core a full replacement for one step. The only component today is `asr`, the Whisper backend used by the lyrics pipeline.
+
+Your replacement must expose all four methods below. Core checks for them when it resolves your plugin and falls back to the built-in backend - with a warning naming your plugin and what is missing - if any is absent.
+
+| Method | Must return                                                                                                                                                                                         |
+|---|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `load_whisper_model()` | Nothing. Load the model, or raise an exception if it cannot be loaded.                                                                                                                              |
+| `transcribe(wav, sr, language=None)` | A dict: `text` (the transcript, `''` when nothing was recognised), `language` (the detected language code, for example `en`), `avg_logprob` (the mean per-token log probability of the transcript). |
+| `is_loaded()` | `True` while your model is in memory. Core uses it to decide whether a memory cleanup is needed after an album.                                                                                     |
+| `unload()` | Nothing. Free the model and its memory.                                                                                                                                                             |
+
+`avg_logprob` is the one that quietly matters: core drops a transcript whose confidence is below `LYRICS_ASR_MIN_AVG_LOGPROB` (and below `LYRICS_ASR_NON_ENGLISH_MIN_LOGPROB` for non-English) and treats the song as instrumental. Return the real value if your backend has one. If it does not, leave the key out entirely - core then skips the confidence check instead of reading a missing value as the worst possible score. Do not return a placeholder like `0.0`, which disables the check without saying so. An explicit `None` or a NaN is treated like a missing key. Infinities are not: they pass through as real values, so `-inf` - the built-in backend's own worst-confidence value - always drops the transcript. A value of `0.0` or above is kept but flagged in the log, because no real mean log-probability is positive.
+
+```python
+def register(ctx):
+    ctx.register_analysis_provider('asr', _load_backend)
+
+
+def _load_backend():
+    # Return None to stand down (no GPU here, say); core logs it and keeps the
+    # built-in backend. Returning the module keeps it loaded for a whole album.
+    from . import my_whisper
+    return my_whisper
+```
 
 ## Who can see and manage plugins
 

lyrics/__init__.py

@@ -82,10 +82,13 @@ def is_lyrics_loaded() -> bool:
     if not _LYRICS_ENABLED:
         return False
     try:
-        from . import whisper_onnx, gte_onnx, silero_onnx
+        from ._asr_backend import get_asr_backend
+        from . import gte_onnx, silero_onnx
+
+        whisper_mod = get_asr_backend()
     except Exception:
         return False
-    for mod in (whisper_onnx, gte_onnx, silero_onnx):
+    for mod in (whisper_mod, gte_onnx, silero_onnx):
         try:
             if mod.is_loaded():
                 return True
@@ -100,12 +103,13 @@ def unload_lyrics_models() -> bool:
     released_any = False
     try:
         try:
-            from . import whisper_onnx
+            from ._asr_backend import get_asr_backend
 
-            if whisper_onnx.is_loaded():
-                released_any = bool(_safe_call('whisper_onnx.unload', whisper_onnx.unload))
+            whisper_mod = get_asr_backend()
+            if whisper_mod.is_loaded():
+                released_any = bool(_safe_call('whisper.unload', whisper_mod.unload))
         except Exception as exc:
-            _logger.warning("Lyrics whisper_onnx release failed: %s", exc)
+            _logger.warning("Lyrics whisper release failed: %s", exc)
 
         try:
             from . import gte_onnx

lyrics/_asr_backend.py

@@ -0,0 +1,83 @@
+# AudioMuse-AI - https://github.com/NeptuneHub/AudioMuse-AI
+# Copyright (C) 2025 NeptuneHub
+# SPDX-License-Identifier: AGPL-3.0-only
+#
+# This program is free software: you can redistribute it and/or modify it under
+# the terms of the GNU Affero General Public License v3.0. See the LICENSE file
+# in the project root or <https://github.com/NeptuneHub/AudioMuse-AI/blob/main/LICENSE>
+
+"""Resolves the Whisper/ASR backend for the lyrics pipeline.
+
+The default is the built-in ONNX backend (whisper_onnx). A plugin can replace it
+by registering an alternative with ``ctx.register_analysis_provider('asr', ...)``
+- used, for example, by an AMD/ROCm plugin that swaps in faster-whisper because
+MIGraphX cannot run the ONNX Whisper decoder. The replacement must expose the
+same public surface: ``load_whisper_model()``, ``transcribe(wav, sr, language=None)``,
+``is_loaded()`` and ``unload()``.
+
+``transcribe`` must return a dict with ``text`` (the transcript), ``language``
+(the detected language code) and ``avg_logprob`` (the mean per-token log
+probability, used to gate transcript quality). A backend that cannot report a
+confidence should leave ``avg_logprob`` out entirely: the quality gate is then
+skipped instead of reading a value that would drop every transcript.
+
+Main Features:
+* Consults the loaded plugins for an 'asr' analysis provider before the built-in
+* Falls back to the built-in whisper_onnx backend when no plugin registered one,
+  or when the replacement is missing part of the required surface
+* Rejection of an incomplete replacement is warned on every resolution, naming
+  the owning plugin so the user knows which one to fix or remove
+* Swallows any plugin-resolution error so a broken plugin never breaks lyrics
+"""
+
+from __future__ import annotations
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+REQUIRED_METHODS = ('load_whisper_model', 'transcribe', 'is_loaded', 'unload')
+
+
+def get_asr_backend():
+    override = _plugin_asr_backend()
+    if override is not None and _has_required_surface(override):
+        return override
+    from . import whisper_onnx
+
+    return whisper_onnx
+
+
+def _plugin_asr_backend():
+    try:
+        from plugin.manager import plugin_manager
+
+        return plugin_manager.get_analysis_provider('asr')
+    except Exception:
+        return None
+
+
+def _provider_owner():
+    try:
+        from plugin.manager import plugin_manager
+
+        return plugin_manager.analysis_provider_owner('asr')
+    except Exception:
+        return None
+
+
+# Rejects a backend that cannot stand in for whisper_onnx. A missing method would
+# surface far from here: a broken is_loaded makes is_lyrics_loaded() report True
+# forever, so every album pays a full memory cleanup for nothing, and a broken
+# transcribe fails every song with no way back to the built-in.
+def _has_required_surface(backend):
+    missing = [
+        name for name in REQUIRED_METHODS if not callable(getattr(backend, name, None))
+    ]
+    if not missing:
+        return True
+    logger.warning(
+        'Plugin %s: its ASR backend %r is missing %s; using the built-in whisper instead',
+        _provider_owner() or 'unknown', backend, ', '.join(missing),
+    )
+    return False

lyrics/gte_onnx.py

@@ -98,11 +98,11 @@ def load_gte_model():
         sess_options.intra_op_num_threads = max(1, (os.cpu_count() or 2) // 2)
         sess_options.inter_op_num_threads = 1
         try:
-            from tasks.analysis.song import create_onnx_session
+            from tasks.analysis.song import create_onnx_session, resolve_providers
 
             session = create_onnx_session(
                 onnx_path,
-                provider_options=[('CPUExecutionProvider', {})],
+                provider_options=resolve_providers(label='gte', cpu_only_default=True),
                 sess_options=sess_options,
                 label='gte',
             )

lyrics/lyrics_transcriber.py

@@ -28,6 +28,7 @@
 from __future__ import annotations
 
 import logging
+import math
 import os
 import re
 import signal
@@ -218,9 +219,9 @@ def _apply_thread_env(num_threads: int) -> None:
 def load_asr_model(num_threads: Optional[int] = None):
     threads = num_threads or get_lyrics_threads()
     _apply_thread_env(threads)
-    from .whisper_onnx import load_whisper_model as _load
+    from ._asr_backend import get_asr_backend
 
-    return _load()
+    return get_asr_backend().load_whisper_model()
 
 
 def load_topic_embedding_model(model_name: Optional[str] = None):
@@ -704,9 +705,9 @@ def _transcribe(
 ) -> Dict[str, object]:
     if audio is None or len(audio) == 0:
         return {'text': '', 'language': language or '', 'duration': 0.0}
-    from .whisper_onnx import transcribe as _whisper_transcribe
+    from ._asr_backend import get_asr_backend
 
-    return _whisper_transcribe(audio, sr, language=language)
+    return get_asr_backend().transcribe(audio, sr, language=language)
 
 
 def _embed_text(text: str, tokenizer, model) -> Optional[np.ndarray]:
@@ -833,16 +834,25 @@ def _lyrics_result(
     }
 
 
+# Decides whether an ASR transcript is too poor to keep. asr_avg_logprob is None
+# when the backend reported no confidence at all (see _asr_confidence): the
+# confidence thresholds are then skipped rather than treated as the worst possible
+# score, so a plugin backend that omits the key loses no transcript. The language
+# checks still apply.
 def _asr_should_drop(
-    raw_text: str, whisper_raw_len: int, asr_lang: str, asr_avg_logprob: float
+    raw_text: str, whisper_raw_len: int, asr_lang: str, asr_avg_logprob: Optional[float]
 ) -> bool:
     if not raw_text or whisper_raw_len <= 0:
         return False
-    if asr_avg_logprob < ASR_MIN_AVG_LOGPROB:
+    if asr_avg_logprob is not None and asr_avg_logprob < ASR_MIN_AVG_LOGPROB:
         return True
     if asr_lang in _ASR_NULL_LANGS:
         return True
-    if asr_lang not in _ASR_ENGLISH_LANGS and asr_avg_logprob < ASR_NON_ENGLISH_MIN_LOGPROB:
+    if (
+        asr_lang not in _ASR_ENGLISH_LANGS
+        and asr_avg_logprob is not None
+        and asr_avg_logprob < ASR_NON_ENGLISH_MIN_LOGPROB
+    ):
         return True
     return False
 
@@ -975,17 +985,45 @@ def _alarm_handler(signum, frame):
             signal.signal(signal.SIGALRM, _old_handler)
 
 
-def _postprocess_asr(transcription: Dict[str, object]) -> Tuple[str, int, str, float, str]:
+# Reads the backend's avg_logprob, or None when it reported none. The built-in
+# backend always returns it, but a plugin ASR backend registered with
+# register_analysis_provider may not. None means "unknown confidence" and switches
+# off the confidence gates in _asr_should_drop; assuming the worst score here would
+# silently turn a whole library into instrumentals.
+def _asr_confidence(transcription: Dict[str, object]) -> Optional[float]:
+    value = transcription.get('avg_logprob')
+    if value is None:
+        return None
+    try:
+        confidence = float(value)
+    except (TypeError, ValueError, OverflowError):
+        logger.warning('ASR backend returned a non-numeric avg_logprob %r - ignoring it', value)
+        return None
+    if math.isnan(confidence):
+        logger.warning('ASR backend returned a NaN avg_logprob - treating confidence as unknown')
+        return None
+    if confidence >= 0:
+        logger.warning(
+            'ASR backend returned avg_logprob %s but a real mean log-probability is negative - '
+            'the confidence gate cannot drop this transcript',
+            confidence,
+        )
+    return confidence
+
+
+def _postprocess_asr(
+    transcription: Dict[str, object]
+) -> Tuple[str, int, str, Optional[float], str]:
     raw_text = _sanitize_lyrics_text((transcription.get('text') or '').strip())
     whisper_raw_len = len(raw_text)
     asr_lang = (transcription.get('language') or '').strip().lower()
-    asr_avg_logprob = float(transcription.get('avg_logprob', float('-inf')))
+    asr_avg_logprob = _asr_confidence(transcription)
     detected_lang = asr_lang or 'en'
     logger.info(
-        'STEP 5 end: transcript length=%s chars / asr_lang=%r / avg_logprob=%.2f',
+        'STEP 5 end: transcript length=%s chars / asr_lang=%r / avg_logprob=%s',
         len(raw_text),
         asr_lang,
-        asr_avg_logprob,
+        'n/a' if asr_avg_logprob is None else '%.2f' % asr_avg_logprob,
     )
     logger.info('STEP 5 raw ASR output: %s', raw_text or '<empty>')
     _resolved, _script, _reject = _resolve_lang_and_quality(raw_text, asr_lang)
@@ -1118,7 +1156,7 @@ def analyze_lyrics(
     used_seconds = 0.0
     detected_lang = 'en'
     asr_lang = 'en'
-    asr_avg_logprob = 0.0
+    asr_avg_logprob = None
     whisper_raw_len = 0
 
     normalized_moods, vocal_prior = _compute_vocal_prior(top_moods)
@@ -1152,7 +1190,9 @@ def analyze_lyrics(
 
     if _asr_should_drop(raw_text, whisper_raw_len, asr_lang, asr_avg_logprob):
         logger.info(
-            'STEP 7: dropping ASR transcript (lang=%r, logprob=%.2f)', asr_lang, asr_avg_logprob
+            'STEP 7: dropping ASR transcript (lang=%r, logprob=%s)',
+            asr_lang,
+            'n/a' if asr_avg_logprob is None else '%.2f' % asr_avg_logprob,
         )
         raw_text = ''
     logger.info('STEP 7 end: language=%s, kept_text=%s', detected_lang, bool(raw_text))

lyrics/silero_onnx.py

@@ -68,11 +68,11 @@ def _load_session(model_path: Optional[str] = None):
         opts.enable_cpu_mem_arena = False
         opts.enable_mem_pattern = False
         try:
-            from tasks.analysis.song import create_onnx_session
+            from tasks.analysis.song import create_onnx_session, resolve_providers
 
             _session = create_onnx_session(
                 path,
-                provider_options=[('CPUExecutionProvider', {})],
+                provider_options=resolve_providers(label='silero_vad', cpu_only_default=True),
                 sess_options=opts,
                 label='silero_vad',
             )