Bagpiper on the command line and the MCP server - #6810

Codecov Report

❌ Patch coverage is 87.71930% with 42 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.57%. Comparing base (158e495) to head (eeb8b18).

@@            Coverage Diff             @@
##           master    #6810      +/-   ##
==========================================
+ Coverage   73.48%   73.57%   +0.08%     
==========================================
  Files         860      861       +1     
  Lines       80440    80768     +328     
==========================================
+ Hits        59112    59422     +310     
- Misses      21328    21346      +18     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

Merge Risk: 🟡 Moderate · up to 769bd

Fix the checkpoint and served-TTS prompt behavior before merging. The demo also needs a confirmed cleanup policy for rendered files; the incorrect loader hint should be corrected.

Security Architecture Review

Security architecture risk: 🟠 High · up to 769bd

The new integrations can access files using the host process's permissions and forward file contents to a configured service. The safeguards around who can invoke them and where that service runs are not established.

Retained concerns

• High · security · observed: The newly introduced MCP output operations accept caller-selected host paths. In particular, render creates parent directories and writes generated bytes directly to the selected file, with no path boundary, atomic replacement, or recovery after a failed write. A caller able to invoke the tool can overwrite files writable by the MCP process; interruption or concurrent use can also leave a partial result.
• Medium · security · inferred: MCP describe permits a caller to name any file readable by the process; the served inference path reads its raw bytes and forwards them to the configured model endpoint. The file-existence check does not confine inputs to approved recordings. Exposure outside the host depends on the deployment-selected endpoint, which is not established.

• Fix all pre-merge checks with AI

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

_{Comment @coderabbitai help to get the list of available commands.}

Both addressed in 66cde2d — thank you, they were real.

Leaked WAVs. The page now hands gradio the samples (type="numpy") instead of writing a file it then owns, so there is nothing of ours left behind, and gr.Blocks(delete_cache=(600, 3600)) bounds what gradio keeps of the copies it serves.

The served TTS prompt. You are right that the name says nothing: the documented docker run serves every checkpoint as bagpiper, so a mounted bagpiper-tts-sft would have been sent the general model's constant system turn, which its training data does not have. ESPNET_BAGPIPER_TTS_SYSTEM now lets the deployment state it (empty for the TTS checkpoint), the name is only the fallback and the docstring says so, and the Space README says to set it. Covered by a test.

Thank you — taking the three in order.

Release logic. Agreed, and it is where you want it: RELEASES is a dict of three keys with two fields each (what, system), read in exactly two places — from_pretrained picks the generation system turn and refuses the base checkpoint, and download_pretrained uses it to name the right loader. There is no BagpiperRelease, no router, no per-release class, and I will not add one. If a fourth SpeechLM release appears it is one more row, and if it needs behaviour rather than metadata that is the moment to stop and design, not now.

ESPNET_BAGPIPER_TTS_SYSTEM. Fair question. The parameter is the API — from_server(..., tts_system=...) wins over everything, and the env var is only its default, read on one line. It exists because the thing that knows is the deployment and the thing that asks is often not a Python caller: the MCP server and the Space are configured by environment (ESPNET_MCP_*, BAGPIPER_URL) and have no argument to pass, and a person running espnet describe --server ... against a TTS deployment has no flag either. One variable covers all three in one place; moving it into the front ends would mean two reads and still leave the CLI without a way. I am happy to drop it and go explicit-only if you would rather — the cost is that the CLI case becomes unfixable — but I would keep it.

MCP file paths. Agreed that this should not become a Bagpiper-specific sandbox. Worth saying plainly that the PR does not widen the contract: reads go through the same _existing_file every audio tool uses, and render writes exactly as synthesize and enhance already do on master (out.parent.mkdir(parents=True, exist_ok=True), then write). So it is the MCP server's policy, unchanged, and it belongs in a separate issue about the server as a whole. Say the word and I will open one.

