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,8 @@ 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)` | Advanced: offer an extra ONNX Runtime execution provider (for example a GPU) for analysis. Scope it to specific models with `only_models`/`exclude_models` (lists of the session `label`: `musicnn`, `clap`, `whisper_encoder`, ...) when the provider can only parse some graphs. |
+| `register_analysis_provider(component, factory)` | 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. |
 
 ## 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,42 @@
+# 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
+* Swallows any plugin-resolution error so a broken plugin never breaks lyrics
+"""
+
+from __future__ import annotations
+
+
+def get_asr_backend():
+    override = _plugin_asr_backend()
+    if override is not None:
+        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

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]:

plugin/api.py

@@ -200,6 +200,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 = []
@@ -225,8 +226,15 @@ 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):
+        self.onnx_providers.append({
+            'name': name,
+            'options': options or {},
+            'position': position,
+            'only_models': list(only_models) if only_models else None,
+            'exclude_models': list(exclude_models) if exclude_models else None,
+        })
 
     def on_flask_start(self, func):
         self.flask_start.append(func)
@@ -249,9 +257,18 @@ 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):
+        """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 (currently ``asr``). ``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.
+        """
+        self.analysis_providers[component] = factory

plugin/manager.py

@@ -324,6 +324,7 @@ def sync(self, conn=None, role=None):
                 record['settings_endpoint'] = None
                 record['cron_tasks'] = {}
                 record['onnx_providers'] = []
+                record['analysis_providers'] = {}
                 record['song_analyzed_hooks'] = []
                 record['error'] = None
                 records[row['id']] = record
@@ -600,6 +601,7 @@ def load(self, role, flask_app=None):
                     record['settings_endpoint'] = ctx.settings_endpoint
                     record['cron_tasks'] = {**(ctx.tasks or {}), **(ctx.cron_tasks or {})}
                     record['onnx_providers'] = ctx.onnx_providers
+                    record['analysis_providers'] = ctx.analysis_providers
                     record['song_analyzed_hooks'] = ctx.song_analyzed_hooks
                     if role == 'web' and flask_app is not None and ctx.blueprint is not None:
                         if ctx.blueprint.name != plugin_id:
@@ -854,6 +856,25 @@ def get_onnx_providers(self):
                 providers.extend(record.get('onnx_providers', []))
         return providers
 
+    def get_analysis_provider(self, component):
+        """Return the first loaded plugin's replacement for ``component``, or None.
+
+        Resolves ``factory`` to the actual implementation (calling it when it is a
+        zero-arg callable). Used by core to let a plugin swap out a whole analysis
+        step such as the ASR/Whisper backend.
+        """
+        for record in self.records.values():
+            if record.get('load_status') not in _LOADED_STATUSES:
+                continue
+            factory = record.get('analysis_providers', {}).get(component)
+            if factory is None:
+                continue
+            try:
+                return factory() if callable(factory) else factory
+            except Exception:
+                logger.exception('Plugin analysis provider for %r failed to resolve', component)
+        return None
+
     def song_analyzed_hooks(self):
         hooks = []
         for record in self.records.values():

tasks/analysis/song.py

@@ -104,7 +104,7 @@ def sigmoid(x):
     return 1 / (1 + np.exp(-x))
 
 
-def resolve_providers(allow_coreml=False, cuda_options=None):
+def resolve_providers(allow_coreml=False, cuda_options=None, label=None):
     available = ort.get_available_providers()
     chain = []
 
@@ -135,8 +135,17 @@ def resolve_providers(allow_coreml=False, cuda_options=None):
 
     for provider in _plugin_onnx_providers():
         name = provider.get('name')
-        if name and name in available and name not in [p[0] for p in chain]:
-            chain.append((name, provider.get('options') or {}))
+        if not name or name not in available or name in [p[0] for p in chain]:
+            continue
+        # Providers can be scoped to specific models by their session label, so a
+        # plugin can offer an accelerator for the graphs it handles
+        only = provider.get('only_models')
+        exclude = provider.get('exclude_models')
+        if only and label not in only:
+            continue
+        if exclude and label in exclude:
+            continue
+        chain.append((name, provider.get('options') or {}))
 
     chain.append(('CPUExecutionProvider', {}))
     logger.info("ONNX provider chain: %s", [p[0] for p in chain])
@@ -209,8 +218,10 @@ def _default_sess_options():
     return opts
 
 
-def create_onnx_session(model_path, provider_options=None, label="", sess_options=None):
-    opts = provider_options or resolve_providers()
+def create_onnx_session(
+    model_path, provider_options=None, label="", sess_options=None, allow_coreml=False
+):
+    opts = provider_options or resolve_providers(allow_coreml=allow_coreml, label=label)
     if sess_options is None:
         sess_options = _default_sess_options()
     try:
@@ -230,7 +241,7 @@ def create_onnx_session(model_path, provider_options=None, label="", sess_option
 
 
 def load_musicnn_sessions(model_paths):
-    opts = resolve_providers(allow_coreml=False)
+    opts = resolve_providers(allow_coreml=False, label='musicnn')
     try:
         sessions = {n: create_onnx_session(p, opts, label=n) for n, p in model_paths.items()}
         logger.info(f"OK Loaded {len(sessions)} MusiCNN models for album reuse")

tasks/clap_analyzer.py

@@ -49,8 +49,15 @@
 
 _SEGMENT_LENGTH_SAMPLES = 480000
 
+# Providers whose graph compiler needs the CLAP audio model's symbolic
+# time-frame axis pinned to a fixed value before it can compile the graph
+# (neither compiles a dynamic dim).
+_PREPARED_MODEL_PROVIDERS = {'CoreMLExecutionProvider', 'MIGraphXExecutionProvider'}
 
-def _static_shape_model_bytes(model_path):
+
+def _prepared_model_bytes(model_path):
+    """Return the audio model with its symbolic time axis pinned to a static
+    shape, or None when the model has no symbolic dims. See _PREPARED_MODEL_PROVIDERS."""
     import onnx
     from onnxruntime.tools.onnx_model_utils import make_dim_param_fixed
 
@@ -115,6 +122,7 @@ def _load_audio_model():
             'arena_extend_strategy': 'kSameAsRequested',
             'cudnn_conv_algo_search': 'DEFAULT',
         },
+        label='clap',
     )
 
     def _create_session(model_input, providers, provider_opts):
@@ -131,14 +139,14 @@ def _create_session(model_input, providers, provider_opts):
     cpu_opts = [{}]
 
     preferred_model_input = model_path
-    if 'CoreMLExecutionProvider' in preferred_providers:
+    if _PREPARED_MODEL_PROVIDERS.intersection(preferred_providers):
         try:
-            static_bytes = _static_shape_model_bytes(model_path)
-            if static_bytes is not None:
-                preferred_model_input = static_bytes
-                logger.info("CoreML: pinned CLAP audio time axis to a static shape for compilation")
+            prepared_bytes = _prepared_model_bytes(model_path)
+            if prepared_bytes is not None:
+                preferred_model_input = prepared_bytes
+                logger.info("Pinned CLAP audio time axis to a static shape for graph compilation")
         except Exception as e:
-            logger.warning(f"CoreML: could not build static-shape model ({e}); using dynamic model")
+            logger.warning(f"Could not build static-shape model ({e}); using dynamic model")
 
     try:
         session = _create_session(preferred_model_input, preferred_providers, preferred_opts)