diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 410695f4b..1dd5caa66 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -228,6 +228,27 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: "TTS · Ming-Omni-TTS Test" + timeout_in_minutes: 20 + commands: + - | + timeout 20m bash -c " + export VLLM_LOGGING_LEVEL=DEBUG + pytest -s -v tests/e2e/online_serving/test_ming_tts.py tests/e2e/offline_inference/test_ming_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model' + " + agents: + queue: "gpu_1_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: "TTS · Voxtral-TTS Test" source_file_dependencies: - tests/e2e/offline_inference/test_voxtral_tts.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index c47b6a503..49e837772 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -63,6 +63,7 @@ th { | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `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-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | | | `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | | `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | | `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index a0e9b529b..15672b94f 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -17,6 +17,7 @@ list of supported architectures across all modalities, see | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz | | GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz | +| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz | | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz | | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz | | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz | @@ -159,6 +160,46 @@ Streaming requires `async_chunk: true` in the stage config. --- +## Ming-omni-tts + +Dense 0.5B two-stage TTS pipeline (`AR + flow` + audio VAE) at 44.1 kHz. The example covers style, IP voice, music-only generation, text-to-audio events, emotion, dialect, zero-shot cloning, podcast, speech+BGM, and speech+environment-sound cases. + +### Quick start +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case style \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +### Voice cloning +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case zero_shot \ + --ref-audio /path/to/reference.wav \ + --ref-text "在此奉劝大家别乱打美白针。" \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +### Streaming +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case basic \ + --ref-audio /path/to/reference.wav \ + --streaming \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +### Notes +- `style`, `ip`, `bgm`, and `tta` do not require reference audio. +- Reference-audio cases use `--ref-audio`; `zero_shot` also requires `--ref-text`. +- `podcast` uses multiple references via `--ref-audio-paths`. +- Full case details live in [`ming_tts/README.md`](ming_tts/README.md). + +--- + ## Ming-flash-omni-TTS Standalone talker-only deployment of Ming-flash-omni-2.0 at 44.1 kHz. Voice is controlled through caption fields (`风格` / `IP` / `语速`/`基频`/`音量`) rather than reference audio. diff --git a/examples/offline_inference/text_to_speech/ming_tts/README.md b/examples/offline_inference/text_to_speech/ming_tts/README.md new file mode 100644 index 000000000..3fdfaf4d5 --- /dev/null +++ b/examples/offline_inference/text_to_speech/ming_tts/README.md @@ -0,0 +1,248 @@ +# Ming-omni-tts Offline Inference + +`end2end.py` runs Ming dense 0.5B end to end with vLLM-Omni. It uses the in-repo Ming prompt assembly helper directly, so the example request shape matches the real integration instead of a simplified wrapper. + +## Files + +| File | Purpose | +|---|---| +| `end2end.py` | Driver: CLI, case loading, prompt construction, orchestration (~150 lines) | +| `cases.yaml` | All 11 built-in case definitions (prompt, text, instruction, ref-audio flags, flow controls) | +| `runner.py` | Engine management and audio output (streaming + blocking paths) | + +## Model Overview + +Ming dense 0.5B is exposed here as a two-stage offline pipeline: + +- **Stage 0**: Qwen2-based AR generation with Ming prompt formatting and inline flow controls +- **Stage 1**: audio VAE decode to mono 44.1 kHz waveform + +`config_ming_tts.py` adapts the checkpoint's HuggingFace config fields +(LLM, DiT, aggregator, AudioVAE, token ids). `vllm_omni/deploy/ming_tts.yaml` +selects the vLLM-Omni pipeline and stage runtime topology, including +connectors, async chunking, memory limits, and sampling defaults. + +The example supports both: + +- **Blocking eager** via `vllm_omni/deploy/ming_tts.yaml` +- **Async chunk eager** via `vllm_omni/deploy/ming_tts.yaml` (default `async_chunk: true`) + +## Setup + +Install vLLM-Omni with the platform requirements for your accelerator: + +```bash +uv pip install -e . +``` + +The Ming offline example does not require a separate upstream Ming package. +Reference-audio cases use the repo dependencies for audio loading, +resampling, and CampPlus speaker extraction, including `soundfile`, +`torchaudio`, and `onnxruntime-rocm` in the documented ROCm environment. + +The tested ROCm environment is summarized in the +[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md). + +## Supported Cases + +These cases cover the upstream dense 0.5B cookbook surface that maps cleanly onto the current vLLM-Omni example: + +- `style`: zero-speaker style-conditioned speech +- `ip`: zero-speaker IP voice generation +- `bgm`: music-only generation +- `tta`: text-to-audio event generation with FlowLoss controls +- `emotion`: reference-audio speech with emotion control +- `basic`: reference-audio speech with speed / pitch / volume control +- `dialect`: reference-audio speech with dialect control +- `zero_shot`: reference-audio cloning with explicit transcript +- `podcast`: multi-reference dialogue generation with automatic speaker embedding extraction +- `speech_bgm`: speech with background music conditioning +- `speech_sound`: speech with environmental sound conditioning + +## Quick Start + +Run the zero-speaker style example: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case style \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +Run zero-shot cloning with a transcript: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case zero_shot \ + --ref-audio /path/to/10002287-00000094.wav \ + --ref-text "在此奉劝大家别乱打美白针。" \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +Run emotion-controlled speech: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case emotion \ + --ref-audio /path/to/emotion_prompt.wav \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +Run podcast generation with two reference clips: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case podcast \ + --ref-audio-paths /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +The script automatically extracts one 192-d speaker embedding per reference WAV using the Ming model's `campplus.onnx`. + +If you already have precomputed multi-speaker embeddings, you can override extraction with: + +```bash +--speaker-embedding /path/to/podcast_speaker_embeddings.json +``` + +where the JSON is a list of speaker embeddings, one 192-d vector per speaker. + +Run text-to-audio event generation: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case tta \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +Use async_chunk streaming: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case basic \ + --ref-audio /path/to/10002287-00000095.wav \ + --streaming \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager +``` + +`--streaming` uses `AsyncOmni` and the async_chunk stage config. It currently +supports one prompt per process invocation; use blocking mode for +`--num-prompts > 1`. + +Collect runtime stats and a manifest: + +```bash +python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --case style \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --enforce-eager \ + --enable-stats \ + --stats-log-file output_audio/ming_style_pipeline.log \ + --metadata-json output_audio/ming_style_manifest.json +``` + +## Reference Fixtures + +The upstream Ming cookbook uses these public audio fixtures from `inclusionAI/Ming-omni-tts/data/wavs`: + +- `10002287-00000094.wav` for zero-shot cloning +- `10002287-00000095.wav` for `basic` +- `emotion_prompt.wav` for `emotion` +- `yue_prompt.wav` for `dialect` +- `00000309-00000300.wav` for `speech_bgm` and `speech_sound` +- `CTS-CN-F2F-2019-11-11-423-012-A.wav` and `CTS-CN-F2F-2019-11-11-423-012-B.wav` for `podcast` + +## Validation Matrix + +The repo-facing example is intended to cover the same dense TTS workflows used +by the local Ming validation script: + +| Case | Blocking `deploy/ming_tts.yaml` | Async chunk `deploy/ming_tts.yaml` | Extra inputs | +|---|---:|---:|---| +| `style` | Yes | Optional smoke test | none | +| `ip` | Yes | Optional smoke test | none | +| `bgm` | Yes | Optional smoke test | none | +| `tta` | Yes | Optional smoke test | none | +| `emotion` | Yes | Yes | `--ref-audio emotion_prompt.wav` | +| `basic` | Yes | Yes | `--ref-audio 10002287-00000095.wav` | +| `dialect` | Yes | Yes | `--ref-audio yue_prompt.wav` | +| `zero_shot` | Yes | Yes | `--ref-audio 10002287-00000094.wav --ref-text ...` | +| `podcast` | Yes | Yes | two `--ref-audio-paths` | +| `speech_bgm` | Yes | Yes | `--ref-audio 00000309-00000300.wav` | +| `speech_sound` | Yes | Yes | `--ref-audio 00000309-00000300.wav` | + +## Validated Outputs + +The following measurements are retained from an earlier L4 CUDA validation; +they are not ROCm benchmark results. Default async_chunk matched blocking +output frame counts and Stage-1 patch counts for every case: + +| Case | Blocking frames / patches / sec | Async chunk frames / patches / sec | +|---|---:|---:| +| `style` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 | +| `ip` | 183456 / 13 / 4.16 | 183456 / 13 / 4.16 | +| `bgm` | 1326528 / 94 / 30.08 | 1326528 / 94 / 30.08 | +| `tta` | 465696 / 33 / 10.56 | 465696 / 33 / 10.56 | +| `emotion` | 324576 / 23 / 7.36 | 324576 / 23 / 7.36 | +| `basic` | 211680 / 15 / 4.80 | 211680 / 15 / 4.80 | +| `dialect` | 239904 / 17 / 5.44 | 239904 / 17 / 5.44 | +| `zero_shot` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 | +| `podcast` | 437472 / 31 / 9.92 | 437472 / 31 / 9.92 | +| `speech_bgm` | 296352 / 21 / 6.72 | 296352 / 21 / 6.72 | +| `speech_sound` | 352800 / 25 / 8.00 | 352800 / 25 / 8.00 | + +## Key Arguments + +| Argument | Description | +|---|---| +| `--model` | Hugging Face repo or local Ming checkpoint path | +| `--deploy-config` | Deploy config YAML. Use `vllm_omni/deploy/ming_tts.yaml` | +| `--case` | Built-in demo case | +| `--ref-audio` | Single reference wav path for cloning-style cases | +| `--ref-audio-paths` | Multiple reference wav paths, used by `podcast` | +| `--ref-text` | Reference transcript. Required for `zero_shot` | +| `--instructions` | Free-form Ming instruction string | +| `--instruction-json` | Structured Ming instruction JSON | +| `--speaker-embedding` | JSON file containing a 192-d speaker embedding | +| `--extract-speaker-embeddings` | Force CampPlus speaker extraction from the provided reference audio paths | +| `--max-decode-steps` | Override `ming_max_decode_steps` | +| `--num-prompts` | Repeat the same case N times. Outputs are indexed when `N > 1` | +| `--streaming` | Use `AsyncOmni` and async_chunk transport | +| `--enforce-eager` | Recommended for Ming dense; non-eager is out of scope | +| `--enable-stats` / `--log-stats` | Enable vLLM-Omni per-request stats logging | +| `--stats-log-file` | Optional path for the stats log | +| `--metadata-json` | Optional path for the run manifest JSON | +| `--stage-init-timeout` | Per-stage initialization timeout in seconds | +| `--init-timeout` | Total initialization timeout in seconds | +| `--batch-timeout` | Batch timeout in seconds | +| `--worker-backend` | `multi_process` or `ray` | +| `--ray-address` | Ray cluster address when using `--worker-backend ray` | + +## Output + +- The script writes one mono 44.1 kHz WAV file per run +- Default output directory: `output_audio/` +- Default filename: `ming_.wav` +- When `--num-prompts > 1`, outputs are indexed as `ming__00000.wav`, `..._00001.wav`, etc. +- When stats are enabled, the script can also write: + - a stats log file such as `ming_style_pipeline.log` + - a manifest JSON with per-output metadata, stage durations, peak memory info, + and streaming client latency metrics when `--streaming` is used + +## Notes + +- `style` and `ip` are zero-speaker paths and do not require a reference clip +- `emotion`, `basic`, `dialect`, `speech_bgm`, and `speech_sound` require one reference clip +- `zero_shot` requires both `--ref-audio` and `--ref-text` +- `podcast` requires at least two reference clips via `--ref-audio-paths` +- `podcast` automatically extracts one speaker embedding per reference clip +- `--speaker-embedding` may contain either one 192-d vector or a list of 192-d vectors +- `--enforce-eager` was used for the validated runs +- The earlier L4 validation used SDPA for the Ming audio VAE instead of + FlashAttention2, which is the preferred default when available. diff --git a/examples/offline_inference/text_to_speech/ming_tts/cases.yaml b/examples/offline_inference/text_to_speech/ming_tts/cases.yaml new file mode 100644 index 000000000..d17b2761a --- /dev/null +++ b/examples/offline_inference/text_to_speech/ming_tts/cases.yaml @@ -0,0 +1,110 @@ +style: + prompt: "Please generate speech based on the following description.\n" + text: "我会一直在这里陪着你，直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗？" + instruction: + 风格: >- + 这是一种ASMR耳语，属于一种旨在引发特殊感官体验的创意风格。这个女性使用轻柔的普通话进行耳语，声音气音成分重。音量极低，紧贴麦克风，语速极慢，旨在制造触发听者颅内快感的声学刺激。 + use_zero_spk_emb: true + max_decode_steps: 200 + +ip: + prompt: "Please generate speech based on the following description.\n" + text: "这款产品的名字，叫变态坑爹牛肉丸。" + instruction: + IP: "灵小甄" + use_zero_spk_emb: true + max_decode_steps: 200 + +bgm: + prompt: "Please generate music based on the following description.\n" + text: " Genre: 电子舞曲. Mood: 自信 / 坚定. Instrument: 架子鼓. Theme: 节日. Duration: 30s." + instruction: null + use_zero_spk_emb: false + max_decode_steps: 400 + +tta: + prompt: "Please generate audio events based on given text.\n" + text: "Thunder and a gentle rain" + instruction: null + use_zero_spk_emb: false + max_decode_steps: 200 + cfg: 4.5 + sigma: 0.3 + temperature: 2.5 + +emotion: + prompt: "Please generate speech based on the following description.\n" + text: "我竟然抢到了陈奕迅的演唱会门票！太棒了！终于可以现场听一听他的歌声了！" + instruction: + 情感: "高兴" + requires_ref_audio: true + auto_extract_speaker_embeddings: true + max_decode_steps: 200 + +basic: + prompt: "Please generate speech based on the following description.\n" + text: "简单地说，这相当于惠普把消费领域市场拱手相让了。" + instruction: + 语速: "快速" + 基频: "中" + 音量: "中" + requires_ref_audio: true + auto_extract_speaker_embeddings: true + max_decode_steps: 200 + +dialect: + prompt: "Please generate speech based on the following description.\n" + text: "我觉得社会企业同个人都有责任" + instruction: + 方言: "广粤话" + requires_ref_audio: true + auto_extract_speaker_embeddings: true + max_decode_steps: 200 + +zero_shot: + prompt: "Please generate speech based on the following description.\n" + text: "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。" + instruction: null + requires_ref_audio: true + requires_ref_text: true + auto_extract_speaker_embeddings: true + max_decode_steps: 200 + +podcast: + prompt: "Please generate speech based on the following description.\n" + text: " speaker_1:你可以说一下，就大概说一下，可能虽然我也不知道，我看过那部电影没有。\n speaker_2:就是那个叫什么，变相一节课的嘛。\n speaker_1:嗯。\n speaker_2:一部搞笑的电影。\n speaker_1:一部搞笑的。\n" + instruction: null + prompt_text: " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试，做个，当服务员还要去笔试了\n speaker_2:对啊，这真的很奇怪，就是 单纯的因，单纯自己工资不高，只是因为可能人家那个店比较出名一点，就对你苛刻要求\n" + requires_ref_audio_count: 2 + auto_extract_speaker_embeddings: true + max_decode_steps: 200 + +speech_bgm: + prompt: "Please generate speech based on the following description.\n" + text: "此次业绩下滑原因，可归结为企业停止服务某些品牌，而带来的负面影响。" + instruction: + BGM: + Genre: "当代古典音乐." + Mood: "温暖 / 友善." + Instrument: "电吉他" + Theme: "节日." + SNR: 10.0 + ENV: null + requires_ref_audio: true + auto_extract_speaker_embeddings: true + max_decode_steps: 200 + +speech_sound: + prompt: "Please generate speech based on the following description.\n" + text: "此次业绩下滑原因，可归结为企业停止服务某些品牌，而带来的负面影响。" + instruction: + BGM: + ENV: "Birds chirping" + SNR: 10.0 + Genre: null + Mood: null + Instrument: null + Theme: null + requires_ref_audio: true + auto_extract_speaker_embeddings: true + max_decode_steps: 200 diff --git a/examples/offline_inference/text_to_speech/ming_tts/end2end.py b/examples/offline_inference/text_to_speech/ming_tts/end2end.py new file mode 100644 index 000000000..1a6991e40 --- /dev/null +++ b/examples/offline_inference/text_to_speech/ming_tts/end2end.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Offline inference demo for Ming-omni-tts via vLLM Omni.""" + +import json +import os +from pathlib import Path + +import soundfile as sf +import torch +import torchaudio +import yaml +from transformers import AutoTokenizer +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( + KEY_CFG, + KEY_MAX_DECODE_STEPS, + KEY_SIGMA, + KEY_TEMPERATURE, + SAMPLE_RATE, +) +from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt +from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor + +try: + from .runner import ( + build_manifest, + build_sampling_params, + resolve_metadata_json, + resolve_stats_log_file, + run_generation, + ) +except ImportError: + from runner import ( + build_manifest, + build_sampling_params, + resolve_metadata_json, + resolve_stats_log_file, + run_generation, + ) + +_DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B" +_REPO_ROOT = Path(__file__).resolve().parents[4] +_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni/deploy/ming_tts.yaml") +_CASES_FILE = Path(__file__).with_name("cases.yaml") + +CASE_DEFAULTS = yaml.safe_load(_CASES_FILE.read_text(encoding="utf-8")) or {} +if not CASE_DEFAULTS: + raise RuntimeError(f"Empty or missing case definitions in {_CASES_FILE}") + + +def _build_parser() -> FlexibleArgumentParser: + p = FlexibleArgumentParser(description="Offline Ming-omni-tts example") + p.add_argument("--model", default=_DEFAULT_MODEL, help="Model name or local path") + p.add_argument("--deploy-config", default=None, help="Deploy config YAML; auto-selected when omitted") + p.add_argument("--case", choices=sorted(CASE_DEFAULTS), default="style", help="Built-in demo case") + p.add_argument("--text", default=None, help="Override case text") + p.add_argument("--prompt", default=None, help="Override the system prompt prefix") + p.add_argument("--instructions", default=None, help="Free-form Ming instruction string") + p.add_argument( + "--instruction-json", default=None, help='Structured Ming instruction JSON, e.g. \'{"方言":"广粤话"}\'' + ) + p.add_argument("--ref-audio", default=None, help="Single reference audio path for cloning-style cases") + p.add_argument("--ref-audio-paths", nargs="+", default=None, help="Multiple reference audio paths (podcast)") + p.add_argument("--ref-text", default=None, help="Reference transcript; required for zero_shot") + p.add_argument("--speaker-embedding", default=None, help="Path to a JSON speaker embedding file") + p.add_argument( + "--extract-speaker-embeddings", + action="store_true", + help="Extract speaker embeddings from ref audio via campplus.onnx", + ) + p.add_argument("--max-decode-steps", type=int, default=None, help="Override ming_max_decode_steps") + p.add_argument("--output-dir", default="output_audio", help="Directory for output wav files") + p.add_argument("--output-name", default=None, help="Output wav filename") + p.add_argument("--num-prompts", type=int, default=1, help="Repeat the same prompt N times") + p.add_argument("--streaming", action="store_true", help="Use AsyncOmni with async_chunk streaming") + p.add_argument("--trust-remote-code", action="store_true") + p.add_argument("--enforce-eager", action="store_true") + p.add_argument( + "--log-stats", "--enable-stats", dest="log_stats", action="store_true", help="Enable Omni stats logging" + ) + p.add_argument("--stats-log-file", default=None, help="Path for the Omni stats log file") + p.add_argument("--metadata-json", default=None, help="Path for the run manifest JSON") + p.add_argument("--stage-init-timeout", type=int, default=300, help="Per-stage init timeout (s)") + p.add_argument("--init-timeout", type=int, default=600, help="Total init timeout (s)") + p.add_argument("--batch-timeout", type=int, default=5, help="Batch timeout (s)") + p.add_argument("--shm-threshold-bytes", type=int, default=65536) + p.add_argument("--worker-backend", default="multi_process", choices=["multi_process", "ray"]) + p.add_argument("--ray-address", default=None, help="Ray cluster address (--worker-backend ray)") + return p + + +def _finalize_args(args) -> None: + if args.instructions is not None and args.instruction_json is not None: + raise RuntimeError("Use either --instructions or --instruction-json, not both") + if args.num_prompts < 1: + raise RuntimeError("--num-prompts must be at least 1") + if args.streaming and args.num_prompts != 1: + raise RuntimeError("--streaming currently supports exactly one prompt") + if args.deploy_config is None: + args.deploy_config = _DEFAULT_DEPLOY_CONFIG + + +def _load_waveform(path: str) -> torch.Tensor: + samples, sr = sf.read(path, dtype="float32") + wav = torch.as_tensor(samples, dtype=torch.float32) + if wav.ndim == 2: + wav = wav.mean(dim=1) + wav = wav.reshape(1, -1) + if int(sr) != SAMPLE_RATE: + wav = torchaudio.functional.resample(wav, int(sr), SAMPLE_RATE) + return wav + + +def _load_speaker_embedding(path: str) -> torch.Tensor: + return torch.as_tensor(json.loads(Path(path).read_text(encoding="utf-8")), dtype=torch.float32) + + +def _ref_audio_paths(args) -> list[str]: + if args.ref_audio is not None and args.ref_audio_paths is not None: + raise RuntimeError("Use either --ref-audio or --ref-audio-paths, not both") + if args.ref_audio_paths is not None: + return list(args.ref_audio_paths) + return [args.ref_audio] if args.ref_audio else [] + + +def _resolve_reference_inputs(args, case: dict, paths: list[str]): + required = int(case.get("requires_ref_audio_count", 0)) + if required > 0 and len(paths) < required: + raise RuntimeError(f"Case '{args.case}' needs {required} ref audio paths via --ref-audio-paths") + if required <= 0 and case.get("requires_ref_audio") and not paths: + raise RuntimeError(f"--ref-audio required for case '{args.case}'") + if not paths: + return None + wavs = [_load_waveform(p) for p in paths] + return wavs[0] if len(wavs) == 1 else wavs + + +def _resolve_speaker_embedding(args, case: dict, paths: list[str]): + if args.speaker_embedding: + return _load_speaker_embedding(args.speaker_embedding) + if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths: + return None + embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths) + if not embs: + raise RuntimeError("Speaker extraction produced no embeddings") + return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0) + + +def _build_prompt_payload(tokenizer, args): + case = CASE_DEFAULTS[args.case] + paths = _ref_audio_paths(args) + prompt_text = args.ref_text if args.ref_text is not None else case.get("prompt_text") + if case.get("requires_ref_text") and not prompt_text: + raise RuntimeError(f"--ref-text required for case '{args.case}'") + prompt_waveform = _resolve_reference_inputs(args, case, paths) if prompt_text is not None else None + speaker_embedding = _resolve_speaker_embedding(args, case, paths) + use_zero_spk_emb = bool(case.get("use_zero_spk_emb")) and prompt_waveform is None and speaker_embedding is None + runtime_controls = {KEY_MAX_DECODE_STEPS: args.max_decode_steps or case["max_decode_steps"]} + for key, field in [(KEY_CFG, "cfg"), (KEY_SIGMA, "sigma"), (KEY_TEMPERATURE, "temperature")]: + if field in case: + runtime_controls[key] = case[field] + instruction = ( + json.loads(args.instruction_json) if args.instruction_json else (args.instructions or case.get("instruction")) + ) + return build_ming_dense_prompt( + tokenizer, + prompt=args.prompt or case["prompt"], + text=args.text or case["text"], + runtime_controls=runtime_controls, + instruction=instruction, + prompt_text=prompt_text, + prompt_waveform=prompt_waveform, + speaker_embedding=speaker_embedding, + use_zero_spk_emb=use_zero_spk_emb, + ) + + +def main(): + args = _build_parser().parse_args() + _finalize_args(args) + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=False) + prompt_payload = _build_prompt_payload(tokenizer, args) + case = CASE_DEFAULTS[args.case] + sampling_params_list = build_sampling_params(args.max_decode_steps or case["max_decode_steps"]) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + stats_log_file = resolve_stats_log_file(args) + summaries = run_generation(args, prompt_payload, sampling_params_list, output_dir, stats_log_file) + metadata_json = resolve_metadata_json(args) + manifest = build_manifest(args, prompt_payload, stats_log_file, summaries) + if metadata_json is not None: + Path(metadata_json).write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8") + print(f"Saved run manifest to {metadata_json}") + + +if __name__ == "__main__": + os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + main() diff --git a/examples/offline_inference/text_to_speech/ming_tts/runner.py b/examples/offline_inference/text_to_speech/ming_tts/runner.py new file mode 100644 index 000000000..35852ce04 --- /dev/null +++ b/examples/offline_inference/text_to_speech/ming_tts/runner.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +import json +import time +import uuid +import wave +from pathlib import Path + +import torch +from vllm import SamplingParams + +from vllm_omni import AsyncOmni, Omni +from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( + KEY_SPEAKER_EMBEDDING, + SAMPLE_RATE, + TEXT_EOS_TOKEN_ID, +) + + +def coerce_audio_tensor(audio, *, async_chunk: bool) -> torch.Tensor: + if isinstance(audio, list): + parts = [] + for item in audio: + tensor = torch.as_tensor(item, dtype=torch.float32).reshape(-1) + if tensor.numel() > 0: + parts.append(tensor) + if not parts: + return torch.zeros((0,), dtype=torch.float32) + return torch.cat(parts, dim=0) + + return torch.as_tensor(audio, dtype=torch.float32).reshape(-1) + + +def resolve_sr(sr) -> int: + if isinstance(sr, list): + sr = sr[-1] + if hasattr(sr, "item"): + return int(sr.item()) + return int(sr) + + +def extract_sample_rate(multimodal_output: dict) -> int: + sr = multimodal_output.get("sr") + if sr is None: + raise RuntimeError("Expected multimodal_output['sr']") + return resolve_sr(sr) + + +def write_wav(path: str, audio: torch.Tensor, sample_rate: int) -> None: + audio = audio.clamp(-1.0, 1.0) + pcm16 = (audio * 32767.0).round().to(torch.int16).cpu().numpy() + with wave.open(path, "wb") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(2) + wav_file.setframerate(int(sample_rate)) + wav_file.writeframes(pcm16.tobytes()) + + +def request_index(request_id: str | None, fallback: int) -> int: + try: + return int(request_id) + except (TypeError, ValueError): + if isinstance(request_id, str): + head = request_id.split("_", 1)[0] + if head.isdigit(): + return int(head) + return fallback + + +def audio_summary(audio: torch.Tensor, sample_rate: int) -> dict: + waveform = audio.detach().cpu().reshape(-1).to(torch.float32) + return { + "sample_rate": int(sample_rate), + "num_samples": int(waveform.numel()), + "duration_seconds": float(waveform.numel()) / float(sample_rate), + "max_abs_amplitude": float(waveform.abs().max().item()) if waveform.numel() > 0 else 0.0, + } + + +def resolve_output_name(output_name: str | None, case: str, index: int, total: int) -> str: + if total == 1: + return output_name or f"ming_{case}.wav" + base = Path(output_name or f"ming_{case}.wav") + return f"{base.stem}_{index:05d}{base.suffix or '.wav'}" + + +def resolve_stats_log_file(args) -> str | None: + if not args.log_stats: + return None + if args.stats_log_file: + return args.stats_log_file + base = Path(args.output_name or f"ming_{args.case}.wav").stem + return str(Path(args.output_dir) / f"{base}_pipeline.log") + + +def resolve_metadata_json(args) -> str | None: + if args.metadata_json: + return args.metadata_json + if args.log_stats: + base = Path(args.output_name or f"ming_{args.case}.wav").stem + return str(Path(args.output_dir) / f"{base}_manifest.json") + return None + + +def build_manifest(args, prompt_payload, stats_log_file: str | None, outputs: list[dict]) -> dict: + additional_information = {} + if isinstance(prompt_payload, dict): + additional_information = dict(prompt_payload.get("additional_information", {})) + return { + "model": args.model, + "case": args.case, + "streaming": bool(args.streaming), + "deploy_config": args.deploy_config, + "enforce_eager": bool(args.enforce_eager), + "num_prompts": int(args.num_prompts), + "log_stats": bool(args.log_stats), + "stats_log_file": stats_log_file, + "prompt_text": additional_information.get("prompt_text"), + "instruction": additional_information.get("instruction"), + "speaker_embedding_shape": ( + list(additional_information[KEY_SPEAKER_EMBEDDING].shape) + if KEY_SPEAKER_EMBEDDING in additional_information + and hasattr(additional_information[KEY_SPEAKER_EMBEDDING], "shape") + else None + ), + "outputs": outputs, + "generated_at_unix": time.time(), + } + + +def build_engine_kwargs(args, stats_log_file: str | None) -> dict: + kwargs = { + "model": args.model, + "deploy_config": args.deploy_config, + "enforce_eager": args.enforce_eager, + "trust_remote_code": args.trust_remote_code, + "log_stats": args.log_stats, + "stage_init_timeout": args.stage_init_timeout, + "init_timeout": args.init_timeout, + "batch_timeout": args.batch_timeout, + "shm_threshold_bytes": args.shm_threshold_bytes, + "worker_backend": args.worker_backend, + } + if stats_log_file is not None: + kwargs["log_file"] = stats_log_file + if args.ray_address is not None: + kwargs["ray_address"] = args.ray_address + return kwargs + + +def build_sampling_params(max_decode_steps: int) -> list[SamplingParams]: + return [ + SamplingParams( + temperature=0.0, + max_tokens=max_decode_steps + 1, + stop_token_ids=[int(TEXT_EOS_TOKEN_ID)], + ), + SamplingParams(temperature=0.0, max_tokens=1), + ] + + +async def run_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None): + engine = AsyncOmni(**build_engine_kwargs(args, stats_log_file)) + try: + all_audio_chunks = [] + accumulated_samples = 0 + chunk_idx = 0 + start_time = time.time() + chunk_times = [] + ttfp_seconds = None + final_stage_output = None + async for stage_output in engine.generate( + prompt=prompt_payload, + request_id=str(uuid.uuid4()), + sampling_params_list=sampling_params_list, + ): + final_stage_output = stage_output + multimodal_output = stage_output.multimodal_output or {} + audio = multimodal_output.get("audio") + if audio is None: + continue + + finished = stage_output.finished + if isinstance(audio, torch.Tensor): + if finished: + audio_chunk = audio[accumulated_samples:].float().detach().cpu() + else: + audio_chunk = audio.float().detach().cpu() + elif isinstance(audio, list): + audio_chunk = torch.as_tensor(audio[chunk_idx], dtype=torch.float32).reshape(-1).cpu() + else: + audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu() + + accumulated_samples += int(audio_chunk.numel()) + chunk_idx += 1 + if audio_chunk.numel() > 0: + now = time.time() + if ttfp_seconds is None: + ttfp_seconds = now - start_time + chunk_times.append(now) + all_audio_chunks.append(audio_chunk) + + if not all_audio_chunks: + raise RuntimeError("Streaming Ming example produced no audio chunks") + + waveform = torch.cat(all_audio_chunks, dim=0) + output_name = resolve_output_name(args.output_name, args.case, 0, 1) + output_path = str(output_dir / output_name) + write_wav(output_path, waveform, SAMPLE_RATE) + summary = { + "request_id": getattr(final_stage_output, "request_id", None), + "stage_id": getattr(final_stage_output, "stage_id", None), + "output_path": output_path, + "stage_durations": getattr(final_stage_output, "stage_durations", {}), + "peak_memory_mb": getattr(final_stage_output, "peak_memory_mb", 0.0), + "ttfp_seconds": ttfp_seconds, + "mean_inter_chunk_seconds": ( + sum(t1 - t0 for t0, t1 in zip(chunk_times, chunk_times[1:])) / (len(chunk_times) - 1) + if len(chunk_times) > 1 + else None + ), + } + summary.update(audio_summary(waveform, SAMPLE_RATE)) + print(f"Saved streaming output to {output_path}") + print(json.dumps(summary, ensure_ascii=False, indent=2)) + return [summary] + finally: + engine.shutdown() + + +def run_non_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None): + engine = Omni(**build_engine_kwargs(args, stats_log_file)) + try: + outputs = engine.generate( + prompts=[prompt_payload for _ in range(args.num_prompts)], + sampling_params_list=sampling_params_list, + py_generator=False, + ) + summaries = [] + for fallback_index, output in enumerate(outputs): + if output.final_output_type != "audio": + continue + multimodal_output = output.multimodal_output or {} + waveform = coerce_audio_tensor(multimodal_output.get("audio"), async_chunk=False) + sample_rate = extract_sample_rate(multimodal_output) + output_name = resolve_output_name( + args.output_name, + args.case, + request_index(output.request_id, fallback_index), + args.num_prompts, + ) + output_path = str(output_dir / output_name) + write_wav(output_path, waveform, sample_rate) + summary = { + "request_id": output.request_id, + "stage_id": output.stage_id, + "output_path": output_path, + "stage_durations": output.stage_durations, + "peak_memory_mb": output.peak_memory_mb, + } + summary.update(audio_summary(waveform, sample_rate)) + summaries.append(summary) + print(f"Saved output to {output_path}") + print(json.dumps(summary, ensure_ascii=False, indent=2)) + if not summaries: + raise RuntimeError("Non-streaming Ming example produced no audio outputs") + return summaries + finally: + engine.close() + + +def run_generation(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None): + if args.streaming: + return asyncio.run(run_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file)) + return run_non_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file) diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index a3b27b722..c9498596e 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -16,6 +16,7 @@ For the full list of supported architectures across all modalities, see |---|---|---|---|---|---| | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ | +| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | ✓ (`ref_audio` / `speaker_embedding`) | ✓ (PCM stream) | IP labels + structured `instructions` | — | | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — | | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ | | OmniVoice | `k2-fsa/OmniVoice` | ✓ | — | — | — | @@ -202,6 +203,51 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server --- +## Ming-omni-tts + +Dense 0.5B two-stage TTS served through `/v1/audio/speech`. Ming uses the standard speech endpoint plus structured controls in `instructions`, `voice`, `language`, `ref_audio`, `ref_text`, and `speaker_embedding`. + +### Launch +```bash +bash examples/online_serving/text_to_speech/ming_tts/run_server.sh +``` +Equivalent manual command: +```bash +vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --host 0.0.0.0 --port 8091 \ + --enforce-eager --omni +``` + +### Sending requests +```bash +python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \ + --text "你好，这是 Ming 在线语音合成测试。" +``` + +Structured dialect control: +```bash +python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \ + --text "我觉得社会企业同个人都有责任" \ + --instruction-json '{"方言":"广粤话"}' \ + --ref-audio /path/to/yue_prompt.wav +``` + +Zero-shot cloning: +```bash +python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \ + --text "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。" \ + --ref-audio /path/to/10002287-00000094.wav \ + --ref-text "在此奉劝大家别乱打美白针。" +``` + +### Notes +- `run_curl.sh` keeps a small sanity subset; use the Ming README for the broader request cookbook. +- Online serving is speech-shaped today; music-only `bgm` and text-to-audio `tta` remain offline examples. +- Full request details live in [`ming_tts/README.md`](ming_tts/README.md). + +--- + ## Ming-flash-omni-TTS Standalone talker-only deployment of Ming-flash-omni-2.0. Voice is controlled through caption text passed via `instructions`. diff --git a/examples/online_serving/text_to_speech/ming_tts/README.md b/examples/online_serving/text_to_speech/ming_tts/README.md new file mode 100644 index 000000000..d0e0b0194 --- /dev/null +++ b/examples/online_serving/text_to_speech/ming_tts/README.md @@ -0,0 +1,113 @@ +# Ming-omni-tts Online Serving + +Serve the dense `inclusionAI/Ming-omni-tts-0.5B` two-stage TTS model through +the OpenAI-compatible `/v1/audio/speech` endpoint. + +## Start Server + +```bash +vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --omni \ + --port 8091 \ + --enforce-eager +``` + +Or: + +```bash +cd examples/online_serving/text_to_speech/ming_tts +./run_server.sh +``` + +The tested ROCm environment is summarized in the +[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md). + +## Send Requests + +The Python client targets `http://localhost:8091/v1` with `api_key=EMPTY`; it +does not call OpenAI's hosted API. + +```bash +python openai_speech_client.py \ + --text "你好，这是 Ming 在线语音合成测试。" \ + --max-new-tokens 200 +``` + +Style or dialect controls can be plain text or Ming JSON. The upstream +dialect example also uses `yue_prompt.wav` for speaker conditioning: + +```bash +python openai_speech_client.py \ + --text "我觉得社会企业同个人都有责任" \ + --instruction-json '{"方言":"广粤话"}' \ + --ref-audio /path/to/yue_prompt.wav \ + --max-new-tokens 200 +``` + +When `--ref-audio` is supplied without `--ref-text`, the server extracts the +Ming speaker embedding, matching upstream `use_spk_emb=True`, without using the +audio as a zero-shot prompt. + +Reference-audio cloning: + +```bash +python openai_speech_client.py \ + --text "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。" \ + --ref-audio /path/to/10002287-00000094.wav \ + --ref-text "在此奉劝大家别乱打美白针。" \ + --max-new-tokens 200 +``` + +Podcast-style multi-speaker prompt: + +```bash +python openai_speech_client.py \ + --text " speaker_1:你可以说一下，就大概说一下，可能虽然我也不知道，我看过那部电影没有。 + speaker_2:就是那个叫什么，变相一节课的嘛。 + speaker_1:嗯。 + speaker_2:一部搞笑的电影。 + speaker_1:一部搞笑的。" \ + --ref-audio /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav \ + --ref-audio /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \ + --ref-text " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试，做个，当服务员还要去笔试了 + speaker_2:对啊，这真的很奇怪，就是 单纯的因，单纯自己工资不高，只是因为可能人家那个店比较出名一点，就对你苛刻要求" +``` + +Streaming PCM: + +```bash +python openai_speech_client.py \ + --text "你好，这是流式输出测试。" \ + --stream \ + --output ming_output.pcm +``` + +`run_curl.sh` keeps small smoke checks: + +```bash +./run_curl.sh basic +REF_AUDIO=/path/to/reference.wav REF_TEXT="在此奉劝大家别乱打美白针。" ./run_curl.sh zero_shot +./run_curl.sh stream +``` + +## Request Fields + +| Field | Ming meaning | +|-------|--------------| +| `input` | target text | +| `instructions` | plain style text, or JSON object for structured Ming controls | +| `voice` | Ming IP voice label unless it resolves to an uploaded speaker | +| `language` | Ming `方言` control | +| `ref_audio` | speaker reference; with `ref_text`, also supplies the prompt waveform | +| `ref_text` | transcript enabling zero-shot or podcast prompt-latent conditioning | +| `speaker_embedding` | 192-d Ming speaker embedding | +| `max_new_tokens` | Ming `max_decode_steps` | + +## Notes + +- `ref_audio` accepts local paths through the client, remote URLs, `file://`, + or `data:` URLs. +- Non-streaming responses return WAV bytes; streaming responses return PCM. +- Music-only `bgm` generation is offline-only until the API exposes Ming + prompt-mode selection. diff --git a/examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py b/examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py new file mode 100644 index 000000000..536bfb048 --- /dev/null +++ b/examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""OpenAI-compatible client for Ming-omni-tts via /v1/audio/speech. + +Examples: + python openai_speech_client.py --text "你好，世界" + python openai_speech_client.py --text "我觉得社会企业同个人都有责任" \ + --instruction-json '{"方言":"广粤话"}' --ref-audio yue_prompt.wav --max-new-tokens 200 + python openai_speech_client.py \ + --text "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。" \ + --ref-audio 10002287-00000094.wav --ref-text "在此奉劝大家别乱打美白针。" \ + --max-new-tokens 200 + python openai_speech_client.py --text "speaker_1:你好。 speaker_2:你好。" \ + --ref-audio speaker_1.wav --ref-audio speaker_2.wav --ref-text "speaker_1:你好。 speaker_2:你好。" + python openai_speech_client.py --text "你好，这是流式输出测试。" \ + --stream --output ming_output.pcm +""" + +import argparse +import base64 +import json +import os + +import httpx + +DEFAULT_API_BASE = "http://localhost:8091" +DEFAULT_API_KEY = "EMPTY" +DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B" +EXPECTED_SPEAKER_EMBEDDING_DIM = 192 + + +def encode_audio_to_base64(audio_path: str) -> str: + """Encode a local audio file to a base64 data URL.""" + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + ext = audio_path.lower().rsplit(".", 1)[-1] + mime_map = { + "wav": "audio/wav", + "mp3": "audio/mpeg", + "flac": "audio/flac", + "ogg": "audio/ogg", + "aac": "audio/aac", + } + mime_type = mime_map.get(ext, "audio/wav") + with open(audio_path, "rb") as f: + audio_b64 = base64.b64encode(f.read()).decode("utf-8") + return f"data:{mime_type};base64,{audio_b64}" + + +def load_speaker_embedding(path: str) -> list[float]: + """Load and validate a 192-d Ming speaker embedding JSON file.""" + with open(path, encoding="utf-8") as f: + data = json.load(f) + + if not isinstance(data, list): + raise ValueError("speaker_embedding file must contain a JSON list") + if len(data) != EXPECTED_SPEAKER_EMBEDDING_DIM: + raise ValueError( + f"Ming dense speaker_embedding must have {EXPECTED_SPEAKER_EMBEDDING_DIM} values, got {len(data)}" + ) + + values = [] + for index, value in enumerate(data): + try: + values.append(float(value)) + except (TypeError, ValueError) as exc: + raise ValueError(f"speaker_embedding[{index}] must be a number, got {value!r}") from exc + return values + + +def build_instruction_payload(args) -> str | None: + """Return a string payload for the API `instructions` field.""" + if args.instructions and args.instruction_json: + raise ValueError("Use either --instructions or --instruction-json, not both") + if args.instruction_json: + parsed = json.loads(args.instruction_json) + return json.dumps(parsed, ensure_ascii=False) + return args.instructions + + +def validate_args(args) -> None: + """Fail fast on invalid combinations before hitting the server.""" + if args.ref_text and not args.ref_audio: + raise ValueError("--ref-audio is required when --ref-text is provided") + if args.speaker_embedding and args.ref_audio and len(args.ref_audio) > 1: + raise ValueError("--speaker-embedding cannot be combined with multiple --ref-audio values") + + +def run_tts(args) -> None: + """Generate speech via the OpenAI-compatible /v1/audio/speech API.""" + validate_args(args) + + payload = { + "model": args.model, + "input": args.text, + "response_format": args.response_format, + } + + if args.voice: + payload["voice"] = args.voice + if args.task_type: + payload["task_type"] = args.task_type + if args.dialect: + payload["language"] = args.dialect + + instructions = build_instruction_payload(args) + if instructions: + payload["instructions"] = instructions + + if args.ref_audio: + ref_audio = [] + for audio in args.ref_audio: + if audio.startswith(("http://", "https://", "data:")): + ref_audio.append(audio) + else: + ref_audio.append(encode_audio_to_base64(audio)) + payload["ref_audio"] = ref_audio[0] if len(ref_audio) == 1 else ref_audio + if args.ref_text: + payload["ref_text"] = args.ref_text + if args.speaker_embedding: + payload["speaker_embedding"] = load_speaker_embedding(args.speaker_embedding) + if args.max_new_tokens: + payload["max_new_tokens"] = args.max_new_tokens + if args.stream: + payload["stream"] = True + payload["response_format"] = "pcm" + + api_url = f"{args.api_base}/v1/audio/speech" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {args.api_key}", + } + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + print(f"Payload keys: {sorted(payload)}") + + if args.stream: + output_path = args.output or "ming_output.pcm" + with httpx.Client(timeout=300.0) as client: + with client.stream("POST", api_url, json=payload, headers=headers) as response: + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.read().decode()) + return + with open(output_path, "wb") as f: + for chunk in response.iter_bytes(): + f.write(chunk) + print(f"Streamed PCM audio to: {output_path}") + return + + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.text) + return + + try: + text = response.content.decode("utf-8") + if text.startswith('{"error"'): + print(f"Error: {text}") + return + except UnicodeDecodeError: + pass + + output_path = args.output or "ming_output.wav" + with open(output_path, "wb") as f: + f.write(response.content) + print(f"Audio saved to: {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="OpenAI-compatible client for Ming-omni-tts via /v1/audio/speech") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name or path") + parser.add_argument("--text", required=True, help="Text to synthesize") + parser.add_argument( + "--task-type", + default=None, + choices=["CustomVoice", "VoiceDesign", "Base"], + help="Optional compatibility task type. Ming accepts the same field but primarily uses prompt metadata.", + ) + parser.add_argument( + "--voice", + default=None, + help="Maps to Ming `IP` when using built-in character voices, or to an uploaded voice sample name", + ) + parser.add_argument("--dialect", default=None, help="Maps to Ming `方言`") + parser.add_argument("--instructions", default=None, help="Free-form Ming instruction string") + parser.add_argument( + "--instruction-json", + default=None, + help='Structured Ming instruction JSON, for example \'{"情感":"高兴"}\'', + ) + parser.add_argument( + "--ref-audio", + action="append", + default=None, + help="Reference audio path, URL, or data URL. Repeat for podcast-style multi-speaker prompts.", + ) + parser.add_argument("--ref-text", default=None, help="Reference transcript for cloning") + parser.add_argument( + "--speaker-embedding", default=None, help="Path to a JSON file containing a 192-d speaker embedding" + ) + parser.add_argument("--max-new-tokens", type=int, default=None, help="Override ming_max_decode_steps") + parser.add_argument("--stream", action="store_true", help="Enable streaming PCM output") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "mp3", "flac", "pcm", "aac", "opus"], + help="Audio format when not streaming", + ) + parser.add_argument("--output", "-o", default=None, help="Output file path") + args = parser.parse_args() + try: + run_tts(args) + except (FileNotFoundError, ValueError, json.JSONDecodeError) as exc: + raise SystemExit(f"Error: {exc}") from exc + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/text_to_speech/ming_tts/run_curl.sh b/examples/online_serving/text_to_speech/ming_tts/run_curl.sh new file mode 100755 index 000000000..10fcaf59a --- /dev/null +++ b/examples/online_serving/text_to_speech/ming_tts/run_curl.sh @@ -0,0 +1,76 @@ +#!/bin/bash +set -euo pipefail + +MODE="${1:-basic}" +HOST="${HOST:-localhost}" +PORT="${PORT:-8091}" +MODEL="${MODEL:-inclusionAI/Ming-omni-tts-0.5B}" +API_URL="http://${HOST}:${PORT}/v1/audio/speech" +TEXT="${TEXT:-你好，这是 Ming 在线语音合成测试。}" +OUTPUT="${OUTPUT:-ming_output.wav}" +STREAM_OUTPUT="${STREAM_OUTPUT:-ming_output.pcm}" +REF_AUDIO="${REF_AUDIO:-}" +REF_TEXT="${REF_TEXT:-}" + +post_json() { + local payload="$1" + local output_path="$2" + curl -X POST "$API_URL" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer EMPTY" \ + -d "$payload" \ + --output "$output_path" +} + +case "$MODE" in + basic) + post_json "{ + \"model\": \"${MODEL}\", + \"input\": \"${TEXT}\", + \"response_format\": \"wav\" + }" "$OUTPUT" + ;; + zero_shot) + if [ -z "$REF_AUDIO" ] || [ -z "$REF_TEXT" ]; then + echo "zero_shot requires REF_AUDIO and REF_TEXT" >&2 + exit 1 + fi + python - <<'PY' > /tmp/ming_zero_shot_payload.json +import base64 +import json +import mimetypes +import os +from pathlib import Path + +path = Path(os.environ["REF_AUDIO"]) +mime_type = mimetypes.guess_type(path.name)[0] or "audio/wav" +payload = { + "model": os.environ["MODEL"], + "input": os.environ["TEXT"], + "ref_audio": f"data:{mime_type};base64,{base64.b64encode(path.read_bytes()).decode('utf-8')}", + "ref_text": os.environ["REF_TEXT"], + "response_format": "wav", +} +print(json.dumps(payload, ensure_ascii=False)) +PY + curl -X POST "$API_URL" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer EMPTY" \ + --data-binary @/tmp/ming_zero_shot_payload.json \ + --output "$OUTPUT" + rm -f /tmp/ming_zero_shot_payload.json + ;; + stream) + post_json "{ + \"model\": \"${MODEL}\", + \"input\": \"${TEXT}\", + \"stream\": true, + \"response_format\": \"pcm\" + }" "$STREAM_OUTPUT" + ;; + *) + echo "Unknown mode: $MODE" >&2 + echo "Supported sanity checks: basic, zero_shot, stream" >&2 + exit 1 + ;; +esac diff --git a/examples/online_serving/text_to_speech/ming_tts/run_server.sh b/examples/online_serving/text_to_speech/ming_tts/run_server.sh new file mode 100755 index 000000000..7d294d950 --- /dev/null +++ b/examples/online_serving/text_to_speech/ming_tts/run_server.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# Launch vLLM-Omni server for Ming-omni-tts. +# +# Usage: +# ./run_server.sh +# PORT=8000 ./run_server.sh + +set -e + +DIR="$(cd "$(dirname "$0")" && pwd)" +ROOT="$(cd "$DIR/../../../.." && pwd)" + +MODEL="${MODEL:-inclusionAI/Ming-omni-tts-0.5B}" +PORT="${PORT:-8091}" +DEPLOY_CONFIG="${DEPLOY_CONFIG:-$ROOT/vllm_omni/deploy/ming_tts.yaml}" + +echo "Starting Ming-omni-tts server with model: $MODEL" +echo "Deploy config: $DEPLOY_CONFIG" + +vllm-omni serve "$MODEL" \ + --deploy-config "$DEPLOY_CONFIG" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --enforce-eager \ + --omni diff --git a/recipes/README.md b/recipes/README.md index a0833827a..2bbc518df 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -33,6 +33,7 @@ recipes/ | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB | | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 | | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB | +| [`inclusionAI/Ming-omni-tts-0.5B.md`](./inclusionAI/Ming-omni-tts-0.5B.md) | Offline + online dense Ming TTS/audio generation | 1x AMD `gfx942` (ROCm 7.2) | | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB | | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) | | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB | diff --git a/recipes/inclusionAI/Ming-omni-tts-0.5B.md b/recipes/inclusionAI/Ming-omni-tts-0.5B.md new file mode 100644 index 000000000..cd1728034 --- /dev/null +++ b/recipes/inclusionAI/Ming-omni-tts-0.5B.md @@ -0,0 +1,82 @@ +# Ming-omni-tts 0.5B + +> Offline and online dense Ming text-to-speech on ROCm + +## Summary + +- Vendor: inclusionAI +- Model: `inclusionAI/Ming-omni-tts-0.5B` +- Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls +- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API; offline inference +- Maintainer: Community + +## References + +- [Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B) +- [Upstream repository](https://github.com/inclusionAI/Ming-omni-tts) +- [Offline example](../../examples/offline_inference/text_to_speech/ming_tts/) +- [Online example](../../examples/online_serving/text_to_speech/ming_tts/) + +## Hardware Support + +## ROCm + +### 1x AMD `gfx942` + +#### Environment + +- OS: Ubuntu 22.04.5 LTS, x86_64 +- Python: 3.12.13 +- ROCm / HIP: 7.2.53211 +- vLLM version: 0.22.0 +- vLLM-Omni version or commit: 0.1.dev1873 / `99c35c410` +- Docker image: `vllm/vllm-omni-rocm:v0.22.0` + +#### Command + +From the vLLM-Omni repository root: + +```bash +docker run --rm \ + --group-add=video \ + --ipc=host \ + --cap-add=SYS_PTRACE \ + --security-opt seccomp=unconfined \ + --device /dev/kfd \ + --device /dev/dri \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + -v "$PWD":/app/vllm-omni \ + -w /app/vllm-omni \ + -e VLLM_ROCM_USE_AITER=0 \ + -p 8091:8091 \ + vllm/vllm-omni-rocm:v0.22.0 \ + --model inclusionAI/Ming-omni-tts-0.5B \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --omni \ + --port 8091 \ + --enforce-eager +``` + +#### Verification + +```bash +python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \ + --text "我觉得社会企业同个人都有责任" \ + --instruction-json '{"方言":"广粤话"}' \ + --ref-audio /path/to/yue_prompt.wav \ + --max-new-tokens 200 \ + --output dialect.wav +``` + +`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text` +for the dialect case. + +## Notes + +- The official ROCm image includes the platform dependencies. +- See the [ROCm installation guide](../../docs/getting_started/installation/gpu.md) + for interactive and source-build workflows. +- The [offline example](../../examples/offline_inference/text_to_speech/ming_tts/) + can be run from an interactive container. +- The tested environment uses `--enforce-eager`. +- Output is mono 44.1 kHz audio. diff --git a/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py b/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py index a81cfe6d4..19cf6fa4c 100644 --- a/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py +++ b/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py @@ -51,7 +51,7 @@ test_params = [(m, c) for m in models for c in [get_eager_config()]] @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) def test_text_to_text(omni_runner, omni_runner_handler) -> None: """ @@ -67,7 +67,7 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) def test_image_to_text(omni_runner, omni_runner_handler) -> None: """ @@ -84,7 +84,7 @@ def test_image_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) def test_audio_to_text(omni_runner, omni_runner_handler) -> None: """ @@ -103,7 +103,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) def test_video_to_text(omni_runner, omni_runner_handler) -> None: """ @@ -120,7 +120,7 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) def test_mixed_to_text(omni_runner, omni_runner_handler) -> None: """ @@ -140,7 +140,7 @@ def test_mixed_to_text(omni_runner, omni_runner_handler) -> None: @pytest.mark.advanced_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) def test_text_to_audio(omni_runner, omni_runner_handler) -> None: """ @@ -156,7 +156,7 @@ def test_text_to_audio(omni_runner, omni_runner_handler) -> None: @pytest.mark.advanced_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=4) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) def test_image_to_audio(omni_runner, omni_runner_handler) -> None: """ diff --git a/tests/e2e/offline_inference/test_ming_tts.py b/tests/e2e/offline_inference/test_ming_tts.py new file mode 100644 index 000000000..ca76c9afe --- /dev/null +++ b/tests/e2e/offline_inference/test_ming_tts.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""End-to-end offline inference tests for Ming-omni-tts.""" + +import asyncio +import uuid + +import numpy as np +import pytest +import torch +from transformers import AutoTokenizer +from vllm import SamplingParams +from vllm.inputs import tokens_input + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni import AsyncOmni +from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( + KEY_MAX_DECODE_STEPS, + SAMPLE_RATE, + TEXT_EOS_TOKEN_ID, +) +from vllm_omni.model_executor.models.ming_tts.prompt_assembly import DEFAULT_PROMPT, build_ming_dense_prompt + +MODEL = "inclusionAI/Ming-omni-tts-0.5B" +DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml") +TEST_TEXT = "我会一直在这里陪着你，直到你慢慢地沉入那个最温柔的梦里。" +TEST_INSTRUCTION = "轻柔的ASMR耳语，慢速，贴近麦克风" +MIN_AUDIO_SAMPLES = 1000 + + +@pytest.fixture(scope="module") +def ming_tokenizer(): + return AutoTokenizer.from_pretrained(MODEL, trust_remote_code=False) + + +@pytest.fixture +def ming_engine(): + with OmniRunner( + MODEL, + deploy_config=DEPLOY_CONFIG, + stage_init_timeout=300, + enforce_eager=True, + ) as runner: + yield runner.omni + + +@pytest.fixture +def async_omni_engine(): + engine = AsyncOmni( + model=MODEL, + deploy_config=DEPLOY_CONFIG, + stage_init_timeout=300, + enforce_eager=True, + ) + yield engine + engine.shutdown() + + +def _build_prompt( + tokenizer, + *, + text: str = TEST_TEXT, + instruction=TEST_INSTRUCTION, + use_zero_spk_emb: bool = True, +) -> dict: + prompt_dict = build_ming_dense_prompt( + tokenizer, + prompt=DEFAULT_PROMPT, + text=text, + instruction=instruction, + runtime_controls={KEY_MAX_DECODE_STEPS: 200}, + use_zero_spk_emb=use_zero_spk_emb, + ) + prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"]) + prompt["prompt"] = prompt_dict["prompt"] + prompt["text"] = prompt_dict["text"] + prompt["additional_information"] = prompt_dict["additional_information"] + return prompt + + +def _sampling_params_list() -> list[SamplingParams]: + return [ + SamplingParams( + temperature=0.0, + max_tokens=201, + stop_token_ids=[int(TEXT_EOS_TOKEN_ID)], + ), + SamplingParams(temperature=0.0, max_tokens=1), + ] + + +def _flatten_audio(audio) -> torch.Tensor: + if isinstance(audio, list): + parts = [torch.as_tensor(item, dtype=torch.float32).reshape(-1).cpu() for item in audio] + parts = [item for item in parts if item.numel() > 0] + if not parts: + return torch.zeros((0,), dtype=torch.float32) + return torch.cat(parts, dim=0) + return torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu() + + +def _extract_audio(multimodal_output: dict) -> torch.Tensor: + assert isinstance(multimodal_output, dict), f"Expected dict, got {type(multimodal_output)}" + audio = multimodal_output.get("audio", multimodal_output.get("model_outputs")) + assert audio is not None, f"No audio output found, keys={list(multimodal_output.keys())}" + waveform = _flatten_audio(audio) + if waveform.numel() == 0: + raise RuntimeError("Generated audio waveform is empty") + return waveform + + +def _extract_sample_rate(multimodal_output: dict) -> int: + sample_rate = multimodal_output.get("sr") + if sample_rate is None: + raise RuntimeError("Expected multimodal_output['sr']") + if isinstance(sample_rate, list): + sample_rate = sample_rate[-1] + if hasattr(sample_rate, "item"): + sample_rate = sample_rate.item() + return int(sample_rate) + + +def _extract_final_audio_outputs(outputs): + final_outputs = [] + for item in outputs: + if getattr(item, "final_output_type", None) == "audio": + final_outputs.append(item) + continue + request_output = getattr(item, "request_output", None) + if request_output is None: + continue + multimodal_output = getattr(request_output, "multimodal_output", None) + if isinstance(multimodal_output, dict): + final_outputs.append(item) + continue + completions = getattr(request_output, "outputs", None) or [] + if any(isinstance(getattr(completion, "multimodal_output", None), dict) for completion in completions): + final_outputs.append(item) + return final_outputs + + +def _extract_multimodal_output(output) -> dict: + multimodal_output = getattr(output, "multimodal_output", None) + if isinstance(multimodal_output, dict): + return multimodal_output + + request_output = getattr(output, "request_output", None) + if request_output is not None: + multimodal_output = getattr(request_output, "multimodal_output", None) + if isinstance(multimodal_output, dict): + return multimodal_output + completions = getattr(request_output, "outputs", None) or [] + for completion in completions: + multimodal_output = getattr(completion, "multimodal_output", None) + if isinstance(multimodal_output, dict): + return multimodal_output + + raise AssertionError("No multimodal audio output found in Ming generate results") + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_ming_tts_offline_basic(ming_engine, ming_tokenizer) -> None: + """Test blocking Ming generation through Omni.""" + outputs = ming_engine.generate( + prompts=[_build_prompt(ming_tokenizer)], + sampling_params_list=_sampling_params_list(), + py_generator=False, + ) + final_outputs = _extract_final_audio_outputs(outputs) + assert len(final_outputs) == 1, f"Expected one final audio output, got {len(final_outputs)}" + multimodal_output = _extract_multimodal_output(final_outputs[0]) + waveform = _extract_audio(multimodal_output) + sample_rate = _extract_sample_rate(multimodal_output) + assert waveform.ndim == 1 + assert waveform.shape[0] == waveform.numel() + assert waveform.numel() > MIN_AUDIO_SAMPLES + assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent" + assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}" + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_ming_tts_speaker_conditioning_differs(ming_engine, ming_tokenizer) -> None: + """Test that different Ming speaker controls produce different waveform outputs.""" + style_outputs = ming_engine.generate( + prompts=[_build_prompt(ming_tokenizer)], + sampling_params_list=_sampling_params_list(), + py_generator=False, + ) + ip_outputs = ming_engine.generate( + prompts=[_build_prompt(ming_tokenizer, text=TEST_TEXT, instruction={"IP": "灵小甄"}, use_zero_spk_emb=True)], + sampling_params_list=_sampling_params_list(), + py_generator=False, + ) + + style_final_outputs = _extract_final_audio_outputs(style_outputs) + ip_final_outputs = _extract_final_audio_outputs(ip_outputs) + assert len(style_final_outputs) == 1, "No style audio output produced" + assert len(ip_final_outputs) == 1, "No IP audio output produced" + + style_waveform = _extract_audio(_extract_multimodal_output(style_final_outputs[0])) + ip_waveform = _extract_audio(_extract_multimodal_output(ip_final_outputs[0])) + assert style_waveform.numel() > MIN_AUDIO_SAMPLES + assert ip_waveform.numel() > MIN_AUDIO_SAMPLES + assert np.max(np.abs(style_waveform.numpy())) > 0.01, "Style audio appears silent" + assert np.max(np.abs(ip_waveform.numpy())) > 0.01, "IP audio appears silent" + + overlap = min(int(style_waveform.numel()), int(ip_waveform.numel())) + mean_abs_diff = torch.mean(torch.abs(style_waveform[:overlap] - ip_waveform[:overlap])).item() + assert style_waveform.shape != ip_waveform.shape or mean_abs_diff > 1e-4, ( + "Speaker-conditioned outputs should differ, but style and IP waveforms were effectively identical" + ) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_ming_tts_multiple_prompts_queued(ming_engine, ming_tokenizer) -> None: + """Regression: supported max_num_seqs=1 config must still drain queued prompts.""" + prompts = [ + _build_prompt( + ming_tokenizer, + text="第一条语音用于验证 Ming dense 队列中的长请求可以完成。", + instruction="平静自然的旁白，语速中等", + ), + _build_prompt( + ming_tokenizer, + text="第二条短语音也必须生成完成。", + instruction="清晰明亮的女声", + ), + ] + outputs = ming_engine.generate( + prompts=prompts, + sampling_params_list=_sampling_params_list(), + py_generator=False, + ) + final_outputs = _extract_final_audio_outputs(outputs) + assert len(final_outputs) == len(prompts), f"Expected {len(prompts)} audio outputs, got {len(final_outputs)}" + for i, output in enumerate(final_outputs): + waveform = _extract_audio(_extract_multimodal_output(output)) + duration_s = waveform.numel() / SAMPLE_RATE + assert 0.1 < duration_s < 30.0, f"Request {i} audio duration out of range: {duration_s:.2f}s" + assert np.max(np.abs(waveform.numpy())) > 0.01, f"Request {i} audio appears silent" + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_ming_tts_offline_streaming(async_omni_engine, ming_tokenizer) -> None: + """Test async_chunk streaming Ming generation through AsyncOmni.""" + + async def _run() -> None: + all_audio_chunks = [] + accumulated_samples = 0 + chunk_idx = 0 + sample_rate = None + async for stage_output in async_omni_engine.generate( + prompt=_build_prompt(ming_tokenizer), + request_id=str(uuid.uuid4()), + sampling_params_list=_sampling_params_list(), + ): + multimodal_output = stage_output.multimodal_output or {} + audio = multimodal_output.get("audio") + if "sr" in multimodal_output: + sample_rate = _extract_sample_rate(multimodal_output) + if audio is None: + continue + finished = stage_output.finished + if isinstance(audio, torch.Tensor): + if finished: + audio_chunk = audio[accumulated_samples:].float().detach().cpu() + else: + audio_chunk = audio.float().detach().cpu() + elif isinstance(audio, list): + audio_chunk = _flatten_audio(audio) + else: + audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu() + accumulated_samples += int(audio_chunk.numel()) + chunk_idx += 1 + if audio_chunk.numel() > 0: + all_audio_chunks.append(audio_chunk) + assert all_audio_chunks, "No streaming audio chunks received" + waveform = torch.cat(all_audio_chunks, dim=0) + assert waveform.numel() > MIN_AUDIO_SAMPLES + assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent" + assert sample_rate is not None, "Streaming path did not return a sample rate" + assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}" + + asyncio.run(_run()) diff --git a/tests/e2e/online_serving/test_ming_tts.py b/tests/e2e/online_serving/test_ming_tts.py new file mode 100644 index 000000000..03982cc4a --- /dev/null +++ b/tests/e2e/online_serving/test_ming_tts.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E online serving tests for Ming-omni-tts (dense) model. +Tests text-to-audio via /v1/audio/speech endpoint. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +pytestmark = [pytest.mark.advanced_model, pytest.mark.omni] + +MODEL = "inclusionAI/Ming-omni-tts-0.5B" +DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml") + +no_async_chunk_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=DEPLOY_CONFIG, + server_args=["--enforce-eager", "--no-async-chunk"], + ), + id="no_async_chunk", + ) +] + +async_chunk_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=DEPLOY_CONFIG, + server_args=["--enforce-eager"], + ), + id="async_chunk", + ) +] + + +def get_prompt(prompt_type="zh"): + prompts = { + "zh": "今天天气真不错，适合出去散散步。", + "zh_short": "这款产品的名字，叫变态坑爹牛肉丸。", + } + return prompts.get(prompt_type, prompts["zh"]) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", no_async_chunk_params, indirect=True) +def test_text_to_audio_non_streaming_001(omni_server, openai_client) -> None: + """ + Deploy Setting: ming_tts.yaml with --no-async-chunk + Input Modal: text + Output Modal: audio + Input Setting: stream=False + Datasets: two concurrent requests + """ + request_config = { + "model": omni_server.model, + "input": get_prompt("zh"), + "stream": False, + "response_format": "wav", + "timeout": 300.0, + } + openai_client.send_audio_speech_request(request_config, request_num=2) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", async_chunk_params, indirect=True) +def test_text_to_audio_streaming_001(omni_server, openai_client) -> None: + """ + Deploy Setting: ming_tts.yaml (async_chunk=true) + Input Modal: text + voice + Output Modal: audio (streamed) + Input Setting: stream=True + Datasets: single request + """ + request_config = { + "model": omni_server.model, + "input": get_prompt("zh_short"), + "voice": "灵小甄", + "stream": True, + "response_format": "wav", + "timeout": 300.0, + } + openai_client.send_audio_speech_request(request_config) diff --git a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py new file mode 100644 index 000000000..7aab6133b --- /dev/null +++ b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig +from vllm_omni.model_executor.models.common.ming.fm import Solver +from vllm_omni.model_executor.models.ming_tts.constants import ( + AGGREGATOR_HIDDEN_SIZE, + HISTORY_PATCH_SIZE, + LATENT_DIM, + LLM_HIDDEN_SIZE, + LLM_VOCAB_SIZE, + PATCH_SIZE, + SAMPLE_RATE, + VAE_PATCH_SIZE, +) +from vllm_omni.model_executor.models.ming_tts.flowloss_head import FlowLoss +from vllm_omni.model_executor.models.ming_tts.ming_tts_llm import MingLLMModel +from vllm_omni.model_executor.models.ming_tts.validation import validate_ming_tts_config +from vllm_omni.model_executor.models.output_templates import OmniOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts] + + +def test_ming_tts_audio_vae_uses_common_config(): + """AudioVAEConfig is shared by Ming dense and Ming flash modules.""" + cfg = AudioVAEConfig(sample_rate=16000, patch_size=-1) + + assert cfg.sample_rate == 16000 + assert cfg.patch_size == -1 + + +def test_ming_tts_cfm_solver_uses_common_implementation(): + """Ming dense imports the shared solver implementation directly.""" + assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.fm" + + +def test_ming_tts_flowloss_preserves_checkpoint_prefix(): + flowloss = FlowLoss(z_channels=4, llm_cond_dim=8, hidden_size=16, depth=1, num_heads=2) + + assert any(name.startswith("cfm.model.") for name in flowloss.state_dict()) + + +def test_ming_dense_validation_rejects_semantic_audio_vae_config(): + """Dense 0.5B validation rejects semantic AudioVAE configs.""" + cfg = SimpleNamespace( + audio_dummy_token_id=151705, + audio_eos_token_id=151704, + text_eos_token_id=151669, + audio_tokenizer_config=AudioVAEConfig( + sample_rate=SAMPLE_RATE, + patch_size=VAE_PATCH_SIZE, + semantic_module_kwargs={"whisper_encoder": {}}, + enc_kwargs={"latent_dim": LATENT_DIM, "input_dim": 882, "hop_size": 882}, + dec_kwargs={"latent_dim": LATENT_DIM, "output_dim": 882}, + ), + latent_dim=LATENT_DIM, + patch_size=PATCH_SIZE, + history_patch_size=HISTORY_PATCH_SIZE, + llm_hidden_size=LLM_HIDDEN_SIZE, + llm_vocab_size=LLM_VOCAB_SIZE, + sample_rate=SAMPLE_RATE, + vae_patch_size=VAE_PATCH_SIZE, + llm_config={"hidden_size": LLM_HIDDEN_SIZE}, + aggregator_config={"hidden_size": AGGREGATOR_HIDDEN_SIZE}, + ditar_config={"hidden_size": AGGREGATOR_HIDDEN_SIZE}, + latent_chunk_size=1, + latent_left_context=0, + max_decode_steps=1, + stop_head_threshold=0.5, + stop_head_min_steps=0, + ) + + with pytest.raises(ValueError, match="semantic_module_kwargs"): + validate_ming_tts_config(cfg) + + +def test_ming_instruction_parser_preserves_dense_and_flash_defaults(): + """Ming dense and Ming flash keep distinct instruction defaults.""" + serving = object.__new__(OmniOpenAIServingSpeech) + serving.uploaded_speakers = {"uploaded": {}} + + dense_plain = serving._parse_ming_instruction(SimpleNamespace(instructions="calm", language=None, voice=None)) + assert dense_plain == "calm" + + dense_with_fields = serving._parse_ming_instruction( + SimpleNamespace(instructions="calm", language="Auto", voice="灵小甄") + ) + assert dense_with_fields == {"IP": "灵小甄", "风格": "calm"} + + flash_fields = serving._parse_ming_instruction_fields( + SimpleNamespace(instructions="calm", language="粤语", voice="灵小甄") + ) + assert flash_fields == {"风格": "calm"} + + +def _make_ming_logits_model(vocab_size=8): + model = object.__new__(MingLLMModel) + model.ming_config = SimpleNamespace( + llm_vocab_size=vocab_size, + max_decode_steps=1, + stop_head_min_steps=0, + text_eos_token_id=7, + ) + model._last_text_mode = False + model._last_ming_next_token_ids = None + return model + + +def test_ming_compute_logits_uses_cached_forced_next_token_ids(): + model = _make_ming_logits_model() + model._last_ming_next_token_ids = [2, 5] + + logits = MingLLMModel.compute_logits(model, torch.zeros((2, 4)), SimpleNamespace()) + + assert logits.shape == (2, 8) + assert logits[0, 2].item() == 0.0 + assert logits[1, 5].item() == 0.0 + assert torch.isneginf(logits[0, [0, 1, 3, 4, 5, 6, 7]]).all() + assert torch.isneginf(logits[1, [0, 1, 2, 3, 4, 6, 7]]).all() + assert model._last_ming_next_token_ids is None + + +def test_ming_compute_logits_falls_back_to_dummy_token_id(): + model = _make_ming_logits_model() + + logits = MingLLMModel.compute_logits(model, torch.zeros((1, 4)), SimpleNamespace()) + assert logits.shape == (1, 8) + assert logits[0, 7].item() == 0.0 + assert torch.isneginf(logits[0, [0, 1, 2, 3, 4, 5, 6]]).all() + + +def test_ming_forward_non_decode_return_clears_cached_forced_next_token_ids(): + class FakeBackbone: + def __call__(self, **kwargs): + return kwargs["inputs_embeds"] + + model = _make_ming_logits_model() + model.model = FakeBackbone() + model._last_ming_next_token_ids = [2] + + output = MingLLMModel.forward( + model, + input_ids=torch.tensor([1]), + positions=torch.tensor([0]), + inputs_embeds=torch.zeros((1, 4)), + model_intermediate_buffer=[], + ) + + assert isinstance(output, OmniOutput) + assert output.multimodal_outputs is None + assert model._last_ming_next_token_ids is None + + +def test_ming_compute_logits_rejects_forced_token_batch_mismatch(): + model = _make_ming_logits_model() + model._last_ming_next_token_ids = [2] + + with pytest.raises(RuntimeError, match="batch mismatch"): + MingLLMModel.compute_logits(model, torch.zeros((2, 4)), SimpleNamespace()) + + +def test_ming_compute_logits_text_mode_delegates_to_backbone(): + class FakeBackbone: + def __init__(self): + self.hidden_states = None + + def compute_logits(self, hidden_states): + self.hidden_states = hidden_states + return torch.ones((hidden_states.shape[0], 3)) + + model = _make_ming_logits_model(vocab_size=3) + model._last_text_mode = True + model.model = FakeBackbone() + hidden_states = torch.zeros((2, 4)) + + logits = MingLLMModel.compute_logits(model, hidden_states, SimpleNamespace()) + + assert torch.equal(model.model.hidden_states, hidden_states) + assert torch.equal(logits, torch.ones((2, 3))) diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py index 4b2bdf3fe..52ca29905 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py @@ -921,6 +921,7 @@ def test_cosyvoice3_full_payload_replace_keys_present() -> None: def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None: """Smoke: ming_flash_omni token-only carries voice metadata.""" from vllm_omni.model_executor.stage_input_processors.ming_flash_omni import ( + thinker2talker, thinker2talker_token_only, ) @@ -938,14 +939,15 @@ def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None: src = [_Wrap("hello world")] prompt = _Prompt({"voice_name": "ZH_FEMALE", "prompt_text": "ref text"}) - out = thinker2talker_token_only(src, prompt=prompt) - assert len(out) == 1 - assert out[0]["prompt_token_ids"] == [0] # talker self-tokenizes; dummy id - info = out[0]["additional_information"] - assert info["text"] == "hello world" - assert info["voice_name"] == "ZH_FEMALE" - assert info["prompt_text"] == "ref text" - assert info["ming_task"] == "omni" + for func in (thinker2talker, thinker2talker_token_only): + out = func(src, prompt=prompt) + assert len(out) == 1 + assert out[0]["prompt_token_ids"] == [0] + info = out[0]["additional_information"] + assert info["text"] == "hello world" + assert info["voice_name"] == "ZH_FEMALE" + assert info["prompt_text"] == "ref text" + assert info["ming_task"] == "omni" def test_qwen2_5_omni_thinker2talker_token_only_smoke() -> None: diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index a12e6bcd6..106b2bf12 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -101,6 +101,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = { "vllm_omni.model_executor.models.mimo_audio.pipeline", "MIMO_AUDIO_PIPELINE", ), + "ming_tts": ( + "vllm_omni.model_executor.models.ming_tts.pipeline", + "MING_TTS_PIPELINE", + ), "voxtral_tts": ( "vllm_omni.model_executor.models.voxtral_tts.pipeline", "VOXTRAL_TTS_PIPELINE", diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 43e00a142..1bff5ef94 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -1128,6 +1128,19 @@ class StageConfigFactory: continue return cls._create_from_registry(registered.model_type, cli_overrides, deploy_config_path) + # An explicit deploy config can select a pipeline when the model's HF + # metadata is too generic for registry auto-detection. + if deploy_config_path is not None: + deploy_path = Path(deploy_config_path) + if deploy_path.exists(): + deploy_config = load_deploy_config(deploy_path) + if deploy_config.pipeline in _PIPELINE_REGISTRY: + return cls._create_from_registry( + deploy_config.pipeline, + cli_overrides, + deploy_config_path, + ) + # Not in the pipeline registry — let the caller fall back to the # legacy ``stage_configs/*.yaml`` path (resolve_model_config_path). return None diff --git a/vllm_omni/deploy/ming_tts.yaml b/vllm_omni/deploy/ming_tts.yaml new file mode 100644 index 000000000..1a2dc99ec --- /dev/null +++ b/vllm_omni/deploy/ming_tts.yaml @@ -0,0 +1,55 @@ +# Ming-omni-tts deploy: stage-0 LLM+flow -> stage-1 audio VAE. +# Verified legacy settings migrated from ming_tts.yaml and ming_tts_async_chunk.yaml. +# Default mode is async-chunk streaming on a single GPU. +# pipeline: required — HF config reports model_type="dense" with no architectures, +# so auto-detection falls through. This key routes the deploy-config probe in +# StageConfigFactory.create_from_model to the correct registry entry. +pipeline: ming_tts +async_chunk: true +trust_remote_code: false +dtype: bfloat16 +enable_prefix_caching: false + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + latent_chunk_size: 25 + latent_left_context: 0 + +stages: + - stage_id: 0 + max_num_seqs: 1 + gpu_memory_utilization: 0.45 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 8192 + max_model_len: 8192 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 512 + detokenize: true + + - stage_id: 1 + max_num_seqs: 1 + gpu_memory_utilization: 0.25 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 8192 + max_model_len: 8192 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 1 + detokenize: false diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 25244c462..ac52de953 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -35,6 +35,7 @@ def _register_omni_hf_configs() -> None: try: from transformers import AutoConfig + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingDenseConfig from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( Qwen3TTSConfig, ) @@ -55,6 +56,7 @@ def _register_omni_hf_configs() -> None: _CONFIG_REGISTRY = None for model_type, config_cls in [ + ("dense", MingDenseConfig), ("qwen3_tts", Qwen3TTSConfig), ("cosyvoice3", CosyVoice3Config), ("glm_tts", GLMTTSConfig), diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index a18687ca2..1dce789ce 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -7,6 +7,43 @@ from pydantic import AliasChoices, BaseModel, Field, field_validator, model_vali _MAX_EMBEDDING_DIM = 8192 +def _normalize_ref_audio_value(value): + if value is None: + return None + if isinstance(value, str): + return value + if isinstance(value, (list, tuple)): + items = [] + for item in value: + if not isinstance(item, str): + raise TypeError("'ref_audio' list entries must be strings") + items.append(item) + if not items: + raise ValueError("'ref_audio' list cannot be empty") + return items + raise TypeError("'ref_audio' must be a string or list of strings") + + +def _normalize_speaker_embedding_value(value): + if value is None: + return None + if not isinstance(value, (list, tuple)): + raise TypeError("'speaker_embedding' must be a list of numbers or list of embedding vectors") + if not value: + return [] + + first = value[0] + if isinstance(first, (list, tuple)): + embeddings = [] + for item in value: + if not isinstance(item, (list, tuple)): + raise TypeError("'speaker_embedding' must not mix flat and nested values") + embeddings.append([float(x) for x in item]) + return embeddings + + return [float(x) for x in value] + + class OpenAICreateSpeechRequest(BaseModel): input: str model: str | None = None @@ -46,7 +83,7 @@ class OpenAICreateSpeechRequest(BaseModel): default=None, description="Language code (e.g., 'Chinese', 'English', 'Auto')", ) - ref_audio: str | None = Field( + ref_audio: str | list[str] | None = Field( default=None, description="Reference audio for voice cloning (Base task). URL, base64, or file URI.", ) @@ -73,7 +110,7 @@ class OpenAICreateSpeechRequest(BaseModel): default=None, description="Use speaker embedding only without in-context learning (Base task)", ) - speaker_embedding: list[float] | None = Field( + speaker_embedding: list[float] | list[list[float]] | None = Field( default=None, max_length=_MAX_EMBEDDING_DIM, description="Pre-computed speaker embedding vector (1024-dim for 0.6B, " @@ -108,10 +145,29 @@ class OpenAICreateSpeechRequest(BaseModel): raise ValueError("'sse' is not a supported stream_format yet. Please use 'audio'.") return v + @field_validator("ref_audio", mode="before") + @classmethod + def normalize_ref_audio(cls, v): + return _normalize_ref_audio_value(v) + @field_validator("speaker_embedding") @classmethod - def validate_speaker_embedding(cls, v: list[float] | None) -> list[float] | None: - if v is not None and not all(math.isfinite(x) for x in v): + def validate_speaker_embedding( + cls, v: list[float] | list[list[float]] | None + ) -> list[float] | list[list[float]] | None: + v = _normalize_speaker_embedding_value(v) + if v is None: + return None + if not v: + return [] + if isinstance(v[0], list): + for item in v: + if not item: + raise ValueError("'speaker_embedding' nested vectors must be non-empty") + if not all(math.isfinite(x) for x in item): + raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)") + return v + if not all(math.isfinite(x) for x in v): raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)") return v @@ -210,7 +266,7 @@ class OpenAICreateSpeechRequest(BaseModel): @model_validator(mode="after") def validate_embedding_constraints(self) -> "OpenAICreateSpeechRequest": if self.speaker_embedding is not None: - if self.ref_audio is not None: + if self.ref_audio is not None and not isinstance(self.ref_audio, list): raise ValueError("'speaker_embedding' and 'ref_audio' are mutually exclusive") return self diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index d24181a0f..9127ac7db 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -66,6 +66,7 @@ from vllm_omni.utils.speaker_cache import ( logger = init_logger(__name__) # TTS Configuration +_MING_TTS_MODEL_ARCHS = {"MingTTSForConditionalGeneration"} _VOXTRAL_TTS_MODEL_STAGES = {"audio_generation"} _QWEN3_TTS_MODEL_STAGES = {"qwen3_tts"} _FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"} @@ -126,6 +127,7 @@ _QWEN3_TTS_REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key" _TTS_MAX_INSTRUCTIONS_LENGTH = 500 _TTS_MAX_NEW_TOKENS_MIN = 1 _TTS_MAX_NEW_TOKENS_MAX = 4096 +_MING_DEFAULT_PROMPT = MING_DEFAULT_PROMPT def _create_wav_header(sample_rate: int, num_channels: int = 1, bits_per_sample: int = 16) -> bytes: @@ -528,6 +530,32 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _load_codec_frame_rate(self) -> float | None: """Load codec frame rate from speech tokenizer config for prompt length estimation.""" + if self._tts_model_type == "ming_tts": + try: + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingTTSConfig + + hf_config = self.engine_client.model_config.hf_config + ming_cfg = MingTTSConfig.from_hf_config(hf_config) + patch_size = int(ming_cfg.patch_size) + audio_frame_hop = int(ming_cfg.audio_frame_hop) + sample_rate = int(ming_cfg.sample_rate) + if patch_size <= 0 or audio_frame_hop <= 0 or sample_rate <= 0: + raise ValueError( + "Ming config has invalid tokenizer timing values: " + f"patch_size={patch_size}, audio_frame_hop={audio_frame_hop}, sample_rate={sample_rate}" + ) + rate = float(sample_rate) / float(audio_frame_hop * patch_size) + logger.info( + "Derived Ming codec frame rate: %.1f Hz (sample_rate=%s, audio_frame_hop=%s, patch_size=%s)", + rate, + sample_rate, + audio_frame_hop, + patch_size, + ) + return rate + except Exception as e: + logger.warning(f"Failed to derive Ming codec frame rate from hf_config: {e}") + try: model_path = self.engine_client.model_config.model st_config_path = os.path.join(model_path, "speech_tokenizer", "config.json") @@ -572,7 +600,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _find_tts_stage(self): """Find and return the TTS stage config, or None if not found.""" for stage in self.engine_client.stage_configs: - if stage.engine_args.model_stage in _TTS_MODEL_STAGES: + engine_args = stage.engine_args + model_stage = engine_args.model_stage + model_arch = getattr(engine_args, "model_arch", None) + worker_type = getattr(engine_args, "worker_type", None) + if model_stage in _TTS_MODEL_STAGES: + return stage + # Ming dense identifies its AR entry stage by architecture because + # it does not use a dedicated TTS model_stage value. + if model_arch in _MING_TTS_MODEL_ARCHS and worker_type == "ar": return stage return None @@ -602,6 +638,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return "voxcpm2" if model_stage in _MING_TTS_MODEL_STAGES: return "ming_flash_omni_tts" + if model_arch in _MING_TTS_MODEL_ARCHS: + return "ming_tts" if model_stage in _MOSS_TTS_MODEL_STAGES: return "moss_tts_nano" if model_stage in _MOSS_TTS_FULL_MODEL_STAGES: @@ -682,6 +720,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): try: if self._tts_model_type == "glm_tts": return set() + if self._tts_model_type == "ming_tts": + return set() if self._tts_model_type == "voxcpm2": return {"default"} if self._tts_model_type == "voxtral_tts": @@ -1214,9 +1254,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)") emb_dim = len(embedding) - dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim) - if dim_err is not None: - raise ValueError(dim_err) + if self._tts_model_type == "ming_tts": + if emb_dim != 192: + logger.warning("speaker_embedding has %d dimensions; Ming dense expects 192", emb_dim) + else: + dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim) + if dim_err is not None: + raise ValueError(dim_err) async with self._upload_lock: voice_name_lower = name.lower() @@ -1303,7 +1347,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _is_tts_model(self) -> bool: """Check if the current model is a supported TTS model.""" - return any(stage.engine_args.model_stage in _TTS_MODEL_STAGES for stage in self.engine_client.stage_configs) + return self._find_tts_stage() is not None def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate TTS request parameters. Returns error message or None.""" @@ -1316,6 +1360,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if self._tts_model_type == "voxcpm2": return self._validate_voxcpm2_request(request) if self._tts_model_type == "ming_flash_omni_tts": + return self._validate_ming_flash_omni_tts_request(request) + if self._tts_model_type == "ming_tts": return self._validate_ming_tts_request(request) if self._tts_model_type in ("moss_tts_nano", "moss_tts"): return self._validate_moss_tts_request(request) @@ -1362,7 +1408,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ids = self._voxcpm2_tokenizer.encode(text, add_special_tokens=True) return split_multichar_chinese(ids, self._voxcpm2_split_map) - def _validate_ming_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: + def _validate_ming_flash_omni_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate Ming-flash-omni standalone-talker request parameters.""" if not request.input or not request.input.strip(): return "Input text cannot be empty" @@ -1399,6 +1445,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _validate_ref_audio_format(self, ref_audio: str) -> str | None: """Validate ref_audio is a supported URI format. Returns error or None.""" + if not isinstance(ref_audio, str): + return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)" if not ( ref_audio.startswith(("http://", "https://")) or ref_audio.startswith("data:") @@ -2066,6 +2114,99 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return None + def _validate_ming_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: + """Validate Ming TTS request parameters. Returns error message or None.""" + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + if isinstance(request.ref_audio, list): + return self._validate_ming_tts_podcast_request(request) + return self._validate_ming_tts_single_speaker_request(request) + + def _validate_ming_tts_single_speaker_request(self, request: OpenAICreateSpeechRequest) -> str | None: + if request.ref_audio is not None: + fmt_err = self._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + + if request.speaker_embedding is not None: + if not request.speaker_embedding: + return "'speaker_embedding' must be a non-empty list of floats" + emb_len = len(request.speaker_embedding) + if emb_len != 192: + logger.warning( + "speaker_embedding has %d dimensions; Ming dense expects 192. " + "Wrong dimensions will likely fail or degrade output.", + emb_len, + ) + + voice_lower = request.voice.lower() if isinstance(request.voice, str) else None + uploaded_voice = bool(voice_lower and voice_lower in self.uploaded_speakers) + clone_source_present = request.ref_audio is not None or request.speaker_embedding is not None or uploaded_voice + + if request.task_type == "Base" and not clone_source_present: + return "Base task requires 'ref_audio', 'speaker_embedding', or an uploaded voice sample" + + if request.ref_audio is not None and request.ref_text is not None and not request.ref_text.strip(): + return "'ref_text' must be non-empty when provided with 'ref_audio'" + + # Ming offline ref-audio cases use prompt_waveform without prompt_text; + # keep the transcript requirement for other TTS models. + if request.ref_audio is not None and request.speaker_embedding is None and not self._is_ming_tts_model(): + uploaded_ref_text = self.uploaded_speakers[voice_lower].get("ref_text") if uploaded_voice else None + if not (request.ref_text and request.ref_text.strip()) and not uploaded_ref_text: + return "Reference-audio cloning requires non-empty 'ref_text'" + + if request.ref_text is not None and request.ref_audio is None and not uploaded_voice: + return "'ref_text' requires 'ref_audio' or an uploaded voice sample" + + if request.instructions and len(request.instructions) > self._max_instructions_length: + return f"Instructions too long (max {self._max_instructions_length} characters)" + + if request.max_new_tokens is not None: + if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN: + return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}" + if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX: + return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}" + + return None + + def _validate_ming_tts_podcast_request(self, request: OpenAICreateSpeechRequest) -> str | None: + if len(request.ref_audio) < 2: + return "Podcast-style Ming requests require at least two 'ref_audio' clips" + + for ref_audio in request.ref_audio: + fmt_err = self._validate_ref_audio_format(ref_audio) + if fmt_err: + return fmt_err + + if not request.ref_text or not request.ref_text.strip(): + return "Podcast-style Ming requests require non-empty 'ref_text'" + + if request.speaker_embedding is not None: + embeddings = request.speaker_embedding + embedding_count = len(embeddings) if embeddings and isinstance(embeddings[0], list) else 1 + if embedding_count != len(request.ref_audio): + return ( + "Podcast-style Ming requests require one speaker embedding per ref_audio clip; " + f"got {embedding_count} embeddings for {len(request.ref_audio)} clips" + ) + if embeddings and isinstance(embeddings[0], list): + for item in embeddings: + if len(item) != 192: + return "Podcast-style Ming speaker embeddings must each have 192 dimensions" + + if request.instructions and len(request.instructions) > self._max_instructions_length: + return f"Instructions too long (max {self._max_instructions_length} characters)" + + if request.max_new_tokens is not None: + if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN: + return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}" + if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX: + return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}" + + return None + async def _resolve_ref_audio(self, ref_audio_str: str) -> tuple[list[float], int]: """Resolve ref_audio to (wav_samples, sample_rate). @@ -2192,6 +2333,142 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _discard_ref_audio_artifact_warmup(self, request_id: str) -> None: self._request_ref_audio_artifact_keys.pop(request_id, None) + async def _resolve_ref_audio_many(self, ref_audio_list: list[str]) -> list[tuple[list[float], int]]: + resolved = [] + for ref_audio in ref_audio_list: + resolved.append(await self._resolve_ref_audio(ref_audio)) + return resolved + + # ---- Ming TTS helpers ---- + + def _is_ming_tts_model(self) -> bool: + return self._tts_model_type == "ming_tts" + + def _coerce_ming_prompt_waveform(self, wav_samples, sample_rate): + from torchaudio.functional import resample as resample_audio + + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import SAMPLE_RATE + + waveform = torch.as_tensor(wav_samples, dtype=torch.float32).reshape(1, -1) + if int(sample_rate) != SAMPLE_RATE: + waveform = resample_audio(waveform, int(sample_rate), SAMPLE_RATE) + return waveform + + def _build_ming_prompt_waveform( + self, + ref_audio_data: tuple[list[float], int] | list[tuple[list[float], int]] | None, + ): + if isinstance(ref_audio_data, list): + return torch.cat( + [self._coerce_ming_prompt_waveform(item[0], item[1]) for item in ref_audio_data], + dim=-1, + ) + if ref_audio_data is not None: + return self._coerce_ming_prompt_waveform(ref_audio_data[0], ref_audio_data[1]) + return None + + def _extract_ming_speaker_embeddings_from_ref_audio( + self, + ref_audio_data_list: list[tuple[list[float], int]], + ) -> list[list[float]]: + from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor + + extractor = MingSpeakerEmbeddingExtractor(self.engine_client.model_config.model, target_sr=16000) + embeddings = [] + for wav_samples, sr in ref_audio_data_list: + waveform = torch.as_tensor(wav_samples, dtype=torch.float32).reshape(1, -1) + embedding = extractor.extract_from_waveform(waveform, int(sr)) + flat = embedding.detach().reshape(-1).to(torch.float32).cpu() + if int(flat.numel()) != 192: + raise ValueError(f"Ming speaker extractor returned {int(flat.numel())} dims; expected 192") + embeddings.append(flat.tolist()) + return embeddings + + def _parse_ming_instruction_fields( + self, + request, + *, + include_language=False, + include_voice=False, + plain_text_passthrough=False, + ): + instruction_text = request.instructions.strip() if isinstance(request.instructions, str) else None + instruction_dict: dict[str, Any] = {} + + if include_language and request.language not in (None, "", "Auto"): + instruction_dict["方言"] = request.language + + voice_lower = request.voice.lower() if isinstance(request.voice, str) else None + if include_voice and request.voice and not (voice_lower and voice_lower in self.uploaded_speakers): + instruction_dict["IP"] = request.voice + + if instruction_text: + try: + parsed = json.loads(instruction_text) + except json.JSONDecodeError: + parsed = None + if isinstance(parsed, dict): + instruction_dict.update(parsed) + elif instruction_dict or not plain_text_passthrough: + instruction_dict["风格"] = instruction_text + else: + return instruction_text + + return instruction_dict or None + + def _parse_ming_instruction(self, request: OpenAICreateSpeechRequest) -> Any: + """Build a Ming instruction payload from OpenAI speech fields.""" + return self._parse_ming_instruction_fields( + request, + include_language=True, + include_voice=True, + plain_text_passthrough=True, + ) + + def _build_ming_dense_prompt( + self, + request: OpenAICreateSpeechRequest, + *, + ref_audio_data: tuple[list[float], int] | list[tuple[list[float], int]] | None = None, + ) -> dict[str, Any]: + """Build a Ming dense prompt directly from the OpenAI speech request.""" + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import KEY_MAX_DECODE_STEPS + from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt + + if self._tts_tokenizer is None: + model_name = self.engine_client.model_config.model + self._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=False) + + ref_text = request.ref_text + prompt_waveform = self._build_ming_prompt_waveform(ref_audio_data) if ref_text is not None else None + speaker_embedding = request.speaker_embedding + use_zero_spk_emb = prompt_waveform is None and speaker_embedding is None + + runtime_controls = {} + if request.max_new_tokens is not None: + runtime_controls[KEY_MAX_DECODE_STEPS] = request.max_new_tokens + + prompt_dict = build_ming_dense_prompt( + self._tts_tokenizer, + # bgm / music-prompt mode not supported online; + # requires prompt_mode API extension (deferred). + prompt=_MING_DEFAULT_PROMPT, + text=request.input, + runtime_controls=runtime_controls or None, + instruction=self._parse_ming_instruction(request), + prompt_text=ref_text, + prompt_waveform=prompt_waveform, + speaker_embedding=speaker_embedding, + use_zero_spk_emb=use_zero_spk_emb, + ) + prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"]) + prompt["prompt"] = prompt_dict["prompt"] + prompt["text"] = prompt_dict["text"] + prompt["additional_information"] = prompt_dict["additional_information"] + return prompt + async def _generate_audio_chunks( self, generator, @@ -2796,25 +3073,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # ---- Ming-flash-omni standalone-talker (TTS) helpers ---- - def _build_ming_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: + def _build_ming_flash_omni_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: # request.instructions accepts two forms: # 1. Plain text: mapped to the caption's 风格 (style) field # 2. JSON object: parsed and splatted into the caption. Unlocks # Unknown keys are dropped by `ming_create_instruction`. - caption_fields: dict[str, Any] = {} - if request.instructions: - stripped = request.instructions.strip() - if stripped.startswith("{"): - try: - parsed = json.loads(stripped) - except json.JSONDecodeError: - parsed = None - if isinstance(parsed, dict): - caption_fields.update(parsed) - else: - caption_fields["风格"] = request.instructions - else: - caption_fields["风格"] = request.instructions + caption_fields = self._parse_ming_instruction_fields(request) or {} has_spk_emb = request.speaker_embedding is not None @@ -2982,8 +3246,29 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): prompt = await self._build_glm_tts_prompt(request, has_inline_ref_audio=has_inline_ref_audio) tts_params = {} elif self._tts_model_type == "ming_flash_omni_tts": - prompt = self._build_ming_prompt(request) + prompt = self._build_ming_flash_omni_prompt(request) tts_params = {} + elif self._tts_model_type == "ming_tts": + ref_audio_source = request.ref_audio + voice_lower = request.voice.lower() if isinstance(request.voice, str) else None + if ref_audio_source is None and voice_lower in self.uploaded_speakers: + ref_audio_source = self._get_uploaded_audio_data(request.voice) + if request.ref_text is None: + request.ref_text = self.uploaded_speakers[voice_lower].get("ref_text") + ref_audio_data = None + if isinstance(ref_audio_source, list): + ref_audio_data = await self._resolve_ref_audio_many(ref_audio_source) + if request.speaker_embedding is None: + request.speaker_embedding = self._extract_ming_speaker_embeddings_from_ref_audio(ref_audio_data) + elif ref_audio_source is not None and isinstance(ref_audio_source, str): + wav_list, sr = await self._resolve_ref_audio(ref_audio_source) + ref_audio_data = (wav_list, sr) + if request.speaker_embedding is None: + request.speaker_embedding = self._extract_ming_speaker_embeddings_from_ref_audio( + [ref_audio_data] + )[0] + prompt = self._build_ming_dense_prompt(request, ref_audio_data=ref_audio_data) + tts_params = prompt.get("additional_information", {}) elif self._tts_model_type in ("moss_tts_nano", "moss_tts"): tts_params = await self._build_moss_tts_params(request) if request.voice: @@ -3069,6 +3354,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_type = "voxcpm2" elif self._tts_model_type == "ming_flash_omni_tts": model_type = "ming_flash_omni_tts" + elif self._tts_model_type == "ming_tts": + model_type = "ming_tts" elif self._tts_model_type == "moss_tts_nano": model_type = "moss_tts_nano" elif self._tts_model_type == "moss_tts": @@ -3177,6 +3464,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): getattr(sampling_params_list[0], "min_tokens", 0), request.max_new_tokens, ) + elif self._tts_model_type == "ming_tts" and sampling_params_list: + import copy + + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import TEXT_EOS_TOKEN_ID + + sampling_params_list = copy.deepcopy(sampling_params_list) + sampling_params_list[0].stop_token_ids = [int(TEXT_EOS_TOKEN_ID)] + if request.max_new_tokens is not None: + # Ming emits TEXT_EOS after the latent decode budget is exhausted, so + # Stage-0 needs one extra token beyond ming_max_decode_steps. + sampling_params_list[0].max_tokens = int(request.max_new_tokens) + 1 if request.seed is not None and sampling_params_list: if sampling_params_list is self.engine_client.default_sampling_params_list: diff --git a/vllm_omni/model_executor/models/common/ming/__init__.py b/vllm_omni/model_executor/models/common/ming/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/model_executor/models/common/ming/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/common/ming/audio_dsp.py b/vllm_omni/model_executor/models/common/ming/audio_dsp.py new file mode 100644 index 000000000..be8cd489c --- /dev/null +++ b/vllm_omni/model_executor/models/common/ming/audio_dsp.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/audio_tokenizer/istft.py + +import torch +import torch.nn as nn + + +class ISTFT(nn.Module): + """ + Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with + windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges. + See issue: https://github.com/pytorch/pytorch/issues/62323 + Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs. + The NOLA constraint is met as we trim padded samples anyway. + + Args: + n_fft (int): Size of Fourier transform. + hop_length (int): The distance between neighboring sliding window frames. + win_length (int): The size of window frame and STFT filter. + padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". + """ + + def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"): + super().__init__() + if padding not in ["center", "same"]: + raise ValueError("Padding must be 'center' or 'same'.") + self.padding = padding + self.n_fft = n_fft + self.hop_length = hop_length + self.win_length = win_length + window = torch.hann_window(win_length) + self.register_buffer("window", window) + + self.buffer_len = self.win_length - self.hop_length + + def _buffer_process( + self, + x: torch.Tensor, + buffer: torch.Tensor | None, + pad: int, + last_chunk: bool = False, + streaming: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if streaming: + if buffer is None: + # first chunk + x = x[:, pad:] + if buffer is not None: + # next chunk + x[:, : self.buffer_len] += buffer + buffer = x[:, -self.buffer_len :] + if not last_chunk: + x = x[:, : -self.buffer_len] + else: + x = x[:, :-pad] + else: + x = x[:, pad:-pad] + + return x, buffer + + def forward( + self, + spec: torch.Tensor, + audio_buffer: torch.Tensor | None = None, + window_buffer: torch.Tensor | None = None, + streaming: bool = False, + last_chunk: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """ + Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram. + + Args: + spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size, + N is the number of frequency bins, and T is the number of time frames. + audio_buffer (Tensor): [Streaming Input/State] The audio overlap buffer from the previous chunk. + Shape: (B, win_length - hop_length) + window_buffer (Tensor): [Streaming Input/State] The window overlap buffer from the previous chunk. + streaming: If `True`, the function operates in streaming mode, processing `spec` as a single chunk. + last_chunk: When `streaming=True` and `last_chunk=True`, the function can perform final "flush" operations + + Returns: + Reconstructed signal, plus streaming buffers when `padding="same"`. + """ + if self.padding == "center": + # Fallback to pytorch native implementation + return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True) + elif self.padding == "same": + pad = (self.win_length - self.hop_length) // 2 + else: + raise ValueError("Padding must be 'center' or 'same'.") + + if spec.dim() != 3: + raise ValueError(f"Expected spec rank-3 [Batch, Freq, Time], got {tuple(spec.shape)}") + _, _, T = spec.shape + + # Inverse FFT + ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") + ifft = ifft * self.window[None, :, None] + + # Overlap and Add + output_size = (T - 1) * self.hop_length + self.win_length + y = torch.nn.functional.fold( + ifft, + output_size=(1, output_size), + kernel_size=(1, self.win_length), + stride=(1, self.hop_length), + )[:, 0, 0, :] + + y, audio_buffer = self._buffer_process(y, audio_buffer, pad, last_chunk=last_chunk, streaming=streaming) + + # Window envelope + window_sq = self.window.square().expand(1, T, -1).transpose(1, 2) + window_envelope = ( + torch.nn.functional.fold( + window_sq, + output_size=(1, output_size), + kernel_size=(1, self.win_length), + stride=(1, self.hop_length), + ) + .squeeze(0) + .squeeze(0) + ) + + window_envelope, window_buffer = self._buffer_process( + window_envelope, window_buffer, pad, last_chunk=last_chunk, streaming=streaming + ) + window_envelope = window_envelope.squeeze() + + # Normalize + if not (window_envelope > 1e-11).all(): + raise RuntimeError("ISTFT window envelope underflowed; invalid overlap-add state.") + y = y / window_envelope + + return y, audio_buffer, window_buffer + + +class FourierHead(nn.Module): + """Base class for inverse fourier modules.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, + L is the sequence length, and H denotes the model dimension. + + Returns: + Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. + """ + raise NotImplementedError("Subclasses must implement the forward method.") + + +class ISTFTHead(FourierHead): + """ + ISTFT Head module for predicting STFT complex coefficients. + + Args: + dim (int): Hidden dimension of the model. + n_fft (int): Size of Fourier transform. + hop_length (int): The distance between neighboring sliding window frames, which should align with + the resolution of the input features. + padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". + """ + + def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"): + super().__init__() + out_dim = n_fft + 2 + self.out = torch.nn.Linear(dim, out_dim) + self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding) + + def forward( + self, + x: torch.Tensor, + audio_buffer: torch.Tensor | None = None, + window_buffer: torch.Tensor | None = None, + streaming: bool = False, + last_chunk: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """ + Forward pass of the ISTFTHead module. + + Args: + x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, + L is the sequence length, and H denotes the model dimension. + + Returns: + Audio, predicted spectrogram coefficients, and streaming buffers. + """ + x_pred = self.out(x) + x_pred = x_pred.transpose(1, 2) + mag, p = x_pred.chunk(2, dim=1) + mag = torch.exp(mag) + mag = torch.clip(mag, max=1e2) # safeguard to prevent excessively large magnitudes + # wrapping happens here. These two lines produce real and imaginary value + x = torch.cos(p) + y = torch.sin(p) + # recalculating phase here does not produce anything new + # only costs time + # phase = torch.atan2(y, x) + # S = mag * torch.exp(phase * 1j) + # better directly produce the complex value + S = mag * (x + 1j * y) + audio, audio_buffer, window_buffer = self.istft( + S, audio_buffer=audio_buffer, window_buffer=window_buffer, streaming=streaming, last_chunk=last_chunk + ) + return audio.unsqueeze(1), x_pred, audio_buffer, window_buffer + + +__all__ = ["FourierHead", "ISTFT", "ISTFTHead"] diff --git a/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py b/vllm_omni/model_executor/models/common/ming/audio_vae.py similarity index 64% rename from vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py rename to vllm_omni/model_executor/models/common/ming/audio_vae.py index 3137a878e..6924ae548 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py +++ b/vllm_omni/model_executor/models/common/ming/audio_vae.py @@ -1,10 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# Copyright 2025 The vLLM-Omni team. -# Copyright (c) Ant Group. All rights reserved. -# Adapted from: -# https://github.com/inclusionAI/Ming/tree/e58533db227031990c5a6864dcf5f08fb53ed0d2/AudioVAE - -from __future__ import annotations +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch import torch.nn as nn @@ -13,6 +8,8 @@ from transformers import PretrainedConfig, PreTrainedModel, Qwen2Config, Qwen2Mo from transformers.utils import is_flash_attn_2_available from vllm.logger import init_logger +from .audio_dsp import ISTFTHead + logger = init_logger(__name__) @@ -21,114 +18,50 @@ class AudioVAEConfig(PretrainedConfig): def __init__( self, - sample_rate: int = 44100, - enc_kwargs: dict | None = None, - dec_kwargs: dict | None = None, - init_method: str = "kaiming", - patch_size: int = 4, + sample_rate=44100, + enc_kwargs=None, + semantic_module_kwargs=None, + dec_kwargs=None, + hifi_gan_disc_kwargs=None, + spec_disc_kwargs=None, + lambda_disc=1.0, + lambda_mel_loss=15, + lambda_adv=1.0, + lambda_feat_match_loss=1.0, + lambda_semantic=5.0, + init_method="kaiming", + patch_size=4, **kwargs, ): self.sample_rate = sample_rate self.enc_kwargs = enc_kwargs or {} + self.semantic_module_kwargs = semantic_module_kwargs self.dec_kwargs = dec_kwargs or {} + self.hifi_gan_disc_kwargs = hifi_gan_disc_kwargs + self.spec_disc_kwargs = spec_disc_kwargs + self.lambda_disc = lambda_disc + self.lambda_mel_loss = lambda_mel_loss + self.lambda_adv = lambda_adv + self.lambda_feat_match_loss = lambda_feat_match_loss + self.lambda_semantic = lambda_semantic self.init_method = init_method self.patch_size = patch_size super().__init__(**kwargs) -class ISTFT(nn.Module): - def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"): - super().__init__() - if padding not in ["center", "same"]: - raise ValueError("Padding must be 'center' or 'same'.") - self.padding = padding - self.n_fft = n_fft - self.hop_length = hop_length - self.win_length = win_length - window = torch.hann_window(win_length) - self.register_buffer("window", window) - self.buffer_len = self.win_length - self.hop_length - - def _buffer_process(self, x, buffer, pad, last_chunk=False, streaming=False): - if streaming: - if buffer is None: - x = x[:, pad:] - if buffer is not None: - x[:, : self.buffer_len] += buffer - buffer = x[:, -self.buffer_len :] - if not last_chunk: - x = x[:, : -self.buffer_len] - else: - x = x[:, :-pad] - else: - x = x[:, pad:-pad] - return x, buffer - - def forward(self, spec, audio_buffer=None, window_buffer=None, streaming=False, last_chunk=False): - if self.padding == "center": - return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True) - elif self.padding == "same": - pad = (self.win_length - self.hop_length) // 2 - else: - raise ValueError("Padding must be 'center' or 'same'.") - - B, N, T = spec.shape - ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") - ifft = ifft * self.window[None, :, None] - - output_size = (T - 1) * self.hop_length + self.win_length - y = torch.nn.functional.fold( - ifft, - output_size=(1, output_size), - kernel_size=(1, self.win_length), - stride=(1, self.hop_length), - )[:, 0, 0, :] - - y, audio_buffer = self._buffer_process(y, audio_buffer, pad, last_chunk=last_chunk, streaming=streaming) - - window_sq = self.window.square().expand(1, T, -1).transpose(1, 2) - window_envelope = ( - torch.nn.functional.fold( - window_sq, - output_size=(1, output_size), - kernel_size=(1, self.win_length), - stride=(1, self.hop_length), - ) - .squeeze(0) - .squeeze(0) - ) - - window_envelope, window_buffer = self._buffer_process( - window_envelope, window_buffer, pad, last_chunk=last_chunk, streaming=streaming - ) - window_envelope = window_envelope.squeeze() - - assert (window_envelope > 1e-11).all() - y = y / window_envelope - - return y, audio_buffer, window_buffer - - -class ISTFTHead(nn.Module): - def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"): - super().__init__() - out_dim = n_fft + 2 - self.out = nn.Linear(dim, out_dim) - self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding) - - def forward(self, x, audio_buffer=None, window_buffer=None, streaming=False, last_chunk=False): - x_pred = self.out(x) - x_pred = x_pred.transpose(1, 2) - mag, p = x_pred.chunk(2, dim=1) - mag = torch.exp(mag) - mag = torch.clip(mag, max=1e2) - x = torch.cos(p) - y = torch.sin(p) - S = mag * (x + 1j * y) - audio, audio_buffer, window_buffer = self.istft( - S, audio_buffer=audio_buffer, window_buffer=window_buffer, streaming=streaming, last_chunk=last_chunk - ) - return audio.unsqueeze(1), x_pred, audio_buffer, window_buffer +def _qwen2_config(backbone): + config = Qwen2Config.from_dict(config_dict=backbone) + if (getattr(config, "_attn_implementation", None) or getattr(config, "attn_implementation", None)) not in ( + None, + "flash_attention_2", + ): + return config + if is_flash_attn_2_available(): + config._attn_implementation_autoset = True + config._attn_implementation = "flash_attention_2" + else: + config._attn_implementation = "sdpa" + return config class StreamingLinearUpsample(nn.Module): @@ -138,6 +71,8 @@ class StreamingLinearUpsample(nn.Module): self.upsampler = nn.Upsample(scale_factor=scale_factor, mode="linear", align_corners=False) def forward(self, x, state=None, is_last=False): + if x is None and is_last and (state is None or state.get("prev_chunk") is None): + raise ValueError("Received end-of-stream without any latent chunk to upsample.") if state is None: state = {"prev_chunk": None, "history_last": None, "is_first": True} @@ -188,16 +123,72 @@ class StreamingLinearUpsample(nn.Module): return final_out, state +class Encoder(nn.Module): + def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1): + super().__init__() + config = _qwen2_config(encoder_args) + logger.info("AudioVAE Encoder: using attn_implementation=%r", config._attn_implementation) + self.encoder = Qwen2Model(config) + self.input_dim = input_dim + self.hop_size = hop_size + self.latent_dim = latent_dim + self.fc1 = nn.Linear(input_dim, config.hidden_size, bias=False) + self.fc2 = nn.Linear(config.hidden_size, config.hidden_size) + self.fc3 = nn.Linear(config.hidden_size, latent_dim * 2) + self.norm = nn.LayerNorm(config.hidden_size) + self.patch_size = patch_size + if patch_size != -1: + aggregator_config = _qwen2_config({**encoder_args, "num_hidden_layers": 4}) + self.aggregator = Qwen2Model(aggregator_config) + self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size)) + self.cls_embed.data.normal_(0, 0.02) + + def get_frames(self, x): + num_frames_total = (x.size(-1) + self.hop_size - 1) // self.hop_size + expected_len = (num_frames_total - 1) * self.hop_size + self.input_dim + padding_needed = expected_len - x.size(-1) + waveform = F.pad(x, (0, padding_needed), value=0.0) + frames = waveform.unfold(dimension=-1, size=self.input_dim, step=self.hop_size) + return frames + + def pad_patch_insert_cls(self, x): + bsz, _, dim = x.size() + num_frame = x.size(1) + r = num_frame % self.patch_size + pad_num = self.patch_size - r if r else 0 + x = F.pad(x, (0, 0, 0, pad_num), value=0.0) + # [Batch, Time, Dimension] -> [Batch*PatchGroups, Patch, Dimension]. + x = x.reshape(-1, self.patch_size, dim) + x = torch.cat((x, self.cls_embed.expand(x.size(0), -1, -1)), dim=1) + # [Batch*PatchGroups, Patch+1, Dimension] -> [Batch, Time, Dimension]. + x = x.reshape(bsz, -1, dim) + return x + + def forward(self, waveform): + x = self.get_frames(waveform) + x = self.fc1(x) + x = self.fc2(x) + x = self.encoder(inputs_embeds=x) + x = x.last_hidden_state + + if self.patch_size != -1: + x = self.pad_patch_insert_cls(x) + x = self.aggregator(inputs_embeds=x) + x = x.last_hidden_state + bsz, _, dim = x.size() + # [Batch, Time, Dimension] -> [Batch*PatchGroups, Patch+1, Dimension]. + x = x.reshape(-1, self.patch_size + 1, dim) + # [Batch*PatchGroups, 1, Dimension] -> [Batch, PatchGroups, Dimension]. + x = x[:, -1:, :].reshape(bsz, -1, dim) + + x = self.fc3(x) + return x, waveform.unsqueeze(1) + + class Decoder(nn.Module): def __init__(self, decoder_args, output_dim=320, latent_dim=64, patch_size=-1): super().__init__() - config = Qwen2Config.from_dict(config_dict=decoder_args) - if is_flash_attn_2_available(): - config._attn_implementation_autoset = True - config._attn_implementation = "flash_attention_2" - else: - config._attn_implementation = "sdpa" - + config = _qwen2_config(decoder_args) logger.info("AudioVAE Decoder: using attn_implementation=%r", config._attn_implementation) self.decoder = Qwen2Model(config) self.output_dim = output_dim @@ -211,7 +202,21 @@ class Decoder(nn.Module): if self.patch_size != -1: self.upsampling = StreamingLinearUpsample(scale_factor=patch_size) + def forward(self, x): + x = self.fc1(x) + + if self.patch_size != -1: + # [Batch, Time, Dimension] -> [Batch, Dimension, Time] -> [Batch, Time, Dimension]. + x = self.upsampling.upsampler(x.transpose(1, 2)).transpose(1, 2) + + x = self.decoder(inputs_embeds=x) + x = x.last_hidden_state + x, _ = self.head(x) + return x, None + def low_level_reconstruct(self, x, past_key_values=None, use_cache=False, stream_state=None, last_chunk=False): + if stream_state is None: + stream_state = (None, None, None) upsample_state, audio_buffer, window_buffer = stream_state bsz, device, dtype = x.size(0), x.device, x.dtype x = self.fc1(x) @@ -222,6 +227,7 @@ class Decoder(nn.Module): stream_state = (upsample_state, audio_buffer, window_buffer) return torch.empty(bsz, 1, 0, device=device, dtype=dtype), stream_state, past_key_values else: + # [Batch, Time, Dimension] -> [Batch, Dimension, Time] -> [Batch, Time, Dimension]. x = self.upsampling.upsampler(x.transpose(1, 2)).transpose(1, 2) hidden_states_list = [] @@ -239,7 +245,6 @@ class Decoder(nn.Module): past_len = 0 curr_len = x.shape[1] - if past_len < target_len and (past_len + curr_len) >= sw_size: fill_len = target_len - past_len x_fill = x[:, :fill_len, :] @@ -269,86 +274,39 @@ class Decoder(nn.Module): return x_out, stream_state, past_key_values -class Encoder(nn.Module): - def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1): - super().__init__() - config = Qwen2Config.from_dict(config_dict=encoder_args) - if is_flash_attn_2_available(): - config._attn_implementation_autoset = True - config._attn_implementation = "flash_attention_2" - else: - config._attn_implementation = "sdpa" - - logger.info("AudioVAE Encoder: using attn_implementation=%r", config._attn_implementation) - self.encoder = Qwen2Model(config) - self.input_dim = input_dim - self.hop_size = hop_size - self.latent_dim = latent_dim - self.fc1 = nn.Linear(input_dim, config.hidden_size, bias=False) - self.fc2 = nn.Linear(config.hidden_size, config.hidden_size) - self.fc3 = nn.Linear(config.hidden_size, latent_dim * 2) - self.norm = nn.LayerNorm(config.hidden_size) - self.patch_size = patch_size - if patch_size != -1: - config.num_hidden_layers = 4 - self.aggregator = Qwen2Model(config) - self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size)) - self.cls_embed.data.normal_(0, 0.02) - - def get_frames(self, x): - num_frames_total = (x.size(-1) + self.hop_size - 1) // self.hop_size - expected_len = (num_frames_total - 1) * self.hop_size + self.input_dim - padding_needed = expected_len - x.size(-1) - waveform = F.pad(x, (0, padding_needed), value=0.0) - frames = waveform.unfold(dimension=-1, size=self.input_dim, step=self.hop_size) - return frames - - def pad_patch_insert_cls(self, x): - bsz, _, dim = x.size() - num_frame = x.size(1) - r = num_frame % self.patch_size - pad_num = self.patch_size - r if r else 0 - x = F.pad(x, (0, 0, 0, pad_num), value=0.0) - x = x.reshape(-1, self.patch_size, dim) - x = torch.cat((x, self.cls_embed.expand(x.size(0), -1, -1)), dim=1) - x = x.reshape(bsz, -1, dim) - return x - - def forward(self, waveform): - x = self.get_frames(waveform) - x = self.fc1(x) - x = self.fc2(x) - x = self.encoder(inputs_embeds=x) - x = x.last_hidden_state - - if self.patch_size != -1: - x = self.pad_patch_insert_cls(x) - x = self.aggregator(inputs_embeds=x) - x = x.last_hidden_state - bsz, _, dim = x.size() - x = x.reshape(-1, self.patch_size + 1, dim) - x = x[:, -1:, :].reshape(bsz, -1, dim) - - x = self.fc3(x) - return x, waveform.unsqueeze(1) - - class AudioVAE(PreTrainedModel): config_class = AudioVAEConfig - def __init__(self, config: AudioVAEConfig): + def __init__(self, config): super().__init__(config) + self.config = config + enc_kwargs = config.enc_kwargs + dec_kwargs = config.dec_kwargs + for key in ("backbone", "input_dim", "latent_dim"): + if key not in enc_kwargs: + raise ValueError(f"AudioVAE.enc_kwargs missing required key: {key}") + for key in ("backbone", "output_dim", "latent_dim"): + if key not in dec_kwargs: + raise ValueError(f"AudioVAE.dec_kwargs missing required key: {key}") + + hop_size = enc_kwargs.get("hop_size", enc_kwargs["input_dim"]) + if enc_kwargs["input_dim"] != hop_size: + raise ValueError(f"AudioVAE encoder input_dim ({enc_kwargs['input_dim']}) != hop_size ({hop_size}).") + if hop_size != dec_kwargs["output_dim"]: + raise ValueError( + f"AudioVAE encoder hop_size ({hop_size}) != decoder output_dim ({dec_kwargs['output_dim']})." + ) self.encoder = Encoder( - encoder_args=config.enc_kwargs["backbone"], - input_dim=config.enc_kwargs["input_dim"], - hop_size=config.enc_kwargs.get("hop_size", 320), - latent_dim=config.enc_kwargs["latent_dim"], + encoder_args=enc_kwargs["backbone"], + input_dim=enc_kwargs["input_dim"], + hop_size=hop_size, + latent_dim=enc_kwargs["latent_dim"], patch_size=config.patch_size, ) self.decoder = Decoder( - decoder_args=config.dec_kwargs["backbone"], - output_dim=config.dec_kwargs["output_dim"], - latent_dim=config.dec_kwargs["latent_dim"], + decoder_args=dec_kwargs["backbone"], + output_dim=dec_kwargs["output_dim"], + latent_dim=dec_kwargs["latent_dim"], patch_size=config.patch_size, ) self.post_init() @@ -367,21 +325,50 @@ class AudioVAE(PreTrainedModel): if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() + @torch.inference_mode() def encode_latent(self, waveform, waveform_length): - from diffusers.models.autoencoders.autoencoder_oobleck import OobleckDiagonalGaussianDistribution + if waveform.ndim != 2: + raise ValueError(f"Expected waveform rank-2 [Batch, Time], got {tuple(waveform.shape)}") + if waveform_length.ndim != 1: + raise ValueError(f"Expected waveform_length rank-1 [Batch], got {tuple(waveform_length.shape)}") + if waveform.shape[0] != waveform_length.shape[0]: + raise ValueError( + "Batch mismatch: " + f"waveform batch={waveform.shape[0]} vs " + f"waveform_length batch={waveform_length.shape[0]}" + ) + if torch.any(waveform_length <= 0): + raise ValueError("waveform_length must be strictly positive.") frame_num = torch.ceil(waveform_length / self.config.enc_kwargs["input_dim"]).to(torch.int32) if self.config.patch_size != -1: frame_num = torch.ceil(frame_num / self.config.patch_size) - h, y = self.encoder(waveform) + h, _ = self.encoder(waveform) h = h.transpose(1, 2) - posterior = OobleckDiagonalGaussianDistribution(h) - latent = posterior.sample() + mean, logvar = torch.chunk(h, 2, dim=1) + logvar = torch.clamp(logvar, -30.0, 20.0) + std = torch.exp(0.5 * logvar) + latent = mean + std * torch.randn_like(mean) latent = latent.transpose(1, 2) return latent, frame_num + @torch.inference_mode() def decode(self, latent, past_key_values=None, use_cache=False, stream_state=(None, None, None), last_chunk=False): + if latent.dim() != 3: + raise ValueError(f"Expected latent rank-3 [B,T,D], got shape={tuple(latent.shape)}") + if latent.shape[0] <= 0: + raise ValueError("latent batch size must be positive.") + + target_dtype = next(self.decoder.parameters()).dtype + target_device = next(self.decoder.parameters()).device + if latent.dtype != target_dtype or latent.device != target_device: + latent = latent.to(device=target_device, dtype=target_dtype) + + expected_latent_dim = self.config.dec_kwargs["latent_dim"] + if latent.shape[-1] != expected_latent_dim: + raise ValueError(f"Latent dim mismatch in decode(): got {latent.shape[-1]}, expected {expected_latent_dim}") + waveform, stream_state, past_key_values = self.decoder.low_level_reconstruct( latent, past_key_values=past_key_values, diff --git a/vllm_omni/model_executor/models/common/ming/dit.py b/vllm_omni/model_executor/models/common/ming/dit.py new file mode 100644 index 000000000..3b4ab1122 --- /dev/null +++ b/vllm_omni/model_executor/models/common/ming/dit.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch +import torch.nn.functional as F +from torch import nn +from x_transformers.x_transformers import apply_rotary_pos_emb + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.to(self.weight.dtype) + x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps) + return x + + +class FeedForward(nn.Module): + def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate="none"): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + activation = nn.GELU(approximate=approximate) + project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation) + self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)) + + def forward(self, x): + return self.ff(x) + + +class Attention(nn.Module): + def __init__( + self, + dim, + heads=8, + dim_head=64, + dropout=0.0, + qk_norm=None, + pe_attn_head=None, + attn_mask_enabled=True, + ): + super().__init__() + self.dim = dim + self.heads = heads + self.inner_dim = dim_head * heads + self.dropout = dropout + self.to_q = nn.Linear(dim, self.inner_dim) + self.to_k = nn.Linear(dim, self.inner_dim) + self.to_v = nn.Linear(dim, self.inner_dim) + if qk_norm is None: + self.q_norm = None + self.k_norm = None + elif qk_norm == "rms_norm": + self.q_norm = RMSNorm(dim_head) + self.k_norm = RMSNorm(dim_head) + else: + raise ValueError(f"Unimplemented qk_norm: {qk_norm}") + + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(self.inner_dim, dim)) + self.to_out.append(nn.Dropout(dropout)) + self.pe_attn_head = pe_attn_head + self.attn_mask_enabled = attn_mask_enabled + + def forward(self, x, mask=None, rope=None): + batch_size = x.shape[0] + query = self.to_q(x) + key = self.to_k(x) + value = self.to_v(x) + + inner_dim = key.shape[-1] + head_dim = inner_dim // self.heads + query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + + if self.q_norm is not None: + query = self.q_norm(query) + if self.k_norm is not None: + key = self.k_norm(key) + + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + if self.pe_attn_head is not None: + on = self.pe_attn_head + query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale) + key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale) + else: + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + if self.attn_mask_enabled and mask is not None: + valid_sample_indices = mask.any(dim=1) + final_output = torch.zeros_like(query).to(query.device) + attn_mask = mask[valid_sample_indices] + query = query[valid_sample_indices] + key = key[valid_sample_indices] + value = value[valid_sample_indices] + attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) + attn_mask = attn_mask.expand(valid_sample_indices.sum().item(), self.heads, query.shape[-2], key.shape[-2]) + else: + attn_mask = None + + x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + if self.attn_mask_enabled and mask is not None: + final_output[valid_sample_indices] = x + x = final_output + + x = x.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) + x = x.to(query.dtype) + x = self.to_out[0](x) + x = self.to_out[1](x) + + if mask is not None: + mask = mask.unsqueeze(-1) + x = x.masked_fill(~mask, 0.0) + + return x + + +class DiTBlock(nn.Module): + def __init__( + self, + hidden_size, + num_heads, + mlp_ratio=4.0, + dropout=0.1, + qk_norm=None, + pe_attn_head=None, + attn_mask_enabled=True, + **kwargs, + ): + super().__init__() + self.norm1 = RMSNorm(hidden_size) + self.attn = Attention( + dim=hidden_size, + heads=num_heads, + dim_head=hidden_size // num_heads, + dropout=dropout, + qk_norm=qk_norm, + pe_attn_head=pe_attn_head, + attn_mask_enabled=attn_mask_enabled, + ) + self.norm2 = RMSNorm(hidden_size) + self.mlp = FeedForward(dim=hidden_size, mult=mlp_ratio, dropout=dropout, approximate="tanh") + + def forward(self, x, mask, rope): + x = x + self.attn(self.norm1(x), mask=mask, rope=rope) + x = x + self.mlp(self.norm2(x)) + return x + + +class FinalLayer(nn.Module): + def __init__(self, hidden_size, out_channels): + super().__init__() + self.norm_final = RMSNorm(hidden_size) + self.linear = nn.Linear(hidden_size, out_channels, bias=True) + + def forward(self, x): + x = self.norm_final(x) + x = self.linear(x) + return x + + +class CondEmbedder(nn.Module): + def __init__(self, input_feature_size, hidden_size): + super().__init__() + self.cond_embedder = nn.Linear(input_feature_size, hidden_size) + + def forward(self, llm_cond): + return self.cond_embedder(llm_cond) + + +def get_epss_timesteps(n, device, dtype): + dt = 1 / 32 + predefined_timesteps = { + 5: [0, 2, 4, 8, 16, 32], + 6: [0, 2, 4, 6, 8, 16, 32], + 7: [0, 2, 4, 6, 8, 16, 24, 32], + 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32], + 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32], + 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32], + } + t = predefined_timesteps.get(n, []) + if not t: + return torch.linspace(0, 1, n + 1, device=device, dtype=dtype) + return dt * torch.tensor(t, device=device, dtype=dtype) diff --git a/vllm_omni/model_executor/models/common/ming/fm.py b/vllm_omni/model_executor/models/common/ming/fm.py new file mode 100644 index 000000000..d2cbc6211 --- /dev/null +++ b/vllm_omni/model_executor/models/common/ming/fm.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from .dit import get_epss_timesteps + + +def apply_sway_sampling(t, sway_sampling_coef): + if sway_sampling_coef is None: + return t + return t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t) + + +def build_timesteps(steps, device, dtype, use_epss=True, sway_sampling_coef=-1.0): + if steps <= 0: + raise ValueError(f"steps must be positive, got {steps}") + if use_epss: + t = get_epss_timesteps(steps, device=device, dtype=dtype) + else: + t = torch.linspace(0, 1, steps + 1, device=device, dtype=dtype) + return apply_sway_sampling(t, sway_sampling_coef) + + +class Solver: + def __init__(self, func, y0, sigma=0.25, temperature=1.5) -> None: + self.func = func + self.y0 = y0 + self.sigma = sigma + self.temperature = temperature + + def integrate(self, t): + solution = torch.empty(len(t), *self.y0.shape, dtype=self.y0.dtype, device=self.y0.device) + solution[0] = self.y0 + + j = 1 + y0 = self.y0 + for t0, t1 in zip(t[:-1], t[1:]): + dt = t1 - t0 + f0 = self.func(t0, y0) + dy = dt * f0 + y1 = y0 + dy + + while j < len(t) and t1 >= t[j]: + solution[j] = self._linear_interp(t0, t1, y0, y1, t[j]) + j += 1 + + noise = torch.randn_like(y0) + shift = self.sigma * (self.temperature**0.5) * (abs(dt) ** 0.5) * noise + y0 = y1 + shift + + return solution + + def _linear_interp(self, t0, t1, y0, y1, t): + if t == t0: + return y0 + if t == t1: + return y1 + slope = (t - t0) / (t1 - t0) + return y0 + slope * (y1 - y0) + + +def integrate_cfm_steps(fn, y0, t, sde_args, sde_rnd, steps): + for step in range(steps): + dt = t[step + 1] - t[step] + y0 = y0 + fn(t[step], y0) * dt + y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step] + return y0 diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py index 08ed9e854..e7b305b1e 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py @@ -26,10 +26,10 @@ from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific +from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE, AudioVAEConfig from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig -from .audio_vae import AudioVAE, AudioVAEConfig from .prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT from .talker_module import CFM, Aggregator, DiT, MingAudioGenerator, build_tts_input from .text_processing import segment_and_normalize diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py index e0a49935d..8c527b8d9 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py @@ -697,6 +697,9 @@ class MingFlashOmniThinkerForConditionalGeneration( f"image_gen_scales={list(self.query_tokens_dict.keys())}" ) + def get_language_model(self): + return self.language_model + def _load_image_gen_query_tokens(self, model_path: str) -> None: """Load learnable query tokens used for image generation. diff --git a/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py b/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py index 53c130b90..18deb4600 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py @@ -110,12 +110,18 @@ BASE_CAPTION_TEMPLATE: dict[str, Any] = { } -def create_instruction(user_input: dict[str, Any]) -> str: +def create_instruction(user_input: Any) -> str | None: """Return a JSON caption string for ``audio_sequence[0]``. Only keys already present on the base template are merged in; unknown keys are silently ignored to keep the output schema stable. """ + if user_input is None: + return None + if isinstance(user_input, str): + return user_input + if not isinstance(user_input, dict): + raise ValueError(f"Ming instruction must be str, dict, or None; got {type(user_input).__name__}") caption = copy.deepcopy(BASE_CAPTION_TEMPLATE) item = caption["audio_sequence"][0] for key, value in user_input.items(): diff --git a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py index 6b2c6b2f2..cfb1090c2 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py @@ -27,223 +27,19 @@ from threading import Lock import torch import torch.nn as nn -import torch.nn.functional as F from transformers import PreTrainedTokenizerBase, Qwen2Config, Qwen2Model, StaticCache from vllm.logger import init_logger from vllm.platforms import current_platform -from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb +from x_transformers.x_transformers import RotaryEmbedding from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding - -from .audio_vae import AudioVAE +from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE +from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer, get_epss_timesteps +from vllm_omni.model_executor.models.common.ming.fm import apply_sway_sampling, integrate_cfm_steps logger = init_logger(__name__) -######################################################################## -# DiT Modules -# Ported from: -# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/modules.py -# Ported from: -# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/dit.py -######################################################################## - - -class RMSNorm(nn.Module): - def __init__(self, dim: int, eps: float = 1e-6): - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.ones(dim)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - if self.weight.dtype in [torch.float16, torch.bfloat16]: - x = x.to(self.weight.dtype) - x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps) - return x - - -class FeedForward(nn.Module): - def __init__( - self, dim: int, dim_out: int | None = None, mult: float = 4, dropout: float = 0.0, approximate: str = "none" - ): - super().__init__() - inner_dim = int(dim * mult) - dim_out = dim_out if dim_out is not None else dim - - activation = nn.GELU(approximate=approximate) - project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation) - self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.ff(x) - - -class Attention(nn.Module): - def __init__( - self, - dim: int, - heads: int = 8, - dim_head: int = 64, - dropout: float = 0.0, - qk_norm: str | None = None, - pe_attn_head: int | None = None, - attn_mask_enabled: bool = True, - ): - super().__init__() - self.dim = dim - self.heads = heads - self.inner_dim = dim_head * heads - self.dropout = dropout - - self.to_q = nn.Linear(dim, self.inner_dim) - self.to_k = nn.Linear(dim, self.inner_dim) - self.to_v = nn.Linear(dim, self.inner_dim) - if qk_norm is None: - self.q_norm = None - self.k_norm = None - elif qk_norm == "rms_norm": - self.q_norm = RMSNorm(dim_head) - self.k_norm = RMSNorm(dim_head) - else: - raise ValueError(f"Unimplemented qk_norm: {qk_norm}") - - self.to_out = nn.ModuleList([]) - self.to_out.append(nn.Linear(self.inner_dim, dim)) - self.to_out.append(nn.Dropout(dropout)) - - self.pe_attn_head = pe_attn_head - self.attn_mask_enabled = attn_mask_enabled - - def forward( - self, - x: torch.Tensor, - mask: torch.Tensor | None = None, - rope: tuple[torch.Tensor, torch.Tensor | None] | None = None, - ) -> torch.Tensor: - batch_size = x.shape[0] - - query = self.to_q(x) - key = self.to_k(x) - value = self.to_v(x) - - inner_dim = key.shape[-1] - head_dim = inner_dim // self.heads - query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) - value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) - - if self.q_norm is not None: - query = self.q_norm(query) - if self.k_norm is not None: - key = self.k_norm(key) - - if rope is not None: - freqs, xpos_scale = rope - q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) - - if self.pe_attn_head is not None: - on = self.pe_attn_head - query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale) - key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale) - else: - query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) - key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) - - if self.attn_mask_enabled and mask is not None: - valid_sample_indices = mask.any(dim=1) - final_output = torch.zeros_like(query).to(query.device) - - attn_mask = mask[valid_sample_indices] - query = query[valid_sample_indices] - key = key[valid_sample_indices] - value = value[valid_sample_indices] - attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) - attn_mask = attn_mask.expand(valid_sample_indices.sum().item(), self.heads, query.shape[-2], key.shape[-2]) - else: - attn_mask = None - - x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) - if self.attn_mask_enabled and mask is not None: - final_output[valid_sample_indices] = x - x = final_output - - x = x.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) - x = x.to(query.dtype) - - x = self.to_out[0](x) - x = self.to_out[1](x) - - if mask is not None: - mask = mask.unsqueeze(-1) - x = x.masked_fill(~mask, 0.0) - - return x - - -class DiTBlock(nn.Module): - """A DiT block with pre-norm and residual connections.""" - - def __init__( - self, - hidden_size: int, - num_heads: int, - mlp_ratio: float = 4.0, - dropout: float = 0.1, - qk_norm: str | None = None, - pe_attn_head: int | None = None, - attn_mask_enabled: bool = True, - **kwargs, - ): - super().__init__() - self.norm1 = RMSNorm(hidden_size) - self.attn = Attention( - dim=hidden_size, - heads=num_heads, - dim_head=hidden_size // num_heads, - dropout=dropout, - qk_norm=qk_norm, - pe_attn_head=pe_attn_head, - attn_mask_enabled=attn_mask_enabled, - ) - self.norm2 = RMSNorm(hidden_size) - self.mlp = FeedForward(dim=hidden_size, mult=mlp_ratio, dropout=dropout, approximate="tanh") - - def forward( - self, - x: torch.Tensor, - mask: torch.Tensor | None, - rope: tuple[torch.Tensor, torch.Tensor | None] | None, - ) -> torch.Tensor: - x = x + self.attn(self.norm1(x), mask=mask, rope=rope) - x = x + self.mlp(self.norm2(x)) - return x - - -class FinalLayer(nn.Module): - """The final layer of DiT.""" - - def __init__(self, hidden_size: int, out_channels: int): - super().__init__() - self.norm_final = RMSNorm(hidden_size) - self.linear = nn.Linear(hidden_size, out_channels, bias=True) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.norm_final(x) - x = self.linear(x) - return x - - -class CondEmbedder(nn.Module): - """Embeds LLM hidden states with optional CFG dropout.""" - - def __init__(self, input_feature_size: int, hidden_size: int): - super().__init__() - self.cond_embedder = nn.Linear(input_feature_size, hidden_size) - - def forward(self, llm_cond: torch.Tensor) -> torch.Tensor: - return self.cond_embedder(llm_cond) - - class DiT(nn.Module): """Diffusion model with a Transformer backbone for audio latent generation.""" @@ -324,29 +120,6 @@ class DiT(nn.Module): return model_out[:, -x.shape[1] :, :] -######################################################################################### -# CFM -# Ported from: -# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/cfm.py -######################################################################################### - - -def get_epss_timesteps(n, device, dtype): - dt = 1 / 32 - predefined_timesteps = { - 5: [0, 2, 4, 8, 16, 32], - 6: [0, 2, 4, 6, 8, 16, 32], - 7: [0, 2, 4, 6, 8, 16, 24, 32], - 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32], - 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32], - 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32], - } - t = predefined_timesteps.get(n, []) - if not t: - return torch.linspace(0, 1, n + 1, device=device, dtype=dtype) - return dt * torch.tensor(t, device=device, dtype=dtype) - - class CFM(nn.Module): """Conditional Flow Matching module for audio latent generation.""" @@ -391,15 +164,8 @@ class CFM(nn.Module): pred, null_pred = torch.chunk(pred_cfg, 2, dim=0) return pred + (pred - null_pred) * sde_args[0] - if self.sway_sampling_coef is not None: - t = t + self.sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t) - - for step in range(self.steps): - dt = t[step + 1] - t[step] - y0 = y0 + fn(t[step], y0) * dt - y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step] - - return y0 + t = apply_sway_sampling(t, self.sway_sampling_coef) + return integrate_cfm_steps(fn, y0, t, sde_args, sde_rnd, self.steps) class CFMGraphExecutor: diff --git a/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py b/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py index 5f54687c0..0a096df86 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py @@ -22,7 +22,8 @@ from .spk_embedding import SpkembExtractor from .talker_module import resample if TYPE_CHECKING: - from .audio_vae import AudioVAE + from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE + from .talker_module import Aggregator logger = init_logger(__name__) diff --git a/vllm_omni/model_executor/models/ming_tts/__init__.py b/vllm_omni/model_executor/models/ming_tts/__init__.py new file mode 100644 index 000000000..35afa2ce5 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/__init__.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from .config_ming_tts import MingDenseConfig +from .ming_tts import MingTTSForConditionalGeneration +from .ming_tts_audio_vae import MingAudioVAEModel +from .ming_tts_llm import MingLLMModel + +__all__ = [ + "MingDenseConfig", + "MingTTSForConditionalGeneration", + "MingLLMModel", + "MingAudioVAEModel", +] diff --git a/vllm_omni/model_executor/models/ming_tts/aggregator.py b/vllm_omni/model_executor/models/ming_tts/aggregator.py new file mode 100644 index 000000000..3233c9659 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/aggregator.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/fm/dit.py +from __future__ import annotations + +import torch +import torch.nn as nn +from x_transformers.x_transformers import RotaryEmbedding + +from vllm_omni.model_executor.models.common.ming.dit import DiTBlock, FinalLayer + + +class Aggregator(nn.Module): + def __init__( + self, + in_channels=4, + hidden_size=1152, + depth=28, + num_heads=16, + mlp_ratio=4.0, + llm_input_dim=896, + **kwargs, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels + self.num_heads = num_heads + + self.word_embedder = nn.Embedding(1, hidden_size) + self.x_embedder = nn.Linear(in_channels, hidden_size) + self.hidden_size = hidden_size + self.rotary_embed = RotaryEmbedding(hidden_size // num_heads) + self.blocks = nn.ModuleList( + [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)] + ) + self.final_layer = FinalLayer(hidden_size, llm_input_dim) + + def forward(self, x, mask=None): + if x.ndim != 3: + raise ValueError(f"Expected x rank-3 [Batch, Time, Dimension], got {tuple(x.shape)}") + if x.shape[-1] != self.in_channels: + raise ValueError(f"x feature dim mismatch: got {x.shape[-1]}, expected {self.in_channels}") + + # [Batch, Time, Dimension] -> [Batch, Time, Hidden]. + x = self.x_embedder(x) + cls_embed = self.word_embedder(torch.zeros((x.shape[0], 1), dtype=torch.long, device=x.device)) + # Prepend a learned CLS token: [Batch, Time, Hidden] -> [Batch, Time + 1, Hidden]. + x = torch.cat([cls_embed, x], dim=1) + + rope = self.rotary_embed.forward_from_seq_len(x.shape[1]) + if mask is not None: + if mask.ndim != 2: + raise ValueError(f"Expected mask rank-2 [Batch, Time], got {tuple(mask.shape)}") + if mask.shape[0] != x.shape[0] or mask.shape[1] != x.shape[1] - 1: + raise ValueError( + f"Mask shape mismatch: got {tuple(mask.shape)}, expected {(x.shape[0], x.shape[1] - 1)}" + ) + mask_pad = mask.clone().detach()[:, :1] + mask = torch.cat([mask_pad, mask], dim=-1) + for block in self.blocks: + x = block(x, mask, rope) + x = self.final_layer(x) + # Keep the CLS projection only: [Batch, Time + 1, Hidden] -> [Batch, 1, Hidden]. + return x[:, :1, :] + + +__all__ = ["Aggregator"] diff --git a/vllm_omni/model_executor/models/ming_tts/audio_prep.py b/vllm_omni/model_executor/models/ming_tts/audio_prep.py new file mode 100644 index 000000000..b63e70cc8 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/audio_prep.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +import math +from io import BytesIO +from typing import Any + +import torch + +from .config_ming_tts import ( + AUDIO_FRAME_HOP, + LATENT_DIM, + PATCH_SIZE, + SAMPLE_RATE, + VAE_PATCH_SIZE, + VISION_START_TOKEN_ID, + MingTTSConfig, +) + + +def pad_prompt_waveform( + waveform: Any, + *, + patch_size: int = PATCH_SIZE, + sample_rate: int = SAMPLE_RATE, +) -> torch.Tensor: + tensor = coerce_prompt_waveform(waveform) + pad_align = int((float(sample_rate) / 12.5) * int(patch_size)) + new_len = ((int(tensor.shape[-1]) + pad_align - 1) // pad_align) * pad_align + if new_len == int(tensor.shape[-1]): + return tensor + padded = torch.zeros((1, new_len), dtype=tensor.dtype, device=tensor.device) + padded[:, : tensor.shape[-1]] = tensor + return padded + + +def coerce_prompt_waveform(value: Any) -> torch.Tensor: + if value is None: + raise ValueError("prompt waveform cannot be None") + if isinstance(value, torch.Tensor): + tensor = value.detach() + if tensor.ndim == 1: + return tensor.unsqueeze(0).to(torch.float32) + if tensor.ndim == 2: + if tensor.shape[0] != 1: + return tensor.reshape(1, -1).to(torch.float32) + return tensor.to(torch.float32) + raise ValueError(f"Unsupported Ming prompt waveform rank: {tuple(tensor.shape)}") + if isinstance(value, (list, tuple)): + parts = [coerce_prompt_waveform(item) for item in value if item is not None] + if not parts: + raise ValueError("prompt waveform list was empty") + return torch.cat(parts, dim=-1) + return coerce_prompt_waveform(torch.as_tensor(value)) + + +def coerce_speaker_embeddings(value: Any, *, use_zero_spk_emb: bool = False) -> list[torch.Tensor] | None: + if value is None: + return [torch.zeros((192,), dtype=torch.float32)] if use_zero_spk_emb else None + if isinstance(value, torch.Tensor): + tensor = value.detach() + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0) + if tensor.ndim != 2: + raise ValueError(f"Unsupported Ming speaker embedding shape: {tuple(tensor.shape)}") + items = [row.reshape(-1).to(torch.float32).cpu() for row in tensor] + elif isinstance(value, (list, tuple)): + if value and all(not isinstance(item, (list, tuple, torch.Tensor)) for item in value): + items = [torch.as_tensor(value).detach().reshape(-1).to(torch.float32).cpu()] + else: + items = [] + for item in value: + if item is None: + continue + if not isinstance(item, torch.Tensor): + item = torch.as_tensor(item) + items.append(item.detach().reshape(-1).to(torch.float32).cpu()) + else: + return coerce_speaker_embeddings(torch.as_tensor(value), use_zero_spk_emb=use_zero_spk_emb) + if not items: + return [torch.zeros((192,), dtype=torch.float32)] if use_zero_spk_emb else None + for item in items: + if int(item.numel()) != 192: + raise ValueError(f"Ming speaker embedding must have 192 dims, got {int(item.numel())}") + return items + + +def count_prompt_latent_patches( + value: Any, + *, + patch_size: int = PATCH_SIZE, + latent_dim: int = LATENT_DIM, +) -> int: + if value is None: + return 0 + if not isinstance(value, torch.Tensor): + value = torch.as_tensor(value) + latents = value.detach() + if latents.ndim == 3 and latents.shape[0] == 1: + latents = latents.squeeze(0) + if latents.ndim == 3 and latents.shape[-2:] == (patch_size, latent_dim): + return int(latents.shape[0]) + if latents.ndim != 2 or latents.shape[-1] != latent_dim: + raise ValueError(f"Unsupported Ming prompt_latents shape: {tuple(latents.shape)}") + if latents.shape[0] % patch_size != 0: + raise ValueError( + f"Ming prompt_latents frame count must be divisible by patch_size={patch_size}, " + f"got frames={int(latents.shape[0])}" + ) + return int(latents.shape[0] // patch_size) + + +def count_prompt_waveform_patches( + value: Any, + *, + patch_size: int = PATCH_SIZE, + frame_hop: int = AUDIO_FRAME_HOP, + vae_patch_size: int = VAE_PATCH_SIZE, +) -> int: + if value is None: + return 0 + waveform = pad_prompt_waveform(value, patch_size=patch_size) + frame_count = int(math.ceil(float(waveform.shape[-1]) / float(frame_hop))) + latent_frames = int(math.ceil(float(frame_count) / float(vae_patch_size))) + if latent_frames % int(patch_size) != 0: + raise ValueError( + f"Ming prompt waveform produced latent frame count not divisible by patch_size={patch_size}: " + f"frames={latent_frames}" + ) + return int(latent_frames // int(patch_size)) + + +def _normalize_prompt_waveform(value: Any, *, target_sr: int) -> torch.Tensor: + if isinstance(value, bytes): + import torchaudio + + waveform, sr = torchaudio.load(BytesIO(value)) + waveform = waveform[:1].to(torch.float32) + if int(sr) != int(target_sr): + from torchaudio.functional import resample as resample_audio + + waveform = resample_audio(waveform, int(sr), int(target_sr)) + return waveform + + if isinstance(value, tuple) and len(value) == 2 and isinstance(value[1], int): + waveform = coerce_prompt_waveform(value[0]) + if int(value[1]) != int(target_sr): + from torchaudio.functional import resample as resample_audio + + waveform = resample_audio(waveform, int(value[1]), int(target_sr)) + return waveform + + if isinstance(value, dict): + samples = value.get("samples", value.get("array", value.get("waveform"))) + sr = value.get("sample_rate", value.get("sr", target_sr)) + return _normalize_prompt_waveform((samples, int(sr)), target_sr=target_sr) + + return coerce_prompt_waveform(value) + + +def _coerce_prompt_latents( + value: Any, + *, + patch_size: int, + latent_dim: int, +) -> dict[str, torch.Tensor] | None: + if value is None: + return None + if not isinstance(value, torch.Tensor): + value = torch.as_tensor(value) + + latents = value.detach() + if latents.ndim == 3 and latents.shape[0] == 1: + latents = latents.squeeze(0) + + if latents.ndim == 3 and latents.shape[-2:] == (patch_size, latent_dim): + patches = latents + # [Patch, Time, Dimension] -> [Frame, Dimension] for history seeding. + frames = patches.reshape(-1, latent_dim) + return {"patches": patches, "frames": frames} + + if latents.ndim != 2 or latents.shape[-1] != latent_dim: + raise ValueError(f"Unsupported prompt latent shape: {tuple(latents.shape)}") + if latents.shape[0] % patch_size != 0: + raise ValueError( + f"Prompt latent frame count must be divisible by patch_size={patch_size}, " + f"got frames={int(latents.shape[0])}" + ) + # [Frame, Dimension] -> [Patch, Time, Dimension] for Aggregator prompt slots. + patches = latents.reshape(-1, patch_size, latent_dim) if latents.shape[0] > 0 else None + return {"patches": patches, "frames": latents} + + +def _initial_history( + frames: torch.Tensor | None, + *, + history_size: int, + latent_dim: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + history = torch.zeros((history_size, latent_dim), device=device, dtype=dtype) + if frames is None or frames.numel() == 0: + return history + frames = frames.to(device=device, dtype=dtype) + take = min(history_size, int(frames.shape[0])) + history[-take:] = frames[-take:] + return history + + +def _take_scalar(value: Any, idx: int) -> float | None: + if not isinstance(value, torch.Tensor) or value.numel() == 0: + return None + return float(value.reshape(-1)[idx].item()) + + +def _find_audio_placeholder_positions(input_ids: torch.Tensor, cfg: MingTTSConfig) -> torch.Tensor: + dummy_pos = (input_ids == cfg.audio_dummy_token_id).nonzero(as_tuple=True)[0] + if dummy_pos.numel() == 0: + return dummy_pos + + audio_start_pos = (input_ids == cfg.audio_start_token_id).nonzero(as_tuple=True)[0] + audio_end_pos = (input_ids == cfg.audio_end_token_id).nonzero(as_tuple=True)[0] + if audio_start_pos.numel() == 0: + return dummy_pos + + start = int(audio_start_pos[0].item()) + end = int(audio_end_pos[0].item()) if audio_end_pos.numel() > 0 else int(input_ids.shape[0]) + keep = (dummy_pos > start) & (dummy_pos < end) + filtered = dummy_pos[keep] + return filtered if filtered.numel() > 0 else dummy_pos + + +def _find_speaker_placeholder_positions(input_ids: torch.Tensor, hf_config: Any) -> list[int]: + vision_start_token_id = getattr(hf_config, "vision_start_token_id", VISION_START_TOKEN_ID) + vision_start_pos = (input_ids == int(vision_start_token_id)).nonzero(as_tuple=True)[0] + if vision_start_pos.numel() == 0: + return [] + + slots = [] + for pos in vision_start_pos: + slot = int(pos.item()) + 1 + if slot < int(input_ids.shape[0]): + slots.append(slot) + return slots + + +__all__ = [ + "coerce_prompt_waveform", + "coerce_speaker_embeddings", + "count_prompt_latent_patches", + "count_prompt_waveform_patches", + "pad_prompt_waveform", + "_coerce_prompt_latents", + "_find_audio_placeholder_positions", + "_find_speaker_placeholder_positions", + "_initial_history", + "_normalize_prompt_waveform", + "_take_scalar", +] diff --git a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py new file mode 100644 index 000000000..b42c4fdca --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ming dense checkpoint config adapters.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from transformers import PretrainedConfig, Qwen2Config + +from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig + +from .constants import ( + AGGREGATOR_HIDDEN_SIZE, + AUDIO_DUMMY_TOKEN_ID, + AUDIO_END_TOKEN_ID, + AUDIO_EOS_TOKEN_ID, + AUDIO_FRAME_HOP, + AUDIO_START_TOKEN_ID, + DEFAULT_CFG, + DEFAULT_SIGMA, + DEFAULT_TEMPERATURE, + HISTORY_PATCH_SIZE, + KEY_CFG, + KEY_CHUNK_ID, + KEY_DECODE_STEP, + KEY_LAST_STOP_PROB, + KEY_LATENT_HISTORY, + KEY_MAX_DECODE_STEPS, + KEY_MIN_DECODE_STEPS, + KEY_NEXT_EMBEDS, + KEY_PROMPT_LATENT_TAIL, + KEY_PROMPT_LATENTS, + KEY_REQUEST_ID, + KEY_SIGMA, + KEY_SPEAKER_EMBEDDING, + KEY_TEMPERATURE, + KEY_TEXT_MODE, + LATENT_CHUNK_SIZE, + LATENT_DIM, + LATENT_LEFT_CONTEXT, + LLM_HIDDEN_SIZE, + LLM_VOCAB_SIZE, + MAX_DECODE_STEPS, + PATCH_SIZE, + SAMPLE_RATE, + STOP_HEAD_MIN_STEPS, + STOP_HEAD_THRESHOLD, + TEXT_EOS_TOKEN_ID, + VAE_PATCH_SIZE, + VISION_START_TOKEN_ID, +) +from .validation import _coerce_audio_vae_config, _nested_get, _to_plain_dict, validate_ming_tts_config + + +def _coerce_qwen2_config(value: Any) -> Qwen2Config: + if isinstance(value, Qwen2Config): + return value + if isinstance(value, PretrainedConfig): + return Qwen2Config.from_dict(value.to_dict()) + if isinstance(value, dict): + return Qwen2Config.from_dict(dict(value)) + raise TypeError(f"Unsupported llm_config type for Ming dense config: {type(value)!r}") + + +class MingDenseConfig(PretrainedConfig): + # The upstream checkpoint declares model_type="dense". Keep it for HF + # config compatibility; deploy/ming_tts.yaml selects the vLLM-Omni + # pipeline via pipeline: ming_tts. + model_type = "dense" + + def __init__( + self, + llm_config: Qwen2Config | dict[str, Any] | None = None, + ditar_config: dict[str, Any] | None = None, + aggregator_config: dict[str, Any] | None = None, + audio_tokenizer_config: AudioVAEConfig | dict[str, Any] | None = None, + architectures: list[str] | None = None, + **kwargs: Any, + ) -> None: + self.llm_config = _coerce_qwen2_config(llm_config or {}) + self.ditar_config = dict(ditar_config or {}) + self.aggregator_config = dict(aggregator_config or {}) + self.audio_tokenizer_config = _coerce_audio_vae_config(audio_tokenizer_config) + super().__init__(architectures=architectures, **kwargs) + + def get_text_config(self, decoder: bool = False, **kwargs: Any) -> Qwen2Config: + del decoder, kwargs + return self.llm_config + + +@dataclass +class MingTTSConfig: + """Flat config object shared by Stage-1 and Stage-2. Build via from_hf_config().""" + + llm_hidden_size: int = LLM_HIDDEN_SIZE + llm_vocab_size: int = LLM_VOCAB_SIZE + llm_config: dict[str, Any] = field(default_factory=dict) + + latent_dim: int = LATENT_DIM + patch_size: int = PATCH_SIZE + history_patch_size: int = HISTORY_PATCH_SIZE + + ditar_config: dict[str, Any] = field(default_factory=dict) + aggregator_config: dict[str, Any] = field(default_factory=dict) + + audio_tokenizer_config: AudioVAEConfig | None = None + vae_patch_size: int = VAE_PATCH_SIZE + sample_rate: int = SAMPLE_RATE + audio_frame_hop: int = AUDIO_FRAME_HOP + + cfg: float = DEFAULT_CFG + sigma: float = DEFAULT_SIGMA + temperature: float = DEFAULT_TEMPERATURE + stop_head_min_steps: int = STOP_HEAD_MIN_STEPS + stop_head_threshold: float = STOP_HEAD_THRESHOLD + max_decode_steps: int = MAX_DECODE_STEPS + + latent_chunk_size: int = LATENT_CHUNK_SIZE + latent_left_context: int = LATENT_LEFT_CONTEXT + + text_eos_token_id: int = TEXT_EOS_TOKEN_ID + audio_dummy_token_id: int = AUDIO_DUMMY_TOKEN_ID + audio_start_token_id: int = AUDIO_START_TOKEN_ID + audio_end_token_id: int = AUDIO_END_TOKEN_ID + audio_eos_token_id: int = AUDIO_EOS_TOKEN_ID + + @classmethod + def from_hf_config(cls, hf_config: PretrainedConfig) -> MingTTSConfig: + llm_raw = getattr(hf_config, "llm_config", {}) or {} + ditar_raw = getattr(hf_config, "ditar_config", {}) or {} + agg_raw = getattr(hf_config, "aggregator_config", {}) or {} + atc_raw = getattr(hf_config, "audio_tokenizer_config", None) + + llm_dict = _to_plain_dict(llm_raw) + ditar = _to_plain_dict(ditar_raw) + agg = _to_plain_dict(agg_raw) + ditar.setdefault("attn_backend", "torch") + + atc = _coerce_audio_vae_config(atc_raw) + atc_enc_latent_dim = _nested_get(atc, "enc_kwargs", "latent_dim", default=LATENT_DIM) + atc_patch_size = _nested_get(atc, "patch_size", default=VAE_PATCH_SIZE) + atc_sample_rate = _nested_get(atc, "sample_rate", default=SAMPLE_RATE) + + enc_hop_size = _nested_get(atc, "enc_kwargs", "hop_size", default=AUDIO_FRAME_HOP) + + cfg = cls( + llm_hidden_size=llm_dict.get("hidden_size", LLM_HIDDEN_SIZE), + llm_vocab_size=llm_dict.get("vocab_size", LLM_VOCAB_SIZE), + llm_config=llm_dict, + latent_dim=atc_enc_latent_dim, + patch_size=ditar.get("patch_size", PATCH_SIZE), + history_patch_size=ditar.get("history_patch_size", HISTORY_PATCH_SIZE), + ditar_config=ditar, + aggregator_config=agg, + audio_tokenizer_config=atc, + vae_patch_size=atc_patch_size, + sample_rate=atc_sample_rate, + audio_frame_hop=enc_hop_size if enc_hop_size is not None else AUDIO_FRAME_HOP, + ) + return cfg + + def validate(self) -> None: + validate_ming_tts_config(self) + + +__all__ = [ + "AGGREGATOR_HIDDEN_SIZE", + "AUDIO_DUMMY_TOKEN_ID", + "AUDIO_END_TOKEN_ID", + "AUDIO_EOS_TOKEN_ID", + "AUDIO_FRAME_HOP", + "AUDIO_START_TOKEN_ID", + "DEFAULT_CFG", + "DEFAULT_SIGMA", + "DEFAULT_TEMPERATURE", + "HISTORY_PATCH_SIZE", + "KEY_CFG", + "KEY_CHUNK_ID", + "KEY_DECODE_STEP", + "KEY_LAST_STOP_PROB", + "KEY_LATENT_HISTORY", + "KEY_MAX_DECODE_STEPS", + "KEY_MIN_DECODE_STEPS", + "KEY_NEXT_EMBEDS", + "KEY_PROMPT_LATENT_TAIL", + "KEY_PROMPT_LATENTS", + "KEY_REQUEST_ID", + "KEY_SIGMA", + "KEY_SPEAKER_EMBEDDING", + "KEY_TEMPERATURE", + "KEY_TEXT_MODE", + "LATENT_CHUNK_SIZE", + "LATENT_DIM", + "LATENT_LEFT_CONTEXT", + "LLM_HIDDEN_SIZE", + "LLM_VOCAB_SIZE", + "MAX_DECODE_STEPS", + "MingDenseConfig", + "MingTTSConfig", + "PATCH_SIZE", + "SAMPLE_RATE", + "STOP_HEAD_MIN_STEPS", + "STOP_HEAD_THRESHOLD", + "TEXT_EOS_TOKEN_ID", + "VAE_PATCH_SIZE", + "VISION_START_TOKEN_ID", +] diff --git a/vllm_omni/model_executor/models/ming_tts/constants.py b/vllm_omni/model_executor/models/ming_tts/constants.py new file mode 100644 index 000000000..2f59e2094 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/constants.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +# --------------------------------------------------------------------------- +# Token IDs (confirmed from tokenizer_config.json) +# --------------------------------------------------------------------------- + +AUDIO_DUMMY_TOKEN_ID = 151705 # +AUDIO_START_TOKEN_ID = 151706 # +AUDIO_EOS_TOKEN_ID = 151704 # +VISION_START_TOKEN_ID = 151652 # <|vision_start|> + +TEXT_EOS_TOKEN_ID = 151669 # + + +# --------------------------------------------------------------------------- +# Architectural constants (confirmed from original config.json) +# --------------------------------------------------------------------------- + +LATENT_DIM = 64 +PATCH_SIZE = 4 +HISTORY_PATCH_SIZE = 32 +LLM_HIDDEN_SIZE = 896 +LLM_VOCAB_SIZE = 151936 +AGGREGATOR_HIDDEN_SIZE = 1024 +VAE_PATCH_SIZE = 4 +SAMPLE_RATE = 44100 + +# AudioVAE frame/hop geometry (confirmed) +AUDIO_FRAME_HOP = 882 # enc input_dim / hop_size / dec output_dim + +# stop_head defaults +STOP_HEAD_MIN_STEPS = 3 +STOP_HEAD_THRESHOLD = 0.5 + +# FlowLoss sampling defaults +DEFAULT_CFG = 2.0 +DEFAULT_SIGMA = 0.25 +DEFAULT_TEMPERATURE = 0.0 + +# Connector / Stage-2 streaming defaults (runtime tuning) +LATENT_CHUNK_SIZE = 25 +LATENT_LEFT_CONTEXT = 0 +MAX_DECODE_STEPS = 200 + + +# --------------------------------------------------------------------------- +# seq_data.extra_data keys +# --------------------------------------------------------------------------- + +KEY_LATENT_HISTORY = "ming_latent_history" +KEY_DECODE_STEP = "ming_decode_step" +KEY_LAST_STOP_PROB = "ming_last_stop_prob" +KEY_NEXT_EMBEDS = "ming_next_embeds" +KEY_PROMPT_LATENTS = "ming_prompt_latents" +KEY_PROMPT_LATENT_TAIL = "ming_prompt_latent_tail" +KEY_SPEAKER_EMBEDDING = "ming_speaker_embedding" +KEY_REQUEST_ID = "ming_request_id" +KEY_CHUNK_ID = "ming_chunk_id" +KEY_CFG = "ming_cfg" +KEY_SIGMA = "ming_sigma" +KEY_TEMPERATURE = "ming_temperature" +KEY_MAX_DECODE_STEPS = "ming_max_decode_steps" +KEY_MIN_DECODE_STEPS = "ming_min_decode_steps" +KEY_TEXT_MODE = "ming_text_mode" diff --git a/vllm_omni/model_executor/models/ming_tts/flowloss_head.py b/vllm_omni/model_executor/models/ming_tts/flowloss_head.py new file mode 100644 index 000000000..3f300cec0 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/flowloss_head.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/fm/flowloss.py + +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.common.ming.fm import Solver, build_timesteps + +from .fm.dit import DiT + + +class FlowLoss(nn.Module): + """Diffusion Loss""" + + def __init__(self, z_channels, llm_cond_dim, **kwargs): + super().__init__() + self.z_channels = z_channels + # Preserve checkpoint parameter names under flowloss.cfm.model.* while + # keeping the inference-only CFM sampling logic local to FlowLoss. + self.cfm = nn.Module() + self.cfm.add_module("model", DiT(in_channels=z_channels, llm_cond_dim=llm_cond_dim, **kwargs)) + + @torch.no_grad() + def sample(self, z, latent_history, cfg=2.0, patch_size=1, sigma=0.25, temperature=0): + # z: [Batch, Time=1, Dimension] + # latent_history: [Batch, History_Time, Dimension] + if z.ndim != 3: + raise ValueError(f"Expected conditioning rank-3 [Batch, Time, Dimension], got {tuple(z.shape)}") + if latent_history.ndim != 3: + raise ValueError( + f"Expected latent_history rank-3 [Batch, Time, Dimension], got {tuple(latent_history.shape)}" + ) + if z.shape[0] != latent_history.shape[0]: + raise ValueError( + f"Batch mismatch across conditioning and latent_history: {z.shape[0]}, {latent_history.shape[0]}" + ) + dtype = z.dtype + noise = torch.randn(z.shape[0], self.z_channels, patch_size, device=z.device, dtype=torch.float32) + + def fn(t, x): + x = x.to(dtype) + t = t.to(dtype) + if cfg < 1e-5: + if t.ndim == 0: + t = t.repeat(x.shape[0]) + pred = self.cfm.model(x=x, t=t, c=torch.zeros_like(z), latent_history=latent_history) + return pred[:, -patch_size:, :].float() + + pred_cfg = self.cfm.model.forward_with_cfg( + x=x, t=t, c=z, latent_history=latent_history, cfg_scale=cfg, patch_size=patch_size + ) + pred, null_pred = torch.chunk(pred_cfg, 2, dim=0) + return (pred + (pred - null_pred) * cfg).float() + + t = build_timesteps( + 10, + device=noise.device, + dtype=noise.dtype, + use_epss=True, + sway_sampling_coef=-1.0, + ) + return Solver(fn, noise.transpose(1, 2), sigma=sigma, temperature=temperature).integrate(t)[-1].to(dtype) diff --git a/vllm_omni/model_executor/models/ming_tts/fm/__init__.py b/vllm_omni/model_executor/models/ming_tts/fm/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/fm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/ming_tts/fm/dit.py b/vllm_omni/model_executor/models/ming_tts/fm/dit.py new file mode 100644 index 000000000..9aa05d4a2 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/fm/dit.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/fm/dit.py + +import torch +import torch.nn as nn +from x_transformers.x_transformers import RotaryEmbedding + +from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding +from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer + + +class DiT(nn.Module): + def __init__( + self, + in_channels=4, + hidden_size=1024, + depth=28, + num_heads=16, + mlp_ratio=4.0, + llm_cond_dim=896, + **kwargs, + ): + super().__init__() + + self.in_channels = in_channels + self.out_channels = in_channels + self.num_heads = num_heads + self.t_embedder = DiTTimestepEmbedding(hidden_size) + self.x_embedder = nn.Linear(in_channels, hidden_size) + self.c_embedder = CondEmbedder(llm_cond_dim, hidden_size) + self.hidden_size = hidden_size + self.rotary_embed = RotaryEmbedding(hidden_size // num_heads) + self.blocks = nn.ModuleList( + [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)] + ) + self.final_layer = FinalLayer(hidden_size, self.out_channels) + + def forward(self, x, t, c, latent_history, mask=None): + if x.ndim != 3: + raise ValueError(f"Expected x rank-3 [Batch, Time, Dimension], got {tuple(x.shape)}") + if latent_history.ndim != 3: + raise ValueError( + f"Expected latent_history rank-3 [Batch, Time, Dimension], got {tuple(latent_history.shape)}" + ) + if c.ndim != 3: + raise ValueError(f"Expected conditioning rank-3 [Batch, Time, Dimension], got {tuple(c.shape)}") + if x.shape[0] != latent_history.shape[0] or x.shape[0] != c.shape[0]: + raise ValueError( + "Batch mismatch across x, conditioning, and latent_history: " + f"{x.shape[0]}, {c.shape[0]}, {latent_history.shape[0]}" + ) + if x.shape[-1] != self.in_channels: + raise ValueError(f"x feature dim mismatch: got {x.shape[-1]}, expected {self.in_channels}") + if latent_history.shape[-1] != self.in_channels: + raise ValueError( + f"latent_history feature dim mismatch: got {latent_history.shape[-1]}, expected {self.in_channels}" + ) + if t.ndim == 0: + t = t.reshape(1) + if t.ndim != 1: + raise ValueError(f"Expected timestep rank-1 [Batch], got {tuple(t.shape)}") + if t.shape[0] != x.shape[0]: + raise ValueError(f"Timestep batch mismatch: got {t.shape[0]}, expected {x.shape[0]}") + + t = self.t_embedder(t).unsqueeze(1) + x_now = self.x_embedder(x) + x_history = self.x_embedder(latent_history) + x = torch.cat([x_history, x_now], dim=1) + c = self.c_embedder(c) + y = t + c + x = torch.cat([y, x], dim=1) + rope = self.rotary_embed.forward_from_seq_len(x.shape[1]) + + if mask is not None: + if mask.ndim != 2: + raise ValueError(f"Expected mask rank-2 [Batch, Time], got {tuple(mask.shape)}") + if mask.shape[0] != x_now.shape[0] or mask.shape[1] != x_now.shape[1]: + raise ValueError( + f"Mask shape mismatch: got {tuple(mask.shape)}, expected {(x_now.shape[0], x_now.shape[1])}" + ) + mask_pad = mask.clone().detach()[:, :1].expand(-1, x_history.shape[1] + c.shape[1]) + mask = torch.cat([mask_pad, mask], dim=-1) + for block in self.blocks: + x = block(x, mask, rope) + x = self.final_layer(x) + return x + + def forward_with_cfg(self, x, t, c, cfg_scale, latent_history, patch_size): + if not cfg_scale == 1: + x = torch.cat([x, x], dim=0) + latent_history = torch.cat([latent_history, latent_history], dim=0) + fake_latent = torch.zeros_like(c) + c = torch.cat([c, fake_latent], dim=0) + if t.ndim == 0: + t = t.repeat(x.shape[0]) + model_out = self.forward(x, t, c, latent_history) + return model_out[:, -patch_size:, :] diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts.py b/vllm_omni/model_executor/models/ming_tts/ming_tts.py new file mode 100644 index 000000000..667aded27 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/ming_tts.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from collections.abc import Iterable +from functools import cached_property +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.model_executor.models import SupportsPP +from vllm.model_executor.models.utils import init_vllm_registered_model +from vllm.v1.sample.sampler import Sampler + +from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin + +from .audio_prep import ( + _coerce_prompt_latents, + _find_audio_placeholder_positions, + _find_speaker_placeholder_positions, + _initial_history, + _take_scalar, + coerce_speaker_embeddings, +) +from .config_ming_tts import ( + AUDIO_START_TOKEN_ID, + KEY_CFG, + KEY_DECODE_STEP, + KEY_LAST_STOP_PROB, + KEY_LATENT_HISTORY, + KEY_MAX_DECODE_STEPS, + KEY_MIN_DECODE_STEPS, + KEY_NEXT_EMBEDS, + KEY_PROMPT_LATENT_TAIL, + KEY_REQUEST_ID, + KEY_SIGMA, + KEY_SPEAKER_EMBEDDING, + KEY_TEMPERATURE, + KEY_TEXT_MODE, + MingTTSConfig, +) +from .patch_emission import MING_STOP_REASON_KEY +from .prompt_encoder import _resolve_prompt_latents + + +class _ModelSampleAdapter(nn.Module): + def __init__(self, model: nn.Module): + super().__init__() + self.model = model + + def forward(self, logits, sampling_metadata): + return self.model.sample(logits, sampling_metadata) + + +class MingTTSForConditionalGeneration(nn.Module, SupportsPP, CustomProcessMixin): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + del prefix + self.vllm_config = vllm_config + self.ming_config = MingTTSConfig.from_hf_config(vllm_config.model_config.hf_config) + self.ming_config.validate() + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.requires_raw_input_tokens = False + self.model_stage = vllm_config.model_config.model_stage + self._prompt_encoder = None + + if self.model_stage == "llm": + self.model = init_vllm_registered_model(vllm_config=vllm_config, architectures=["MingLLMModel"]) + self.has_preprocess = True + self.has_postprocess = True + self.set_custom_preprocess(self.preprocess) + self.set_custom_postprocess(self.postprocess) + elif self.model_stage == "audio_vae": + self.model = init_vllm_registered_model(vllm_config=vllm_config, architectures=["MingAudioVAEModel"]) + self.requires_raw_input_tokens = True + else: + raise ValueError(f"Invalid Ming model_stage={self.model_stage}") + self.make_empty_intermediate_tensors = getattr(self.model, "make_empty_intermediate_tensors", lambda: None) + + @cached_property + def sampler(self): + if hasattr(self.model, "sample"): + return _ModelSampleAdapter(self.model) + if hasattr(self.model, "sampler"): + return self.model.sampler + return Sampler() + + def embed_input_ids(self, input_ids: torch.Tensor, **kwargs: Any) -> torch.Tensor: + return self.model.embed_input_ids(input_ids=input_ids, **kwargs) + + def forward(self, *args: Any, **kwargs: Any): + return self.model(*args, **kwargs) + + def compute_logits(self, hidden_states, sampling_metadata=None): + return self.model.compute_logits(hidden_states, sampling_metadata=sampling_metadata) + + def sample(self, logits, sampling_metadata): + return self.model.sample(logits, sampling_metadata) if hasattr(self.model, "sample") else None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + weights = list(weights) + if self.model_stage == "llm": + allowed = ("model.", "linear_proj_audio.", "flowloss.", "stop_head.", "spk_head.") + llm_weights = [(k, v) for k, v in weights if k.startswith(allowed)] + if not llm_weights: + raise RuntimeError( + "Ming Stage-0 received no loadable checkpoint weights. " + "Expected prefixes: model.*, linear_proj_audio.*, flowloss.*, stop_head.*, spk_head.*" + ) + loaded = self.model.load_weights(llm_weights) + return {f"model.{name}" for name in loaded} + + audio_weights = [(k, v) for k, v in weights if k.startswith("audio.")] + if not audio_weights: + raise RuntimeError("Ming Stage-1 received no loadable checkpoint weights. Expected prefix: audio.*") + loaded = self.model.load_weights(audio_weights) + return {f"model.{name}" for name in loaded} + + def _resolve_prompt_latents(self, info_dict: dict[str, Any]): + return _resolve_prompt_latents(self, info_dict) + + def preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor | None, **info_dict: Any): + if self.model_stage != "llm": + return input_ids, input_embeds, {} + input_embeds = self.model.embed_input_ids(input_ids).clone() + return ( + self._prefill_preprocess(input_ids, input_embeds, **info_dict) + if int(input_ids.shape[0]) > 1 + else self._decode_preprocess(input_ids, input_embeds, **info_dict) + ) + + def postprocess(self, hidden_states: torch.Tensor, **info_dict: Any) -> dict[str, Any]: + if self.model_stage != "llm" or hidden_states.numel() == 0: + return {} + req_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id")) + pending = self.model.pop_postprocess_update(req_id) + if not pending or not isinstance(pending.get("ming_latent_patch"), torch.Tensor): + return {} + + update = { + KEY_LATENT_HISTORY: pending[KEY_LATENT_HISTORY].detach().to("cpu").contiguous(), + KEY_NEXT_EMBEDS: pending[KEY_NEXT_EMBEDS].detach().to("cpu").contiguous(), + KEY_DECODE_STEP: int(info_dict.get(KEY_DECODE_STEP, 0)) + 1, + } + stop_prob = _take_scalar(pending.get("ming_stop_prob"), 0) + if stop_prob is not None: + update[KEY_LAST_STOP_PROB] = stop_prob + stop_reason = pending.get(MING_STOP_REASON_KEY) + if isinstance(stop_reason, str): + update[MING_STOP_REASON_KEY] = stop_reason + if isinstance(req_id, str): + update[KEY_REQUEST_ID] = req_id + return update + + def _prefill_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: Any): + if bool(info_dict.get(KEY_TEXT_MODE, False)): + update: dict[str, Any] = {KEY_TEXT_MODE: True} + request_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id")) + if request_id is not None: + update[KEY_REQUEST_ID] = request_id + if int(input_ids.shape[0]) > 1 and int(input_ids[-1].item()) == AUDIO_START_TOKEN_ID: + return input_ids[:-1], input_embeds[:-1], update + return input_ids, input_embeds, update + + update: dict[str, Any] = {KEY_DECODE_STEP: int(info_dict.get(KEY_DECODE_STEP, 0))} + prompt_latents = _resolve_prompt_latents(self, info_dict) + history = _initial_history( + prompt_latents["frames"] if prompt_latents is not None else None, + history_size=self.ming_config.history_patch_size, + latent_dim=self.ming_config.latent_dim, + device=input_embeds.device, + dtype=torch.float32, + ) + update[KEY_LATENT_HISTORY] = history.detach().to("cpu").contiguous() + update[KEY_PROMPT_LATENT_TAIL] = update[KEY_LATENT_HISTORY] + + speaker_embedding = info_dict.get(KEY_SPEAKER_EMBEDDING, info_dict.get("speaker_embedding")) + speaker_embeddings = None + if speaker_embedding is not None: + speaker_embeddings = coerce_speaker_embeddings( + speaker_embedding, + use_zero_spk_emb=bool(info_dict.get("use_zero_spk_emb", False)), + ) + if speaker_embeddings is not None and len(speaker_embeddings) > 0: + speaker_slots = _find_speaker_placeholder_positions(input_ids, self.vllm_config.model_config.hf_config) + if len(speaker_slots) < len(speaker_embeddings): + raise RuntimeError( + "Could not locate enough speaker placeholder slots: " + f"found {len(speaker_slots)}, need {len(speaker_embeddings)}" + ) + for speaker_slot, spk in zip(speaker_slots, speaker_embeddings): + spk_proj = self.model.project_speaker_embedding( + spk.to(device=input_embeds.device, dtype=input_embeds.dtype).unsqueeze(0) + ).squeeze(0) + input_embeds[speaker_slot] = spk_proj + + if prompt_latents is not None and prompt_latents["patches"] is not None: + prompt_patches = prompt_latents["patches"].to(dtype=getattr(self.model, "fm_dtype", torch.float32)) + prompt_embeds = self.model.linear_proj_audio(prompt_patches).squeeze(1) + placeholder_pos = _find_audio_placeholder_positions(input_ids, self.ming_config) + take = min(int(placeholder_pos.numel()), int(prompt_embeds.shape[0])) + if take > 0: + input_embeds[placeholder_pos[:take]] = prompt_embeds[:take].to(dtype=input_embeds.dtype) + + request_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id")) + if request_id is not None: + update[KEY_REQUEST_ID] = request_id + _copy_runtime_controls(update, info_dict) + return input_ids, input_embeds, update + + def _decode_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: Any): + if bool(info_dict.get(KEY_TEXT_MODE, False)): + update: dict[str, Any] = {KEY_TEXT_MODE: True} + request_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id")) + if request_id is not None: + update[KEY_REQUEST_ID] = request_id + return input_ids, input_embeds, update + + update: dict[str, Any] = {KEY_DECODE_STEP: int(info_dict.get(KEY_DECODE_STEP, 0))} + history = info_dict.get(KEY_LATENT_HISTORY) + if isinstance(history, torch.Tensor): + update[KEY_LATENT_HISTORY] = history.detach().to("cpu").contiguous() + else: + zero_history = torch.zeros( + (self.ming_config.history_patch_size, self.ming_config.latent_dim), + device=input_embeds.device, + dtype=torch.float32, + ) + update[KEY_LATENT_HISTORY] = zero_history.detach().to("cpu").contiguous() + + next_embeds = info_dict.get(KEY_NEXT_EMBEDS) + if isinstance(next_embeds, torch.Tensor) and input_ids.numel() == 1: + if not torch.isfinite(next_embeds).all(): + raise RuntimeError("Non-finite next_embeds before decode preprocess write.") + input_embeds[0] = ( + next_embeds.detach() + .reshape(-1, self.ming_config.llm_hidden_size)[0] + .to( + device=input_embeds.device, + dtype=input_embeds.dtype, + ) + ) + if not torch.isfinite(input_embeds[0]).all(): + raise RuntimeError("Non-finite backbone input_embeds after decode preprocess write.") + + request_id = info_dict.get(KEY_REQUEST_ID, info_dict.get("request_id")) + if request_id is not None: + update[KEY_REQUEST_ID] = request_id + _copy_runtime_controls(update, info_dict) + return input_ids, input_embeds, update + + +def _copy_runtime_controls(update: dict[str, Any], info_dict: dict[str, Any]) -> None: + for key in (KEY_CFG, KEY_SIGMA, KEY_TEMPERATURE, KEY_MAX_DECODE_STEPS, KEY_MIN_DECODE_STEPS): + if key in info_dict: + update[key] = info_dict[key] + + +__all__ = [ + "MingTTSForConditionalGeneration", + "_ModelSampleAdapter", + "_coerce_prompt_latents", + "_find_audio_placeholder_positions", + "_find_speaker_placeholder_positions", + "_initial_history", +] diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts_audio_vae.py b/vllm_omni/model_executor/models/ming_tts/ming_tts_audio_vae.py new file mode 100644 index 000000000..8b029ffba --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/ming_tts_audio_vae.py @@ -0,0 +1,389 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/audio_tokenizer/modeling_audio_vae.py +from __future__ import annotations + +import os +import time +import warnings +from collections import OrderedDict +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .config_ming_tts import KEY_CHUNK_ID, KEY_REQUEST_ID, MingTTSConfig +from .patch_emission import MING_STOP_REASON_KEY + +logger = init_logger(__name__) + +MING_FINAL_DECODE_STEP_KEY = "ming_final_decode_step" +_STREAM_STATE_TTL_SECONDS = 15 * 60 +_MAX_STREAM_STATES = 1024 + + +class MingAudioVAEModel(nn.Module): + input_modalities = "audio" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + del prefix + self.vllm_config = vllm_config + self.ming_config = MingTTSConfig.from_hf_config(vllm_config.model_config.hf_config) + if self.ming_config.audio_tokenizer_config is None: + raise ValueError("MingAudioVAEModel requires audio_tokenizer_config") + self.ming_config.validate() + + self.audio = AudioVAE(self.ming_config.audio_tokenizer_config) + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.enable_update_additional_information = True + self.requires_raw_input_tokens = True + + self._past_key_values: dict[str, Any] = {} + self._stream_state: dict[str, tuple[Any, Any, Any]] = {} + self._patch_totals: dict[str, int] = {} + self._sample_totals: dict[str, int] = {} + self._state_access_times: OrderedDict[str, float] = OrderedDict() + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + hidden_size = int(self.ming_config.llm_hidden_size) + if input_ids is None or input_ids.numel() == 0: + return torch.empty((0, hidden_size), device=self.vllm_config.device_config.device, dtype=torch.float32) + return torch.zeros((input_ids.shape[0], hidden_size), device=input_ids.device, dtype=torch.float32) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None: + del hidden_states, sampling_metadata + return None + + def chunked_decode_streaming( + self, + latent_chunk: torch.Tensor, + *, + request_id: str, + finished: bool, + ) -> tuple[torch.Tensor, Any, Any, bool, bool]: + had_past_key_values = request_id in self._past_key_values + had_stream_state = _has_stream_state(self._stream_state.get(request_id)) + stream_state = self._stream_state.get(request_id, (None, None, None)) + past_key_values = self._past_key_values.get(request_id) + waveform_parts: list[torch.Tensor] = [] + + patch_count = int(latent_chunk.shape[0]) + for patch_idx in range(patch_count): + # [Batch, Time, Dimension] = [1, patch_size, latent_dim] + latent_patch = latent_chunk[patch_idx : patch_idx + 1] + is_last_patch = finished and patch_idx == patch_count - 1 + waveform, stream_state, past_key_values = self.audio.decode( + latent_patch, + past_key_values=past_key_values, + use_cache=True, + stream_state=stream_state, + last_chunk=is_last_patch, + ) + waveform_parts.append(waveform.reshape(-1).to(torch.float32)) + + waveform_flat = ( + torch.cat(waveform_parts, dim=0) + if waveform_parts + else torch.zeros((0,), dtype=torch.float32, device=latent_chunk.device) + ) + return waveform_flat, stream_state, past_key_values, had_past_key_values, had_stream_state + + @torch.inference_mode() + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + model_intermediate_buffer: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + del input_ids, positions, intermediate_tensors, inputs_embeds + + info_list = model_intermediate_buffer if isinstance(model_intermediate_buffer, list) else None + if info_list is None: + runtime_infos = kwargs.get("runtime_additional_information") + info_list = runtime_infos if isinstance(runtime_infos, list) else None + info_list = info_list or [{}] + num_reqs = max(len(info_list), 1) + sr_tensor = torch.tensor(int(self.ming_config.sample_rate), dtype=torch.int32) + empty = torch.zeros((0,), dtype=torch.float32, device=self.vllm_config.device_config.device) + + outputs: list[torch.Tensor] = [] + srs: list[torch.Tensor] = [] + + for idx in range(num_reqs): + info = info_list[idx] if idx < len(info_list) and isinstance(info_list[idx], dict) else {} + meta = info.get("meta") or {} + kv_meta = info.get("kv_metadata") or {} + has_ming_context = _has_ming_context(info) + if has_ming_context and KEY_REQUEST_ID not in info and KEY_REQUEST_ID not in kv_meta: + raise RuntimeError( + f"Ming Stage-2 received a payload without {KEY_REQUEST_ID}. keys={sorted(info.keys())}" + ) + request_id = _resolve_request_id(info, idx) + self._touch_request_state(request_id) + chunk_id = _coerce_optional_int(_first_present(info.get(KEY_CHUNK_ID), kv_meta.get(KEY_CHUNK_ID))) + finished = _coerce_finished( + _first_present( + meta.get("stream_finished"), + meta.get("finished"), + info.get("stream_finished"), + kv_meta.get("stream_finished"), + torch.tensor(True), + ) + ) + latent = _first_present( + info.get("latent"), + info.get("ming_latent_patches"), + kv_meta.get("ming_latent_patches"), + ) + stripped = bool( + _first_present( + meta.get("_ming_payload_stripped"), + info.get("_ming_payload_stripped"), + kv_meta.get("_ming_payload_stripped"), + False, + ) + ) + if stripped: + raise RuntimeError( + "Ming Stage-2 payload was stripped before model entry. " + f"request_id={request_id} chunk_id={chunk_id} keys={sorted(info.keys())}" + ) + if latent is None: + if has_ming_context and not finished: + raise RuntimeError( + "Ming Stage-2 received no latent chunk for an unfinished request. " + f"request_id={request_id} chunk_id={chunk_id} keys={sorted(info.keys())}" + ) + if finished: + self._clear_request_state(request_id) + outputs.append(empty) + srs.append(sr_tensor) + continue + + latent_tensor = _coerce_latent_chunk( + latent, + device=self.vllm_config.device_config.device, + dtype=next(self.audio.parameters()).dtype, + latent_dim=self.ming_config.latent_dim, + patch_size=self.ming_config.patch_size, + ) + if latent_tensor is None or latent_tensor.numel() == 0: + if not finished: + raise RuntimeError( + "Ming Stage-2 received an empty latent chunk before final flush. " + f"request_id={request_id} chunk_id={chunk_id} latent_shape={_shape_of(latent_tensor)}" + ) + if finished: + self._clear_request_state(request_id) + outputs.append(empty) + srs.append(sr_tensor) + continue + + patch_count = int(latent_tensor.shape[0]) + waveform_flat, stream_state, past_key_values, had_past_key_values, had_stream_state = ( + self.chunked_decode_streaming( + latent_tensor, + request_id=request_id, + finished=finished, + ) + ) + total_patch_count = self._patch_totals.get(request_id, 0) + patch_count + total_waveform_numel = self._sample_totals.get(request_id, 0) + int(waveform_flat.numel()) + self._patch_totals[request_id] = total_patch_count + self._sample_totals[request_id] = total_waveform_numel + if (had_past_key_values or had_stream_state) and waveform_flat.numel() == 0 and not finished: + raise RuntimeError( + "Ming Stage-2 produced an empty waveform after cached streaming state already existed. " + f"request_id={request_id} chunk_id={chunk_id} latent_shape={tuple(latent_tensor.shape)} " + f"had_past_key_values={had_past_key_values} had_stream_state={had_stream_state}" + ) + + if finished: + if os.environ.get("MING_TTS_STAGE1_FINAL_LOG") == "1": + logger.info( + "MING_STAGE1_FINAL %s", + { + "request_id": request_id, + "chunk_id": chunk_id, + "stop_reason": _first_present( + info.get(MING_STOP_REASON_KEY), + kv_meta.get(MING_STOP_REASON_KEY), + ), + "final_decode_step": _coerce_optional_int( + _first_present( + info.get(MING_FINAL_DECODE_STEP_KEY), + kv_meta.get(MING_FINAL_DECODE_STEP_KEY), + ) + ), + "final_chunk_patch_count": patch_count, + "total_patch_count": total_patch_count, + "final_chunk_waveform_numel": int(waveform_flat.numel()), + "total_waveform_numel": total_waveform_numel, + }, + ) + self._clear_request_state(request_id) + else: + self._past_key_values[request_id] = past_key_values + self._stream_state[request_id] = stream_state + + outputs.append(waveform_flat) + srs.append(sr_tensor) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "model_outputs": outputs, + "sr": srs, + }, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + weights = list(weights) + if not weights: + raise RuntimeError("MingAudioVAEModel received no checkpoint weights.") + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded: set[str] = set() + skipped: list[str] = [] + + for name, loaded_weight in weights: + if name not in params_dict: + skipped.append(name) + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded.add(name) + + missing = {name for name in params_dict if name.startswith("audio.")} - loaded + if missing: + raise RuntimeError(f"MingAudioVAEModel: {len(missing)} params not loaded. First few: {sorted(missing)[:5]}") + if skipped: + warnings.warn( + f"MingAudioVAEModel: skipped {len(skipped)} checkpoint keys during load. First few: {skipped[:8]}", + stacklevel=2, + ) + return loaded + + def _clear_request_state(self, request_id: str) -> None: + self._past_key_values.pop(request_id, None) + self._stream_state.pop(request_id, None) + self._patch_totals.pop(request_id, None) + self._sample_totals.pop(request_id, None) + self._state_access_times.pop(request_id, None) + + def _touch_request_state(self, request_id: str) -> None: + now = time.monotonic() + self._evict_expired_request_states(now) + self._state_access_times[request_id] = now + self._state_access_times.move_to_end(request_id) + while len(self._state_access_times) > _MAX_STREAM_STATES: + evicted_request_id, _ = self._state_access_times.popitem(last=False) + self._clear_request_state(evicted_request_id) + + def _evict_expired_request_states(self, now: float) -> None: + cutoff = now - _STREAM_STATE_TTL_SECONDS + while self._state_access_times: + request_id, last_access = next(iter(self._state_access_times.items())) + if last_access >= cutoff: + break + self._clear_request_state(request_id) + + +def _coerce_finished(value: Any) -> bool: + if isinstance(value, torch.Tensor): + if value.numel() == 0: + return False + return bool(value.reshape(-1)[0].item()) + return bool(value) + + +def _first_present(*values: Any) -> Any: + for value in values: + if value is not None: + return value + return None + + +def _coerce_optional_int(value: Any, default: int | None = None) -> int | None: + if value is None: + return default + if isinstance(value, list): + value = value[0] if value else default + if isinstance(value, torch.Tensor): + if value.numel() == 0: + return default + return int(value.reshape(-1)[0].item()) + try: + return int(value) + except (TypeError, ValueError): + return default + + +def _resolve_request_id(info: dict[str, Any], idx: int) -> str: + kv_meta = info.get("kv_metadata") or {} + request_id = info.get(KEY_REQUEST_ID) or kv_meta.get(KEY_REQUEST_ID) + if request_id is None: + return str(idx) + if not isinstance(request_id, str) or not request_id: + raise RuntimeError(f"Ming Stage-2 received invalid request id: {request_id!r}") + return request_id + + +def _has_ming_context(info: dict[str, Any]) -> bool: + kv_meta = info.get("kv_metadata") or {} + keys = ( + KEY_REQUEST_ID, + KEY_CHUNK_ID, + "ming_latent_patches", + "_ming_payload_stripped", + ) + return any(key in info or key in kv_meta for key in keys) or "latent" in info + + +def _shape_of(value: Any) -> tuple[int, ...] | None: + if isinstance(value, torch.Tensor): + return tuple(value.shape) + return None + + +def _has_stream_state(value: Any) -> bool: + if not isinstance(value, tuple): + return False + return any(item is not None for item in value) + + +def _coerce_latent_chunk( + value: Any, + *, + device: torch.device, + dtype: torch.dtype, + latent_dim: int, + patch_size: int, +) -> torch.Tensor | None: + if not isinstance(value, torch.Tensor): + value = torch.as_tensor(value) + + latent = value.detach() + if latent.ndim == 2: + latent = latent.unsqueeze(0) + if latent.ndim != 3: + raise ValueError(f"Expected latent chunk rank-3 [Batch, Time, Dimension], got {tuple(latent.shape)}") + if latent.shape[-2] != patch_size: + raise ValueError(f"Latent patch size mismatch: got {latent.shape[-2]}, expected {patch_size}") + if latent.shape[-1] != latent_dim: + raise ValueError(f"Latent dim mismatch: got {latent.shape[-1]}, expected {latent_dim}") + return latent.to(device=device, dtype=dtype) diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py new file mode 100644 index 000000000..92daad967 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/modeling_bailingmm.py +from __future__ import annotations + +import warnings +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.qwen2 import Qwen2Model +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper, maybe_prefix +from vllm.sequence import IntermediateTensors +from vllm.v1.outputs import SamplerOutput +from vllm.v1.sample.metadata import SamplingMetadata + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .aggregator import Aggregator +from .config_ming_tts import ( + KEY_CFG, + KEY_DECODE_STEP, + KEY_LATENT_HISTORY, + KEY_MAX_DECODE_STEPS, + KEY_MIN_DECODE_STEPS, + KEY_NEXT_EMBEDS, + KEY_REQUEST_ID, + KEY_SIGMA, + KEY_TEMPERATURE, + KEY_TEXT_MODE, + MingTTSConfig, +) +from .flowloss_head import FlowLoss +from .patch_emission import ( + MING_STOP_REASON_CODES, + MING_STOP_REASON_CONTINUE, + MING_STOP_REASON_KEY, + MING_STOP_REASON_MAX_DECODE_STEPS, + MING_STOP_REASON_STOP_HEAD, + _coerce_latent_history, + _get_request_token_counts, + _normalize_request_infos, + _resolve_ming_stop_decision, + _resolve_optional_runtime_int, + _resolve_runtime_float, + _resolve_runtime_int, + _validate_ming_decode_window, +) + +logger = init_logger(__name__) + + +class MingLLMModel(nn.Module): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.model.": "model.", + } + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.ming_config = MingTTSConfig.from_hf_config(vllm_config.model_config.hf_config) + self.ming_config.validate() + self.vllm_config = vllm_config + self.prefix = prefix + self.fm_dtype = _resolve_ming_runtime_dtype(vllm_config) + self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + self.linear_proj_audio = Aggregator( + in_channels=self.ming_config.latent_dim, + llm_input_dim=self.ming_config.llm_hidden_size, + **self.ming_config.aggregator_config, + ) + self.flowloss = FlowLoss( + z_channels=self.ming_config.latent_dim, + llm_cond_dim=self.ming_config.llm_hidden_size, + **self.ming_config.ditar_config, + ) + self.stop_head = nn.Linear(self.ming_config.llm_hidden_size, 2, bias=True) + self.spk_head = nn.Linear(192, self.ming_config.llm_hidden_size, bias=True) + self.flowloss.to(dtype=self.fm_dtype) + self.linear_proj_audio.to(dtype=self.fm_dtype) + self.stop_head.to(dtype=self.fm_dtype) + self.spk_head.to(dtype=self.fm_dtype) + self._pending_postprocess_updates: dict[str, dict[str, Any]] = {} + self._last_ming_next_token_ids = None + self._last_text_mode = False + + def embed_input_ids( + self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor | None = None, **_: Any + ) -> torch.Tensor: + if inputs_embeds is not None: + return inputs_embeds + return self.model.embed_input_ids(input_ids) + + def project_speaker_embedding(self, spk_emb: torch.Tensor) -> torch.Tensor: + return self.spk_head(spk_emb) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + latent_history: torch.Tensor | None = None, + model_intermediate_buffer: list[dict[str, Any]] | None = None, + seq_token_counts: list[int] | None = None, + **kwargs: object, + ) -> OmniOutput | IntermediateTensors | torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + if model_intermediate_buffer is None: + model_intermediate_buffer = kwargs.get("runtime_additional_information") + request_infos = _normalize_request_infos(model_intermediate_buffer) + _validate_ming_decode_window( + request_infos, + min_stop_step=int(self.ming_config.stop_head_min_steps), + default_max_decode_steps=self.ming_config.max_decode_steps, + ) + backbone_out = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + if isinstance(backbone_out, IntermediateTensors): + self._last_ming_next_token_ids = None + return backbone_out + hidden_states = _extract_hidden_states(backbone_out) + token_counts = _get_request_token_counts(hidden_states, request_infos, seq_token_counts) + text_mode = bool(request_infos) and all(bool(info.get(KEY_TEXT_MODE, False)) for info in request_infos) + if request_infos and any(bool(info.get(KEY_TEXT_MODE, False)) for info in request_infos) and not text_mode: + raise RuntimeError("Mixed Ming text/audio modes in one Stage-0 batch are unsupported.") + if text_mode: + self._last_ming_next_token_ids = None + self._last_text_mode = True + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs={KEY_TEXT_MODE: True}, + intermediate_tensors=intermediate_tensors, + ) + self._last_text_mode = False + if latent_history is None and not token_counts: + self._last_ming_next_token_ids = None + return OmniOutput( + text_hidden_states=hidden_states, multimodal_outputs=None, intermediate_tensors=intermediate_tensors + ) + if latent_history is not None and not token_counts: + token_counts, request_infos = [hidden_states.shape[0]], [{KEY_LATENT_HISTORY: latent_history}] + + total_tokens = hidden_states.shape[0] + latent_patch_tokens = next_embed_tokens = new_history_tokens = None + decode_step_tokens = has_patch = stop_reason_code_tokens = None + next_token_ids = [] + pending_updates: dict[str, dict[str, Any]] = {} + cursor = 0 + any_decode = False + for req_idx, token_count in enumerate(token_counts): + end = min(cursor + token_count, total_tokens) + if end <= cursor: + continue + req_info = request_infos[req_idx] if req_idx < len(request_infos) else {} + req_id = req_info.get(KEY_REQUEST_ID) + req_history = _coerce_latent_history( + req_info.get(KEY_LATENT_HISTORY), device=hidden_states.device, dtype=self.fm_dtype, cfg=self.ming_config + ) + if req_history is None: + cursor = end + continue + decode_step = int(req_info.get(KEY_DECODE_STEP, req_info.get("generated_len", 0))) + decode_hidden, output_index = ( + (hidden_states[cursor:end], cursor) if token_count == 1 else (hidden_states[end - 1 : end], end - 1) + ) + sampled_token_latent, next_embeds, new_history, stop_probs = self._decode_one_step( + hidden_states=decode_hidden, + latent_history=req_history, + cfg_scale=_resolve_runtime_float(req_info, KEY_CFG, self.ming_config.cfg), + sigma=_resolve_runtime_float(req_info, KEY_SIGMA, self.ming_config.sigma), + temperature=_resolve_runtime_float(req_info, KEY_TEMPERATURE, self.ming_config.temperature), + ) + req_max_decode_steps = _resolve_runtime_int( + req_info, KEY_MAX_DECODE_STEPS, self.ming_config.max_decode_steps + ) + req_min_decode_steps = _resolve_optional_runtime_int(req_info, KEY_MIN_DECODE_STEPS, 0) + if latent_patch_tokens is None: + latent_patch_tokens = sampled_token_latent.new_zeros( + (total_tokens, self.ming_config.patch_size, self.ming_config.latent_dim) + ) + next_embed_tokens = next_embeds.new_zeros((total_tokens, 1, self.ming_config.llm_hidden_size)) + new_history_tokens = new_history.new_zeros( + (total_tokens, self.ming_config.history_patch_size, self.ming_config.latent_dim) + ) + decode_step_tokens = torch.zeros((total_tokens,), dtype=torch.int32, device=hidden_states.device) + has_patch = torch.zeros((total_tokens,), dtype=torch.bool, device=hidden_states.device) + stop_reason_code_tokens = torch.zeros((total_tokens,), dtype=torch.int32, device=hidden_states.device) + latent_patch_tokens[output_index : output_index + 1] = sampled_token_latent + next_embed_tokens[output_index : output_index + 1] = next_embeds + new_history_tokens[output_index : output_index + 1] = new_history + decode_step_tokens[output_index : output_index + 1] = decode_step + has_patch[output_index : output_index + 1] = True + stop_reason, _, _, _, next_token_id = _resolve_ming_stop_decision( + step=decode_step, + stop_prob=float(stop_probs.reshape(-1)[0].item()), + stop_threshold=float(self.ming_config.stop_head_threshold), + min_stop_step=int(self.ming_config.stop_head_min_steps), + min_decode_steps=req_min_decode_steps, + max_decode_steps=req_max_decode_steps, + audio_dummy_token_id=int(self.ming_config.audio_dummy_token_id), + text_eos_token_id=int(self.ming_config.text_eos_token_id), + ) + stop_reason_code_tokens[output_index : output_index + 1] = MING_STOP_REASON_CODES[stop_reason] + next_token_ids.append(int(next_token_id)) + if isinstance(req_id, str): + pending_updates[req_id] = { + KEY_LATENT_HISTORY: new_history, + KEY_NEXT_EMBEDS: next_embeds, + "ming_latent_patch": sampled_token_latent, + "ming_stop_prob": stop_probs, + MING_STOP_REASON_KEY: stop_reason, + } + any_decode = True + cursor = end + + self._pending_postprocess_updates = pending_updates + if not any_decode: + self._last_ming_next_token_ids = None + return OmniOutput( + text_hidden_states=hidden_states, multimodal_outputs=None, intermediate_tensors=intermediate_tensors + ) + self._last_ming_next_token_ids = next_token_ids + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs={ + "ming_latent_patch": latent_patch_tokens, + "ming_next_embeds": next_embed_tokens, + "ming_new_history": new_history_tokens, + "ming_decode_step": decode_step_tokens, + "ming_has_patch": has_patch, + MING_STOP_REASON_KEY: stop_reason_code_tokens, + }, + intermediate_tensors=intermediate_tensors, + ) + + def pop_postprocess_update(self, req_id: str) -> dict[str, Any]: + return self._pending_postprocess_updates.pop(req_id, {}) if isinstance(req_id, str) else {} + + def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: SamplingMetadata) -> torch.Tensor | None: + del sampling_metadata + if self._last_text_mode: + return ( + None + if hidden_states is None or hidden_states.numel() == 0 + else self.model.compute_logits(hidden_states) + ) + if hidden_states is None or hidden_states.numel() == 0: + return None + if hidden_states.dim() != 2: + raise RuntimeError( + f"Expected hidden_states rank-2 [B,H] in compute_logits, got {tuple(hidden_states.shape)}" + ) + batch_size = hidden_states.shape[0] + next_token_ids = self._last_ming_next_token_ids + self._last_ming_next_token_ids = None + if next_token_ids is None: + logger.debug( + "Missing Ming forced next-token ids before compute_logits. " + "Using dummy next-token IDs (likely during a profiling or dummy run)." + ) + next_token_ids = [int(self.ming_config.text_eos_token_id)] * batch_size + if len(next_token_ids) != batch_size: + raise RuntimeError( + "Ming forced next-token batch mismatch: " + f"got {len(next_token_ids)} ids for hidden_states batch {batch_size}." + ) + logits = torch.full( + (batch_size, self.ming_config.llm_vocab_size), + float("-inf"), + device=hidden_states.device, + dtype=torch.float32, + ) + for i, next_token_id in enumerate(next_token_ids): + logits[i, int(next_token_id)] = 0.0 + return logits + + def sample(self, logits, sampling_metadata): + if logits is None: + return None + if self._last_text_mode: + return self.model.sample(logits, sampling_metadata) + del sampling_metadata + return SamplerOutput( + sampled_token_ids=logits.argmax(dim=-1, keepdim=True).to(dtype=torch.int32), + logprobs_tensors=None, + ) + + def _decode_one_step( + self, + *, + hidden_states: torch.Tensor, + latent_history: torch.Tensor, + cfg_scale: float, + sigma: float, + temperature: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if hidden_states.dim() != 2: + raise RuntimeError(f"Expected decode hidden_states rank-2 [B,H], got {tuple(hidden_states.shape)}") + if latent_history.dim() != 3: + raise RuntimeError(f"Expected latent_history rank-3 [B,T,D], got {tuple(latent_history.shape)}") + if hidden_states.shape[0] != latent_history.shape[0]: + raise RuntimeError( + "Batch mismatch: " + f"hidden_states B={hidden_states.shape[0]} " + f"vs latent_history B={latent_history.shape[0]}" + ) + # [Batch, Hidden] -> [Batch, Time, Hidden] = [B, 1, H] for FlowLoss conditioning. + z_diff_cond = hidden_states.to(dtype=self.fm_dtype).unsqueeze(1) + if not torch.isfinite(z_diff_cond).all(): + raise RuntimeError("Non-finite z_diff_cond before FlowLoss.sample().") + sampled_token_latent = self.flowloss.sample( + z=z_diff_cond, + latent_history=latent_history, + cfg=cfg_scale, + patch_size=self.ming_config.patch_size, + sigma=sigma, + temperature=temperature, + ) + expected_shape = (hidden_states.shape[0], self.ming_config.patch_size, self.ming_config.latent_dim) + if tuple(sampled_token_latent.shape) != expected_shape: + raise RuntimeError( + f"FlowLoss output shape mismatch: got {tuple(sampled_token_latent.shape)}, expected {expected_shape}" + ) + new_history = torch.cat([latent_history[:, self.ming_config.patch_size :, :], sampled_token_latent], dim=1) + # Aggregator expects [Batch, Time, Dimension] = [B, 4, 64] and returns [B, 1, H]. + next_embeds = self.linear_proj_audio(sampled_token_latent) + stop_probs = self.stop_head(hidden_states.to(dtype=self.fm_dtype)).softmax(dim=-1)[:, 1] + if not torch.isfinite(sampled_token_latent).all(): + raise RuntimeError("Non-finite sampled_token_latent in Ming decode step.") + if not torch.isfinite(next_embeds).all(): + raise RuntimeError("Non-finite next_embeds in Ming decode step.") + if not torch.isfinite(stop_probs).all(): + raise RuntimeError("Non-finite stop_probs in Ming decode step.") + return sampled_token_latent, next_embeds, new_history, stop_probs + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params = AutoWeightsLoader(self).load_weights(weights, mapper=self.hf_to_vllm_mapper) + _warn_missing_prefix("flowloss", params_dict, loaded_params, prefix="flowloss.", fatal=True) + _warn_missing_prefix("linear_proj_audio", params_dict, loaded_params, prefix="linear_proj_audio.", fatal=True) + _warn_missing_prefix("stop_head", params_dict, loaded_params, prefix="stop_head.", fatal=True) + _warn_missing_prefix("spk_head", params_dict, loaded_params, prefix="spk_head.", fatal=True) + return loaded_params + + +def _extract_hidden_states(backbone_out: object) -> torch.Tensor: + if isinstance(backbone_out, torch.Tensor): + return backbone_out + if hasattr(backbone_out, "last_hidden_state"): + return backbone_out.last_hidden_state + if isinstance(backbone_out, (tuple, list)) and backbone_out and isinstance(backbone_out[0], torch.Tensor): + return backbone_out[0] + raise TypeError(f"Unsupported backbone forward output type: {type(backbone_out)}") + + +def _resolve_ming_runtime_dtype(vllm_config: VllmConfig) -> torch.dtype: + dtype = getattr(vllm_config.model_config, "dtype", None) + if isinstance(dtype, torch.dtype): + return dtype + if isinstance(dtype, str): + normalized = dtype.strip().lower() + if normalized in ("float16", "half", "torch.float16"): + return torch.float16 + if normalized in ("bfloat16", "bf16", "torch.bfloat16"): + return torch.bfloat16 + if normalized in ("float32", "fp32", "torch.float32"): + return torch.float32 + return torch.float32 + + +def _warn_missing_prefix( + module_name: str, + params_dict: dict[str, nn.Parameter], + loaded_params: set[str], + prefix: str, + fatal: bool = False, +) -> None: + missing = {key for key in params_dict if key.startswith(prefix)} - loaded_params + if not missing: + return + msg = ( + f"MingLLMModel: {len(missing)} {module_name} params not loaded " + f"(prefix={prefix}). First few: {sorted(missing)[:5]}" + ) + if fatal: + raise RuntimeError(msg) + warnings.warn(msg, stacklevel=3) + + +__all__ = [ + "Aggregator", + "FlowLoss", + "MING_STOP_REASON_CODES", + "MING_STOP_REASON_CONTINUE", + "MING_STOP_REASON_KEY", + "MING_STOP_REASON_MAX_DECODE_STEPS", + "MING_STOP_REASON_STOP_HEAD", + "MingLLMModel", + "_coerce_latent_history", + "_extract_hidden_states", + "_resolve_ming_runtime_dtype", + "_resolve_ming_stop_decision", + "_warn_missing_prefix", +] diff --git a/vllm_omni/model_executor/models/ming_tts/patch_emission.py b/vllm_omni/model_executor/models/ming_tts/patch_emission.py new file mode 100644 index 000000000..c0a25a1c1 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/patch_emission.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/modeling_bailingmm.py +from __future__ import annotations + +from typing import Any + +import torch +from vllm.forward_context import get_forward_context, is_forward_context_available + +from .config_ming_tts import KEY_MAX_DECODE_STEPS, KEY_MIN_DECODE_STEPS, KEY_REQUEST_ID, KEY_TEXT_MODE, MingTTSConfig + +MING_STOP_REASON_CONTINUE = "continue" +MING_STOP_REASON_STOP_HEAD = "stop_head" +MING_STOP_REASON_MAX_DECODE_STEPS = "max_decode_steps" +MING_STOP_REASON_KEY = "ming_stop_reason" +MING_STOP_REASON_CODES = { + MING_STOP_REASON_CONTINUE: 0, + MING_STOP_REASON_STOP_HEAD: 1, + MING_STOP_REASON_MAX_DECODE_STEPS: 2, +} + + +def _normalize_request_infos(model_intermediate_buffer: object) -> list[dict[str, Any]]: + if not isinstance(model_intermediate_buffer, list): + return [] + infos: list[dict[str, Any]] = [] + for item in model_intermediate_buffer: + infos.append(item if isinstance(item, dict) else {}) + return infos + + +def _get_request_token_counts( + hidden_states: torch.Tensor, + request_infos: list[dict[str, Any]], + seq_token_counts: list[int] | None, +) -> list[int]: + if seq_token_counts: + return [int(x) for x in seq_token_counts] + + if is_forward_context_available(): + slices = getattr(get_forward_context(), "ubatch_slices", None) + if slices is not None and len(slices) > 0: + counts: list[int] = [] + for item in slices: + if isinstance(item, int): + counts.append(int(item)) + elif hasattr(item, "stop") and hasattr(item, "start"): + counts.append(int(item.stop) - int(item.start)) + if counts: + return counts + + if request_infos: + if len(request_infos) == hidden_states.shape[0]: + return [1] * hidden_states.shape[0] + return [hidden_states.shape[0]] + + return [] + + +def _coerce_latent_history( + value: object, + *, + device: torch.device, + dtype: torch.dtype, + cfg: MingTTSConfig, +) -> torch.Tensor | None: + if value is None: + return None + if not isinstance(value, torch.Tensor): + value = torch.as_tensor(value) + + history = value.detach() + if history.ndim == 2: + history = history.unsqueeze(0) + if history.ndim != 3: + raise RuntimeError(f"Expected latent_history rank-3 [B,T,D], got {tuple(history.shape)}") + if history.shape[1] != cfg.history_patch_size or history.shape[2] != cfg.latent_dim: + raise RuntimeError( + f"latent_history shape mismatch: got {tuple(history.shape)}, " + f"expected [B,{cfg.history_patch_size},{cfg.latent_dim}]" + ) + return history.to(device=device, dtype=dtype) + + +def _resolve_runtime_float(req_info: dict[str, Any], key: str, default_value: float) -> float: + raw = req_info.get(key, default_value) + try: + value = float(raw) + except (TypeError, ValueError) as exc: + raise RuntimeError(f"Invalid {key}: expected float-like value, got {raw!r}") from exc + if not value >= 0.0: + raise RuntimeError(f"Invalid {key}: expected non-negative value, got {value}") + return value + + +def _resolve_runtime_int(req_info: dict[str, Any], key: str, default_value: int) -> int: + raw = req_info.get(key, default_value) + try: + value = int(raw) + except (TypeError, ValueError) as exc: + raise RuntimeError(f"Invalid {key}: expected int-like value, got {raw!r}") from exc + if value <= 0: + raise RuntimeError(f"Invalid {key}: expected positive value, got {value}") + return value + + +def _resolve_optional_runtime_int(req_info: dict[str, Any], key: str, default_value: int) -> int: + raw = req_info.get(key, default_value) + try: + value = int(raw) + except (TypeError, ValueError) as exc: + raise RuntimeError(f"Invalid {key}: expected int-like value, got {raw!r}") from exc + if value < 0: + raise RuntimeError(f"Invalid {key}: expected non-negative value, got {value}") + return value + + +def _validate_ming_decode_window( + request_infos: list[dict[str, Any]], + *, + min_stop_step: int, + default_max_decode_steps: int, +) -> None: + for i, info in enumerate(request_infos): + if info.get(KEY_TEXT_MODE): + continue + max_steps = _resolve_runtime_int(info, KEY_MAX_DECODE_STEPS, default_max_decode_steps) + min_steps = _resolve_optional_runtime_int(info, KEY_MIN_DECODE_STEPS, 0) + min_required = max(min_stop_step + 1, min_steps) + if max_steps < min_required: + req_id = info.get(KEY_REQUEST_ID, f"idx={i}") + raise ValueError( + f"Ming request {req_id!r}: max_decode_steps={max_steps} < " + f"min_required_decode_steps={min_required} " + f"(min_stop_step={min_stop_step}, min_decode_steps={min_steps})" + ) + + +def _resolve_ming_stop_decision( + *, + step: int, + stop_prob: float, + stop_threshold: float, + min_stop_step: int, + min_decode_steps: int, + max_decode_steps: int, + audio_dummy_token_id: int, + text_eos_token_id: int, +) -> tuple[str, bool, bool, int, int]: + min_required_decode_steps = max(min_stop_step + 1, min_decode_steps) + if max_decode_steps < min_required_decode_steps: + raise RuntimeError( + "Invalid Ming decode window: " + f"max_decode_steps={max_decode_steps} is smaller than " + f"min_required_decode_steps={min_required_decode_steps}" + ) + should_force_stop = (step + 1) >= max_decode_steps + should_stop_head = ((step + 1) >= min_required_decode_steps) and stop_prob > stop_threshold + + if should_force_stop: + return ( + MING_STOP_REASON_MAX_DECODE_STEPS, + True, + True, + min_required_decode_steps, + text_eos_token_id, + ) + if should_stop_head: + return ( + MING_STOP_REASON_STOP_HEAD, + True, + False, + min_required_decode_steps, + text_eos_token_id, + ) + return ( + MING_STOP_REASON_CONTINUE, + False, + False, + min_required_decode_steps, + audio_dummy_token_id, + ) diff --git a/vllm_omni/model_executor/models/ming_tts/pipeline.py b/vllm_omni/model_executor/models/ming_tts/pipeline.py new file mode 100644 index 000000000..31ec98e5a --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/pipeline.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ming TTS pipeline: Stage-0 LLM+flow -> Stage-1 audio VAE.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.ming_tts" + +MING_TTS_PIPELINE = PipelineConfig( + model_type="ming_tts", + model_arch="MingTTSForConditionalGeneration", + hf_architectures=("MingTTSForConditionalGeneration",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="llm", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + hf_config_name="llm_config", + engine_output_type="latent", + async_chunk_process_next_stage_input_func=(f"{_PROC}.llm2audio_vae_async_chunk"), + sampling_constraints={ + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 512, + "detokenize": True, + }, + ), + StagePipelineConfig( + stage_id=1, + model_stage="audio_vae", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + hf_config_name="llm_config", + engine_output_type="audio", + sync_process_input_func=f"{_PROC}.llm2audio_vae", + sampling_constraints={ + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 1, + "detokenize": False, + }, + ), + ), +) diff --git a/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py b/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py new file mode 100644 index 000000000..c732e07a0 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +import copy +import json +import math +import re +from typing import Any + +import torch + +from .audio_prep import ( + coerce_speaker_embeddings, + count_prompt_latent_patches, + count_prompt_waveform_patches, + pad_prompt_waveform, +) +from .config_ming_tts import ( + AUDIO_FRAME_HOP, + KEY_MAX_DECODE_STEPS, + KEY_MIN_DECODE_STEPS, + KEY_PROMPT_LATENTS, + KEY_REQUEST_ID, + KEY_SPEAKER_EMBEDDING, + LATENT_DIM, + PATCH_SIZE, + SAMPLE_RATE, + VAE_PATCH_SIZE, +) + +DEFAULT_PROMPT = "Please generate speech based on the following description.\n" +BASE_CAPTION_TEMPLATE: dict[str, Any] = { + "audio_sequence": [ + { + "序号": 1, + "说话人": "speaker_1", + "方言": None, + "风格": None, + "语速": None, + "基频": None, + "音量": None, + "情感": None, + "BGM": { + "Genre": None, + "Mood": None, + "Instrument": None, + "Theme": None, + "ENV": None, + "SNR": None, + }, + "IP": None, + } + ] +} +_DURATION_SECONDS_RE = re.compile(r"Duration:\s*([0-9]+(?:\.[0-9]+)?)\s*s\b", re.IGNORECASE) + + +def create_instruction(user_input: Any) -> str | None: + if user_input is None: + return None + if isinstance(user_input, str): + return user_input + if not isinstance(user_input, dict): + raise ValueError(f"Ming instruction must be str, dict, or None; got {type(user_input).__name__}") + caption = copy.deepcopy(BASE_CAPTION_TEMPLATE) + item = caption["audio_sequence"][0] + for key, value in user_input.items(): + if key in item: + item[key] = value + return json.dumps(caption, ensure_ascii=False) + + +def parse_duration_seconds(text: str | None) -> float | None: + if not isinstance(text, str): + return None + match = _DURATION_SECONDS_RE.search(text) + if match is None: + return None + try: + value = float(match.group(1)) + except ValueError: + return None + if value <= 0.0: + return None + return value + + +def estimate_decode_steps_for_duration( + duration_seconds: float, + *, + sample_rate: int = SAMPLE_RATE, + frame_hop: int = AUDIO_FRAME_HOP, + patch_size: int = PATCH_SIZE, + vae_patch_size: int = VAE_PATCH_SIZE, +) -> int: + if duration_seconds <= 0.0: + return 0 + samples_per_decode_step = int(frame_hop) * int(patch_size) * int(vae_patch_size) + required_samples = float(duration_seconds) * float(sample_rate) + return max(1, int(math.ceil(required_samples / float(samples_per_decode_step)))) + + +def estimate_decode_step_window_for_duration(duration_seconds: float) -> tuple[int, int]: + target_steps = estimate_decode_steps_for_duration(duration_seconds) + min_steps = max(1, target_steps - 3) + max_steps = max(min_steps, target_steps + 3) + return min_steps, max_steps + + +def resolve_effective_runtime_controls( + *, + text: str, + runtime_controls: dict[str, Any] | None = None, +) -> dict[str, Any]: + controls = {} if runtime_controls is None else dict(runtime_controls) + has_explicit_min = KEY_MIN_DECODE_STEPS in controls and controls[KEY_MIN_DECODE_STEPS] is not None + has_explicit_max = KEY_MAX_DECODE_STEPS in controls and controls[KEY_MAX_DECODE_STEPS] is not None + if has_explicit_min or has_explicit_max: + return controls + duration_seconds = parse_duration_seconds(text) + if duration_seconds is None: + return controls + min_decode_steps, max_decode_steps = estimate_decode_step_window_for_duration(duration_seconds) + controls[KEY_MIN_DECODE_STEPS] = min_decode_steps + controls[KEY_MAX_DECODE_STEPS] = max_decode_steps + return controls + + +def build_dense_prompt_token_ids( + tokenizer: Any, + *, + prompt: str, + text: str, + instruction: str | None = None, + prompt_text: str | None = None, + speaker_count: int = 0, + prompt_patch_count: int = 0, +) -> list[int]: + speaker_prompt = [] + for idx in range(int(speaker_count)): + speaker_prompt.extend( + tokenizer.encode(f" speaker_{idx + 1}:") + + tokenizer.encode("<|vision_start|>") + + tokenizer.encode("<|vision_pad|>") + + tokenizer.encode("<|vision_end|>\n") + ) + instruction_prompt = ( + tokenizer.encode(instruction) + tokenizer.encode("<|endoftext|>") if instruction is not None else [] + ) + prompt_text_tokens = ( + tokenizer.encode(prompt_text) if int(prompt_patch_count) > 0 and prompt_text is not None else [] + ) + audio_patch_token_id = tokenizer.convert_tokens_to_ids("") + if audio_patch_token_id == tokenizer.unk_token_id: + raise ValueError("Ming tokenizer is missing required token.") + prompt_latent_tokens = [audio_patch_token_id] * int(prompt_patch_count) + text_input_prefix = ( + [] + if all(token in text for token in ("Genre: ", "Mood: ", "Instrument: ", "Theme: ", "Duration: ")) + else tokenizer.encode(" Text input:\n") + ) + return ( + tokenizer.encode("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n") + + tokenizer.encode("<|im_start|>user\n") + + tokenizer.encode(prompt) + + speaker_prompt + + text_input_prefix + + prompt_text_tokens + + tokenizer.encode(text) + + tokenizer.encode("<|im_end|>\n") + + tokenizer.encode("<|im_start|>assistant\n") + + instruction_prompt + + tokenizer.encode("