[Model] Add YuE2-3B text-to-music (model + /v1/audio/speech adapter) (#7886)

Signed-off-by: Youyuan Li <105433145+Darcy-Lee@users.noreply.github.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Co-authored-by: princepride <wangzhipeng628@gmail.com>

docs/models/supported_models.md

@@ -83,6 +83,7 @@ th {
 | `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
+| `Yue2ForCausalLM` | YuE2-3B (text-to-music with ABC plan) | `m-a-p/YuE2-3B` | ✅︎ | | | | — |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |

examples/offline_inference/yue2/end2end.py

@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""YuE2-3B offline end-to-end text-to-music.
+
+One song is one or two engine requests against the single-stage pipeline:
+``cot=full|melody`` first asks the model to write an ABC score (the abc phase,
+tokens only), then the semantic phase generates codec tokens and the model
+solves the acoustic ODE + decodes 48 kHz stereo as the request finishes.
+``cot=off`` and a user-supplied ``--abc-file`` skip the first request.
+Prefix caching lets the second request reuse the first one's blocks.
+
+Example:
+    python end2end.py --model /models/YuE2-3B --vae /models/YuE2-Vae \
+        --style "gentle lo-fi hip-hop, 85 BPM, F minor" \
+        --lyrics "[Verse]\nHumming on a quiet street\n[Chorus]\nStay a while" \
+        --cot full --seed 831001 --max-frames 200 --output song.wav
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+from pathlib import Path
+
+import soundfile as sf
+import torch
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.yue2.yue2 import (
+    ABC_SAMPLING,
+    KEY_MAX_AUDIO_FRAMES,
+    KEY_MIN_TOKENS,
+    KEY_PENALTY_WINDOW,
+    KEY_PHASE,
+    KEY_PREFIX_IDS,
+    KEY_REPETITION_PENALTY,
+    KEY_SEED,
+    KEY_SKIP_SYNTHESIS,
+    KEY_TEMPERATURE,
+    KEY_TOP_K,
+    KEY_TOP_P,
+    SEMANTIC_SAMPLING,
+    STOP_TOKEN_IDS,
+)
+from vllm_omni.tokenizers.yue2_prompt import (
+    abc_ids_from_generated,
+    abc_prefix_ids,
+    semantic_prefix_ids,
+)
+from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer
+
+
+def sampling_params(engine, *, phase, seed, max_frames, prompt_ids):
+    """Deep-copy the stage defaults, then set the yue2_* keys for one phase."""
+    params = copy.deepcopy(engine.resolve_sampling_params_list(None))[0]
+    preset = ABC_SAMPLING if phase == "abc" else SEMANTIC_SAMPLING
+    params.extra_args = {
+        **(params.extra_args or {}),
+        KEY_PHASE: phase,
+        KEY_SEED: seed,
+        KEY_TEMPERATURE: preset["temperature"],
+        KEY_TOP_P: preset["top_p"],
+        KEY_TOP_K: preset["top_k"],
+        KEY_REPETITION_PENALTY: preset["repetition_penalty"],
+        KEY_PENALTY_WINDOW: preset["penalty_window"],
+        KEY_MIN_TOKENS: preset["min_tokens"],
+        KEY_MAX_AUDIO_FRAMES: max_frames,
+        KEY_SKIP_SYNTHESIS: phase == "abc",
+        # Full prompt ids: under a KV prefix-cache hit the engine schedules
+        # only the uncached tail, so the model cannot rebuild its NAR
+        # conditioning prefix from the scheduled tokens alone.
+        KEY_PREFIX_IDS: list(prompt_ids),
+    }
+    params.max_tokens = preset["max_tokens"] if phase == "abc" else max_frames + 1
+    params.stop_token_ids = list(STOP_TOKEN_IDS)
+    params.detokenize = False
+    return params
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", required=True, help="m-a-p/YuE2-3B checkpoint dir or repo id")
+    parser.add_argument("--vae", default=None, help="m-a-p/YuE2-Vae dir or repo id (or $YUE2_VAE)")
+    parser.add_argument("--style", required=True)
+    parser.add_argument("--lyrics", required=True)
+    parser.add_argument("--cot", choices=["off", "melody", "full"], default="full")
+    parser.add_argument("--seed", type=int, default=831001)
+    parser.add_argument("--max-frames", type=int, default=200, help="semantic frame budget (25 frames = 1 s)")
+    parser.add_argument("--abc-file", default=None, help="external ABC score (requires cot=melody/full)")
+    parser.add_argument("--output", default="yue2_song.wav")
+    parser.add_argument(
+        "--dump-tokens",
+        default=None,
+        metavar="PATH",
+        help="write prompt + generated token ids to this json (default: <output>.tokens.json)",
+    )
+    parser.add_argument(
+        "--gpu-memory-utilization",
+        type=float,
+        default=None,
+        help="override the deploy yaml's per-stage gpu_memory_utilization "
+        "(necessary on cards shared with other tenants: the yaml fraction is "
+        "of the WHOLE card, not of the free remainder)",
+    )
+    args = parser.parse_args()
+
+    if args.vae:
+        import os
+
+        os.environ["YUE2_VAE"] = args.vae
+
+    tokenizer = YuE2TextTokenizer(Path(args.model) / "qwen.tiktoken")
+    engine_kwargs: dict = {}
+    if args.gpu_memory_utilization is not None:
+        engine_kwargs["gpu_memory_utilization"] = args.gpu_memory_utilization
+    engine = Omni(model=args.model, **engine_kwargs)
+
+    abc_ids = None
+    if args.abc_file is not None:
+        if args.cot == "off":
+            raise SystemExit("--abc-file requires --cot melody|full")
+        abc_ids = tokenizer.encode(Path(args.abc_file).read_text())
+        print(f"Using supplied ABC score ({len(abc_ids)} tokens)")
+    elif args.cot != "off":
+        prompt_ids = abc_prefix_ids(tokenizer.encode, args.style, args.lyrics, args.cot)
+        prompt = {"prompt_token_ids": prompt_ids}
+        params = sampling_params(engine, phase="abc", seed=args.seed, max_frames=args.max_frames, prompt_ids=prompt_ids)
+        outputs = engine.generate([prompt], [params])
+        generated = list(outputs[0].outputs[0].token_ids)
+        abc_ids = abc_ids_from_generated(generated)
+        print(f"Generated ABC score ({len(abc_ids)} tokens):\n{tokenizer.decode(abc_ids)}\n")
+
+    prompt_ids = semantic_prefix_ids(tokenizer.encode, args.style, args.lyrics, args.cot, abc_ids=abc_ids)
+    prompt = {"prompt_token_ids": prompt_ids}
+    params = sampling_params(
+        engine, phase="semantic", seed=args.seed, max_frames=args.max_frames, prompt_ids=prompt_ids
+    )
+    outputs = engine.generate([prompt], [params])
+    output = outputs[0].outputs[0]
+    generated_ids = list(output.token_ids)
+    if args.dump_tokens:
+        import json
+
+        payload = {
+            "prompt_token_ids": prompt_ids,
+            "generated_token_ids": generated_ids,
+            "seed": args.seed,
+            "cot": args.cot,
+            "max_frames": args.max_frames,
+        }
+        dump_path = args.dump_tokens if args.dump_tokens.endswith(".json") else args.output + ".tokens.json"
+        Path(dump_path).write_text(json.dumps(payload))
+        print(f"Dumped {len(generated_ids)} generated token ids to {dump_path}")
+    mm = output.multimodal_output or {}
+    audio = mm.get("audio")
+    if audio is None and "model_outputs" in mm:
+        audio = mm["model_outputs"]
+    if audio is None:
+        raise SystemExit("No audio in the model output; check server logs")
+    sr = int(mm.get("sr").item()) if mm.get("sr") is not None else 48000
+    truncated = False
+    meta = mm.get("meta") or {}
+    if "truncated" in meta:
+        truncated = bool(torch.as_tensor(meta["truncated"]).reshape(-1)[-1].item())
+
+    waveform = audio.float()
+    # soundfile wants [frames, channels] float in [-1, 1]; torchaudio has no
+    # wheel matching vllm 0.29.0's torch pin, so the driver uses libsndfile.
+    sf.write(args.output, waveform.T.contiguous(), sr)
+    print(
+        f"Saved {args.output}: {waveform.shape[-1] / sr:.1f}s @ {sr} Hz stereo, "
+        f"truncated={truncated}, generated {len(output.token_ids)} semantic tokens"
+    )
+
+
+if __name__ == "__main__":
+    main()

recipes/m-a-p/YuE2-3B.md

@@ -0,0 +1,85 @@
+# YuE2-3B
+
+> Text-to-music: style + lyrics (optionally an ABC score) in, a 48 kHz stereo song out
+
+## Summary
+
+- Vendor: m-a-p (Multimodal Art Projection)
+- Model: `m-a-p/YuE2-3B` + `m-a-p/YuE2-Vae`
+- Task: Text-to-music generation with an editable symbolic plan (ABC notation)
+- Mode: Offline generation and online `/v1/audio/speech` serving (adapter `tts_adapters/yue2.py`)
+- Integration issue: [#7661](https://github.com/vllm-project/vllm-omni/issues/7661)
+
+## When to use this recipe
+
+Use it to generate songs from a style description and lyrics. Unlike MiniMax
+Music 3, YuE2-3B first writes an ABC score (chords and melody, `cot=full` /
+`cot=melody`) that you can inspect, edit and resubmit (`--abc-file`), then
+renders it with the lyrics. `cot=off` skips the score and generates music
+directly. Weights are **CC BY-NC 4.0** (non-commercial); this recipe loads
+user-downloaded weights and bundles none.
+
+## References
+
+- Model card: [m-a-p/YuE2-3B](https://huggingface.co/m-a-p/YuE2-3B)
+- Upstream: [multimodal-art-projection/YuE](https://github.com/multimodal-art-projection/YuE) (pinned `bd90e4c`)
+- Offline example: `examples/offline_inference/yue2/end2end.py`
+- Deploy config: `vllm_omni/deploy/yue2.yaml` (auto-discovered for `model_type=yue2`)
+
+## Architecture
+
+| Component | Spec |
+| --- | --- |
+| Backbone | Qwen3-1.7B-class decoder, 28 layers, hidden 2048, GQA 16/8, q/k-norm, vocab 184,704, context 24,576 |
+| MoT | NAR path (`nar_self_attn`/`nar_mlp` per layer) shares embed_tokens, final norm and lm_head with the AR path |
+| Frames | Single codebook: one codec token per frame, 25 frames/s, span [151853, 184521) |
+| Sampling | temperature/top-k/top-p + windowed repetition penalty, seeded; model-owned (phase-masked vocab) |
+| Guidance | Off by default (1.0 full/melody, 1.01 off); not exercised |
+| Acoustic | 32-step midpoint flow matching over 64-dim latents on the NAR path, chunked ~12k frames |
+| VAE | `m-a-p/YuE2-Vae` decoder-only, tiled (1024 core + 16 halo frames), FP32 |
+| Output | 48 kHz stereo, whole song delivered when the semantic request finishes |
+
+Generation is one or two requests on a single AR stage: the abc phase writes
+the score, the semantic phase generates codec frames and, on its last step,
+solves the ODE and decodes the whole song in-engine (weights load once; the
+VAE loads at startup from `$YUE2_VAE` or the hub id).
+
+## Running
+
+```bash
+export YUE2_VAE=/models/YuE2-Vae   # or omit and let it resolve from the hub
+python examples/offline_inference/yue2/end2end.py \
+    --model /models/YuE2-3B \
+    --style "A melancholic lo-fi hip-hop track at 85 BPM in F minor" \
+    --lyrics "[Verse]
+Walking down the empty street at midnight
+[Chorus]
+And I keep on walking" \
+    --cot full --seed 831001 --max-frames 200 --output song.wav
+```
+
+`--max-frames 200` caps the song at 8 s (25 frames/s). The end token is
+masked for the preset's first 200 steps, so a 200-frame budget always
+truncates; a longer budget lets the song end on its own end token, and
+`truncated` in the report line says which happened. Serving follows the
+same rule: omit `max_new_tokens` (or set it above 200) for a natural
+ending.
+
+## Notes
+
+- **Memory (RTX 4090, 24 GB, `gpu_memory_utilization: 0.70`):** the engine
+  reserves ~15.8 GiB after startup; with 4 concurrent requests pinned to the
+  9000-frame cap (worst case), peak usage during the terminal NAR/VAE
+  finishing pass reaches ~20.9 GiB, leaving ~3 GiB of headroom. At 0.85 the
+  same workload OOMs inside the finishing pass, so the deploy yaml pins 0.70.
+- **Known limitations:** 4 concurrent full-length requests are the verified
+  shape (`max_num_seqs: 4`); CUDA graph capture is a
+  follow-up; the abc phase runs eagerly
+  after prefill (its tokens are the product, not audio).
+- Audio is not expected to match the upstream torch reference bit-for-bit
+  (fused vs eager kernels). Measured against the upstream torch reference
+  (bd90e4c, same style/lyrics/seed 831001, cot=off, 200 frames): prompt token
+  ids 63/63 bit-identical, and the sampled semantic stream agrees
+  bit-for-bit for the first 55 frames before kernel numerics diverge — well
+  past the "first frames agree" acceptance bar of the MiniMax Music 3
+  precedent.

tests/assets/yue2/twinkle_golden.abc

@@ -0,0 +1,6 @@
+X:1
+T:Twinkle Twinkle Little Star (SheetSage2 transcription of the Wikimedia piano rendition, single voice)
+M:4/4
+L:1/16
+K:C
+C4C4G4G4|A4A4G8|F4F4E4E4|D4D4C8|G4G4F4F4|E4E4D8|G4G4F4F4|E4E4D8|C4C4G4G4|A4A4G8|F4F4E4E4|D4D4C8|

tests/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 32,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 74,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 75,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
         EnvironmentVariableCategory.INTERNAL: 6,
     }
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
     )
     assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
-        ModelEnvironmentVariableDisposition.PROMOTE: 45,
+        ModelEnvironmentVariableDisposition.PROMOTE: 46,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 17,

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -1048,3 +1048,83 @@ def test_minimax_music3_speech_invalid_field_values(
     online_client.send_audio_speech_http_request(
         {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
     )
+
+
+# ─── POST /v1/audio/speech · YuE2-3B ───
+
+# Text-to-music like MiniMax Music 3, with one twist: the ABC score span is
+# request-controlled via extra_params.cot/abc and the two must agree, and
+# sampling is pinned by the checkpoint preset so sampling knobs are rejected.
+_YUE2_SPEECH = [
+    pytest.param(
+        OmniServerParams(
+            model="m-a-p/YuE2-3B",
+            stage_config_path=get_deploy_config_path("yue2.yaml"),
+            server_args=_SPEECH_SERVER_ARGS,
+        ),
+        id="yue2",
+        marks=hardware_marks(res={"cuda": "H100"}),
+    ),
+]
+
+_YUE2_LYRICS = "一闪一闪亮晶晶\n满天都是小星星"
+_YUE2_CAPTION = "Chinese heavy metal, distorted electric guitars, 140 BPM, key of C."
+_YUE2_ABC = "X:1\nT:Twinkle\nM:4/4\nL:1/4\nK:C\nC C G G|A A G2|"
+# 25 frames per second; the semantic preset bounds the frame budget.
+_YUE2_MAX_FRAMES = 9000
+
+
+@pytest.mark.parametrize(
+    "overrides, err_message",
+    [
+        pytest.param({"instructions": None}, ("instructions",), id="caption_missing"),
+        pytest.param({"instructions": "   "}, ("instructions",), id="caption_blank"),
+        pytest.param({"input": "   "}, ("input",), id="lyrics_blank"),
+        pytest.param({"voice": "alloy"}, ("voice",), id="voice_rejected"),
+        pytest.param({"speed": 1.5}, ("speed",), id="speed_rejected"),
+        pytest.param(
+            {"max_new_tokens": _YUE2_MAX_FRAMES + 1},
+            ("max_new_tokens", str(_YUE2_MAX_FRAMES)),
+            id="max_new_tokens_over_cap",
+        ),
+        pytest.param({"max_new_tokens": 199}, ("max_new_tokens",), id="max_new_tokens_below_min"),
+        pytest.param(
+            {"extra_params": {"cot": "melody"}},
+            ("abc",),
+            id="cot_melody_without_abc",
+        ),
+        pytest.param(
+            {"extra_params": {"abc": _YUE2_ABC}},
+            ("cot",),
+            id="abc_without_cot",
+        ),
+        pytest.param(
+            {"extra_params": {"temperature": 0.3}},
+            ("temperature",),
+            id="temperature_rejected",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _YUE2_SPEECH, indirect=True)
+def test_yue2_speech_invalid_field_values(
+    omni_server: OmniServer,
+    online_client: OnlineOmniClient,
+    overrides: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    body: dict[str, Any] = {
+        "model": omni_server.model,
+        "input": _YUE2_LYRICS,
+        "instructions": _YUE2_CAPTION,
+        "seed": 1,
+        "max_new_tokens": 250,
+        "response_format": "wav",
+    }
+    for key, value in overrides.items():
+        if value is None:
+            body.pop(key, None)
+        else:
+            body[key] = value
+    online_client.send_audio_speech_http_request(
+        {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )

tests/e2e/offline_inference/test_yue2.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""End-to-end regressions for YuE2 offline WAV export."""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import soundfile as sf
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import get_model_prefix
+from vllm_omni.transformers_utils.repo_utils import hf_api
+
+pytestmark = [
+    pytest.mark.full_model,
+    pytest.mark.tts,
+    *hardware_marks(res={"cuda": ["L4", "H100", "B200"]}, num_cards=1),
+]
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+
+
+@pytest.fixture(scope="module")
+def vae_path() -> Path:
+    path = Path(os.environ.get("YUE2_VAE", get_model_prefix() + "m-a-p/YuE2-Vae"))
+    if not path.is_dir():
+        path = Path(hf_api().snapshot_download(str(path), allow_patterns=["*.json", "*.safetensors"]))
+    return path
+
+
+@pytest.mark.parametrize("max_frames,truncated", [(200, True), (9000, False)], ids=["budget", "natural-eos"])
+def test_offline_example_saves_audio(vae_path: Path, tmp_path: Path, max_frames: int, truncated: bool) -> None:
+    model = Path(os.environ.get("YUE2_MODEL_DIR", get_model_prefix() + "m-a-p/YuE2-3B"))
+    if not model.is_dir():
+        model = Path(
+            hf_api().snapshot_download(str(model), allow_patterns=["*.json", "*.safetensors", "qwen.tiktoken"])
+        )
+    output = tmp_path / "song.wav"
+    env = os.environ.copy()
+    env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+    env["PYTHONPATH"] = str(REPO_ROOT) + os.pathsep + env.get("PYTHONPATH", "")
+    completed = subprocess.run(
+        [
+            sys.executable,
+            str(REPO_ROOT / "examples/offline_inference/yue2/end2end.py"),
+            "--model",
+            str(model),
+            "--vae",
+            str(vae_path),
+            "--style",
+            "gentle acoustic folk, clear English vocals, 90 BPM",
+            "--lyrics",
+            "[Verse]\nMorning light across the bay\nCarry all my dreams away\n[Outro]",
+            "--cot",
+            "off",
+            "--seed",
+            "831001",
+            "--max-frames",
+            str(max_frames),
+            "--output",
+            str(output),
+        ],
+        cwd=REPO_ROOT,
+        env=env,
+        capture_output=True,
+        text=True,
+        timeout=900,
+    )
+    assert completed.returncode == 0, completed.stdout + completed.stderr
+    assert f"truncated={truncated}" in completed.stdout
+    audio, sample_rate = sf.read(output, always_2d=True)
+    assert sample_rate == 48000 and audio.shape[1] == 2
+    assert np.isfinite(audio).all() and np.abs(audio).max() > 0
+    duration = len(audio) / sample_rate
+    if truncated:
+        assert duration == pytest.approx(max_frames / 25, abs=0.1)
+    else:
+        assert 8 <= duration < max_frames / 25

tests/e2e/online_serving/test_yue2.py

@@ -0,0 +1,237 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""E2E online tests for YuE2-3B text-to-music on ``/v1/audio/speech``.
+
+Scenario: heavy-metal "Twinkle Twinkle Little Star" — Chinese lyrics plus a
+metal style caption plus an ABC score of the melody (``extra_params.cot =
+"melody"``). The golden score is a SheetSage2 transcription of the Wikimedia
+piano rendition, reduced to a single voice (``tests/assets/yue2``).
+
+Two programmatic content checks, both calibrated against a same-seed control
+request with ``cot="off"`` (no ABC conditioning):
+
+- Style: CLAP zero-shot probability of "heavy metal" must clear 0.9.
+- Melody: the generated song is transcribed back with SheetSage2 and its
+  pitch-class sequence is LCS-matched against the golden score (max over
+  major/minor variants and 12 transpositions, tolerating pedal-tone
+  insertions). SheetSage2's transcription of distorted metal vocals carries
+  roughly ±0.05 noise, so the bar sits at 0.80; the ABC-conditioned request
+  must also beat the control.
+
+The content checks run in a dedicated verification interpreter because
+SheetSage2 wants a dependency set (mir_eval, pretty_midi, numpy<2) that must
+not leak into the serving environment. They are enabled by three env vars:
+
+- ``YUE2_E2E_VERIFY_PYTHON``: python with sheetsage2 + CLAP deps installed
+- ``YUE2_E2E_CLAP_DIR``: local laion/clap-htsat-unfused snapshot (~0.6 GB)
+- ``YUE2_E2E_SHEETSAGE2_DIR``: local m-a-p/SheetSage2 snapshot (~0.25 GB,
+  plus its MERT-v2 backbone ~2.5 GB from the HF cache)
+
+Without them the test still asserts the structural floor (non-empty payload
+above the byte floor) and passes on that alone, following the MiniMax Music 3
+precedent.
+
+Weekly rather than per-PR: each request is a 16 s song and the verification
+checkpoints add ~5.5 GB of downloads. Validated on a single RTX 4090
+(24 GB); the hardware mark targets an H100 because the CI fleet has no
+consumer cards.
+"""
+
+import concurrent.futures
+import io
+import json
+import os
+import subprocess
+import wave
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import httpx
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "m-a-p/YuE2-3B"
+DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 900.0
+
+# 400 frames = 16 s of song. The byte floor sits far below the observed
+# payload (~3.0 MB of 48 kHz stereo 16-bit) to leave room for an early EOS.
+_FRAMES = 400
+# Concurrent pair stays short: it exists to exercise row re-indexing, not
+# to double the wall time.
+_CONCURRENT_FRAMES = 250
+_MIN_BYTES = 1_000_000
+
+LYRICS = "一闪一闪亮晶晶\n满天都是小星星\n挂在天空放光明\n好像千万小眼睛"
+CAPTION = (
+    "Chinese heavy metal, distorted electric guitars, aggressive double-kick drums, powerful vocals, 140 BPM, key of C"
+)
+SEED = 831001
+
+_VERIFY_PYTHON = os.environ.get("YUE2_E2E_VERIFY_PYTHON")
+_CLAP_DIR = os.environ.get("YUE2_E2E_CLAP_DIR")
+_SHEETSAGE2_DIR = os.environ.get("YUE2_E2E_SHEETSAGE2_DIR")
+_CONTENT_CHECKS_ENABLED = all([_VERIFY_PYTHON, _CLAP_DIR, _SHEETSAGE2_DIR])
+_VERIFY_HELPER = os.path.join(os.path.dirname(__file__), "yue2_metal_verify.py")
+
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("yue2.yaml"),
+            server_args=["--disable-log-stats"],
+        ),
+        id="yue2",
+    )
+]
+
+
+def _http_speech(omni_server, *, seed: int, frames: int) -> bytes:
+    """Raw-HTTP speech request returning (and structurally checking) the WAV.
+
+    The OpenAI SDK client swallows response headers, so the truncation flag is
+    cross-checked here: ``X-Audio-Truncated`` must agree with the audio — the
+    full frame budget (a 25 frames/s song of exactly frames/25 seconds) means
+    truncated, a shorter song means a natural end token.
+    """
+    response = httpx.post(
+        f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech",
+        json={
+            "model": omni_server.model,
+            "input": LYRICS,
+            "instructions": CAPTION,
+            "seed": seed,
+            "max_new_tokens": frames,
+            "stream": False,
+            "response_format": "wav",
+            "extra_params": {"cot": "off"},
+        },
+        timeout=DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+    )
+    assert response.status_code == 200, response.text[:500]
+    truncated = response.headers.get("X-Audio-Truncated")
+    assert truncated in ("true", "false"), f"X-Audio-Truncated missing/invalid: {truncated!r}"
+    audio = response.content
+    with wave.open(io.BytesIO(audio)) as wav:
+        assert wav.getnchannels() == 2
+        assert wav.getframerate() == 48000
+        duration = wav.getnframes() / wav.getframerate()
+    hit_budget = duration >= frames / 25 - 0.5
+    assert truncated == ("true" if hit_budget else "false"), (
+        f"X-Audio-Truncated={truncated} but duration={duration:.2f}s vs {frames / 25:.1f}s budget"
+    )
+    return audio
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None:
+    """Heavy-metal Twinkle: ABC-conditioned request vs same-seed control.
+
+    Deploy Setting: yue2.yaml
+    Input Modal: text (lyrics) + text (style caption) + ABC score (extra_params)
+    Output Modal: audio, 48 kHz stereo
+    Input Setting: stream=False, seed pinned, 400 frames
+    """
+    golden_abc = get_asset_path("yue2/twinkle_golden.abc").read_text()
+
+    def _request(extra_params: dict) -> bytes:
+        responses = openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": LYRICS,
+                "instructions": CAPTION,
+                "seed": SEED,
+                "max_new_tokens": _FRAMES,
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_BYTES,
+                # Sung metal vocals are not whisper-transcribable; skip the
+                # full_model transcript check (our content checks do the work).
+                "transcript_expected_text": "",
+                "extra_params": extra_params,
+            }
+        )
+        audio = responses[0].audio_bytes
+        assert audio, "empty audio payload"
+        return audio
+
+    metal_path = tmp_path / "metal.wav"
+    control_path = tmp_path / "control.wav"
+    metal_bytes = _request({"cot": "melody", "abc": golden_abc})
+    control_bytes = _request({"cot": "off"})
+    metal_path.write_bytes(metal_bytes)
+    control_path.write_bytes(control_bytes)
+
+    # CI runs without the verification assets, so the structural checks below
+    # are the only guard against the two real bugs this PR fixed during review:
+    # mono/half-speed output (WAV header + duration bounds) and the int32
+    # crash under concurrent row re-indexing (two simultaneous requests).
+    # Duration bounds: the end token is masked for the first 200 steps
+    # (min_tokens), and the 400-frame budget caps the song, so a healthy
+    # request lands in [8, 16.5] s; the half-speed bug produced ~32 s.
+    for label, audio in (("metal", metal_bytes), ("control", control_bytes)):
+        with wave.open(io.BytesIO(audio)) as wav:
+            assert wav.getnchannels() == 2, f"{label}: {wav.getnchannels()} channels, want stereo"
+            assert wav.getframerate() == 48000, f"{label}: {wav.getframerate()} Hz, want 48 kHz"
+            duration = wav.getnframes() / wav.getframerate()
+            assert _FRAMES / 50 <= duration <= _FRAMES / 25 + 0.5, (
+                f"{label}: {duration:.2f}s outside [{_FRAMES / 50}, {_FRAMES / 25 + 0.5}]s"
+            )
+
+    # A pinned-seed request must flag truncation consistently with the audio
+    # (see _http_speech); the OpenAI SDK response does not expose headers.
+    _http_speech(omni_server, seed=SEED, frames=_FRAMES)
+
+    # Two concurrent requests with different seeds: both must succeed and
+    # produce different songs (the model-owned sampler re-indexes rows across
+    # concurrent requests; a dtype mismatch here crashed the engine before).
+    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
+        futures = [pool.submit(_http_speech, omni_server, seed=SEED + i, frames=_CONCURRENT_FRAMES) for i in range(2)]
+        pair = [f.result() for f in futures]
+    assert pair[0] != pair[1], "different seeds produced identical audio"
+
+    # Without the verification assets the structural assertions above
+    # (non-empty payload, byte floor) are the whole test — pass on them.
+    if not _CONTENT_CHECKS_ENABLED:
+        return
+    assert _VERIFY_PYTHON is not None
+    assert _CLAP_DIR is not None
+    assert _SHEETSAGE2_DIR is not None
+
+    proc = subprocess.run(
+        [
+            _VERIFY_PYTHON,
+            _VERIFY_HELPER,
+            "--clap-dir",
+            _CLAP_DIR,
+            "--sheetsage2-dir",
+            _SHEETSAGE2_DIR,
+            "--work-dir",
+            str(tmp_path),
+            str(metal_path),
+            str(control_path),
+        ],
+        capture_output=True,
+        text=True,
+        timeout=1800,
+    )
+    assert proc.returncode == 0, f"verification helper failed:\n{proc.stderr}"
+    report = json.loads(proc.stdout.strip().splitlines()[-1])
+
+    metal_prob = report["metal_probability"]
+    assert metal_prob >= 0.9, f"style check: P(heavy metal)={metal_prob:.3f} < 0.9"
+
+    melody_score = report["melody_score"]
+    control_score = report["control_score"]
+    assert melody_score >= 0.80, f"melody check: score={melody_score:.3f} < 0.80"
+    assert melody_score > control_score, (
+        f"calibration: ABC-conditioned {melody_score:.3f} not above control {control_score:.3f}"
+    )

tests/e2e/online_serving/yue2_metal_verify.py

@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Content-gate helper for the YuE2 heavy-metal Twinkle e2e test.
+
+Run as ``__main__`` by a dedicated verification interpreter (see
+``YUE2_E2E_VERIFY_PYTHON`` in test_yue2.py), NOT by the serving environment:
+SheetSage2 wants its own dependency set (mir_eval, pretty_midi, numpy<2),
+which must not leak into the serving venv.
+
+Usage:
+    python yue2_metal_verify.py --clap-dir DIR --sheetsage2-dir DIR \
+        metal.wav control.wav
+
+Prints one JSON object on stdout:
+    {"metal_probability": float, "melody_score": float, "control_score": float}
+
+Gates (applied by the caller):
+- metal_probability >= 0.9  (CLAP zero-shot heavy-metal style)
+- melody_score >= 0.80 and melody_score > control_score
+  (SheetSage2 re-transcription, pitch-class LCS vs the golden score, max over
+  major/minor variants and 12 transpositions, pedal-tone tolerant)
+"""
+
+import argparse
+import json
+import re
+import wave
+from typing import Any
+
+import numpy as np
+
+_NOTE_BASE = {"C": 0, "D": 2, "E": 4, "F": 5, "G": 7, "A": 9, "B": 11}
+_GOLDEN_MAJOR = "C4C4G4G4|A4A4G8|F4F4E4E4|D4D4C8|G4G4F4F4|E4E4D8|G4G4F4F4|E4E4D8|C4C4G4G4|A4A4G8|F4F4E4E4|D4D4C8|"
+# Metal arrangements commonly minor-ize the melody; accept both modes.
+_GOLDEN_MINOR = (
+    "C4C4G4G4|_A4_A4G8|F4F4_E4_E4|D4D4C8|G4G4F4F4|_E4_E4D8|G4G4F4F4|_E4_E4D8|C4C4G4G4|_A4_A4G8|F4F4_E4_E4|D4D4C8|"
+)
+
+_STYLE_CANDIDATES = [
+    "heavy metal music with distorted guitars and aggressive drums",
+    "a soft children's lullaby",
+    "classical piano music",
+    "pop music",
+]
+
+
+def _abc_note_to_midi(token):
+    m = re.match(r"([\^_=]*)([A-Ga-g])([,']*)(\d*)", token)
+    if not m:
+        return None
+    acc, letter, oct_marks, _ = m.groups()
+    pc = _NOTE_BASE[letter.upper()] + acc.count("^") - acc.count("_")
+    octave = 5 if letter.isupper() else 6
+    octave += oct_marks.count("'") - oct_marks.count(",")
+    return 12 * octave + pc
+
+
+def _collapse(seq):
+    out: list[Any] = []
+    for x in seq:
+        if not out or x != out[-1]:
+            out.append(x)
+    return out
+
+
+def _golden_events(body):
+    return _collapse([_abc_note_to_midi(t) for t in re.findall(r"[\^_=]*[A-G]\d*", body)])
+
+
+def _melody_events(score_abc):
+    """Note events of the voice carrying the most notes in a SheetSage2 score."""
+    voices: dict[str | None, list[int]] = {}
+    current = None
+    for line in score_abc.splitlines():
+        line = line.strip()
+        if line.startswith("V:"):
+            parts = line.split()
+            current = parts[1] if len(parts) > 1 else None
+            continue
+        if current is None or not line or line.startswith(("%", "X:", "T:", "M:", "L:", "Q:", "K:")):
+            continue
+        for tok in re.findall(r"[\^_=]*[A-Ga-g][,']*\d*", line.replace("|", " ")):
+            midi = _abc_note_to_midi(tok)
+            if midi is not None:
+                voices.setdefault(current, []).append(midi)
+    return _collapse(max(voices.values(), key=len, default=[]))
+
+
+def _lcs(a, b):
+    dp = [0] * (len(b) + 1)
+    for x in a:
+        prev = 0
+        for j, y in enumerate(b, 1):
+            tmp = dp[j]
+            dp[j] = prev + 1 if x == y else max(dp[j], dp[j - 1])
+            prev = tmp
+    return dp[-1]
+
+
+def melody_score(score_abc):
+    gen = [m % 12 for m in _melody_events(score_abc)]
+    if len(gen) < 4:
+        return 0.0
+    best = 0
+    for body in (_GOLDEN_MAJOR, _GOLDEN_MINOR):
+        gold = [m % 12 for m in _golden_events(body)]
+        best = max(best, max(_lcs(gen, [(p + r) % 12 for p in gold]) for r in range(12)))
+    return best / len(_golden_events(_GOLDEN_MAJOR))
+
+
+def _wav_to_mono(path):
+    with wave.open(path) as wf:
+        sr = wf.getframerate()
+        channels = wf.getnchannels()
+        frames = wf.readframes(wf.getnframes())
+    data = np.frombuffer(frames, dtype=np.int16).astype(np.float32) / 32768.0
+    if channels > 1:
+        data = data.reshape(-1, channels).mean(axis=1)
+    return data, sr
+
+
+def metal_probability(wav_path, clap_dir):
+    import torch
+    from transformers import ClapModel, ClapProcessor
+
+    model = ClapModel.from_pretrained(clap_dir)
+    proc = ClapProcessor.from_pretrained(clap_dir)
+    audio, sr = _wav_to_mono(wav_path)
+    if sr != 48000:
+        n = int(len(audio) * 48000 / sr)
+        audio = np.interp(np.linspace(0, len(audio), n), np.arange(len(audio)), audio).astype(np.float32)
+    audio = audio[: 48000 * 10]
+    inputs = proc(
+        text=_STYLE_CANDIDATES,
+        audios=[audio] * len(_STYLE_CANDIDATES),
+        sampling_rate=48000,
+        return_tensors="pt",
+        padding=True,
+    )
+    with torch.no_grad():
+        probs = model(**inputs).logits_per_audio[0].softmax(dim=-1)
+    return float(probs[0])
+
+
+def transcribe_melody_abc(wav_path, sheetsage2_dir, output_dir):
+    import torch
+    from transformers import AutoModel
+
+    audio, sr = _wav_to_mono(wav_path)
+    model = AutoModel.from_pretrained(sheetsage2_dir, trust_remote_code=True)
+    result = model.transcribe(
+        torch.from_numpy(audio[None, :]),
+        sampling_rate=sr,
+        output_dir=str(output_dir),
+        melody_only=True,
+        dtype="fp32",
+    )
+    abc = result.get("abc")
+    if not abc:
+        raise RuntimeError(f"SheetSage2 produced no melody ABC: {result.get('abc_error')}")
+    return abc
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--clap-dir", required=True)
+    parser.add_argument("--sheetsage2-dir", required=True)
+    parser.add_argument("metal_wav")
+    parser.add_argument("control_wav")
+    parser.add_argument("--work-dir", default=None)
+    args = parser.parse_args()
+
+    work = args.work_dir or "."
+    report = {"metal_probability": metal_probability(args.metal_wav, args.clap_dir)}
+    report["melody_score"] = melody_score(
+        transcribe_melody_abc(args.metal_wav, args.sheetsage2_dir, f"{work}/metal_score")
+    )
+    report["control_score"] = melody_score(
+        transcribe_melody_abc(args.control_wav, args.sheetsage2_dir, f"{work}/control_score")
+    )
+    print(json.dumps(report))
+
+
+if __name__ == "__main__":
+    main()

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1105,6 +1105,38 @@ class TestTTSMethods:
         assert isinstance(audio_obj, CreateAudio)
         assert audio_obj.speed == 1.0
 
+    @pytest.mark.asyncio
+    async def test_audio_synthesis_error_flag_raises_tts_generation_error(
+        self,
+        speech_server,
+        mocker: MockerFixture,
+    ):
+        """A model-flagged synthesis failure surfaces as a non-retryable
+        TTSGenerationError, so create_speech answers 500 instead of shipping a
+        zero-length WAV or crashing on an unpacked Response."""
+
+        async def mock_generate():
+            yield create_mock_audio_output_for_test()
+
+        mocker.patch.object(
+            speech_server,
+            "_prepare_speech_generation",
+            new=mocker.AsyncMock(return_value=("speech-synth-err", mock_generate(), {})),
+        )
+        adapter = mocker.MagicMock()
+        adapter.collect_response_metadata = lambda _audio_output, collect: collect.__setitem__(
+            "audio_synthesis_error", True
+        )
+        mocker.patch.object(speech_server, "_get_tts_adapter", return_value=adapter)
+
+        with pytest.raises(TTSGenerationError, match="failed to synthesize audio") as exc_info:
+            await speech_server._generate_audio_bytes(
+                OpenAICreateSpeechRequest(input="Hello"),
+                collect={},
+            )
+
+        assert exc_info.value.retryable is False
+
     def test_is_tts_detection_with_tts_stage(self, mocker: MockerFixture):
         """Test TTS model detection when TTS stage exists."""
         mock_engine_client = mocker.MagicMock()

tests/entrypoints/openai_api/test_tts_detection.py

@@ -183,6 +183,7 @@ _PIPELINE_STAGES = [
     "token2wav",
     "tts",
     "encoder",
+    "yue2",
 ]
 
 _STAGES = [*_PIPELINE_STAGES, None, "vae", "not_a_real_stage"]

tests/worker/test_gpu_ar_model_runner.py

@@ -15,6 +15,7 @@ import numpy as np
 import pytest
 import torch
 from vllm.sampling_params import SamplingParams
+from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.worker import gpu_input_batch
 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 
@@ -103,6 +104,8 @@ def test_post_sample_talker_mtp_uses_current_sample_and_hidden() -> None:
     )
 
     assert received["req_ids"] == ["ready"]
+    assert isinstance(received["input_ids"], torch.Tensor)
+    assert isinstance(received["hidden_states"], torch.Tensor)
     assert received["input_ids"].tolist() == [101]
     assert received["hidden_states"].tolist() == [[1.0, 2.0]]
     assert received["req_infos"] == [{"duplex": {"data_plane": True}}]
@@ -139,6 +142,8 @@ def test_post_sample_talker_mtp_uses_gpu_token_with_async_scheduling() -> None:
     )
 
     assert received["req_ids"] == ["ready"]
+    assert isinstance(received["input_ids"], torch.Tensor)
+    assert isinstance(received["hidden_states"], torch.Tensor)
     assert received["input_ids"].tolist() == [101]
     assert received["hidden_states"].tolist() == [[1.0, 2.0]]
     assert received["req_infos"] == [{"duplex": {"data_plane": True}}]
@@ -275,7 +280,12 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
     )
     input_batch.sampling_metadata = input_batch._make_sampling_metadata()
 
-    received = []
+    received: list[SamplingMetadata] = []
+
+    def sample(logits: torch.Tensor, metadata: SamplingMetadata) -> str:
+        received.append(metadata)
+        return "model-sampler"
+
     runner = object.__new__(GPUARModelRunner)
     runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True)
     runner._pooler_payload_include_hidden_flag = True
@@ -283,7 +293,7 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
     runner.model = SimpleNamespace(
         prefer_model_sampler=True,
         skips_model_sampler_output_token_history=True,
-        sample=lambda logits, metadata: received.append(metadata) or "model-sampler",
+        sample=sample,
     )
     runner.sampler = SimpleNamespace()
     logits = torch.zeros((1, 4))
@@ -1670,7 +1680,6 @@ class TestMergeModelKvTransferMetadata:
         # ...while the engine-shared originals are untouched.
         assert original["r1"]["custom_metadata"] is original_meta
         assert original_meta == {"a": 1}
-        assert "talker_codes" not in original["r1"]["custom_metadata"]
 
     def test_merge_without_model_meta_passes_entry_through(self):
         class _Model:
@@ -1713,7 +1722,7 @@ class TestDownstreamPayloadMemoization:
         return runner
 
     def test_missing_marker_defaults_true_without_memoizing(self):
-        stages = {"r1": None}
+        stages: dict[str, int | None] = {"r1": None}
         runner = self._runner(stages)
 
         assert runner._request_needs_downstream_stage_payload("r1") is True
@@ -1861,6 +1870,7 @@ class TestPreferModelSamplerNoneFallback:
             "minicpmo_4_5",
             "minimax_music3",
             "nemotron_voicechat",
+            "yue2",  # Always supplies SamplerOutput, including empty/prefill steps.
         }
         assert declarers == expected, (
             "The set of models declaring `prefer_model_sampler` changed:\n"

vllm_omni/config/environment_variable_inventory.py

@@ -189,6 +189,7 @@ _MODEL_PROMOTE = (
     "VLLM_OMNI_SEEDVR2_SHARDED_CLIP_PIXELS",
     "VLLM_OMNI_SEEDVR2_SHARDED_FRAME_PIXELS",
     "VLLM_OMNI_SENSENOVA_PAGED_DECODE",
+    "YUE2_VAE",  # VAE checkpoint selection belongs in typed model configuration.
 )
 
 _MODEL_REQUEST_SCOPE = (

vllm_omni/config/pipeline_registry.py

@@ -125,6 +125,7 @@ from vllm_omni.model_executor.models.step_audio2.pipeline import (
 from vllm_omni.model_executor.models.voxcpm2.pipeline import VOXCPM2_PIPELINE
 from vllm_omni.model_executor.models.voxtral_tts.pipeline import VOXTRAL_TTS_PIPELINE
 from vllm_omni.model_executor.models.wan2_2.pipeline import WAN2_2_TI2V_PIPELINE
+from vllm_omni.model_executor.models.yue2.pipeline import YUE2_PIPELINE
 
 logger = init_logger(__name__)
 
@@ -208,6 +209,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "voxcpm2": VOXCPM2_PIPELINE,
     "voxtral_tts": VOXTRAL_TTS_PIPELINE,
     "wan2_2_ti2v": WAN2_2_TI2V_PIPELINE,
+    "yue2": YUE2_PIPELINE,
 }
 
 

vllm_omni/deploy/yue2.yaml

@@ -0,0 +1,55 @@
+# YuE2-3B deploy: single-stage native-AR text-to-music (~3.6B bf16 + VAE).
+# Target 1x 24 GB card; verified on a single RTX 4090.
+#
+# The architectures pin is required: the checkpoint's config.json names
+# "YuE2ForCausalLM" with an auto_map that only makes sense under
+# trust_remote_code, so the deploy pins vLLM to this package's model class.
+# model_type stays "yue2" so the registered Yue2Config is used; no
+# trust_remote_code is needed (the registered config/model win over the
+# checkpoint's auto_map).
+pipeline: yue2
+async_chunk: false
+dtype: bfloat16
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 4
+    # Measured on a 24 GB RTX 4090 with 4 concurrent requests pinned to the
+    # 9000-frame cap (worst case): the engine reserves ~15.8 GiB at 0.70 and
+    # peak usage during the terminal NAR/VAE finishing pass reaches 20.9 GiB,
+    # leaving ~3 GiB of headroom. 0.85 reserves ~19.4 GiB and the same
+    # workload OOMs inside the finishing pass.
+    gpu_memory_utilization: 0.70
+    # Phase 1 keeps the backbone eager: the model-owned sampler runs after the
+    # forward, and the terminal NAR/VAE pass runs inside the finishing step.
+    # CUDA-graph capture of the standard Qwen3 backbone is a perf follow-up.
+    enforce_eager: true
+    # Prefix caching stays off: the omni prefix-cache save snapshots the
+    # step's mm payload before the model-owned sampler runs, but YuE2
+    # appends the decoded waveform inside sample() (the terminal NAR/VAE
+    # pass), so the cached payload would ship empty audio. Re-enable once
+    # the runner saves post-sample.
+    enable_prefix_caching: false
+    max_num_batched_tokens: 24576
+    max_model_len: 24576
+    devices: "0"
+    # The checkpoint ships no HF tokenizer files (only qwen.tiktoken, which
+    # the offline driver reads itself); every prompt arrives as token ids,
+    # so the engine builds no tokenizer. Placeholder tokenizer dir follows
+    # the indextts2 precedent.
+    skip_tokenizer_init: true
+    tokenizer: gpt2
+    hf_overrides:
+      architectures: ["Yue2ForCausalLM"]
+    default_sampling_params:
+      # Placeholder values: the model's own sampler overrides everything from
+      # extra_args. Their presence turns on has_sampling_extra_args, which is
+      # what threads the per-request yue2_* keys into the model's forward.
+      extra_args:
+        yue2_phase: "semantic"
+        yue2_seed: 831001
+        yue2_max_audio_frames: 200
+      detokenize: false
+      # Engine-side backstop; the model stops on its own end token first.
+      max_tokens: 9000
+      temperature: 1.0

vllm_omni/entrypoints/openai/serving_speech.py

@@ -2275,6 +2275,22 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 if ts is not None:
                     collect["word_timestamps"] = ts
 
+            # Let the adapter fold engine-side metadata (e.g. YuE2's
+            # meta.truncated) into ``collect`` for response headers.
+            if collect is not None and (adapter := self._get_tts_adapter()) is not None:
+                adapter.collect_response_metadata(audio_output, collect)
+
+            # A model can flag a per-request synthesis failure through the
+            # adapter (e.g. YuE2's terminal NAR/VAE pass OOMing on one
+            # request); answer 500 instead of shipping a zero-length WAV.
+            # Raising (not returning a Response) keeps this function's
+            # tuple contract; create_speech maps TTSGenerationError to 500.
+            if collect is not None and collect.get("audio_synthesis_error"):
+                raise TTSGenerationError(
+                    "The model failed to synthesize audio for this request",
+                    retryable=False,
+                )
+
             audio_tensor = audio_output[audio_key]
             sr_raw = audio_output.get("sr", 24000)
             sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
@@ -2715,6 +2731,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         "(use the WebSocket streaming path for long transcripts)",
                         len(ts_json),
                     )
+            if collect.get("audio_truncated") is not None:
+                headers["X-Audio-Truncated"] = "true" if collect["audio_truncated"] else "false"
             return Response(content=audio_bytes, media_type=media_type, headers=headers)
 
         except asyncio.CancelledError:

vllm_omni/entrypoints/openai/tts_adapters/__init__.py

@@ -124,6 +124,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import (  # noqa: E402,F401
     step_audio2,
     voxcpm2,
     voxtral,
+    yue2,
 )
 
 __all__ = [

vllm_omni/entrypoints/openai/tts_adapters/base.py

@@ -320,6 +320,15 @@ class TTSModelAdapter(ABC):
         to unrelated TTS models.
         """
 
+    def collect_response_metadata(self, audio_output: Mapping[str, Any], collect: dict) -> None:
+        """Fold engine-side metadata from the mm payload into ``collect``.
+
+        Non-streaming only. The orchestrator surfaces select ``collect``
+        entries as response headers (precedent: ``word_timestamps`` ->
+        ``X-Word-Timestamps``, ``audio_truncated`` -> ``X-Audio-Truncated``);
+        adapters opt in by setting the keys it knows. Default: no-op.
+        """
+
     async def warmup(self) -> None:
         return
 

vllm_omni/entrypoints/openai/tts_adapters/yue2.py

@@ -0,0 +1,267 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""YuE2-3B serving adapter for ``/v1/audio/speech``.
+
+This is a text-to-music model on a speech endpoint, so most of the speech
+contract does not apply. ``input`` carries the lyrics and ``instructions`` the
+style caption (genre, instrumentation, tempo, mood). There is no speaker to
+select, no reference audio, and no sampling knobs: the request-local preset is
+pinned by the checkpoint's reference config. Model-specific controls travel in
+``extra_params``:
+
+- ``cot``: ``"off"`` (default) | ``"melody"`` | ``"full"`` — how the ABC
+  score span in the prompt is handled.
+- ``abc``: an ABC score as text, required for ``cot=melody|full``. Serving v1
+  does not run the model-generated ABC phase (that is two sequential engine
+  requests; use the offline example for it).
+
+Unsupported parameters are rejected rather than silently ignored, following
+the MiniMax Music 3 precedent.
+"""
+
+import secrets
+from collections.abc import Mapping
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import (
+    ARTTSAdapter,
+    PreparedRequest,
+    resolve_stage_model_path,
+)
+from vllm_omni.model_executor.models.yue2.yue2 import (
+    CONTEXT,
+    FRAMES_PER_SECOND,
+    KEY_MAX_AUDIO_FRAMES,
+    KEY_MIN_TOKENS,
+    KEY_PENALTY_WINDOW,
+    KEY_PHASE,
+    KEY_PREFIX_IDS,
+    KEY_REPETITION_PENALTY,
+    KEY_SEED,
+    KEY_SKIP_SYNTHESIS,
+    KEY_TEMPERATURE,
+    KEY_TOP_K,
+    KEY_TOP_P,
+    SEMANTIC_SAMPLING,
+    STOP_TOKEN_IDS,
+)
+from vllm_omni.tokenizers.yue2_prompt import semantic_prefix_ids
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+_COT_MODES = ("off", "melody", "full")
+
+# Without max_new_tokens the budget is the preset ceiling, so the model's
+# own end token sets the song length (upstream reference behavior).
+_DEFAULT_MAX_FRAMES = int(SEMANTIC_SAMPLING["max_tokens"])
+
+# Sampling is fixed by the checkpoint's reference preset. Accepting these
+# would imply a control the model does not honor.
+_UNSUPPORTED_SAMPLING_PARAMS = ("temperature", "top_p", "top_k", "repetition_penalty")
+
+
+@register_tts_adapter
+class Yue2Adapter(ARTTSAdapter):
+    """Lyrics plus a style caption in, one 48 kHz stereo song out."""
+
+    name = "yue2"
+    stage_keys = frozenset({"yue2"})
+    model_archs = frozenset({"Yue2ForCausalLM"})
+    max_new_tokens_min = int(SEMANTIC_SAMPLING["min_tokens"])
+    max_new_tokens_max = int(SEMANTIC_SAMPLING["max_tokens"])
+
+    def __init__(self, ctx: Any) -> None:
+        super().__init__(ctx)
+        self._cached_tokenizer: Any = None
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        if not request.input or not request.input.strip():
+            return "YuE2 requires non-empty 'input' carrying the lyrics"
+
+        instructions = request.instructions
+        if not isinstance(instructions, str) or not instructions.strip():
+            return (
+                "YuE2 requires 'instructions' describing the music "
+                "(genre, instrumentation, tempo, mood); it is what decides the arrangement"
+            )
+
+        extra: dict[str, Any] = request.extra_params or {}
+        cot = extra.get("cot", "off")
+        if cot not in _COT_MODES:
+            return f"extra_params.cot must be one of {_COT_MODES}; got {cot!r}"
+        abc = extra.get("abc")
+        if cot == "off":
+            if abc is not None:
+                return "extra_params.abc requires cot=melody|full"
+        elif not isinstance(abc, str) or not abc.strip():
+            return (
+                f"cot={cot} requires extra_params.abc carrying an ABC score. "
+                "Model-generated ABC is two engine requests and is not served "
+                "in v1; use examples/offline_inference/yue2 for that."
+            )
+
+        if request.max_new_tokens is not None:
+            if request.max_new_tokens < self.max_new_tokens_min:
+                return f"'max_new_tokens' must be at least {self.max_new_tokens_min}"
+            if request.max_new_tokens > self.max_new_tokens_max:
+                return (
+                    f"'max_new_tokens' counts audio frames at {FRAMES_PER_SECOND} per second "
+                    f"and cannot exceed {self.max_new_tokens_max}"
+                )
+
+        voice = (request.voice or "").strip().lower()
+        if voice not in ("", "default"):
+            return (
+                "YuE2 has no speaker to select; the vocal comes from "
+                f"'instructions'. Omit 'voice' or pass 'default'; got {request.voice!r}"
+            )
+        if request.ref_audio is not None or request.ref_text is not None:
+            return "YuE2 does not support reference-audio conditioning"
+        if request.language is not None:
+            return "YuE2 has no language tag; the language follows the lyrics"
+        if request.task_type is not None:
+            return "YuE2 does not support 'task_type'"
+        if request.speed is not None and float(request.speed) != 1.0:
+            return "YuE2 only supports speed=1.0; put the tempo in 'instructions', e.g. 'at 120 BPM'"
+        if request.is_streaming():
+            return (
+                "YuE2 does not support streaming: the song is decoded in one "
+                "terminal NAR/VAE pass after AR decoding finishes. Use stream=false."
+            )
+
+        rejected = sorted(key for key in _UNSUPPORTED_SAMPLING_PARAMS if extra.get(key) is not None)
+        if rejected:
+            return "YuE2 has fixed sampling and does not accept: " + ", ".join(rejected)
+        return None
+
+    async def warmup(self) -> None:
+        # Resolve (and possibly download) qwen.tiktoken at startup: build()
+        # runs on the request path and must not block on a hub download, and a
+        # bad path should surface before the first request.
+        self._tokenizer()
+
+    async def build(
+        self,
+        request: "OpenAICreateSpeechRequest",
+        sampling_params_list: list,
+        has_inline_ref_audio: bool,
+    ) -> PreparedRequest:
+        del sampling_params_list, has_inline_ref_audio
+        extra: dict[str, Any] = request.extra_params or {}
+        cot = extra.get("cot", "off")
+        encode = self._tokenizer().encode
+        abc_ids = encode(extra["abc"]) if cot != "off" else None
+        # validate() has already rejected blank lyrics/caption.
+        instructions = request.instructions or ""
+        lyrics = request.input or ""
+        prefix = semantic_prefix_ids(encode, instructions, lyrics, cot, abc_ids=abc_ids)
+        max_frames = int(request.max_new_tokens or _DEFAULT_MAX_FRAMES)
+        if len(prefix) + max_frames > CONTEXT:
+            return_err = (
+                f"YuE2 prompt is {len(prefix)} tokens and the {max_frames}-frame budget "
+                f"exceeds the {CONTEXT}-token context. Shorten the lyrics/style or "
+                "lower max_new_tokens."
+            )
+            raise ValueError(return_err)
+        return PreparedRequest(
+            prompt={"prompt_token_ids": prefix},
+            tts_params={"max_audio_frames": [max_frames]},
+            model_type=self.name,
+        )
+
+    def apply_sampling_overrides(
+        self,
+        sampling_params_list: list,
+        request: "OpenAICreateSpeechRequest",
+        prompt: dict[str, Any] | None = None,
+        request_id: str | None = None,
+    ) -> list:
+        import copy
+
+        sampling_params_list = copy.deepcopy(sampling_params_list)
+        params = sampling_params_list[0]
+        max_frames = int(request.max_new_tokens or _DEFAULT_MAX_FRAMES)
+        # A service must not give every caller the same song; without an
+        # explicit seed, draw one per request.
+        seed = request.seed if request.seed is not None else secrets.randbelow(2**31)
+        assert prompt is not None  # the engine always passes the scheduled prompt
+        params.extra_args = {
+            **(params.extra_args or {}),
+            KEY_PHASE: "semantic",
+            KEY_SEED: seed,
+            KEY_TEMPERATURE: SEMANTIC_SAMPLING["temperature"],
+            KEY_TOP_P: SEMANTIC_SAMPLING["top_p"],
+            KEY_TOP_K: SEMANTIC_SAMPLING["top_k"],
+            KEY_REPETITION_PENALTY: SEMANTIC_SAMPLING["repetition_penalty"],
+            KEY_PENALTY_WINDOW: SEMANTIC_SAMPLING["penalty_window"],
+            KEY_MIN_TOKENS: SEMANTIC_SAMPLING["min_tokens"],
+            KEY_MAX_AUDIO_FRAMES: max_frames,
+            KEY_SKIP_SYNTHESIS: False,
+            # Full prompt ids: under a KV prefix-cache hit the engine schedules
+            # only the uncached tail, so the model cannot rebuild its NAR
+            # conditioning prefix from the scheduled tokens alone.
+            KEY_PREFIX_IDS: list(prompt["prompt_token_ids"]),
+        }
+        # The model needs one decode step beyond the last frame to deliver the
+        # terminal NAR/VAE pass, so do not cap the engine at exactly max_frames.
+        params.max_tokens = max_frames + 1
+        params.stop_token_ids = list(STOP_TOKEN_IDS)
+        params.detokenize = False
+        return sampling_params_list
+
+    @staticmethod
+    def _last_meta_flag(audio_output: Mapping[str, Any], name: str) -> Any:
+        """Read a per-request 0/1 meta flag in all three wire shapes.
+
+        The runner flattens nested meta.* to dotted keys (flatten_payload)
+        and the wire payload is tensor-only (_ensure_tensor_values), so by
+        serving time the flag is usually a 0-d "0"/"1" tensor under a dotted
+        key; accept the nested/offline shapes too.
+        """
+        flags = audio_output.get(f"meta.{name}")
+        if flags is None:
+            meta = audio_output.get("meta")
+            if isinstance(meta, dict):
+                flags = meta.get(name)
+        if isinstance(flags, torch.Tensor):
+            flags = flags.reshape(-1)[-1].item() if flags.numel() else None
+        elif isinstance(flags, (list, tuple)):
+            flags = flags[-1] if flags else None
+        return flags
+
+    def collect_response_metadata(self, audio_output: Mapping[str, Any], collect: dict) -> None:
+        """Surface meta.truncated / meta.error as response headers."""
+        truncated = self._last_meta_flag(audio_output, "truncated")
+        if truncated is not None:
+            collect["audio_truncated"] = truncated is True or str(truncated) == "1"
+        # meta.error=1 marks a failed NAR/VAE pass (empty clip); serving turns
+        # it into a 500 via this generic key. Absent/0 means success.
+        error = self._last_meta_flag(audio_output, "error")
+        if error is not None and (error is True or str(error) == "1"):
+            collect["audio_synthesis_error"] = True
+
+    def _tokenizer(self) -> Any:
+        """The checkpoint-native tiktoken BPE, loaded once per process."""
+        if self._cached_tokenizer is None:
+            from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer
+
+            model_path = resolve_stage_model_path(self.ctx.engine_client)
+            if model_path is None:
+                raise RuntimeError("YuE2 tokenizer needs a resolvable stage model path")
+            merge_file = Path(model_path).expanduser() / "qwen.tiktoken"
+            if not merge_file.is_file():
+                # model_path is an HF repo id: pull just the merge file into
+                # the hub cache (vae.py does the same for the VAE snapshot).
+                from vllm_omni.transformers_utils.repo_utils import hf_api
+
+                merge_file = Path(hf_api().hf_hub_download(model_path, "qwen.tiktoken"))
+            self._cached_tokenizer = YuE2TextTokenizer(merge_file)
+        return self._cached_tokenizer
+
+
+__all__ = ["Yue2Adapter"]

vllm_omni/model_executor/models/registry.py

@@ -483,6 +483,12 @@ _OMNI_MODELS = {
         "acoustic",
         "MiniMaxMusic3AcousticForConditionalGeneration",
     ),
+    ## YuE2-3B (text-to-music; single-stage native-AR with terminal NAR/VAE)
+    "Yue2ForCausalLM": (
+        "yue2",
+        "yue2",
+        "Yue2ForCausalLM",
+    ),
 }
 
 

vllm_omni/model_executor/models/yue2/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""YuE2-3B text-to-music model package."""
+
+from .pipeline import YUE2_PIPELINE
+
+__all__ = ["YUE2_PIPELINE"]

vllm_omni/model_executor/models/yue2/pipeline.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""YuE2-3B pipeline topology.
+
+Single-stage native-AR text-to-music (gepard pattern). The whole
+Mixture-of-Transformers checkpoint — AR path, NAR path, projection heads —
+lives in one ``model.safetensors``, so both the autoregressive sampling and
+the terminal flow-matching + VAE pass run inside one stage and no weights are
+duplicated across stages. The VAE decoder is a separate repository
+(``m-a-p/YuE2-Vae``), loaded in ``load_weights`` from ``$YUE2_VAE`` or the
+default hub id.
+
+Sampling is model-owned. ``sampling_constraints`` pins ``detokenize`` and the
+union of both phase end tokens; ``extra_args`` carry the phase and the
+request-local sampling preset, and must be present in the stage's
+``default_sampling_params`` so ``has_sampling_extra_args`` turns on and the
+per-request args reach the model's forward.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+from .yue2 import STOP_TOKEN_IDS
+
+YUE2_PIPELINE = PipelineConfig(
+    model_type="yue2",
+    default_deploy_config_name="yue2.yaml",
+    model_arch="Yue2ForCausalLM",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="yue2",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="audio",
+            owns_tokenizer=True,
+            engine_output_type="audio",
+            # Whole-song audio ships once, on the finishing step; the audio
+            # accumulator concatenates over time and a single tensor just
+            # passes through.
+            sampling_constraints={
+                "detokenize": False,
+                "stop_token_ids": list(STOP_TOKEN_IDS),
+            },
+        ),
+    ),
+)
+
+__all__ = ["YUE2_PIPELINE"]

vllm_omni/tokenizers/yue2_prompt.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Prompt construction for YuE2-3B, mirroring upstream ``protocol.py``.
+
+The driver (offline example or, later, the serving adapter) builds token-id
+prefixes here and submits them as ``prompt_token_ids``. Both generation phases
+share one prefix head (instruction + tags + lyrics); the abc phase continues
+from ``ABC_START`` and stops at ``ABC_END``, the semantic phase embeds the ABC
+score (generated or user-supplied) and continues from ``MUSIC_START``.
+"""
+
+from __future__ import annotations
+
+from vllm_omni.model_executor.models.yue2.yue2 import (
+    ABC_END,
+    ABC_START,
+    CODEC_OFFSET,
+    EOD,
+    MUSIC_END,
+    MUSIC_START,
+)
+
+INSTRUCTIONS = {
+    "off": "Generate music with codec tokens from the given conditions.",
+    "melody": (
+        "Generate a melody-only ABC transcription without chord symbols, then "
+        "generate music with codec tokens from the given conditions."
+    ),
+    "full": (
+        "Generate a chord-annotated ABC transcription, then generate music with codec tokens from the given conditions."
+    ),
+}
+
+
+def song_text(style: str, lyrics: str, cot: str) -> str:
+    return f"{INSTRUCTIONS[cot]}\n[Tags]\n{style}\n[Lyrics]\n{lyrics}\n"
+
+
+def abc_prefix_ids(encode, style: str, lyrics: str, cot: str) -> list[int]:
+    """Prompt for the abc phase (cot=full/melody, no external ABC)."""
+    if cot == "off":
+        raise ValueError("cot=off skips the abc phase entirely")
+    return [EOD] + encode(song_text(style, lyrics, cot)) + [ABC_START]
+
+
+def semantic_prefix_ids(
+    encode,
+    style: str,
+    lyrics: str,
+    cot: str,
+    abc_ids: list[int] | None = None,
+) -> list[int]:
+    """Prompt for the semantic phase.
+
+    ``abc_ids`` are the token ids of the (generated or user-supplied) ABC
+    score. ``cot=off`` passes ``abc_ids=None`` and skips the span, matching
+    upstream ``token_prefixes``.
+    """
+    base = [EOD] + encode(song_text(style, lyrics, cot))
+    if cot == "off":
+        return base + [ABC_START, ABC_END, MUSIC_START]
+    if abc_ids is None:
+        raise ValueError("cot=melody/full requires abc ids")
+    if any(not 0 <= t < EOD for t in abc_ids):
+        raise ValueError("ABC ids must stay inside the ordinary text vocabulary")
+    return base + [ABC_START] + list(abc_ids) + [ABC_END, MUSIC_START]
+
+
+def abc_ids_from_generated(generated: list[int]) -> list[int]:
+    """ABC token ids generated by the abc phase (end token stripped)."""
+    ids = list(generated)
+    if ids and ids[-1] == ABC_END:
+        ids.pop()
+    if ABC_START in ids:
+        ids = ids[ids.index(ABC_START) + 1 :]
+    if any(t >= EOD for t in ids):
+        raise ValueError("abc phase produced tokens outside the text vocabulary")
+    return ids
+
+
+def semantic_frames(generated: list[int]) -> list[int]:
+    """Codec frame values (0..CODEC_SIZE-1) from semantic-phase output."""
+
+    ids = list(generated)
+    if ids and ids[-1] == MUSIC_END:
+        ids.pop()
+    return [t - CODEC_OFFSET for t in ids]
+
+
+__all__ = [
+    "INSTRUCTIONS",
+    "abc_ids_from_generated",
+    "abc_prefix_ids",
+    "semantic_frames",
+    "semantic_prefix_ids",
+    "song_text",
+]

vllm_omni/tokenizers/yue2_tokenizer.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Minimal YuE2 text tokenizer (frozen tiktoken BPE), vendored from upstream.
+
+Used by the offline driver to build token-id prompts; the audio codec tokens
+never pass through it.
+"""
+
+from __future__ import annotations
+
+import unicodedata
+from pathlib import Path
+
+import pybase64 as base64
+
+
+class YuE2TextTokenizer:
+    def __init__(self, merge_file):
+        import tiktoken
+
+        self.merge_file = Path(merge_file)
+        ranks = {
+            base64.b64decode(t): int(r)
+            for t, r in (line.split() for line in self.merge_file.read_bytes().splitlines() if line)
+        }
+        if len(ranks) != 151643:
+            raise ValueError("Expected checkpoint-native qwen.tiktoken (151643 ordinary tokens)")
+        specials = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "<R>", "<S>", "<X>", "<mask>", "<sep>"]
+        specials += [f"<extra_{i}>" for i in range(200)]
+        specials[204:206] = ["<abc>", "</abc>"]
+        pattern = (
+            r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}"
+            r"| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
+        )
+        self._enc = tiktoken.Encoding(
+            "YuE2",
+            pat_str=pattern,
+            mergeable_ranks=ranks,
+            special_tokens={s: i + len(ranks) for i, s in enumerate(specials)},
+        )
+
+    def encode(self, text):
+        return self._enc.encode_ordinary(unicodedata.normalize("NFC", text))
+
+    def decode(self, ids):
+        return self._enc.decode([int(i) for i in ids if 0 <= i < self._enc.n_vocab], errors="replace")
+
+
+__all__ = ["YuE2TextTokenizer"]

vllm_omni/transformers_utils/configs/__init__.py

@@ -26,6 +26,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
     "VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts",
     "CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3",
     "MiniMaxMusic3Config": "vllm_omni.transformers_utils.configs.minimax_music3",
+    "Yue2Config": "vllm_omni.transformers_utils.configs.yue2",
     "OmniVoiceConfig": "vllm_omni.transformers_utils.configs.omnivoice",
     "BailingMoeV2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni",
     "BailingMM2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni",
@@ -56,6 +57,7 @@ __all__ = [
     "VoxtralTTSConfig",
     "CosyVoice3Config",
     "MiniMaxMusic3Config",
+    "Yue2Config",
     "OmniVoiceConfig",
     "BailingMoeV2Config",
     "BailingMM2Config",
@@ -96,3 +98,4 @@ from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice  # noqa
 from vllm_omni.transformers_utils.configs import sensenova_u1 as _sensenova_u1  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import voxcpm2 as _voxcpm2  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import voxtral_tts as _voxtral_tts  # noqa: F401, E402
+from vllm_omni.transformers_utils.configs import yue2 as _yue2  # noqa: F401, E402

vllm_omni/transformers_utils/configs/yue2.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""HuggingFace config for m-a-p/YuE2-3B.
+
+The checkpoint's own ``config.json`` already carries every field the AR
+backbone needs (it is Qwen3-1.7B-shaped with q_norm/k_norm and an extended
+184,704-entry vocabulary), plus the acoustic-side fields (``latent_dim``,
+``max_latent_frames``, ``timestep_shift``). Registering this class lets vLLM
+build a ``ModelConfig`` without ``trust_remote_code``: the package's own
+``Yue2ForCausalLM`` replaces the repository's auto_map modeling code, and the
+deploy YAML pins ``hf_overrides.architectures`` to it.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from transformers import AutoConfig
+from transformers.configuration_utils import PretrainedConfig
+
+
+class Yue2Config(PretrainedConfig):
+    """Config for the YuE2-3B AR–NAR Mixture-of-Transformers checkpoint."""
+
+    model_type = "yue2"
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_hidden_layers: int = 28,
+        num_attention_heads: int = 16,
+        num_key_value_heads: int = 8,
+        head_dim: int = 128,
+        intermediate_size: int = 6144,
+        vocab_size: int = 184704,
+        rms_norm_eps: float = 1e-6,
+        rope_theta: float = 1000000.0,
+        max_position_embeddings: int = 24576,
+        tie_word_embeddings: bool = False,
+        latent_dim: int = 64,
+        max_latent_frames: int = 24576,
+        timestep_shift: float = 1.0,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.intermediate_size = intermediate_size
+        self.vocab_size = vocab_size
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.latent_dim = latent_dim
+        self.max_latent_frames = max_latent_frames
+        self.timestep_shift = timestep_shift
+        # The checkpoint's config.json (like upstream's qwen_config derivative,
+        # fast.py) omits the standard Qwen3 fields the AR backbone expects;
+        # synthesize the same defaults upstream does so vLLM's Qwen3Model and
+        # MLP construct without AttributeError.
+        self.hidden_act = str(kwargs.pop("hidden_act", "silu"))
+        self.attention_bias = bool(kwargs.pop("attention_bias", False))
+        self.attention_dropout = float(kwargs.pop("attention_dropout", 0.0))
+        self.use_sliding_window = bool(kwargs.pop("use_sliding_window", False))
+        self.sliding_window = kwargs.pop("sliding_window", None)
+        self.max_window_layers = int(kwargs.pop("max_window_layers", 0))
+
+
+AutoConfig.register("yue2", Yue2Config)
+
+__all__ = ["Yue2Config"]