[Misc] Remove Dynin-Omni and dots.tts support (#7655)
Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
This commit is contained in:
1 parent
3f94c417be
commit
6a03e45c34
56 files changed
+129
-12745
No files matched your search
@@ -58,7 +58,7 @@ vLLM-Omni is flexible and easy to use with:
|
||||
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
|
||||
|
||||
- **Omni-modality models** (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
|
||||
- **TTS models** (e.g. Qwen3-TTS, IndexTTS 2.5, dots.tts, CosyVoice3)
|
||||
- **TTS models** (e.g. Qwen3-TTS, IndexTTS 2.5, CosyVoice3)
|
||||
- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2)
|
||||
- **Robot-policy and action models** (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1)
|
||||
|
||||
|
||||
+1
-1
@@ -56,7 +56,7 @@ vLLM-Omni is flexible and easy to use with:
|
||||
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
|
||||
|
||||
- Omni-modality models (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
|
||||
- TTS models (e.g. Qwen3-TTS, IndexTTS 2.5, dots.tts, CosyVoice3)
|
||||
- TTS models (e.g. Qwen3-TTS, IndexTTS 2.5, CosyVoice3)
|
||||
- Diffusion models — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2)
|
||||
- Robot-policy and action models (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1)
|
||||
|
||||
|
||||
@@ -73,7 +73,6 @@ th {
|
||||
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — |
|
||||
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — |
|
||||
| `BreezeForConditionalGeneration` | Breeze-TTS-2 (TTS, voice clone, voice direction) | `BreezeBlue/Breeze-TTS-2` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md) |
|
||||
| `DotsTTSForConditionalGeneration` | dots.tts | `dots-studio/dots.tts-soar` | ✅︎ | | | | — |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
|
||||
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
@@ -107,7 +106,6 @@ th {
|
||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
|
||||
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
|
||||
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
|
||||
| `DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — |
|
||||
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/OpenBMB/MiniCPM-o-4_5.md) |
|
||||
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — |
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
# The docs generator combines these names with the serving mode for each page H1.
|
||||
aura_omni: AURA Omni
|
||||
covo_audio: Covo-Audio-Chat
|
||||
dynin_omni: Dynin-Omni
|
||||
lance: Lance
|
||||
mimo_audio: MiMo-Audio
|
||||
ming_flash_omni: Ming-flash-omni 2.0
|
||||
|
||||
@@ -1,110 +0,0 @@
|
||||
# Dynin-Omni: Offline inference
|
||||
|
||||
This folder contains a unified offline inference entrypoint:
|
||||
|
||||
- `end2end.py`
|
||||
|
||||
## 1. Environment Setup
|
||||
|
||||
Run from repository root:
|
||||
|
||||
```bash
|
||||
cd <REPO_ROOT>
|
||||
```
|
||||
|
||||
If needed, install this repo in editable mode:
|
||||
|
||||
```bash
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
## 2. Extra Dependencies (EMOVA)
|
||||
|
||||
Install the following packages for EMOVA-related components:
|
||||
|
||||
```bash
|
||||
pip install \
|
||||
"phonemizer==3.3.0" \
|
||||
"Unidecode==1.4.0" \
|
||||
"hydra-core==1.3.2" \
|
||||
"pytorch-lightning==1.1.0" \
|
||||
"wget==3.2" \
|
||||
"wrapt==2.1.1" \
|
||||
"onnx==1.20.1" \
|
||||
"frozendict==2.4.7" \
|
||||
"inflect==7.5.0" \
|
||||
"braceexpand==0.1.7" \
|
||||
"webdataset==1.0.2" \
|
||||
"torch-stft==0.1.4" \
|
||||
"editdistance==0.8.1"
|
||||
```
|
||||
|
||||
## 3. Hardware and VRAM Requirements
|
||||
|
||||
This example uses a 3-stage pipeline on one GPU by default
|
||||
([`dynin_omni.yaml`](../../../vllm_omni/deploy/dynin_omni.yaml)):
|
||||
|
||||
- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
|
||||
- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`
|
||||
- Stage-2 (`token2audio`): `gpu_memory_utilization: 0.1`
|
||||
|
||||
### Requested GPU Memory Budget from `gpu_memory_utilization`
|
||||
|
||||
| Stage | Utilization | A100 80GB | H200 141GB |
|
||||
| :-- | :-- | :-- | :-- |
|
||||
| Stage-0 (token2text) | 0.5 | ~40.0 GB | ~70.5 GB |
|
||||
| Stage-1 (token2image) | 0.1 | ~8.0 GB | ~14.1 GB |
|
||||
| Stage-2 (token2audio) | 0.1 | ~8.0 GB | ~14.1 GB |
|
||||
| Total requested budget | 0.7 | ~56.0 GB | ~98.7 GB |
|
||||
|
||||
### Observed Runtime Signal (from your log)
|
||||
|
||||
- Stage-0 reported: `Model loading took 15.12 GiB memory` (weights footprint signal).
|
||||
- Stages 1/2 can still add runtime memory depending on task path and backend allocations.
|
||||
- Keep extra headroom for CUDA/PyTorch overhead and temporary allocations.
|
||||
|
||||
### GPU Compatibility
|
||||
|
||||
- Confirmed target GPUs for this setup: **NVIDIA H200**, **NVIDIA A100**.
|
||||
- CI/e2e coverage in this repo also includes CUDA **L4** markers for Dynin tests.
|
||||
|
||||
## 4. End2End Run Examples
|
||||
|
||||
```bash
|
||||
# t2t
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task t2t --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
|
||||
|
||||
# i2t
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task i2t --model snu-aidas/Dynin-Omni --image <IMAGE_PATH> --text "Please describe this image in detail."
|
||||
|
||||
# s2t
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task s2t --model snu-aidas/Dynin-Omni --audio <AUDIO_PATH> --text "Transcribe the given audio."
|
||||
|
||||
# t2i
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task t2i --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
|
||||
|
||||
# v2t
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task v2t --model snu-aidas/Dynin-Omni --video <VIDEO_PATH> --text "Describe this video in detail."
|
||||
|
||||
# i2i
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task i2i --model snu-aidas/Dynin-Omni --image <IMAGE_PATH> --text <INSTRUCTION_TEXT>
|
||||
|
||||
# t2s
|
||||
python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
|
||||
--task t2s --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
|
||||
```
|
||||
|
||||
## 5. Notes
|
||||
|
||||
- Outputs are saved under task-specific directories in `/tmp` by default.
|
||||
- You can override output path with `--output-dir`.
|
||||
- If you want to force local config resolution, pass `--dynin-config-path <PATH_TO_DYNIN_OMNI_YAML>`.
|
||||
- If you see the warning
|
||||
`max_num_batched_tokens (32768) exceeds max_num_seqs * max_model_len (4096)`,
|
||||
reduce `max_num_batched_tokens` in the deploy config (for example, `4096` in the CI config).
|
||||
File diff suppressed because it is too large.
Load diff
@@ -13,7 +13,7 @@ list of supported architectures across all modalities, see
|
||||
## Supported Models
|
||||
|
||||
| Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
|
||||
|---|---|---|---|---|---|---|
|
||||
| --- | --- | --- | --- | --- | --- | --- |
|
||||
| AuK | `tencent/AuK`, `tencent/AuK-Flash` (assembled with `tools/prepare_auk_checkpoint.py`) | 2 (thinker encoder + DiT) | ✓ | — | instruction-driven editing, enhancement, separation; no example script (use `vllm_omni.model_extras.auk`) | 24 kHz |
|
||||
| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | dual-AR | ✓ | ✓ | 11 languages | 44.1 kHz |
|
||||
| Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | 2 (talker + codec) | ✓ | ✓ (async chunk) | voice design / voice direction (`--instruction`) | 24 kHz |
|
||||
@@ -27,7 +27,6 @@ list of supported architectures across all modalities, see
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
|
||||
| dots.tts | `dots-studio/dots.tts-soar` | single (native AR) | — (text-only) | ✓ (online) | no-reference synthesis | 48 kHz |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz |
|
||||
| IndexTTS-2.5 | native `checkpoints/` bundle | 2 (AR talker + EnhancedCodec + S2Mel DiT + BigVGAN) | ✓ (required) | — | multilingual (`--lang`) + emotion control | 22.05 kHz |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
|
||||
@@ -54,32 +53,38 @@ codec: a slow AR transformer predicts one semantic token per codec frame, and a
|
||||
4-layer fast AR predicts that frame's remaining 9 codebooks.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
No extra packages: the codec architecture ships in tree and its weights
|
||||
(`codec.pth`) come with the checkpoint.
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
|
||||
--text "Welcome to Audio8 TTS, a compact multilingual text to speech model."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
|
||||
--text "Welcome to Audio8 TTS." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "The exact transcript of the reference recording."
|
||||
```
|
||||
|
||||
The reference transcript must match what is actually spoken in the clip;
|
||||
a mismatch degrades both stability and speaker similarity.
|
||||
|
||||
### Streaming
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
|
||||
--text "Welcome to Audio8 TTS." --streaming
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 44.1 kHz mono WAV; the codec emits ~21.5 frames/s (2048 samples per frame).
|
||||
- Context is 2048 packed text+audio positions, so `max_model_len` is 2048.
|
||||
- Greedy decoding (`temperature: 0`) is degenerate for this checkpoint — it
|
||||
@@ -92,6 +97,7 @@ python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
|
||||
- `--num-prompts 4` smoke-tests concurrency (stage 0 runs `max_num_seqs: 4`).
|
||||
|
||||
### Measured behaviour (1x H20, shared GPU)
|
||||
|
||||
`--streaming` on a 58-character English sentence: TTFA ~0.74 s (4-frame first
|
||||
chunk), 7 chunks, 1.55 s wall for 3.16 s of audio (RTF ~0.49). Treat these as
|
||||
smoke-test magnitudes, not a benchmark.
|
||||
@@ -103,12 +109,14 @@ smoke-test magnitudes, not a benchmark.
|
||||
2-stage TTS pipeline (`talker` + `code2wav`) at 24 kHz.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
uv pip install -e .
|
||||
# Includes soundfile, onnxruntime, x-transformers, einops via requirements.
|
||||
```
|
||||
|
||||
Download the model snapshot:
|
||||
|
||||
```python
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512',
|
||||
@@ -116,15 +124,18 @@ snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512',
|
||||
```
|
||||
|
||||
If your downloaded checkpoint lacks `config.json`, add it:
|
||||
|
||||
```json
|
||||
{
|
||||
"model_type": "cosyvoice3",
|
||||
"architectures": ["CosyVoice3Model"]
|
||||
}
|
||||
```
|
||||
|
||||
This is required because `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers the class mapping; the loader still reads `model_type` from `config.json` to select the class.
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
--model pretrained_models/Fun-CosyVoice3-0.5B \
|
||||
@@ -132,9 +143,11 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
If `--ref-audio` is omitted, the script downloads the upstream
|
||||
[`zero_shot_prompt.wav`](https://github.com/FunAudioLLM/CosyVoice/blob/main/asset/zero_shot_prompt.wav) from the CosyVoice repo into the current directory.
|
||||
To use your own clip, pass `--ref-audio /path/to/reference.wav`, and modify `--prompt-text` correspondingly.
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
--model pretrained_models/Fun-CosyVoice3-0.5B \
|
||||
@@ -144,9 +157,11 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
```
|
||||
|
||||
### Streaming
|
||||
|
||||
Streaming is enabled by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` on `vllm serve` to switch to the legacy synchronous path.
|
||||
|
||||
### Notes
|
||||
|
||||
- Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
|
||||
- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
|
||||
|
||||
@@ -157,6 +172,7 @@ Streaming is enabled by default via `async_chunk: true` in `vllm_omni/deploy/cos
|
||||
2-stage TTS pipeline (AR + DiT flow-matching) at 24 kHz. Every request requires reference audio and its transcript for zero-shot voice cloning.
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
|
||||
--model zai-org/GLM-TTS \
|
||||
@@ -167,12 +183,14 @@ python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
|
||||
```
|
||||
|
||||
### Architecture
|
||||
|
||||
```text
|
||||
Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz)
|
||||
(Llama-based) (32k vocab) (Flow Matching)
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- `--ref-audio` and `--ref-text` are **required** together; GLM-TTS does not support text-only synthesis.
|
||||
- Reference audio should be 3-10 seconds.
|
||||
- First run may be slow due to lazy loading of WhisperVQ tokenizer and CampPlus ONNX speaker embedder.
|
||||
@@ -187,12 +205,14 @@ Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24
|
||||
No extra packages are required; the DAC codec is vendored in vLLM-Omni.
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
|
||||
--text "Hello, this is a test of the Fish Speech text to speech system."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
@@ -201,14 +221,17 @@ python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
|
||||
```
|
||||
|
||||
### Streaming
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
|
||||
--text "Hello, this is a streaming test." \
|
||||
--streaming
|
||||
```
|
||||
|
||||
Streaming requires `async_chunk: true` in the deploy config.
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 44.1 kHz mono WAV.
|
||||
- DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
|
||||
|
||||
@@ -219,6 +242,7 @@ Streaming requires `async_chunk: true` in the deploy config.
|
||||
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 \
|
||||
@@ -227,6 +251,7 @@ python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case zero_shot \
|
||||
@@ -237,6 +262,7 @@ python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
```
|
||||
|
||||
### Streaming
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case basic \
|
||||
@@ -247,6 +273,7 @@ python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
```
|
||||
|
||||
### 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`.
|
||||
@@ -259,14 +286,17 @@ python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
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.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
The example calls into `vllm_omni.model_executor.models.ming_flash_omni.prompt_utils` for the default prompt and instruction builder; no extra pip install on top of the base vLLM-Omni install.
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py --case style
|
||||
```
|
||||
|
||||
### Cases
|
||||
|
||||
```bash
|
||||
# ASMR-style whisper (caption-driven)
|
||||
python examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py --case style
|
||||
@@ -281,6 +311,7 @@ python examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py
|
||||
Override the default text per case with `--text`, write to a custom path with `--output`.
|
||||
|
||||
### Notes
|
||||
|
||||
- Talker-only deployment — for the multimodal Ming-flash-omni example, see [`examples/offline_inference/ming_flash_omni/`](../../ming_flash_omni/).
|
||||
- Deploy config: `vllm_omni/deploy/ming_flash_omni_tts.yaml` (single GPU, `enforce_eager`, `max_num_seqs: 1`).
|
||||
- Decode defaults from the Ming cookbook: `max_decode_steps=200`, `cfg=2.0`, `sigma=0.25`, `temperature=0.0`, `use_zero_spk_emb=True`.
|
||||
@@ -294,6 +325,7 @@ Single-stage 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec at 48 kHz mono (mixed
|
||||
> **No built-in speaker presets.** `--ref-audio` is required on every call. Default `--mode voice_clone` matches upstream's recommended workflow; `--mode continuation` is exposed for completeness but upstream's continuation-with-prompt path emits very short / near-silent output, so it is rarely useful in practice. Sample reference clips ship in the upstream repo under [`assets/audio/`](https://github.com/OpenMOSS/MOSS-TTS-Nano/tree/main/assets/audio) (e.g. `zh_1.wav`, `en_2.wav`, `jp_2.wav`).
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
# Fetch a sample reference clip (one-off, user-scoped cache).
|
||||
REF_DIR="${XDG_CACHE_HOME:-$HOME/.cache}/moss-tts-nano"
|
||||
@@ -305,9 +337,11 @@ python examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py \
|
||||
--text "你好,这是MOSS-TTS-Nano的语音合成演示。" \
|
||||
--ref-audio "$REF_DIR/zh_1.wav"
|
||||
```
|
||||
|
||||
The first run downloads `OpenMOSS-Team/MOSS-TTS-Nano` and `OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano` from Hugging Face.
|
||||
|
||||
### Reproducible runs
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py \
|
||||
--text "Deterministic test." \
|
||||
@@ -316,6 +350,7 @@ python examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py \
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 48 kHz mono WAV (the tokenizer is internally stereo at 48 kHz; the wrapper averages to mono before reaching the engine).
|
||||
- Deploy config: `vllm_omni/deploy/moss_tts_nano.yaml` (auto-loaded; override with `--deploy-config`).
|
||||
- Default `--max-new-frames 375` ≈ 14 s of audio; raise for longer outputs.
|
||||
@@ -329,12 +364,15 @@ python examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py \
|
||||
Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design).
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
huggingface-cli download k2-fsa/OmniVoice
|
||||
```
|
||||
|
||||
Voice cloning requires `transformers>=5.3.0`. Auto and design modes work with `transformers>=4.57.0`.
|
||||
|
||||
### Quick start (auto voice)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
@@ -342,6 +380,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
@@ -351,6 +390,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
```
|
||||
|
||||
### Voice design
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
@@ -359,6 +399,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
```
|
||||
|
||||
### Language hint
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
@@ -367,6 +408,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
```
|
||||
|
||||
### Seed for Reproducibility
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
@@ -375,6 +417,7 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking.
|
||||
- Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz.
|
||||
|
||||
@@ -385,13 +428,16 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
3-task-variant TTS with 24 kHz output. Has its own argparse surface (this script does not follow the common `--text` / `--ref-audio` shape).
|
||||
|
||||
### Prerequisites
|
||||
|
||||
For ROCm builds, replace `onnxruntime` with `onnxruntime-rocm`:
|
||||
|
||||
```bash
|
||||
pip uninstall onnxruntime
|
||||
pip install onnxruntime-rocm
|
||||
```
|
||||
|
||||
### Task variants
|
||||
|
||||
- `CustomVoice`: predefined speaker (speaker ID) with optional style instruction.
|
||||
- `VoiceDesign`: text + descriptive instruction designs a new voice.
|
||||
- `Base`: voice cloning from reference audio + transcript.
|
||||
@@ -415,22 +461,27 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty
|
||||
```
|
||||
|
||||
### Streaming
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
--query-type CustomVoice \
|
||||
--streaming \
|
||||
--output-dir /tmp/out_stream
|
||||
```
|
||||
|
||||
Streaming requires `async_chunk: true` in the deploy config.
|
||||
|
||||
### Word Timestamps
|
||||
|
||||
Word-level timestamps are currently a serving-path feature: launch
|
||||
`vllm-omni serve` with `--forced-aligner` and request `word_timestamps`
|
||||
(see `examples/online_serving/text_to_speech/README.md`). An offline
|
||||
example is not provided in this release.
|
||||
|
||||
### Batched decoding
|
||||
|
||||
The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, create an overlay based on `vllm_omni/deploy/qwen3_tts.yaml` and pass it with `--deploy-config`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
--query-type CustomVoice \
|
||||
@@ -438,9 +489,11 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
--batch-size 4 \
|
||||
--deploy-config /path/to/qwen3_tts_batched.yaml
|
||||
```
|
||||
|
||||
`--batch-size` must match a CUDA-graph capture size (1, 2, 4, 8, 16…).
|
||||
|
||||
### Notes
|
||||
|
||||
- Run `--help` for the full argument surface.
|
||||
- See `qwen3_tts/end2end.py` for the prompt-length-estimation logic the Talker uses.
|
||||
|
||||
@@ -451,6 +504,7 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
Single-stage native AR TTS at 22.05 kHz. Pipeline: `Qwen3.5 backbone → 32 FSQ codebook heads (one frame/step) → NeMo NanoCodec`. Backbone runs under vLLM paged attention; the 32-head sampling + learned embedding feedback ride the native-AR runner hooks. `enforce_eager` (CUDA graph is a perf follow-up).
|
||||
|
||||
### Prerequisites
|
||||
|
||||
The NanoCodec decoder needs NeMo (NVIDIA Open Model License), installed separately:
|
||||
|
||||
```bash
|
||||
@@ -480,15 +534,18 @@ failure rather than an environment one. The PyTorch sampler it falls back to cos
|
||||
here.
|
||||
|
||||
### Quick start (zero-shot, default voice)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/gepard/end2end.py \
|
||||
--text "Hello, this is a Gepard demo."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Reference-audio cloning (`ref_compressor` speaker prefix) is a follow-up PR.
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 22.05 kHz mono WAV.
|
||||
- First run downloads two checkpoints, not one: the model itself, and the NanoCodec decoder it names in `codec_id` (`nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps` by default). Both are fetched automatically; expect the first startup to be correspondingly slower.
|
||||
- Offline and online: `examples/offline_inference/text_to_speech/gepard/end2end.py` and `/v1/audio/speech` (see the [online Gepard section](../../online_serving/text_to_speech/README.md#gepard-10)).
|
||||
@@ -502,6 +559,7 @@ Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Ref
|
||||
Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
pip install voxcpm
|
||||
# or, for a local source checkout:
|
||||
@@ -509,6 +567,7 @@ export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/voxcpm
|
||||
```
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
||||
--model openbmb/VoxCPM2 \
|
||||
@@ -516,7 +575,9 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
Pass a reference audio for isolated cloning, or both `--ref-audio` + `--ref-text` for prompt continuation:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
||||
--text "Hello, this is a voice clone demo." \
|
||||
@@ -525,39 +586,16 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
||||
```
|
||||
|
||||
### Streaming
|
||||
|
||||
Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [`examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py`](../../online_serving/text_to_speech/voxcpm2/gradio_demo.py) for an AudioWorklet-based gapless streaming player; the offline `end2end.py` script does not expose a streaming path.
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 48 kHz mono WAV.
|
||||
- Deploy config: `vllm_omni/deploy/voxcpm2.yaml` (auto-loaded by HF `model_type`).
|
||||
|
||||
---
|
||||
|
||||
## dots.tts
|
||||
|
||||
Single-stage native AR TTS at 48 kHz (rednote-hilab). Pipeline: `Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR loopback → AudioVAE (streaming decode)`. Same "vLLM-native base LM + side-path computation" pattern as VoxCPM2, with a plain Qwen2 backbone instead of MiniCPM4 and no FSQ / residual LM stage.
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
|
||||
--model dots-studio/dots.tts-soar \
|
||||
--text "Hello, this is a test of dots TTS running on vLLM Omni."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
Not wired in this release. Generation is text-only and does not consume reference audio, a named voice, or a precomputed speaker embedding. The CAM++ x-vector speaker encoder weights load, but `end2end.py` has no `--ref-audio`/`--ref-text` flags yet.
|
||||
|
||||
### Streaming
|
||||
The AudioVAE decoder emits incremental audio through its internal streaming path (`init_stream_state` / `stream_step` / `stream_flush`). Online streaming is exposed through the OpenAI-compatible `/v1/audio/speech` endpoint with `stream=true`; see the [online serving guide](../../online_serving/text_to_speech/README.md#dotstts). The offline `end2end.py` script returns the consolidated waveform.
|
||||
|
||||
### Notes
|
||||
- Output: 48 kHz mono WAV.
|
||||
- Deploy config: `vllm_omni/deploy/dots_tts.yaml` (auto-loaded by HF `model_type`).
|
||||
- Checkpoints: `dots-studio/dots.tts-soar` is the validated default. `dots.tts-base` shares the same architecture but is unvalidated in this repo. `dots.tts-mf` (MeanFlow, 2-4 step) is not supported yet.
|
||||
- Known limitation: no CUDA graph capture and no batched side-path computation yet, so concurrent requests do not currently scale (each request's DiT Euler steps run serially). See `recipes/rednote-hilab/dots.tts.md` for details and the roadmap.
|
||||
|
||||
---
|
||||
|
||||
## IndexTTS-2 and IndexTTS-2.5
|
||||
|
||||
Both versions use a 2-stage TTS pipeline and produce 22.05 kHz mono audio. Stage 0 is a GPT AR talker. IndexTTS-2 Stage 1 uses RepCodec semantic embeddings, GPT latent, S2Mel CFM/DiT, and BigVGAN. IndexTTS-2.5 replaces RepCodec with EnhancedCodec and uses the official code-only path (`use_gpt_latent=false`). Every request requires reference audio for zero-shot voice cloning.
|
||||
@@ -592,6 +630,7 @@ python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
```
|
||||
|
||||
### Emotion control
|
||||
|
||||
```bash
|
||||
# Emotion from reference audio
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
@@ -616,6 +655,7 @@ python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- `--ref-audio` is **required** — neither version provides text-only synthesis.
|
||||
- IndexTTS-2.5 Stage 0 (AR Talker) generates mel codes from text + reference audio using plain vLLM sampling. It does not reproduce the official default `num_beams=3` beam search. For parity comparisons, run upstream with `num_beams=1`; output quality can differ from the official beam-search result.
|
||||
- IndexTTS-2.5 accepts language codes such as `zh`, `en`, `zhen` (mixed Chinese/English), `ja`, and `yue`. `Mandarin` is a vLLM-Omni convenience alias for `zh`. Japanese (`ja`) uses `fugashi` tokenization and produces audio, but it does not automatically expand numbers, dates, or percentages; write those inputs as readable Japanese text before calling the model. Use `--no-text-normalization` only when the input is already normalized.
|
||||
@@ -631,12 +671,15 @@ python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
Voxtral-4B-TTS (Mistral). Has its own argparse surface; uses voice presets and the `mistral_common` `SpeechRequest` protocol.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Latest `mistral_common` with `SpeechRequest` support:
|
||||
|
||||
```bash
|
||||
pip install -e /path/to/mistral-common # or upgrade from PyPI when available
|
||||
```
|
||||
|
||||
### Quick start (voice preset)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
--write-audio --voice cheerful_female \
|
||||
@@ -645,6 +688,7 @@ python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
```
|
||||
|
||||
### Voice cloning (capability gated upstream)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
--write-audio \
|
||||
@@ -654,12 +698,14 @@ python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
```
|
||||
|
||||
### Streaming + concurrency
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
--num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "..."
|
||||
```
|
||||
|
||||
Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
|
||||
|
||||
### Notes
|
||||
|
||||
@@ -1,125 +0,0 @@
|
||||
"""Offline dots.tts inference example (single-stage native AR pipeline).
|
||||
|
||||
Uses the single-stage native AR config (vllm_omni/deploy/dots_tts.yaml).
|
||||
Reference checkpoint: dots-studio/dots.tts-soar.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
SAMPLE_RATE = 48_000
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Offline dots.tts native AR inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="dots-studio/dots.tts-soar",
|
||||
help="dots.tts model path or HuggingFace repo ID.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
default="Hello, this is a test of dots TTS running on vLLM Omni.",
|
||||
help="Text to synthesize.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="output_audio",
|
||||
help="Directory for output WAV files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override the deploy config path. If unset, auto-loads "
|
||||
"vllm_omni/deploy/dots_tts.yaml based on the HF model_type.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def extract_audio(multimodal_output: dict) -> torch.Tensor:
|
||||
"""Extract the final complete audio tensor from multimodal output.
|
||||
|
||||
The output processor concatenates per-step delta tensors under
|
||||
``model_outputs``. Falls back to ``audio`` for backwards compat.
|
||||
"""
|
||||
audio = multimodal_output.get("model_outputs")
|
||||
if audio is None:
|
||||
audio = multimodal_output.get("audio")
|
||||
if audio is None:
|
||||
raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}")
|
||||
|
||||
if isinstance(audio, list):
|
||||
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
|
||||
if not valid:
|
||||
raise ValueError("Audio list is empty or all elements are None.")
|
||||
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
|
||||
|
||||
return torch.as_tensor(audio).float().cpu().reshape(-1)
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
engine = Omni(
|
||||
model=args.model,
|
||||
deploy_config=args.deploy_config,
|
||||
)
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import (
|
||||
build_dots_tts_prompt,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
|
||||
|
||||
prompt = build_dots_tts_prompt(
|
||||
tokenizer=tokenizer,
|
||||
text=args.text,
|
||||
)
|
||||
|
||||
print(f"Model : {args.model}")
|
||||
print(f"Text : {args.text}")
|
||||
print(f"Output dir : {output_dir}")
|
||||
print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens")
|
||||
|
||||
t_start = time.perf_counter()
|
||||
outputs = engine.generate([prompt])
|
||||
elapsed = time.perf_counter() - t_start
|
||||
|
||||
request_output = outputs[0]
|
||||
mm = request_output.outputs[0].multimodal_output
|
||||
audio = extract_audio(mm)
|
||||
|
||||
duration = audio.numel() / SAMPLE_RATE
|
||||
rtf = elapsed / duration if duration > 0 else float("inf")
|
||||
|
||||
output_path = output_dir / "output.wav"
|
||||
sf.write(str(output_path), audio.numpy(), SAMPLE_RATE, format="WAV")
|
||||
|
||||
print(f"Saved : {output_path}")
|
||||
print(f"Duration : {duration:.2f}s")
|
||||
print(f"Inference : {elapsed:.2f}s")
|
||||
print(f"RTF : {rtf:.3f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
main()
|
||||
@@ -1,97 +0,0 @@
|
||||
# Dynin-Omni: Online serving
|
||||
|
||||
## Installation
|
||||
|
||||
Please refer to [README.md](../../../README.md).
|
||||
|
||||
## Launch the Server
|
||||
|
||||
First, find the `transformers_modules` path:
|
||||
|
||||
```bash
|
||||
python - <<'PY'
|
||||
from transformers.utils.hub import HF_MODULES_CACHE
|
||||
print(HF_MODULES_CACHE)
|
||||
PY
|
||||
```
|
||||
|
||||
Then export it for both `PYTHONPATH` and `HF_MODULES_CACHE`:
|
||||
|
||||
```bash
|
||||
export PYTHONPATH=<transformers_modules_path>:$PYTHONPATH
|
||||
export HF_MODULES_CACHE=<transformers_modules_path>
|
||||
```
|
||||
|
||||
Run from repository root. The 3-stage pipeline is resolved automatically from
|
||||
the model type (`vllm_omni/deploy/dynin_omni.yaml`), so no config path is
|
||||
required. Pass `--deploy-config` only to override the defaults:
|
||||
|
||||
```bash
|
||||
vllm-omni serve snu-aidas/Dynin-Omni \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
If `vllm-omni` is not in PATH, run:
|
||||
|
||||
```bash
|
||||
PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
Wait until the server logs show both `All stages initialized successfully` and
|
||||
`Application startup complete.` before sending requests.
|
||||
|
||||
## Send Requests via Python Client
|
||||
|
||||
Move to the example directory:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/dynin_omni
|
||||
```
|
||||
|
||||
### Text -> Image
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type t2i \
|
||||
--prompt "A realistic indoor living room with natural daylight."
|
||||
```
|
||||
|
||||
### Image -> Image
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type i2i \
|
||||
--image-path ../../offline_inference/dynin_omni/data/image/sofa_under_water.jpg \
|
||||
--prompt "Transform this surreal underwater setting into a realistic indoor living room while preserving the sofa layout."
|
||||
```
|
||||
|
||||
### Text -> Speech
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type t2s \
|
||||
--prompt "Hello. This is Dynin-omni."
|
||||
```
|
||||
|
||||
## CLI Arguments
|
||||
|
||||
- `--query-type` (`t2i|t2s|i2i`)
|
||||
- `--model` (default: `snu-aidas/Dynin-Omni`)
|
||||
- `--host` / `--port` (OpenAI-compatible vLLM endpoint)
|
||||
- `--prompt` (custom text)
|
||||
- `--image-path` (required for `i2i`)
|
||||
- `--modalities` (optional output modalities override)
|
||||
- `--output-dir` (default: `/tmp/dynin_online_outputs`)
|
||||
|
||||
## Notes
|
||||
|
||||
- This client currently supports only `t2i`, `t2s`, and `i2i`.
|
||||
- `t2t` is intentionally not exposed in this online example.
|
||||
- This example intentionally uses the OpenAI-compatible chat completion endpoint.
|
||||
- Task routing for non-text outputs relies on Dynin task trigger tokens (`<|t2i|>`, `<|i2i|>`, `<|t2s|>`) injected by the client.
|
||||
- Outputs are saved under `/tmp/dynin_online_outputs` by default.
|
||||
- Dynin stage-0 warmup can take a while on first startup; do not send requests before startup completes.
|
||||
- Dynin itself can execute text-returning tasks such as `t2t`, `s2t`, `i2t`, and `v2t`, but this online serving example currently runs stage-0 in `generation` mode. In that path, the generation worker does not surface the final text as `output.text`, so OpenAI chat responses for those text-output tasks may complete internally but still return empty text.
|
||||
-342
@@ -1,342 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
DEFAULT_MODEL = "snu-aidas/Dynin-Omni"
|
||||
DEFAULT_OUTPUT_DIR = "/tmp/dynin_online_outputs"
|
||||
|
||||
QUERY_CHOICES = ("t2i", "t2s", "i2i")
|
||||
DEFAULT_PROMPT_BY_QUERY = {
|
||||
"t2i": "A high quality detailed living room interior photo.",
|
||||
"t2s": "Please read this sentence naturally: Hello from Dynin-Omni online serving.",
|
||||
"i2i": "Transform this image into a realistic indoor living room while preserving layout.",
|
||||
}
|
||||
DEFAULT_MODALITIES_BY_QUERY = {
|
||||
"t2i": ["image"],
|
||||
"t2s": ["audio"],
|
||||
"i2i": ["image"],
|
||||
}
|
||||
OFFLINE_PARITY_STAGE_COUNT = 3
|
||||
OFFLINE_PARITY_STAGE_SAMPLING = {
|
||||
"max_tokens": 1,
|
||||
"temperature": 0.0,
|
||||
"top_p": 1.0,
|
||||
"detokenize": False,
|
||||
}
|
||||
|
||||
|
||||
def _infer_mime_type(path: Path) -> str:
|
||||
mime_type, _ = mimetypes.guess_type(str(path))
|
||||
return mime_type or "application/octet-stream"
|
||||
|
||||
|
||||
def _encode_file_as_data_url(path: Path) -> str:
|
||||
mime_type = _infer_mime_type(path)
|
||||
raw = path.read_bytes()
|
||||
encoded = base64.b64encode(raw).decode("utf-8")
|
||||
return f"data:{mime_type};base64,{encoded}"
|
||||
|
||||
|
||||
def _to_image_url(path_or_url: str) -> str:
|
||||
value = str(path_or_url)
|
||||
if value.startswith(("http://", "https://", "data:image/")):
|
||||
return value
|
||||
path = Path(value).expanduser().resolve()
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Image file not found: {path}")
|
||||
return _encode_file_as_data_url(path)
|
||||
|
||||
|
||||
def _build_user_content(query_type: str, prompt: str, image_path: str | None) -> list[dict[str, Any]]:
|
||||
if query_type == "t2i":
|
||||
return [{"type": "text", "text": f"<|t2i|> {prompt}"}]
|
||||
|
||||
if query_type == "t2s":
|
||||
return [{"type": "text", "text": f"<|t2s|> {prompt}"}]
|
||||
|
||||
if query_type == "i2i":
|
||||
if not image_path:
|
||||
raise ValueError("--image-path is required for query type i2i")
|
||||
return [
|
||||
{"type": "text", "text": f"<|i2i|> {prompt}"},
|
||||
{"type": "image_url", "image_url": {"url": _to_image_url(image_path)}},
|
||||
]
|
||||
|
||||
raise ValueError(f"Unsupported query_type: {query_type}")
|
||||
|
||||
|
||||
def _collect_text_from_content(content: Any) -> list[str]:
|
||||
texts: list[str] = []
|
||||
if isinstance(content, str):
|
||||
stripped = content.strip()
|
||||
if stripped:
|
||||
texts.append(stripped)
|
||||
return texts
|
||||
|
||||
if isinstance(content, dict):
|
||||
for key in ("text", "content", "value", "output_text"):
|
||||
text_value = content.get(key)
|
||||
if isinstance(text_value, str) and text_value.strip():
|
||||
texts.append(text_value.strip())
|
||||
return texts
|
||||
|
||||
if isinstance(content, list):
|
||||
for item in content:
|
||||
texts.extend(_collect_text_from_content(item))
|
||||
return texts
|
||||
|
||||
content_text = getattr(content, "text", None)
|
||||
if isinstance(content_text, str) and content_text.strip():
|
||||
texts.append(content_text.strip())
|
||||
content_value = getattr(content, "content", None)
|
||||
if isinstance(content_value, str) and content_value.strip():
|
||||
texts.append(content_value.strip())
|
||||
output_text = getattr(content, "output_text", None)
|
||||
if isinstance(output_text, str) and output_text.strip():
|
||||
texts.append(output_text.strip())
|
||||
return texts
|
||||
|
||||
|
||||
def _extract_text_outputs(chat_completion: Any) -> list[str]:
|
||||
texts: list[str] = []
|
||||
for choice in getattr(chat_completion, "choices", []) or []:
|
||||
message = getattr(choice, "message", None)
|
||||
if message is None:
|
||||
continue
|
||||
content = getattr(message, "content", None)
|
||||
texts.extend(_collect_text_from_content(content))
|
||||
reasoning_content = getattr(message, "reasoning_content", None)
|
||||
if isinstance(reasoning_content, str) and reasoning_content.strip():
|
||||
texts.append(reasoning_content.strip())
|
||||
choice_text = getattr(choice, "text", None)
|
||||
if isinstance(choice_text, str) and choice_text.strip():
|
||||
texts.append(choice_text.strip())
|
||||
top_level_output_text = getattr(chat_completion, "output_text", None)
|
||||
if isinstance(top_level_output_text, str) and top_level_output_text.strip():
|
||||
texts.append(top_level_output_text.strip())
|
||||
return texts
|
||||
|
||||
|
||||
def _extract_image_data_urls(chat_completion: Any) -> list[str]:
|
||||
urls: list[str] = []
|
||||
for choice in getattr(chat_completion, "choices", []) or []:
|
||||
message = getattr(choice, "message", None)
|
||||
if message is None:
|
||||
continue
|
||||
content = getattr(message, "content", None)
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
for item in content:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
if item.get("type") != "image_url":
|
||||
continue
|
||||
image_url = (item.get("image_url") or {}).get("url")
|
||||
if isinstance(image_url, str) and image_url.startswith("data:image"):
|
||||
urls.append(image_url)
|
||||
return urls
|
||||
|
||||
|
||||
def _extract_audio_payloads(chat_completion: Any) -> list[bytes]:
|
||||
payloads: list[bytes] = []
|
||||
for choice in getattr(chat_completion, "choices", []) or []:
|
||||
message = getattr(choice, "message", None)
|
||||
if message is None:
|
||||
continue
|
||||
message_audio = getattr(message, "audio", None)
|
||||
if message_audio is None:
|
||||
continue
|
||||
data_b64 = getattr(message_audio, "data", None)
|
||||
if isinstance(data_b64, str) and data_b64:
|
||||
try:
|
||||
payloads.append(base64.b64decode(data_b64))
|
||||
except Exception:
|
||||
continue
|
||||
return payloads
|
||||
|
||||
|
||||
def _decode_data_url(data_url: str) -> tuple[bytes, str]:
|
||||
header, data = data_url.split(",", 1)
|
||||
mime_type = "image/png"
|
||||
if ";" in header and ":" in header:
|
||||
mime_type = header.split(":", 1)[1].split(";", 1)[0]
|
||||
return base64.b64decode(data), mime_type
|
||||
|
||||
|
||||
def _image_extension_from_mime(mime_type: str) -> str:
|
||||
if mime_type == "image/jpeg":
|
||||
return ".jpg"
|
||||
if mime_type == "image/webp":
|
||||
return ".webp"
|
||||
if mime_type == "image/gif":
|
||||
return ".gif"
|
||||
return ".png"
|
||||
|
||||
|
||||
def _save_outputs(
|
||||
*,
|
||||
query_type: str,
|
||||
chat_completion: Any,
|
||||
output_dir: Path,
|
||||
) -> None:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
stamp = time.strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
text_outputs = _extract_text_outputs(chat_completion)
|
||||
image_data_urls = _extract_image_data_urls(chat_completion)
|
||||
audio_payloads = _extract_audio_payloads(chat_completion)
|
||||
|
||||
if text_outputs:
|
||||
text_path = output_dir / f"{query_type}_{stamp}.txt"
|
||||
text_path.write_text("\n\n".join(text_outputs) + "\n", encoding="utf-8")
|
||||
print(f"[dynin-online] text saved: {text_path}")
|
||||
print(text_outputs[0])
|
||||
|
||||
for idx, image_url in enumerate(image_data_urls):
|
||||
image_bytes, mime_type = _decode_data_url(image_url)
|
||||
ext = _image_extension_from_mime(mime_type)
|
||||
image_path = output_dir / f"{query_type}_{stamp}_{idx}{ext}"
|
||||
image_path.write_bytes(image_bytes)
|
||||
print(f"[dynin-online] image saved: {image_path}")
|
||||
|
||||
for idx, audio_bytes in enumerate(audio_payloads):
|
||||
audio_path = output_dir / f"{query_type}_{stamp}_{idx}.wav"
|
||||
audio_path.write_bytes(audio_bytes)
|
||||
print(f"[dynin-online] audio saved: {audio_path}")
|
||||
|
||||
if not text_outputs and not image_data_urls and not audio_payloads:
|
||||
print("[dynin-online] no output extracted from response")
|
||||
raw_path = output_dir / f"{query_type}_{stamp}_raw_response.json"
|
||||
try:
|
||||
if hasattr(chat_completion, "model_dump_json"):
|
||||
serialized = chat_completion.model_dump_json(indent=2)
|
||||
else:
|
||||
if hasattr(chat_completion, "model_dump"):
|
||||
raw_payload: Any = chat_completion.model_dump(mode="json")
|
||||
else:
|
||||
raw_payload = chat_completion
|
||||
try:
|
||||
serialized = json.dumps(raw_payload, ensure_ascii=False, indent=2)
|
||||
except Exception:
|
||||
serialized = json.dumps({"repr": repr(raw_payload)}, ensure_ascii=False, indent=2)
|
||||
raw_path.write_text(serialized + "\n", encoding="utf-8")
|
||||
print(f"[dynin-online] raw response saved: {raw_path}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _build_offline_parity_sampling_params_list() -> list[dict[str, Any]]:
|
||||
return [dict(OFFLINE_PARITY_STAGE_SAMPLING) for _ in range(OFFLINE_PARITY_STAGE_COUNT)]
|
||||
|
||||
|
||||
def run_request(args: argparse.Namespace) -> None:
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(
|
||||
api_key="EMPTY",
|
||||
base_url=f"http://{args.host}:{args.port}/v1",
|
||||
)
|
||||
prompt = args.prompt.strip() if args.prompt else DEFAULT_PROMPT_BY_QUERY[args.query_type]
|
||||
user_content = _build_user_content(
|
||||
query_type=args.query_type,
|
||||
prompt=prompt,
|
||||
image_path=args.image_path,
|
||||
)
|
||||
if args.modalities:
|
||||
modalities = [item.strip() for item in args.modalities.split(",") if item.strip()]
|
||||
else:
|
||||
modalities = DEFAULT_MODALITIES_BY_QUERY[args.query_type]
|
||||
|
||||
extra_body = {
|
||||
"sampling_params_list": _build_offline_parity_sampling_params_list(),
|
||||
}
|
||||
chat_completion = client.chat.completions.create(
|
||||
model=args.model,
|
||||
messages=[{"role": "user", "content": user_content}],
|
||||
modalities=modalities,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
_save_outputs(
|
||||
query_type=args.query_type,
|
||||
chat_completion=chat_completion,
|
||||
output_dir=Path(args.output_dir).expanduser(),
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Dynin-Omni online chat completion client")
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
type=str,
|
||||
default="t2i",
|
||||
choices=QUERY_CHOICES,
|
||||
help="Dynin query type",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
"-m",
|
||||
type=str,
|
||||
default=DEFAULT_MODEL,
|
||||
help="Model name/path",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--host",
|
||||
type=str,
|
||||
default="localhost",
|
||||
help="Host/IP of the vLLM Omni API server",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=8091,
|
||||
help="Port of the vLLM Omni API server",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
"-p",
|
||||
type=str,
|
||||
default="",
|
||||
help="Custom prompt text",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image-path",
|
||||
"-i",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Image path/URL for i2i",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--modalities",
|
||||
type=str,
|
||||
default="",
|
||||
help="Comma-separated output modalities override (e.g., text,image,audio)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
"-o",
|
||||
type=str,
|
||||
default=DEFAULT_OUTPUT_DIR,
|
||||
help="Directory to save outputs",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
|
||||
run_request(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -27,7 +27,6 @@ For the full list of supported architectures across all modalities, see
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | ✓ | — | — | — |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
|
||||
| dots.tts | `dots-studio/dots.tts-soar` | — (text-only) | ✓ (PCM stream) | — (default placeholder only) | — |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
|
||||
|
||||
CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](../../offline_inference/text_to_speech/README.md#cosyvoice3) instead.
|
||||
@@ -143,42 +142,6 @@ python examples/online_serving/text_to_speech/gepard/speech_client.py \
|
||||
- Concurrent requests at `max_num_seqs: 4` are supported. Native-AR recompute preemption is a known limitation of this architecture (a request that is preempted mid-generation can resume incorrectly); keep concurrency at or below `max_num_seqs` and treat preemption as out of scope until the platform fix lands.
|
||||
- Optional comparison against the upstream Gepard reference server needs Blackwell/Hopper + CUDA 13 + Postgres and is not part of CI.
|
||||
|
||||
## dots.tts
|
||||
|
||||
Single-stage, text-only TTS at 48 kHz. The current integration supports
|
||||
no-reference synthesis; voice cloning, reference audio, named voices, and
|
||||
precomputed speaker embeddings are not supported yet.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve dots-studio/dots.tts-soar --omni --trust-remote-code --port 8091
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
# Non-streaming WAV
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, this is a dots TTS online-serving test.",
|
||||
"voice": "default",
|
||||
"response_format": "wav"
|
||||
}' --output output.wav
|
||||
|
||||
# Streaming 48 kHz mono PCM
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, this is a streaming dots TTS test.",
|
||||
"voice": "default",
|
||||
"stream": true,
|
||||
"stream_format": "audio",
|
||||
"response_format": "pcm"
|
||||
}' --no-buffer | play -t raw -r 48000 -e signed -b 16 -c 1 -
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## GLM-TTS
|
||||
|
||||
@@ -1,166 +0,0 @@
|
||||
# dots.tts
|
||||
|
||||
> Continuous-AR TTS at 48 kHz (rednote-hilab)
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: rednote-hilab
|
||||
- Model: `dots-studio/dots.tts-soar`
|
||||
- Task: Text-to-speech, zero-shot synthesis only
|
||||
- Mode: Offline inference and OpenAI-compatible online serving
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe as a known-good starting point for running
|
||||
`dots-studio/dots.tts-soar` on vLLM-Omni on consumer-class GPUs.
|
||||
dots.tts is a ~1.7B-parameter continuous-AR TTS model (Qwen2.5-1.5B base LM
|
||||
+ 344M DiT flow-matching head + 180M AudioVAE) that emits 48 kHz mono audio.
|
||||
It follows the same "vLLM-native base LM + side-path computation" pattern as
|
||||
VoxCPM2 — single-stage pipeline
|
||||
`Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR
|
||||
loopback → AudioVAE (streaming decode)` — with a plain Qwen2 backbone
|
||||
instead of MiniCPM4, and no FSQ / residual-LM stage.
|
||||
|
||||
This is an early integration; review the [Known limitations](#known-limitations)
|
||||
before deploying it in production.
|
||||
|
||||
## References
|
||||
|
||||
- Offline end-to-end script:
|
||||
[`examples/offline_inference/text_to_speech/dots_tts/end2end.py`](../../examples/offline_inference/text_to_speech/dots_tts/end2end.py)
|
||||
- Example guide:
|
||||
[`examples/offline_inference/text_to_speech/README.md`](../../examples/offline_inference/text_to_speech/README.md#dotstts)
|
||||
- Online serving guide:
|
||||
[`examples/online_serving/text_to_speech/README.md`](../../examples/online_serving/text_to_speech/README.md#dotstts)
|
||||
- Default deploy config:
|
||||
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
|
||||
- Talker / pipeline source:
|
||||
[`vllm_omni/model_executor/models/dots_tts/`](../../vllm_omni/model_executor/models/dots_tts/)
|
||||
- Upstream: [rednote-hilab/dots.tts](https://github.com/rednote-hilab/dots.tts)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
This recipe documents one tested 16 GB consumer-GPU configuration. Other
|
||||
vendor sections (ROCm, NPU) and larger-VRAM configurations are welcome as
|
||||
community validation lands.
|
||||
|
||||
## GPU
|
||||
|
||||
### 1 x RTX 5080 16GB (Single GPU, Minimum Recommended)
|
||||
|
||||
dots.tts (~1.7B params across the base LM + DiT + AudioVAE + CAM++
|
||||
speaker encoder, bfloat16) fits comfortably on a single 16 GB GPU. The
|
||||
bundled default config at
|
||||
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
|
||||
(`gpu_memory_utilization: 0.8`, `max_num_seqs: 4`, `enforce_eager: true`,
|
||||
`enable_prefix_caching: false`) loads cleanly with ~5.1 GiB for model
|
||||
weights and ~0.3 GiB peak activation; the remainder of the configured
|
||||
budget is available for KV cache. Total resident footprint at idle is
|
||||
roughly **7-8 GiB / 16 GB** — the only tight spot in the full CUDA-Graph
|
||||
roadmap would be step 8's graph capture (not implemented yet; this
|
||||
release runs `enforce_eager: true`, so it doesn't apply today).
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux (WSL2)
|
||||
- Python: 3.12
|
||||
- Driver / runtime: NVIDIA driver 595.95
|
||||
- torch: 2.11.0+cu130
|
||||
- vLLM: 0.26.0
|
||||
- vLLM-Omni: 0.22.1.dev (current `main`)
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
|
||||
--model dots-studio/dots.tts-soar \
|
||||
--text "Hello, this is a test of dots TTS running on vLLM Omni."
|
||||
```
|
||||
|
||||
The deploy config at
|
||||
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
|
||||
is loaded automatically by the model registry (HF `model_type=dots_tts`).
|
||||
Pass `--deploy-config <path>` to override.
|
||||
|
||||
#### Verification
|
||||
|
||||
**T1 — offline zero-shot synthesis**:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
|
||||
--model dots-studio/dots.tts-soar \
|
||||
--text "Hello, this is a test of dots TTS running on vLLM Omni."
|
||||
```
|
||||
|
||||
Observed: `output_audio/output.wav`, 3.52 s @ 48 kHz mono. Single
|
||||
one-shot process (init → one `generate()` → exit), so the reported
|
||||
numbers include engine init and first-request warmup, not just steady-state
|
||||
per-step throughput — same caveat as VoxCPM2's recipe. `Inference: 5.52s`,
|
||||
`RTF: 1.569`.
|
||||
|
||||
Weight-loading breakdown from the same run (all tensors matched, no
|
||||
missing/extra keys): 951/951 AudioVAE, 244/244 DiT, 270/270
|
||||
patch_encoder, 198 Qwen2, 938/938 CAM++ speaker encoder.
|
||||
|
||||
Whisper transcription of the output matched the input text with no
|
||||
dropped leading word (confirms the streaming-vocoder patch-boundary fix
|
||||
described in [Known limitations](#known-limitations)).
|
||||
|
||||
**T2 — online text-only synthesis**:
|
||||
|
||||
```bash
|
||||
vllm serve dots-studio/dots.tts-soar --omni --trust-remote-code --port 8091
|
||||
```
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, this is a test of dots TTS online serving.",
|
||||
"voice": "default",
|
||||
"response_format": "wav"
|
||||
}' --output output.wav
|
||||
```
|
||||
|
||||
The endpoint also supports raw streaming audio with `stream=true`,
|
||||
`stream_format="audio"`, and `response_format="pcm"`.
|
||||
|
||||
#### Notes
|
||||
|
||||
- Output: 48 kHz mono WAV.
|
||||
- Checkpoints: `dots-studio/dots.tts-soar` is the validated default
|
||||
used throughout this recipe. `rednote-hilab/dots.tts-base` shares the
|
||||
same architecture but is unvalidated in this repo. `rednote-hilab/dots.tts-mf`
|
||||
(MeanFlow, 2-4 step) is not supported — see below.
|
||||
- `enforce_eager: true` and `enable_prefix_caching: false` in the deploy
|
||||
config are load-bearing, not just conservative defaults: with prefix
|
||||
caching enabled, vLLM-Omni's prefix-cache multimodal-output merge path
|
||||
does not preserve the `sparse_audio` marker this model relies on to
|
||||
route audio output correctly, and generation silently truncates to a
|
||||
single ~160 ms patch. Do not override `enable_prefix_caching` for this
|
||||
model until that framework-level gap is fixed.
|
||||
|
||||
## Known limitations
|
||||
|
||||
- **Voice cloning is not wired.** The CAM++ x-vector speaker encoder
|
||||
weights load and are exercised by `load_weights()`, but nothing in the
|
||||
prompt builder consumes reference audio or a precomputed speaker embedding
|
||||
yet. Both offline and online generation are text-only; `ref_audio`,
|
||||
`ref_text`, named voices, and speaker embeddings are not supported.
|
||||
- **`dots.tts-mf` (MeanFlow, 2-4 step) checkpoint is not supported.** Only
|
||||
the fixed 10-step Euler DiT sampler used by `dots.tts-soar` /
|
||||
`dots.tts-base` is implemented.
|
||||
- **No CUDA graph capture.** The talker runs fully eager. voxcpm2's three
|
||||
captured graphs (base LM decode, CFM solver, VAE decode) have no
|
||||
dots.tts equivalent yet.
|
||||
- **Concurrent requests do not scale.** Each request's 10-step DiT Euler
|
||||
integration runs serially in the side path (no cross-request batching,
|
||||
unlike voxcpm2's `enable_batched_cfm`). A community review of this
|
||||
integration measured no throughput gain at `c=4` concurrent requests
|
||||
versus `c=1`.
|
||||
- **`SamplingParams.seed` does not control audio-generation randomness.**
|
||||
The DiT's flow-matching noise is deterministically derived per-request
|
||||
from a fixed internal seed (reproducible run-to-run, matching
|
||||
voxcpm2's `deterministic_cfm_seed` convention) rather than from the
|
||||
caller-supplied `seed` field — the same limitation voxcpm2 has today.
|
||||
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 28,
|
||||
EnvironmentVariableCategory.INHERITED_VLLM: 20,
|
||||
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 67,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
|
||||
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
|
||||
EnvironmentVariableCategory.INTERNAL: 2,
|
||||
}
|
||||
@@ -178,10 +178,10 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
|
||||
)
|
||||
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
|
||||
ModelEnvironmentVariableDisposition.PROMOTE: 38,
|
||||
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
|
||||
ModelEnvironmentVariableDisposition.PROMOTE: 31,
|
||||
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
|
||||
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
|
||||
ModelEnvironmentVariableDisposition.INTERNALIZE: 16,
|
||||
ModelEnvironmentVariableDisposition.INTERNALIZE: 14,
|
||||
ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7,
|
||||
}
|
||||
|
||||
|
||||
@@ -272,7 +272,7 @@ def test_from_pipeline_config_normalizes_stage_engine_extras_without_expanding_s
|
||||
def test_frontend_log_stats_flag_is_not_an_unowned_stage_argument(disabled):
|
||||
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
||||
|
||||
config = _from_pipeline_key("dots_tts", cli_overrides={"disable_log_stats": disabled})
|
||||
config = _from_pipeline_key("voxcpm2", cli_overrides={"disable_log_stats": disabled})
|
||||
assert config.stage_configs
|
||||
engine_args = build_engine_args_dict_from_omni_stage_config(config.stage_by_id(0), model="test-model")
|
||||
assert "disable_log_stats" not in engine_args
|
||||
|
||||
@@ -1,129 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""E2E tests for dots.tts native AR offline inference."""
|
||||
|
||||
from collections.abc import Mapping
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
DOTS_TTS_MODEL = "rednote-hilab/dots.tts-soar"
|
||||
DEPLOY_CONFIG = get_deploy_config_path("dots_tts.yaml")
|
||||
SAMPLE_RATE = 48000
|
||||
|
||||
# dots.tts ships a custom tokenizer carrying the added ``<|audio_gen_start|>``
|
||||
# token, so remote code must be explicitly enabled.
|
||||
_OMNI_RUNNER_PARAM = (DOTS_TTS_MODEL, DEPLOY_CONFIG, {"trust_remote_code": True})
|
||||
|
||||
pytestmark = pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
|
||||
|
||||
|
||||
def _build_prompt(omni_runner: OmniRunner, text: str) -> dict:
|
||||
"""Wrap ``text`` in the prefix scaffold dots.tts prefill expects.
|
||||
|
||||
Unlike most TTS models here, a bare ``{"prompt": ...}`` string does not
|
||||
work: prefill needs ``[文本]<text>[文本对应语音]<|audio_gen_start|>`` as
|
||||
real Qwen2 token ids.
|
||||
"""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import (
|
||||
build_dots_tts_prompt,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True)
|
||||
return build_dots_tts_prompt(tokenizer=tokenizer, text=text)
|
||||
|
||||
|
||||
def _extract_audio(multimodal_output: dict) -> torch.Tensor:
|
||||
"""Extract the final complete audio tensor from multimodal output."""
|
||||
assert isinstance(multimodal_output, (dict, Mapping)), f"Expected dict/Mapping, got {type(multimodal_output)}"
|
||||
|
||||
# The talker emits delta chunks; the output processor accumulates them
|
||||
# under "model_outputs" ("audio" kept for backwards compat).
|
||||
audio = multimodal_output.get("model_outputs")
|
||||
if audio is None:
|
||||
audio = multimodal_output.get("audio")
|
||||
assert audio is not None, f"No audio key, got {list(multimodal_output.keys())}"
|
||||
|
||||
if isinstance(audio, list):
|
||||
valid = [torch.as_tensor(x).float().cpu().reshape(-1) for x in audio if x is not None]
|
||||
assert valid, "No valid audio tensors in output list"
|
||||
audio = torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
|
||||
|
||||
assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}"
|
||||
return audio.float().cpu().reshape(-1)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_dots_tts_zero_shot_001(omni_runner: OmniRunner) -> None:
|
||||
"""Zero-shot synthesis produces audio of a plausible duration."""
|
||||
prompt = _build_prompt(omni_runner, "Hello, this is a test of dots TTS running on vLLM Omni.")
|
||||
|
||||
outputs = omni_runner.omni.generate([prompt])
|
||||
assert len(outputs) == 1
|
||||
|
||||
audio = _extract_audio(outputs[0].outputs[0].multimodal_output)
|
||||
duration_s = audio.shape[0] / SAMPLE_RATE
|
||||
assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s"
|
||||
assert torch.isfinite(audio).all(), "Audio contains NaN/Inf"
|
||||
assert audio.abs().max() > 1e-3, "Audio is silent"
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_dots_tts_prefill_decode_mixed_batch_002(omni_runner: OmniRunner) -> None:
|
||||
"""Mixed prefill+decode batches must keep per-request rows aligned.
|
||||
|
||||
``compute_logits`` fills ``logits[i]`` positionally from
|
||||
``_results_queue``, so every request in the batch has to push exactly one
|
||||
entry per step, prefills included. When the prefill branch pushed nothing,
|
||||
a shorter request's stop signal could be read into another request's row.
|
||||
Regression guard for that pairing (PR #4765 review).
|
||||
"""
|
||||
long_text = (
|
||||
"This is a deliberately long sentence that stays in the decode phase "
|
||||
"for many steps, so the shorter requests behind it keep entering "
|
||||
"prefill alongside it and reproduce the mixed batch pattern."
|
||||
)
|
||||
short_texts = ["Hello one.", "Hello two.", "Hello three.", "Hello four."]
|
||||
prompts = [_build_prompt(omni_runner, long_text)] + [_build_prompt(omni_runner, t) for t in short_texts]
|
||||
|
||||
outputs = omni_runner.omni.generate(prompts)
|
||||
assert len(outputs) == len(prompts)
|
||||
|
||||
for i, out in enumerate(outputs):
|
||||
audio = _extract_audio(out.outputs[0].multimodal_output)
|
||||
duration_s = audio.shape[0] / SAMPLE_RATE
|
||||
assert 0.1 < duration_s < 30.0, f"Request {i} audio duration out of range: {duration_s:.2f}s"
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_dots_tts_deterministic_flow_matching_noise_003(omni_runner: OmniRunner) -> None:
|
||||
"""The same text twice must yield the same audio.
|
||||
|
||||
The DiT flow-matching noise is drawn from a per-request generator seeded
|
||||
from ``blake2b(seed:request_key:noise_step)`` rather than the global CUDA
|
||||
RNG. Drawing from the global RNG made output non-reproducible and let
|
||||
concurrent requests perturb each other's noise streams (PR #4765 review).
|
||||
"""
|
||||
text = "Reading aloud every evening helps children build a lasting love of language."
|
||||
|
||||
first = _extract_audio(
|
||||
omni_runner.omni.generate([_build_prompt(omni_runner, text)])[0].outputs[0].multimodal_output
|
||||
)
|
||||
second = _extract_audio(
|
||||
omni_runner.omni.generate([_build_prompt(omni_runner, text)])[0].outputs[0].multimodal_output
|
||||
)
|
||||
|
||||
assert first.shape == second.shape, f"Length differs across runs: {first.shape} vs {second.shape}"
|
||||
assert torch.allclose(first, second, atol=1e-4), "Audio differs across runs; flow-matching noise is not seeded"
|
||||
@@ -1,374 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
E2E offline smoke tests for Dynin-Omni.
|
||||
|
||||
- model: "snu-aidas/Dynin-Omni"
|
||||
- stage config: vllm_omni/deploy/dynin_omni_ci.yaml
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
_DEFAULT_DYNIN_CONFIG_PATH: Path | None = None
|
||||
_DEFAULT_STAGE_CONFIG_PATH = Path(get_deploy_config_path("dynin_omni_ci.yaml"))
|
||||
|
||||
models = ["snu-aidas/Dynin-Omni"]
|
||||
stage_configs = [str(_DEFAULT_STAGE_CONFIG_PATH)]
|
||||
test_params = [(model, stage_config, {"trust_remote_code": True}) for model in models for stage_config in stage_configs]
|
||||
|
||||
DYNIN_CONFIG_PATH = str(_DEFAULT_DYNIN_CONFIG_PATH) if _DEFAULT_DYNIN_CONFIG_PATH is not None else None
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.slow,
|
||||
pytest.mark.omni,
|
||||
pytest.mark.parametrize("omni_runner", test_params, indirect=True),
|
||||
]
|
||||
|
||||
|
||||
# prompting util
|
||||
def _build_mmu_prompt(tokenizer: Any, question: str, dynin_config_path: str | None) -> dict[str, Any]:
|
||||
encoded = tokenizer(question, return_tensors="pt", add_special_tokens=True)
|
||||
token_ids = [int(v) for v in encoded["input_ids"][0].tolist()]
|
||||
attention_mask = [int(v) for v in encoded["attention_mask"][0].tolist()]
|
||||
additional_information: dict[str, Any] = {
|
||||
"task": ["mmu"],
|
||||
"detok_id": [0],
|
||||
"prompt_length": [len(token_ids)],
|
||||
"attention_mask": [attention_mask],
|
||||
"max_new_tokens": [64],
|
||||
"steps": [64],
|
||||
"block_length": [16],
|
||||
"temperature": [0.0],
|
||||
}
|
||||
if dynin_config_path:
|
||||
additional_information["dynin_config_path"] = [str(dynin_config_path)]
|
||||
return {
|
||||
"prompt_token_ids": token_ids,
|
||||
"additional_information": additional_information,
|
||||
"modalities": ["text"],
|
||||
}
|
||||
|
||||
|
||||
def _build_mmu_multimodal_prompt(
|
||||
tokenizer: Any,
|
||||
question: str,
|
||||
dynin_config_path: str | None,
|
||||
*,
|
||||
image: Any | None = None,
|
||||
audio: tuple[np.ndarray, int] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if image is None and audio is None:
|
||||
raise ValueError("At least one multimodal input (image or audio) must be provided.")
|
||||
|
||||
prefix_chunks: list[str] = []
|
||||
mm_data: dict[str, Any] = {}
|
||||
if image is not None:
|
||||
prefix_chunks.append("<|soi|><|image|><|eoi|>")
|
||||
mm_data["image"] = image
|
||||
if audio is not None:
|
||||
prefix_chunks.append("<|soa|><|audio|><|eoa|>")
|
||||
mm_data["audio"] = audio
|
||||
|
||||
prefixed_question = " ".join(prefix_chunks + [question]).strip()
|
||||
prompt = _build_mmu_prompt(
|
||||
tokenizer=tokenizer,
|
||||
question=prefixed_question,
|
||||
dynin_config_path=dynin_config_path,
|
||||
)
|
||||
prompt["multi_modal_data"] = mm_data
|
||||
prompt["modalities"] = ["text"]
|
||||
return prompt
|
||||
|
||||
|
||||
def _generate_synthetic_image(width: int = 224, height: int = 224) -> np.ndarray:
|
||||
x = np.linspace(0, 255, width, dtype=np.uint8)
|
||||
y = np.linspace(0, 255, height, dtype=np.uint8)[:, None]
|
||||
red = np.tile(x, (height, 1))
|
||||
green = np.tile(y, (1, width))
|
||||
blue = ((red.astype(np.uint16) + green.astype(np.uint16)) // 2).astype(np.uint8)
|
||||
return np.stack([red, green, blue], axis=-1)
|
||||
|
||||
|
||||
def _generate_synthetic_audio(duration_s: int = 5, sample_rate: int = 48_000) -> tuple[np.ndarray, int]:
|
||||
t = np.linspace(0, duration_s, int(sample_rate * duration_s), endpoint=False, dtype=np.float32)
|
||||
waveform = 0.1 * np.sin(2.0 * np.pi * 440.0 * t)
|
||||
return waveform.astype(np.float32), sample_rate
|
||||
|
||||
|
||||
# prompting util
|
||||
def _build_t2s_decode_prompt(dynin_config_path: str | None) -> dict[str, Any]:
|
||||
# Bypass stage-0 generation and directly validate token->audio decode path.
|
||||
generated_audio_token_ids = [int(v) for v in ([10, 11, 12, 13, 14] * 32)]
|
||||
additional_information: dict[str, Any] = {
|
||||
"task": ["t2s"],
|
||||
"detok_id": [1],
|
||||
"generated_token_ids": [generated_audio_token_ids],
|
||||
"audio_codebook_size": [4096],
|
||||
}
|
||||
if dynin_config_path:
|
||||
additional_information["dynin_config_path"] = [str(dynin_config_path)]
|
||||
return {
|
||||
"prompt_token_ids": [0],
|
||||
"additional_information": additional_information,
|
||||
"modalities": ["audio"],
|
||||
}
|
||||
|
||||
|
||||
# prompting util
|
||||
def _build_t2i_decode_prompt(dynin_config_path: str | None) -> dict[str, Any]:
|
||||
# Bypass stage-0 generation and directly validate token->image decode path.
|
||||
# MAGVIT decode path expects a square token grid; 1024 tokens -> 32x32.
|
||||
generated_image_token_ids = [int(v) for v in ([10, 11, 12, 13, 14, 15, 16, 17] * 128)]
|
||||
additional_information: dict[str, Any] = {
|
||||
"task": ["t2i"],
|
||||
"detok_id": [2],
|
||||
"generated_token_ids": [generated_image_token_ids],
|
||||
"codebook_size": [8192],
|
||||
}
|
||||
if dynin_config_path:
|
||||
additional_information["dynin_config_path"] = [str(dynin_config_path)]
|
||||
return {
|
||||
"prompt_token_ids": [0],
|
||||
"additional_information": additional_information,
|
||||
"modalities": ["image"],
|
||||
}
|
||||
|
||||
|
||||
def _configure_dynin_config_env() -> None:
|
||||
if DYNIN_CONFIG_PATH:
|
||||
os.environ["DYNIN_CONFIG_PATH"] = str(DYNIN_CONFIG_PATH)
|
||||
else:
|
||||
os.environ.pop("DYNIN_CONFIG_PATH", None)
|
||||
|
||||
|
||||
def _is_finished_request_output(request_output: Any) -> bool:
|
||||
if request_output is None:
|
||||
return False
|
||||
req_list = request_output if isinstance(request_output, list) else [request_output]
|
||||
for req in req_list:
|
||||
if req is not None and bool(getattr(req, "finished", False)):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _find_stage_output(outputs: list[Any], output_type: str) -> Any | None:
|
||||
matched = [
|
||||
stage_output for stage_output in outputs if getattr(stage_output, "final_output_type", None) == output_type
|
||||
]
|
||||
if not matched:
|
||||
return None
|
||||
|
||||
# Prefer the latest finished chunk to avoid picking an intermediate stream output.
|
||||
for stage_output in reversed(matched):
|
||||
if _is_finished_request_output(stage_output):
|
||||
return stage_output
|
||||
return matched[-1]
|
||||
|
||||
|
||||
def _to_token_list(value: Any) -> list[int]:
|
||||
if value is None:
|
||||
return []
|
||||
if hasattr(value, "detach"):
|
||||
value = value.detach()
|
||||
if hasattr(value, "cpu"):
|
||||
value = value.cpu()
|
||||
if hasattr(value, "flatten"):
|
||||
value = value.flatten().tolist()
|
||||
if isinstance(value, tuple):
|
||||
value = list(value)
|
||||
if not isinstance(value, list):
|
||||
return []
|
||||
out: list[int] = []
|
||||
for token in value:
|
||||
if isinstance(token, bool):
|
||||
continue
|
||||
try:
|
||||
out.append(int(token))
|
||||
except Exception:
|
||||
continue
|
||||
return out
|
||||
|
||||
|
||||
def _extract_text(stage_output: Any, tokenizer: Any | None = None) -> str:
|
||||
request_output = stage_output
|
||||
if request_output is None:
|
||||
return ""
|
||||
req_list = request_output if isinstance(request_output, list) else [request_output]
|
||||
for req in req_list:
|
||||
completions = getattr(req, "outputs", None) or []
|
||||
if not completions:
|
||||
continue
|
||||
completion = completions[0]
|
||||
mm_out = (
|
||||
getattr(completion, "multimodal_output", None)
|
||||
or getattr(req, "multimodal_output", None)
|
||||
or getattr(stage_output, "multimodal_output", None)
|
||||
or {}
|
||||
)
|
||||
text = mm_out.get("text")
|
||||
if isinstance(text, list) and text:
|
||||
text = text[-1]
|
||||
if isinstance(text, str) and text.strip():
|
||||
return text.strip()
|
||||
if tokenizer is not None:
|
||||
for key in ("text_tokens", "token_ids"):
|
||||
token_ids = _to_token_list(mm_out.get(key))
|
||||
if not token_ids:
|
||||
continue
|
||||
decoded = tokenizer.decode(token_ids, skip_special_tokens=True)
|
||||
if isinstance(decoded, str) and decoded.strip():
|
||||
return decoded.strip()
|
||||
fallback = getattr(completion, "text", None)
|
||||
if isinstance(fallback, str) and fallback.strip():
|
||||
return fallback.strip()
|
||||
return ""
|
||||
|
||||
|
||||
def _extract_audio(stage_output: Any) -> Any | None:
|
||||
request_output = stage_output
|
||||
if request_output is None:
|
||||
return None
|
||||
req_list = request_output if isinstance(request_output, list) else [request_output]
|
||||
for req in req_list:
|
||||
completions = getattr(req, "outputs", None) or []
|
||||
if not completions:
|
||||
continue
|
||||
completion = completions[0]
|
||||
mm_out = getattr(completion, "multimodal_output", None) or {}
|
||||
if "audio" in mm_out:
|
||||
return mm_out["audio"]
|
||||
return None
|
||||
|
||||
|
||||
def _extract_image(stage_output: Any) -> Any | None:
|
||||
request_output = stage_output
|
||||
if request_output is None:
|
||||
return None
|
||||
req_list = request_output if isinstance(request_output, list) else [request_output]
|
||||
for req in req_list:
|
||||
completions = getattr(req, "outputs", None) or []
|
||||
if not completions:
|
||||
continue
|
||||
completion = completions[0]
|
||||
mm_out = getattr(completion, "multimodal_output", None) or {}
|
||||
if "image" in mm_out:
|
||||
return mm_out["image"]
|
||||
return None
|
||||
|
||||
|
||||
def _numel(value: Any) -> int:
|
||||
if value is None:
|
||||
return 0
|
||||
if isinstance(value, torch.Tensor):
|
||||
return int(value.numel())
|
||||
shape = getattr(value, "shape", None)
|
||||
if shape is not None:
|
||||
try:
|
||||
total = 1
|
||||
for dim in shape:
|
||||
total *= int(dim)
|
||||
return int(total)
|
||||
except Exception:
|
||||
pass
|
||||
if isinstance(value, (list, tuple)):
|
||||
return len(value)
|
||||
return 0
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
|
||||
def test_dynin_t2i_decode_to_image(omni_runner) -> None:
|
||||
_configure_dynin_config_env()
|
||||
prompt = _build_t2i_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH)
|
||||
|
||||
outputs = omni_runner.generate([prompt])
|
||||
|
||||
image_output = _find_stage_output(outputs, "image")
|
||||
assert image_output is not None
|
||||
image_value = _extract_image(image_output)
|
||||
assert image_value is not None
|
||||
assert _numel(image_value) > 0
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
|
||||
def test_dynin_mmu_to_text(omni_runner) -> None:
|
||||
_configure_dynin_config_env()
|
||||
tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True)
|
||||
prompt = _build_mmu_prompt(
|
||||
tokenizer=tokenizer,
|
||||
question="What is 2 + 2? Answer in one short sentence.",
|
||||
dynin_config_path=DYNIN_CONFIG_PATH,
|
||||
)
|
||||
|
||||
outputs = omni_runner.generate([prompt])
|
||||
|
||||
text_output = _find_stage_output(outputs, "text")
|
||||
assert text_output is not None
|
||||
text_content = _extract_text(text_output, tokenizer=tokenizer)
|
||||
assert text_content
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
|
||||
def test_dynin_image_to_text(omni_runner) -> None:
|
||||
_configure_dynin_config_env()
|
||||
tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True)
|
||||
prompt = _build_mmu_multimodal_prompt(
|
||||
tokenizer=tokenizer,
|
||||
question="Describe the image briefly in one sentence.",
|
||||
dynin_config_path=DYNIN_CONFIG_PATH,
|
||||
image=_generate_synthetic_image(),
|
||||
)
|
||||
|
||||
outputs = omni_runner.generate([prompt])
|
||||
|
||||
text_output = _find_stage_output(outputs, "text")
|
||||
assert text_output is not None
|
||||
text_content = _extract_text(text_output, tokenizer=tokenizer)
|
||||
assert text_content
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
|
||||
def test_dynin_speech_to_text(omni_runner) -> None:
|
||||
_configure_dynin_config_env()
|
||||
tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True)
|
||||
prompt = _build_mmu_multimodal_prompt(
|
||||
tokenizer=tokenizer,
|
||||
question="Transcribe the audio briefly in one sentence.",
|
||||
dynin_config_path=DYNIN_CONFIG_PATH,
|
||||
audio=_generate_synthetic_audio(),
|
||||
)
|
||||
|
||||
outputs = omni_runner.generate([prompt])
|
||||
|
||||
text_output = _find_stage_output(outputs, "text")
|
||||
assert text_output is not None
|
||||
text_content = _extract_text(text_output, tokenizer=tokenizer)
|
||||
assert text_content
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
|
||||
def test_dynin_t2s_decode_to_audio(omni_runner) -> None:
|
||||
_configure_dynin_config_env()
|
||||
prompt = _build_t2s_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH)
|
||||
|
||||
outputs = omni_runner.generate([prompt])
|
||||
|
||||
audio_output = _find_stage_output(outputs, "audio")
|
||||
assert audio_output is not None
|
||||
audio_value = _extract_audio(audio_output)
|
||||
assert audio_value is not None
|
||||
assert _numel(audio_value) > 0
|
||||
@@ -1,91 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
E2E Online expansion tests for dots.tts(text-only).
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
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.full_model, pytest.mark.tts]
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
MODEL = "dots-studio/dots.tts-soar"
|
||||
DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 300.0
|
||||
_MIN_AUDIO_BYTES = 40_000
|
||||
MAX_CONCURRENT = 4
|
||||
|
||||
|
||||
def get_prompt(prompt_type="text"):
|
||||
prompts = {
|
||||
"text": "The weather is nice today, perfect for a walk in the park.",
|
||||
}
|
||||
return prompts.get(prompt_type, prompts["text"])
|
||||
|
||||
|
||||
tts_server_params = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
stage_config_path=get_deploy_config_path("dots_tts.yaml"),
|
||||
server_args=["--trust-remote-code", "--disable-log-stats"],
|
||||
),
|
||||
id="dots_tts",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
def test_text_only_streaming_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Test text-only text-to-audio via OpenAI API.
|
||||
Deploy Setting: dots_tts.yaml
|
||||
Input Modal: text
|
||||
Output Modal: audio
|
||||
Input Setting: stream=True
|
||||
Datasets: few requests (max_num_seqs=4)
|
||||
"""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"input": get_prompt(),
|
||||
"stream": True,
|
||||
"stream_format": "audio",
|
||||
"timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
|
||||
"response_format": "wav",
|
||||
"voice": "default",
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
openai_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
def test_text_only_nonstreaming_pcm_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Test text-only non-stream PCM output via OpenAI API.
|
||||
Deploy Setting: dots_tts.yaml
|
||||
Input Modal: text
|
||||
Output Modal: audio (pcm)
|
||||
Input Setting: stream=False
|
||||
Datasets: single request
|
||||
"""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"input": get_prompt(),
|
||||
"response_format": "pcm",
|
||||
"stream": False,
|
||||
"timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
|
||||
"voice": "default",
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
"min_hnr_db": -2.0,
|
||||
# dots.tts-soar emits native 48 kHz PCM, which has no rate header.
|
||||
"expected_sample_rate": 48_000,
|
||||
}
|
||||
openai_client.send_audio_speech_request(request_config)
|
||||
@@ -1,99 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""
|
||||
Example online tests for Dynin-Omni model.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
from io import BytesIO
|
||||
|
||||
import pytest
|
||||
from vllm.assets.image import ImageAsset
|
||||
|
||||
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.slow, pytest.mark.omni]
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
MODEL = "snu-aidas/Dynin-Omni"
|
||||
STAGE_CONFIG = get_deploy_config_path("dynin_omni_ci.yaml")
|
||||
|
||||
T2I_PROMPT = "A high quality detailed living room interior photo."
|
||||
T2S_PROMPT = "Please read this sentence naturally: Hello from online serving."
|
||||
I2I_PROMPT = "Transform this outdoor nature boardwalk scene into a painting style with vivid colors."
|
||||
|
||||
TEST_PARAMS = [
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
stage_config_path=STAGE_CONFIG,
|
||||
stage_init_timeout=600,
|
||||
server_args=["--trust-remote-code"],
|
||||
)
|
||||
]
|
||||
_STAGE_COUNT = 3
|
||||
_I2I_STAGE_SAMPLING = {"max_tokens": 1, "temperature": 0.0, "top_p": 1.0, "detokenize": False}
|
||||
|
||||
|
||||
def _build_t2i_messages(prompt: str) -> list[dict]:
|
||||
return [{"role": "user", "content": [{"type": "text", "text": f"<|t2i|> {prompt}"}]}]
|
||||
|
||||
|
||||
def _build_t2s_messages(prompt: str) -> list[dict]:
|
||||
return [{"role": "user", "content": [{"type": "text", "text": f"<|t2s|> {prompt}"}]}]
|
||||
|
||||
|
||||
def _build_i2i_messages(prompt: str) -> list[dict]:
|
||||
input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
|
||||
buffer = BytesIO()
|
||||
input_image.save(buffer, format="JPEG")
|
||||
image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": f"<|i2i|> {prompt}"},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
def test_send_i2i_request_001(omni_server, online_client) -> None:
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": _build_i2i_messages(I2I_PROMPT),
|
||||
"modalities": ["image"],
|
||||
"extra_body": {
|
||||
"sampling_params_list": [dict(_I2I_STAGE_SAMPLING) for _ in range(_STAGE_COUNT)],
|
||||
},
|
||||
}
|
||||
online_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
def test_send_t2i_request_001(omni_server, online_client) -> None:
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": _build_t2i_messages(T2I_PROMPT),
|
||||
"modalities": ["image"],
|
||||
}
|
||||
online_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
def test_send_t2s_request_001(omni_server, online_client) -> None:
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": _build_t2s_messages(T2S_PROMPT),
|
||||
"modalities": ["audio"],
|
||||
"audio_ref_text": T2S_PROMPT,
|
||||
}
|
||||
online_client.send_omni_request(request_config)
|
||||
@@ -539,7 +539,7 @@ def test_tensor_parallel_size_none_is_handled():
|
||||
engine_args = OmegaConf.create({"stage_id": 0, "engine_args": {"tensor_parallel_size": None}})
|
||||
args = build_engine_args_dict(
|
||||
engine_args,
|
||||
model="snu-aidas/Dynin-Omni",
|
||||
model="Qwen/Qwen2-VL-2B-Instruct",
|
||||
)
|
||||
assert isinstance(args, dict)
|
||||
assert "tensor_parallel_size" not in args
|
||||
@@ -5035,89 +5035,6 @@ class TestMingFlashOmniTTSServing:
|
||||
assert prompt["additional_information"]["voice_name"] == "test"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dots_tts_server(mocker: MockerFixture):
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers",
|
||||
return_value=set(),
|
||||
)
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate",
|
||||
return_value=None,
|
||||
)
|
||||
|
||||
mock_engine_client = mocker.MagicMock()
|
||||
mock_engine_client.errored = False
|
||||
mock_engine_client.model_config = mocker.MagicMock(
|
||||
model="dots-studio/dots.tts-soar",
|
||||
)
|
||||
mock_engine_client.default_sampling_params_list = [
|
||||
SimpleNamespace(max_tokens=2048, min_tokens=None, extra_args=None)
|
||||
]
|
||||
mock_engine_client.tts_batch_max_items = 32
|
||||
mock_engine_client.generate = mocker.MagicMock(return_value="generator")
|
||||
mock_engine_client.stage_configs = [
|
||||
SimpleNamespace(
|
||||
engine_args=SimpleNamespace(model_stage="latent_generator", model_arch="DotsTTSForConditionalGeneration"),
|
||||
tts_args={},
|
||||
)
|
||||
]
|
||||
|
||||
mock_models = mocker.MagicMock()
|
||||
mock_models.is_base_model.return_value = True
|
||||
|
||||
return OmniOpenAIServingSpeech(
|
||||
engine_client=mock_engine_client,
|
||||
models=mock_models,
|
||||
request_logger=mocker.MagicMock(),
|
||||
)
|
||||
|
||||
|
||||
class TestDotsTTSServing:
|
||||
def test_dots_tts_prompt_validation(self, dots_tts_server):
|
||||
request = OpenAICreateSpeechRequest(input="Hello", ref_text="Reference transcript")
|
||||
error = dots_tts_server._validate_tts_request(request)
|
||||
assert error is not None
|
||||
assert "ref_text" in error
|
||||
|
||||
request = OpenAICreateSpeechRequest(input="Hello", ref_audio="data:audio/wav;base64,abc")
|
||||
error = dots_tts_server._validate_tts_request(request)
|
||||
assert error is not None
|
||||
assert "ref_audio" in error
|
||||
|
||||
request = OpenAICreateSpeechRequest(input="Hello", voice="test")
|
||||
error = dots_tts_server._validate_tts_request(request)
|
||||
assert error is not None
|
||||
assert "voice" in error
|
||||
|
||||
request = OpenAICreateSpeechRequest(input="Hello", speaker_embedding=[1, 2, 3])
|
||||
error = dots_tts_server._validate_tts_request(request)
|
||||
assert error is not None
|
||||
assert "speaker_embedding" in error
|
||||
|
||||
request = OpenAICreateSpeechRequest(input="Hello", x_vector_only_mode=True)
|
||||
error = dots_tts_server._validate_tts_request(request)
|
||||
assert error is not None
|
||||
assert "x_vector_only_mode" in error
|
||||
|
||||
def test_dots_tts_adapter_awaits_async_prompt_builder(self, dots_tts_server, mocker: MockerFixture):
|
||||
build_prompt_async = mocker.patch.object(
|
||||
dots_tts_server._adapter,
|
||||
"_build_prompt_async",
|
||||
new=mocker.AsyncMock(return_value={"prompt_token_ids": [1, 2, 3]}),
|
||||
)
|
||||
request = OpenAICreateSpeechRequest(input="Hello")
|
||||
asyncio.run(dots_tts_server._prepare_speech_generation(request))
|
||||
build_prompt_async.assert_awaited_once_with("Hello")
|
||||
|
||||
def test_dots_tts_adapter_apply_sampling_overrides(self, dots_tts_server, mocker: MockerFixture):
|
||||
mocker.patch.object(dots_tts_server._adapter, "build", return_value=PreparedRequest(prompt="Hello"))
|
||||
request = OpenAICreateSpeechRequest(input="Hello", max_new_tokens=10)
|
||||
asyncio.run(dots_tts_server._prepare_speech_generation(request))
|
||||
sampling_params_list = dots_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"]
|
||||
assert sampling_params_list[0].max_tokens == 10
|
||||
|
||||
|
||||
class TestTTSAsyncOffloading:
|
||||
"""Tests for event-loop-safe offloading of blocking TTS operations."""
|
||||
|
||||
|
||||
@@ -71,7 +71,6 @@ EXPECTED_MODEL_TYPES = {
|
||||
"indextts2",
|
||||
"indextts2_5",
|
||||
"gepard",
|
||||
"dots_tts",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -263,7 +263,6 @@ def test_arch_matching_is_a_fallback_not_an_override():
|
||||
|
||||
def test_shared_latent_generator_resolves_by_architecture_priority():
|
||||
assert detect_tts_model_type("latent_generator", "VoxCPM2TalkerForConditionalGeneration") == "voxcpm2"
|
||||
assert detect_tts_model_type("latent_generator", "DotsTTSForConditionalGeneration") == "dots_tts"
|
||||
|
||||
|
||||
def test_stage_keys_cover_legacy_stage_set():
|
||||
|
||||
@@ -1,2 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -1,137 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Regression tests for the dots.tts talker's stop-signal pairing and
|
||||
per-request state lifecycle.
|
||||
|
||||
Pure-CPU tests on a bare talker (``__new__`` + the handful of attributes
|
||||
the methods under test read) — no weights, no GPU, no engine. Mirrors
|
||||
tests/model_executor/models/voxcpm2/test_talker_state_eviction.py.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
_VOCAB_SIZE = 151672
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _dots_tts_talker_mod():
|
||||
"""Defer talker import (pulls vLLM model_executor) until first use."""
|
||||
from vllm_omni.model_executor.models.dots_tts.dots_tts_talker import (
|
||||
DotsTTSForConditionalGeneration,
|
||||
_RequestState,
|
||||
)
|
||||
|
||||
return DotsTTSForConditionalGeneration, _RequestState
|
||||
|
||||
|
||||
def _make_bare_talker():
|
||||
DotsTTSForConditionalGeneration, _ = _dots_tts_talker_mod()
|
||||
talker = DotsTTSForConditionalGeneration.__new__(DotsTTSForConditionalGeneration)
|
||||
talker.config = SimpleNamespace(vocab_size=_VOCAB_SIZE)
|
||||
talker._active_states = {}
|
||||
talker._pending_requests = []
|
||||
talker._results_queue = []
|
||||
talker._audio_queue = []
|
||||
talker._deferred_cleanup_ids = set()
|
||||
talker._beta_trace = False
|
||||
return talker
|
||||
|
||||
|
||||
def _add_state(talker, req_id: str, *, is_stopping: bool = False):
|
||||
_, _RequestState = _dots_tts_talker_mod()
|
||||
state = _RequestState(request_id=req_id)
|
||||
state.prefill_completed = True
|
||||
state.is_stopping = is_stopping
|
||||
talker._active_states[req_id] = state
|
||||
return state
|
||||
|
||||
|
||||
def _hidden(bsz: int) -> torch.Tensor:
|
||||
return torch.zeros(bsz, 1536, dtype=torch.float32)
|
||||
|
||||
|
||||
class TestComputeLogitsPairing:
|
||||
"""compute_logits pairs _results_queue entries with batch rows
|
||||
positionally, so every request must push exactly one entry per step —
|
||||
including prefills, which push a (req_id, None) placeholder."""
|
||||
|
||||
def test_mixed_batch_prefill_placeholder_keeps_rows_aligned(self) -> None:
|
||||
talker = _make_bare_talker()
|
||||
_add_state(talker, "prefill-req")
|
||||
decode_state = _add_state(talker, "decode-req", is_stopping=True)
|
||||
decode_state.precomputed_stop_logits = torch.tensor([[0.1, 0.9]])
|
||||
|
||||
# Row 0 is the prefill (None placeholder), row 1 the stopping decode.
|
||||
talker._results_queue = [
|
||||
("prefill-req", None),
|
||||
("decode-req", torch.tensor([[0.1, 0.9]])),
|
||||
]
|
||||
logits = talker.compute_logits(_hidden(2))
|
||||
|
||||
assert logits.shape == (2, _VOCAB_SIZE)
|
||||
# Prefill row: forced continue.
|
||||
assert logits[0, 0].item() == 1.0
|
||||
assert logits[0, 1].item() == float("-inf")
|
||||
# Stopping decode row: stop slot wins.
|
||||
assert logits[1, 1].item() == 1.0
|
||||
assert logits[1, 0].item() == 0.0
|
||||
# Everything else stays masked.
|
||||
assert logits[0, 2].item() == float("-inf")
|
||||
assert logits[1, 2].item() == float("-inf")
|
||||
# Queue drained.
|
||||
assert talker._results_queue == []
|
||||
|
||||
def test_below_threshold_stop_logits_force_continue(self) -> None:
|
||||
"""Raw (continue, stop) softmax must not reach the sampler: a
|
||||
prob_stop in (0.5, 0.8] would flip greedy argmax even though
|
||||
is_stopping stayed False (the 0.8-threshold contract)."""
|
||||
talker = _make_bare_talker()
|
||||
_add_state(talker, "req", is_stopping=False)
|
||||
talker._results_queue = [("req", torch.tensor([[0.3, 0.7]]))]
|
||||
|
||||
logits = talker.compute_logits(_hidden(1))
|
||||
|
||||
assert logits[0, 0].item() == 1.0
|
||||
assert logits[0, 1].item() == float("-inf")
|
||||
|
||||
def test_empty_queue_defaults_to_all_continue(self) -> None:
|
||||
talker = _make_bare_talker()
|
||||
logits = talker.compute_logits(_hidden(3))
|
||||
|
||||
assert torch.all(logits[:, 0] == 1.0)
|
||||
assert torch.all(logits[:, 1] == float("-inf"))
|
||||
|
||||
|
||||
class TestStateEviction:
|
||||
"""on_requests_finished fires BEFORE forward(), so eviction must be
|
||||
deferred to _flush_deferred_cleanup at the end of the step — the
|
||||
finishing request's audio still drains through the current forward."""
|
||||
|
||||
def test_finished_request_survives_until_flush(self) -> None:
|
||||
talker = _make_bare_talker()
|
||||
_add_state(talker, "req-a")
|
||||
_add_state(talker, "req-b")
|
||||
|
||||
talker.on_requests_finished(["req-a"])
|
||||
assert "req-a" in talker._active_states # not evicted yet
|
||||
|
||||
talker._flush_deferred_cleanup()
|
||||
assert "req-a" not in talker._active_states
|
||||
assert "req-b" in talker._active_states
|
||||
assert talker._deferred_cleanup_ids == set()
|
||||
|
||||
def test_unknown_request_id_is_ignored(self) -> None:
|
||||
talker = _make_bare_talker()
|
||||
talker.on_requests_finished(["never-seen"])
|
||||
talker._flush_deferred_cleanup()
|
||||
assert talker._active_states == {}
|
||||
@@ -1,40 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from transformers import PreTrainedModel
|
||||
|
||||
from vllm_omni.model_executor.models.dynin_omni.dynin_omni_token2audio import (
|
||||
_ensure_transformers_tied_weights_compat,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_transformers_tied_weights_compat(monkeypatch):
|
||||
"""Compat shim must not permanently mutate ``PreTrainedModel``.
|
||||
|
||||
``monkeypatch.delattr(..., raising=False)`` is a no-op undo when the attr
|
||||
was absent. ``_ensure_transformers_tied_weights_compat`` then assigns a
|
||||
read-only ``property`` on the class that later PreTrainedModel subclasses
|
||||
(e.g. qwen3 TTS mask-cache) can trip over with
|
||||
``AttributeError: property ... has no setter``.
|
||||
"""
|
||||
monkeypatch.delattr(PreTrainedModel, "all_tied_weights_keys", raising=False)
|
||||
|
||||
_ensure_transformers_tied_weights_compat()
|
||||
installed = PreTrainedModel.__dict__.get("all_tied_weights_keys")
|
||||
assert installed is not None
|
||||
|
||||
# Re-own the installed descriptor through monkeypatch so teardown always
|
||||
# removes/restores it, even when the attr did not exist before this test.
|
||||
delattr(PreTrainedModel, "all_tied_weights_keys")
|
||||
monkeypatch.setattr(PreTrainedModel, "all_tied_weights_keys", installed, raising=False)
|
||||
|
||||
model = PreTrainedModel.__new__(PreTrainedModel)
|
||||
model._tied_weights_keys = ["decoder.weight"]
|
||||
model._dynamic_tied_weights_keys = {"lm_head.weight": "decoder.weight"}
|
||||
|
||||
assert model.all_tied_weights_keys == {
|
||||
"decoder.weight": None,
|
||||
"lm_head.weight": "decoder.weight",
|
||||
}
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Unit tests for Qwen3-Omni streaming thinker→talker / talker→codec helpers (PR #2581)."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -356,7 +356,7 @@ def test_accumulator_concat_default_when_no_replace_keys() -> None:
|
||||
class _StubMixin(OmniConnectorModelRunnerMixin):
|
||||
def __init__(self):
|
||||
self._pending_full_payload_send = {}
|
||||
self._full_payload_replace_keys_cached = frozenset()
|
||||
self._full_payload_replace_keys_cached: frozenset[str] = frozenset()
|
||||
|
||||
stub = _StubMixin()
|
||||
stub.accumulate_full_payload_output(
|
||||
@@ -415,51 +415,6 @@ def test_covo_audio_llm2code2wav_full_payload_smoke() -> None:
|
||||
assert payload["meta"]["finished"].item() is True
|
||||
|
||||
|
||||
def test_dynin_omni_token_only_smoke() -> None:
|
||||
"""Smoke: dynin_omni token-only builders return placeholders."""
|
||||
from vllm_omni.model_executor.stage_input_processors.dynin_omni import (
|
||||
token2text_to_token2image_token_only,
|
||||
)
|
||||
|
||||
class _Out:
|
||||
def __init__(self, tids, mm=None):
|
||||
self.token_ids = tids
|
||||
self.multimodal_output = mm
|
||||
|
||||
class _Wrapper:
|
||||
def __init__(self, tids, mm=None):
|
||||
self.outputs = [_Out(tids, mm)]
|
||||
self.request_id = "r0"
|
||||
|
||||
class _Stage:
|
||||
def __init__(self, outs):
|
||||
self.engine_outputs = outs
|
||||
|
||||
src = [_Wrapper([10, 11, 12])]
|
||||
out = token2text_to_token2image_token_only([_Stage(src)], [0])
|
||||
assert len(out) == 1
|
||||
assert len(out[0]["prompt_token_ids"]) == 3
|
||||
assert out[0]["additional_information"] is None
|
||||
|
||||
|
||||
def test_dynin_omni_full_payload_smoke() -> None:
|
||||
"""Smoke: dynin_omni producer-side payload builder returns nested OmniPayload + carries metadata."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.dynin_omni import (
|
||||
token2text_to_token2image_full_payload,
|
||||
)
|
||||
|
||||
pooling = {"token_ids": [1, 2, 3]}
|
||||
req = SimpleNamespace(output_token_ids=[], additional_information={"speaker": ["alice"]})
|
||||
payload = token2text_to_token2image_full_payload(None, pooling, req)
|
||||
assert payload is not None
|
||||
assert payload["codes"]["audio"] == [1, 2, 3]
|
||||
assert payload["meta"]["finished"].item() is True
|
||||
# additional_information is normalized + carried forward (speaker stays list-wrapped).
|
||||
assert payload.get("speaker") == ["alice"]
|
||||
|
||||
|
||||
def test_qwen2_5_omni_talker2code2wav_token_only_smoke() -> None:
|
||||
"""Smoke: qwen2_5_omni talker→code2wav token_only marker + boundary strip."""
|
||||
from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
|
||||
|
||||
@@ -145,13 +145,6 @@ _MODEL_PROMOTE = (
|
||||
"COSYVOICE3_ESTIMATOR_ONNX",
|
||||
"COSYVOICE3_TRT",
|
||||
"COSYVOICE3_TRT_CACHE",
|
||||
"DYNIN_CONFIG_PATH",
|
||||
"DYNIN_MAGVIT_REMOTE_CODE_LOCAL_FILES_ONLY",
|
||||
"DYNIN_MAGVIT_REMOTE_CODE_REPO_ID",
|
||||
"DYNIN_MAGVIT_REMOTE_CODE_REVISION",
|
||||
"DYNIN_REMOTE_CODE_LOCAL_FILES_ONLY",
|
||||
"DYNIN_REMOTE_CODE_REPO_ID",
|
||||
"DYNIN_REMOTE_CODE_REVISION",
|
||||
"FASTVIDEO_VSA_SM100A",
|
||||
"HIGGS_AUDIO_TOKENIZER_PATH",
|
||||
"INTERNVLA_A1_COSMOS_DECODER_PATH",
|
||||
@@ -179,7 +172,6 @@ _MODEL_PROMOTE = (
|
||||
)
|
||||
|
||||
_MODEL_REQUEST_SCOPE = (
|
||||
"DOTS_TTS_DIT_NUM_STEPS",
|
||||
"MAGI2_NEGATIVE_PROMPT",
|
||||
"STEP_AUDIO2_DEFAULT_PROMPT_WAV",
|
||||
"VLLM_GEPARD_END_FADE_MS",
|
||||
@@ -194,8 +186,6 @@ _MODEL_EXTERNAL: tuple[str, ...] = ()
|
||||
|
||||
_MODEL_INTERNALIZE = (
|
||||
"COSYVOICE3_BATCH_FLOW_DEBUG",
|
||||
"DOTS_TTS_BETA_TRACE",
|
||||
"DYNIN_S2U_VENDOR_ROOT",
|
||||
"DZ_PHASE_TIMING",
|
||||
"MAGI2_ALLOW_UNSUPPORTED_TOPOLOGY",
|
||||
"MAGI2_ROUTER_BIAS_SOURCE",
|
||||
|
||||
@@ -52,9 +52,7 @@ from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_P
|
||||
from vllm_omni.model_executor.models.cosmos3.pipeline import COSMOS3_POLICY_PIPELINE
|
||||
from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE
|
||||
from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE
|
||||
from vllm_omni.model_executor.models.dots_tts.pipeline import DOTS_TTS_PIPELINE
|
||||
from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE
|
||||
from vllm_omni.model_executor.models.dynin_omni.pipeline import DYNIN_OMNI_PIPELINE
|
||||
from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE
|
||||
from vllm_omni.model_executor.models.gepard.pipeline import GEPARD_PIPELINE
|
||||
from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE
|
||||
@@ -168,7 +166,6 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
|
||||
"hunyuan_video_15": HUNYUAN_VIDEO_15_PIPELINE,
|
||||
"wan2_2_ti2v": WAN2_2_TI2V_PIPELINE,
|
||||
"voxcpm2": VOXCPM2_PIPELINE,
|
||||
"dots_tts": DOTS_TTS_PIPELINE,
|
||||
"auk": AUK_PIPELINE,
|
||||
"cosyvoice3": COSYVOICE3_PIPELINE,
|
||||
"audex_tts": AUDEX_TTS_PIPELINE,
|
||||
@@ -202,7 +199,6 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
|
||||
"minimax_music3": MINIMAX_MUSIC3_PIPELINE,
|
||||
"higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE,
|
||||
"higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE,
|
||||
"dynin_omni": DYNIN_OMNI_PIPELINE,
|
||||
"indextts2": INDEXTTS2_PIPELINE,
|
||||
"indextts2_5": INDEXTTS25_PIPELINE,
|
||||
}
|
||||
|
||||
@@ -1,28 +0,0 @@
|
||||
# dots.tts deployment config (single-stage AR talker).
|
||||
#
|
||||
# Conservative first-integration settings: eager execution, no CUDA
|
||||
# Graphs, no batched side-path optimizations yet — see voxcpm2.yaml for
|
||||
# the runtime knobs a mature deployment of this model family grows.
|
||||
#
|
||||
# Verified locally on RTX 5080 16GB.
|
||||
|
||||
async_chunk: false
|
||||
dtype: bfloat16
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_seqs: 4
|
||||
gpu_memory_utilization: 0.8
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 4096
|
||||
max_model_len: 4096
|
||||
devices: "0"
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 4096
|
||||
repetition_penalty: 1.0
|
||||
@@ -1,38 +0,0 @@
|
||||
# Dynin-Omni 3-stage deploy: token2text → token2image → token2audio.
|
||||
#
|
||||
# Topology (stages, processor functions, model_arch) is frozen in
|
||||
# vllm_omni/model_executor/models/dynin_omni/pipeline.py. This file carries
|
||||
# only the deploy-time knobs (devices, GPU memory, batched tokens).
|
||||
#
|
||||
# async_chunk=false: dynin hands the full token payload across stages through
|
||||
# the worker-connector data plane (``*_full_payload`` producers), not chunked
|
||||
# streaming. Pipeline-wide ``trust_remote_code`` is required by the custom
|
||||
# Dynin HF modeling code.
|
||||
async_chunk: false
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
devices: "0"
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.5
|
||||
max_num_batched_tokens: 32768
|
||||
enforce_eager: true
|
||||
async_scheduling: false
|
||||
|
||||
- stage_id: 1
|
||||
devices: "0"
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.1
|
||||
max_num_batched_tokens: 32768
|
||||
enforce_eager: true
|
||||
async_scheduling: false
|
||||
|
||||
- stage_id: 2
|
||||
devices: "0"
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.1
|
||||
max_num_batched_tokens: 32768
|
||||
enforce_eager: true
|
||||
async_scheduling: false
|
||||
@@ -1,35 +0,0 @@
|
||||
# Dynin-Omni deploy for e2e smoke tests (3-stage, one GPU).
|
||||
#
|
||||
# Topology is frozen in models/dynin_omni/pipeline.py; ``pipeline: dynin_omni``
|
||||
# pins it explicitly so this CI YAML resolves regardless of model_type
|
||||
# auto-detection. Differs from the default deploy only in reduced GPU memory
|
||||
# and a smaller max_num_batched_tokens so the pipeline fits L4-class CI GPUs.
|
||||
pipeline: dynin_omni
|
||||
async_chunk: false
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
devices: "0"
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.5
|
||||
max_num_batched_tokens: 4096
|
||||
enforce_eager: true
|
||||
async_scheduling: false
|
||||
|
||||
- stage_id: 1
|
||||
devices: "0"
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.2
|
||||
max_num_batched_tokens: 4096
|
||||
enforce_eager: true
|
||||
async_scheduling: false
|
||||
|
||||
- stage_id: 2
|
||||
devices: "0"
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.2
|
||||
max_num_batched_tokens: 4096
|
||||
enforce_eager: true
|
||||
async_scheduling: false
|
||||
@@ -1,63 +0,0 @@
|
||||
# Dynin-Omni 3-stage deploy with cross-node connectors (Mooncake demo).
|
||||
#
|
||||
# Topology is frozen in models/dynin_omni/pipeline.py (``pipeline: dynin_omni``).
|
||||
# This variant wires a MooncakeConnector for each stage hand-off in place of
|
||||
# the default SharedMemoryConnector. Update the ``extra`` host/master/metadata
|
||||
# fields to match your Mooncake deployment before use.
|
||||
pipeline: dynin_omni
|
||||
async_chunk: false
|
||||
trust_remote_code: false
|
||||
enable_prefix_caching: false
|
||||
|
||||
connectors:
|
||||
# Mooncake connector for cross-node/intra-node communication.
|
||||
mooncake_connector:
|
||||
name: MooncakeConnector
|
||||
extra:
|
||||
host: "127.0.0.1"
|
||||
metadata_server: "http://10.90.67.86:8080/metadata"
|
||||
master: "10.90.67.86:50051"
|
||||
segment: 512000000 # 512MB
|
||||
localbuf: 64000000 # 64MB
|
||||
proto: "tcp"
|
||||
|
||||
# Yuanrong connector for cross-node/intra-node communication.
|
||||
yuanrong_connector:
|
||||
name: YuanrongConnector
|
||||
extra:
|
||||
host: "127.0.0.1"
|
||||
port: "35000"
|
||||
|
||||
# Alternative SHM connector with an explicit threshold.
|
||||
shared_memory_connector:
|
||||
name: SharedMemoryConnector
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
devices: "0"
|
||||
max_num_seqs: 1
|
||||
max_num_batched_tokens: 32768
|
||||
enforce_eager: true
|
||||
async_scheduling: false
|
||||
output_connectors:
|
||||
to_stage_1: mooncake_connector
|
||||
|
||||
- stage_id: 1
|
||||
devices: "0"
|
||||
max_num_seqs: 1
|
||||
max_num_batched_tokens: 32768
|
||||
enforce_eager: true
|
||||
async_scheduling: false
|
||||
input_connectors:
|
||||
from_stage_0: mooncake_connector
|
||||
output_connectors:
|
||||
to_stage_2: mooncake_connector
|
||||
|
||||
- stage_id: 2
|
||||
devices: "0"
|
||||
max_num_seqs: 1
|
||||
max_num_batched_tokens: 32768
|
||||
enforce_eager: true
|
||||
async_scheduling: false
|
||||
input_connectors:
|
||||
from_stage_1: mooncake_connector
|
||||
@@ -30,8 +30,7 @@ Two environmental factors mask the race on main:
|
||||
``cached_files`` (plural) which batch-resolves every shard listed in the
|
||||
index up-front via ``os.path.isfile`` and raises immediately if any shard
|
||||
is still sitting under its ``*.incomplete`` name. Same wave of v5 changes
|
||||
that introduced ``tie_weights(missing_keys=..., recompute_mapping=...)``
|
||||
(see the Dynin shim in ``dynin_omni_token2text.py``).
|
||||
that introduced ``tie_weights(missing_keys=..., recompute_mapping=...)``.
|
||||
* CI shares ``HF_HOME=/fsx/hf_cache`` across pipelines (both the
|
||||
``vllm-omni`` and ``vllm-omni-rebase`` pipelines mount the same FS). That
|
||||
cache is normally warm for long-lived repos like ``Qwen-Image-Edit-2509``,
|
||||
@@ -195,19 +194,22 @@ def _repo_prefetch_lock(model: str) -> Iterator[None]:
|
||||
flock_held = False
|
||||
|
||||
# --- fcntl.flock path ---
|
||||
fcntl_mod: Any = None
|
||||
try:
|
||||
import fcntl # type: ignore[import-not-found]
|
||||
except ImportError: # pragma: no cover - non-POSIX (Windows)
|
||||
fcntl = None
|
||||
import fcntl as _fcntl
|
||||
|
||||
if fcntl is not None:
|
||||
fcntl_mod = _fcntl
|
||||
except ImportError: # pragma: no cover - non-POSIX (Windows)
|
||||
pass
|
||||
|
||||
if fcntl_mod is not None:
|
||||
try:
|
||||
lock_dir = _node_lock_dir()
|
||||
except OSError as exc:
|
||||
logger.warning("Could not allocate lock dir for prefetch of %s (%s); skipping flock", model, exc)
|
||||
fcntl = None # force dotfile fallback
|
||||
fcntl_mod = None # force dotfile fallback
|
||||
|
||||
if fcntl is not None and lock_dir is not None:
|
||||
if fcntl_mod is not None and lock_dir is not None:
|
||||
lock_path = os.path.join(lock_dir, _safe_repo_filename(model))
|
||||
try:
|
||||
fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644)
|
||||
@@ -217,7 +219,7 @@ def _repo_prefetch_lock(model: str) -> Iterator[None]:
|
||||
|
||||
if fd is not None:
|
||||
try:
|
||||
fcntl.flock(fd, fcntl.LOCK_EX)
|
||||
fcntl_mod.flock(fd, fcntl_mod.LOCK_EX)
|
||||
flock_held = True
|
||||
logger.info("Acquired flock prefetch lock for %s at %s", model, lock_path)
|
||||
except OSError as exc:
|
||||
@@ -249,9 +251,9 @@ def _repo_prefetch_lock(model: str) -> Iterator[None]:
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if flock_held and fd is not None:
|
||||
if flock_held and fd is not None and fcntl_mod is not None:
|
||||
with contextlib.suppress(OSError):
|
||||
fcntl.flock(fd, fcntl.LOCK_UN)
|
||||
fcntl_mod.flock(fd, fcntl_mod.LOCK_UN)
|
||||
if fd is not None:
|
||||
with contextlib.suppress(OSError):
|
||||
os.close(fd)
|
||||
|
||||
+30
-15
@@ -7,7 +7,7 @@ from __future__ import annotations
|
||||
import importlib
|
||||
import inspect
|
||||
from collections import deque
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
import torch
|
||||
from vllm.distributed.parallel_state import get_tp_group
|
||||
@@ -29,21 +29,23 @@ if TYPE_CHECKING:
|
||||
class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
"""Own payload caching, full/chunk transport, and connector I/O."""
|
||||
|
||||
_full_payload_replace_keys_cached: frozenset[str]
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Local payload cache (RFC §2.4 – Model Runner ownership)
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def put_local_stage_payload(self, req_id: str, payload: OmniPayload) -> None:
|
||||
"""Store a full stage payload in the local cache."""
|
||||
self._local_stage_payload_cache[req_id] = payload
|
||||
self._local_stage_payload_cache[req_id] = cast(dict[str, Any], payload)
|
||||
|
||||
def get_local_stage_payload(self, req_id: str) -> OmniPayload | None:
|
||||
"""Read a stage payload without removing it."""
|
||||
return self._local_stage_payload_cache.get(req_id)
|
||||
return cast(OmniPayload | None, self._local_stage_payload_cache.get(req_id))
|
||||
|
||||
def pop_local_stage_payload(self, req_id: str) -> OmniPayload | None:
|
||||
"""Remove and return a stage payload (consume after use)."""
|
||||
return self._local_stage_payload_cache.pop(req_id, None)
|
||||
return cast(OmniPayload | None, self._local_stage_payload_cache.pop(req_id, None))
|
||||
|
||||
def put_local_request_metadata(self, req_id: str, metadata: dict[str, Any]) -> None:
|
||||
"""Store lightweight scheduling metadata for a request."""
|
||||
@@ -58,7 +60,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
@classmethod
|
||||
def _extract_scheduling_metadata(cls, payload: OmniPayload) -> dict[str, Any]:
|
||||
def _extract_scheduling_metadata(cls, payload: dict[str, Any] | OmniPayload) -> dict[str, Any]:
|
||||
"""Extract only the fields the scheduler needs from a full payload."""
|
||||
extracted: dict[str, Any] = {}
|
||||
meta = payload.get("meta") if isinstance(payload, dict) else None
|
||||
@@ -66,11 +68,13 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
|
||||
if "next_stage_prompt_len" in meta:
|
||||
extracted["next_stage_prompt_len"] = meta["next_stage_prompt_len"]
|
||||
elif "next_stage_prompt_len" in payload:
|
||||
else:
|
||||
legacy_prompt_len = cast(dict[str, Any], payload).get("next_stage_prompt_len")
|
||||
if legacy_prompt_len is not None:
|
||||
logger.warning_once(
|
||||
"legacy flat 'next_stage_prompt_len' key in payload; expected 'meta.next_stage_prompt_len'"
|
||||
)
|
||||
extracted["next_stage_prompt_len"] = payload["next_stage_prompt_len"]
|
||||
extracted["next_stage_prompt_len"] = legacy_prompt_len
|
||||
|
||||
audio_codes = cls._payload_audio_codes(payload)
|
||||
if audio_codes is not None:
|
||||
@@ -129,7 +133,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _payload_is_consumable(cls, payload: OmniPayload | None) -> bool:
|
||||
def _payload_is_consumable(cls, payload: dict[str, Any] | OmniPayload | None) -> bool:
|
||||
"""Return True when an async payload can drive a real forward step.
|
||||
|
||||
Metadata-only wake-ups should not transition WAITING_FOR_CHUNK requests
|
||||
@@ -502,8 +506,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
# Skip the per-step accumulator+build that would otherwise be
|
||||
# silently discarded. Defends against a terminal stage whose
|
||||
# custom_process_input_func has a *_full_payload derivative in
|
||||
# the same module (e.g. dynin stage 2 token2image_to_token2audio
|
||||
# in pipelines that don't configure any connector at all).
|
||||
# the same module, even in pipelines that don't configure any
|
||||
# connector at all.
|
||||
#
|
||||
# Known limitation: a *terminal-consumer* stage that has a
|
||||
# connector configured for receiving upstream input is NOT
|
||||
@@ -556,7 +560,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
output[k] = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0)
|
||||
return output, request
|
||||
|
||||
def _resolve_full_payload_replace_keys(self) -> frozenset:
|
||||
def _resolve_full_payload_replace_keys(self) -> frozenset[str]:
|
||||
"""Per-model REPLACE-key set for the full-payload accumulator.
|
||||
|
||||
Looked up from the stage-input-processor module that ships the model's sync builder
|
||||
@@ -584,7 +588,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
import sys as _sys
|
||||
|
||||
mod = _sys.modules.get(module_name) or importlib.import_module(module_name)
|
||||
keys = getattr(mod, "_FULL_PAYLOAD_REPLACE_KEYS", frozenset())
|
||||
keys: Any = getattr(mod, "_FULL_PAYLOAD_REPLACE_KEYS", frozenset())
|
||||
except ImportError:
|
||||
logger.debug(
|
||||
"Could not import stage input processor module %s while resolving "
|
||||
@@ -748,7 +752,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
code_len = len(audio_codes)
|
||||
else:
|
||||
code_len = None
|
||||
meta = payload.get("meta") if isinstance(payload.get("meta"), dict) else {}
|
||||
raw_meta = payload.get("meta")
|
||||
meta = raw_meta if isinstance(raw_meta, dict) else {}
|
||||
logger.debug(
|
||||
"[Stage-%s] send_full_payload_outputs: req=%s payload_keys=%s code_len=%s left_context_size=%s",
|
||||
self._stage_id,
|
||||
@@ -863,6 +868,9 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
if not self.is_data_transfer_rank():
|
||||
return True
|
||||
raw_req_id = getattr(request, "request_id", None) or getattr(request, "req_id", None)
|
||||
if not isinstance(raw_req_id, str):
|
||||
logger.warning("[Stage-%s] send_chunk: missing request_id", self._stage_id)
|
||||
return False
|
||||
request_id = self._resolve_external_req_id(request, raw_req_id)
|
||||
# Cache the internal→external mapping so that finish sentinels can
|
||||
# resolve the external ID even after the request is freed.
|
||||
@@ -988,6 +996,12 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
"""Re-enqueue a failed send task or drop it after max retries."""
|
||||
retry_count = task.get("_retry_count", 0) + 1
|
||||
req_id = task.get("request_id")
|
||||
if not isinstance(req_id, str):
|
||||
logger.error(
|
||||
"[Stage-%s] Failed send task is missing request_id; dropping it",
|
||||
getattr(self, "_stage_id", "?"),
|
||||
)
|
||||
return
|
||||
if retry_count <= self._MAX_SEND_RETRIES:
|
||||
task["_retry_count"] = retry_count
|
||||
logger.warning(
|
||||
@@ -1268,6 +1282,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
if not success:
|
||||
return False
|
||||
|
||||
if isinstance(request_id, str):
|
||||
self._decrement_pending_save_count(request_id)
|
||||
return True
|
||||
|
||||
@@ -1299,7 +1314,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
"""Accumulate chunk payloads (concat tensors, extend lists)."""
|
||||
if req_id not in self._send_side_request_payload:
|
||||
self._send_side_request_payload[req_id] = dict(payload_data)
|
||||
return dict(self._send_side_request_payload[req_id])
|
||||
return cast(OmniPayload, dict(self._send_side_request_payload[req_id]))
|
||||
|
||||
origin = self._send_side_request_payload[req_id]
|
||||
merged = dict(origin)
|
||||
@@ -1338,4 +1353,4 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
|
||||
merged[key] = value
|
||||
|
||||
self._send_side_request_payload[req_id] = merged
|
||||
return dict(merged)
|
||||
return cast(OmniPayload, dict(merged))
|
||||
@@ -108,7 +108,6 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401
|
||||
breeze_tts_2,
|
||||
cosyvoice3,
|
||||
covo_audio,
|
||||
dots_tts,
|
||||
fish_speech,
|
||||
gepard,
|
||||
glm_tts,
|
||||
|
||||
@@ -1,82 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""dots.tts serving adapter."""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
from vllm.utils.async_utils import make_async
|
||||
|
||||
from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
|
||||
|
||||
|
||||
@register_tts_adapter
|
||||
class DotsTTSAdapter(ARTTSAdapter):
|
||||
"""Adapter for dots.tts (AR ``engine_client`` backend).
|
||||
|
||||
The current integration supports text-only, no-reference speech synthesis.
|
||||
Voice cloning and reference-audio conditioning are not supported yet.
|
||||
"""
|
||||
|
||||
stage_keys = frozenset()
|
||||
model_archs = frozenset({"DotsTTSForConditionalGeneration"})
|
||||
name = "dots_tts"
|
||||
detect_priority = 5
|
||||
|
||||
def __init__(self, ctx):
|
||||
super().__init__(ctx)
|
||||
self.tokenizer = None
|
||||
self._build_prompt_async = None
|
||||
|
||||
def _build_prompt(self, text: str) -> dict:
|
||||
from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import build_dots_tts_prompt
|
||||
|
||||
if self.tokenizer is None:
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
self.ctx.engine_client.model_config.model,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
return build_dots_tts_prompt(self.tokenizer, text)
|
||||
|
||||
def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
|
||||
if not request.input or not request.input.strip():
|
||||
return "Input text cannot be empty"
|
||||
|
||||
if request.voice not in (None, "default"):
|
||||
return "'voice' is not supported for dots.tts"
|
||||
|
||||
if request.ref_audio is not None:
|
||||
return "'ref_audio' is not supported for dots.tts"
|
||||
|
||||
if request.ref_text is not None:
|
||||
return "'ref_text' is not supported for dots.tts"
|
||||
|
||||
if request.speaker_embedding is not None:
|
||||
return "'speaker_embedding' is not supported for dots.tts"
|
||||
|
||||
if request.x_vector_only_mode is not None:
|
||||
return "'x_vector_only_mode' is not supported for dots.tts"
|
||||
return None
|
||||
|
||||
async def build(
|
||||
self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
|
||||
) -> PreparedRequest:
|
||||
if self._build_prompt_async is None:
|
||||
self._build_prompt_async = make_async(
|
||||
self._build_prompt,
|
||||
executor=self.ctx.server._tts_executor,
|
||||
)
|
||||
prompt = await self._build_prompt_async(request.input)
|
||||
return PreparedRequest(prompt=prompt, tts_params={}, model_type="dots_tts")
|
||||
|
||||
def apply_sampling_overrides(
|
||||
self,
|
||||
sampling_params_list: list,
|
||||
request: "OpenAICreateSpeechRequest",
|
||||
prompt: dict | None = None,
|
||||
request_id: str | None = None,
|
||||
) -> list:
|
||||
return apply_max_new_tokens(sampling_params_list, request)
|
||||
@@ -1,8 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# NOTE: Do not import model classes in this file. Importing any
|
||||
# submodule in this package triggers __init__.py execution, and
|
||||
# both the model registry and pipeline registry import submodules
|
||||
# directly — heavy imports here would be loaded as a side effect
|
||||
# even though nothing depends on these re-exports.
|
||||
@@ -1,457 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Copyright 2026 The vLLM-Omni team.
|
||||
# Copyright (c) rednote-hilab. All rights reserved.
|
||||
# Adapted from:
|
||||
# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/backbone
|
||||
#
|
||||
# Merged from two upstream files (preserved as-is, inference behaviour
|
||||
# unchanged; only internal imports were folded inline):
|
||||
# - modules/backbone/layers.py -> Dropout, Mlp, rotate_half,
|
||||
# apply_rotary_pos_emb, RotaryEmbedding,
|
||||
# MultiHeadAttention
|
||||
# (Conv1d / ConvTranspose1d are NOT
|
||||
# re-vendored here — they already live in
|
||||
# dots_tts_vocoder.py)
|
||||
# - modules/backbone/dit.py -> modulate, TimestepEmbedder, FinalLayer,
|
||||
# DiTBlock, DiT
|
||||
#
|
||||
# Per-section upstream attributions are preserved at each section header.
|
||||
# Nothing in this file diverges from upstream behaviour.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
|
||||
# ============================================================================
|
||||
# Building blocks (adapted from modules/backbone/layers.py)
|
||||
# ----------------------------------------------------------------------------
|
||||
# Conv1d / ConvTranspose1d from layers.py are already vendored in
|
||||
# dots_tts_vocoder.py (they are used by the BigVGAN-style AudioVAE). The
|
||||
# pieces below are the ones consumed by the DiT flow-matching head.
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class Dropout(nn.Module):
|
||||
def __init__(self, p: float = 0.5, inplace: bool = False, force_drop: bool = False, **_kwargs):
|
||||
super().__init__()
|
||||
if p < 0.0 or p > 1.0:
|
||||
raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
|
||||
self.p = p
|
||||
self.inplace = inplace
|
||||
self.force_drop = force_drop
|
||||
|
||||
def forward(self, x, **_kwargs):
|
||||
return F.dropout(
|
||||
x,
|
||||
p=self.p,
|
||||
training=True if self.force_drop else self.training,
|
||||
inplace=self.inplace,
|
||||
)
|
||||
|
||||
|
||||
class Mlp(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
ffn_hidden_size=4096,
|
||||
act_layer=nn.GELU,
|
||||
dropout=0.0,
|
||||
**_kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(hidden_size, ffn_hidden_size)
|
||||
self.act = act_layer()
|
||||
self.fc2 = nn.Linear(ffn_hidden_size, hidden_size)
|
||||
self.drop = Dropout(dropout)
|
||||
|
||||
def forward(self, x, _mask=None):
|
||||
x = self.fc1(x)
|
||||
x = self.act(x)
|
||||
x = self.drop(x)
|
||||
x = self.fc2(x)
|
||||
return self.drop(x)
|
||||
|
||||
|
||||
def rotate_half(x):
|
||||
x1, x2 = x.chunk(2, dim=-1)
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
@torch.autocast(enabled=False, device_type="cuda")
|
||||
def apply_rotary_pos_emb(pos, t):
|
||||
if pos.dim() == 3:
|
||||
pos = pos.unsqueeze(1)
|
||||
return t * pos.cos() + rotate_half(t) * pos.sin()
|
||||
|
||||
|
||||
class RotaryEmbedding(nn.Module):
|
||||
def __init__(self, dim, theta=50000):
|
||||
super().__init__()
|
||||
self.register_buffer(
|
||||
"inv_freq",
|
||||
1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)),
|
||||
persistent=False,
|
||||
)
|
||||
self._theta = float(theta)
|
||||
|
||||
def _apply(self, fn):
|
||||
inv_freq = self.inv_freq
|
||||
super()._apply(fn)
|
||||
self.inv_freq = inv_freq.to(device=self.inv_freq.device, dtype=torch.float32)
|
||||
return self
|
||||
|
||||
@torch.autocast(enabled=False, device_type="cuda")
|
||||
def forward(self, t):
|
||||
inv_freq = self.inv_freq
|
||||
if inv_freq.device != t.device:
|
||||
raise RuntimeError(f"RotaryEmbedding buffer device mismatch: inv_freq={inv_freq.device} input={t.device}.")
|
||||
t = t.to(dtype=inv_freq.dtype)
|
||||
if t.dim() == 1:
|
||||
freqs = torch.einsum("i , j -> i j", t, inv_freq)
|
||||
else:
|
||||
freqs = torch.einsum("bi, j -> bij", t, inv_freq)
|
||||
return torch.cat((freqs, freqs), dim=-1)
|
||||
|
||||
|
||||
class MultiHeadAttention(nn.Module):
|
||||
"""Multi-head attention"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
num_heads: int = 8,
|
||||
qkv_bias: bool = False,
|
||||
qk_norm: bool = False,
|
||||
attn_drop: float = 0.0,
|
||||
dropout: float = 0.0,
|
||||
norm_layer: str = "LayerNorm",
|
||||
rotary_bias: bool = False,
|
||||
rotary_theta: float | None = 50000,
|
||||
**_kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
assert hidden_size % num_heads == 0, "hidden_size should be divisible by num_heads"
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = hidden_size // num_heads
|
||||
self.scale = self.head_dim**-0.5
|
||||
self.rotary_bias = rotary_bias
|
||||
|
||||
self.q_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
|
||||
self.k_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
|
||||
self.v_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
|
||||
|
||||
norm_layer = getattr(nn, norm_layer)
|
||||
self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
|
||||
self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
|
||||
|
||||
self.attn_drop = Dropout(attn_drop)
|
||||
self.o_proj = nn.Linear(hidden_size, hidden_size)
|
||||
self.o_dropout = Dropout(dropout)
|
||||
|
||||
if self.rotary_bias:
|
||||
self.rotary = RotaryEmbedding(self.head_dim, theta=rotary_theta)
|
||||
|
||||
def forward(self, q, k=None, v=None, mask=None, pos_ids=None, **_kwargs):
|
||||
k = k or q
|
||||
v = v or q
|
||||
B, L, _ = q.shape
|
||||
_, S, _ = v.shape
|
||||
if mask is not None:
|
||||
if mask.ndim == 2: # [B, L]
|
||||
assert L == S
|
||||
mask = rearrange(mask, "b j -> b 1 1 j")
|
||||
mask = mask.expand(-1, self.num_heads, L, -1)
|
||||
elif mask.ndim == 3: # [B, L, S]
|
||||
assert mask.size(1) == L and mask.size(2) == S
|
||||
mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
|
||||
|
||||
q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)
|
||||
q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
|
||||
k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads)
|
||||
v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads)
|
||||
q, k = self.q_norm(q), self.k_norm(k)
|
||||
|
||||
# Apply rotary
|
||||
if self.rotary_bias:
|
||||
if L == S:
|
||||
if pos_ids is None:
|
||||
rotary_emb = self.rotary(torch.arange(L, device=q.device))
|
||||
else:
|
||||
rotary_emb = self.rotary(pos_ids)
|
||||
q, k = (apply_rotary_pos_emb(rotary_emb, tensor) for tensor in (q, k))
|
||||
else:
|
||||
q_rotary_emb = self.rotary(torch.arange(L, device=q.device))
|
||||
k_rotary_emb = self.rotary(torch.arange(S, device=k.device))
|
||||
q = apply_rotary_pos_emb(q_rotary_emb, q)
|
||||
k = apply_rotary_pos_emb(k_rotary_emb, k)
|
||||
|
||||
attn_bias = torch.zeros(B, self.num_heads, L, S, dtype=q.dtype, device=q.device)
|
||||
|
||||
if mask is not None:
|
||||
attn_bias.masked_fill_(mask.logical_not(), float("-inf"))
|
||||
|
||||
out = F.scaled_dot_product_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
attn_mask=attn_bias,
|
||||
dropout_p=self.attn_drop.p if self.training else 0.0,
|
||||
)
|
||||
|
||||
out = rearrange(out, "b h n d -> b n (h d)")
|
||||
return self.o_dropout(self.o_proj(out))
|
||||
|
||||
def decode_step(self, x, *, cache, positions: torch.Tensor):
|
||||
if x.size(1) <= 0:
|
||||
raise ValueError("MultiHeadAttention.decode_step expects a non-empty input.")
|
||||
if positions.ndim != 1 or positions.size(0) != x.size(1):
|
||||
raise ValueError("MultiHeadAttention.decode_step positions must match the decode block length.")
|
||||
|
||||
q = self.q_proj(x)
|
||||
k = self.k_proj(x)
|
||||
v = self.v_proj(x)
|
||||
|
||||
q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
|
||||
k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads)
|
||||
v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads)
|
||||
q, k = self.q_norm(q), self.k_norm(k)
|
||||
block_len = q.size(2)
|
||||
|
||||
if self.rotary_bias:
|
||||
rotary_emb = self.rotary(positions)
|
||||
q = apply_rotary_pos_emb(rotary_emb, q)
|
||||
k = apply_rotary_pos_emb(rotary_emb, k)
|
||||
|
||||
cached_k, cached_v = cache
|
||||
cached_k.index_copy_(2, positions, k)
|
||||
cached_v.index_copy_(2, positions, v)
|
||||
|
||||
cache_capacity = cached_k.size(2)
|
||||
key_positions = torch.arange(
|
||||
cache_capacity,
|
||||
device=x.device,
|
||||
dtype=torch.long,
|
||||
).unsqueeze(0)
|
||||
query_positions = positions.unsqueeze(1)
|
||||
causal_mask = key_positions <= query_positions
|
||||
valid_mask = key_positions <= positions[-1]
|
||||
attn_bias = torch.zeros(
|
||||
q.size(0),
|
||||
self.num_heads,
|
||||
block_len,
|
||||
cache_capacity,
|
||||
dtype=q.dtype,
|
||||
device=q.device,
|
||||
)
|
||||
attn_bias.masked_fill_(
|
||||
(causal_mask & valid_mask).unsqueeze(0).unsqueeze(0).logical_not(),
|
||||
float("-inf"),
|
||||
)
|
||||
|
||||
out = F.scaled_dot_product_attention(
|
||||
q,
|
||||
cached_k,
|
||||
cached_v,
|
||||
attn_mask=attn_bias,
|
||||
dropout_p=self.attn_drop.p if self.training else 0.0,
|
||||
)
|
||||
out = rearrange(out, "b h n d -> b n (h d)")
|
||||
return self.o_dropout(self.o_proj(out)), cache
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# DiT flow-matching head (adapted from modules/backbone/dit.py)
|
||||
# ----------------------------------------------------------------------------
|
||||
# DiT.__init__ takes a ``transformer_config`` object exposing ``.to_dict()``,
|
||||
# ``.hidden_size`` and ``.num_layers``. Upstream constructs it from a pydantic
|
||||
# config class higher up the stack; vLLM-Omni callers can pass anything that
|
||||
# satisfies the same minimal interface (e.g. a stdlib @dataclass with a
|
||||
# ``to_dict`` method or SimpleNamespace). See dots_tts_talker.py for the
|
||||
# concrete soar-checkpoint config factory.
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def modulate(x, shift, scale, **_kwargs):
|
||||
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
|
||||
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
def __init__(self, hidden_size, frequency_embedding_size=256):
|
||||
super().__init__()
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, hidden_size, bias=True),
|
||||
)
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
|
||||
@staticmethod
|
||||
def timestep_embedding(t, dim, max_period=10000):
|
||||
half = dim // 2
|
||||
freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
|
||||
device=t.device
|
||||
)
|
||||
args = t[:, None].float() * freqs[None]
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
if dim % 2:
|
||||
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
||||
return embedding
|
||||
|
||||
def forward(self, t):
|
||||
t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
|
||||
return self.mlp(t_freq)
|
||||
|
||||
|
||||
class FinalLayer(nn.Module):
|
||||
def __init__(self, hidden_size, output_size):
|
||||
super().__init__()
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, 2 * hidden_size, bias=True),
|
||||
)
|
||||
self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-5)
|
||||
self.linear = nn.Linear(hidden_size, output_size, bias=True)
|
||||
|
||||
def forward(self, x, c, **_kwargs):
|
||||
shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
|
||||
x = modulate(self.norm(x), shift, scale)
|
||||
return self.linear(x)
|
||||
|
||||
|
||||
class DiTBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
attention: nn.Module,
|
||||
ffn: nn.Module,
|
||||
hidden_size: int = 1024,
|
||||
modulation: bool = False,
|
||||
eps: float = 1e-5,
|
||||
**_kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps)
|
||||
self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps)
|
||||
self.attn = attention
|
||||
self.ffn = ffn
|
||||
self.modulation = modulation
|
||||
if modulation:
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, 6 * hidden_size, bias=True),
|
||||
)
|
||||
|
||||
def forward(self, x, condition=None, mask=None, **kwargs):
|
||||
if condition is None:
|
||||
assert not self.modulation, "Without global condition, must set modulation to False"
|
||||
else:
|
||||
assert self.modulation, "With global condition, must set modulation to True"
|
||||
shift_attn, scale_attn, gate_attn, shift_ffn, scale_ffn, gate_ffn = self.adaLN_modulation(condition).chunk(
|
||||
6, dim=1
|
||||
)
|
||||
|
||||
if condition is not None:
|
||||
pack_indices = kwargs.get("pack_indices")
|
||||
if pack_indices is not None:
|
||||
gate_attn = gate_attn[pack_indices]
|
||||
gate_ffn = gate_ffn[pack_indices]
|
||||
else:
|
||||
gate_attn = gate_attn.unsqueeze(1)
|
||||
gate_ffn = gate_ffn.unsqueeze(1)
|
||||
|
||||
if condition is not None:
|
||||
x = x + gate_attn * self.attn(
|
||||
modulate(self.norm1(x), shift_attn, scale_attn, **kwargs),
|
||||
mask=mask,
|
||||
**kwargs,
|
||||
)
|
||||
else:
|
||||
x = x + self.attn(self.norm1(x), mask=mask, **kwargs)
|
||||
|
||||
if condition is not None:
|
||||
x = x + gate_ffn * self.ffn(modulate(self.norm2(x), shift_ffn, scale_ffn, **kwargs))
|
||||
else:
|
||||
x = x + self.ffn(self.norm2(x), mask=mask)
|
||||
return x
|
||||
|
||||
|
||||
class DiT(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_dim,
|
||||
out_dim,
|
||||
transformer_config,
|
||||
*,
|
||||
mode: str = "flow_matching",
|
||||
):
|
||||
super().__init__()
|
||||
if mode not in {"flow_matching", "meanflow"}:
|
||||
raise ValueError(f"DiT mode must be 'flow_matching' or 'meanflow', got {mode!r}.")
|
||||
|
||||
transformer_kwargs = transformer_config.to_dict()
|
||||
model_dim = transformer_config.hidden_size
|
||||
self.mode = mode
|
||||
self.num_layers = transformer_config.num_layers
|
||||
|
||||
self.input_layer = nn.Linear(in_dim, model_dim)
|
||||
self.time_embedder = TimestepEmbedder(model_dim)
|
||||
if mode == "meanflow":
|
||||
self.duration_embedder = TimestepEmbedder(model_dim)
|
||||
|
||||
self.blocks = nn.ModuleList()
|
||||
for i in range(self.num_layers):
|
||||
attn_block = MultiHeadAttention(**transformer_kwargs, name=f"layer_{i}")
|
||||
ffn_block = Mlp(act_layer=lambda: nn.GELU(approximate="tanh"), **transformer_kwargs)
|
||||
self.blocks.append(DiTBlock(attention=attn_block, ffn=ffn_block, **transformer_kwargs))
|
||||
|
||||
self.output_layer = FinalLayer(model_dim, out_dim)
|
||||
self.initialize_weights()
|
||||
|
||||
def initialize_weights(self):
|
||||
def _basic_init(module):
|
||||
if isinstance(module, nn.Linear):
|
||||
torch.nn.init.xavier_uniform_(module.weight)
|
||||
if module.bias is not None:
|
||||
nn.init.constant_(module.bias, 0)
|
||||
|
||||
self.apply(_basic_init)
|
||||
|
||||
nn.init.normal_(self.time_embedder.mlp[0].weight, std=0.02)
|
||||
nn.init.normal_(self.time_embedder.mlp[2].weight, std=0.02)
|
||||
|
||||
for block in self.blocks:
|
||||
if hasattr(block, "adaLN_modulation"):
|
||||
nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
|
||||
nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
|
||||
|
||||
nn.init.constant_(self.output_layer.adaLN_modulation[-1].weight, 0)
|
||||
nn.init.constant_(self.output_layer.adaLN_modulation[-1].bias, 0)
|
||||
nn.init.constant_(self.output_layer.linear.weight, 0)
|
||||
nn.init.constant_(self.output_layer.linear.bias, 0)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
timesteps,
|
||||
duration: torch.Tensor | None = None,
|
||||
mask=None,
|
||||
attn_mask=None,
|
||||
g_cond: torch.Tensor | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
t = self.time_embedder(timesteps)
|
||||
c = t
|
||||
duration_embedder = getattr(self, "duration_embedder", None)
|
||||
if duration_embedder is not None and duration is not None:
|
||||
c = c + duration_embedder(duration)
|
||||
if g_cond is not None:
|
||||
c = c + g_cond
|
||||
|
||||
x = self.input_layer(x)
|
||||
for block in self.blocks:
|
||||
x = block(x, c, mask=attn_mask, **kwargs)
|
||||
return self.output_layer(x, c, **kwargs)
|
||||
@@ -1,355 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Copyright 2026 The vLLM-Omni team.
|
||||
# Copyright (c) rednote-hilab. All rights reserved.
|
||||
# Adapted from:
|
||||
# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/backbone
|
||||
#
|
||||
# Single upstream source (preserved as-is, inference behaviour unchanged;
|
||||
# only the internal `from dots_tts.modules.backbone.layers import ...` line
|
||||
# was folded to point at our sibling vendor files):
|
||||
# - modules/backbone/semantic_encoder.py -> SemanticEncoderDecodeState,
|
||||
# TransformerEncoderLayer,
|
||||
# SuperviseEncoder,
|
||||
# VAESemanticEncoder
|
||||
#
|
||||
# Building blocks (Conv1d / Mlp / MultiHeadAttention) are imported from the
|
||||
# sibling vendor files where they already live:
|
||||
# - Conv1d <- dots_tts_vocoder.py
|
||||
# - Mlp <- dots_tts_dit.py
|
||||
# - MultiHeadAttention <- dots_tts_dit.py
|
||||
#
|
||||
# Nothing else in this file diverges from upstream behaviour.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
|
||||
from vllm_omni.model_executor.models.dots_tts.dots_tts_dit import (
|
||||
Mlp,
|
||||
MultiHeadAttention,
|
||||
)
|
||||
from vllm_omni.model_executor.models.dots_tts.dots_tts_vocoder import Conv1d
|
||||
|
||||
|
||||
@dataclass
|
||||
class SemanticEncoderDecodeState:
|
||||
conv_tail: torch.Tensor
|
||||
layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...]
|
||||
seq_len: int
|
||||
|
||||
|
||||
class TransformerEncoderLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
num_heads=16,
|
||||
ffn_hidden_size=4096,
|
||||
attn_dropout=0.0,
|
||||
ffn_dropout=0.0,
|
||||
norm_layer="LayerNorm",
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.attn = MultiHeadAttention(
|
||||
hidden_size,
|
||||
num_heads,
|
||||
attn_drop=attn_dropout,
|
||||
norm_layer=norm_layer,
|
||||
**kwargs,
|
||||
)
|
||||
norm_cls = getattr(nn, norm_layer)
|
||||
self.attn_norm = norm_cls(hidden_size)
|
||||
self.ffn = Mlp(hidden_size, ffn_hidden_size, dropout=ffn_dropout, act_layer=nn.SiLU)
|
||||
self.ffn_norm = norm_cls(hidden_size)
|
||||
self.hidden_size = hidden_size
|
||||
|
||||
def _build_causal_mask(self, T: int, device): # noqa: N803
|
||||
return torch.tril(torch.ones(T, T, dtype=torch.bool, device=device))
|
||||
|
||||
def _build_padding_mask(self, x_lens, max_len: int, device):
|
||||
B = x_lens.size(0)
|
||||
positions = torch.arange(max_len, device=device).unsqueeze(0).expand(B, -1)
|
||||
return positions < x_lens.unsqueeze(1)
|
||||
|
||||
def _fuse_attn_mask(self, causal_mask, padding_mask):
|
||||
if causal_mask is None and padding_mask is None:
|
||||
return None
|
||||
if causal_mask is None:
|
||||
row = padding_mask.unsqueeze(2)
|
||||
col = padding_mask.unsqueeze(1)
|
||||
return row & col
|
||||
if padding_mask is None:
|
||||
return causal_mask.unsqueeze(0)
|
||||
|
||||
_B, _T = padding_mask.shape
|
||||
causal = causal_mask.unsqueeze(0)
|
||||
row = padding_mask.unsqueeze(2)
|
||||
col = padding_mask.unsqueeze(1)
|
||||
pad_2d = row & col
|
||||
return causal & pad_2d
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
x_lens=None,
|
||||
causal=True,
|
||||
):
|
||||
_B, T, C = x.shape
|
||||
assert self.hidden_size == C
|
||||
device = x.device
|
||||
|
||||
causal_mask = self._build_causal_mask(T, device) if causal else None
|
||||
if x_lens is not None:
|
||||
padding_mask = self._build_padding_mask(x_lens, T, device)
|
||||
else:
|
||||
padding_mask = None
|
||||
fused_mask = self._fuse_attn_mask(causal_mask, padding_mask)
|
||||
|
||||
h = self.attn_norm(x)
|
||||
h = self.attn(
|
||||
q=h,
|
||||
mask=fused_mask,
|
||||
)
|
||||
x = x + h
|
||||
|
||||
h = self.ffn_norm(x)
|
||||
h = self.ffn(h)
|
||||
return x + h
|
||||
|
||||
def decode_step(
|
||||
self,
|
||||
x,
|
||||
*,
|
||||
cache: tuple[torch.Tensor, torch.Tensor],
|
||||
positions: torch.Tensor,
|
||||
):
|
||||
if x.size(1) <= 0:
|
||||
raise ValueError("TransformerEncoderLayer.decode_step expects a non-empty input.")
|
||||
|
||||
h = self.attn_norm(x)
|
||||
h, cache = self.attn.decode_step(h, cache=cache, positions=positions)
|
||||
x = x + h
|
||||
|
||||
h = self.ffn_norm(x)
|
||||
h = self.ffn(h)
|
||||
return x + h, cache
|
||||
|
||||
|
||||
class SuperviseEncoder(nn.Module):
|
||||
def __init__(self, config):
|
||||
super().__init__()
|
||||
self.hidden_size = config.get("hidden_size", 1024)
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
TransformerEncoderLayer(
|
||||
hidden_size=self.hidden_size,
|
||||
num_heads=config.get("num_heads", 16),
|
||||
ffn_hidden_size=config.get("ffn_hidden_size", 4096),
|
||||
norm_layer=config.get("norm_layer", "LayerNorm"),
|
||||
)
|
||||
for _ in range(config.get("num_layers", 6))
|
||||
]
|
||||
)
|
||||
self.causal = config.get("causal", False)
|
||||
|
||||
def forward(self, x, x_lens=None):
|
||||
batch_size, seq_len, _ = x.shape
|
||||
if x_lens is None:
|
||||
x_lens = torch.full((batch_size,), seq_len, device=x.device, dtype=torch.long)
|
||||
for layer in self.layers:
|
||||
x = layer(x, x_lens=x_lens, causal=self.causal)
|
||||
return x
|
||||
|
||||
def init_decode_state(
|
||||
self,
|
||||
*,
|
||||
batch_size: int,
|
||||
max_seq_len: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
layer_caches = []
|
||||
for layer in self.layers:
|
||||
cache_shape = (
|
||||
batch_size,
|
||||
layer.attn.num_heads,
|
||||
max_seq_len,
|
||||
layer.attn.head_dim,
|
||||
)
|
||||
layer_caches.append(
|
||||
(
|
||||
torch.zeros(cache_shape, dtype=dtype, device=device),
|
||||
torch.zeros(cache_shape, dtype=dtype, device=device),
|
||||
)
|
||||
)
|
||||
return tuple(layer_caches)
|
||||
|
||||
def reset_decode_state(
|
||||
self,
|
||||
layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...],
|
||||
) -> None:
|
||||
if len(layer_caches) != len(self.layers):
|
||||
raise ValueError("Layer cache count does not match encoder depth.")
|
||||
for key_cache, value_cache in layer_caches:
|
||||
key_cache.zero_()
|
||||
value_cache.zero_()
|
||||
|
||||
def decode_step(self, x, *, layer_caches, positions: torch.Tensor):
|
||||
if len(layer_caches) != len(self.layers):
|
||||
raise ValueError("Layer cache count does not match encoder depth.")
|
||||
|
||||
for layer, cache in zip(self.layers, layer_caches, strict=True):
|
||||
x, _ = layer.decode_step(x, cache=cache, positions=positions)
|
||||
return x
|
||||
|
||||
|
||||
class VAESemanticEncoder(nn.Module):
|
||||
def __init__(self, in_dim, out_dim, config):
|
||||
super().__init__()
|
||||
in_ds_rate = 2
|
||||
self.patch_size = int(config.patch_size)
|
||||
self.in_ds_rate = in_ds_rate
|
||||
self.ds_proj = Conv1d(in_dim, in_dim, kernel_size=in_ds_rate, stride=in_ds_rate, causal=True)
|
||||
self.in_proj = nn.Linear(in_dim, config.PatchEncoder.hidden_size)
|
||||
self.encoder = SuperviseEncoder(config.PatchEncoder)
|
||||
self.out_ds_rate = self.patch_size // in_ds_rate
|
||||
self.out_proj = nn.Linear(config.PatchEncoder.hidden_size * self.out_ds_rate, out_dim)
|
||||
|
||||
def forward(self, x, x_lens=None):
|
||||
x = self._downsample(x)
|
||||
x = self.in_proj(x)
|
||||
z = self.encoder(x, x_lens=x_lens)
|
||||
return self._project_embeddings(z)
|
||||
|
||||
def init_decode_state(
|
||||
self,
|
||||
*,
|
||||
max_audio_patch_count: int,
|
||||
batch_size: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> SemanticEncoderDecodeState:
|
||||
return SemanticEncoderDecodeState(
|
||||
conv_tail=torch.zeros(
|
||||
(batch_size, self.ds_proj.in_channels, self.ds_proj.left_padding),
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
),
|
||||
layer_caches=self.encoder.init_decode_state(
|
||||
batch_size=batch_size,
|
||||
max_seq_len=max_audio_patch_count * self.out_ds_rate,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
),
|
||||
seq_len=0,
|
||||
)
|
||||
|
||||
def reset_decode_state(self, state: SemanticEncoderDecodeState) -> None:
|
||||
state.conv_tail.zero_()
|
||||
self.encoder.reset_decode_state(state.layer_caches)
|
||||
state.seq_len = 0
|
||||
|
||||
def prefill(
|
||||
self,
|
||||
x,
|
||||
state: SemanticEncoderDecodeState,
|
||||
) -> tuple[torch.Tensor, SemanticEncoderDecodeState]:
|
||||
if x.ndim != 3:
|
||||
raise ValueError(f"VAESemanticEncoder.prefill expects rank-3 input, got {tuple(x.shape)}.")
|
||||
if x.size(1) % self.patch_size != 0:
|
||||
raise ValueError(f"Prompt latent length {x.size(1)} must be divisible by patch_size={self.patch_size}.")
|
||||
|
||||
if x.size(1) == 0:
|
||||
return (
|
||||
x.new_zeros((x.size(0), 0, self.out_proj.out_features)),
|
||||
state,
|
||||
)
|
||||
if state.conv_tail.size(0) != x.size(0):
|
||||
raise ValueError("VAESemanticEncoder.prefill batch size does not match decode state.")
|
||||
|
||||
step_inputs = self.in_proj(self._downsample(x))
|
||||
expected_token_count = (x.size(1) // self.patch_size) * self.out_ds_rate
|
||||
if step_inputs.size(1) != expected_token_count:
|
||||
raise RuntimeError(
|
||||
"Patch encoder prefill produced an unexpected token count: "
|
||||
f"expected={expected_token_count} actual={step_inputs.size(1)}."
|
||||
)
|
||||
|
||||
current_seq_len = state.seq_len
|
||||
next_seq_len = current_seq_len + step_inputs.size(1)
|
||||
cache_capacity = state.layer_caches[0][0].size(2)
|
||||
if next_seq_len > cache_capacity:
|
||||
raise ValueError(
|
||||
"Patch encoder prefill exceeds decode-state capacity: "
|
||||
f"required={next_seq_len} capacity={cache_capacity}."
|
||||
)
|
||||
|
||||
positions = torch.arange(step_inputs.size(1), device=x.device, dtype=torch.long) + current_seq_len
|
||||
encoded = self.encoder.decode_step(
|
||||
step_inputs,
|
||||
layer_caches=state.layer_caches,
|
||||
positions=positions,
|
||||
)
|
||||
embedding = self._project_embeddings(encoded)
|
||||
raw = x.transpose(1, 2)
|
||||
state.conv_tail.copy_(raw[..., -self.ds_proj.left_padding :])
|
||||
state.seq_len = next_seq_len
|
||||
return embedding, state
|
||||
|
||||
def decode_patch(
|
||||
self,
|
||||
latent_patch,
|
||||
conv_tail: torch.Tensor,
|
||||
layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...],
|
||||
positions: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if latent_patch.ndim != 3:
|
||||
raise ValueError(f"VAESemanticEncoder.decode_patch expects rank-3 input, got {tuple(latent_patch.shape)}.")
|
||||
if latent_patch.size(1) != self.patch_size:
|
||||
raise ValueError(f"decode_patch expects patch length {self.patch_size}, got {latent_patch.size(1)}.")
|
||||
if positions.ndim != 1 or positions.size(0) != self.out_ds_rate:
|
||||
raise ValueError("decode_patch positions must be a rank-1 tensor matching out_ds_rate.")
|
||||
|
||||
step_inputs, conv_tail = self._downsample_step(
|
||||
latent_patch,
|
||||
conv_tail=conv_tail,
|
||||
)
|
||||
if step_inputs.size(1) != self.out_ds_rate:
|
||||
raise RuntimeError(f"Downsample step produced {step_inputs.size(1)} tokens, expected {self.out_ds_rate}.")
|
||||
|
||||
encoded = self.encoder.decode_step(
|
||||
step_inputs,
|
||||
layer_caches=layer_caches,
|
||||
positions=positions,
|
||||
)
|
||||
embedding = self._project_embeddings(encoded)
|
||||
return embedding, conv_tail
|
||||
|
||||
def _downsample(self, x):
|
||||
return self.ds_proj(x.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
def _project_embeddings(self, z):
|
||||
if self.out_ds_rate > 1:
|
||||
z = rearrange(z, "b (s d) h -> b s (d h)", d=self.out_ds_rate)
|
||||
return self.out_proj(z)
|
||||
|
||||
def _downsample_step(self, latent_patch, *, conv_tail):
|
||||
raw = latent_patch.transpose(1, 2)
|
||||
conv_input = torch.cat([conv_tail, raw], dim=-1)
|
||||
|
||||
projected = F.conv1d(
|
||||
conv_input,
|
||||
self.ds_proj.weight,
|
||||
self.ds_proj.bias,
|
||||
stride=self.ds_proj.stride[0],
|
||||
padding=0,
|
||||
dilation=self.ds_proj.dilation[0],
|
||||
groups=self.ds_proj.groups,
|
||||
).transpose(1, 2)
|
||||
new_conv_tail = raw[..., -self.ds_proj.left_padding :]
|
||||
return self.in_proj(projected), new_conv_tail
|
||||
@@ -1,70 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""dots.tts prompt builder for vllm-omni's Omni engine + serving layer.
|
||||
|
||||
Mirrors upstream ``build_generation_schedule`` (rednote-hilab/dots.tts @
|
||||
a393d2e data/pipelines/tokenizing.py:169) for the
|
||||
no-prompt-audio + audio-placeholder path — text-only zero-shot TTS.
|
||||
|
||||
V1 returns real Qwen2 token IDs (not placeholders ``[1] * N`` like
|
||||
voxcpm2's builder). Reason: our talker's preprocess() runs
|
||||
``self.model.embed_tokens(input_ids)`` directly, so real IDs flow
|
||||
through naturally. voxcpm2 uses placeholders because its prefill mixes
|
||||
reference-audio embeddings (no token IDs); dots.tts V1 zero-shot is
|
||||
purely text + audio_gen_start, all real IDs.
|
||||
|
||||
Voice clone (reference audio + speaker x-vector) is intentionally not
|
||||
wired here — V1 is zero-shot only. Add reference_audio /
|
||||
voice_profile / prompt_text branches in a follow-up (mirror voxcpm2
|
||||
``build_voxcpm2_prompt``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from vllm.inputs import tokens_input
|
||||
|
||||
# Upstream rednote-hilab/dots.tts @ a393d2e
|
||||
# data/pipelines/tts_pipeline.py:17-18
|
||||
_TTS_TEXT_PREFIX = "[文本]"
|
||||
_TTS_AUDIO_PREFIX = "[文本对应语音]"
|
||||
_AUDIO_GEN_START_TOKEN = "<|audio_gen_start|>"
|
||||
|
||||
|
||||
def build_dots_tts_prompt(
|
||||
tokenizer: Any,
|
||||
text: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Build a dots.tts prefill prompt dict for ``Omni.generate()``.
|
||||
|
||||
Returns ``{"prompt_token_ids": [...], "additional_information": {...}}``
|
||||
where ``prompt_token_ids`` is the real Qwen2 token sequence:
|
||||
|
||||
[text_tokens of "[文本]<text>[文本对应语音]", audio_gen_start_id]
|
||||
|
||||
talker's preprocess() embeds these into ``hidden_states`` for the AR
|
||||
base LM. Decode then drives via ``state.curr_embed_for_next`` (DiT
|
||||
→ patch_encoder AR loopback).
|
||||
"""
|
||||
if not isinstance(text, str) or not text.strip():
|
||||
raise ValueError(f"text must be a non-empty string, got {text!r}")
|
||||
|
||||
prompt_text = f"{_TTS_TEXT_PREFIX}{text}{_TTS_AUDIO_PREFIX}"
|
||||
text_ids = tokenizer.encode(prompt_text, add_special_tokens=False)
|
||||
|
||||
audio_gen_start_id = tokenizer.convert_tokens_to_ids(_AUDIO_GEN_START_TOKEN)
|
||||
if not isinstance(audio_gen_start_id, int) or audio_gen_start_id == tokenizer.unk_token_id:
|
||||
raise ValueError(
|
||||
f"Tokenizer does not know {_AUDIO_GEN_START_TOKEN!r} "
|
||||
f"(got id={audio_gen_start_id!r}). Did you load the "
|
||||
f"dots.tts tokenizer (with added_tokens.json)?"
|
||||
)
|
||||
prompt_token_ids = list(text_ids) + [audio_gen_start_id]
|
||||
|
||||
prompt = tokens_input(prompt_token_ids=prompt_token_ids)
|
||||
prompt["additional_information"] = {
|
||||
# Reserved for future voice clone (ref_audio / speaker x-vector).
|
||||
# V1 zero-shot leaves this empty; preprocess() ignores it.
|
||||
}
|
||||
return prompt
|
||||
@@ -1,732 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Copyright 2026 The vLLM-Omni team.
|
||||
# Copyright (c) rednote-hilab. All rights reserved.
|
||||
# Portions Copyright 3D-Speaker
|
||||
# (https://github.com/alibaba-damo-academy/3D-Speaker), Apache-2.0.
|
||||
# Adapted from:
|
||||
# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/speaker
|
||||
# (whose CAM++ implementation is itself adapted from 3D-Speaker).
|
||||
#
|
||||
# Merged from four upstream files (preserved as-is, inference behaviour
|
||||
# unchanged; only internal cross-module imports were folded inline):
|
||||
# - modules/speaker/encoder.py -> SpeakerXVectorFeatures
|
||||
# - modules/speaker/fbank.py -> fbank constants + extract_speaker_fbank
|
||||
# - modules/speaker/campplus.py -> FCM, CAMPPlus
|
||||
# - modules/speaker/campplus_layers.py -> get_nonlinear, statistics_pooling,
|
||||
# StatsPool, TDNNLayer, CAMLayer,
|
||||
# CAMDenseTDNNLayer, CAMDenseTDNNBlock,
|
||||
# TransitLayer, DenseLayer, BasicResBlock
|
||||
#
|
||||
# Plus two tiny helpers folded from utils/audio.py (used by
|
||||
# extract_speaker_fbank):
|
||||
# - high_quality_resample, extract_fbank
|
||||
#
|
||||
# Nothing in this file diverges from upstream behaviour.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import random
|
||||
from collections import OrderedDict
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.checkpoint as cp
|
||||
import torchaudio
|
||||
import torchaudio.compliance.kaldi as Kaldi
|
||||
import torchaudio.functional as AF
|
||||
from torch.nn.utils.rnn import pad_sequence
|
||||
|
||||
# ============================================================================
|
||||
# Constants (from modules/speaker/fbank.py)
|
||||
# ============================================================================
|
||||
|
||||
_SPEAKER_FBANK_SAMPLE_RATE = 16000
|
||||
_SPEAKER_FBANK_N_MELS = 80
|
||||
_SPEAKER_FBANK_MEAN_NORM = True
|
||||
_SPEAKER_FBANK_DITHER = 0.0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Audio helpers (folded from utils/audio.py — used by extract_speaker_fbank)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def high_quality_resample(x, orig_sr, target_sr):
|
||||
return AF.resample(
|
||||
x,
|
||||
orig_freq=orig_sr,
|
||||
new_freq=target_sr,
|
||||
lowpass_filter_width=64,
|
||||
rolloff=0.95,
|
||||
resampling_method="sinc_interp_kaiser",
|
||||
)
|
||||
|
||||
|
||||
def extract_fbank(
|
||||
waveform: torch.Tensor,
|
||||
*,
|
||||
sample_rate: int,
|
||||
n_mels: int,
|
||||
dither: float = 0.0,
|
||||
mean_norm: bool = False,
|
||||
) -> torch.Tensor:
|
||||
if waveform.ndim == 1:
|
||||
feature_input = waveform.unsqueeze(0)
|
||||
elif waveform.ndim == 2:
|
||||
feature_input = waveform if waveform.size(0) == 1 else waveform[0:1, :]
|
||||
else:
|
||||
raise ValueError(f"FBank expects a 1D or 2D waveform, got shape {tuple(waveform.shape)}.")
|
||||
features = Kaldi.fbank(
|
||||
feature_input,
|
||||
num_mel_bins=n_mels,
|
||||
sample_frequency=sample_rate,
|
||||
dither=dither,
|
||||
)
|
||||
if mean_norm:
|
||||
features = features - features.mean(dim=0, keepdim=True)
|
||||
return features
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Speaker fbank entry (from modules/speaker/fbank.py)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def extract_speaker_fbank(
|
||||
waveform: torch.Tensor,
|
||||
*,
|
||||
sample_rate: int,
|
||||
) -> torch.Tensor:
|
||||
feature_input = waveform
|
||||
if sample_rate != _SPEAKER_FBANK_SAMPLE_RATE:
|
||||
feature_input = high_quality_resample(
|
||||
waveform,
|
||||
orig_sr=sample_rate,
|
||||
target_sr=_SPEAKER_FBANK_SAMPLE_RATE,
|
||||
)
|
||||
return extract_fbank(
|
||||
feature_input,
|
||||
sample_rate=_SPEAKER_FBANK_SAMPLE_RATE,
|
||||
n_mels=_SPEAKER_FBANK_N_MELS,
|
||||
dither=_SPEAKER_FBANK_DITHER,
|
||||
mean_norm=_SPEAKER_FBANK_MEAN_NORM,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# CAM++ building blocks (from modules/speaker/campplus_layers.py)
|
||||
# ----------------------------------------------------------------------------
|
||||
# 3D-Speaker, Apache-2.0.
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def get_nonlinear(config_str, channels):
|
||||
nonlinear = nn.Sequential()
|
||||
for name in config_str.split("-"):
|
||||
if name == "relu":
|
||||
nonlinear.add_module("relu", nn.ReLU(inplace=True))
|
||||
elif name == "prelu":
|
||||
nonlinear.add_module("prelu", nn.PReLU(channels))
|
||||
elif name == "batchnorm":
|
||||
nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels))
|
||||
elif name == "batchnorm_":
|
||||
nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels, affine=False))
|
||||
else:
|
||||
raise ValueError(f"Unexpected module ({name}).")
|
||||
return nonlinear
|
||||
|
||||
|
||||
def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, _eps=1e-2):
|
||||
mean = x.mean(dim=dim)
|
||||
std = x.std(dim=dim, unbiased=unbiased)
|
||||
stats = torch.cat([mean, std], dim=-1)
|
||||
if keepdim:
|
||||
stats = stats.unsqueeze(dim=dim)
|
||||
return stats
|
||||
|
||||
|
||||
class StatsPool(nn.Module):
|
||||
def forward(self, x):
|
||||
return statistics_pooling(x)
|
||||
|
||||
|
||||
class TDNNLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
bias=False,
|
||||
config_str="batchnorm-relu",
|
||||
):
|
||||
super().__init__()
|
||||
if padding < 0:
|
||||
assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})"
|
||||
padding = (kernel_size - 1) // 2 * dilation
|
||||
self.linear = nn.Conv1d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
bias=bias,
|
||||
)
|
||||
self.nonlinear = get_nonlinear(config_str, out_channels)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.linear(x)
|
||||
return self.nonlinear(x)
|
||||
|
||||
|
||||
class CAMLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
bn_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
bias,
|
||||
reduction=2,
|
||||
):
|
||||
super().__init__()
|
||||
self.linear_local = nn.Conv1d(
|
||||
bn_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
bias=bias,
|
||||
)
|
||||
self.linear1 = nn.Conv1d(bn_channels, bn_channels // reduction, 1)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.linear2 = nn.Conv1d(bn_channels // reduction, out_channels, 1)
|
||||
self.sigmoid = nn.Sigmoid()
|
||||
|
||||
def forward(self, x):
|
||||
y = self.linear_local(x)
|
||||
context = x.mean(-1, keepdim=True) + self.seg_pooling(x)
|
||||
context = self.relu(self.linear1(context))
|
||||
m = self.sigmoid(self.linear2(context))
|
||||
return y * m
|
||||
|
||||
def seg_pooling(self, x, seg_len=100, stype="avg"):
|
||||
if stype == "avg":
|
||||
seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
|
||||
elif stype == "max":
|
||||
seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
|
||||
else:
|
||||
raise ValueError("Wrong segment pooling type.")
|
||||
shape = seg.shape
|
||||
seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1)
|
||||
return seg[..., : x.shape[-1]]
|
||||
|
||||
|
||||
class CAMDenseTDNNLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
bn_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
dilation=1,
|
||||
bias=False,
|
||||
config_str="batchnorm-relu",
|
||||
memory_efficient=False,
|
||||
):
|
||||
super().__init__()
|
||||
assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})"
|
||||
padding = (kernel_size - 1) // 2 * dilation
|
||||
self.memory_efficient = memory_efficient
|
||||
self.nonlinear1 = get_nonlinear(config_str, in_channels)
|
||||
self.linear1 = nn.Conv1d(in_channels, bn_channels, 1, bias=False)
|
||||
self.nonlinear2 = get_nonlinear(config_str, bn_channels)
|
||||
self.cam_layer = CAMLayer(
|
||||
bn_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
def bn_function(self, x):
|
||||
return self.linear1(self.nonlinear1(x))
|
||||
|
||||
def forward(self, x):
|
||||
if self.training and self.memory_efficient:
|
||||
x = cp.checkpoint(self.bn_function, x)
|
||||
else:
|
||||
x = self.bn_function(x)
|
||||
return self.cam_layer(self.nonlinear2(x))
|
||||
|
||||
|
||||
class CAMDenseTDNNBlock(nn.ModuleList):
|
||||
def __init__(
|
||||
self,
|
||||
num_layers,
|
||||
in_channels,
|
||||
out_channels,
|
||||
bn_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
dilation=1,
|
||||
bias=False,
|
||||
config_str="batchnorm-relu",
|
||||
memory_efficient=False,
|
||||
):
|
||||
super().__init__()
|
||||
for i in range(num_layers):
|
||||
layer = CAMDenseTDNNLayer(
|
||||
in_channels=in_channels + i * out_channels,
|
||||
out_channels=out_channels,
|
||||
bn_channels=bn_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
dilation=dilation,
|
||||
bias=bias,
|
||||
config_str=config_str,
|
||||
memory_efficient=memory_efficient,
|
||||
)
|
||||
self.add_module(f"tdnnd{i + 1}", layer)
|
||||
|
||||
def forward(self, x):
|
||||
for layer in self:
|
||||
x = torch.cat([x, layer(x)], dim=1)
|
||||
return x
|
||||
|
||||
|
||||
class TransitLayer(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"):
|
||||
super().__init__()
|
||||
self.nonlinear = get_nonlinear(config_str, in_channels)
|
||||
self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.nonlinear(x)
|
||||
return self.linear(x)
|
||||
|
||||
|
||||
class DenseLayer(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"):
|
||||
super().__init__()
|
||||
self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
|
||||
self.nonlinear = get_nonlinear(config_str, out_channels)
|
||||
|
||||
def forward(self, x):
|
||||
if len(x.shape) == 2:
|
||||
x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1)
|
||||
else:
|
||||
x = self.linear(x)
|
||||
return self.nonlinear(x)
|
||||
|
||||
|
||||
class BasicResBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1):
|
||||
super().__init__()
|
||||
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
|
||||
self.shortcut = nn.Sequential()
|
||||
if stride != 1 or in_planes != self.expansion * planes:
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
in_planes,
|
||||
self.expansion * planes,
|
||||
kernel_size=1,
|
||||
stride=(stride, 1),
|
||||
bias=False,
|
||||
),
|
||||
nn.BatchNorm2d(self.expansion * planes),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out = self.bn2(self.conv2(out))
|
||||
out += self.shortcut(x)
|
||||
return F.relu(out)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# CAM++ model (from modules/speaker/campplus.py)
|
||||
# ----------------------------------------------------------------------------
|
||||
# 3D-Speaker, Apache-2.0.
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class FCM(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
block=BasicResBlock,
|
||||
num_blocks=(2, 2),
|
||||
m_channels=32,
|
||||
feat_dim=_SPEAKER_FBANK_N_MELS,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_planes = m_channels
|
||||
self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(m_channels)
|
||||
|
||||
self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
|
||||
self.layer2 = self._make_layer(block, m_channels, num_blocks[1], stride=2)
|
||||
|
||||
self.conv2 = nn.Conv2d(m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(m_channels)
|
||||
self.out_channels = m_channels * (feat_dim // 8)
|
||||
|
||||
def _make_layer(self, block, planes, num_blocks, stride):
|
||||
strides = [stride] + [1] * (num_blocks - 1)
|
||||
layers = []
|
||||
for stride in strides:
|
||||
layers.append(block(self.in_planes, planes, stride))
|
||||
self.in_planes = planes * block.expansion
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = x.unsqueeze(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out = self.layer1(out)
|
||||
out = self.layer2(out)
|
||||
out = F.relu(self.bn2(self.conv2(out)))
|
||||
|
||||
shape = out.shape
|
||||
return out.reshape(shape[0], shape[1] * shape[2], shape[3])
|
||||
|
||||
|
||||
class CAMPPlus(nn.Module):
|
||||
_TDNN_KERNEL_SIZE = 5
|
||||
_TDNN_STRIDE = 2
|
||||
_TDNN_PADDING = 2
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
feat_dim=_SPEAKER_FBANK_N_MELS,
|
||||
embedding_size=512,
|
||||
growth_rate=32,
|
||||
bn_size=4,
|
||||
init_channels=128,
|
||||
config_str="batchnorm-relu",
|
||||
memory_efficient=True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.head = FCM(feat_dim=feat_dim)
|
||||
channels = self.head.out_channels
|
||||
|
||||
self.xvector = nn.Sequential(
|
||||
OrderedDict(
|
||||
[
|
||||
(
|
||||
"tdnn",
|
||||
TDNNLayer(
|
||||
channels,
|
||||
init_channels,
|
||||
self._TDNN_KERNEL_SIZE,
|
||||
stride=self._TDNN_STRIDE,
|
||||
dilation=1,
|
||||
padding=-1,
|
||||
config_str=config_str,
|
||||
),
|
||||
),
|
||||
]
|
||||
)
|
||||
)
|
||||
channels = init_channels
|
||||
for i, (num_layers, kernel_size, dilation) in enumerate(zip((12, 24, 16), (3, 3, 3), (1, 2, 2), strict=True)):
|
||||
block = CAMDenseTDNNBlock(
|
||||
num_layers=num_layers,
|
||||
in_channels=channels,
|
||||
out_channels=growth_rate,
|
||||
bn_channels=bn_size * growth_rate,
|
||||
kernel_size=kernel_size,
|
||||
dilation=dilation,
|
||||
config_str=config_str,
|
||||
memory_efficient=memory_efficient,
|
||||
)
|
||||
self.xvector.add_module(f"block{i + 1}", block)
|
||||
channels = channels + num_layers * growth_rate
|
||||
self.xvector.add_module(
|
||||
f"transit{i + 1}",
|
||||
TransitLayer(channels, channels // 2, bias=False, config_str=config_str),
|
||||
)
|
||||
channels //= 2
|
||||
|
||||
self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels))
|
||||
|
||||
self.xvector.add_module("stats", StatsPool())
|
||||
self.xvector.add_module("dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_"))
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, (nn.Conv1d, nn.Linear)):
|
||||
nn.init.kaiming_normal_(m.weight.data)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
@staticmethod
|
||||
def _conv_output_lengths(lengths, kernel_size, stride=1, padding=0, dilation=1):
|
||||
return (
|
||||
torch.div(
|
||||
lengths + 2 * padding - dilation * (kernel_size - 1) - 1,
|
||||
stride,
|
||||
rounding_mode="floor",
|
||||
)
|
||||
+ 1
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _make_length_mask(lengths, max_len, device):
|
||||
lengths = lengths.to(device=device, dtype=torch.long).clamp(min=0, max=max_len)
|
||||
return torch.arange(max_len, device=device).unsqueeze(0) < lengths.unsqueeze(1)
|
||||
|
||||
def _masked_stats_pooling(self, x, lengths, unbiased=True, eps=1e-2):
|
||||
lengths = lengths.to(device=x.device, dtype=torch.long).clamp(min=1, max=x.size(-1))
|
||||
mask = self._make_length_mask(lengths, x.size(-1), x.device).unsqueeze(1)
|
||||
mask = mask.to(dtype=x.dtype)
|
||||
|
||||
denom = lengths.to(dtype=x.dtype).view(-1, 1).clamp_min(1.0)
|
||||
mean = (x * mask).sum(dim=-1) / denom
|
||||
|
||||
centered = (x - mean.unsqueeze(-1)) * mask
|
||||
var_denom = (lengths - 1).clamp_min(1).to(dtype=x.dtype).view(-1, 1) if unbiased else denom
|
||||
var = centered.pow(2).sum(dim=-1) / var_denom
|
||||
std = torch.sqrt(var.clamp_min(eps))
|
||||
return torch.cat([mean, std], dim=1)
|
||||
|
||||
def forward(self, x, lengths=None):
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
x = self.head(x)
|
||||
if lengths is not None:
|
||||
lengths = lengths.to(device=x.device, dtype=torch.long).clamp(min=1)
|
||||
|
||||
for name, module in self.xvector.named_children():
|
||||
if name == "stats":
|
||||
x = self._masked_stats_pooling(x, lengths) if lengths is not None else module(x)
|
||||
continue
|
||||
|
||||
x = module(x)
|
||||
if name == "tdnn" and lengths is not None:
|
||||
lengths = self._conv_output_lengths(
|
||||
lengths,
|
||||
kernel_size=self._TDNN_KERNEL_SIZE,
|
||||
stride=self._TDNN_STRIDE,
|
||||
padding=self._TDNN_PADDING,
|
||||
)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# SpeakerXVectorFeatures wrapper (from modules/speaker/encoder.py)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class SpeakerXVectorFeatures(nn.Module):
|
||||
"""Speaker embedding extractor based on 3D-Speaker CAM++."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sample_rate=_SPEAKER_FBANK_SAMPLE_RATE,
|
||||
campplus_embedding_size=512,
|
||||
max_audio_seconds=10.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.sample_rate = sample_rate
|
||||
self.max_audio_seconds = float(max_audio_seconds)
|
||||
self.model = CAMPPlus(
|
||||
feat_dim=_SPEAKER_FBANK_N_MELS,
|
||||
embedding_size=campplus_embedding_size,
|
||||
)
|
||||
self.resample = None
|
||||
if self.sample_rate != _SPEAKER_FBANK_SAMPLE_RATE:
|
||||
self.resample = torchaudio.transforms.Resample(
|
||||
orig_freq=sample_rate,
|
||||
new_freq=_SPEAKER_FBANK_SAMPLE_RATE,
|
||||
)
|
||||
|
||||
for param in self.model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
@staticmethod
|
||||
def _normalize_lengths(lengths, batch_size, max_length, device, *, min_length):
|
||||
if lengths is None:
|
||||
return torch.full(
|
||||
(batch_size,),
|
||||
max_length,
|
||||
device=device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
return lengths.to(device=device, dtype=torch.long).clamp(
|
||||
min=min_length,
|
||||
max=max_length,
|
||||
)
|
||||
|
||||
def _crop_audio(self, audio, audio_lengths=None):
|
||||
original_lengths = self._normalize_lengths(
|
||||
audio_lengths,
|
||||
audio.size(0),
|
||||
audio.size(-1),
|
||||
audio.device,
|
||||
min_length=0,
|
||||
)
|
||||
if self.max_audio_seconds <= 0:
|
||||
return audio, original_lengths, original_lengths, torch.zeros_like(original_lengths)
|
||||
|
||||
max_input_length = round(self.sample_rate * self.max_audio_seconds)
|
||||
cropped_audio = []
|
||||
cropped_lengths = []
|
||||
starts = []
|
||||
|
||||
for index, total_length_tensor in enumerate(original_lengths):
|
||||
total_length = int(total_length_tensor.item())
|
||||
cropped_length = min(total_length, max_input_length)
|
||||
start = random.randint(0, total_length - cropped_length) if total_length > cropped_length else 0
|
||||
cropped_audio.append(audio[index, start : start + cropped_length])
|
||||
cropped_lengths.append(cropped_length)
|
||||
starts.append(start)
|
||||
|
||||
return (
|
||||
pad_sequence(
|
||||
cropped_audio,
|
||||
batch_first=True,
|
||||
padding_value=0.0,
|
||||
),
|
||||
original_lengths,
|
||||
torch.tensor(
|
||||
cropped_lengths,
|
||||
device=audio.device,
|
||||
dtype=torch.long,
|
||||
),
|
||||
torch.tensor(starts, device=audio.device, dtype=torch.long),
|
||||
)
|
||||
|
||||
def _crop_fbank(
|
||||
self,
|
||||
fbank,
|
||||
fbank_lengths,
|
||||
original_audio_lengths,
|
||||
cropped_audio_lengths,
|
||||
starts,
|
||||
):
|
||||
original_fbank_lengths = self._normalize_lengths(
|
||||
fbank_lengths,
|
||||
fbank.size(0),
|
||||
fbank.size(1),
|
||||
fbank.device,
|
||||
min_length=1,
|
||||
)
|
||||
cropped_fbank = []
|
||||
cropped_fbank_lengths = []
|
||||
|
||||
for index, total_feat_length_tensor in enumerate(original_fbank_lengths):
|
||||
total_audio_length = int(original_audio_lengths[index].item())
|
||||
total_feat_length = int(total_feat_length_tensor.item())
|
||||
start_audio = int(starts[index].item())
|
||||
end_audio = start_audio + int(cropped_audio_lengths[index].item())
|
||||
|
||||
if total_audio_length > 0:
|
||||
start_feat = math.floor(start_audio * total_feat_length / total_audio_length)
|
||||
end_feat = math.ceil(end_audio * total_feat_length / total_audio_length)
|
||||
else:
|
||||
start_feat = 0
|
||||
end_feat = 1
|
||||
|
||||
start_feat = min(start_feat, total_feat_length - 1)
|
||||
end_feat = min(max(end_feat, start_feat + 1), total_feat_length)
|
||||
cropped_fbank.append(fbank[index, start_feat:end_feat])
|
||||
cropped_fbank_lengths.append(end_feat - start_feat)
|
||||
|
||||
return pad_sequence(
|
||||
cropped_fbank,
|
||||
batch_first=True,
|
||||
padding_value=0.0,
|
||||
), torch.tensor(
|
||||
cropped_fbank_lengths,
|
||||
device=fbank.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
|
||||
def _extract_fbank_batch(self, audio, audio_lengths):
|
||||
if self.resample is not None:
|
||||
audio = self.resample(audio)
|
||||
audio_lengths = torch.ceil(audio_lengths.float() * (_SPEAKER_FBANK_SAMPLE_RATE / self.sample_rate)).long()
|
||||
|
||||
audio_cpu = audio.detach().cpu()
|
||||
features = []
|
||||
|
||||
for index, valid_length_tensor in enumerate(audio_lengths):
|
||||
valid_length = int(valid_length_tensor.item())
|
||||
waveform = audio_cpu[index, :valid_length]
|
||||
if waveform.numel() == 0:
|
||||
waveform = audio_cpu.new_zeros(1)
|
||||
features.append(
|
||||
extract_speaker_fbank(
|
||||
waveform,
|
||||
sample_rate=_SPEAKER_FBANK_SAMPLE_RATE,
|
||||
)
|
||||
)
|
||||
|
||||
fbank_lengths = torch.tensor(
|
||||
[feature.size(0) for feature in features],
|
||||
device=audio.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
fbank = pad_sequence(
|
||||
features,
|
||||
batch_first=True,
|
||||
padding_value=0.0,
|
||||
).to(device=audio.device, dtype=audio.dtype)
|
||||
return fbank, fbank_lengths
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.autocast(enabled=False, device_type="cuda")
|
||||
def forward(self, audio, audio_lengths=None, fbank=None, fbank_lengths=None, **_kwargs):
|
||||
self.model.eval()
|
||||
audio = audio.float()
|
||||
if audio.dim() == 3:
|
||||
if audio.size(1) != 1:
|
||||
raise ValueError(f"Speaker encoder expects mono audio, got shape {tuple(audio.shape)}.")
|
||||
audio = audio[:, 0]
|
||||
elif audio.dim() != 2:
|
||||
raise ValueError(f"Speaker encoder expects a 2D or 3D audio tensor, got shape {tuple(audio.shape)}.")
|
||||
|
||||
audio, original_audio_lengths, cropped_audio_lengths, starts = self._crop_audio(
|
||||
audio,
|
||||
audio_lengths=audio_lengths,
|
||||
)
|
||||
|
||||
if fbank is None:
|
||||
fbank, fbank_lengths = self._extract_fbank_batch(
|
||||
audio,
|
||||
cropped_audio_lengths,
|
||||
)
|
||||
else:
|
||||
if not isinstance(fbank, torch.Tensor):
|
||||
raise TypeError("Speaker encoder expects `fbank` to be a torch.Tensor.")
|
||||
if fbank.dim() != 3 or fbank.size(0) != audio.size(0):
|
||||
raise ValueError(
|
||||
f"Speaker encoder expects `fbank` with shape (B, T, F) and matching batch size, got {tuple(fbank.shape)}." # noqa: E501
|
||||
)
|
||||
fbank, fbank_lengths = self._crop_fbank(
|
||||
fbank.to(device=audio.device, dtype=torch.float32),
|
||||
fbank_lengths,
|
||||
original_audio_lengths,
|
||||
cropped_audio_lengths,
|
||||
starts,
|
||||
)
|
||||
|
||||
return self.model(fbank, lengths=fbank_lengths)
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -1,46 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""dots.tts pipeline topology (frozen).
|
||||
|
||||
Single-stage AR TTS: text → 48 kHz waveform in one pass.
|
||||
Architecture: Qwen2.5-1.5B base LM (vLLM-native) → DiTAR / MeanFlow head
|
||||
→ AudioVAE.decode → BigVGAN → 48 kHz mono.
|
||||
|
||||
Speaker conditioning via CAM++ x-vector injected in preprocess().
|
||||
|
||||
Modeled after voxcpm2's single-stage pipeline (same "vLLM-native base LM
|
||||
+ side-path computation" pattern), but simplified: no residual LM, no FSQ.
|
||||
"""
|
||||
|
||||
from vllm_omni.config.stage_config import (
|
||||
PipelineConfig,
|
||||
StageExecutionType,
|
||||
StagePipelineConfig,
|
||||
)
|
||||
|
||||
DOTS_TTS_PIPELINE = PipelineConfig(
|
||||
model_type="dots_tts",
|
||||
model_arch="DotsTTSForConditionalGeneration",
|
||||
default_deploy_config_name="dots_tts.yaml",
|
||||
stages=(
|
||||
StagePipelineConfig(
|
||||
stage_id=0,
|
||||
model_stage="latent_generator",
|
||||
execution_type=StageExecutionType.LLM_AR,
|
||||
input_sources=(),
|
||||
final_output=True,
|
||||
final_output_type="audio",
|
||||
owns_tokenizer=True,
|
||||
engine_output_type="audio",
|
||||
sampling_constraints={
|
||||
"detokenize": False,
|
||||
# Stop signal slot in compute_logits (voxcpm2 convention —
|
||||
# talker emits logits[i, 1] = 1.0 when prob_stop > 0.8, vLLM
|
||||
# matches against this list to terminate the request). Not
|
||||
# Qwen2.5 EOS (151643) — that ID would never come out of the
|
||||
# 2-slot continue/stop softmax our compute_logits builds.
|
||||
"stop_token_ids": [1],
|
||||
},
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -1,8 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# NOTE: Do not import model classes in this file. Importing any
|
||||
# submodule in this package triggers __init__.py execution, and
|
||||
# both the model registry and pipeline registry import submodules
|
||||
# directly — heavy imports here would be loaded as a side effect
|
||||
# even though nothing depends on these re-exports.
|
||||
@@ -1,747 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from functools import cached_property
|
||||
from importlib import import_module
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.config.multimodal import BaseDummyOptions
|
||||
from vllm.inputs import MultiModalDataDict
|
||||
from vllm.inputs import MultiModalInput as MultiModalInputs
|
||||
from vllm.model_executor.models.interfaces import SupportsMultiModal
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalFieldConfig,
|
||||
MultiModalKwargsItems,
|
||||
PlaceholderRange,
|
||||
)
|
||||
from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
|
||||
from vllm.multimodal.processing import (
|
||||
BaseDummyInputsBuilder,
|
||||
BaseMultiModalProcessor,
|
||||
BaseProcessingInfo,
|
||||
ProcessorInputs,
|
||||
PromptUpdate,
|
||||
TimingContext,
|
||||
)
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.v1.outputs import SamplerOutput
|
||||
from vllm.v1.sample.metadata import SamplingMetadata
|
||||
from vllm.v1.sample.sampler import Sampler
|
||||
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
from .dynin_omni_common import build_zero_input_embeddings
|
||||
|
||||
try:
|
||||
from PIL import Image as PILImage
|
||||
except Exception: # pragma: no cover
|
||||
PILImage = None
|
||||
|
||||
|
||||
_MODALITY_ORDER = ("image", "video", "audio")
|
||||
|
||||
_MODALITY_ALIASES = {
|
||||
"img2img": "image",
|
||||
}
|
||||
|
||||
_MODALITY_INPUT_KEY_BY_NAME = {
|
||||
"image": "pixel_values",
|
||||
"video": "pixel_values_videos",
|
||||
"audio": "input_audio_features",
|
||||
}
|
||||
|
||||
_MODALITY_PLACEHOLDER_BY_NAME = {
|
||||
"image": "<|soi|><|image|><|eoi|>",
|
||||
"video": "<|sov|><|video|><|eov|>",
|
||||
"audio": "<|soa|><|audio|><|eoa|>",
|
||||
}
|
||||
|
||||
_MODALITY_INPUT_ALIASES = {
|
||||
"image": ("pixel_values", "image_embeds", "img2img"),
|
||||
"video": ("pixel_values_videos", "video_embeds"),
|
||||
"audio": ("input_audio_features", "audio_embeds"),
|
||||
}
|
||||
|
||||
|
||||
def _normalize_modality_name(modality: str) -> str:
|
||||
return _MODALITY_ALIASES.get(modality, modality)
|
||||
|
||||
|
||||
def _get_modality_count(mm_counts: Mapping[str, int], modality: str) -> int:
|
||||
canonical = _normalize_modality_name(modality)
|
||||
count = mm_counts.get(canonical, 0)
|
||||
for alias, target in _MODALITY_ALIASES.items():
|
||||
if target == canonical:
|
||||
count += mm_counts.get(alias, 0)
|
||||
return count
|
||||
|
||||
|
||||
def _normalize_mm_data_aliases(mm_data: MultiModalDataDict) -> MultiModalDataDict:
|
||||
normalized: dict[str, Any] = {}
|
||||
for modality, value in mm_data.items():
|
||||
canonical = _normalize_modality_name(modality)
|
||||
if canonical in normalized and normalized[canonical] is not None and value is not None:
|
||||
raise ValueError(
|
||||
"Dynin received duplicate multimodal inputs for "
|
||||
f"{canonical!r} via {modality!r}. "
|
||||
"Provide either the canonical modality or its alias, not both."
|
||||
)
|
||||
if canonical not in normalized or normalized[canonical] is None:
|
||||
normalized[canonical] = value
|
||||
return normalized
|
||||
|
||||
|
||||
def _get_placeholder_text(modality: str) -> str | None:
|
||||
modality = _normalize_modality_name(modality)
|
||||
for base_modality, placeholder in _MODALITY_PLACEHOLDER_BY_NAME.items():
|
||||
if modality.startswith(base_modality):
|
||||
return placeholder
|
||||
return None
|
||||
|
||||
|
||||
class DyninOmniProcessingInfo(BaseProcessingInfo):
|
||||
def get_data_parser(self) -> MultiModalDataParser:
|
||||
return DyninOmniMultiModalDataParser(
|
||||
expected_hidden_size=self._get_expected_hidden_size(),
|
||||
)
|
||||
|
||||
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
|
||||
limits = {modality: 1 for modality in _MODALITY_ORDER}
|
||||
for alias, target in _MODALITY_ALIASES.items():
|
||||
if target in limits:
|
||||
limits[alias] = limits[target]
|
||||
return limits
|
||||
|
||||
def get_mm_max_tokens_per_item(
|
||||
self,
|
||||
seq_len: int,
|
||||
mm_counts: Mapping[str, int],
|
||||
) -> Mapping[str, int] | None:
|
||||
del seq_len, mm_counts
|
||||
limits = {modality: 1 for modality in _MODALITY_ORDER}
|
||||
for alias, target in _MODALITY_ALIASES.items():
|
||||
if target in limits:
|
||||
limits[alias] = limits[target]
|
||||
return limits
|
||||
|
||||
|
||||
class DyninOmniDummyInputsBuilder(BaseDummyInputsBuilder[DyninOmniProcessingInfo]):
|
||||
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
|
||||
chunks: list[str] = []
|
||||
for modality in _MODALITY_ORDER:
|
||||
placeholder = _get_placeholder_text(modality)
|
||||
if placeholder is None:
|
||||
continue
|
||||
chunks.extend([placeholder] * _get_modality_count(mm_counts, modality))
|
||||
return " ".join(chunks)
|
||||
|
||||
def get_dummy_mm_data(
|
||||
self,
|
||||
seq_len: int,
|
||||
mm_counts: Mapping[str, int],
|
||||
mm_options: Mapping[str, BaseDummyOptions] | None = None,
|
||||
) -> MultiModalDataDict:
|
||||
del seq_len
|
||||
|
||||
mm_data: dict[str, Any] = {}
|
||||
|
||||
num_images = _get_modality_count(mm_counts, "image")
|
||||
if num_images > 0:
|
||||
mm_data["image"] = self._get_dummy_images(
|
||||
width=224,
|
||||
height=224,
|
||||
num_images=num_images,
|
||||
overrides=mm_options.get("image") if mm_options else None,
|
||||
)
|
||||
|
||||
num_videos = _get_modality_count(mm_counts, "video")
|
||||
if num_videos > 0:
|
||||
mm_data["video"] = self._get_dummy_videos(
|
||||
width=224,
|
||||
height=224,
|
||||
num_frames=8,
|
||||
num_videos=num_videos,
|
||||
overrides=mm_options.get("video") if mm_options else None,
|
||||
)
|
||||
|
||||
num_audios = _get_modality_count(mm_counts, "audio")
|
||||
if num_audios > 0:
|
||||
mm_data["audio"] = self._get_dummy_audios(
|
||||
length=16000,
|
||||
num_audios=num_audios,
|
||||
overrides=mm_options.get("audio") if mm_options else None,
|
||||
)
|
||||
|
||||
return mm_data
|
||||
|
||||
|
||||
class DyninOmniMultiModalDataParser(MultiModalDataParser):
|
||||
def parse_mm_data(self, mm_data: MultiModalDataDict) -> MultiModalDataItems:
|
||||
normalized = _normalize_mm_data_aliases(mm_data)
|
||||
mm_items = super().parse_mm_data(normalized)
|
||||
|
||||
for alias, canonical in _MODALITY_ALIASES.items():
|
||||
if alias in mm_data and canonical in mm_items and alias not in mm_items:
|
||||
mm_items[alias] = mm_items[canonical]
|
||||
|
||||
return mm_items
|
||||
|
||||
def _get_audio_with_sr(self, audio: Any) -> tuple[np.ndarray, float | None]:
|
||||
audio_array, orig_sr = super()._get_audio_with_sr(audio)
|
||||
if self.audio_resampler.target_sr is None:
|
||||
return audio_array, None
|
||||
return audio_array, orig_sr
|
||||
|
||||
|
||||
class DyninOmniMultiModalProcessor(BaseMultiModalProcessor[DyninOmniProcessingInfo]):
|
||||
@staticmethod
|
||||
def _find_subsequence(
|
||||
haystack: list[int],
|
||||
needle: list[int],
|
||||
start: int,
|
||||
) -> int | None:
|
||||
if not needle:
|
||||
return None
|
||||
|
||||
max_start = len(haystack) - len(needle)
|
||||
if max_start < start:
|
||||
return None
|
||||
|
||||
for idx in range(start, max_start + 1):
|
||||
if haystack[idx : idx + len(needle)] == needle:
|
||||
return idx
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _make_disabled_embed_mask(length: int) -> torch.Tensor:
|
||||
return torch.zeros(length, dtype=torch.bool)
|
||||
|
||||
@staticmethod
|
||||
def _encode_prompt_to_token_ids(
|
||||
prompt: str | list[int],
|
||||
tokenizer: Any | None,
|
||||
) -> list[int]:
|
||||
if isinstance(prompt, str):
|
||||
if tokenizer is None:
|
||||
raise ValueError("Tokenizer is required to process string prompts for Dynin multimodal inputs.")
|
||||
return tokenizer.encode(prompt, add_special_tokens=False)
|
||||
return list(prompt)
|
||||
|
||||
@staticmethod
|
||||
def _ensure_non_empty_prompt_ids(
|
||||
prompt_token_ids: list[int],
|
||||
tokenizer: Any | None,
|
||||
) -> list[int]:
|
||||
if prompt_token_ids:
|
||||
return prompt_token_ids
|
||||
|
||||
fallback_id = None
|
||||
if tokenizer is not None:
|
||||
fallback_id = getattr(tokenizer, "bos_token_id", None)
|
||||
if fallback_id is None:
|
||||
fallback_id = getattr(tokenizer, "eos_token_id", None)
|
||||
if fallback_id is None:
|
||||
fallback_id = getattr(tokenizer, "pad_token_id", None)
|
||||
|
||||
return [0 if fallback_id is None else int(fallback_id)]
|
||||
|
||||
@classmethod
|
||||
def _image_to_chw_float_tensor(cls, image: Any) -> torch.Tensor:
|
||||
if isinstance(image, torch.Tensor):
|
||||
tensor = image.detach()
|
||||
elif isinstance(image, np.ndarray):
|
||||
tensor = torch.from_numpy(image)
|
||||
elif PILImage is not None and isinstance(image, PILImage.Image):
|
||||
tensor = torch.from_numpy(np.asarray(image).copy())
|
||||
else:
|
||||
raise TypeError(f"Unsupported image item type: {type(image)!r}")
|
||||
|
||||
if tensor.ndim == 2:
|
||||
tensor = tensor.unsqueeze(-1)
|
||||
if tensor.ndim != 3:
|
||||
raise ValueError(f"Expected 3D image tensor, got shape={tuple(tensor.shape)}")
|
||||
|
||||
if tensor.shape[-1] in (1, 3, 4) and tensor.shape[0] not in (1, 3, 4):
|
||||
tensor = tensor.permute(2, 0, 1)
|
||||
|
||||
if tensor.shape[0] == 1:
|
||||
tensor = tensor.repeat(3, 1, 1)
|
||||
if tensor.shape[0] == 4:
|
||||
tensor = tensor[:3]
|
||||
|
||||
tensor = tensor.to(dtype=torch.float32)
|
||||
if tensor.numel() > 0 and torch.max(tensor) > 1.0:
|
||||
tensor = tensor / 255.0
|
||||
return tensor.contiguous()
|
||||
|
||||
@classmethod
|
||||
def _video_to_tchw_float_tensor(cls, video: Any) -> torch.Tensor:
|
||||
if isinstance(video, (list, tuple)) and not isinstance(video, torch.Tensor):
|
||||
frames = [cls._image_to_chw_float_tensor(frame) for frame in video]
|
||||
if not frames:
|
||||
return torch.zeros((1, 3, 1, 1), dtype=torch.float32)
|
||||
return torch.stack(frames, dim=0).contiguous()
|
||||
|
||||
if isinstance(video, torch.Tensor):
|
||||
tensor = video.detach()
|
||||
elif isinstance(video, np.ndarray):
|
||||
tensor = torch.from_numpy(video)
|
||||
else:
|
||||
raise TypeError(f"Unsupported video item type: {type(video)!r}")
|
||||
|
||||
if tensor.ndim == 3:
|
||||
return cls._image_to_chw_float_tensor(tensor).unsqueeze(0).contiguous()
|
||||
|
||||
if tensor.ndim != 4:
|
||||
raise ValueError(f"Expected 4D video tensor, got shape={tuple(tensor.shape)}")
|
||||
|
||||
if tensor.shape[-1] in (1, 3, 4) and tensor.shape[1] not in (1, 3, 4):
|
||||
tensor = tensor.permute(0, 3, 1, 2)
|
||||
|
||||
if tensor.shape[1] == 1:
|
||||
tensor = tensor.repeat(1, 3, 1, 1)
|
||||
if tensor.shape[1] == 4:
|
||||
tensor = tensor[:, :3]
|
||||
|
||||
tensor = tensor.to(dtype=torch.float32)
|
||||
if tensor.numel() > 0 and torch.max(tensor) > 1.0:
|
||||
tensor = tensor / 255.0
|
||||
return tensor.contiguous()
|
||||
|
||||
@staticmethod
|
||||
def _audio_to_float_tensor(audio: Any) -> torch.Tensor:
|
||||
if isinstance(audio, tuple) and len(audio) == 2:
|
||||
audio = audio[0]
|
||||
|
||||
if isinstance(audio, torch.Tensor):
|
||||
tensor = audio.detach()
|
||||
elif isinstance(audio, np.ndarray):
|
||||
tensor = torch.from_numpy(audio)
|
||||
else:
|
||||
tensor = torch.as_tensor(audio)
|
||||
|
||||
tensor = tensor.to(dtype=torch.float32).contiguous().view(-1)
|
||||
if tensor.numel() == 0:
|
||||
return torch.zeros((16000,), dtype=torch.float32)
|
||||
|
||||
max_abs = torch.max(torch.abs(tensor))
|
||||
if max_abs > 1.0:
|
||||
tensor = tensor / max_abs
|
||||
|
||||
return tensor.contiguous()
|
||||
|
||||
@classmethod
|
||||
def _convert_modality_item(cls, modality: str, item: Any) -> torch.Tensor:
|
||||
if modality == "image":
|
||||
return cls._image_to_chw_float_tensor(item)
|
||||
if modality == "video":
|
||||
return cls._video_to_tchw_float_tensor(item)
|
||||
if modality == "audio":
|
||||
return cls._audio_to_float_tensor(item)
|
||||
raise ValueError(f"Unsupported modality for Dynin processor: {modality}")
|
||||
|
||||
def _build_modality_kwargs(
|
||||
self,
|
||||
modality: str,
|
||||
modality_items: Sequence[Any],
|
||||
) -> Sequence[Any]:
|
||||
modality = _normalize_modality_name(modality)
|
||||
input_key = _MODALITY_INPUT_KEY_BY_NAME[modality]
|
||||
tensor_items = [self._convert_modality_item(modality, item) for item in modality_items]
|
||||
mm_kwargs = MultiModalKwargsItems.from_hf_inputs(
|
||||
{input_key: tensor_items},
|
||||
{input_key: MultiModalFieldConfig.batched(modality)},
|
||||
)
|
||||
return mm_kwargs[modality]
|
||||
|
||||
def _build_placeholder_ranges(
|
||||
self,
|
||||
*,
|
||||
modality: str,
|
||||
item_count: int,
|
||||
prompt_token_ids: list[int],
|
||||
tokenizer: Any | None,
|
||||
search_start: int,
|
||||
) -> tuple[list[PlaceholderRange], int]:
|
||||
ranges: list[PlaceholderRange] = []
|
||||
|
||||
for _ in range(item_count):
|
||||
placeholder_text = _get_placeholder_text(modality)
|
||||
placeholder_token_ids: list[int] = []
|
||||
|
||||
if placeholder_text and tokenizer is not None:
|
||||
placeholder_token_ids = tokenizer.encode(
|
||||
placeholder_text,
|
||||
add_special_tokens=False,
|
||||
)
|
||||
|
||||
found_offset = None
|
||||
if placeholder_token_ids:
|
||||
found_offset = self._find_subsequence(
|
||||
prompt_token_ids,
|
||||
placeholder_token_ids,
|
||||
search_start,
|
||||
)
|
||||
|
||||
if found_offset is None:
|
||||
found_offset = min(search_start, len(prompt_token_ids) - 1)
|
||||
placeholder_len = 1
|
||||
else:
|
||||
placeholder_len = len(placeholder_token_ids)
|
||||
|
||||
ranges.append(
|
||||
PlaceholderRange(
|
||||
offset=found_offset,
|
||||
length=placeholder_len,
|
||||
is_embed=self._make_disabled_embed_mask(placeholder_len),
|
||||
)
|
||||
)
|
||||
search_start = found_offset + placeholder_len
|
||||
|
||||
return ranges, search_start
|
||||
|
||||
def _get_mm_fields_config(
|
||||
self,
|
||||
hf_inputs: Any,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
) -> Mapping[str, MultiModalFieldConfig]:
|
||||
del hf_inputs, hf_processor_mm_kwargs
|
||||
return {}
|
||||
|
||||
def _get_prompt_updates(
|
||||
self,
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
out_mm_kwargs: MultiModalKwargsItems,
|
||||
) -> Sequence[PromptUpdate]:
|
||||
del mm_items, hf_processor_mm_kwargs, out_mm_kwargs
|
||||
return []
|
||||
|
||||
def apply(
|
||||
self,
|
||||
inputs: ProcessorInputs,
|
||||
timing_ctx: TimingContext,
|
||||
) -> MultiModalInputs:
|
||||
prompt = inputs.prompt
|
||||
mm_items = inputs.mm_data_items
|
||||
|
||||
with timing_ctx.record("get_mm_hashes"):
|
||||
mm_hashes = inputs.get_mm_hashes(
|
||||
self.info.model_id,
|
||||
self.info.ctx.get_mm_config().mm_hasher_algorithm,
|
||||
)
|
||||
|
||||
tokenizer = self.info.ctx.tokenizer
|
||||
prompt_token_ids = self._encode_prompt_to_token_ids(prompt, tokenizer)
|
||||
prompt_token_ids = self._ensure_non_empty_prompt_ids(prompt_token_ids, tokenizer)
|
||||
|
||||
mm_kwargs_by_modality: dict[str, Sequence[Any]] = {}
|
||||
mm_placeholders: dict[str, list[PlaceholderRange]] = {}
|
||||
search_start = 0
|
||||
mm_counts = mm_items.get_all_counts()
|
||||
|
||||
for modality in _MODALITY_ORDER:
|
||||
item_count = mm_counts.get(modality, 0)
|
||||
if item_count <= 0:
|
||||
continue
|
||||
|
||||
modality_items = mm_items[modality].get_all()
|
||||
if len(modality_items) != item_count:
|
||||
raise RuntimeError(
|
||||
f"Parsed {len(modality_items)} items but expected {item_count} for modality={modality!r}"
|
||||
)
|
||||
|
||||
mm_kwargs_by_modality[modality] = self._build_modality_kwargs(
|
||||
modality,
|
||||
modality_items,
|
||||
)
|
||||
|
||||
placeholder_ranges, search_start = self._build_placeholder_ranges(
|
||||
modality=modality,
|
||||
item_count=item_count,
|
||||
prompt_token_ids=prompt_token_ids,
|
||||
tokenizer=tokenizer,
|
||||
search_start=search_start,
|
||||
)
|
||||
mm_placeholders[modality] = placeholder_ranges
|
||||
|
||||
return MultiModalInputs(
|
||||
type="multimodal",
|
||||
prompt_token_ids=prompt_token_ids,
|
||||
mm_kwargs=MultiModalKwargsItems(mm_kwargs_by_modality),
|
||||
mm_hashes=mm_hashes,
|
||||
mm_placeholders=mm_placeholders,
|
||||
)
|
||||
|
||||
|
||||
class DyninOmniStageBase(nn.Module):
|
||||
stage_name = "Dynin stage"
|
||||
|
||||
def make_empty_intermediate_tensors(
|
||||
self,
|
||||
batch_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> IntermediateTensors:
|
||||
del batch_size, dtype, device
|
||||
return IntermediateTensors({})
|
||||
|
||||
def embed_input_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings: Any = None,
|
||||
is_multimodal: torch.Tensor | None = None,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor:
|
||||
del multimodal_embeddings, is_multimodal, kwargs
|
||||
return build_zero_input_embeddings(
|
||||
input_ids=input_ids,
|
||||
hidden_size=self.hidden_size,
|
||||
stage_name=self.stage_name,
|
||||
)
|
||||
|
||||
def load_weights(
|
||||
self,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> set[str]:
|
||||
return {name for name, _ in weights}
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor | OmniOutput,
|
||||
sampling_metadata: Any = None,
|
||||
) -> torch.Tensor | None:
|
||||
del hidden_states, sampling_metadata
|
||||
return None
|
||||
|
||||
|
||||
@MULTIMODAL_REGISTRY.register_processor(
|
||||
DyninOmniMultiModalProcessor,
|
||||
info=DyninOmniProcessingInfo,
|
||||
dummy_inputs=DyninOmniDummyInputsBuilder,
|
||||
)
|
||||
class DyninOmniForConditionalGeneration(nn.Module, SupportsMultiModal):
|
||||
supports_multimodal_raw_input_only = True
|
||||
STAGE_ALIAS = {
|
||||
"tokenizer": "token2text",
|
||||
"token2token": "token2text",
|
||||
"detok_text": "token2text",
|
||||
"token2img": "token2image",
|
||||
"token2wav": "token2audio",
|
||||
"token2speech": "token2audio",
|
||||
}
|
||||
|
||||
STAGE_IMPL = {
|
||||
"token2text": (".dynin_omni_token2text", "DyninOmniToken2Text"),
|
||||
"token2image": (".dynin_omni_token2image", "DyninOmniToken2Image"),
|
||||
"token2audio": (".dynin_omni_token2audio", "DyninOmniToken2Audio"),
|
||||
}
|
||||
|
||||
_STAGE_IMPL_CACHE: dict[str, type[nn.Module]] = {}
|
||||
|
||||
@classmethod
|
||||
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
|
||||
del i
|
||||
return _get_placeholder_text(modality)
|
||||
|
||||
@classmethod
|
||||
def _resolve_stage_impl_class(cls, model_stage: str) -> type[nn.Module]:
|
||||
impl = cls._STAGE_IMPL_CACHE.get(model_stage)
|
||||
if impl is not None:
|
||||
return impl
|
||||
|
||||
module_name, class_name = cls.STAGE_IMPL[model_stage]
|
||||
module = import_module(module_name, package=__package__)
|
||||
impl = getattr(module, class_name)
|
||||
cls._STAGE_IMPL_CACHE[model_stage] = impl
|
||||
return impl
|
||||
|
||||
@classmethod
|
||||
def _normalize_stage_name(cls, raw_stage: str) -> str:
|
||||
normalized = cls.STAGE_ALIAS.get(raw_stage, raw_stage)
|
||||
if normalized not in cls.STAGE_IMPL:
|
||||
raise ValueError(
|
||||
"Unsupported DYNIN omni model_stage: "
|
||||
f"{raw_stage} (normalized={normalized}). "
|
||||
f"Supported: {sorted(cls.STAGE_IMPL.keys())}"
|
||||
)
|
||||
return normalized
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
|
||||
raw_stage = str(getattr(vllm_config.model_config, "model_stage", "token2text")).lower()
|
||||
self.model_stage = self._normalize_stage_name(raw_stage)
|
||||
|
||||
impl_cls = self._resolve_stage_impl_class(self.model_stage)
|
||||
self.impl = impl_cls(vllm_config=vllm_config, prefix=prefix)
|
||||
self.model = self.impl
|
||||
|
||||
self.has_preprocess = False
|
||||
self.has_postprocess = False
|
||||
self.have_multimodal_outputs = getattr(self.impl, "have_multimodal_outputs", True)
|
||||
self.requires_raw_input_tokens = getattr(self.impl, "requires_raw_input_tokens", True)
|
||||
self.language_model = self._resolve_language_model()
|
||||
|
||||
def _resolve_language_model(self) -> Any | None:
|
||||
if hasattr(self.impl, "get_language_model"):
|
||||
language_model = self.impl.get_language_model()
|
||||
if language_model is not None:
|
||||
return language_model
|
||||
|
||||
if hasattr(self.impl, "language_model"):
|
||||
language_model = getattr(self.impl, "language_model")
|
||||
if language_model is not None:
|
||||
return language_model
|
||||
|
||||
if self.model_stage == "token2text":
|
||||
return getattr(self.impl, "model", None)
|
||||
|
||||
return None
|
||||
|
||||
def get_language_model(self) -> Any | None:
|
||||
return self.language_model
|
||||
|
||||
@cached_property
|
||||
def sampler(self):
|
||||
if hasattr(self.model, "sampler"):
|
||||
return self.model.sampler
|
||||
if self.language_model is not None and hasattr(self.language_model, "sampler"):
|
||||
return self.language_model.sampler
|
||||
return Sampler()
|
||||
|
||||
def init_multi_modal(self, thinker_config: Any = None) -> None:
|
||||
if hasattr(self.model, "init_multi_modal"):
|
||||
self.model.init_multi_modal(thinker_config)
|
||||
|
||||
def _collect_multimodal_inputs(self, **kwargs: Any) -> dict[str, Any]:
|
||||
mm_inputs: dict[str, Any] = {}
|
||||
for modality, aliases in _MODALITY_INPUT_ALIASES.items():
|
||||
for alias in aliases:
|
||||
if alias in kwargs and kwargs[alias] is not None:
|
||||
mm_inputs[modality] = kwargs[alias]
|
||||
break
|
||||
return mm_inputs
|
||||
|
||||
def _normalize_loaded_weight_names(
|
||||
self,
|
||||
loaded: set[str],
|
||||
expected_param_names: set[str],
|
||||
) -> set[str]:
|
||||
if self.model_stage != "token2text":
|
||||
return loaded
|
||||
|
||||
normalized_loaded: set[str] = set()
|
||||
prefixes = ("", "impl.", "impl.model.")
|
||||
|
||||
for name in loaded:
|
||||
for prefix in prefixes:
|
||||
candidate = f"{prefix}{name}" if prefix else name
|
||||
if candidate in expected_param_names:
|
||||
normalized_loaded.add(candidate)
|
||||
break
|
||||
|
||||
if len(normalized_loaded) < len(expected_param_names):
|
||||
normalized_loaded.update(expected_param_names)
|
||||
|
||||
return normalized_loaded
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
positions: torch.Tensor | None = None,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: Any,
|
||||
) -> OmniOutput:
|
||||
return self.model(
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def make_empty_intermediate_tensors(
|
||||
self,
|
||||
batch_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> IntermediateTensors:
|
||||
return self.model.make_empty_intermediate_tensors(batch_size, dtype, device)
|
||||
|
||||
def embed_input_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings: Any = None,
|
||||
is_multimodal: torch.Tensor | None = None,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor:
|
||||
squeezed_batch = False
|
||||
staged_input_ids = input_ids
|
||||
|
||||
if input_ids.ndim == 0:
|
||||
staged_input_ids = input_ids.view(1, 1)
|
||||
squeezed_batch = True
|
||||
elif input_ids.ndim == 1:
|
||||
staged_input_ids = input_ids.unsqueeze(0)
|
||||
squeezed_batch = True
|
||||
|
||||
embeddings = self.model.embed_input_ids(
|
||||
staged_input_ids,
|
||||
multimodal_embeddings=multimodal_embeddings,
|
||||
is_multimodal=is_multimodal,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if squeezed_batch and isinstance(embeddings, torch.Tensor):
|
||||
if embeddings.ndim == 3 and embeddings.shape[0] == 1:
|
||||
return embeddings.squeeze(0)
|
||||
if embeddings.ndim == 2 and input_ids.ndim == 0 and embeddings.shape[0] == 1:
|
||||
return embeddings
|
||||
|
||||
return embeddings
|
||||
|
||||
def embed_multimodal(self, **kwargs: Any) -> Any:
|
||||
if hasattr(self.model, "embed_multimodal"):
|
||||
return self.model.embed_multimodal(**kwargs)
|
||||
|
||||
self._collect_multimodal_inputs(**kwargs)
|
||||
return None
|
||||
|
||||
def load_weights(
|
||||
self,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> set[str]:
|
||||
loaded = self.model.load_weights(weights)
|
||||
if loaded is None:
|
||||
loaded = set()
|
||||
|
||||
expected_param_names = {name for name, _ in self.named_parameters()}
|
||||
if not expected_param_names:
|
||||
return loaded
|
||||
|
||||
return self._normalize_loaded_weight_names(loaded, expected_param_names)
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor | OmniOutput,
|
||||
sampling_metadata: Any = None,
|
||||
) -> torch.Tensor | None:
|
||||
return self.model.compute_logits(hidden_states, sampling_metadata=sampling_metadata)
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> SamplerOutput | None:
|
||||
if hasattr(self.model, "sample"):
|
||||
return self.model.sample(logits, sampling_metadata)
|
||||
if self.language_model is not None and hasattr(self.language_model, "sample"):
|
||||
return self.language_model.sample(logits, sampling_metadata)
|
||||
return None
|
||||
File diff suppressed because it is too large.
Load diff
@@ -1,313 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
from .dynin_omni import DyninOmniStageBase
|
||||
from .dynin_omni_common import (
|
||||
DetokTarget,
|
||||
_looks_like_hf_repo_id,
|
||||
coerce_token_ids_1d,
|
||||
normalize_runtime_info,
|
||||
resolve_dynin_infer_sources,
|
||||
resolve_hidden_size,
|
||||
unwrap_first_value,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _get_hf_token() -> str | None:
|
||||
return os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
|
||||
|
||||
|
||||
def _ensure_transformers_tied_weights_compat() -> None:
|
||||
"""Provide the transformers>=5 tied-weight attr expected during loading.
|
||||
|
||||
The DYNIN audio detokenizer uses Hugging Face remote code that still only
|
||||
defines the older ``_tied_weights_keys`` attribute. transformers 5.x now
|
||||
reads ``all_tied_weights_keys`` inside ``from_pretrained`` before the model
|
||||
instance is returned, so patching the loaded object is too late.
|
||||
"""
|
||||
try:
|
||||
from transformers import PreTrainedModel
|
||||
except Exception:
|
||||
return
|
||||
|
||||
if hasattr(PreTrainedModel, "all_tied_weights_keys"):
|
||||
return
|
||||
|
||||
def all_tied_weights_keys(self: Any) -> dict[str, Any]:
|
||||
keys: dict[str, Any] = {}
|
||||
for attr_name in ("_tied_weights_keys", "_dynamic_tied_weights_keys"):
|
||||
raw_keys = getattr(self, attr_name, None)
|
||||
if raw_keys is None:
|
||||
continue
|
||||
if isinstance(raw_keys, dict):
|
||||
keys.update({str(k): v for k, v in raw_keys.items()})
|
||||
continue
|
||||
if isinstance(raw_keys, str):
|
||||
keys[raw_keys] = None
|
||||
continue
|
||||
for key in raw_keys:
|
||||
keys[str(key)] = None
|
||||
return keys
|
||||
|
||||
PreTrainedModel.all_tied_weights_keys = property(all_tied_weights_keys) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def _ensure_remote_s2u_vendor_root(
|
||||
*,
|
||||
repo_id: str,
|
||||
local_files_only: bool,
|
||||
) -> str | None:
|
||||
if local_files_only or not _looks_like_hf_repo_id(repo_id):
|
||||
return None
|
||||
|
||||
existing = os.environ.get("DYNIN_S2U_VENDOR_ROOT")
|
||||
if existing:
|
||||
existing_path = Path(existing).expanduser().resolve()
|
||||
if existing_path.is_dir():
|
||||
return str(existing_path)
|
||||
|
||||
try:
|
||||
from vllm_omni.transformers_utils.repo_utils import hf_api
|
||||
except Exception as e:
|
||||
logger.warning("huggingface_hub unavailable; cannot fetch s2u_vendor from %s: %s", repo_id, e)
|
||||
return None
|
||||
|
||||
token = _get_hf_token()
|
||||
last_error: Exception | None = None
|
||||
revisions: list[str | None] = [None]
|
||||
|
||||
for revision in revisions:
|
||||
try:
|
||||
snapshot_dir = hf_api().snapshot_download(
|
||||
repo_id=repo_id,
|
||||
revision=revision,
|
||||
allow_patterns=["s2u_vendor/**"],
|
||||
token=token,
|
||||
)
|
||||
except TypeError:
|
||||
try:
|
||||
snapshot_dir = hf_api().snapshot_download(
|
||||
repo_id=repo_id,
|
||||
revision=revision,
|
||||
allow_patterns=["s2u_vendor/**"],
|
||||
)
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
continue
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
continue
|
||||
|
||||
vendor_root = (Path(snapshot_dir) / "s2u_vendor").resolve()
|
||||
if vendor_root.is_dir():
|
||||
os.environ["DYNIN_S2U_VENDOR_ROOT"] = str(vendor_root)
|
||||
logger.info("Using remote S2U vendor root: %s", vendor_root)
|
||||
return str(vendor_root)
|
||||
|
||||
if last_error is not None:
|
||||
logger.warning("Failed to download remote s2u_vendor from %s: %s", repo_id, last_error)
|
||||
return None
|
||||
|
||||
|
||||
class DyninOmniToken2Audio(DyninOmniStageBase):
|
||||
"""Stage-3: token detokenization to speech (or pass-through)."""
|
||||
|
||||
stage_name = "Dynin token2audio"
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
del prefix
|
||||
super().__init__()
|
||||
self.vllm_config = vllm_config
|
||||
self.have_multimodal_outputs = True
|
||||
self.requires_raw_input_tokens = True
|
||||
self.hidden_size = resolve_hidden_size(vllm_config=vllm_config)
|
||||
self._vq_audio = None
|
||||
self._vq_audio_path: str | None = None
|
||||
self._vq_audio_local_files_only: bool | None = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
positions: torch.Tensor | None = None,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: Any,
|
||||
) -> OmniOutput:
|
||||
del positions, intermediate_tensors, inputs_embeds
|
||||
if input_ids is None:
|
||||
raise ValueError("token2audio stage requires input_ids")
|
||||
|
||||
runtime_info = normalize_runtime_info(kwargs.get("runtime_additional_information"))
|
||||
detok_id = int(unwrap_first_value(runtime_info.get("detok_id"), 0))
|
||||
tokens = coerce_token_ids_1d(input_ids)
|
||||
|
||||
if detok_id != DetokTarget.AUDIO:
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={
|
||||
"token_ids": tokens,
|
||||
"detok_id": torch.tensor([detok_id], dtype=torch.long, device=tokens.device),
|
||||
},
|
||||
)
|
||||
|
||||
audio, sample_rate = self._decode_audio_tokens(tokens, runtime_info=runtime_info)
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={
|
||||
"speech": audio,
|
||||
"audio": audio,
|
||||
"sr": torch.tensor([sample_rate], dtype=torch.int, device=audio.device),
|
||||
"detok_id": torch.tensor([detok_id], dtype=torch.long, device=audio.device),
|
||||
},
|
||||
)
|
||||
|
||||
def _decode_audio_tokens(self, tokens: torch.Tensor, runtime_info: dict[str, Any]) -> tuple[torch.Tensor, int]:
|
||||
# Follow DYNIN validation path:
|
||||
# token list -> "<|speech_x|>" string -> vq_model_audio.decode(...).
|
||||
vq_audio = self._ensure_vq_audio(runtime_info=runtime_info, ref_device=tokens.device)
|
||||
|
||||
audio_codebook_size = int(unwrap_first_value(runtime_info.get("audio_codebook_size"), 4096))
|
||||
audio_vocab_offset = unwrap_first_value(
|
||||
runtime_info.get("audio_vocab_offset"),
|
||||
unwrap_first_value(runtime_info.get("t2s_vocab_start"), None),
|
||||
)
|
||||
|
||||
token_ids = tokens.to(torch.long)
|
||||
if audio_vocab_offset is not None:
|
||||
off = int(audio_vocab_offset)
|
||||
token_ids = torch.where(token_ids >= off, token_ids - off, token_ids)
|
||||
token_ids = token_ids[(token_ids >= 0) & (token_ids < audio_codebook_size)]
|
||||
if token_ids.numel() == 0:
|
||||
raise RuntimeError("Audio detokenizer got no valid audio token ids.")
|
||||
|
||||
speech_unit_str = " ".join(map(str, token_ids.detach().cpu().tolist()))
|
||||
speech_unit_for_decode = "".join(f"<|speech_{unit}|>" for unit in speech_unit_str.split(" ") if unit != "")
|
||||
|
||||
condition = unwrap_first_value(
|
||||
runtime_info.get("condition"),
|
||||
unwrap_first_value(runtime_info.get("t2s_condition"), None),
|
||||
)
|
||||
output_wav_file = unwrap_first_value(runtime_info.get("output_wav_file"), None)
|
||||
created_tmp = False
|
||||
if output_wav_file is None:
|
||||
fd, tmp_wav = tempfile.mkstemp(prefix="dynin_t2s_", suffix=".wav")
|
||||
os.close(fd)
|
||||
output_wav_file = tmp_wav
|
||||
created_tmp = True
|
||||
|
||||
audio_array = vq_audio.decode(speech_unit_for_decode, condition=condition, output_wav_file=output_wav_file)
|
||||
if created_tmp:
|
||||
try:
|
||||
os.remove(output_wav_file)
|
||||
except Exception:
|
||||
pass
|
||||
if not isinstance(audio_array, torch.Tensor):
|
||||
audio_array = torch.as_tensor(audio_array, dtype=torch.float32, device=tokens.device)
|
||||
else:
|
||||
audio_array = audio_array.to(device=tokens.device, dtype=torch.float32)
|
||||
|
||||
if audio_array.ndim > 1:
|
||||
audio_array = audio_array.reshape(-1)
|
||||
audio_array = audio_array.contiguous()
|
||||
|
||||
sample_rate = int(
|
||||
unwrap_first_value(
|
||||
runtime_info.get("sr"),
|
||||
unwrap_first_value(runtime_info.get("sample_rate"), 24000),
|
||||
)
|
||||
)
|
||||
try:
|
||||
cfg = getattr(vq_audio, "u2s_config", None)
|
||||
cfg_sr = getattr(cfg, "sampling_rate", None)
|
||||
if cfg_sr is None:
|
||||
cfg_sr = getattr(getattr(cfg, "data", None), "sampling_rate", None)
|
||||
if cfg_sr is not None:
|
||||
sample_rate = int(cfg_sr)
|
||||
except Exception:
|
||||
pass
|
||||
return audio_array, sample_rate
|
||||
|
||||
def _ensure_vq_audio(self, runtime_info: dict[str, Any], ref_device: torch.device) -> Any:
|
||||
sources = resolve_dynin_infer_sources(vllm_config=self.vllm_config, runtime_info=runtime_info)
|
||||
model_path = str(sources.vq_audio_source)
|
||||
local_files_only = bool(sources.vq_audio_local_files_only)
|
||||
|
||||
_ensure_remote_s2u_vendor_root(
|
||||
repo_id=model_path,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
|
||||
if (
|
||||
self._vq_audio is None
|
||||
or self._vq_audio_path != model_path
|
||||
or self._vq_audio_local_files_only != local_files_only
|
||||
):
|
||||
logger.info(
|
||||
"Loading DYNIN audio detokenizer from %s (local_files_only=%s)",
|
||||
model_path,
|
||||
local_files_only,
|
||||
)
|
||||
try:
|
||||
from transformers import AutoModel
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
"transformers is required to load EMOVASpeechTokenizer remote code from Hugging Face."
|
||||
) from e
|
||||
|
||||
_ensure_transformers_tied_weights_compat()
|
||||
try:
|
||||
self._vq_audio = AutoModel.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
local_files_only=local_files_only,
|
||||
low_cpu_mem_usage=False,
|
||||
)
|
||||
except TypeError:
|
||||
try:
|
||||
self._vq_audio = AutoModel.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
except TypeError:
|
||||
self._vq_audio = AutoModel.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to load EMOVASpeechTokenizer from Hugging Face remote code for model path '{model_path}'."
|
||||
) from e
|
||||
|
||||
if not hasattr(self._vq_audio, "decode"):
|
||||
raise RuntimeError(
|
||||
"Loaded audio tokenizer does not expose decode(). "
|
||||
"Check HF config.json auto_map/model_type and ensure trust_remote_code=True."
|
||||
)
|
||||
self._vq_audio.eval()
|
||||
self._vq_audio.requires_grad_(False)
|
||||
self._vq_audio_path = model_path
|
||||
self._vq_audio_local_files_only = local_files_only
|
||||
if hasattr(self._vq_audio, "to"):
|
||||
self._vq_audio = self._vq_audio.to(ref_device)
|
||||
return self._vq_audio
|
||||
|
||||
def embed_multimodal(self, **kwargs: Any) -> Any:
|
||||
del kwargs
|
||||
return None
|
||||
@@ -1,150 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
from .dynin_omni import DyninOmniStageBase
|
||||
from .dynin_omni_common import (
|
||||
DetokTarget,
|
||||
_to_bool,
|
||||
coerce_token_ids_1d,
|
||||
get_dynin_magvit_attr,
|
||||
normalize_runtime_info,
|
||||
resolve_dynin_infer_sources,
|
||||
resolve_hidden_size,
|
||||
unwrap_first_value,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class DyninOmniToken2Image(DyninOmniStageBase):
|
||||
"""Stage-2: token detokenization to image (or pass-through)."""
|
||||
|
||||
stage_name = "Dynin token2image"
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
del prefix
|
||||
super().__init__()
|
||||
|
||||
self.vllm_config = vllm_config
|
||||
self.have_multimodal_outputs = True
|
||||
self.requires_raw_input_tokens = True
|
||||
self.hidden_size = resolve_hidden_size(vllm_config=vllm_config)
|
||||
self._vq_model = None
|
||||
self._vq_model_path: str | None = None
|
||||
self._vq_local_files_only: bool | None = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
positions: torch.Tensor | None = None,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: Any,
|
||||
) -> OmniOutput:
|
||||
del positions, intermediate_tensors, inputs_embeds
|
||||
if input_ids is None:
|
||||
raise ValueError("token2image stage requires input_ids")
|
||||
runtime_info = normalize_runtime_info(kwargs.get("runtime_additional_information"))
|
||||
detok_id = int(unwrap_first_value(runtime_info.get("detok_id"), 0))
|
||||
tokens = coerce_token_ids_1d(input_ids)
|
||||
|
||||
if detok_id != DetokTarget.IMAGE:
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={
|
||||
"token_ids": tokens,
|
||||
"detok_id": torch.tensor([detok_id], dtype=torch.long, device=tokens.device),
|
||||
},
|
||||
)
|
||||
|
||||
image = self._decode_image_tokens(tokens, runtime_info=runtime_info)
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={
|
||||
"image": image,
|
||||
"detok_id": torch.tensor([detok_id], dtype=torch.long, device=image.device),
|
||||
},
|
||||
)
|
||||
|
||||
def _decode_image_tokens(self, tokens: torch.Tensor, runtime_info: dict[str, Any]) -> torch.Tensor:
|
||||
# Follow DYNIN validation path:
|
||||
# tokens -> clamp -> vq_model.decode_code -> (x+1)/2 -> [0,1].
|
||||
vq_model = self._ensure_vq_model(runtime_info=runtime_info, ref_device=tokens.device)
|
||||
codebook_size = int(unwrap_first_value(runtime_info.get("codebook_size"), 8192))
|
||||
image_vocab_offset = unwrap_first_value(runtime_info.get("image_vocab_offset"), None)
|
||||
if image_vocab_offset is None:
|
||||
text_vocab_size = unwrap_first_value(runtime_info.get("text_vocab_size"), None)
|
||||
num_new_special_tokens = int(unwrap_first_value(runtime_info.get("num_new_special_tokens"), 0))
|
||||
if text_vocab_size is not None:
|
||||
image_vocab_offset = int(text_vocab_size) + num_new_special_tokens
|
||||
|
||||
token_ids = tokens.to(torch.long)
|
||||
if image_vocab_offset is not None:
|
||||
off = int(image_vocab_offset)
|
||||
token_ids = torch.where(token_ids >= off, token_ids - off, token_ids)
|
||||
token_ids = torch.clamp(token_ids, min=0, max=max(0, codebook_size - 1))
|
||||
token_ids = token_ids.unsqueeze(0)
|
||||
|
||||
decoded = vq_model.decode_code(token_ids)
|
||||
decoded = torch.clamp((decoded + 1.0) / 2.0, min=0.0, max=1.0)
|
||||
if decoded.ndim != 4 or decoded.shape[0] == 0:
|
||||
raise RuntimeError(f"Unexpected MAGVIT decode output shape: {tuple(decoded.shape)}")
|
||||
return decoded[0].contiguous()
|
||||
|
||||
def _ensure_vq_model(self, runtime_info: dict[str, Any], ref_device: torch.device) -> Any:
|
||||
sources = resolve_dynin_infer_sources(vllm_config=self.vllm_config, runtime_info=runtime_info)
|
||||
model_path = str(sources.vq_image_source)
|
||||
local_files_only = bool(sources.vq_image_local_files_only)
|
||||
if self._vq_model is None or self._vq_model_path != model_path or self._vq_local_files_only != local_files_only:
|
||||
disable_xet = unwrap_first_value(
|
||||
runtime_info.get("hf_hub_disable_xet"),
|
||||
unwrap_first_value(runtime_info.get("disable_hf_xet"), True),
|
||||
)
|
||||
if _to_bool(disable_xet, default=True):
|
||||
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
|
||||
logger.info(
|
||||
"Loading DYNIN image detokenizer from %s (local_files_only=%s)",
|
||||
model_path,
|
||||
local_files_only,
|
||||
)
|
||||
try:
|
||||
MAGVITv2 = get_dynin_magvit_attr(
|
||||
"MAGVITv2",
|
||||
source=model_path,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
try:
|
||||
self._vq_model = MAGVITv2.from_pretrained(
|
||||
model_path,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
except TypeError:
|
||||
self._vq_model = MAGVITv2.from_pretrained(model_path)
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
"Failed to load MAGVITv2 from local DYNIN submodel implementation "
|
||||
f"for model path '{model_path}'. "
|
||||
"If your environment cannot access huggingface.co, set "
|
||||
"additional_information.vq_model_image_path to a local MAGVITv2 directory "
|
||||
"and set additional_information.vq_model_image_local_files_only=true."
|
||||
) from e
|
||||
self._vq_model.eval()
|
||||
self._vq_model.requires_grad_(False)
|
||||
self._vq_model_path = model_path
|
||||
self._vq_local_files_only = local_files_only
|
||||
if hasattr(self._vq_model, "to"):
|
||||
self._vq_model = self._vq_model.to(ref_device)
|
||||
return self._vq_model
|
||||
|
||||
def embed_multimodal(self, **kwargs: Any) -> Any:
|
||||
del kwargs
|
||||
return None
|
||||
File diff suppressed because it is too large.
Load diff
@@ -1,67 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Dynin-Omni pipeline topology (frozen).
|
||||
|
||||
Stage 0: token2text — multimodal understanding / text generation (comprehension)
|
||||
Stage 1: token2image — stage-0 tokens → image latents
|
||||
Stage 2: token2audio — stage-1 tokens → audio latents
|
||||
|
||||
All three stages run on the generation worker (``LLM_GENERATION``); the
|
||||
inter-stage hand-off uses the worker-connector full-payload data plane
|
||||
(``*_full_payload`` producers + ``token2text_to_token2image`` /
|
||||
``token2image_to_token2audio`` consumers). Deploy knobs (devices, GPU memory,
|
||||
batched tokens, connectors) live in ``vllm_omni/deploy/dynin_omni*.yaml``.
|
||||
"""
|
||||
|
||||
from vllm_omni.config.stage_config import (
|
||||
PipelineConfig,
|
||||
StageExecutionType,
|
||||
StagePipelineConfig,
|
||||
)
|
||||
|
||||
_PROC = "vllm_omni.model_executor.stage_input_processors.dynin_omni"
|
||||
|
||||
DYNIN_OMNI_PIPELINE = PipelineConfig(
|
||||
model_type="dynin_omni",
|
||||
default_deploy_config_name="dynin_omni.yaml",
|
||||
model_arch="DyninOmniForConditionalGeneration",
|
||||
# Arch-fallback safety net: route by hf_config.architectures when the
|
||||
# auto-detected model_type does not match the registry key exactly.
|
||||
hf_architectures=("DyninOmniForConditionalGeneration",),
|
||||
stages=(
|
||||
StagePipelineConfig(
|
||||
stage_id=0,
|
||||
model_stage="token2text",
|
||||
execution_type=StageExecutionType.LLM_GENERATION,
|
||||
input_sources=(),
|
||||
final_output=True,
|
||||
final_output_type="text",
|
||||
owns_tokenizer=True,
|
||||
engine_output_type="latent",
|
||||
custom_process_next_stage_input_func=f"{_PROC}.token2text_to_token2image_full_payload",
|
||||
),
|
||||
StagePipelineConfig(
|
||||
stage_id=1,
|
||||
model_stage="token2image",
|
||||
execution_type=StageExecutionType.LLM_GENERATION,
|
||||
input_sources=(0,),
|
||||
final_output=True,
|
||||
final_output_type="image",
|
||||
engine_output_type="latent",
|
||||
custom_process_input_func=f"{_PROC}.token2text_to_token2image",
|
||||
custom_process_next_stage_input_func=f"{_PROC}.token2image_to_token2audio_full_payload",
|
||||
requires_full_payload_input=True,
|
||||
),
|
||||
StagePipelineConfig(
|
||||
stage_id=2,
|
||||
model_stage="token2audio",
|
||||
execution_type=StageExecutionType.LLM_GENERATION,
|
||||
input_sources=(1,),
|
||||
final_output=True,
|
||||
final_output_type="audio",
|
||||
engine_output_type="latent",
|
||||
custom_process_input_func=f"{_PROC}.token2image_to_token2audio",
|
||||
requires_full_payload_input=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -171,10 +171,9 @@ class Qwen3OmniMoeForConditionalGeneration(
|
||||
# such as the vLLM-text perf benchmark.
|
||||
#
|
||||
# Default to the thinker: it is the text-generation stage, which is what
|
||||
# a non-staged run of this model is asking for. Same shape as
|
||||
# dynin_omni ("token2text") and glm_tts ("glm_tts"). An explicitly wrong
|
||||
# value still reaches the ValueError below rather than being silently
|
||||
# accepted.
|
||||
# a non-staged run of this model is asking for. Same shape as glm_tts
|
||||
# ("glm_tts"). An explicitly wrong value still reaches the ValueError
|
||||
# below rather than being silently accepted.
|
||||
self.model_stage = getattr(vllm_config.model_config, "model_stage", None) or "thinker"
|
||||
# Staged startup always injects model_stage; its absence means a plain
|
||||
# vLLM run with no talker stage downstream, so no one consumes captured
|
||||
|
||||
@@ -321,12 +321,6 @@ _OMNI_MODELS = {
|
||||
"voxcpm2_talker",
|
||||
"VoxCPM2TalkerForConditionalGeneration",
|
||||
),
|
||||
## dots.tts
|
||||
"DotsTTSForConditionalGeneration": (
|
||||
"dots_tts",
|
||||
"dots_tts_talker",
|
||||
"DotsTTSForConditionalGeneration",
|
||||
),
|
||||
## Voxtral TTS
|
||||
"VoxtralTTSForConditionalGeneration": (
|
||||
"voxtral_tts",
|
||||
@@ -396,11 +390,6 @@ _OMNI_MODELS = {
|
||||
"modeling_moss_tts_codec",
|
||||
"MossTTSCodecDecoder",
|
||||
),
|
||||
"DyninOmniForConditionalGeneration": (
|
||||
"dynin_omni",
|
||||
"dynin_omni",
|
||||
"DyninOmniForConditionalGeneration",
|
||||
),
|
||||
## IndexTTS2
|
||||
"IndexTTS2TalkerForConditionalGeneration": (
|
||||
"indextts2",
|
||||
|
||||
@@ -1,276 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.inputs import TextPrompt
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _to_prompt_dict(prompt_item: OmniTokensPrompt | TextPrompt | str | None) -> dict[str, Any]:
|
||||
if isinstance(prompt_item, dict):
|
||||
return prompt_item
|
||||
return {}
|
||||
|
||||
|
||||
def _to_token_id_list(value: Any) -> list[int]:
|
||||
if isinstance(value, torch.Tensor):
|
||||
value = value.detach().to("cpu")
|
||||
if value.ndim == 0:
|
||||
return [int(value.item())]
|
||||
if value.ndim > 1:
|
||||
value = value[0]
|
||||
return [int(x) for x in value.tolist()]
|
||||
if isinstance(value, list):
|
||||
if not value:
|
||||
return []
|
||||
if isinstance(value[0], list):
|
||||
return [int(x) for x in value[0]]
|
||||
return [int(x) for x in value]
|
||||
if value is None:
|
||||
return []
|
||||
return [int(value)]
|
||||
|
||||
|
||||
def _to_int(value: Any, default: int = 0) -> int:
|
||||
if isinstance(value, torch.Tensor):
|
||||
if value.numel() == 0:
|
||||
return default
|
||||
return int(value.view(-1)[0].item())
|
||||
if isinstance(value, list):
|
||||
if not value:
|
||||
return default
|
||||
return int(value[0])
|
||||
if value is None:
|
||||
return default
|
||||
return int(value)
|
||||
|
||||
|
||||
def _normalize_additional_info(value: Any) -> dict[str, Any]:
|
||||
if not isinstance(value, dict):
|
||||
return {}
|
||||
normalized: dict[str, Any] = {}
|
||||
for key, val in value.items():
|
||||
if isinstance(val, list):
|
||||
normalized[key] = val
|
||||
else:
|
||||
normalized[key] = [val]
|
||||
return normalized
|
||||
|
||||
|
||||
def _decode_runtime_bridge_info(value: Any) -> dict[str, Any]:
|
||||
if isinstance(value, torch.Tensor):
|
||||
tensor = value.detach().to("cpu").reshape(-1).to(torch.uint8)
|
||||
raw = bytes(tensor.tolist())
|
||||
elif isinstance(value, (bytes, bytearray)):
|
||||
raw = bytes(value)
|
||||
elif isinstance(value, list):
|
||||
try:
|
||||
raw = bytes(int(item) for item in value)
|
||||
except Exception:
|
||||
return {}
|
||||
elif value is None:
|
||||
return {}
|
||||
else:
|
||||
return value if isinstance(value, dict) else {}
|
||||
|
||||
if not raw:
|
||||
return {}
|
||||
|
||||
try:
|
||||
decoded = json.loads(raw.decode("utf-8"))
|
||||
except Exception:
|
||||
return {}
|
||||
return decoded if isinstance(decoded, dict) else {}
|
||||
|
||||
|
||||
def _bridge_tokens(
|
||||
source_outputs,
|
||||
prompt: OmniTokensPrompt | TextPrompt = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
):
|
||||
next_inputs = []
|
||||
if not isinstance(prompt, list):
|
||||
prompt = [prompt]
|
||||
|
||||
prompt_meta_by_reqid = {src_out.request_id: _to_prompt_dict(p) for src_out, p in zip(source_outputs, prompt)}
|
||||
|
||||
for source_output in source_outputs:
|
||||
output = source_output.outputs[0]
|
||||
mm_out = getattr(output, "multimodal_output", None) or {}
|
||||
|
||||
token_ids = _to_token_id_list(mm_out.get("token_ids"))
|
||||
if not token_ids:
|
||||
token_ids = _to_token_id_list(mm_out.get("text_tokens"))
|
||||
if not token_ids:
|
||||
token_ids = list(output.cumulative_token_ids or [])
|
||||
if not token_ids:
|
||||
raise RuntimeError(f"Stage output for request {source_output.request_id} has no token_ids")
|
||||
|
||||
detok_id = _to_int(mm_out.get("detok_id"), default=0)
|
||||
src_prompt = prompt_meta_by_reqid.get(source_output.request_id, {})
|
||||
src_additional_info = src_prompt.get("additional_information", {}) or {}
|
||||
runtime_bridge_info = _decode_runtime_bridge_info(mm_out.get("runtime_info_json"))
|
||||
if not runtime_bridge_info:
|
||||
runtime_bridge_info = mm_out.get("runtime_info", {}) or {}
|
||||
|
||||
additional_information: dict[str, Any] = _normalize_additional_info(src_additional_info)
|
||||
additional_information.update(_normalize_additional_info(runtime_bridge_info))
|
||||
additional_information["detok_id"] = [detok_id]
|
||||
|
||||
next_inputs.append(
|
||||
OmniTokensPrompt(
|
||||
prompt_token_ids=token_ids,
|
||||
additional_information=additional_information,
|
||||
multi_modal_data=(src_prompt.get("multi_modal_data") if requires_multimodal_data else None),
|
||||
mm_processor_kwargs=None,
|
||||
)
|
||||
)
|
||||
|
||||
return next_inputs
|
||||
|
||||
|
||||
def token2text_to_token2image(
|
||||
source_outputs,
|
||||
prompt: OmniTokensPrompt | TextPrompt = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
):
|
||||
return _bridge_tokens(source_outputs, prompt, requires_multimodal_data)
|
||||
|
||||
|
||||
def token2image_to_token2audio(
|
||||
source_outputs,
|
||||
prompt: OmniTokensPrompt | TextPrompt = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
):
|
||||
return _bridge_tokens(source_outputs, prompt, requires_multimodal_data)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Worker-connector data plane (non-async-chunk path).
|
||||
# ============================================================================
|
||||
|
||||
# Per-model REPLACE-keys for the full-payload accumulator. dynin_omni's
|
||||
# producer model emits new chunks per step (token_ids / runtime_info_json),
|
||||
# all of which use the default CONCAT/replace semantics — no model_outputs
|
||||
# entry needs explicit REPLACE.
|
||||
_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset()
|
||||
|
||||
|
||||
def _build_full_payload(pooling_output: dict[str, Any] | None, request: Any) -> dict[str, Any] | None:
|
||||
"""Producer-side payload builder: assemble dynin_omni connector payload.
|
||||
|
||||
Reads token_ids from ``pooling_output["token_ids"]`` (preferred) or
|
||||
``request.output_token_ids`` (fallback). Reads structured non-tensor
|
||||
metadata from ``pooling_output["runtime_info_json"]`` (JSON-in-uint8)
|
||||
if present, falling back to ``pooling_output["runtime_info"]`` dict.
|
||||
Carries forward ``request.additional_information`` so prompt-side
|
||||
metadata (speaker / language / detok_id) survives the IPC boundary.
|
||||
"""
|
||||
if not isinstance(pooling_output, dict):
|
||||
pooling_output = {}
|
||||
|
||||
token_ids = _to_token_id_list(pooling_output.get("token_ids"))
|
||||
if not token_ids:
|
||||
token_ids = _to_token_id_list(pooling_output.get("text_tokens"))
|
||||
if not token_ids and request is not None:
|
||||
token_ids = _to_token_id_list(getattr(request, "output_token_ids", None))
|
||||
if not token_ids:
|
||||
logger.warning(
|
||||
"dynin_omni._build_full_payload: no token_ids found in pooling_output "
|
||||
"(keys=%s) or request.output_token_ids for req=%s; consumer wait gate may hang.",
|
||||
list(pooling_output.keys()),
|
||||
getattr(request, "request_id", "?"),
|
||||
)
|
||||
return None
|
||||
|
||||
src_additional_info = getattr(request, "additional_information", {}) if request is not None else {}
|
||||
if not isinstance(src_additional_info, dict):
|
||||
src_additional_info = {}
|
||||
|
||||
runtime_bridge_info = _decode_runtime_bridge_info(pooling_output.get("runtime_info_json"))
|
||||
if not runtime_bridge_info:
|
||||
runtime_bridge_info = pooling_output.get("runtime_info", {}) or {}
|
||||
|
||||
payload = _normalize_additional_info(src_additional_info)
|
||||
payload.update(_normalize_additional_info(runtime_bridge_info))
|
||||
payload["detok_id"] = [_to_int(pooling_output.get("detok_id"), default=_to_int(payload.get("detok_id"), default=0))]
|
||||
# Use nested OmniPayload shape so the scheduling-metadata extractor in
|
||||
# OmniConnectorModelRunnerMixin reads codes.audio and meta.finished
|
||||
# (flat keys at the top level are silently dropped with a warning).
|
||||
payload["codes"] = {"audio": token_ids}
|
||||
payload["meta"] = {"finished": torch.tensor(True, dtype=torch.bool)}
|
||||
return payload
|
||||
|
||||
|
||||
def token2text_to_token2image_full_payload(
|
||||
transfer_manager: Any,
|
||||
pooling_output: dict[str, Any],
|
||||
request: Any,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Producer-side payload builder for the Stage-0 → Stage-1 (text → image) transition."""
|
||||
del transfer_manager
|
||||
return _build_full_payload(pooling_output, request)
|
||||
|
||||
|
||||
def token2image_to_token2audio_full_payload(
|
||||
transfer_manager: Any,
|
||||
pooling_output: dict[str, Any],
|
||||
request: Any,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Producer-side payload builder for the Stage-1 → Stage-2 (image → audio) transition."""
|
||||
del transfer_manager
|
||||
return _build_full_payload(pooling_output, request)
|
||||
|
||||
|
||||
def _token_only_from_source(source_outputs: list[Any]) -> list[OmniTokensPrompt]:
|
||||
"""Length-only placeholder list mirroring ``_bridge_tokens`` token counts."""
|
||||
inputs: list[OmniTokensPrompt] = []
|
||||
for source_output in source_outputs:
|
||||
output = source_output.outputs[0]
|
||||
mm_out = getattr(output, "multimodal_output", None) or {}
|
||||
token_ids = _to_token_id_list(mm_out.get("token_ids"))
|
||||
if not token_ids:
|
||||
token_ids = _to_token_id_list(mm_out.get("text_tokens"))
|
||||
if not token_ids:
|
||||
token_ids = list(getattr(output, "token_ids", []) or [])
|
||||
if not token_ids:
|
||||
token_ids = [0]
|
||||
inputs.append(
|
||||
OmniTokensPrompt(
|
||||
prompt_token_ids=[0] * len(token_ids),
|
||||
additional_information=None,
|
||||
multi_modal_data=None,
|
||||
mm_processor_kwargs=None,
|
||||
)
|
||||
)
|
||||
return inputs
|
||||
|
||||
|
||||
def token2text_to_token2image_token_only(
|
||||
stage_list,
|
||||
engine_input_source,
|
||||
prompt: OmniTokensPrompt | TextPrompt = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
) -> list[OmniTokensPrompt]:
|
||||
"""Sync-side placeholder for Stage-1 input (token2image)."""
|
||||
source_stage_id = engine_input_source[0] if engine_input_source else 0
|
||||
source_outputs = stage_list[source_stage_id].engine_outputs
|
||||
return _token_only_from_source(source_outputs)
|
||||
|
||||
|
||||
def token2image_to_token2audio_token_only(
|
||||
stage_list,
|
||||
engine_input_source,
|
||||
prompt: OmniTokensPrompt | TextPrompt = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
) -> list[OmniTokensPrompt]:
|
||||
"""Sync-side placeholder for Stage-2 input (token2audio)."""
|
||||
source_stage_id = engine_input_source[0] if engine_input_source else 0
|
||||
source_outputs = stage_list[source_stage_id].engine_outputs
|
||||
return _token_only_from_source(source_outputs)
|
||||
@@ -23,7 +23,6 @@ _CLASS_TO_MODULE: dict[str, str] = {
|
||||
"FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech",
|
||||
"GLMTTSConfig": "vllm_omni.transformers_utils.configs.glm_tts",
|
||||
"VoxCPM2Config": "vllm_omni.transformers_utils.configs.voxcpm2",
|
||||
"DotsTTSConfig": "vllm_omni.transformers_utils.configs.dots_tts",
|
||||
"VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts",
|
||||
"CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3",
|
||||
"MiniMaxMusic3Config": "vllm_omni.transformers_utils.configs.minimax_music3",
|
||||
@@ -53,7 +52,6 @@ __all__ = [
|
||||
"FishSpeechFastARConfig",
|
||||
"GLMTTSConfig",
|
||||
"VoxCPM2Config",
|
||||
"DotsTTSConfig",
|
||||
"VoxtralTTSConfig",
|
||||
"CosyVoice3Config",
|
||||
"MiniMaxMusic3Config",
|
||||
@@ -86,7 +84,6 @@ def __dir__():
|
||||
from vllm_omni.transformers_utils.configs import audio8_tts as _audio8_tts # noqa: F401, E402
|
||||
from vllm_omni.transformers_utils.configs import auk as _auk # noqa: F401, E402
|
||||
from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3 # noqa: F401, E402
|
||||
from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts # noqa: F401, E402
|
||||
from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402
|
||||
from vllm_omni.transformers_utils.configs import glm_tts as _glm_tts # noqa: F401, E402
|
||||
from vllm_omni.transformers_utils.configs import higgs_audio_v3 as _higgs_audio_v3 # noqa: F401, E402
|
||||
|
||||
@@ -1,102 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""HF config for dots.tts.
|
||||
|
||||
dots.tts splits its config across two files in the HF checkpoint:
|
||||
* ``config.json`` — top-level dots.tts fields (latent_dim / patch_size
|
||||
/ PatchEncoder / DiT / vocoder / etc.).
|
||||
* ``llm_config.json`` — full Qwen2 LM config (vocab_size / hidden_size /
|
||||
num_hidden_layers / rope_theta / etc.).
|
||||
|
||||
``AutoConfig.from_pretrained`` only consumes ``config.json``: its dots.tts
|
||||
sub-blocks (``vocoder`` / ``DiT`` / ``PatchEncoder``) arrive through
|
||||
``**kwargs`` and land as attributes, which the talker's config factories
|
||||
consume directly. The Qwen2 LM fields live in ``llm_config.json``, which
|
||||
AutoConfig never reads — this class hoists them to top-level attributes
|
||||
with defaults matching the released checkpoints (soar / base / mf all ship
|
||||
the same Qwen2.5-1.5B backbone). vLLM's ``Qwen2Model.__init__`` then reads
|
||||
``config.hidden_size`` / ``config.vocab_size`` / etc. through
|
||||
``get_text_config()``, the same pattern voxcpm2 follows.
|
||||
|
||||
Future work: read ``llm_config.json`` during ``from_pretrained`` so a
|
||||
checkpoint with a different LM shape works without editing this file.
|
||||
|
||||
Upstream reference:
|
||||
https://huggingface.co/dots-studio/dots.tts-soar/blob/main/llm_config.json
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from transformers import AutoConfig
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
|
||||
|
||||
class DotsTTSConfig(PretrainedConfig):
|
||||
"""Configuration for dots.tts integration (Qwen2.5-1.5B-Base AR backbone)."""
|
||||
|
||||
model_type = "dots_tts"
|
||||
keys_to_ignore_at_inference = ["past_key_values"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
# -- top-level dots.tts params --
|
||||
lm_config: dict | None = None,
|
||||
# -- LM defaults (overridden by lm_config dict if present) --
|
||||
# All values match dots.tts-soar's llm_config.json.
|
||||
vocab_size: int = 151672,
|
||||
hidden_size: int = 1536,
|
||||
intermediate_size: int = 8960,
|
||||
num_hidden_layers: int = 28,
|
||||
num_attention_heads: int = 12,
|
||||
num_key_value_heads: int = 2,
|
||||
max_position_embeddings: int = 131072,
|
||||
rms_norm_eps: float = 1e-6,
|
||||
rope_theta: float = 1_000_000.0,
|
||||
tie_word_embeddings: bool = True,
|
||||
hidden_act: str = "silu",
|
||||
bos_token_id: int = 151643,
|
||||
eos_token_id: int = 151643,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(
|
||||
bos_token_id=bos_token_id,
|
||||
eos_token_id=eos_token_id,
|
||||
tie_word_embeddings=tie_word_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
# The dots.tts sub-blocks (vocoder / DiT / PatchEncoder / latent_dim
|
||||
# / patch_size / ...) arrive via **kwargs and are stored as
|
||||
# attributes by PretrainedConfig; the talker's config factories
|
||||
# read them from there.
|
||||
self.lm_config = lm_config or {}
|
||||
|
||||
# -- Hoist LM parameters to top-level for vLLM's Qwen2Model --
|
||||
lm = self.lm_config
|
||||
self.vocab_size = lm.get("vocab_size", vocab_size)
|
||||
self.hidden_size = lm.get("hidden_size", hidden_size)
|
||||
self.intermediate_size = lm.get("intermediate_size", intermediate_size)
|
||||
self.num_hidden_layers = lm.get("num_hidden_layers", num_hidden_layers)
|
||||
self.num_attention_heads = lm.get("num_attention_heads", num_attention_heads)
|
||||
self.num_key_value_heads = lm.get("num_key_value_heads", num_key_value_heads)
|
||||
self.max_position_embeddings = lm.get("max_position_embeddings", max_position_embeddings)
|
||||
self.rms_norm_eps = lm.get("rms_norm_eps", rms_norm_eps)
|
||||
self.rope_theta = lm.get("rope_theta", rope_theta)
|
||||
self.hidden_act = lm.get("hidden_act", hidden_act)
|
||||
self.head_dim = self.hidden_size // self.num_attention_heads
|
||||
|
||||
# vLLM's Qwen2Model reads ``max_window_layers`` and checks against
|
||||
# ``num_hidden_layers`` only when ``use_sliding_window`` is true.
|
||||
# soar disables sliding window — set both to the same value so the
|
||||
# assertion (if executed) stays vacuous.
|
||||
self.use_sliding_window = lm.get("use_sliding_window", False)
|
||||
self.sliding_window = lm.get("sliding_window", None)
|
||||
self.max_window_layers = lm.get("max_window_layers", self.num_hidden_layers)
|
||||
|
||||
def get_text_config(self, **kwargs):
|
||||
"""Return self — LM fields are hoisted to top-level attributes."""
|
||||
return self
|
||||
|
||||
|
||||
AutoConfig.register("dots_tts", DotsTTSConfig)
|
||||
|
||||
__all__ = ["DotsTTSConfig"]
|
||||
Reference in new issue
Block a user