Skip to content

Plugins: Allow exchanging analysis providers and inference providers alike - #782

Merged
NeptuneHub merged 13 commits into
NeptuneHub:mainfrom
Schaka:main
Jul 29, 2026
Merged

Plugins: Allow exchanging analysis providers and inference providers alike#782
NeptuneHub merged 13 commits into
NeptuneHub:mainfrom
Schaka:main

Conversation

@Schaka

@Schaka Schaka commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

This solves issue #755

It was not feasible to expand the plugin installer scripts to be able to install system dependencies into the image.
Even if it was, it would create a messy image and probably take absolutely forever, because we also need to compile dependencies and set up PATH anyway.

Once you merge these changes, I will be able to provide my own plugin to enable MiGraphX AMD GPU (plugin already written and tested against these changes). I have a base image ready in the form of https://github.com/Schaka/rocm-migraphx-ort-builder and will instruct users to either use my image (containing your Audiomuse-AI code) to provide system dependencies for MiGraphX to work, or I will provide a more clear cut base image and create another PR to your CI that will build a tag for AMD, similar to Nvidia (just BASE_IMAGE swapped).

Let me know if you need any more changes made to this.
Looking forward to having this upstream so I can publish the plugin.

Edit: I kept whisper as being named whisper. Let me know if you prefer renaming to ASR entirely

Adds two provider-agnostic seams so an accelerator (GPU, NPU, ...) can be
shipped as an out-of-tree plugin instead of forking core:

* register_onnx_provider() gains only_models/exclude_models so a plugin can
  offer an ONNX execution provider for just the graphs it can compile, scoped
  by the session label passed to create_onnx_session (musicnn, clap,
  whisper_encoder, ...). resolve_providers() honors the scope and leaves
  unmatched models on their default chain.
* register_analysis_provider(component, factory) is wired end to end (it was a
  logging-only stub). A plugin can replace a whole analysis component; core
  consults the registered provider first and falls back to the built-in. The
  lyrics pipeline resolves its ASR/Whisper backend through lyrics/_asr_backend
  so the built-in ONNX Whisper can be swapped without touching call sites.

clap_analyzer pins the CLAP audio model's symbolic time axis to a static shape
for any provider whose graph compiler needs it, generalizing the existing
CoreML-only path.
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@NeptuneHub

Copy link
Copy Markdown
Owner

Which test did you run on this PR, can you share more details ?

My expectation is that on container image (normal image, Nvidia, novx) and native image (Windows, Linux and Macos) on both Intel and Arm without plugin installed this is a NO-OP.

With plugin installed, the plugin should be able to:

  • Receive the list of Model using during the analysis
  • Eventually add other model (this should be already in palce with the analysis hooc)
  • Add an additional (any) execution provider PER MODEL (there is some model, for example, that work good only on CPU)
  • Install off course all the dependencies of that specific provider if needed

Can you share more about the above points? Thanks.

@Schaka

Schaka commented Jul 25, 2026

Copy link
Copy Markdown
Contributor Author

I have run the changed code locally with and without the plugin. Without the plugin, no additional ExecutionProvider is added. The only real change is the asr_backend, so that for whisper we can switch out the model. Since just switching the model itself is not enough, I made it interface-like where any plugin can supply its own implementation (in my case, I will use faster-whisper, since default whisper will not run on MiGraphX).

without plugin installed this is a NO-OP

