diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 5b8bf3464..336b41d8b 100644 --- a/docs/models/supported_models.md +++ b/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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | diff --git a/examples/offline_inference/yue2/end2end.py b/examples/offline_inference/yue2/end2end.py new file mode 100644 index 000000000..39fddd371 --- /dev/null +++ b/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: .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() diff --git a/recipes/m-a-p/YuE2-3B.md b/recipes/m-a-p/YuE2-3B.md new file mode 100644 index 000000000..8ad3b008d --- /dev/null +++ b/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. diff --git a/tests/assets/yue2/twinkle_golden.abc b/tests/assets/yue2/twinkle_golden.abc new file mode 100644 index 000000000..9d2ec067e --- /dev/null +++ b/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| diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index b66c930fc..6f004b6b6 100644 --- a/tests/config/test_environment_variables.py +++ b/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, diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py index 98b4184d9..6cf93526a 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py +++ b/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} + ) diff --git a/tests/e2e/offline_inference/test_yue2.py b/tests/e2e/offline_inference/test_yue2.py new file mode 100644 index 000000000..996f42e9f --- /dev/null +++ b/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 diff --git a/tests/e2e/online_serving/test_yue2.py b/tests/e2e/online_serving/test_yue2.py new file mode 100644 index 000000000..35b25b1b0 --- /dev/null +++ b/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}" + ) diff --git a/tests/e2e/online_serving/yue2_metal_verify.py b/tests/e2e/online_serving/yue2_metal_verify.py new file mode 100644 index 000000000..8c7553e73 --- /dev/null +++ b/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() diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 6504359e4..dc49a7e21 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/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() diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index 3ef55f673..ca27c6e37 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/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"] diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index ef9712eb9..79d8f610d 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/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" diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index aa94a6879..69d9306a7 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/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 = ( diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 686ef1bf0..f016b7f7d 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/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, } diff --git a/vllm_omni/deploy/yue2.yaml b/vllm_omni/deploy/yue2.yaml new file mode 100644 index 000000000..502d950c5 --- /dev/null +++ b/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 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index b6c495a64..3a9520574 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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: diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index f98fa87df..d6e7ff186 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/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__ = [ diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index b2297d826..6ad7dcdcc 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/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 diff --git a/vllm_omni/entrypoints/openai/tts_adapters/yue2.py b/vllm_omni/entrypoints/openai/tts_adapters/yue2.py new file mode 100644 index 000000000..bcacc34e7 --- /dev/null +++ b/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"] diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index a0a812343..ca0fba7f8 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/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", + ), } diff --git a/vllm_omni/model_executor/models/yue2/__init__.py b/vllm_omni/model_executor/models/yue2/__init__.py new file mode 100644 index 000000000..9db16ee2d --- /dev/null +++ b/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"] diff --git a/vllm_omni/model_executor/models/yue2/pipeline.py b/vllm_omni/model_executor/models/yue2/pipeline.py new file mode 100644 index 000000000..0a1fa51bd --- /dev/null +++ b/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"] diff --git a/vllm_omni/model_executor/models/yue2/yue2.py b/vllm_omni/model_executor/models/yue2/yue2.py new file mode 100644 index 000000000..02549cf84 --- /dev/null +++ b/vllm_omni/model_executor/models/yue2/yue2.py @@ -0,0 +1,1750 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""YuE2-3B text-to-music: single-stage native-AR model, ported from upstream +YuE2 (bd90e4c, Apache-2.0). + +One file per the single-stage layout (see auk/): constants, the AR model with +its model-owned sampler, the terminal NAR flow-matching pass, weight routing +and the tiled VAE decoder all live here because they all run inside stage 0. +The text tokenizer and token-id prompt construction live in +``vllm_omni/tokenizers/`` (driver-side only, never used by the model). +""" + +from __future__ import annotations + +import json +import math +import os +import time +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass, field +from numbers import Integral +from pathlib import Path +from typing import Any, Literal + +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn.utils import weight_norm +from transformers import PretrainedConfig, PreTrainedModel +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix +from vllm.v1.outputs import SamplerOutput + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +# -------------------- token/protocol constants -------------------- + + +"""Token/protocol constants for YuE2-3B, pinned by the checkpoint. + +Values mirror the upstream ``yue2/protocol.py`` (bd90e4c) exactly; changing any +of them desynchronizes generation from the reference implementation. The model +is single-codebook: one codec token per latent frame, 25 frames per second, +so ``max_tokens`` in engine terms equals frames/25 seconds of music. +""" + +EOD = 151643 +ABC_START, ABC_END = 151847, 151848 +MUSIC_START, MUSIC_END = 151851, 151852 +CODEC_OFFSET, CODEC_SIZE = 151853, 32768 +LATENT_START, LATENT_END, LATENT_PAD = 184621, 184622, 184623 +VOCAB_SIZE, CONTEXT = 184704, 24576 + +# Engine-side stop ids: the union of both phase ends. The model's own sampler +# only ever emits the end token of the request's current phase, so a request +# never stops on the other phase's end by accident. +STOP_TOKEN_IDS = [ABC_END, MUSIC_END] + +SAMPLE_RATE = 48000 +AUDIO_CHANNELS = 2 +FRAMES_PER_SECOND = 25 +SAMPLES_PER_FRAME = 1920 +LATENT_DIM = 64 + +# Reference sampling presets (upstream yue2_generation_config.json). +ABC_SAMPLING = { + "temperature": 0.7, + "top_p": 0.9, + "top_k": 30, + "repetition_penalty": 1.005, + "penalty_window": 100, + "min_tokens": 32, + "max_tokens": 4096, +} +SEMANTIC_SAMPLING = { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 100, + "repetition_penalty": 1.2, + "penalty_window": 50, + "min_tokens": 200, + "max_tokens": 9000, +} + +ODE_STEPS = 32 +ODE_METHOD = "midpoint" +VAE_CORE_FRAMES = 1024 +VAE_HALO_FRAMES = 16 +DEFAULT_VAE_ID = "m-a-p/YuE2-Vae" + +# extra_args keys the model reads per request. +KEY_PHASE = "yue2_phase" # "abc" | "semantic" +KEY_TEMPERATURE = "yue2_temperature" +KEY_TOP_P = "yue2_top_p" +KEY_TOP_K = "yue2_top_k" +KEY_REPETITION_PENALTY = "yue2_repetition_penalty" +KEY_PENALTY_WINDOW = "yue2_penalty_window" +KEY_MIN_TOKENS = "yue2_min_tokens" +KEY_MAX_AUDIO_FRAMES = "yue2_max_audio_frames" +KEY_SEED = "yue2_seed" +KEY_SKIP_SYNTHESIS = "yue2_skip_synthesis" # abc phase: tokens only, no audio +# Full prompt token ids, shipped by the driver. The NAR conditioning needs the +# whole prefix, but under a KV prefix-cache hit the engine schedules only the +# uncached tail, so the scheduled input_ids slice cannot rebuild it. +KEY_PREFIX_IDS = "yue2_prefix_ids" + + +# -------------------- checkpoint weight routing -------------------- + + +"""AR/side routing for YuE2's single-file MoT checkpoint. + +The checkpoint interleaves AR-path, NAR-path and projection tensors under one +namespace. AR tensors go through the vLLM backbone loader unchanged; NAR and +projection tensors are hand-loaded, with ``model.layers.N.nar_*`` names +remapped onto the ``nar_layers.N.*`` module layout. +""" + +_SIDE_TOP_LEVEL = {"vae2llm", "llm2vae", "time_embedder"} + + +def partition_checkpoint_weights( + weights: Iterable[tuple[str, torch.Tensor]], +) -> tuple[list[tuple[str, torch.Tensor]], list[tuple[str, torch.Tensor]]]: + """Split checkpoint tensors into (backbone AR pairs, remapped side pairs).""" + ar: list[tuple[str, torch.Tensor]] = [] + side: list[tuple[str, torch.Tensor]] = [] + for name, tensor in weights: + if name == "latent_pos_embed.pe": + continue # deterministic sinusoid, rebuilt in __init__ + if ".nar_" in name or name.split(".", 1)[0] in _SIDE_TOP_LEVEL: + new = name + if name.startswith("model.layers."): + rest = name[len("model.layers.") :] + layer_no, _, tail = rest.partition(".") + if tail.startswith("nar_"): + tail = tail[len("nar_") :] + new = f"nar_layers.{layer_no}.{tail}" + side.append((new, tensor)) + continue + ar.append((name, tensor)) + return ar, side + + +# -------------------- model-owned sampler -------------------- + + +"""Model-owned sampling for YuE2-3B, ported from upstream ``sampling.py``. + +The model claims ``prefer_model_sampler`` and reproduces the reference +request-local arithmetic exactly: phase masking (abc = text vocabulary + its +end; semantic = codec span + its end), windowed repetition penalty over the +request's own history, then temperature/top-k/top-p with a seeded multinomial. +Default CFG is off (guidance 1.0 for full/melody, 1.01 for off); the upstream +torch backend treats those the same, so Phase 1 samples a single row. +""" + + +def window_penalty(logits: torch.Tensor, recent_ids: list[int], penalty: float) -> torch.Tensor: + """Upweight/downweight ids seen in the window, upstream arithmetic.""" + if penalty == 1.0 or not recent_ids: + return logits + recent = torch.as_tensor(recent_ids, dtype=torch.long, device=logits.device) + if logits.dim() > 1: + recent = recent.reshape(1, -1) + freq = torch.zeros_like(logits) + freq.scatter_add_(-1, recent, torch.ones_like(recent, dtype=logits.dtype)) + alpha = penalty**freq + return torch.where(logits < 0, logits * alpha, logits / alpha) + + +def distribution( + logits: torch.Tensor, + *, + temperature: float, + top_p: float, + top_k: int, + repetition_penalty: float, + penalty_window: int, + history: list[int], + step: int, + min_tokens: int, + phase: str, +) -> torch.Tensor: + """Masked/penalized/shaped scores for one row (float32 path). + + Accepts either one row ``[vocab]`` (what the model's sampler sees) or a + batch ``[rows, vocab]`` and returns the same shape it was given. + """ + single = logits.dim() == 1 + logits = logits.unsqueeze(0) if single else logits + scores = logits.float().clone() + end = ABC_END if phase == "abc" else MUSIC_END + allowed = torch.full_like(scores, float("-inf")) + if phase == "abc": + allowed[..., :EOD] = 0 + else: + allowed[..., CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0 + allowed[..., end] = 0 + scores = scores + allowed + if step < min_tokens: + scores[..., end] = -torch.inf + scores = window_penalty(scores, history[-penalty_window:], repetition_penalty) + if temperature == 0: + return scores.squeeze(0) if single else scores + if temperature != 1: + scores = scores / temperature + threshold = scores.topk(min(top_k, scores.shape[-1])).values[..., -1, None] + scores = scores.masked_fill(scores < threshold, -torch.inf) + if top_p < 1: + values, indices = scores.sort(descending=True) + probabilities = values.softmax(-1) + removed = probabilities.cumsum(-1) - probabilities > top_p + removed[..., :1] = False + values = values.masked_fill(removed, -torch.inf) + scores = values.scatter(-1, indices, values) + return scores.squeeze(0) if single else scores + + +def sample_row(scores: torch.Tensor, generator: torch.Generator, *, greedy: bool = False) -> int: + """Draw one token id from prepared scores with the request's generator.""" + if not torch.isfinite(scores).any(): + raise RuntimeError("sampling scores are all -inf; check phase masking") + if greedy: + return int(scores.argmax().item()) + probabilities = scores.softmax(-1) + next_id = torch.multinomial(probabilities, 1, generator=generator) + return int(next_id.item()) + + +# -------------------- NAR flow-matching pass -------------------- + + +"""Acoustic flow matching for YuE2-3B, ported from upstream ``nar.py``. + +One AR prefill per original chunk, then 32 midpoint ODE steps that walk only +the NAR path. The AR prefill re-uses the vLLM backbone's own projections +(fused ``qkv_proj`` split by head counts, ``q_norm``/``k_norm``, the layer's +``o_proj``/``mlp``), with RoPE and SDPA computed here exactly like the +reference: NAR positions attend the full AR prefix plus all NAR positions +bidirectionally, which a per-chunk prefill + concatenated K/V reproduces. +""" + + +def chunk_ranges(frames: int, prefix_tokens: int, context: int = CONTEXT): + size = min((context - prefix_tokens - 3) // 2, CONTEXT) + if frames < 1 or size < 1: + raise ValueError("Empty codec or prefix leaves no acoustic context") + return [(a, min(a + size, frames)) for a in range(0, frames, size)] + + +@dataclass +class Chunk: + ar_tokens: list[int] + noise: torch.Tensor + + +def song_chunks(prefix: Sequence[int], codec: Sequence[int], seed: int, context: int = CONTEXT): + """Draw the whole-song noise once, then cut it at the historical boundaries.""" + prefix = [int(v) for v in prefix] + codec = [int(v) for v in codec] + if min(prefix) < 0 or not codec or min(codec) < 0 or max(codec) >= CODEC_SIZE: + raise ValueError("Token IDs are outside their allowed vocabulary") + ranges = chunk_ranges(len(codec), len(prefix), int(context)) + generator = torch.Generator(device="cpu").manual_seed(int(seed)) + noise = torch.randn((len(codec), 64), dtype=torch.float32, device="cpu", generator=generator) + return [ + Chunk( + prefix + [value + CODEC_OFFSET for value in codec[a:b]] + [MUSIC_END], + noise[a:b], + ) + for a, b in ranges + ] + + +def _linear(module, x): + out = module(x) + return out[0] if isinstance(out, tuple) else out + + +def _attention(q, k, v, *, causal: bool, query_chunk_size: int | None = None): + """SDPA over [tokens, heads, dim]; query tiling only bounds temporaries.""" + if q.ndim != 3 or k.ndim != 3 or v.shape != k.shape or q.shape[-1] != k.shape[-1]: + raise ValueError("Expected Q/K/V [tokens, heads, dim] with matching K/V") + if min(q.shape) < 1 or min(k.shape) < 1 or q.shape[1] % k.shape[1]: + raise ValueError("Invalid attention lengths or grouped-query head count") + if causal and len(q) != len(k): + raise ValueError("Causal prefill requires matching Q/K sequence lengths") + block = query_chunk_size or (len(q) if q.device.type == "cuda" else 256) + query = q.transpose(0, 1).unsqueeze(0) + key = k.transpose(0, 1).unsqueeze(0) + value = v.transpose(0, 1).unsqueeze(0) + grouped = query.shape[1] != key.shape[1] + outputs = [] + for start in range(0, len(q), block): + end = min(start + block, len(q)) + used_key = key[..., :end, :] if causal else key + used_value = value[..., :end, :] if causal else value + mask = None + if causal and start: + mask = torch.arange(end, device=q.device)[None, :] <= torch.arange(start, end, device=q.device)[:, None] + outputs.append( + F.scaled_dot_product_attention( + query[..., start:end, :], + used_key, + used_value, + attn_mask=mask, + is_causal=causal and start == 0, + enable_gqa=grouped, + ) + ) + return torch.cat(outputs, dim=-2)[0].transpose(0, 1) + + +class CachedNAR: + """One acoustic chunk: AR prefix K/V prefilled once, ODE walks the NAR path.""" + + def __init__(self, model, chunk: Chunk, query_chunk_size: int | None = None): + self.model = model + self.chunk = chunk + self.query_chunk_size = query_chunk_size + weight = next(model.vae2llm.parameters()) + self.device, self.dtype = weight.device, weight.dtype + if chunk.noise.ndim != 2 or chunk.noise.shape[1] != 64 or len(chunk.noise) < 1: + raise ValueError("Expected nonempty acoustic noise [frames,64]") + self.ar_length, self.nar_length = len(chunk.ar_tokens), len(chunk.noise) + 2 + if self.ar_length < 1 or max(chunk.ar_tokens) >= model.vocab_size: + raise ValueError("AR prefix is empty or outside the model vocabulary") + if self.ar_length + self.nar_length > model.max_position_embeddings: + raise ValueError("Acoustic chunk exceeds the model context") + positions = torch.arange(self.ar_length, self.ar_length + self.nar_length, device=self.device)[None] + self.cos, self.sin = model.rotary(positions) + local = torch.arange(self.nar_length, device=self.device).clamp(max=model.max_latent_frames - 1) + self.pos_emb = model.latent_pos_embed(local)[None] + self.cache: list[tuple[torch.Tensor, torch.Tensor]] = [] + self._prefill() + + def _prefill(self): + backbone = self.model.model + ids = torch.tensor([self.chunk.ar_tokens], dtype=torch.long, device=self.device) + positions = torch.arange(self.ar_length, device=self.device)[None] + cos, sin = self.model.rotary(positions) + x = backbone.embed_tokens(ids)[0] # [T, H] + for layer in backbone.layers: + # Upstream: layer.self_attn.project_qkv(layer.input_layernorm(x), ...) + # The layernorm is load-bearing: without it the cached AR K/V + # conditioning is garbage and the ODE latents blow up (clipped audio). + q, k, v = self.model.ar_project(layer, layer.input_layernorm(x), cos, sin) + self.cache.append((k, v)) + h = _attention(q, k, v, causal=True, query_chunk_size=self.query_chunk_size) + x = x + _linear(layer.self_attn.o_proj, h.flatten(1)) + x = x + _linear(layer.mlp, layer.post_attention_layernorm(x)) + + @torch.inference_mode() + def velocity(self, state: torch.Tensor, raw_t: float) -> torch.Tensor: + model = self.model + x_nar = F.pad(state, (0, 0, 1, 1)) + shifted = model.shift_t(raw_t, self.device, self.dtype) + x = model.vae2llm(x_nar[None]) + x = x + model.time_embedder(shifted.expand(self.nar_length))[None] + x = x + self.pos_emb + for layer, (ar_k, ar_v) in zip(model.nar_layers, self.cache): + q, k, v = layer.self_attn.project_qkv(layer.input_layernorm(x), self.cos, self.sin) + k, v = torch.cat((ar_k, k[0])), torch.cat((ar_v, v[0])) + h = _attention(q[0], k, v, causal=False, query_chunk_size=self.query_chunk_size) + x = x + _linear(layer.self_attn.o_proj, h.flatten(1)[None]) + x = x + layer.mlp(layer.pre_mlp_layernorm(x)) + return model.llm2vae(model.model.norm(x))[0, 1:-1] + + @torch.inference_mode() + def solve(self, steps: int = 32) -> torch.Tensor: + if isinstance(steps, bool) or not isinstance(steps, Integral) or steps < 1: + raise ValueError("steps must be a positive integer") + state = self.chunk.noise.to(device=self.device, dtype=self.dtype) + dt = 1.0 / steps + for step in range(steps): + t = 1.0 - step * dt + raw = torch.logit(torch.tensor(t, dtype=torch.float64, device="cpu")).clamp(-20, 20).item() + first = self.velocity(state, raw) + mid = state - first * (dt / 2) + raw_mid = torch.logit(torch.tensor(t - dt / 2, dtype=torch.float64, device="cpu")).clamp(-20, 20).item() + state = state - self.velocity(mid, raw_mid) * dt + result = state.float().cpu() + if not torch.isfinite(result).all(): + raise FloatingPointError("Acoustic flow matching produced non-finite latents") + return result + + def close(self): + self.cache.clear() + self.cos = self.sin = self.pos_emb = None + + +@torch.inference_mode() +def synthesize(model, prefix: Sequence[int], codec: Sequence[int], seed: int, steps: int = 32) -> torch.Tensor: + """CPU FP32 latents [frames,64] for the whole song, chunk by chunk.""" + if model.training: + raise ValueError("synthesize requires model.eval()") + output = [] + for chunk in song_chunks(prefix, codec, seed): + engine = CachedNAR(model, chunk) + try: + output.append(engine.solve(steps)) + finally: + engine.close() + del engine + return torch.cat(output, dim=0) + + +# -------------------- tiled VAE decoder -------------------- + + +"""Portable FP32 YuE2 Oobleck VAE with exact-boundary tiled decoding. + +Oobleck and SnakeBeta derived from stable-audio-tools a6ae0cdf8b2eb1567a4b42ceadddec3712d99d45. +Copyright (c) 2023 Stability AI; Copyright (c) 2022 NVIDIA CORPORATION. +MIT: see THIRD_PARTY_NOTICES.md shipped with this module/model repository. +""" + + +def checkpoint(function, *args, **kwargs): + from torch.utils.checkpoint import checkpoint as torch_checkpoint + + kwargs.setdefault("use_reentrant", False) + return torch_checkpoint(function, *args, **kwargs) + + +def WNConv1d(*args, **kwargs): + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +def WNConvTranspose1d(*args, **kwargs): + return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) + + +def snake_beta(x, alpha, beta): + return x + (1.0 / (beta + 0.000000001)) * torch.pow(torch.sin(x * alpha), 2) + + +class SnakeBeta(nn.Module): + def __init__( + self, + in_features, + alpha=1.0, + alpha_trainable=True, + alpha_logscale=True, + ): + super().__init__() + self.in_features = in_features + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: + self.alpha = nn.Parameter(torch.zeros(in_features) * alpha) + self.beta = nn.Parameter(torch.zeros(in_features) * alpha) + else: + self.alpha = nn.Parameter(torch.ones(in_features) * alpha) + self.beta = nn.Parameter(torch.ones(in_features) * alpha) + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + return snake_beta(x, alpha, beta) + + +def get_activation(activation: Literal["elu", "snake", "none"], channels=None) -> nn.Module: + if activation == "elu": + return nn.ELU() + if activation == "snake": + return SnakeBeta(channels) + if activation == "none": + return nn.Identity() + raise ValueError(f"Unknown activation {activation}") + + +class ResidualUnit(nn.Module): + def __init__(self, in_channels, out_channels, dilation, act_type): + super().__init__() + self.dilation = dilation + padding = (dilation * (7 - 1)) // 2 + self.layers = nn.Sequential( + get_activation(act_type, channels=out_channels), + WNConv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=7, + dilation=dilation, + padding=padding, + ), + get_activation(act_type, channels=out_channels), + WNConv1d( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=1, + ), + ) + + def forward(self, x): + residual = x + if self.training: + x = checkpoint(self.layers, x) + else: + x = self.layers(x) + return x + residual + + +class EncoderBlock(nn.Module): + def __init__(self, in_channels, out_channels, stride, act_type): + super().__init__() + self.layers = nn.Sequential( + ResidualUnit(in_channels, in_channels, 1, act_type), + ResidualUnit(in_channels, in_channels, 3, act_type), + ResidualUnit(in_channels, in_channels, 9, act_type), + get_activation(act_type, channels=in_channels), + WNConv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ), + ) + + def forward(self, x): + return self.layers(x) + + +class DecoderBlock(nn.Module): + def __init__(self, in_channels, out_channels, stride, act_type): + super().__init__() + upsample_layer = WNConvTranspose1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ) + self.layers = nn.Sequential( + get_activation(act_type, channels=in_channels), + upsample_layer, + ResidualUnit(out_channels, out_channels, 1, act_type), + ResidualUnit(out_channels, out_channels, 3, act_type), + ResidualUnit(out_channels, out_channels, 9, act_type), + ) + + def forward(self, x): + return self.layers(x) + + +class OobleckEncoder(nn.Module): + def __init__( + self, + in_channels=2, + channels=128, + latent_dim=32, + c_mults=(1, 2, 4, 8), + strides=(2, 4, 8, 8), + use_snake=False, + antialias_activation=False, + ): + super().__init__() + if antialias_activation: + raise ValueError("The released encoder does not use antialias_activation") + self.in_channels = in_channels + c_mults = [1] + list(c_mults) + self.depth = len(c_mults) + layers = [ + WNConv1d( + in_channels=in_channels, + out_channels=c_mults[0] * channels, + kernel_size=7, + padding=3, + ) + ] + act_type = "snake" if use_snake else "elu" + for i in range(self.depth - 1): + layers.append( + EncoderBlock( + in_channels=c_mults[i] * channels, + out_channels=c_mults[i + 1] * channels, + stride=strides[i], + act_type=act_type, + ) + ) + layers.extend( + [ + get_activation(act_type, channels=c_mults[-1] * channels), + WNConv1d( + in_channels=c_mults[-1] * channels, + out_channels=latent_dim, + kernel_size=3, + padding=1, + ), + ] + ) + self.layers = nn.Sequential(*layers) + + def forward(self, x): + return self.layers(x) + + +class OobleckDecoder(nn.Module): + def __init__( + self, + out_channels=2, + channels=128, + latent_dim=32, + c_mults=(1, 2, 4, 8), + strides=(2, 4, 8, 8), + use_snake=False, + snake_type="vanilla", + antialias_activation=False, + use_nearest_upsample=False, + use_filter=False, + final_tanh=True, + ): + super().__init__() + if antialias_activation or use_nearest_upsample or use_filter: + raise ValueError("Unsupported option for the released decoder") + if use_snake and snake_type != "vanilla": + raise ValueError("The released decoder uses vanilla SnakeBeta") + self.out_channels = out_channels + c_mults = [1] + list(c_mults) + self.depth = len(c_mults) + layers = [ + WNConv1d( + in_channels=latent_dim, + out_channels=c_mults[-1] * channels, + kernel_size=7, + padding=3, + ) + ] + act_type = "snake" if use_snake else "elu" + for i in range(self.depth - 1, 0, -1): + layers.append( + DecoderBlock( + in_channels=c_mults[i] * channels, + out_channels=c_mults[i - 1] * channels, + stride=strides[i - 1], + act_type=act_type, + ) + ) + layers.extend( + [ + get_activation(act_type, channels=c_mults[0] * channels), + WNConv1d( + in_channels=c_mults[0] * channels, + out_channels=out_channels, + kernel_size=7, + padding=3, + bias=False, + ), + nn.Tanh() if final_tanh else nn.Identity(), + ] + ) + self.layers = nn.Sequential(*layers) + + def forward(self, x): + return self.layers(x) + + +class YuE2VAEConfig(PretrainedConfig): + """Configuration shared by YuE2-Vae and YuE2-Vae-legacy. + + Use ``standard`` for listening and ``legacy`` for the paper metric baseline. + """ + + model_type = "yue2_vae" + + _hf_fields = frozenset( + { + "model_type", + "architectures", + "auto_map", + "transformers_version", + "dtype", + "torch_dtype", + "return_dict", + "output_hidden_states", + "output_attentions", + "use_cache", + "tie_word_embeddings", + "torchscript", + "is_decoder", + "is_encoder_decoder", + "add_cross_attention", + "bos_token_id", + "eos_token_id", + "pad_token_id", + "decoder_start_token_id", + "attn_implementation", + } + ) + + def to_dict(self): + return { + key: value + for key, value in super().to_dict().items() + if key in self._hf_fields or key in self._inference_fields + } + + _inference_fields = frozenset( + [ + "encoder_config", + "decoder_config", + "sample_rate", + "latent_dim", + "downsampling_ratio", + "audio_channels", + "release_variant", + "decode_core_frames", + "decode_halo_frames", + ] + ) + + def __init__( + self, + encoder_config=None, + decoder_config=None, + sample_rate=48000, + latent_dim=64, + downsampling_ratio=1920, + audio_channels=2, + release_variant="standard", + decode_core_frames=1024, + decode_halo_frames=16, + **kwargs, + ): + kwargs = {key: value for key, value in kwargs.items() if key in self._hf_fields} + kwargs.setdefault("architectures", ["YuE2VAE"]) + kwargs.setdefault( + "auto_map", + { + "AutoConfig": "modeling_vae.YuE2VAEConfig", + "AutoModel": "modeling_vae.YuE2VAE", + }, + ) + super().__init__(**kwargs) + self.encoder_config = encoder_config or dict( + in_channels=2, + channels=64, + c_mults=[1, 2, 4, 8, 16, 32], + strides=[2, 2, 4, 4, 5, 6], + latent_dim=128, + use_snake=True, + ) + self.decoder_config = decoder_config or dict( + out_channels=2, + channels=64, + c_mults=[1, 2, 4, 8, 16, 32], + strides=[2, 2, 4, 4, 5, 6], + latent_dim=64, + use_snake=True, + snake_type="vanilla", + use_filter=False, + final_tanh=False, + ) + encoder_fields = { + "in_channels", + "channels", + "latent_dim", + "c_mults", + "strides", + "use_snake", + "antialias_activation", + } + decoder_fields = { + "out_channels", + "channels", + "latent_dim", + "c_mults", + "strides", + "use_snake", + "snake_type", + "antialias_activation", + "use_nearest_upsample", + "use_filter", + "final_tanh", + } + self.encoder_config = {key: value for key, value in self.encoder_config.items() if key in encoder_fields} + self.decoder_config = {key: value for key, value in self.decoder_config.items() if key in decoder_fields} + self.sample_rate = int(sample_rate) + self.latent_dim = int(latent_dim) + self.downsampling_ratio = int(downsampling_ratio) + self.audio_channels = int(audio_channels) + self.release_variant = release_variant + self.decode_core_frames = int(decode_core_frames) + self.decode_halo_frames = int(decode_halo_frames) + if self.decode_core_frames < 1 or self.decode_halo_frames < 0: + raise ValueError("Invalid VAE core/halo configuration") + if math.prod(self.decoder_config["strides"]) != self.downsampling_ratio: + raise ValueError("Decoder strides do not match downsampling_ratio") + if self.decoder_config["latent_dim"] != self.latent_dim: + raise ValueError("Decoder input channels do not match latent_dim") + + +def _dependency_interval(module, low, high): + """Inclusive input support of an output interval; no waveform blending.""" + if isinstance(module, (nn.Sequential, OobleckDecoder, DecoderBlock)): + layers = module if isinstance(module, nn.Sequential) else module.layers + for child in reversed(list(layers)): + low, high = _dependency_interval(child, low, high) + return low, high + if isinstance(module, ResidualUnit): + a, b = _dependency_interval(module.layers, low, high) + return min(a, low), max(b, high) + if isinstance(module, nn.ConvTranspose1d): + s, p, d, k = (module.stride[0], module.padding[0], module.dilation[0], module.kernel_size[0]) + return -(-(low + p - d * (k - 1)) // s), (high + p) // s + if isinstance(module, nn.Conv1d): + s, p, d, k = (module.stride[0], module.padding[0], module.dilation[0], module.kernel_size[0]) + return low * s - p, high * s - p + d * (k - 1) + if isinstance(module, (SnakeBeta, nn.ELU, nn.Identity, nn.Tanh)): + return low, high + raise TypeError(f"No audited support rule for {type(module).__name__}") + + +def _output_length(module, length): + if isinstance(module, (nn.Sequential, OobleckDecoder, DecoderBlock)): + layers = module if isinstance(module, nn.Sequential) else module.layers + for child in layers: + length = _output_length(child, length) + return length + if isinstance(module, nn.ConvTranspose1d): + return ( + (length - 1) * module.stride[0] + - 2 * module.padding[0] + + module.dilation[0] * (module.kernel_size[0] - 1) + + module.output_padding[0] + + 1 + ) + if isinstance(module, nn.Conv1d): + return (length + 2 * module.padding[0] - module.dilation[0] * (module.kernel_size[0] - 1) - 1) // module.stride[ + 0 + ] + 1 + if isinstance(module, (ResidualUnit, SnakeBeta, nn.ELU, nn.Identity, nn.Tanh)): + return length + raise TypeError(f"No audited length rule for {type(module).__name__}") + + +class YuE2VAE(PreTrainedModel): + """Strict EMA model; only the decoder needs to reside on an accelerator. + + ``decode_tiled`` preserves finite receptive-field context and writes cropped + cores to CPU. The mathematical waveform is the full decoder's waveform; + convolution kernel choices may cause small FP32 rounding differences. + """ + + config_class = YuE2VAEConfig + base_model_prefix = "" + main_input_name = "audio" + + def __init__(self, config, decoder_only=False): + super().__init__(config) + self.decoder_only = bool(decoder_only) + if not self.decoder_only: + self.encoder = OobleckEncoder(**config.encoder_config) + self.decoder = OobleckDecoder(**config.decoder_config) + self.eval().requires_grad_(False) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path, + *model_args, + config=None, + decoder_only=False, + device="cpu", + torch_dtype=None, + dtype=None, + revision=None, + token=None, + cache_dir=None, + local_files_only=False, + force_download=False, + subfolder="", + **kwargs, + ): + """Load a local export or Hub repository, selecting decoder tensors. + + Complete exports use unprefixed EMA keys. ``decoder_only=True`` avoids + constructing the encoder or loading encoder tensors into RAM/GPU. + """ + from safetensors import safe_open + + requested_dtype = dtype if dtype is not None else torch_dtype + if requested_dtype not in (None, "auto", "float32", torch.float32): + raise ValueError("The validated VAE requires FP32; quantize the LM separately") + device_map = kwargs.pop("device_map", None) + if device_map is not None: + if device_map == "auto": + device = "cuda" if torch.cuda.is_available() else "cpu" + elif isinstance(device_map, str): + device = device_map + elif isinstance(device_map, dict) and set(device_map) == {""}: + device = device_map[""] + else: + raise ValueError("Use decoder_only=True and a single device for the VAE") + for name in ( + "trust_remote_code", + "low_cpu_mem_usage", + "_from_auto", + "_from_pipeline", + "_commit_hash", + "adapter_kwargs", + "_fast_init", + "weights_only", + "use_safetensors", + ): + kwargs.pop(name, None) + output_loading_info = kwargs.pop("output_loading_info", False) + if kwargs or model_args: + raise TypeError(f"Unsupported VAE loading options: {sorted(kwargs)}") + path = Path(pretrained_model_name_or_path).expanduser() + if not path.is_dir(): + from vllm_omni.transformers_utils.repo_utils import hf_api + + path = Path( + hf_api().snapshot_download( + str(pretrained_model_name_or_path), + revision=revision, + token=token, + cache_dir=cache_dir, + local_files_only=local_files_only, + force_download=force_download, + allow_patterns=[ + f"{subfolder + '/' if subfolder else ''}{pattern}" + for pattern in ("config.json", "*.safetensors", "*.safetensors.index.json") + ], + ) + ) + path = path / subfolder + if config is None: + config = YuE2VAEConfig.from_pretrained(path, local_files_only=True) + # vLLM may construct us under a BF16 default dtype. Widen before copying + # the FP32 checkpoint so load_state_dict does not round its values. + model = cls(config, decoder_only=decoder_only).float() + index = path / "model.safetensors.index.json" + if index.exists(): + mapping = json.loads(index.read_text())["weight_map"] + files = sorted({name for key, name in mapping.items() if not decoder_only or key.startswith("decoder.")}) + else: + files = ["model.safetensors"] + state = {} + for name in files: + with safe_open(path / name, framework="pt", device="cpu") as handle: + for key in handle.keys(): + if not decoder_only or key.startswith("decoder."): + if key in state: + raise ValueError(f"Duplicate VAE tensor: {key}") + state[key] = handle.get_tensor(key) + expected = set(model.state_dict()) + if set(state) != expected: + raise ValueError( + f"VAE tensor mismatch: missing={sorted(expected - set(state))}, " + f"unexpected={sorted(set(state) - expected)}" + ) + if any(value.dtype != torch.float32 for value in state.values()): + raise ValueError("VAE export contains tensors that are not FP32") + model.load_state_dict(state, strict=True) + model.to(device=device, dtype=torch.float32).eval().requires_grad_(False) + if output_loading_info: + return model, dict(missing_keys=[], unexpected_keys=[], mismatched_keys=[], error_msgs=[]) + return model + + def save_pretrained(self, save_directory, *args, **kwargs): + if self.decoder_only: + raise ValueError("Reload decoder_only=False to save a complete VAE repository") + if kwargs.get("safe_serialization", True) is False: + raise ValueError("YuE2 VAE release exports require safetensors") + return super().save_pretrained(save_directory, *args, **kwargs) + + @property + def decoder_device(self): + return next(self.decoder.parameters()).device + + def _latent(self, latent): + latent = torch.as_tensor(latent) + if latent.ndim != 3 or latent.shape[1] != self.config.latent_dim or latent.shape[0] < 1 or latent.shape[-1] < 1: + raise ValueError(f"Expected nonempty [B,{self.config.latent_dim},T] latents") + if not torch.isfinite(latent).all(): + raise ValueError("VAE latents contain non-finite values") + if next(self.decoder.parameters()).dtype != torch.float32: + raise ValueError("VAE decoder weights must remain FP32") + return latent + + @torch.inference_mode() + def encode(self, audio, sample=False, generator=None, return_info=False): + """Encode FP32 stereo audio; posterior mean by default. + + Set ``sample=True`` with a per-request ``torch.Generator`` to reproduce + stochastic posterior sampling. This audio VAE is not the unreleased + semantic audio tokenizer semantic tokenizer. + """ + if self.decoder_only: + raise RuntimeError("Encoder not loaded; reload with decoder_only=False") + audio = torch.as_tensor(audio) + if ( + audio.ndim != 3 + or audio.shape[1] != self.config.audio_channels + or audio.shape[-1] < self.config.downsampling_ratio + ): + raise ValueError("Expected audio [B,2,S] with at least one latent frame") + if not torch.isfinite(audio).all(): + raise ValueError("Audio contains non-finite values") + device = next(self.encoder.parameters()).device + pre = self.encoder(audio.to(device=device, dtype=torch.float32)) + mean, scale = pre.chunk(2, dim=1) + stdev = torch.nn.functional.softplus(scale) + 1e-4 + if sample: + noise = torch.randn(mean.shape, dtype=mean.dtype, device=device, generator=generator) + latent = noise * stdev + mean + else: + latent = mean + if return_info: + return latent, dict(mean=mean, scale=scale, stdev=stdev) + return latent + + @torch.inference_mode() + def decode(self, latent): + """Full waveform, FP32 [B,2,1920*T-64], without clipping.""" + latent = self._latent(latent) + with torch.autocast(device_type=self.decoder_device.type, enabled=False): + return self.decoder(latent.to(device=self.decoder_device, dtype=torch.float32)) + + def natural_output_length(self, frames): + if int(frames) < 1: + raise ValueError("frames must be positive") + return _output_length(self.decoder, int(frames)) + + def required_halo(self, core_frames=None): + core_frames = self.config.decode_core_frames if core_frames is None else core_frames + ratio = self.config.downsampling_ratio + low, high = _dependency_interval(self.decoder, 0, core_frames * ratio - 1) + return max(0, -low, high - core_frames + 1) + + @torch.inference_mode() + def decode_tiled( + self, + latent, + core_frames=None, + halo_frames=None, + output_device="cpu", + on_progress: Callable[[int, int], None] | None = None, + ): + """Decode bounded tiles, retaining exact cores with natural end length. + + Each crop has enough left/right context for every dependency. There is + no crossfade or boundary smoothing, and no zero padding of final audio. + CPU output prevents an entire song from accumulating on the GPU. + ``on_progress(completed, total)`` runs after each existing crop copy; + no extra synchronization is added. With a CUDA output device, queued + work may still be executing. Callback exceptions propagate. + """ + latent = self._latent(latent) + core_frames = self.config.decode_core_frames if core_frames is None else core_frames + halo_frames = self.config.decode_halo_frames if halo_frames is None else halo_frames + if not isinstance(core_frames, int) or core_frames < 1: + raise ValueError("core_frames must be a positive integer") + required = self.required_halo(core_frames) + if not isinstance(halo_frames, int) or halo_frames < required: + raise ValueError(f"halo_frames must be at least {required} for this decoder") + frames = latent.shape[-1] + ratio = self.config.downsampling_ratio + total = self.natural_output_length(frames) + audio = torch.empty( + (latent.shape[0], self.config.audio_channels, total), dtype=torch.float32, device=output_device + ) + tiles = (frames + core_frames - 1) // core_frames + for tile_index, start in enumerate(range(0, frames, core_frames)): + end = min(frames, start + core_frames) + left = max(0, start - halo_frames) + right = min(frames, end + halo_frames) + tile = self.decode(latent[..., left:right]) + out_start, out_end = start * ratio, min(end * ratio, total) + crop_start = (start - left) * ratio + crop = tile[..., crop_start : crop_start + out_end - out_start] + if crop.shape[-1] != out_end - out_start: + raise RuntimeError("VAE tile did not cover its requested output core") + audio[..., out_start:out_end].copy_(crop.to(output_device)) + del tile, crop + if on_progress is not None: + on_progress(tile_index + 1, tiles) + return audio + + def decode_audio(self, latent, chunked=True, **kwargs): + return self.decode_tiled(latent, **kwargs) if chunked else self.decode(latent) + + def forward(self, audio, sample=False, generator=None): + return self.decode(self.encode(audio, sample=sample, generator=generator)) + + +YuE2VAEConfig.register_for_auto_class("AutoConfig") +YuE2VAE.register_for_auto_class("AutoModel") + + +# -------------------- the model -------------------- + + +"""YuE2-3B text-to-music model for vllm-omni: single-stage native-AR. + +The checkpoint is one Mixture-of-Transformers: a Qwen3-shaped AR path (which +is exactly a Qwen3-1.7B-class backbone over an extended vocabulary) plus a +parallel NAR path (``nar_self_attn``/``nar_mlp`` per layer) that a 32-step +midpoint flow-matching solver walks over 64-dim VAE latents. Both paths and +the projection heads live in the same ``model.safetensors``, so a single +vLLM stage loads everything once and no weights are duplicated. + +Sampling is model-owned (``prefer_model_sampler``), reproducing the upstream +request-local arithmetic: per-phase vocabulary masking, windowed repetition +penalty, seeded multinomial, ``min_tokens`` end-token suppression. One song +is one or two engine requests — the abc phase (``cot=full|melody``) first, +then the semantic phase; the driver stitches them. When the semantic phase's +end token is drawn (or its frame budget is hit), the model solves the ODE and +decodes 48 kHz stereo audio inside the step and ships it as the request's +final multimodal payload, so the whole song arrives in one piece. +""" +logger = init_logger(__name__) + +__all__ = ["Yue2ForCausalLM"] + +DEFAULT_SEED = 831001 + + +class _RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight + + +class _RotaryEmbedding(nn.Module): + """Checkpoint-compatible RoPE: cos/sin over [B, T, head_dim/2].""" + + def __init__(self, head_dim: int, base: float = 1000000.0): + super().__init__() + self.head_dim = head_dim + self.base = base + self._inv_freq: torch.Tensor | None = None + + def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + device = position_ids.device + if self._inv_freq is None or self._inv_freq.device != device: + self._inv_freq = 1.0 / ( + self.base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32, device=device) / self.head_dim) + ) + angles = position_ids.float().unsqueeze(-1) * self._inv_freq + return angles.cos(), angles.sin() + + +def _apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + x1, x2 = x[..., :half], x[..., half:] + cos, sin = cos.to(x.dtype), sin.to(x.dtype) + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) + + +class _NARAttention(nn.Module): + """NAR-path attention with checkpoint-native separate projections.""" + + def __init__(self, config): + super().__init__() + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False) + self.q_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) + self.k_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) + + def project_qkv(self, x, cos, sin): + B, T, _ = x.shape + q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim) + k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim) + v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim) + q, k = self.q_norm(q), self.k_norm(k) + rc, rs = cos.unsqueeze(2), sin.unsqueeze(2) + return _apply_rotary(q, rc, rs), _apply_rotary(k, rc, rs), v + + +class _NARMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x): + return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class _NARLayer(nn.Module): + """One NAR path layer; checkpoint ``model.layers.N.nar_*`` remaps onto it.""" + + def __init__(self, config): + super().__init__() + self.input_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.self_attn = _NARAttention(config) + self.pre_mlp_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.mlp = _NARMLP(config) + + +class _TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size), + ) + self.frequency_embedding_size = frequency_embedding_size + + def forward(self, t): + half = self.frequency_embedding_size // 2 + freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device, dtype=torch.float32) / half) + args = t.float().unsqueeze(-1) * freqs.unsqueeze(0) + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + return self.mlp(emb.to(next(self.parameters()).dtype)) + + +class _AudioPositionEmbedding(nn.Module): + def __init__(self, max_frames: int, hidden_size: int): + super().__init__() + pe = torch.zeros(max_frames, hidden_size) + position = torch.arange(0, max_frames, dtype=torch.float32).unsqueeze(1) + div_term = torch.exp(torch.arange(0, hidden_size, 2, dtype=torch.float32) * (-math.log(10000.0) / hidden_size)) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + self.register_buffer("pe", pe, persistent=False) + + def forward(self, position_ids): + return self.pe[position_ids] + + +def _request_key(value: Any) -> str: + if isinstance(value, (list, tuple)): + return str(value[0]) if value else "" + return "" if value is None else str(value) + + +@dataclass +class _RowConstants: + """Per-request values fixed at admission (a prefill step, never replayed).""" + + request_id: str + phase: str + temperature: float + top_p: float + top_k: int + repetition_penalty: float + penalty_window: int + min_tokens: int + max_audio_frames: int + seed: int + skip_synthesis: bool + + +@dataclass +class _RequestState: + """Everything one song request accumulates across decode steps.""" + + request_id: str + constants: _RowConstants + prompt_tokens: int = 0 + prompt_len: int = 0 # full prompt length; a completing prefill row reaches it + prefix_ids: list[int] = field(default_factory=list) + history: list[int] = field(default_factory=list) + generator: torch.Generator | None = None + finished: bool = False + truncated: bool = False + + +class Yue2ForCausalLM(nn.Module): + """YuE2-3B: vLLM-native Qwen3 AR backbone + NAR flow-matching + VAE.""" + + have_multimodal_outputs = True + prefer_model_sampler = True + has_postprocess = False + # The song is decoded in-model and shipped via make_omni_output; per-step + # hidden states must NOT ride the pooler payload, or the output pipeline + # remaps the accumulated "hidden" rows to the audio modality key and the + # driver receives hidden states instead of the decoded waveform + # (single-stage precedent: minimax_music3 talker). + omni_pooler_payload_include_hidden: bool = False + # No downstream stage consumes prefix-cached hidden/mm tensors for YuE2; + # opting out keeps the merged prefix-cache view from rebuilding "hidden" + # payloads (KV-block prefix caching for cot=full stays on). + requires_full_prefix_cached_hidden_states: bool = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + # The NAR pass keeps every KV head on each rank (ar_project splits by + # the per-rank head count), so TP>1 builds a server whose every song + # fails after its full AR decode. Pipeline parallelism has the same + # problem: AR and NAR must live in one stage. Fail at startup instead. + tp = int(getattr(vllm_config.parallel_config, "tensor_parallel_size", 1) or 1) + pp = int(getattr(vllm_config.parallel_config, "pipeline_parallel_size", 1) or 1) + if tp > 1 or pp > 1: + raise ValueError( + f"Yue2 is single-GPU only: the NAR path is replicated, so " + f"tensor_parallel_size={tp} and pipeline_parallel_size={pp} are " + "unsupported; use 1/1." + ) + self.vllm_config = vllm_config + config = vllm_config.model_config.hf_config + self.config = config + hidden = int(config.hidden_size) + + self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + if getattr(config, "tie_word_embeddings", False): + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + int(config.vocab_size), + hidden, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(int(config.vocab_size)) + + self.nar_layers = nn.ModuleList([_NARLayer(config) for _ in range(int(config.num_hidden_layers))]) + self.vae2llm = nn.Linear(int(getattr(config, "latent_dim", LATENT_DIM)), hidden) + self.llm2vae = nn.Linear(hidden, int(getattr(config, "latent_dim", LATENT_DIM))) + self.time_embedder = _TimestepEmbedder(hidden) + self.latent_pos_embed = _AudioPositionEmbedding( + int(getattr(config, "max_latent_frames", config.max_position_embeddings)), + hidden, + ) + self.rotary_emb = _RotaryEmbedding(int(config.head_dim), float(config.rope_theta)) + + self.vocab_size = int(config.vocab_size) + self.max_position_embeddings = int(config.max_position_embeddings) + self.max_latent_frames = int(getattr(config, "max_latent_frames", self.max_position_embeddings)) + self._timestep_shift = float(getattr(config, "timestep_shift", 1.0)) + + self._states: dict[str, _RequestState] = {} + self._row_constants: dict[str, _RowConstants] = {} + self._audio_queue: list[tuple[str, torch.Tensor, bool, bool]] = [] + self._deferred_cleanup_ids: set[str] = set() + self._step_rows: list[tuple[str, int, int]] = [] # (req_id, computed, scheduled) + self._decode_t0: dict[str, float] = {} # first decode-step wall time, for AR tok/s + self._last_mm: dict[str, Any] | None = None # current step's make_omni_output payload + self._vae: YuE2VAE | None = None + self._vae_device: torch.device | None = None + + # ------------------------------------------------------------ sampling + + def shift_t(self, raw_t: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + t_sig = torch.sigmoid(torch.tensor(raw_t, dtype=dtype, device=device)) + shift = self._timestep_shift + return shift * t_sig / (1 + (shift - 1) * t_sig) + + def rotary(self, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return self.rotary_emb(positions) + + def ar_project(self, layer, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + """AR-path Q/K/V from the vLLM backbone's fused projection.""" + attn = layer.self_attn + qkv = attn.qkv_proj(x) + if isinstance(qkv, tuple): + qkv = qkv[0] + num_heads, num_kv_heads, head_dim = attn.num_heads, attn.num_kv_heads, attn.head_dim + q, k, v = torch.split( + qkv, + [num_heads * head_dim, num_kv_heads * head_dim, num_kv_heads * head_dim], + dim=-1, + ) + T = x.shape[-2] + q = attn.q_norm(q.view(T, num_heads, head_dim)) + k = attn.k_norm(k.view(T, num_kv_heads, head_dim)) + rc, rs = cos.unsqueeze(2), sin.unsqueeze(2) + # q/k are [T, heads, head_dim] (3D); rc/rs are [1, T, 1, head_dim/2] + # (4D). Broadcasting them directly promotes q/k to 4D, which the NAR + # attention rejects. Add the batch axis explicitly, then drop it. + q = _apply_rotary(q.unsqueeze(0), rc, rs).squeeze(0) + k = _apply_rotary(k.unsqueeze(0), rc, rs).squeeze(0) + return q, k, v.view(T, num_kv_heads, head_dim) + + # ------------------------------------------------------------ weights + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + ar_pairs, remapped = partition_checkpoint_weights(weights) + missing, unexpected = self.load_state_dict(dict(remapped), strict=False) + missing = [k for k in missing if not k.startswith(("model.", "lm_head"))] + if missing or unexpected: + raise RuntimeError(f"YuE2 NAR/projection weight mismatch: missing={missing}, unexpected={unexpected}") + + loader = AutoWeightsLoader(self) + loaded = loader.load_weights(iter(ar_pairs)) + # Load the VAE now, not lazily at the first finishing request: a bad + # path or a failed download must surface at startup, and the decoder + # weights must sit on the GPU before vLLM's memory profiling sizes + # the KV cache. + device = f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu" + self._vae_model(torch.device(device)) + # The NAR/projection modules were populated by hand above; + # DefaultModelLoader.track_weights_loading diffs every named + # parameter against the returned set, so the side keys must be + # reported too or the loader flags them as uninitialized. + return loaded | {name for name, _ in remapped} + + # ------------------------------------------------------------ hooks + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.model.embed_tokens(input_ids) + + def prepare_runner_inputs( + self, + *, + req_ids: list[str], + num_computed_tokens: Any, + num_scheduled_tokens: Any, + input_ids: torch.Tensor, + positions: torch.Tensor, + **_: Any, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Bind rows to requests for the step; prompt tokens are captured in + ``forward``, where the state is created (see ``_capture_constants``). + Decode rows feed tokens the model itself sampled, so they never + contribute to the prompt prefix.""" + computed = [int(v) for v in num_computed_tokens] + scheduled = [int(v) for v in num_scheduled_tokens] + self._step_rows = [(str(rid), computed[i], scheduled[i]) for i, rid in enumerate(req_ids)] + # make_omni_output runs every step; without this reset the first step's + # payload would stay live forever and the _audio_queue fallback in + # _ship_audio could never trigger. + self._last_mm = None + return input_ids, positions + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Any | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + del intermediate_tensors + self._capture_constants(kwargs, input_ids) + hidden = self.model( + input_ids=input_ids if inputs_embeds is None else None, + positions=positions, + inputs_embeds=inputs_embeds, + ) + if isinstance(hidden, tuple): + hidden = hidden[0] + self._flush_deferred_cleanup() + return hidden + + def _capture_constants(self, kwargs: dict[str, Any], input_ids: torch.Tensor | None) -> None: + """Create per-request state on the request's FIRST scheduled step. + + With KV prefix caching the first step may arrive with ``comp > 0`` + (a shared head was cached by an earlier request), so ``comp == 0`` + cannot be the trigger. The scheduled ``input_ids`` slice then holds + only the uncached tail, so the driver ships the full prompt ids in + the request's extra args (KEY_PREFIX_IDS); they are the source of + truth for the NAR conditioning prefix and for the prompt length the + sampler uses to tell a completing prefill row from a mid-chunk one. + """ + extra_args = kwargs.get("sampling_extra_args") + if extra_args is None or input_ids is None: + return + offset = 0 + for row, (req_id, comp, span) in enumerate(self._step_rows): + if row >= len(extra_args): + break + if req_id in self._states: + offset += span + continue + args = extra_args[row] or {} + prefix = args.get(KEY_PREFIX_IDS) + if isinstance(prefix, (list, tuple)) and prefix: + prefix_ids = [int(v) for v in prefix] + prompt_len = len(prefix_ids) + elif comp == 0 and span >= 1: + # Legacy path (no ids in extra args): only correct when the + # whole prompt is scheduled in one chunk with no cache hit. + prefix_ids = input_ids[offset : offset + span].tolist() + prompt_len = len(prefix_ids) + else: + logger.warning( + "YuE2 request %s arrived with %d cached tokens but no " + "yue2_prefix_ids in its extra args; cannot rebuild its " + "full prompt (prefix-cache hit on the first request?).", + req_id, + comp, + ) + offset += span + continue + phase = str(args.get(KEY_PHASE, "semantic")) + preset = ABC_SAMPLING if phase == "abc" else SEMANTIC_SAMPLING + seed = args.get(KEY_SEED, DEFAULT_SEED) + constants = _RowConstants( + request_id=req_id, + phase=phase, + temperature=float(args.get(KEY_TEMPERATURE, preset["temperature"])), + top_p=float(args.get(KEY_TOP_P, preset["top_p"])), + top_k=int(args.get(KEY_TOP_K, preset["top_k"])), + repetition_penalty=float(args.get(KEY_REPETITION_PENALTY, preset["repetition_penalty"])), + penalty_window=int(args.get(KEY_PENALTY_WINDOW, preset["penalty_window"])), + min_tokens=int(args.get(KEY_MIN_TOKENS, preset["min_tokens"])), + max_audio_frames=int(args.get(KEY_MAX_AUDIO_FRAMES, preset["max_tokens"])), + seed=int(seed), + skip_synthesis=bool(args.get(KEY_SKIP_SYNTHESIS, phase == "abc")), + ) + device = input_ids.device + generator = torch.Generator(device=device if device.type != "mps" else "cpu") + generator.manual_seed(constants.seed) + self._row_constants[req_id] = constants + self._states[req_id] = _RequestState( + request_id=req_id, + constants=constants, + prompt_tokens=prompt_len, + prompt_len=prompt_len, + prefix_ids=prefix_ids, + generator=generator, + ) + offset += span + + def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor: + return self.logits_processor(self.lm_head, hidden_states) + + # ------------------------------------------------------------ sampler + + def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> SamplerOutput | None: + """Draw one token per decoding row with the request's own arithmetic.""" + del sampling_metadata + rows = int(logits.shape[0]) + # The runner's input_ids buffer is int32; long ids crash its scatter + # when concurrent requests re-index rows. + token_ids = torch.zeros((rows, 1), dtype=torch.int32, device=logits.device) + if not self._step_rows: + return SamplerOutput(sampled_token_ids=token_ids, logprobs_tensors=None) + + for row, (req_id, comp, span) in enumerate(self._step_rows): + if row >= rows: + break + state = self._states.get(req_id) + if state is None or state.finished: + continue + if comp + span < state.prompt_len: + # Mid-chunk prefill rows: their sampled token is discarded. + # (A completing prefill row reaches the full prompt length + # even when a prefix-cache hit left comp > 0.) + continue + constants = state.constants + if state.request_id not in self._decode_t0: + self._decode_t0[state.request_id] = time.perf_counter() + scores = distribution( + logits[row], + temperature=constants.temperature, + top_p=constants.top_p, + top_k=constants.top_k, + repetition_penalty=constants.repetition_penalty, + penalty_window=constants.penalty_window, + history=state.history, + step=len(state.history), + min_tokens=constants.min_tokens, + phase=constants.phase, + ) + token = sample_row(scores, state.generator, greedy=constants.temperature == 0) + end = ABC_END if constants.phase == "abc" else MUSIC_END + if token == end: + state.finished = True + state.truncated = False + self._finish_request_safely(state, hit_end=True) + token_ids[row, 0] = end + else: + state.history.append(token) + budget = constants.max_audio_frames + if constants.phase == "semantic" and len(state.history) >= budget: + state.finished = True + state.truncated = True + self._finish_request_safely(state, hit_end=False) + token_ids[row, 0] = end + else: + token_ids[row, 0] = token + return SamplerOutput(sampled_token_ids=token_ids, logprobs_tensors=None) + + # ------------------------------------------------------------ audio + + def _finish_request_safely(self, state: _RequestState, *, hit_end: bool) -> None: + """Keep an NAR/VAE failure from escaping ``sample``. + + The runner calls ``model.sample()`` with no error handling, so an + exception here (OOM in the ODE solve, non-finite VAE output) would + kill Stage-0 and take every live request down with it. Fail only + this request: ship an empty error-flagged clip that serving turns + into a 500. + """ + try: + self._finish_request(state, hit_end=hit_end) + except Exception: + logger.exception("YuE2 synthesis failed for req=%s; failing only this request", state.request_id) + self._ship_audio(state.request_id, torch.zeros((2, 0)), state.truncated, error=True) + + def _finish_request(self, state: _RequestState, *, hit_end: bool) -> None: + """Solve the ODE and decode the song; abc-phase requests skip this.""" + constants = state.constants + if constants.skip_synthesis or not state.history: + return + codec = [t - CODEC_OFFSET for t in state.history] + if min(codec) < 0 or max(codec) >= CODEC_SIZE: + raise RuntimeError("semantic history contains non-codec tokens") + t0 = self._decode_t0.pop(state.request_id, None) + t_nar0 = time.perf_counter() + latents = synthesize(self, state.prefix_ids, codec, constants.seed, steps=ODE_STEPS) + logger.info( + "YUE2_LATENTS frames=%d absmax=%.4f mean=%.5f std=%.4f", + latents.shape[0], + float(latents.abs().max()), + float(latents.mean()), + float(latents.std()), + ) + t_nar = time.perf_counter() - t_nar0 + t_vae0 = time.perf_counter() + audio = self._decode_latents(latents) + t_vae = time.perf_counter() - t_vae0 + t_ar = (t_nar0 - t0) if t0 is not None else float("nan") + logger.info( + "YUE2_TIMING req=%s phase=%s ar_tokens=%d ar_s=%.2f ar_tps=%.1f nar_s=%.2f vae_s=%.2f", + state.request_id, + constants.phase, + len(state.history), + t_ar, + len(state.history) / t_ar if t_ar else float("nan"), + t_nar, + t_vae, + ) + self._ship_audio(state.request_id, audio, state.truncated) + del hit_end + + def _ship_audio(self, req_id: str, audio: torch.Tensor, truncated: bool, *, error: bool = False) -> None: + """Append the finished song to the CURRENT step's mm payload in place. + + make_omni_output already ran for this step (it ships the sparse marker + with empty lists), and the runner assembles per-request payloads only + after sample(), so mutating the same dict here is picked up by the + sparse routing on the request's last step in the output batch. Leaving + the payload for a later step would drop it: the request is gone by + then (gepard_talker flushes at its last step for the same reason). + """ + mm = self._last_mm + if mm is None: + # No forward built a payload this step; fall back to the queue. + self._audio_queue.append((req_id, audio, truncated, error)) + return + mm["model_outputs"].append(audio) + mm["sr"].append(torch.tensor(SAMPLE_RATE, dtype=torch.int32)) + meta = mm["meta"] + meta["req_id"].append(req_id) + # Keep the flags ints, not strs: the wire payload is tensor-only + # (_ensure_tensor_values) and a string would be dropped before serving. + meta["truncated"].append(int(truncated)) + meta["error"].append(int(error)) + + def _vae_model(self, device: torch.device) -> YuE2VAE: + if self._vae is None: + vae_path = os.environ.get("YUE2_VAE", DEFAULT_VAE_ID) + logger.info("Loading YuE2 VAE decoder from %s", vae_path) + # Keep the VAE out of the module tree (__dict__, not an nn.Module + # attribute): the checkpoint has no VAE keys, so a registered + # submodule would trip DefaultModelLoader.track_weights_loading, + # and a model-wide dtype/device pass would change VAE numerics. + self.__dict__["_vae"] = YuE2VAE.from_pretrained(vae_path, decoder_only=True, device="cpu") + if self._vae_device != device: + self._vae.to(device) + self._vae_device = device + return self._vae + + def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: + """[frames, 64] FP32 CPU latents -> channels-first stereo [2, samples].""" + device = f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu" + model = self._vae_model(torch.device(device)) + z = latents.T.unsqueeze(0) # [1, 64, T] + with torch.inference_mode(): + audio = model.decode_tiled( + z, + core_frames=VAE_CORE_FRAMES, + halo_frames=VAE_HALO_FRAMES, + output_device="cpu", + ) + if not torch.isfinite(audio).all(): + raise RuntimeError("VAE produced non-finite audio") + logger.info("YUE2_AUDIO_PRECLAMP absmax=%.4f", float(audio.abs().max())) + # Channels-first [2, T], the same layout MiniMax Music 3 ships in + # model_outputs; serving's create_audio transposes it for soundfile. + return audio[0].float().clamp(-1, 1).contiguous() + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + by_req: dict[str, torch.Tensor] = {} + truncated: dict[str, bool] = {} + failed: dict[str, bool] = {} + for req_id, audio, is_truncated, is_error in self._audio_queue: + by_req[req_id] = audio + truncated[req_id] = is_truncated + failed[req_id] = is_error + self._audio_queue.clear() + ready = list(by_req) + sr = torch.tensor(SAMPLE_RATE, dtype=torch.int32) + # The sparse marker rides EVERY step, even ones that decoded nothing + # (req_id=[] is the legal zero-audio shape): without it the runner + # falls back to the dense pooler payload, whose per-step hidden states + # land under this stage's "audio" key and the driver receives hidden + # states instead of the decoded waveform (gepard_talker precedent). + mm: dict[str, Any] = { + "model_outputs": [by_req[r] for r in ready], + "sr": [sr for _ in ready], + "meta": { + "req_id": ready, + "sparse_audio": ["1"], + "truncated": [int(truncated[r]) for r in ready], + "error": [int(failed[r]) for r in ready], + }, + } + # sample() runs after this each step and appends finished songs into + # this very dict (see _ship_audio); keep the handle for it. + self._last_mm = mm + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=mm) + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + # Fires before forward; defer the free so the in-flight step can read. + for rid in finished_req_ids: + self._deferred_cleanup_ids.add(_request_key(rid)) + + def _flush_deferred_cleanup(self) -> None: + for req_id in self._deferred_cleanup_ids: + state = self._states.pop(req_id, None) + if state is not None and not state.finished: + logger.warning( + "YuE2 request %s ended without an end token or its frame " + "budget (aborted or preempted); no audio was produced.", + req_id, + ) + self._row_constants.pop(req_id, None) + self._deferred_cleanup_ids.clear() diff --git a/vllm_omni/tokenizers/yue2_prompt.py b/vllm_omni/tokenizers/yue2_prompt.py new file mode 100644 index 000000000..97d7b7dea --- /dev/null +++ b/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", +] diff --git a/vllm_omni/tokenizers/yue2_tokenizer.py b/vllm_omni/tokenizers/yue2_tokenizer.py new file mode 100644 index 000000000..35c222821 --- /dev/null +++ b/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|>", "", "", "", "", ""] + specials += [f"" for i in range(200)] + specials[204:206] = ["", ""] + 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"] diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index c8a5c19f5..10c16b322 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/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 diff --git a/vllm_omni/transformers_utils/configs/yue2.py b/vllm_omni/transformers_utils/configs/yue2.py new file mode 100644 index 000000000..5063d5b15 --- /dev/null +++ b/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"]