Tests and coverage. Pushed a small addition: the uncovered lines were not all the 8B path — which release directory, which YAML, which .pt, and the messages when a directory holds more than one, are plain filesystem work. The module is at 85% now, and the remainder is from_pretrained's body, which needs 99 GB of weights and a GPU. The release-specific behaviours you list are each covered directly: bagpiper-sft sends the constant system turn, bagpiper-tts-sft sends none, the base is refused without allow_base=True.

On the docstring-coverage warning: I have fixed what CI enforces (H403/H405) and left the rest.

@jctian98 — asking for your eye on this, since it stands on espnet2/speechlm.

espnet2/bin/speechlm_inference.py is meant to be the single-sample sibling of espnet2/speechlm/bin/inference.py: same job template, same checkpoint, same decoding configs, one dialogue at a time and no manifests. What it does per request is

dialogue = [(role, modality, content), ...]          # system / user turns
batch = preprocessor.collate_fn([(("dialogue", "local", key), {"dialogue": dialogue})])
batch = to_device(batch, device, dtype); batch.pop("keys")
messages, _ = model.inference(config, **batch)

Four things are yours to confirm, and I could not test any of them without a GPU — I have a Bridges-2 H100 job queued to do exactly that:

1. Is that the contract you intended for calling model.inference outside the batch script, in particular popping only keys and leaving data_stats in the kwargs, as inference.py does?
2. The served TTS system turn. The docker run in egs2/bagpiper/speechlm1/README.md serves every checkpoint as bagpiper, so the name cannot say which one is mounted. Bagpiper-TTS was trained with no system turn at all (arXiv:2606.22811 §2.3), so the wrapper needs to be told: ESPNET_BAGPIPER_TTS_SYSTEM= (empty) for that deployment. Does that match how you actually serve it?
3. attn_implementation. The published configs ask for flash_attention_3 in two places, which is Hopper-only, so from_pretrained(..., attn_implementation="sdpa") sets both. Is sdpa an acceptable fallback for inference, or is there a reason the encoder in particular wants FA?
4. Prompting. describe sends [system, audio, text] and render sends [system?, text], with the SFT constant as the generation system turn. If either shape is off distribution, that is the thing most worth catching here.

One measured note for context, since the PR explains why there is no hosted Space: I read the shard headers of Qwen/Qwen3-Omni-30B-A3B-Instruct on Bridges-2 — thinker.audio_tower is 1.21 GiB of the 65.68 GiB the loader instantiates before deleting the rest.

@jctian98 a gentle ping on this one — it is green and mergeable, and waiting only on your review.

If you are short of time, two of the four questions are the ones that actually need you; the rest I can settle myself:

1. The served TTS system turn. The docker run in egs2/bagpiper/speechlm1/README.md serves every checkpoint under the name bagpiper, so the wrapper cannot tell which one is mounted. Bagpiper-TTS was trained with no system turn (arXiv:2606.22811 §2.3), so a deployment serving it sets ESPNET_BAGPIPER_TTS_SYSTEM= (empty). Does that match how you serve it in practice?
2. attn_implementation="sdpa". The published configs ask for FlashAttention-3 in two places, which is Hopper-only, so this sets both. Is sdpa an acceptable fallback for inference, or does the audio encoder in particular want FA?

The other two — whether popping only keys from the collated batch is the contract you intended for calling model.inference outside the batch script, and whether the [system, audio, text] / [system?, text] prompt shapes are in distribution — are the ones I would most like a second pair of eyes on, but a "looks right" is enough.

Still no GPU run of my own to report: the Bridges-2 H100 queue has been 100+ deep. Everything in the PR is verified against stand-in servers and a stand-in model, which is why I am asking rather than asserting.

Hi @sw005320
(1) Yes, the TTS model doesn't need a system prompt, which was intentional. We don't want the users to manually select its work mode.
(2) The modeling file will enforce flash attention, which is because the SDPA cannot solve the packed sequence properly and will cause training data leakage across the examples. This is only applied to training, not inference. Using SDPA for inference is always ok.

