Plugins: Allow exchanging analysis providers and inference providers alike

.gitignore

@@ -29,6 +29,8 @@ env/
 # IDE and editor configuration
 .vscode/
 .idea/
+.tokensave
+.claude/.headroom_wrap_marker.json
 
 # macOS Finder / metadata files
 .DS_Store
@@ -76,3 +78,4 @@ analysis.md
 .claude/
 lyrics_model_gte_vnni.tar.gz
 native-build/windows/vendor/redis/amd64/redis-server.exe
+.venv-test

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,37 @@ 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`. 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 (the CLAP audio model needs this).
+
+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'],
+    )
+```
 
 ## 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,66 @@
+# 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()``.
+
+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
+* 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
+
+
+# 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 ASR backend %r is missing %s; using the built-in whisper instead',
+        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

@@ -218,9 +218,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 +704,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]:

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',
             )

plugin/api.py

@@ -50,6 +50,12 @@
 _ID_RE = re.compile(r'^[a-z][a-z0-9_]{1,63}$')
 _NAME_RE = re.compile(r'^[a-z][a-z0-9_]{0,62}$')
 
+# Analysis steps a plugin may replace wholesale with register_analysis_provider.
+ANALYSIS_COMPONENTS = frozenset({'asr'})
+
+# Where a plugin provider goes in the ONNX chain, see register_onnx_provider.
+ONNX_POSITIONS = frozenset({'before_cuda', 'before_cpu'})
+
 __all__ = [
     'PluginContext', 'config', 'logger', 'get_db', 'save_task_status',
     'get_score_data_by_ids', 'get_tracks_by_ids', 'get_setting', 'set_setting',
@@ -183,6 +189,19 @@ def use_server(server_id):
     return ms_registry.bind({'server_id': server_id})
 
 
+def _model_scope(value):
+    """Normalize an only_models/exclude_models argument to a list or None.
+
+    A single label is accepted as a plain string; without this ``'musicnn'``
+    would become ``['m', 'u', 's', ...]`` and silently match nothing.
+    """
+    if not value:
+        return None
+    if isinstance(value, str):
+        return [value]
+    return list(value)
+
+
 class PluginContext:
     """Registration sink passed to a plugin's ``register(ctx)``.
 
@@ -202,6 +221,7 @@ def __init__(self, plugin_id, role):
         self.cron_tasks = {}
         self.tasks = {}
         self.onnx_providers = []
+        self.analysis_providers = {}
         self.flask_start = []
         self.worker_start = []
         self.song_analyzed_hooks = []
@@ -227,8 +247,32 @@ def add_task(self, name, func, queue='default'):
     def add_cron_task(self, name, func, queue='default'):
         self.cron_tasks[name] = {'dotted': dotted_path(func), 'queue': queue}
 
-    def register_onnx_provider(self, name, options=None, position='before_cpu'):
-        self.onnx_providers.append({'name': name, 'options': options or {}, 'position': position})
+    def register_onnx_provider(self, name, options=None, position='before_cpu',
+                               only_models=None, exclude_models=None,
+                               needs_static_shapes=False):
+        """Offer an extra ONNX Runtime execution provider for the analysis sessions.
+
+        ``only_models``/``exclude_models`` scope the provider to specific session
+        labels (a single label may be given as a plain string). ``position`` is
+        ``'before_cpu'`` (the default) or ``'before_cuda'``. Set
+        ``needs_static_shapes`` when the provider's graph compiler cannot handle
+        symbolic dimensions, so core pins them before building the session.
+        """
+        if position not in ONNX_POSITIONS:
+            logger.warning(
+                "Plugin %s registered ONNX provider %s with unknown position %r; "
+                "using 'before_cpu'. Valid positions: %s",
+                self.plugin_id, name, position, sorted(ONNX_POSITIONS),
+            )
+            position = 'before_cpu'
+        self.onnx_providers.append({
+            'name': name,
+            'options': options or {},
+            'position': position,
+            'only_models': _model_scope(only_models),
+            'exclude_models': _model_scope(exclude_models),
+            'needs_static_shapes': bool(needs_static_shapes),
+        })
 
     def on_flask_start(self, func):
         self.flask_start.append(func)
@@ -251,9 +295,36 @@ def on_song_analyzed(self, func):
     def on_install(self, func):
         self.install_hooks.append(func)
 
-    def register_analysis_provider(self, *args, **kwargs):
-        logger.info(
-            'Plugin %s called register_analysis_provider; alternative analysis models '
-            'are a forward seam and are not wired in this version.',
-            self.plugin_id,
-        )
+    def register_analysis_provider(self, component, factory, cache=True):
+        """Replace a whole analysis component with a plugin-supplied implementation.
+
+        Some accelerators need more than a different ONNX execution provider: they
+        need a different library entirely. MIGraphX, for instance, cannot run the
+        ONNX Whisper decoder at all, so an AMD plugin swaps in faster-whisper.
+
+        ``component`` names the step to replace, one of ANALYSIS_COMPONENTS.
+        ``factory`` is the replacement module/object, or a zero-arg callable
+        returning one. It must match the built-in module's public surface; for
+        ``asr`` that is ``load_whisper_model()``, ``transcribe(wav, sr,
+        language=None)``, ``is_loaded()`` and ``unload()``. Core consults the
+        registered provider first and falls back to the built-in when no plugin
+        registered one, or when the replacement is missing part of that surface.
+
+        A callable ``factory`` is resolved once and the result reused, matching the
+        built-in modules, which stay loaded for a whole album and are freed at its
+        end. Pass ``cache=False`` to be called for every use instead; then the
+        plugin owns unloading whatever it hands out.
+        """
+        if component not in ANALYSIS_COMPONENTS:
+            logger.warning(
+                'Plugin %s registered an analysis provider for unknown component %r; '
+                'ignoring it. Known components: %s',
+                self.plugin_id, component, sorted(ANALYSIS_COMPONENTS),
+            )
+            return
+        if component in self.analysis_providers:
+            logger.warning(
+                'Plugin %s registered two analysis providers for %r; keeping the last one',
+                self.plugin_id, component,
+            )
+        self.analysis_providers[component] = {'factory': factory, 'cache': bool(cache)}