anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Model] Add Ming-omni-tts dense 0.5B pipeline (#2906)

Browse Source 
Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
akshatvishuandHongsheng Liu authored and GitHub committed 2026-06-10 20:11:04 +08:00
1 parent 39ae6e12f8
commit 2fa4c7a4a9
55 files changed
+6263 -549
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/test-merge.yml
+21
View File
@@ -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
docs/models/supported_models.md
+1
View File
@@ -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` | ✅︎ | | | |
examples/offline_inference/text_to_speech/README.md
+41
View File
@@ -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.
examples/offline_inference/text_to_speech/ming_tts/README.md
+248
View File
@@ -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_<case>.wav`
- When `--num-prompts > 1`, outputs are indexed as `ming_<case>_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.
examples/offline_inference/text_to_speech/ming_tts/cases.yaml
+110
View File
@@ -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
examples/offline_inference/text_to_speech/ming_tts/end2end.py
+200
View File
@@ -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()
examples/offline_inference/text_to_speech/ming_tts/runner.py
+277
View File
@@ -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)
examples/online_serving/text_to_speech/README.md
+46
View File
@@ -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`.
examples/online_serving/text_to_speech/ming_tts/README.md
+113
View File
@@ -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.
examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py
+225
View File
@@ -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()
examples/online_serving/text_to_speech/ming_tts/run_curl.sh
Executable
+76
View File
@@ -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
examples/online_serving/text_to_speech/ming_tts/run_server.sh
Executable
+25
View File
@@ -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
recipes/README.md
+1
View File
@@ -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 |
recipes/inclusionAI/Ming-omni-tts-0.5B.md
+82
View File
@@ -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.
tests/e2e/offline_inference/test_ming_flash_omni_expansion.py
+7 -7
View File
@@ -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:
"""
tests/e2e/offline_inference/test_ming_tts.py
+295
View File
@@ -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())
tests/e2e/online_serving/test_ming_tts.py
+92
View File
@@ -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)
tests/model_executor/models/ming_tts/test_ming_shared_modules.py
+185
View File
@@ -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)))
tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py
+10 -8
View File
@@ -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:
vllm_omni/config/pipeline_registry.py
+4
View File
@@ -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",
vllm_omni/config/stage_config.py
+13
View File
@@ -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
vllm_omni/deploy/ming_tts.yaml
+55
View File
@@ -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
vllm_omni/engine/arg_utils.py
+2
View File
@@ -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),
vllm_omni/entrypoints/openai/protocol/audio.py
+61 -5
View File
@@ -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
vllm_omni/entrypoints/openai/serving_speech.py
+320 -22
View File
@@ -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:
vllm_omni/model_executor/models/common/ming/__init__.py
+2
View File
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm_omni/model_executor/models/common/ming/audio_dsp.py
+209
View File
@@ -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"]
vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py → vllm_omni/model_executor/models/common/ming/audio_vae.py
+175 -188
View File
@@ -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,
vllm_omni/model_executor/models/common/ming/dit.py
+193
View File
@@ -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)
vllm_omni/model_executor/models/common/ming/fm.py
+68
View File
@@ -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
vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py
+1 -1
View File
@@ -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
vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py
+3
View File
@@ -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.
vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py
+7 -1
View File
@@ -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():
vllm_omni/model_executor/models/ming_flash_omni/talker_module.py
+6 -240
View File
@@ -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:
vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py
+2 -1
View File
@@ -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__)
vllm_omni/model_executor/models/ming_tts/__init__.py
+13
View File
@@ -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",
]
vllm_omni/model_executor/models/ming_tts/aggregator.py
+67
View File
@@ -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"]
vllm_omni/model_executor/models/ming_tts/audio_prep.py
+261
View File
@@ -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",
]
vllm_omni/model_executor/models/ming_tts/config_ming_tts.py
+209
View File
@@ -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",
]
vllm_omni/model_executor/models/ming_tts/constants.py
+67
View File
@@ -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 # <audioPatch>
AUDIO_START_TOKEN_ID = 151706 # <audio>
AUDIO_END_TOKEN_ID = 151707 # </audio>
AUDIO_EOS_TOKEN_ID = 151704 # <end_of_audio>
VISION_START_TOKEN_ID = 151652 # <|vision_start|>
TEXT_EOS_TOKEN_ID = 151669 # <text_eos>
# ---------------------------------------------------------------------------
# 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"
vllm_omni/model_executor/models/ming_tts/flowloss_head.py
+63
View File
@@ -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)
vllm_omni/model_executor/models/ming_tts/fm/__init__.py
+2
View File
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm_omni/model_executor/models/ming_tts/fm/dit.py
+98
View File
@@ -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:, :]
vllm_omni/model_executor/models/ming_tts/ming_tts.py
+269
View File
@@ -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",
]
vllm_omni/model_executor/models/ming_tts/ming_tts_audio_vae.py
+389
View File
@@ -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)
vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py
+414
View File
@@ -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",
]
vllm_omni/model_executor/models/ming_tts/patch_emission.py
+183
View File
@@ -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,
)
vllm_omni/model_executor/models/ming_tts/pipeline.py
+54
View File
@@ -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,
},
),
),
)
vllm_omni/model_executor/models/ming_tts/prompt_assembly.py
+274
View File
@@ -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("<audioPatch>")
if audio_patch_token_id == tokenizer.unk_token_id:
raise ValueError("Ming tokenizer is missing required <audioPatch> 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("<audio>")
+ prompt_latent_tokens
)
def build_ming_dense_prompt(
tokenizer: Any,
*,
prompt: str,
text: str,
runtime_controls: dict[str, Any] | None = None,
instruction: Any = None,
prompt_text: str | None = None,
prompt_waveform: Any = None,
prompt_latents: Any = None,
speaker_embedding: Any = None,
use_zero_spk_emb: bool = False,
request_id: str | None = None,
) -> dict[str, Any]:
instruction_text = create_instruction(instruction)
speaker_embeddings = coerce_speaker_embeddings(speaker_embedding, use_zero_spk_emb=use_zero_spk_emb)
effective_runtime_controls = resolve_effective_runtime_controls(text=text, runtime_controls=runtime_controls)
prompt_waveform_tensor = None
prompt_patch_count = 0
if prompt_waveform is not None:
prompt_waveform_tensor = pad_prompt_waveform(prompt_waveform)
prompt_patch_count = count_prompt_waveform_patches(prompt_waveform_tensor)
if prompt_waveform_tensor is not None and prompt_latents is not None:
raise ValueError(
"Ming waveform cloning request provided both raw prompt_waveform and explicit prompt_latents. "
"Choose exactly one source of truth."
)
prompt_latent_value = None
if prompt_waveform_tensor is not None and prompt_text is None:
raise ValueError(
"Ming prompt_waveform requires prompt_text for prompt-latent conditioning. "
"Use speaker_embedding for reference-audio-only speaker conditioning."
)
if prompt_latents is not None:
prompt_latent_value = torch.as_tensor(prompt_latents)
prompt_patch_count = count_prompt_latent_patches(
prompt_latent_value, patch_size=PATCH_SIZE, latent_dim=LATENT_DIM
)
prompt_token_ids = build_dense_prompt_token_ids(
tokenizer,
prompt=prompt,
text=text,
instruction=instruction_text,
prompt_text=prompt_text if prompt_patch_count > 0 else None,
speaker_count=0 if speaker_embeddings is None else len(speaker_embeddings),
prompt_patch_count=prompt_patch_count,
)
additional_information = {}
for key, value in effective_runtime_controls.items():
if isinstance(value, torch.Tensor):
additional_information[key] = value
elif key in (KEY_MIN_DECODE_STEPS, KEY_MAX_DECODE_STEPS):
additional_information[key] = torch.tensor(int(value), dtype=torch.int32)
else:
additional_information[key] = torch.tensor(float(value), dtype=torch.float32)
if request_id is not None:
additional_information[KEY_REQUEST_ID] = request_id
if instruction_text is not None:
additional_information["instruction"] = instruction_text
if prompt_text is not None:
additional_information["prompt_text"] = prompt_text
if prompt_waveform_tensor is not None:
additional_information["prompt_waveform"] = prompt_waveform_tensor
additional_information["prompt_waveform_length"] = torch.tensor(
[int(prompt_waveform_tensor.shape[-1])], dtype=torch.int32
)
if prompt_latent_value is not None:
additional_information[KEY_PROMPT_LATENTS] = prompt_latent_value
if speaker_embeddings is not None:
additional_information[KEY_SPEAKER_EMBEDDING] = (
speaker_embeddings[0] if len(speaker_embeddings) == 1 else torch.stack(speaker_embeddings, dim=0)
)
if use_zero_spk_emb:
additional_information["use_zero_spk_emb"] = True
return {
"prompt": prompt,
"text": text,
"prompt_token_ids": prompt_token_ids,
"additional_information": additional_information,
}
__all__ = [
"DEFAULT_PROMPT",
"build_dense_prompt_token_ids",
"build_ming_dense_prompt",
"create_instruction",
"estimate_decode_step_window_for_duration",
"estimate_decode_steps_for_duration",
"parse_duration_seconds",
"resolve_effective_runtime_controls",
]
vllm_omni/model_executor/models/ming_tts/prompt_encoder.py
+174
View File
@@ -0,0 +1,174 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import json
import threading
import time
from pathlib import Path
from typing import Any
import torch
from safetensors import safe_open
from vllm.logger import init_logger
from vllm_omni.engine.stage_init_utils import _resolve_model_to_local_path
from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
from .audio_prep import (
_coerce_prompt_latents,
_normalize_prompt_waveform,
count_prompt_latent_patches,
pad_prompt_waveform,
)
from .config_ming_tts import KEY_PROMPT_LATENTS
logger = init_logger(__name__)
_PROMPT_ENCODER_LOAD_LOCK = threading.Lock()
def _resolve_prompt_latents(wrapper: Any, info_dict: dict[str, Any]) -> dict[str, torch.Tensor] | None:
raw_latents = info_dict.get(KEY_PROMPT_LATENTS, info_dict.get("prompt_latents"))
raw_waveform = info_dict.get("prompt_waveform", info_dict.get("prompt_waveforms"))
if raw_latents is not None and raw_waveform is not None:
raise ValueError(
"Ming waveform cloning request provided both raw prompt_waveform and explicit prompt_latents. "
"Choose exactly one source of truth."
)
direct_latents = _coerce_prompt_latents(
raw_latents,
patch_size=wrapper.ming_config.patch_size,
latent_dim=wrapper.ming_config.latent_dim,
)
if direct_latents is not None:
return direct_latents
if raw_waveform is None:
return None
encode_fn = getattr(wrapper, "_encode_prompt_waveform_to_latents", None)
if callable(encode_fn):
latents = encode_fn(raw_waveform, info_dict.get("prompt_waveform_length"))
else:
latents = _encode_prompt_waveform_to_latents(
wrapper,
raw_waveform,
info_dict.get("prompt_waveform_length"),
)
return _coerce_prompt_latents(
latents,
patch_size=wrapper.ming_config.patch_size,
latent_dim=wrapper.ming_config.latent_dim,
)
def _load_prompt_encoder(wrapper: Any) -> AudioVAE:
if wrapper._prompt_encoder is not None:
return wrapper._prompt_encoder
with _PROMPT_ENCODER_LOAD_LOCK:
if wrapper._prompt_encoder is not None:
return wrapper._prompt_encoder
if wrapper.ming_config.audio_tokenizer_config is None:
raise RuntimeError("Ming Stage-0 requires audio_tokenizer_config to encode prompt audio.")
load_start = time.perf_counter()
encoder = AudioVAE(wrapper.ming_config.audio_tokenizer_config).eval()
state_dict = encoder.state_dict()
loaded = 0
loaded_encoder_params = set()
with torch.no_grad():
for shard_path in _iter_model_safetensors(
_resolve_model_to_local_path(str(wrapper.vllm_config.model_config.model))
):
with safe_open(str(shard_path), framework="pt", device="cpu") as handle:
for key in handle.keys():
if not key.startswith("audio.encoder."):
continue
name = key[len("audio.") :]
if name not in state_dict:
continue
target = state_dict[name]
target.copy_(handle.get_tensor(key).to(device=target.device, dtype=target.dtype))
loaded += 1
loaded_encoder_params.add(name)
if loaded == 0:
raise RuntimeError("Ming prompt encoder received no audio.encoder.* weights from checkpoint.")
expected_encoder_params = {f"encoder.{name}" for name, _ in encoder.encoder.named_parameters()}
missing = expected_encoder_params - loaded_encoder_params
if missing:
raise RuntimeError(
f"Ming prompt encoder: {len(missing)} params not loaded. First few: {sorted(missing)[:5]}"
)
dev = next(wrapper.parameters()).device
try:
del encoder.decoder
encoder.decoder = None
if dev.type != "cpu":
encoder.encoder.to(dev, dtype=getattr(wrapper.model, "fm_dtype", torch.bfloat16))
else:
encoder.encoder.to(dev)
except Exception as exc:
raise RuntimeError(f"Failed to move Ming prompt encoder to {dev}: {exc}") from exc
wrapper._prompt_encoder = encoder
logger.info("Ming prompt encoder cold-loaded in %.3f ms", (time.perf_counter() - load_start) * 1000.0)
return encoder
@torch.inference_mode()
def _encode_prompt_waveform_to_latents(wrapper: Any, waveform: Any, waveform_length: Any = None) -> torch.Tensor:
encoder = _load_prompt_encoder(wrapper)
waveform = _normalize_prompt_waveform(waveform, target_sr=wrapper.ming_config.sample_rate)
waveform = pad_prompt_waveform(
waveform,
patch_size=wrapper.ming_config.patch_size,
sample_rate=wrapper.ming_config.sample_rate,
)
dev = next(encoder.encoder.parameters()).device
waveform = waveform.to(device=dev, dtype=next(encoder.encoder.parameters()).dtype)
if waveform_length is None:
waveform_length = torch.full((waveform.shape[0],), waveform.shape[-1], dtype=torch.int32, device=dev)
elif not isinstance(waveform_length, torch.Tensor):
waveform_length = torch.as_tensor(waveform_length, dtype=torch.int32, device=dev)
else:
waveform_length = waveform_length.to(device=dev, dtype=torch.int32)
latents, _ = encoder.encode_latent(waveform, waveform_length)
if latents.ndim == 3 and latents.shape[0] == 1:
latents = latents.squeeze(0)
count_prompt_latent_patches(
latents,
patch_size=wrapper.ming_config.patch_size,
latent_dim=wrapper.ming_config.latent_dim,
)
return latents.detach().to(dtype=torch.float32).contiguous()
def _iter_model_safetensors(local_model_path: str) -> list[Path]:
model_root = Path(local_model_path)
index_path = model_root / "model.safetensors.index.json"
if index_path.exists():
with index_path.open("r", encoding="utf-8") as handle:
index_data = json.load(handle)
filenames = sorted(set(index_data.get("weight_map", {}).values()))
if not filenames:
raise RuntimeError(f"No checkpoint shards listed in {index_path}")
return [model_root / filename for filename in filenames]
single_file = model_root / "model.safetensors"
if single_file.exists():
return [single_file]
files = sorted(model_root.glob("*.safetensors"))
if not files:
raise RuntimeError(f"No .safetensors checkpoint found under {local_model_path}")
return files
__all__ = [
"_encode_prompt_waveform_to_latents",
"_iter_model_safetensors",
"_load_prompt_encoder",
"_resolve_prompt_latents",
]
vllm_omni/model_executor/models/ming_tts/speaker_extractor.py
+51
View File
@@ -0,0 +1,51 @@
# 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/spkemb_extractor.py
import os
import torch
import torchaudio
from vllm.multimodal.media.audio import load_audio
from vllm_omni.model_executor.models.ming_flash_omni.spk_embedding import SpkembExtractor
def resolve_model_to_local_path(model):
if os.path.isdir(model):
return model
from huggingface_hub import snapshot_download
return snapshot_download(model)
class MingSpeakerEmbeddingExtractor:
def __init__(self, model, target_sr=16000):
local_model_path = resolve_model_to_local_path(model)
campplus_path = os.path.join(local_model_path, "campplus.onnx")
if not os.path.exists(campplus_path):
raise RuntimeError(f"Missing Ming speaker extractor model: {campplus_path}")
self.target_sr = int(target_sr)
self._core = SpkembExtractor(campplus_path, target_sr=self.target_sr)
def extract_from_waveform(self, waveform, sample_rate):
if not isinstance(waveform, torch.Tensor):
waveform = torch.as_tensor(waveform)
tensor = waveform.detach().to(torch.float32)
if tensor.ndim == 1:
tensor = tensor.unsqueeze(0)
if int(sample_rate) != self.target_sr:
tensor = torchaudio.transforms.Resample(orig_freq=int(sample_rate), new_freq=self.target_sr)(tensor)
embedding = self._core._extract_spk_embedding(tensor)
return embedding.squeeze(0).to(dtype=torch.float32)
def extract_from_file(self, audio_path):
audio, sample_rate = load_audio(audio_path, sr=None, mono=False)
waveform = torch.from_numpy(audio).contiguous()
return self.extract_from_waveform(waveform, sample_rate)
def extract_many(self, audio_paths):
return [self.extract_from_file(path) for path in audio_paths]
vllm_omni/model_executor/models/ming_tts/validation.py
+172
View File
@@ -0,0 +1,172 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
from typing import Any
from transformers import PretrainedConfig
from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig
from .constants import (
AGGREGATOR_HIDDEN_SIZE,
HISTORY_PATCH_SIZE,
LATENT_DIM,
LLM_HIDDEN_SIZE,
LLM_VOCAB_SIZE,
PATCH_SIZE,
SAMPLE_RATE,
)
def _to_plain_dict(obj: Any) -> dict[str, Any]:
"""Normalize nested config objects into plain dicts when possible."""
if obj is None:
return {}
if isinstance(obj, dict):
return dict(obj)
if isinstance(obj, PretrainedConfig):
return obj.to_dict()
if hasattr(obj, "to_dict") and callable(obj.to_dict):
try:
return dict(obj.to_dict())
except Exception:
pass
try:
return dict(vars(obj))
except Exception:
return {}
def _coerce_audio_vae_config(atc_raw: Any) -> AudioVAEConfig | None:
if atc_raw is None:
return None
if isinstance(atc_raw, AudioVAEConfig):
return atc_raw
if isinstance(atc_raw, PretrainedConfig):
atc_dict = atc_raw.to_dict()
elif isinstance(atc_raw, dict):
atc_dict = dict(atc_raw)
elif hasattr(atc_raw, "to_dict") and callable(atc_raw.to_dict):
atc_dict = dict(atc_raw.to_dict())
else:
raise TypeError(f"Unsupported audio_tokenizer_config type for Ming dense config: {type(atc_raw)!r}")
return AudioVAEConfig(**atc_dict)
def _nested_get(obj: Any, *keys: str, default: Any = None) -> Any:
"""Safe nested attribute/key access for dicts and config-like objects."""
cur = obj
for key in keys:
if cur is None:
return default
if isinstance(cur, dict):
cur = cur.get(key)
else:
cur = getattr(cur, key, None)
return cur if cur is not None else default
def validate_ming_tts_config(cfg: Any) -> None:
"""Run before GPU allocation/weight loading. Raises ValueError on mismatches."""
if cfg.audio_dummy_token_id != 151705:
raise ValueError(
f"audio_dummy_token_id={cfg.audio_dummy_token_id}, expected 151705 (<audioPatch>). "
"Wrong tokenizer/checkpoint?"
)
if cfg.audio_eos_token_id != 151704:
raise ValueError(
f"audio_eos_token_id={cfg.audio_eos_token_id}, expected 151704 (<end_of_audio>). "
"Wrong tokenizer/checkpoint?"
)
if cfg.text_eos_token_id != 151669:
raise ValueError(
f"text_eos_token_id={cfg.text_eos_token_id}, expected 151669 (<text_eos>). Wrong tokenizer/checkpoint?"
)
if cfg.audio_tokenizer_config is None:
raise ValueError("audio_tokenizer_config is None. Nested AudioVAE config was not deserialized correctly.")
if cfg.latent_dim != LATENT_DIM:
raise ValueError(
f"latent_dim mismatch: got {cfg.latent_dim}, expected {LATENT_DIM}. "
"Check audio_tokenizer_config.enc_kwargs.latent_dim."
)
if cfg.patch_size != PATCH_SIZE:
raise ValueError(
f"patch_size mismatch: got {cfg.patch_size}, expected {PATCH_SIZE}. Check ditar_config.patch_size."
)
if cfg.history_patch_size != HISTORY_PATCH_SIZE:
raise ValueError(
f"history_patch_size mismatch: got {cfg.history_patch_size}, expected {HISTORY_PATCH_SIZE}. "
"Check ditar_config.history_patch_size."
)
if cfg.llm_hidden_size != LLM_HIDDEN_SIZE:
raise ValueError(
f"llm_hidden_size mismatch: got {cfg.llm_hidden_size}, expected {LLM_HIDDEN_SIZE}. "
"Check llm_config.hidden_size."
)
if cfg.llm_vocab_size != LLM_VOCAB_SIZE:
raise ValueError(f"llm_vocab_size mismatch: got {cfg.llm_vocab_size}, expected {LLM_VOCAB_SIZE}.")
if cfg.sample_rate != SAMPLE_RATE:
raise ValueError(f"sample_rate mismatch: got {cfg.sample_rate}, expected {SAMPLE_RATE}.")
if cfg.vae_patch_size != cfg.patch_size:
raise ValueError(f"VAE patch size ({cfg.vae_patch_size}) != flow/DiT patch size ({cfg.patch_size}).")
llm_hidden_from_cfg = cfg.llm_config.get("hidden_size")
if llm_hidden_from_cfg is not None and llm_hidden_from_cfg != cfg.llm_hidden_size:
raise ValueError(f"llm_hidden_size ({cfg.llm_hidden_size}) != llm_config.hidden_size ({llm_hidden_from_cfg}).")
agg_h = cfg.aggregator_config.get("hidden_size")
dit_h = cfg.ditar_config.get("hidden_size")
if agg_h is not None and dit_h is not None and agg_h != dit_h:
raise ValueError(f"aggregator_config.hidden_size ({agg_h}) != ditar_config.hidden_size ({dit_h}).")
if agg_h is not None and agg_h != AGGREGATOR_HIDDEN_SIZE:
raise ValueError(f"aggregator hidden_size mismatch: got {agg_h}, expected {AGGREGATOR_HIDDEN_SIZE}.")
if dit_h is not None and dit_h != AGGREGATOR_HIDDEN_SIZE:
raise ValueError(f"ditar hidden_size mismatch: got {dit_h}, expected {AGGREGATOR_HIDDEN_SIZE}.")
atc = cfg.audio_tokenizer_config
semantic_module_kwargs = _nested_get(atc, "semantic_module_kwargs", default=None)
if semantic_module_kwargs is not None:
raise ValueError("Ming dense 0.5B expects audio_tokenizer_config.semantic_module_kwargs to be null.")
enc_latent = _nested_get(atc, "enc_kwargs", "latent_dim", default=None)
dec_latent = _nested_get(atc, "dec_kwargs", "latent_dim", default=None)
if enc_latent is not None and enc_latent != cfg.latent_dim:
raise ValueError(f"audio enc latent_dim ({enc_latent}) != Ming latent_dim ({cfg.latent_dim}).")
if dec_latent is not None and dec_latent != cfg.latent_dim:
raise ValueError(f"audio dec latent_dim ({dec_latent}) != Ming latent_dim ({cfg.latent_dim}).")
atc_patch = _nested_get(atc, "patch_size", default=None)
if atc_patch is not None and atc_patch != cfg.vae_patch_size:
raise ValueError(f"audio_tokenizer_config.patch_size ({atc_patch}) != vae_patch_size ({cfg.vae_patch_size}).")
atc_sr = _nested_get(atc, "sample_rate", default=None)
if atc_sr is not None and atc_sr != cfg.sample_rate:
raise ValueError(f"audio_tokenizer_config.sample_rate ({atc_sr}) != sample_rate ({cfg.sample_rate}).")
enc_input_dim = _nested_get(atc, "enc_kwargs", "input_dim", default=None)
enc_hop_size = _nested_get(atc, "enc_kwargs", "hop_size", default=None)
dec_output_dim = _nested_get(atc, "dec_kwargs", "output_dim", default=None)
if enc_input_dim is not None and enc_hop_size is not None and enc_input_dim != enc_hop_size:
raise ValueError(f"AudioVAE encoder input_dim ({enc_input_dim}) != hop_size ({enc_hop_size}).")
if enc_hop_size is not None and dec_output_dim is not None and enc_hop_size != dec_output_dim:
raise ValueError(
f"AudioVAE encoder hop_size ({enc_hop_size}) != decoder output_dim ({dec_output_dim}). "
"Expected 882 in this checkpoint family."
)
if cfg.latent_chunk_size <= 0:
raise ValueError(f"latent_chunk_size must be > 0, got {cfg.latent_chunk_size}.")
if cfg.latent_left_context < 0:
raise ValueError(f"latent_left_context must be >= 0, got {cfg.latent_left_context}.")
if cfg.max_decode_steps <= 0:
raise ValueError(f"max_decode_steps must be > 0, got {cfg.max_decode_steps}.")
if not (0.0 <= cfg.stop_head_threshold <= 1.0):
raise ValueError(f"stop_head_threshold must be in [0,1], got {cfg.stop_head_threshold}.")
if cfg.stop_head_min_steps < 0:
raise ValueError(f"stop_head_min_steps must be >= 0, got {cfg.stop_head_min_steps}.")
vllm_omni/model_executor/models/registry.py
+16
View File
@@ -156,6 +156,22 @@ _OMNI_MODELS = {
"mimo_audio_code2wav",
"MiMoAudioToken2WavForConditionalGenerationVLLM",
),
## ming-tts
"MingTTSForConditionalGeneration": (
"ming_tts",
"ming_tts",
"MingTTSForConditionalGeneration",
),
"MingLLMModel": (
"ming_tts",
"ming_tts_llm",
"MingLLMModel",
),
"MingAudioVAEModel": (
"ming_tts",
"ming_tts_audio_vae",
"MingAudioVAEModel",
),
## glm_image
"GlmImageForConditionalGeneration": (
"glm_image",
vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py
+15 -76
View File
@@ -205,18 +205,6 @@ def _resolve_image_patch_token_id(stage: Any) -> int:
return token_id
def _validate_stage_inputs(stage_list, engine_input_source):
if not engine_input_source:
raise ValueError("engine_input_source cannot be empty")
stage_id = engine_input_source[0]
if stage_id >= len(stage_list):
raise IndexError(f"Invalid stage_id: {stage_id}")
stage = stage_list[stage_id]
if stage.engine_outputs is None:
raise RuntimeError(f"Stage {stage_id} has no outputs yet")
return stage, stage.engine_outputs
def _ensure_list(x) -> list[int]:
"""Convert ConstantList / tensor-like to plain list."""
if hasattr(x, "_x"):
@@ -458,20 +446,10 @@ def thinker2imagegen(
return [{"prompt": "", "extra": extra}]
def thinker2talker(
stage_list: list[Any],
engine_input_source: list[int],
def _build_talker_inputs(
source_outputs: list[Any],
prompt: OmniTokensPrompt | TextPrompt | None = None,
requires_multimodal_data: bool = False,
) -> list[OmniTokensPrompt]:
"""Build talker stage inputs from thinker stage outputs.
Extracts the generated text from thinker output and constructs
a talker input prompt with the text and any speaker/instruction info
from the original request.
"""
_, source_outputs = _validate_stage_inputs(stage_list, engine_input_source)
if not isinstance(prompt, list):
prompt = [prompt]
@@ -492,8 +470,6 @@ def thinker2talker(
# the talker's spk_head wants a torch tensor.
spk_emb = additional_info.get("spk_emb", None)
if isinstance(spk_emb, list) and spk_emb and not hasattr(spk_emb[0], "device"):
import torch
spk_emb = torch.tensor(spk_emb, dtype=torch.float32).unsqueeze(0)
# Omni speech path mirrors upstream `omni_audio_generation`:
@@ -535,12 +511,22 @@ def thinker2talker(
return talker_inputs
def thinker2talker(
source_outputs: list[Any],
prompt: OmniTokensPrompt | TextPrompt | None = None,
_requires_multimodal_data: bool = False,
_streaming_context: Any | None = None,
) -> list[OmniTokensPrompt]:
"""Build talker stage inputs from thinker stage outputs."""
return _build_talker_inputs(source_outputs, prompt)
# ming_flash_omni is not in ``_OMNI_CONNECTOR_INIT_ARCHS`` or
# ``_FULL_PAYLOAD_INPUT_STAGES``, so the worker connector is not
# initialised for this arch and the consumer never waits on a connector
# payload. Data flows through ``additional_information`` written by
# ``thinker2talker_token_only`` (wired as ``sync_process_input_func``
# in the pipeline) or the legacy ``thinker2talker`` (wired as
# in the pipeline) or ``thinker2talker`` (wired as
# ``custom_process_input_func``).
@@ -549,55 +535,8 @@ def thinker2talker_token_only(
prompt: OmniTokensPrompt | TextPrompt | None = None,
_requires_multimodal_data: bool = False,
) -> list[OmniTokensPrompt]:
"""Sync-side builder for the non-async-chunk thinker→talker path.
Ports the legacy ``thinker2talker`` body to the new stage-input-processor signature
(``source_outputs`` instead of ``stage_list, engine_input_source``).
Body is otherwise identical: extracts the
generated text from each thinker output and packages it with the
request's voice/speaker additional_information for the talker.
"""
if not isinstance(prompt, list):
prompt = [prompt]
talker_inputs: list[OmniTokensPrompt] = []
for i, source_output in enumerate(source_outputs):
output = source_output.outputs[0]
generated_text = output.text if hasattr(output, "text") and output.text else ""
original_prompt = prompt[i] if i < len(prompt) else None
additional_info: dict[str, Any] = {}
if original_prompt is not None and hasattr(original_prompt, "additional_information"):
additional_info = original_prompt.additional_information or {}
spk_emb = additional_info.get("spk_emb", None)
if isinstance(spk_emb, list) and spk_emb and not hasattr(spk_emb[0], "device"):
import torch
spk_emb = torch.tensor(spk_emb, dtype=torch.float32).unsqueeze(0)
talker_info = {
"ming_task": "omni",
"text": generated_text,
"spk_emb": spk_emb,
"voice_name": additional_info.get("voice_name", "DB30"),
"prompt_text": additional_info.get("prompt_text", None),
"prompt_wav_lat": additional_info.get("prompt_wav_lat", None),
"prompt_wav_emb": additional_info.get("prompt_wav_emb", None),
"max_text_length": additional_info.get("max_text_length", 50),
}
talker_inputs.append(
OmniTokensPrompt(
prompt_token_ids=[0],
additional_information=talker_info,
multi_modal_data=None,
mm_processor_kwargs=None,
)
)
return talker_inputs
"""Sync-side builder for the non-async-chunk thinker→talker path."""
return _build_talker_inputs(source_outputs, prompt)
thinker2talker_token_only._is_sync_input = True
vllm_omni/model_executor/stage_input_processors/ming_tts.py
+296
View File
@@ -0,0 +1,296 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
from typing import Any
import torch
from vllm.inputs import TextPrompt
from vllm.logger import init_logger
from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
from vllm_omni.inputs.data import OmniTokensPrompt
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
KEY_CHUNK_ID,
KEY_REQUEST_ID,
LATENT_CHUNK_SIZE,
LATENT_DIM,
LATENT_LEFT_CONTEXT,
PATCH_SIZE,
)
from vllm_omni.model_executor.models.ming_tts.patch_emission import (
MING_STOP_REASON_CODES,
MING_STOP_REASON_KEY,
)
logger = init_logger(__name__)
MING_EMIT_PATCH_COUNT_KEY = "ming_emit_patch_count"
MING_LATENT_SHAPE_KEY = "ming_latent_shape"
MING_ESTIMATED_BYTES_KEY = "ming_estimated_bytes"
MING_FINAL_FLUSH_KEY = "ming_final_flush"
MING_FINAL_DECODE_STEP_KEY = "ming_final_decode_step"
MING_STOP_REASON_BY_CODE = {code: reason for reason, code in MING_STOP_REASON_CODES.items()}
def _extract_last_patch(pooling_output: dict[str, Any] | None) -> torch.Tensor | None:
if not isinstance(pooling_output, dict):
return None
has_patch = pooling_output.get("ming_has_patch")
patch = pooling_output.get("ming_latent_patch")
if not isinstance(patch, torch.Tensor) or patch.numel() == 0:
return None
if isinstance(has_patch, torch.Tensor) and has_patch.numel() > 0:
active = (has_patch.reshape(-1) > 0).nonzero(as_tuple=True)[0]
if active.numel() == 0:
return None
patch = patch[int(active[-1].item())]
elif patch.ndim == 3:
patch = patch[-1]
if patch.ndim != 2:
raise ValueError(f"Invalid Ming latent patch shape: {tuple(patch.shape)}")
return patch.to(torch.float32).cpu()
def _extract_all_patches(pooling_output: dict[str, Any] | None) -> torch.Tensor | None:
if not isinstance(pooling_output, dict):
return None
has_patch = pooling_output.get("ming_has_patch")
patch = pooling_output.get("ming_latent_patch")
if not isinstance(patch, torch.Tensor) or patch.numel() == 0:
return None
if patch.ndim == 2:
patch = patch.unsqueeze(0)
if patch.ndim != 3:
raise ValueError(f"Invalid Ming latent patch tensor shape: {tuple(patch.shape)}")
if isinstance(has_patch, torch.Tensor) and has_patch.numel() > 0:
active = (has_patch.reshape(-1) > 0).nonzero(as_tuple=True)[0]
if active.numel() == 0:
return None
patch = patch.index_select(0, active.to(device=patch.device))
if patch.numel() == 0:
return None
return patch.to(torch.float32).cpu()
def _extract_last_value(pooling_output: dict[str, Any] | None, key: str) -> Any:
if not isinstance(pooling_output, dict):
return None
value = pooling_output.get(key)
if value is None:
return None
has_patch = pooling_output.get("ming_has_patch")
selected_index = -1
if isinstance(has_patch, torch.Tensor) and has_patch.numel() > 0:
active = (has_patch.reshape(-1) > 0).nonzero(as_tuple=True)[0]
if active.numel() == 0:
return None
selected_index = int(active[-1].item())
if isinstance(value, torch.Tensor):
flat = value.reshape(-1)
if flat.numel() == 0:
return None
return flat[min(selected_index, flat.numel() - 1)].item()
if isinstance(value, (list, tuple)):
if not value:
return None
if selected_index < 0:
return value[-1]
return value[min(selected_index, len(value) - 1)]
return value
def _decode_stop_reason(value: Any) -> str | None:
if value is None:
return None
if isinstance(value, str):
return value
return MING_STOP_REASON_BY_CODE.get(int(value))
def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int]:
connector = getattr(transfer_manager, "connector", None)
raw_cfg = getattr(connector, "config", {}) or {}
cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
if "latent_chunk_size" not in cfg:
logger.warning(
"Ming async chunk config missing latent_chunk_size, using fallback value %s",
LATENT_CHUNK_SIZE,
)
if "latent_left_context" not in cfg:
logger.warning(
"Ming async chunk config missing latent_left_context, using fallback value %s",
LATENT_LEFT_CONTEXT,
)
chunk_size = int(cfg.get("latent_chunk_size", LATENT_CHUNK_SIZE))
left_context = int(cfg.get("latent_left_context", LATENT_LEFT_CONTEXT))
if chunk_size <= 0:
raise ValueError(f"Invalid Ming latent_chunk_size={chunk_size}")
# Stage-2 VAE caches past_key_values and stream_state by request_id.
# Replaying left-context latents would double-feed cached decoder state.
if left_context != 0:
raise ValueError(
"Ming async chunk transport does not support latent_left_context replay. "
"Ming boundary continuity is handled by per-request decoder state cache, not "
f"latent replay. Got latent_left_context={left_context}."
)
return chunk_size, left_context
def _build_chunk_observability(
latent_patches: torch.Tensor | None,
*,
final_flush: bool,
) -> dict[str, Any]:
if latent_patches is None:
emit_patch_count = 0
latent_shape = None
estimated_bytes = 0
else:
emit_patch_count = int(latent_patches.shape[0])
latent_shape = tuple(latent_patches.shape)
estimated_bytes = int(latent_patches.numel() * latent_patches.element_size())
return {
MING_EMIT_PATCH_COUNT_KEY: emit_patch_count,
MING_LATENT_SHAPE_KEY: latent_shape,
MING_ESTIMATED_BYTES_KEY: estimated_bytes,
MING_FINAL_FLUSH_KEY: bool(final_flush),
}
def llm2audio_vae_async_chunk(
transfer_manager: Any,
pooling_output: dict[str, Any] | None,
request: Any,
is_finished: bool = False,
) -> OmniPayloadStruct | None:
request_id = request.external_req_id
chunk_id = int(transfer_manager.put_req_chunk[request_id])
finished = bool(is_finished or request.is_finished())
final_decode_step = _extract_last_value(pooling_output, "ming_decode_step")
stop_reason = _decode_stop_reason(_extract_last_value(pooling_output, MING_STOP_REASON_KEY))
request_payload = transfer_manager.request_payload
request_state = request_payload.get(request_id)
if not isinstance(request_state, dict) or "_ming_async_state" not in request_state:
request_state = {
"_ming_async_state": {
"seen_patch_len": 0,
"terminal_sent": False,
}
}
request_payload[request_id] = request_state
state = request_state["_ming_async_state"]
if bool(state.get("terminal_sent", False)):
return None
patch = _extract_last_patch(pooling_output)
if patch is not None:
transfer_manager.code_prompt_token_ids[request_id].append(patch)
chunk_size, _ = _get_async_chunk_config(transfer_manager)
patches = transfer_manager.code_prompt_token_ids[request_id]
seen_patch_len = int(state.get("seen_patch_len", 0))
new_patches = patches[seen_patch_len:] if seen_patch_len < len(patches) else []
length = len(new_patches)
if length <= 0:
if finished and not bool(state.get("terminal_sent", False)):
observability = _build_chunk_observability(None, final_flush=True)
kv_meta = {
KEY_CHUNK_ID: chunk_id,
KEY_REQUEST_ID: request_id,
**observability,
}
if final_decode_step is not None:
kv_meta[MING_FINAL_DECODE_STEP_KEY] = int(final_decode_step)
if stop_reason is not None:
kv_meta[MING_STOP_REASON_KEY] = stop_reason
state["terminal_sent"] = True
return OmniPayloadStruct(
codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
meta=MetaStruct(
finished=torch.tensor(True, dtype=torch.bool),
stream_finished=torch.tensor(True, dtype=torch.bool),
),
kv_metadata=kv_meta,
)
return None
chunk_length = length % chunk_size
if chunk_length != 0 and not finished:
return None
emit_count = chunk_length if chunk_length != 0 else chunk_size
emit_patches = list(new_patches[:emit_count])
state["seen_patch_len"] = seen_patch_len + len(emit_patches)
latent_patches = torch.stack(emit_patches, dim=0)
observability = _build_chunk_observability(latent_patches, final_flush=finished)
kv_meta = {
KEY_CHUNK_ID: chunk_id,
KEY_REQUEST_ID: request_id,
**observability,
}
if final_decode_step is not None:
kv_meta[MING_FINAL_DECODE_STEP_KEY] = int(final_decode_step)
if stop_reason is not None:
kv_meta[MING_STOP_REASON_KEY] = stop_reason
if finished:
state["terminal_sent"] = True
return OmniPayloadStruct(
codes=CodesStruct(audio=torch.tensor([0], dtype=torch.long)),
meta=MetaStruct(
finished=torch.tensor(finished, dtype=torch.bool),
stream_finished=torch.tensor(finished, dtype=torch.bool),
),
latent=latent_patches,
kv_metadata=kv_meta,
)
def llm2audio_vae(
source_outputs: list[Any],
prompt: OmniTokensPrompt | TextPrompt | None = None,
requires_multimodal_data: bool = False,
streaming_context: Any = None,
) -> list[OmniTokensPrompt]:
del prompt, requires_multimodal_data, streaming_context
if not source_outputs:
raise ValueError("source_outputs cannot be empty")
outputs = []
for stage_output in source_outputs:
finished = bool(getattr(stage_output, "finished", True))
if not finished:
continue
output = stage_output.outputs[0]
patches = _extract_all_patches(output.multimodal_output)
additional_information = {
"ming_latent_patches": patches
if patches is not None
else torch.zeros((0, PATCH_SIZE, LATENT_DIM), dtype=torch.float32),
KEY_REQUEST_ID: getattr(stage_output, "request_id", None),
"finished": torch.tensor(finished, dtype=torch.bool),
}
final_decode_step = _extract_last_value(output.multimodal_output, "ming_decode_step")
stop_reason = _decode_stop_reason(_extract_last_value(output.multimodal_output, MING_STOP_REASON_KEY))
if final_decode_step is not None:
additional_information[MING_FINAL_DECODE_STEP_KEY] = int(final_decode_step)
if stop_reason is not None:
additional_information[MING_STOP_REASON_KEY] = stop_reason
outputs.append(
OmniTokensPrompt(
prompt_token_ids=[0],
multi_modal_data=None,
mm_processor_kwargs=None,
additional_information=additional_information,
)
)
return outputs