Specifically, turn this assertion to train-only
espnet2/speechlm/model/speechlm/lm/parallel.py:95
espnet2/speechlm/model/speechlm/lm/parallel_pp.py:116

On the prompting question, I checked the released Bagpiper SFT Data. The two directions have different message formats:

• General espnet/bagpiper-sft generation examples are [system, user text, assistant text, assistant audio]. Their system text matches this PR's TTS_SYSTEM verbatim, so render's [system, text] prefix is in distribution.
• Open-ended understanding examples use [system, user audio, user text, assistant text]. MMAU uses that shape with a multiple-choice-specific system prompt.

This is a check of the released data's message format; I have not compared response quality with the general 8B checkpoint.

Follow-up from testing the released espnet/bagpiper-tts-sft checkpoint locally. To be precise, this was a CPU + SDPA run; I have not run this PR on a GPU. I temporarily bypassed the two FlashAttention assertions already discussed to reach inference. I found two problems in the underlying SpeechLM inference path exercised by this PR:

1. In espnet2/speechlm/model/speechlm/lm/parallel.py, the unused-modality dummy_forward() loop runs during eval too. A text-only render() request therefore runs the audio encoder. Adding 0 * a finite dummy result does not change the prediction; the problem is the unnecessary encoder call (and the failure it can trigger). This graph-inclusion step is for training and should be skipped in eval.
2. In espnet2/speechlm/model/speechlm/multimodal_io/audio.py, ContinuousAudioIO.encode_batch() calls .split() directly on get_audio_features()'s result. Qwen3 Omni's get_audio_features() returns BaseModelOutputWithPooling under the published Transformers 5.5.4 and installed 5.14.1, so .split() raises instead of producing per-sample audio features. Skipping the dummy path fixes text-only generation, but real audio input still needs this return-type handling fixed before describe() can work.

With the eval dummy path skipped, the local TTS checkpoint produced a 3.76-second, 16 kHz WAV from a user-only request; Whisper tiny transcribed it as “Hello, this is a test.” This validates that CPU text-to-audio path only. It is not a GPU or audio-input end-to-end result.

Thank you — that is a far more useful review than a rubber stamp, and running it on CPU to find the two real bugs was above and beyond. All three code points are now in this PR (f6d1ac7), since it is the path they block.

The Flash Attention assertion is now training's. from_pretrained takes is_train, the job template passes self.is_train, and the message says "Training OpusLM requires Flash Attention". Both sites, parallel.py and parallel_pp.py. I left the default at True on purpose: a training path that forgets the argument should fail loudly rather than leak across packed examples quietly, and the espnet path always passes it. Say the word if you would rather it defaulted the other way.

dummy_forward() is now skipped in eval. The loop is wrapped in if self.training:, with the reason written beside it — the graph it keeps whole is DeepSpeed ZeRO's concern, and in eval it bought a text-only request an audio-encoder call that can fail.

encode_batch() unwraps the model output. get_audio_features() returns BaseModelOutputWithPooling on both 5.5.4 and 5.14.1, so .split() raised and audio input never reached the model. It now takes last_hidden_state when what arrived is not a tensor, and two tests cover both return types. 44 tests pass in an environment with the speechlm stack installed.

On the prompting shapes, thank you for checking the released data — that is exactly the confirmation I could not get myself. It is good to know the general model's generation system turn matches TTS_SYSTEM verbatim, and that understanding is [system, user audio, user text, assistant text], which is the shape describe() sends. I have noted in the code that MMAU uses its own multiple-choice system prompt; a caller can pass system= for that.

And on the TTS checkpoint having no system turn, that is what ESPNET_BAGPIPER_TTS_SYSTEM= (empty) is for, since the documented docker run serves every checkpoint under the name bagpiper and the wrapper cannot tell which is mounted. Does that match how you would want a deployment to declare it?

I still have no GPU run of my own to report — the Bridges-2 H100 queue has been 100+ deep for days — so your CPU result is the only end-to-end evidence this PR has. Worth saying plainly in case anyone reads the thread later.

This pull request is now in conflict :(