This is correct. Without the plugin, nothing changes. Default is always CPUExecutionProvider or CUDAExcecutionProvider if Nvidia image is used (this I cannot test, I don't have Nvidia GPUs, sorry).

Here's an example log file of analysis running purely on CPU:
analysis.log

Receive the list of Model using during the analysis

This is currently not in. Do you want me to implement that? The mechanism is currently inverted, the core calls resolve_providers(label=X) and the plugin, when registering providers can decide (via this label) with inclusion/exclusion list for which models the new provider counts. So I can add MiGraphExecutionProvider for musicnn but not for CLAP or the other way around. Hopefully this answers your question number 3 too.

Install off course all the dependencies of that specific provider if needed

Yes and no. I did not change this mechanism, so just like you implemented before, you can provide a .json to install additional dependencies via PyPi. What is not possible is to install system dependencies. This would be insecure imo.

This is a similar problem to Nvidia. You need system dependencies and a different base image to get everything working perfectly. I have CI running to publish base images for all supported AMD GPUs here: https://github.com/Schaka/rocm-migraphx-ort-builder

One last point:

Eventually add other model (this should be already in palce with the analysis hooc)

Currently, I only did this for whisper. It only made sense for me to do it for whisper, because there are so many speech to text models that all work completely different. For something like CLAP or even musicnn, I don't think there is any alternative so I'm not sure we could benefit from switching out the model.

Does this all make sense to you? Let me know if you need anything more from me.

@NeptuneHub NeptuneHub left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remember that this PR MUST allow to everyone to use different provider for any of the actual model by a plugin. So the core idea should be change this to allow multiple use, not only this specific one.

Unit test is failing, please check:

=========================== short test summary info ============================
FAILED test/unit/test_plugin_analysis_seams.py::TestGetAsrBackend::test_falls_back_to_builtin_when_no_override - AssertionError: assert <module 'lyrics.whisper_onnx' from '/home/runner/work/AudioMuse-AI/AudioMuse-AI/lyrics/whisper_onnx.py'> is <module 'lyrics.whisper_onnx'>
 +  where <module 'lyrics.whisper_onnx' from '/home/runner/work/AudioMuse-AI/AudioMuse-AI/lyrics/whisper_onnx.py'> = <function get_asr_backend at 0x7febd5fc7a60>()
 +    where <function get_asr_backend at 0x7febd5fc7a60> = <module 'lyrics._asr_backend' from '/home/runner/work/AudioMuse-AI/AudioMuse-AI/lyrics/_asr_backend.py'>.get_asr_backend
FAILED test/unit/test_plugin_analysis_seams.py::TestGetAsrBackend::test_manager_error_falls_back_to_builtin - AssertionError: assert <module 'lyrics.whisper_onnx' from '/home/runner/work/AudioMuse-AI/AudioMuse-AI/lyrics/whisper_onnx.py'> is <module 'lyrics.whisper_onnx'>
 +  where <module 'lyrics.whisper_onnx' from '/home/runner/work/AudioMuse-AI/AudioMuse-AI/lyrics/whisper_onnx.py'> = <function get_asr_backend at 0x7febd5fc7a60>()
 +    where <function get_asr_backend at 0x7febd5fc7a60> = <module 'lyrics._asr_backend' from '/home/runner/work/AudioMuse-AI/AudioMuse-AI/lyrics/_asr_backend.py'>.get_asr_backend
====== 2 failed, 2667 passed, 1 skipped, 22 warnings in 154.71s (0:02:34) ======

Code Review point of attention

1. The per-track fallback never gets a label

In tasks/analysis/song.py the MusiCNN models can be loaded in two different ways.

The normal way: at the start of every album the worker loads the two models one time and reuses them for all the tracks of that album. This is ensure_musicnn_sessions calling load_musicnn_sessions, line 244. Your PR gives this one label='musicnn', so here it works.

The fallback way: if that album load fails, load_musicnn_sessions catches the error and returns None. Then analyze_track receives onnx_sessions=None, and _sessions_for_track (line 428) builds a new session for every single track. This one calls resolve_providers() with no argument, so label stays None, and it never gets a label.

With label=None your new filter drops the plugin provider:

if only and label not in only:   # None not in ['musicnn'] -> skip
    continue

So a plugin with only_models=['musicnn'] uses the GPU in the normal way, and goes back to CPU in the fallback way, without any message in the log. What do you think about adding label='musicnn' on line 428 too?

2. Today the ASR model is loaded one time per album, not one time per song

How it works now: the built-in whisper is a module singleton. It is loaded on the first track that needs lyrics, kept in memory for all the tracks of the album, and freed at the end of the album by cleanup_optional_models.

With the new seam this is not guaranteed anymore. get_asr_backend() has no cache, and the manager runs factory() on every call (once per song in _transcribe, and again in is_lyrics_loaded and unload_lyrics_models). So if a plugin gives a callable that builds a new object, the model is loaded again for every song, and the old one is never freed, because is_loaded() on a new object returns False.

What do you think about keeping the same behaviour as today, so the backend is resolved one time and cached? Or, if you prefer, give the choice to the plugin developer, for example a flag on register_analysis_provider, so the plugin can decide if it wants a cached backend or a new one every time.

3. Only 4 labels really work, and the docs miss one

The labels that really arrive to your new filter are only musicnn, clap, whisper_encoder and whisper_decoder. The docs list the first three and miss whisper_decoder, which is exactly the graph you say MIGraphX cannot compile. So someone who writes exclude_models=['whisper_encoder'] still gets the provider on the decoder, and there is no warning.

The other labels (embedding, prediction, gte, silero_vad) never arrive to the filter, because those sessions already receive a ready provider list, so scoping them does nothing.

What do you think about writing the 4 real labels in the docs, and logging a warning when a plugin uses a label that does not exist?

4. Other things, can you check these too?

  • _PREPARED_MODEL_PROVIDERS in tasks/clap_analyzer.py has MIGraphXExecutionProvider written inside core. The idea of the PR is that a new accelerator does not need a core change, but the next provider that needs static shapes will need one. What do you think about a flag at registration time, for example needs_static_shapes=True, so core does not need to know the provider names?
  • only_models='musicnn' written as a plain string becomes ['m','u','s','i','c','n','n'], so it matches nothing and there is no error. Can you accept a string and convert it to a list, like _plugin_targets already does for targets?
  • Nothing checks that a plugin backend really has the 4 methods. If is_loaded is missing, lyrics/__init__.py catches the error and is_lyrics_loaded() returns True forever, so we run a full memory cleanup after every album for nothing. And if transcribe is broken, every song fails and we never go back to the built-in whisper. Can we check the 4 methods one time when we resolve, log a warning and use the built-in if something is missing?
  • register_analysis_provider accepts any component name. If a plugin writes 'ASR' or 'musicnn', the value is stored, the plugin says ok, and nothing happens, with no log. The old stub at least logged a message. Same thing if two plugins register 'asr': the winner depends on dict order, with no warning.
  • position is accepted and it is also in the doc row, but it is never read, so a plugin provider is always added after CUDA. This is true also before your PR, but since you are touching this API, maybe it is a good moment to honour it or to remove it.

5. Half of the models have no label, so a plugin can never reach them

Only 4 sessions ask resolve_providers, so only 4 models can be scoped today. These other sessions pass their own provider list, so they never ask resolve_providers and a plugin provider can never arrive there:

  • gte, the lyrics text embedding, lyrics/gte_onnx.py:103, CPU only
  • silero_vad, lyrics/silero_onnx.py:73, CPU only
  • the CLAP text model, tasks/clap_analyzer.py:209, it builds its own CUDA or CPU list
  • MusiCNN in the per-track path, tasks/analysis/song.py:428, this is my point 1

The reason is that create_onnx_session does provider_options or resolve_providers(...), so when the call site already passes a list, resolve_providers is not called and the label is used only for the log line.

So today the feature works only for the models your own plugin needs. For example someone with OpenVINO would like to say "use it for MusiCNN and for gte, but not for CLAP and not for Whisper", and the gte part is simply not possible.

What do you think about making this general for all the models? An idea: keep CPU as the default for these sessions, but let the plugin opt-in, for example resolve_providers(label='gte', cpu_only_default=True). In this way nothing changes for the normal user, and a plugin can use only_models=['gte'] if it wants. For the CLAP text model please keep the Flask process on CPU, because that one is for thread safety, and open it only on the worker.

@Schaka

Schaka commented Jul 26, 2026

Copy link
Copy Markdown
Contributor Author

Thank you for the extensive review.

Unit test is failing, please check

They were running locally. As far as I understand, it had something to do with my local environment missing a dependency. This came down to my lack of experience with Python.

I initially tried to be minimally invasive with my changes, but I agree the that we should cover all models and cache instantiation of the asr backend. Good catch, honestly. I didn't really think about that because I come from other languages and framework.

I now gated the asr backend and validate that methods exist. Do you think this is enough or is there a way to validate that methods also take the correct parameter maybe?

Hopefully I covered everything now (in a proper Python way)

@Schaka
Schaka requested a review from NeptuneHub July 26, 2026 14:12

@NeptuneHub NeptuneHub left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi,
thanks for your additional update, there is still some small update that I need (the first is probably the most important one):

1. Document what transcribe() must return

Background: after a song is transcribed, core reads a confidence value (avg_logprob) from the result to decide if the transcript is good enough. If the confidence is too low, the transcript is thrown away and the song is treated as instrumental.

The problem: the docs only list the four methods a plugin must provide, they never say what transcribe() must return. If a plugin returns only text and language (a very natural thing to do) and forgets avg_logprob, core assumes the worst confidence and drops every transcript. No crash, no error, the plugin looks like it works, but the user gets zero lyrics for the whole library. This is exactly the trap the faster-whisper plugin this PR is made for could fall into.

Fix asked: document the required return keys (text, language, avg_logprob). If possible, also skip the quality check when avg_logprob is missing instead of assuming the worst, so this failure becomes impossible.

2. Make plugin failures visible in the log

Background: when a plugin cannot provide its component, core falls back to the built-in behavior. This is the right choice, but the user must be able to see it happened.

The problem: two cases are hard to diagnose today. When a plugin factory fails, the error message says which component failed but not which plugin, so the user cannot tell which plugin to fix or remove. And when a factory returns None, nothing is logged at all: the user installs the plugin, core silently keeps using the built-in Whisper, and there is no trace anywhere of why.

Fix asked: add the plugin id to the factory error message, and log a message when a factory returns None.

3. Small doc additions

Background: these are the first questions every plugin author will ask.

The problem: three things are true in the code but not written anywhere: the execution provider must already be present in the onnxruntime build of the image (a plugin cannot add one via pip into its own _lib); needs_static_shapes only affects the CLAP audio model; and if a plugin registers a provider core already manages (like CUDA), core wins and the plugin entry is dropped.

Fix asked: one sentence for each in docs/PLUGIN.md.

4. Drop the .gitignore change

Background: the PR touches .gitignore.

The problem: the added entries look like local development stuff not related to this PR (.claude/ is already ignored, .tokensave is not used anywhere, and .venv-test is added without a final newline).

Fix asked: remove this hunk from the PR.

Thanks again!

@Schaka

Schaka commented Jul 28, 2026

Copy link
Copy Markdown
Contributor Author

Again, thank you so much for the expansive review.
I had kept the .gitignore in because it helped me hide a bunch of stuff I use during local development (including venv to run tests locally). .tokensave is just part of headroom. But I added it all to my local ignore list now.

Hopefully the documentation is enough now. I asked Claude to summarize my changes, so it's not initially written by me, but I went over to correct/shorten things where possible. I'm not great at writing docs, so hopefully it covers everything you'd want it to now.

Document what transcribe() must return

This was a really good catch. In my comment I only wondered about whether we should check method signatures for transcribe() etc, but I never thought to also check the return type. I've worked in statically typed languages for so long.

Let me know if you need anything more and sorry for taking so long to cover all the bases.

@Schaka
Schaka requested a review from NeptuneHub July 28, 2026 16:58
…fidence, name the owning plugin in fallback warnings, fix PLUGIN.md static-shape claim
…arn on impossible confidence, None-init and stale-owner cleanup
@sonarqubecloud

Copy link
Copy Markdown

@NeptuneHub

NeptuneHub commented Jul 29, 2026

Copy link
Copy Markdown
Owner

Hi,
I did a last review and I decided directly to commit the fix against the branch to save time, this is the recap of the change done:

Commit 1 (58a0b83)

  • If a plugin's speech-to-text backend reports a broken confidence value (NaN or a huge number), the transcript is no longer thrown away by mistake: the quality check is simply skipped, with a warning in the log.-
  • When a plugin fails or refuses to provide its backend, the log now tells you which plugin it was, so you know what to fix or remove.
  • Fixed a wrong sentence in the plugin docs about which models get their shapes pinned.

Commit 2 (822d847)

  • When two or more plugins offer the same component, the log now says honestly "trying them in this order" and then names the one actually in use, instead of guessing the winner in advance.
  • If a backend reports a confidence of zero or higher (which is impossible for real data), it is kept but flagged in the log, because such a value would silently disable the quality check.
  • Confidence of minus infinity keeps its old meaning ("worst possible, drop it") and this is now locked in by tests.
  • Small internal cleanups: the "who provides the backend" record is cleared when a plugin stops providing it, and the starting value of the confidence is now "unknown" instead of a fake zero.

I personally did some test with no plugins installed, and for a normal users nothing should change, so I'm going to merge this PR.

PLEASE:

Your plugin will become "the case study" for future developer, so really thanks for your effort.

Important: Support for plugin that add execution provider will be released in the next release, that will be the v3.1.0, so be sure that your plugin have that as minimum AudioMuse-AI version. thanks.

@NeptuneHub
NeptuneHub merged commit 798072e into NeptuneHub:main Jul 29, 2026
25 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants