diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b9686047a85..5af1978f0ab 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -58,6 +58,7 @@ th { | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — | | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — | | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — | +| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index d246f9e45af..64e9b4d4205 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -23,6 +23,7 @@ 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 | `rednote-hilab/dots.tts-soar` | single (native AR) | — (not wired yet) | — | — | 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 | @@ -430,6 +431,31 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [ --- +## 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 rednote-hilab/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 zero-shot only. The CAM++ x-vector speaker encoder weights load, but `end2end.py` has no `--ref-audio`/`--ref-text` flags yet. + +### Streaming +The AudioVAE decoder has an internal streaming path (`init_stream_state` / `stream_step` / `stream_flush`) used to avoid boundary artifacts between 160 ms patches, but it is not yet exposed through an online serving endpoint or example. + +### Notes +- Output: 48 kHz mono WAV. +- Deploy config: `vllm_omni/deploy/dots_tts.yaml` (auto-loaded by HF `model_type`). +- Checkpoints: `rednote-hilab/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. diff --git a/examples/offline_inference/text_to_speech/dots_tts/end2end.py b/examples/offline_inference/text_to_speech/dots_tts/end2end.py new file mode 100644 index 00000000000..8705f03f3f3 --- /dev/null +++ b/examples/offline_inference/text_to_speech/dots_tts/end2end.py @@ -0,0 +1,125 @@ +"""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: rednote-hilab/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="rednote-hilab/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() diff --git a/recipes/rednote-hilab/dots.tts.md b/recipes/rednote-hilab/dots.tts.md new file mode 100644 index 00000000000..d5c7868cba9 --- /dev/null +++ b/recipes/rednote-hilab/dots.tts.md @@ -0,0 +1,152 @@ +# dots.tts + +> Offline continuous-AR TTS at 48 kHz (rednote-hilab) + +## Summary + +- Vendor: rednote-hilab +- Model: `rednote-hilab/dots.tts-soar` +- Task: Text-to-speech, zero-shot synthesis only +- Mode: Offline end-to-end example only (not yet wired for online serving — + see [Known limitations](#known-limitations)) +- Maintainer: Community + +## When to use this recipe + +Use this recipe as a known-good starting point for running +`rednote-hilab/dots.tts-soar` offline 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 (see [Known limitations](#known-limitations) +below before depending on it for anything beyond offline experimentation). + +## 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) +- 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 rednote-hilab/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 ` to override. + +#### Verification + +**T1 — offline zero-shot synthesis**: + +```bash +python examples/offline_inference/text_to_speech/dots_tts/end2end.py \ + --model rednote-hilab/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)). + +#### Notes + +- Output: 48 kHz mono WAV. +- Checkpoints: `rednote-hilab/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 + +- **No online serving yet.** dots.tts is not registered in + `vllm_omni/entrypoints/openai/tts_adapters/` (the framework other TTS + models in this repo use for `/v1/audio/speech`), so `vllm serve + rednote-hilab/dots.tts-soar --omni` will start but the OpenAI-compatible + speech endpoint does not know how to build request params for it yet. + Offline `Omni()` (as used in `end2end.py`) is the only supported path + today. +- **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 a reference audio yet — generation is zero-shot + only. `end2end.py` has no `--ref-audio`/`--ref-text` flags. +- **`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. diff --git a/tests/model_executor/models/dots_tts/__init__.py b/tests/model_executor/models/dots_tts/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/tests/model_executor/models/dots_tts/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/model_executor/models/dots_tts/test_talker.py b/tests/model_executor/models/dots_tts/test_talker.py new file mode 100644 index 00000000000..ed3c6bcaa83 --- /dev/null +++ b/tests/model_executor/models/dots_tts/test_talker.py @@ -0,0 +1,137 @@ +# 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 == {} diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index f02a3578edb..f06259eaa88 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -47,6 +47,7 @@ ) 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 @@ -144,6 +145,7 @@ "hunyuan_video_15": HUNYUAN_VIDEO_15_PIPELINE, "wan2_2_ti2v": WAN2_2_TI2V_PIPELINE, "voxcpm2": VOXCPM2_PIPELINE, + "dots_tts": DOTS_TTS_PIPELINE, "cosyvoice3": COSYVOICE3_PIPELINE, "audex_tts": AUDEX_TTS_PIPELINE, "audex_tta": AUDEX_TTA_PIPELINE, diff --git a/vllm_omni/deploy/dots_tts.yaml b/vllm_omni/deploy/dots_tts.yaml new file mode 100644 index 00000000000..acd28de52ae --- /dev/null +++ b/vllm_omni/deploy/dots_tts.yaml @@ -0,0 +1,28 @@ +# 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 diff --git a/vllm_omni/model_executor/models/dots_tts/__init__.py b/vllm_omni/model_executor/models/dots_tts/__init__.py new file mode 100644 index 00000000000..fb6738a2967 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/__init__.py @@ -0,0 +1,8 @@ +# 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. diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py new file mode 100644 index 00000000000..086ac109ad5 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py @@ -0,0 +1,457 @@ +# 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) diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py new file mode 100644 index 00000000000..39d1a497389 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py @@ -0,0 +1,355 @@ +# 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 diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py new file mode 100644 index 00000000000..a2c4ac9f90b --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py @@ -0,0 +1,70 @@ +# 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 "[文本][文本对应语音]", 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 diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py new file mode 100644 index 00000000000..09961d0ab47 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py @@ -0,0 +1,732 @@ +# 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) diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py new file mode 100644 index 00000000000..368a3fd1e01 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py @@ -0,0 +1,1426 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""dots.tts talker — vLLM-native AR base LM + audio side path. + +Mirrors upstream rednote-hilab/dots.tts (pinned @ a393d2e): + + * Qwen2.5-1.5B base LM runs under vLLM PagedAttention (stock + ``Qwen2Model``); weights come from the checkpoint's ``llm.model.*`` + namespace. ``DotsTTSConfig`` hoists the LM fields to top-level + attributes. + * Everything else is side computation, invisible to the vLLM + scheduler and driven once per step by ``_finish_decode``: + DiT flow-matching head (N-step Euler, fp32 integration state, + deterministic per-request noise) → io_helper denormalize → + patch_encoder AR loopback (produces the next step's LLM input + embedding) + streaming AudioVAE decode (per-request + ``BigVGANStreamState`` sliding window, so every latent patch is + decoded with real left conv context) → 48 kHz wav chunks pushed to + ``_audio_queue`` → eos_proj stop signal via ``_results_queue`` + (threshold 0.8, upstream default; the stop step also drains the + vocoder's 2-frame lookahead tail via ``stream_flush``). + * Cross-step state is isolated per request in ``_RequestState`` + (keyed by request_id); eviction is deferred to the end of + ``forward()`` so a finishing request's audio still drains. + * The CAM++ speaker encoder is constructed and its weights load, but + voice cloning is not exposed yet — generation is zero-shot and the + DiT falls back to its null conditioning (``fm_null_g_cond``). + +Debugging: set ``DOTS_TTS_BETA_TRACE=1`` for an env-gated per-patch +rms/max trace across the full side path (zero overhead when unset). + +Reference implementations: ``vllm_omni/model_executor/models/voxcpm2/`` +(architectural template: vLLM-native base LM + side path + per-request +state) and ``ming_flash_omni`` (AudioVAE weight-loading pattern). +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import os +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.qwen2 import Qwen2Model +from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.dots_tts.dots_tts_dit import DiT +from vllm_omni.model_executor.models.dots_tts.dots_tts_patch_encoder import ( + VAESemanticEncoder, +) +from vllm_omni.model_executor.models.dots_tts.dots_tts_speaker_encoder import ( + SpeakerXVectorFeatures, +) +from vllm_omni.model_executor.models.dots_tts.dots_tts_vocoder import ( + AudioVAE, + AudioVAEConfig, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + + +def _require_config_block(hf_config: Any, name: str) -> dict[str, Any]: + """Fetch a sub-config dict (``vocoder`` / ``DiT`` / ``PatchEncoder``) + from the checkpoint's config.json, failing loudly when absent — + silently falling back to another checkpoint's values would produce + garbage audio, not an error.""" + block = getattr(hf_config, name, None) + if not isinstance(block, dict) or not block: + raise ValueError( + f"dots.tts checkpoint config.json has no '{name}' block — cannot " + "build the audio side path. Expected the upstream dots.tts config " + "layout (e.g. rednote-hilab/dots.tts-soar)." + ) + return block + + +def _build_audio_vae_config(hf_config: Any) -> AudioVAEConfig: + # AudioVAEConfig declares exactly the checkpoint's 17 ``vocoder`` keys, + # so an unknown / missing key raises TypeError instead of guessing. + return AudioVAEConfig(**_require_config_block(hf_config, "vocoder")) + + +# Architecture constants used by the per-request FM workspace math. All +# released dots.tts checkpoints (soar / base / mf) share these values; +# _validate_architecture_constants refuses a checkpoint that disagrees +# instead of silently generating garbage. +_FM_HIDDEN = 1024 # DiT.hidden_size (also LLM↔DiT projection target) +_LATENT_DIM = 128 # AudioVAE.latent_dim +_LATENT_PATCH_SIZE = 4 # config.patch_size — DiT samples 4 latent frames / audio patch +_HIDDEN_PATCH_SIZE = 1 # upstream core.py hardcodes self.hidden_patch_size = 1 +_MAX_AUDIO_PATCHES = 1024 # ~164 s @ 4×1920/48k; bounds per-request FM static buffer +_DIT_NUM_STEPS = int( + os.environ.get("DOTS_TTS_DIT_NUM_STEPS", "10") +) # env-gated; upstream default 10 for fixed-step Euler +_DIT_GUIDANCE_SCALE = 1.2 # upstream default guidance_scale for soar +_DIT_NOISE_SEED = 20260601 # base seed for per-request FM noise (voxcpm2 parity) +_PATCH_ENCODER_OUT_DS_RATE = 2 # patch_size / in_ds_rate = 4 / 2 (VAESemanticEncoder hardcodes in_ds_rate=2) + + +class _IOHelper: + """Latent stats wrapper. Holds (mean, var) over the 128-dim VAE latent + space and normalizes / denormalizes between DiT-internal space (which + operates on normalized latents) and AudioVAE / patch_encoder input space + (which expects raw latents). + + Mirrors upstream IOHelper (rednote-hilab/dots.tts @ a393d2e + models/dots_tts/core.py:684). Plain class (not nn.Module) so it stays + out of state_dict() and load_weights() doesn't have to mute its buffers. + Stats are loaded from latent_stats.pt — independent of safetensors. + """ + + def __init__(self, latent_stats_path: str | None = None) -> None: + if latent_stats_path is None: + self.global_mean: torch.Tensor | None = None + self.global_var: torch.Tensor | None = None + return + stats = torch.load(latent_stats_path, weights_only=False) + self.global_mean = torch.as_tensor(stats["mean"]) + self.global_var = torch.as_tensor(stats["var"]) + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + if self.global_mean is None or self.global_var is None: + return x + mean = self.global_mean.to(device=x.device, dtype=x.dtype) + var = self.global_var.to(device=x.device, dtype=x.dtype) + return (x - mean) / torch.sqrt(var) + + def denormalize(self, x: torch.Tensor) -> torch.Tensor: + if self.global_mean is None or self.global_var is None: + return x + mean = self.global_mean.to(device=x.device, dtype=x.dtype) + var = self.global_var.to(device=x.device, dtype=x.dtype) + return x * torch.sqrt(var) + mean + + +# Per-request state container for AR-loop continuity. All cross-step +# side-path state lives here, keyed by request_id, so concurrent requests +# stay fully isolated (the voxcpm2 _RequestState pattern). +@dataclass +class _RequestState: + request_id: str + # AR loopback: patch_encoder's previous decode output, fed to the LM + # as next-step inputs_embeds. Shape: [1, llm_hidden] = [1, 1536]. + # (patch_encoder collapses its internal _PATCH_ENCODER_OUT_DS_RATE + # positions into one LLM token via out_proj; see + # VAESemanticEncoder._project_embeddings.) _finish_decode writes a + # fresh tensor every call. + curr_embed_for_next: torch.Tensor | None = None + # Stop-signal cache: _finish_decode populates precomputed_stop_logits; + # compute_logits drains it. + precomputed_stop_logits: torch.Tensor | None = None + is_stopping: bool = False + prefill_completed: bool = False + # Per-request FM noise counter: draw #n of this request hashes to a + # deterministic Generator seed (see _run_dit_n_step_euler), so outputs + # are reproducible run-to-run and concurrent requests cannot perturb + # each other's noise streams (voxcpm2 _fill_deterministic_cfm_noise). + noise_step: int = 0 + # Per-request FM static workspace (lazy-allocated by + # _initialize_request_fm_state on first _finish_decode call). Sized for + # _MAX_AUDIO_PATCHES × (_HIDDEN_PATCH_SIZE + _LATENT_PATCH_SIZE) = 1024 × 5 + # = 5120 positions. + fm_sequence: torch.Tensor | None = None # [1, fm_capacity, _FM_HIDDEN] + fm_cfg_sequence: torch.Tensor | None = None # [1, fm_capacity, _FM_HIDDEN] + fm_null_g_cond: torch.Tensor | None = None # [1, _FM_HIDDEN] + fm_seq_len: int = 0 + fm_capacity: int = 0 + # Speaker x-vector after _xvec_proj. Zero-shot generation leaves this + # None → DiT falls back to fm_null_g_cond; voice-clone wiring is a + # follow-up. + g_cond: torch.Tensor | None = None + # AR-loop bookkeeping. + # patch_encoder_state holds conv_tail + per-layer KV caches for the + # patch_encoder's streaming decode (mirrors upstream state.patch_encoder_state). + # Lazy-allocated by _run_patch_encoder_loopback on first audio step. + patch_encoder_state: Any | None = None + # Review M2: streaming vocoder state (BigVGANStreamState — LSTM hidden + # + decoder sliding latent window). Gives each patch real left conv + # context; output lags input by decoder.stream_lookahead (2 frames = + # 3840 samples). Lazy-allocated by _run_vocoder_stream_step. + vocoder_stream_state: Any | None = None + + +# DiT transformer hyperparameters from rednote-hilab/dots.tts-soar config.json +# (the ``DiT`` block). Mirrors the shape of an upstream pydantic config: DiT's +# constructor calls ``transformer_config.to_dict()`` and reads ``.hidden_size`` +# / ``.num_layers``, so this @dataclass plus the ``to_dict`` helper satisfies +# that minimal interface without dragging in upstream's config system. +# +# Note: ``attn_dropout`` in the JSON spells differently from MultiHeadAttention's +# ``attn_drop`` kwarg — the value gets absorbed by ``**_kwargs`` and the MHA +# default 0.0 applies. Same numeric outcome for the soar checkpoint (both 0.0), +# so this is benign; flagging it because it would mask non-default values if +# upstream ever raised this field. +@dataclass +class _DiTConfig: + num_layers: int = 18 + num_heads: int = 16 + hidden_size: int = 1024 + ffn_hidden_size: int = 4096 + modulation: bool = True + qkv_bias: bool = False + qk_norm: bool = True + attn_dropout: float = 0.0 + dropout: float = 0.0 + norm_layer: str = "RMSNorm" + alibi_bias: bool = False + rotary_bias: bool = True + rotary_theta: float = 10000.0 + + def to_dict(self) -> dict[str, Any]: + return dataclasses.asdict(self) + + +def _build_dit_config(hf_config: Any) -> _DiTConfig: + # _DiTConfig declares exactly the checkpoint's 13 ``DiT`` keys, so an + # unknown key raises TypeError instead of being silently dropped. + return _DiTConfig(**_require_config_block(hf_config, "DiT")) + + +# Patch encoder hyperparameters from rednote-hilab/dots.tts-soar config.json +# (the top-level ``patch_size`` and the ``PatchEncoder`` sub-block). +# +# Two layers of config: +# * Outer (``_PatchEncoderConfig``): VAESemanticEncoder reads ``.patch_size`` +# and forwards ``.PatchEncoder`` to its inner SuperviseEncoder. +# * Inner (``_PatchEncoderInner``): the constructor uses BOTH attribute +# access (``config.PatchEncoder.hidden_size``) and dict-style +# ``config.get("hidden_size", 1024)``. Upstream's pydantic ConfigBase +# supports both natively; we add a tiny ``.get()`` helper to a dataclass. +# +# Upstream's ``_EncoderConfig`` also defines qk_norm / rotary_bias / +# rotary_theta / input_dim, but ``SuperviseEncoder.__init__`` only reads six +# keys (num_layers / num_heads / hidden_size / ffn_hidden_size / norm_layer / +# causal). The rest are swallowed and MultiHeadAttention falls back to its +# own defaults — matching upstream behaviour means *not* plumbing them through. +@dataclass +class _PatchEncoderInner: + num_layers: int = 24 + num_heads: int = 16 + hidden_size: int = 1024 + ffn_hidden_size: int = 4096 + norm_layer: str = "RMSNorm" + causal: bool = True + + def get(self, key: str, default: Any = None) -> Any: + return getattr(self, key, default) + + +@dataclass +class _PatchEncoderConfig: + patch_size: int = 4 + PatchEncoder: _PatchEncoderInner = dataclasses.field(default_factory=_PatchEncoderInner) + + +def _build_patch_encoder_config(hf_config: Any) -> _PatchEncoderConfig: + block = _require_config_block(hf_config, "PatchEncoder") + # SuperviseEncoder reads only the six declared fields; upstream's + # pydantic config swallows the other checkpoint keys, so filtering + # here matches upstream behaviour exactly. + inner_fields = {f.name for f in dataclasses.fields(_PatchEncoderInner)} + inner = _PatchEncoderInner(**{k: v for k, v in block.items() if k in inner_fields}) + return _PatchEncoderConfig( + patch_size=int(hf_config.patch_size), + PatchEncoder=inner, + ) + + +def _validate_architecture_constants(hf_config: Any, dit_config: _DiTConfig) -> None: + """The FM workspace helpers size buffers off module-level constants; + a checkpoint that disagrees would corrupt the DiT conditioning, so + refuse it up front (making these fully config-driven is deferred + until a released checkpoint actually differs).""" + checks = [ + ("latent_dim", getattr(hf_config, "latent_dim", None), _LATENT_DIM), + ("patch_size", getattr(hf_config, "patch_size", None), _LATENT_PATCH_SIZE), + ("DiT.hidden_size", dit_config.hidden_size, _FM_HIDDEN), + ] + mismatches = [f"{name}={got!r} (expected {want})" for name, got, want in checks if got != want] + if mismatches: + raise ValueError( + "dots.tts checkpoint config disagrees with this integration's " + "architecture constants: " + ", ".join(mismatches) + ) + + +class DotsTTSForConditionalGeneration(nn.Module): + """dots.tts AR talker: Qwen2.5 base LM + DiT / AudioVAE / CAM++ side path.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + + self.vllm_config = vllm_config + self.config = vllm_config.model_config.hf_config + + # vllm-omni capability flags — gpu_model_runner checks these to decide + # whether to call preprocess() / postprocess() / accept multimodal + # outputs. Without these, the runner bypasses our preprocess hook + # and _pending_requests stays empty (forward never dispatches the + # side-path). Voxcpm2 sets the same three. + self.have_multimodal_outputs = True + self.has_preprocess = True + self.has_postprocess = True + + self.model = Qwen2Model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + + # Side-path module hyperparameters come off the checkpoint's + # config.json blocks (fail loudly if a block is absent); the FM + # workspace constants are cross-checked against the same config. + dit_config = _build_dit_config(self.config) + _validate_architecture_constants(self.config, dit_config) + llm_hidden = self.config.hidden_size + fm_hidden = dit_config.hidden_size + latent_dim = int(self.config.latent_dim) + xvec_dim = int(getattr(self.config, "campplus_embedding_size", 512)) + + self._audio_vae = AudioVAE(_build_audio_vae_config(self.config)) + # Upstream serializes vocoder.safetensors with weight_norm folded into + # plain `weight` on the decoder (encoder kept weight_norm — it's not in + # the synthesis hot path). Match that layout before load_weights runs, + # so checkpoint keys align with our state_dict 1:1. + self._audio_vae.remove_weight_norm() + # AudioVAE / speaker encoder stay fp32 — upstream only casts + # `model.core` to bf16 (runtime.py:81), leaving vocoder + xvector + # extractor in fp32 because their kernels (conv1d / fbank) need + # higher precision. Without this _audio_vae would inherit talker's + # bf16 dtype via vLLM's auto-cast and dtype-mismatch under + # stream_step (input .float() vs bf16 weights). + + # DiT flow-matching head. Dimensions per upstream core.py:101-104: + # in_dim = config.DiT.hidden_size (DiT internal space) + # out_dim = config.latent_dim (AudioVAE input space) + # mode is "flow_matching" for soar/base; "meanflow" only for the mf + # checkpoint (handled by a separate factory in a later step). + self._head = DiT( + in_dim=fm_hidden, + out_dim=latent_dim, + transformer_config=dit_config, + mode="flow_matching", + ) + + # Patch encoder: closes the AR loop by mapping each DiT-emitted audio + # latent patch back into the LLM hidden space, so Qwen2 "hears" what it + # has already generated. Upstream core.py:75 calls this with + # in_dim=latent_dim and out_dim=llm.config.hidden_size. + self._patch_encoder = VAESemanticEncoder( + in_dim=latent_dim, + out_dim=llm_hidden, + config=_build_patch_encoder_config(self.config), + ) + + # Five thin projectors (upstream core.py:82-113). Soar dimensions: + # llm_hidden=1536, fm_hidden=1024, latent_dim=128, xvec_dim=512. + # These are kept as direct attributes (no vendor file) — they are flat + # nn.Linear / nn.Sequential, and live at the top of the checkpoint + # namespace under their own short prefixes (see load_weights). + self._hidden_proj = nn.Linear(llm_hidden, fm_hidden) + self._latent_proj = nn.Linear(latent_dim, fm_hidden) + self._coordinate_proj = nn.Linear(latent_dim, fm_hidden) + self._xvec_proj = nn.Sequential( + nn.Linear(xvec_dim, fm_hidden), + nn.LayerNorm(fm_hidden), + ) + # eos_proj is the stop predictor — feeds LLM hidden through a 2-layer + # MLP and produces a 2-way logit (continue / stop). + self._eos_proj = nn.Sequential( + nn.Linear(llm_hidden, llm_hidden), + nn.SiLU(), + nn.Linear(llm_hidden, 2), + ) + + # CAM++ speaker encoder (3D-Speaker), produces the x-vector from + # reference audio for voice cloning. sample_rate matches AudioVAE's + # input space (48 kHz); the wrapper resamples to 16k internally for + # CAM++. Weights are frozen (upstream does `requires_grad = False`). + self._speaker_encoder = SpeakerXVectorFeatures( + sample_rate=self._audio_vae.sample_rate, + campplus_embedding_size=xvec_dim, + max_audio_seconds=float(getattr(self.config, "xvec_max_audio_seconds", 10.0)), + ) + + # Pin AudioVAE + speaker encoder to fp32 (mirror upstream + # runtime.py:81 where only `model.core` is cast to bf16). vLLM's + # auto-cast would otherwise lower these to bf16, breaking + # stream_step (input .float() vs bf16 weights) and the fbank + # pipeline. + self._audio_vae.float() + self._speaker_encoder.float() + + # Latent stats — independent of safetensors (torch.load on + # latent_stats.pt sitting next to the checkpoint files). Used by + # _finish_decode to bridge between DiT's normalized latent space and + # the raw latent space expected by patch_encoder + AudioVAE. + latent_stats_path = self._resolve_latent_stats_path(vllm_config.model_config.model) + if latent_stats_path is not None: + self._io_helper = _IOHelper(latent_stats_path) + else: + logger.warning( + "latent_stats.pt not resolvable for %s; running with " + "identity normalize/denormalize. Wav will be ~10x quiet " + "and distorted (no sqrt(var) denormalize).", + vllm_config.model_config.model, + ) + self._io_helper = _IOHelper() + + # (β) diagnostic: env-gated per-patch rms/max trace across the full + # side path (LLM hidden → DiT → denormalize → patch_encoder → wav → + # eos). Zero overhead when DOTS_TTS_BETA_TRACE is unset (single + # bool check per step); kept as the model's debugging instrument + # (cf. voxcpm2's enable_profiling runtime knob). + self._beta_trace = bool(os.environ.get("DOTS_TTS_BETA_TRACE")) + + # Per-request state plumbing. + # _active_states: per-request state keyed by request_id (created in + # preprocess prefill, evicted by _flush_deferred_cleanup). + # _pending_requests: per-step list of (req_id, is_prefill, embeds, + # span_len) populated by preprocess(), consumed by forward(). + # _deferred_cleanup_ids: requests vLLM has marked finished but whose + # side-path audio still needs draining in the current forward step; + # evicted from _active_states at the end of forward(). + self._active_states: dict[str, _RequestState] = {} + self._pending_requests: list[tuple[str, bool, torch.Tensor, int]] = [] + self._deferred_cleanup_ids: set[str] = set() + + # Audio side-channel: + # _audio_queue: _finish_decode pushes (req_id, audio_tensor) tuples + # here; make_omni_output drains them into multimodal_outputs each + # step. + # _sample_rate: cached from AudioVAE config (= 48000 for soar). + self._audio_queue: list[tuple[str, torch.Tensor]] = [] + self._sample_rate = self._audio_vae.sample_rate + + # Stop-signal side-channel: + # _results_queue: _finish_decode pushes (req_id, stop_logits) tuples + # here; compute_logits drains them into the [bsz, vocab_size] + # logits tensor each step. + self._results_queue: list[tuple[str, torch.Tensor | None]] = [] + + logger.info( + "DotsTTS talker built (base_lm=Qwen2[28L,12H,1536d], " + "audio_vae=AudioVAE, dit=DiT[18L,16H,1024d], " + "patch_encoder=VAESemanticEncoder[24L,16H,1024d], " + "5 projectors, speaker_encoder=CAM++[7.2M], " + "io_helper=%s, model=%s, dit_steps=%d); " + "_finish_decode: DiT → denormalize → patch_encoder AR + " + "streaming AudioVAE decode → wav + eos_proj → stop signal.", + "loaded" if self._io_helper.global_mean is not None else "identity", + vllm_config.model_config.model, + _DIT_NUM_STEPS, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor | IntermediateTensors: + """Drive the base LM, then the per-request audio side path. + + One vLLM step: base LM forward under PagedAttention, then slice + the scaffold hidden states per request (in _pending_requests + order) and run _finish_decode on each — DiT sampling, streaming + AudioVAE decode, AR loopback, and stop signal all happen there. + vLLM only ever sees the returned hidden states; audio leaves + through the _audio_queue side channel. + """ + output = self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + if isinstance(output, IntermediateTensors): + return output + if isinstance(output, tuple): + output = output[0] + + # Per-request side-path: slice scaffold_hidden by token span_len + # in the order preprocess() appended to _pending_requests. + token_offset = 0 + for req_id, is_prefill, _embeds, span_len in self._pending_requests: + req_hidden = output[token_offset : token_offset + span_len] + token_offset += span_len + self._finish_decode(req_id, req_hidden, is_prefill) + + # End-of-step cleanup + self._pending_requests.clear() + self._flush_deferred_cleanup() + + return output + + # ── vllm-omni protocol methods ── + + @staticmethod + def _resolve_latent_stats_path(model_arg: str) -> str | None: + """Resolve latent_stats.pt to a local filesystem path. + + ``vllm_config.model_config.model`` is typically an HF repo ID like + ``"rednote-hilab/dots.tts-soar"`` — joining it with ``"latent_ + stats.pt"`` gives a non-existent path. Try the local-dir form + first (covers ``--model /local/path``), then fall back to HF + cache lookup with ``local_files_only=True`` so we never trigger + a download from this side (the file was fetched alongside the + safetensors at load time). + """ + local = os.path.join(model_arg, "latent_stats.pt") + if os.path.exists(local): + return local + try: + from huggingface_hub import hf_hub_download + + return hf_hub_download( + repo_id=model_arg, + filename="latent_stats.pt", + local_files_only=True, + ) + except Exception: + return None + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + """Map token ids to embeddings via the base LM's embed_tokens. + + Required by vLLM's ``is_text_generation_model`` Protocol check + (interfaces_base._check_vllm_model_embed_input_ids). Without this + the model fails ``--runner generate`` validation at engine startup. + """ + return self.model.embed_tokens(input_ids) + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None = None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Per-step embedding + per-request state registration. + + Prefill: embed the wrapped prompt + ``[文本][文本对应语音]`` via the LM's + embed_tokens. Caller (serving layer or test harness) is + responsible for the wrap — talker just embeds whatever it gets. + + Decode (span_len == 1): use ``state.curr_embed_for_next`` from + the previous step's AR loopback (DiT → patch_encoder output); + zero fallback only if the loopback hasn't produced one yet. + + Voice clone is intentionally not wired here — zero-shot only. + + Returns ``(input_ids unchanged, embeds [span_len, hidden_size], + {})``. ``forward()`` consumes ``self._pending_requests`` to + slice scaffold_hidden per request. + """ + # voxcpm2 protocol compatibility: flatten additional_information. + additional = info_dict.get("additional_information") + if isinstance(additional, dict): + merged = {k: v for k, v in info_dict.items() if k != "additional_information"} + for k, v in additional.items(): + merged.setdefault(k, v) + info_dict = merged + + span_len = int(input_ids.shape[0]) + dev = input_ids.device + req_id = info_dict.get("request_id", "default") + state = self._active_states.get(req_id) + # Prefill is detected by state lifecycle (not span_len), so decode + # can use span_len=1 with a multi-position curr_embed_for_next. + is_prefill = state is None or not state.prefill_completed + + if is_prefill: + if state is None: + state = _RequestState(request_id=req_id) + self._active_states[req_id] = state + # Reset per-step state on every prefill — previous run is stale. + state.curr_embed_for_next = None + state.precomputed_stop_logits = None + state.is_stopping = False + state.prefill_completed = False + state.noise_step = 0 + # Reset AR-loop state on every prefill. + state.patch_encoder_state = None + state.vocoder_stream_state = None + embeds = self.model.embed_tokens(input_ids) + else: + curr = state.curr_embed_for_next + if curr is not None: + embeds = curr.to(dev) # [1, llm_hidden] + else: + # First decode step before AR loop closed — zero fallback. + embed_dtype = self.model.embed_tokens.weight.dtype + embeds = torch.zeros( + 1, + self.config.hidden_size, + device=dev, + dtype=embed_dtype, + ) + + self._pending_requests.append((req_id, is_prefill, embeds, span_len)) + return input_ids, embeds, {} + + def postprocess(self, *args: Any, **kwargs: Any) -> dict: + return {} + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + """Mark finished requests for deferred eviction. + + vLLM scheduler calls this BEFORE forward() to notify that certain + requests have completed. We can't drop their _RequestState yet — + the current forward step still needs to drain side-path audio + chunks. Real eviction happens at the end of forward() via + _flush_deferred_cleanup. + """ + for req_id in finished_req_ids: + if req_id in self._active_states: + self._deferred_cleanup_ids.add(req_id) + + def _flush_deferred_cleanup(self) -> None: + """Evict states marked by on_requests_finished. Called at the + tail of forward().""" + for req_id in self._deferred_cleanup_ids: + self._active_states.pop(req_id, None) + self._deferred_cleanup_ids.clear() + + def _finish_decode( + self, + req_id: str, + req_hidden: torch.Tensor, + is_prefill: bool, + ) -> None: + """One request's per-step side path: DiT → AR loopback → wav → stop. + + One vLLM step = one audio patch. patch_encoder.decode_patch + collapses its internal _PATCH_ENCODER_OUT_DS_RATE positions into + a single LLM input embedding (see VAESemanticEncoder._project_ + embeddings), so vLLM's span_len=1 decode matches naturally — no + slot bookkeeping needed. + + Per-call flow (prefill or decode): + 1. _hidden_proj(last_hidden) → fm_sequence (+1 position) + 2. DiT N-step Euler → audio_patch [1, 4, 128] (normalized) + 3. _latent_proj(audio_patch) → fm_sequence (+4 positions, normalized) + 4. io_helper.denormalize(audio_patch) → audio_patch_raw + 5. patch_encoder(audio_patch_raw) → next-step embed [1, 1536] + 6. AudioVAE stream_step(audio_patch_raw) → wav chunk pushed to + _audio_queue (prefill's patch #0 included; output lags 2 + frames behind input) + 7. (decode only) eos_proj(last_hidden) → stop_logits pushed to + _results_queue; threshold > 0.8 flips state.is_stopping and + stream_flush drains the decoder's lookahead tail. + """ + state = self._active_states.get(req_id) + if state is None: + return + + if state.fm_sequence is None: + self._initialize_request_fm_state( + state, + device=req_hidden.device, + dtype=req_hidden.dtype, + ) + + # This step's append would overflow the per-request FM buffer. + # Stop gracefully (same path as a model-decided stop below) instead + # of raising past _append_hidden_chunk/_append_history_chunk — + # those raises are engine-fatal, not request-fatal, and would take + # every other in-flight and future request down with this one. + if state.fm_seq_len + _HIDDEN_PATCH_SIZE + _LATENT_PATCH_SIZE > state.fm_capacity: + state.is_stopping = True + stop_logits = torch.tensor([[0.0, 1.0]], device=req_hidden.device, dtype=req_hidden.dtype) + state.precomputed_stop_logits = stop_logits + tail = self._audio_vae.stream_flush(state.vocoder_stream_state) + if tail.size(-1) > 0: + self._audio_queue.append((req_id, tail.reshape(-1))) + self._results_queue.append((req_id, stop_logits)) + return + + # 1. Append last LLM hidden to fm_sequence (+1 position). + last_hidden = req_hidden[-_HIDDEN_PATCH_SIZE:].unsqueeze(0) + + self._append_hidden_chunk(state, last_hidden) + + # 2. DiT N-step Euler → audio latent patch [1, 4, 128] (normalized). + audio_patch = self._run_dit_n_step_euler(state) + + # 3. Append latent to fm_sequence (+4 positions). Stays in + # normalized space — DiT's KV cache lives in that space. + self._append_history_chunk(state, audio_patch) + + # 4. Denormalize for downstream consumers (patch_encoder, AudioVAE). + audio_patch_raw = self._io_helper.denormalize(audio_patch) + + # 5. patch_encoder AR loopback → next-step inputs_embeds [1, 1536]. + next_embeds = self._run_patch_encoder_loopback(state, audio_patch_raw) + state.curr_embed_for_next = next_embeds.squeeze(0).detach() + + # 6. Streaming AudioVAE decode → wav chunk. Runs on prefill too: + # upstream zero-shot emits every DiT patch including the first + # (_generate_latents_stream drops one patch only under prompt + # prefill / voice clone), and our prefill patch #0 corresponds + # to upstream's first decode-loop patch. + wav = self._run_vocoder_stream_step(state, audio_patch_raw) + if wav.size(-1) > 0: + self._audio_queue.append((req_id, wav.reshape(-1))) + + if is_prefill: + if self._beta_trace: + self._beta_trace_log( + state, + is_prefill=True, + last_hidden=last_hidden, + audio_patch=audio_patch, + audio_patch_raw=audio_patch_raw, + next_embeds=next_embeds, + wav=wav if wav.size(-1) > 0 else None, + prob_stop=None, + ) + state.prefill_completed = True + # Keep len(_results_queue) == bsz: compute_logits pairs queue + # entries with batch rows positionally, so every request in the + # batch must push exactly one entry per step (voxcpm2_talker.py + # prefill placeholder pattern). None → forced continue via the + # stop_logits-is-None branch in compute_logits. + self._results_queue.append((req_id, None)) + return + + # 7. eos_proj → stop probability. Mirrors upstream + # model.py:1681 — softmax([continue, stop]) on detached LLM + # hidden, threshold the stop slot at 0.8 (upstream default + # eos_threshold). When stop fires for this step, compute_logits + # drains the high-stop logit on the next dispatch and vLLM + # terminates the request. + eos_logits = self._eos_proj(last_hidden.detach()).softmax(dim=-1) + stop_logits = eos_logits.squeeze(0) + state.precomputed_stop_logits = stop_logits + if eos_logits[0, -1, 1].item() > 0.8: + state.is_stopping = True + # Upstream flushes after its decode loop (model.py:1898); our + # request ends when vLLM sees this step's stop token, so drain + # the decoder's 2-frame lookahead tail in the same step. + # Requests cut by max_tokens/abort never reach this path and + # lose the 40 ms tail — no engine hook fires late enough to + # flush after the final step. + tail = self._audio_vae.stream_flush(state.vocoder_stream_state) + if tail.size(-1) > 0: + self._audio_queue.append((req_id, tail.reshape(-1))) + self._results_queue.append((req_id, stop_logits)) + + if self._beta_trace: + self._beta_trace_log( + state, + is_prefill=False, + last_hidden=last_hidden, + audio_patch=audio_patch, + audio_patch_raw=audio_patch_raw, + next_embeds=next_embeds, + wav=wav if wav.size(-1) > 0 else None, + prob_stop=eos_logits[0, -1, 1].item(), + ) + + def _beta_trace_log( + self, + state: _RequestState, + *, + is_prefill: bool, + last_hidden: torch.Tensor, + audio_patch: torch.Tensor, + audio_patch_raw: torch.Tensor, + next_embeds: torch.Tensor, + wav: torch.Tensor | None, + prob_stop: float | None, + ) -> None: + """(β) diagnostic — patch label matches output.wav indexing in + decode rows; prefill row labels itself ``pre``.""" + + def rmsmax(t: torch.Tensor) -> tuple[float, float]: + f = t.detach().float() + return f.pow(2).mean().sqrt().item(), f.abs().max().item() + + patch_label = "pre" if is_prefill else f"{state.fm_seq_len // 5 - 2:2d}" + req_tail = state.request_id[-8:] + hid_rms, hid_max = rmsmax(last_hidden) + dit_rms, dit_max = rmsmax(audio_patch) + raw_rms, raw_max = rmsmax(audio_patch_raw) + nxt_rms, nxt_max = rmsmax(next_embeds) + if wav is not None: + w_rms, w_max = rmsmax(wav) + wav_part = f"wav:rms={w_rms:.3f} max={w_max:.3f}" + else: + wav_part = "wav:rms= - max= - " + stop_part = f"stop={prob_stop:.3f}" if prob_stop is not None else "stop= - " + logger.info( + f"[β req={req_tail} patch={patch_label}] " + f"hid:rms={hid_rms:.3f} max={hid_max:.3f} | " + f"dit:rms={dit_rms:.3f} max={dit_max:.3f} | " + f"raw:rms={raw_rms:.3f} max={raw_max:.3f} | " + f"next:rms={nxt_rms:.3f} max={nxt_max:.3f} | " + f"{wav_part} | {stop_part}" + ) + + # ── FM helpers (per-request workspace + DiT N-step Euler) ── + + def _initialize_request_fm_state( + self, + state: _RequestState, + *, + device: torch.device, + dtype: torch.dtype, + ) -> None: + """Lazy-allocate per-request FM workspace. + + Mirrors upstream _allocate_generate_state (model.py:413) but per- + request (no workspace sharing) so concurrent requests stay + isolated. + """ + fm_capacity = _MAX_AUDIO_PATCHES * (_HIDDEN_PATCH_SIZE + _LATENT_PATCH_SIZE) + state.fm_capacity = fm_capacity + state.fm_sequence = torch.zeros( + (1, fm_capacity, _FM_HIDDEN), + device=device, + dtype=dtype, + ) + state.fm_cfg_sequence = torch.zeros( + (1, fm_capacity, _FM_HIDDEN), + device=device, + dtype=dtype, + ) + state.fm_null_g_cond = torch.zeros( + (1, _FM_HIDDEN), + device=device, + dtype=dtype, + ) + state.fm_seq_len = 0 + + def _append_hidden_chunk( + self, + state: _RequestState, + hidden_chunk: torch.Tensor, + ) -> None: + """Project last _HIDDEN_PATCH_SIZE LLM hiddens via _hidden_proj into fm_sequence. + + Mirrors upstream _append_hidden_chunk (model.py:1308): cfg buffer + gets the null-projected counterpart (``_hidden_proj`` of zero) so + CFG dropout sees no past hidden, only past audio history. + """ + last_hidden = hidden_chunk[:, -_HIDDEN_PATCH_SIZE:, :] + projected = self._hidden_proj(last_hidden) + null_projected = self._hidden_proj(torch.zeros_like(last_hidden)) + end = state.fm_seq_len + projected.size(1) + if end > state.fm_capacity: + raise RuntimeError(f"FM buffer overflow on hidden append: end={end} cap={state.fm_capacity}") + state.fm_sequence[:, state.fm_seq_len : end].copy_(projected.to(state.fm_sequence.dtype)) + state.fm_cfg_sequence[:, state.fm_seq_len : end].copy_(null_projected.to(state.fm_cfg_sequence.dtype)) + state.fm_seq_len = end + + def _append_history_chunk( + self, + state: _RequestState, + latent_chunk: torch.Tensor, + ) -> None: + """Project audio latent via _latent_proj into fm_sequence (history append). + + Mirrors upstream _append_history_chunk (model.py:1325): cfg + buffer gets the same value (latent history is shared between + conditional and unconditional CFG paths — only the hidden chunk + gets nulled). + """ + history_latent = self._latent_proj(latent_chunk) + end = state.fm_seq_len + history_latent.size(1) + if end > state.fm_capacity: + raise RuntimeError(f"FM buffer overflow on latent append: end={end} cap={state.fm_capacity}") + state.fm_sequence[:, state.fm_seq_len : end].copy_(history_latent.to(state.fm_sequence.dtype)) + state.fm_cfg_sequence[:, state.fm_seq_len : end].copy_(history_latent.to(state.fm_cfg_sequence.dtype)) + state.fm_seq_len = end + + def _build_fm_attn_mask( + self, + state: _RequestState, + total_len: int, + device: torch.device, + ) -> torch.Tensor: + """Build DiT attention mask for one decode step. + + Mirrors upstream _build_fm_attn_mask (model.py:1203). Layout: + * [0, fm_seq_len) — committed history. Causal over the past + up to a non-causal tail block of size _HIDDEN_PATCH_SIZE that + sees itself and the new noise patch (the just-appended hidden). + * [latent_start, total_len) — new noise patch, bidirectional + with the committed history and itself. + """ + attn_mask = torch.zeros((1, total_len, total_len), device=device, dtype=torch.bool) + latent_start = total_len - _LATENT_PATCH_SIZE + block_start = state.fm_seq_len - _HIDDEN_PATCH_SIZE + if block_start > 0: + causal_mask = torch.ones((block_start, block_start), device=device, dtype=torch.bool).triu(1).logical_not() + attn_mask[:, :block_start, :block_start] = causal_mask + attn_mask[:, block_start : state.fm_seq_len, : state.fm_seq_len] = True + attn_mask[:, block_start : state.fm_seq_len, latent_start:] = True + attn_mask[:, latent_start:, : state.fm_seq_len] = True + attn_mask[:, latent_start:, latent_start:] = True + if latent_start > state.fm_seq_len: + padding_indices = torch.arange(state.fm_seq_len, latent_start, device=device) + attn_mask[:, padding_indices, padding_indices] = True + return attn_mask + + def _build_fm_pos_ids( + self, + state: _RequestState, + total_len: int, + device: torch.device, + ) -> torch.Tensor: + """Build DiT position ids for one decode step. + + Mirrors upstream _build_fm_pos_ids (model.py:1236). Position ids + stay fp32 — DiT's RotaryEmbedding does its own dtype handling. + """ + pos_ids = torch.zeros((1, total_len), device=device, dtype=torch.float32) + latent_start = total_len - _LATENT_PATCH_SIZE + if state.fm_seq_len > 0: + pos_ids[:, : state.fm_seq_len] = torch.arange( + state.fm_seq_len, + device=device, + dtype=pos_ids.dtype, + ) + pos_ids[:, latent_start:] = torch.arange( + state.fm_seq_len, + state.fm_seq_len + _LATENT_PATCH_SIZE, + device=device, + dtype=pos_ids.dtype, + ) + return pos_ids + + def _run_dit_n_step_euler( + self, + state: _RequestState, + *, + num_steps: int = _DIT_NUM_STEPS, + guidance_scale: float = _DIT_GUIDANCE_SCALE, + ) -> torch.Tensor: + """DiT N-step Euler integration → audio latent patch [1, 4, 128]. + + Manual Euler loop (upstream uses torchdyn.odeint; we manualize to + avoid the dep and keep control flow explicit). Each step: + 1. z_proj = _coordinate_proj(z), placed into the last + _LATENT_PATCH_SIZE positions of input_sequence / cfg_sequence. + 2. Batch [conditional, unconditional] → DiT velocity field. + 3. CFG blend: v_c + guidance_scale * (v_c - v_u). + 4. Euler: z += v * dt. + + Mirrors upstream fm_solver_step (core.py:295) + _flow_matching_step_fm + (core.py:463), reduced to a fixed-step Euler integrator. + """ + assert state.fm_sequence is not None + assert state.fm_cfg_sequence is not None + assert state.fm_null_g_cond is not None + device = state.fm_sequence.device + dtype = state.fm_sequence.dtype + total_len = state.fm_seq_len + _LATENT_PATCH_SIZE + + # Workspace: committed history + 4 noise slots (overwritten per step). + input_sequence = torch.zeros((1, total_len, _FM_HIDDEN), device=device, dtype=dtype) + input_sequence[:, : state.fm_seq_len] = state.fm_sequence[:, : state.fm_seq_len] + cfg_sequence = torch.zeros((1, total_len, _FM_HIDDEN), device=device, dtype=dtype) + cfg_sequence[:, : state.fm_seq_len] = state.fm_cfg_sequence[:, : state.fm_seq_len] + + attn_mask = self._build_fm_attn_mask(state, total_len, device) + pos_ids = self._build_fm_pos_ids(state, total_len, device) + + g_cond = state.g_cond if state.g_cond is not None else state.fm_null_g_cond + g_cond = g_cond.to(device=device, dtype=dtype) + g_cond_batched = torch.cat([g_cond, torch.zeros_like(g_cond)], dim=0) + + latent_start = total_len - _LATENT_PATCH_SIZE + # fp32 ODE integration (same fix as Ming's CFM sampler, PR #4341): + # bf16's 7-bit mantissa loses the small per-step increments in the + # z += v*dt accumulation (hidden-state cos drift vs upstream past + # step ~7) and quantizes the linspace timesteps. Integration + # state stays fp32; DiT matmuls still run bf16 under the autocast + # block below, and the result is cast back to the sequence dtype + # on return. + # Deterministic per-request noise (voxcpm2 _fill_deterministic_cfm_ + # noise pattern): hash seed:request_key:draw# into a private + # Generator instead of the global CUDA RNG, so outputs reproduce + # run-to-run and concurrent requests cannot perturb each other's + # noise streams. request_key strips the engine's per-run "_" + # uuid suffix so replay across runs keys on the stable batch index. + request_key = state.request_id.split("_", 1)[0] + if not request_key.isdigit(): + request_key = state.request_id + noise_key = f"{_DIT_NOISE_SEED}:{request_key}:{state.noise_step}".encode() + digest = hashlib.blake2b(noise_key, digest_size=8).digest() + gen = torch.Generator(device=device) + gen.manual_seed(int.from_bytes(digest, "little") & 0x7FFF_FFFF_FFFF_FFFF) + z = torch.empty((1, _LATENT_PATCH_SIZE, _LATENT_DIM), device=device, dtype=torch.float32) + z.normal_(generator=gen) + state.noise_step += 1 + dt = 1.0 / num_steps + times = torch.linspace(0.0, 1.0, num_steps + 1, device=device, dtype=torch.float32) + + # Match upstream's autocast wrapper (model.py:280). DiT's + # TimestepEmbedder internally forces fp32 (`.float()` on freqs); + # without autocast the fp32 output hits a bf16 Linear → dtype + # mismatch. autocast bridges this for matmuls inside DiT. + use_amp = device.type == "cuda" and dtype in (torch.float16, torch.bfloat16) + with torch.autocast( + device_type=device.type if device.type == "cuda" else "cuda", + dtype=dtype if use_amp else torch.float32, + enabled=use_amp, + ): + for step in range(num_steps): + t = times[step].reshape(1) + z_proj = self._coordinate_proj(z) + z_c = input_sequence.clone() + z_c[:, latent_start:] = z_proj + z_u = cfg_sequence.clone() + z_u[:, latent_start:] = z_proj + z_batched = torch.cat([z_c, z_u], dim=0) + t_batched = t.repeat(2) + vt = self._head( + x=z_batched, + timesteps=t_batched, + attn_mask=attn_mask, + pos_ids=pos_ids, + g_cond=g_cond_batched, + ) + # Upcast DiT output before CFG blend + Euler update so the + # accumulation arithmetic stays fp32 (autocast only affects + # matmuls; these pointwise ops keep their input dtype). + vt = vt[:, latent_start:].float() + vt_c, vt_u = vt[0:1], vt[1:2] + velocity = vt_c + guidance_scale * (vt_c - vt_u) + z = z + velocity * dt + return z.to(dtype) + + def _run_patch_encoder_loopback( + self, + state: _RequestState, + audio_patch_raw: torch.Tensor, + ) -> torch.Tensor: + """Decode one audio latent patch → 1 LLM input embed. + + Mirrors upstream _consume_audio_patch (model.py:1557) — patch_encoder + side only. Lazy-allocates state.patch_encoder_state on first call. + Returns ``llm_embedding`` of shape ``[1, 1, llm_hidden]`` — the + _PATCH_ENCODER_OUT_DS_RATE encoder positions fold into the feature + axis before ``out_proj``, not into separate output positions. + + ``audio_patch_raw`` must already be denormalized (raw VAE latent + space). Caller (_finish_decode) runs io_helper.denormalize on + DiT's normalized output before passing in. + """ + if state.patch_encoder_state is None: + state.patch_encoder_state = self._patch_encoder.init_decode_state( + max_audio_patch_count=_MAX_AUDIO_PATCHES, + batch_size=1, + device=audio_patch_raw.device, + dtype=audio_patch_raw.dtype, + ) + pe_state = state.patch_encoder_state + positions = ( + torch.arange( + _PATCH_ENCODER_OUT_DS_RATE, + device=audio_patch_raw.device, + dtype=torch.long, + ) + + pe_state.seq_len + ) + llm_embedding, conv_tail = self._patch_encoder.decode_patch( + audio_patch_raw, + pe_state.conv_tail, + pe_state.layer_caches, + positions, + ) + pe_state.conv_tail.copy_(conv_tail) + pe_state.seq_len += _PATCH_ENCODER_OUT_DS_RATE + return llm_embedding + + def _run_vocoder_stream_step( + self, + state: _RequestState, + audio_patch_raw: torch.Tensor, + ) -> torch.Tensor: + """Streaming AudioVAE decode of one latent patch (review M2). + + Mirrors upstream generate_audio_stream (model.py:1881-1900): a + per-request BigVGANStreamState carries the decoder's LSTM hidden + and sliding latent window across steps, so each patch is decoded + with real left context instead of inference_from_latents' + zero-padded isolated boundary. Output lags input by + decoder.stream_lookahead (2 frames): the first call emits 2 frames + (3840 samples), steady state 4 (7680); the stop path in + _finish_decode drains the tail via stream_flush. + + audio_patch_raw: [1, T=4, D=128] denormalized latents. + Returns wav chunk [1, 1, n_samples] (n_samples may be 0). + """ + if state.vocoder_stream_state is None: + state.vocoder_stream_state = self._audio_vae.init_stream_state( + batch_size=1, + chunk_size=_LATENT_PATCH_SIZE, + ) + # stream_step expects [B, latent_dim, T]; the vocoder is pinned + # fp32 (__init__), so cast the (possibly bf16) latents up. + return self._audio_vae.stream_step( + audio_patch_raw.transpose(1, 2).float(), + state.vocoder_stream_state, + ) + + def make_omni_output( + self, + model_outputs: torch.Tensor | OmniOutput, + **kwargs: Any, + ) -> OmniOutput: + """Drain self._audio_queue into the OmniOutput's multimodal_outputs. + + Output contract is **delta** (SKILL I1): each call emits only + audio produced since the previous step; the engine's + _consolidate_multimodal_tensors concatenates across steps for + offline consumers. + + Same-request multi-chunk: chunks for the same req_id within one + step are cat'd together. Multi-request: one entry per req_id + in multimodal_outputs["model_outputs"]. + + Note: ``sr`` is not a declared OmniPayload key but the engine / + serving consumer reads it (voxcpm2 convention). + """ + if isinstance(model_outputs, OmniOutput): + return model_outputs + + mm: dict[str, Any] = {} + if self._audio_queue: + audio_by_req: dict[str, torch.Tensor] = {} + for req_id, audio in self._audio_queue: + if audio is None: + continue + if req_id in audio_by_req: + audio_by_req[req_id] = torch.cat( + [audio_by_req[req_id].reshape(-1), audio.reshape(-1)], + dim=0, + ) + else: + audio_by_req[req_id] = audio + if audio_by_req: + sr = torch.tensor(self._sample_rate, dtype=torch.int32) + ready_req_ids = list(audio_by_req) + chunks = [audio_by_req[req_id].reshape(-1) for req_id in ready_req_ids] + mm["model_outputs"] = chunks + mm["sr"] = [sr for _ in ready_req_ids] + # sparse_audio: ["1"] flags GPUARModelRunner to skip the default + # payload["hidden"] = scaffold_hidden injection (gpu_ar_model_ + # runner.py:1114). Otherwise prefill scaffold (24 hidden × 1536 + # = 29184 numbers observed) leaks into mm["audio"] via the + # output_processor "hidden" → target_key="audio" rename + # (output_processor.py:84), prepending ~0.61 s of noise. + mm["meta"] = {"req_id": ready_req_ids, "sparse_audio": ["1"]} + self._audio_queue.clear() + else: + # Empty-audio step (prefill) still needs the sparse marker so the + # engine doesn't bleed scaffold_hidden into mm["audio"]. + mm["model_outputs"] = [] + mm["sr"] = [] + mm["meta"] = {"req_id": [], "sparse_audio": ["1"]} + + if self._beta_trace: + payload = mm.get("model_outputs") + if payload: + for i, t in enumerate(payload): + logger.info( + f"[β-emit] mm[model_outputs][{i}]: shape={tuple(t.shape)} " + f"dtype={t.dtype} samples={t.numel()} " + f"rms={t.float().pow(2).mean().sqrt().item():.4f} " + f"max={t.abs().max().item():.4f}" + ) + else: + logger.info("[β-emit] mm empty / no model_outputs") + mo_kind = type(model_outputs).__name__ + mo_info = f"shape={tuple(model_outputs.shape)}" if hasattr(model_outputs, "shape") else "no_shape" + logger.info(f"[β-emit] text_hidden_states arg: {mo_kind} {mo_info}") + + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=mm) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: Any = None, + ) -> torch.Tensor | None: + """Per-step stop signal logits (vLLM sampling protocol). + + Encodes continue/stop into ``logits[i, 0]`` / ``logits[i, 1]`` + (voxcpm2 convention — rest of vocab stays at -inf so the sampler + picks between the two slots). Source of stop logits is + ``self._results_queue``, populated by forward()'s _finish_decode. + When the queue is empty (preprocess-only steps), defaults to + all-continue. + """ + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + return None + + bsz = hidden_states.shape[0] + logits = torch.full( + (bsz, self.config.vocab_size), + float("-inf"), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + if self._results_queue: + for i, (req_id, stop_logits) in enumerate(self._results_queue): + if i >= bsz: + break + state = self._active_states.get(req_id) + if stop_logits is not None: + if state is not None and state.is_stopping: + logits[i, 0] = 0.0 + logits[i, 1] = 1.0 + else: + # is_stopping=False means prob_stop <= 0.8 — force + # continue. Feeding the raw softmax (continue, stop) + # pair as logits would let the greedy sampler compare + # them directly, silently lowering the effective stop + # threshold to 0.5 (early stop, swallowed endings). + logits[i, 0] = 1.0 + if state is not None: + state.precomputed_stop_logits = None + else: + # No stop signal pushed this step (prefill placeholder). + # Default to continue — stop is driven by eos_proj. + logits[i, 0] = 1.0 # continue + self._results_queue.clear() + else: + logits[:, 0] = 1.0 # all continue + + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load all dots.tts checkpoint weights — step 6b. + + Routing branches in a single pass over the input iterator (it can only + be consumed once): + + * **AudioVAE**: keys from ``vocoder.safetensors`` (no parent prefix). + Matched via set membership against ``self._audio_vae.state_dict()``; + acts as the final catch-all so it must run last. + * **DiT**: keys from ``model.safetensors`` under ``velocity_field_predictor.``. + * **patch_encoder**: keys from ``model.safetensors`` under ``patch_encoder.``. + * **Five projectors** (``hidden_proj`` / ``latent_proj`` / + ``coordinate_proj`` / ``xvec_proj`` / ``eos_proj``): each lives at + its own top-level prefix; dispatched through a shared table. + * **Qwen2 base LM**: keys under ``llm.model.``, handed off to vLLM's + Qwen2Model fuser (q/k/v_proj -> qkv_proj, gate/up_proj -> + gate_up_proj). ``llm.lm_head.*`` is silently dropped (soar uses + ``tie_word_embeddings=True`` and our ``self.model`` is a Qwen2Model + with no lm_head module). + * **CAM++ speaker encoder**: keys from ``speaker_encoder.safetensors`` + (no parent prefix, same convention as AudioVAE). Routed by set + membership against ``self._speaker_encoder.state_dict()`` (keys are + ``model.*`` / ``resample.kernel``, no distinguishing top-level + prefix). **Placed before the AudioVAE fallback** so its keys don't + fall through to the catch-all. + """ + vae_state_keys = set(self._audio_vae.state_dict().keys()) + dit_state_keys = set(self._head.state_dict().keys()) + patch_state_keys = set(self._patch_encoder.state_dict().keys()) + speaker_state_keys = set(self._speaker_encoder.state_dict().keys()) + + # (top-level prefix in checkpoint, local attr name, local nn.Module) + projector_specs: list[tuple[str, str, nn.Module]] = [ + ("hidden_proj.", "_hidden_proj", self._hidden_proj), + ("latent_proj.", "_latent_proj", self._latent_proj), + ("coordinate_proj.", "_coordinate_proj", self._coordinate_proj), + ("xvec_proj.", "_xvec_proj", self._xvec_proj), + ("eos_proj.", "_eos_proj", self._eos_proj), + ] + projector_state_keys = {prefix: set(mod.state_dict().keys()) for prefix, _, mod in projector_specs} + projector_matched: dict[str, list[tuple[str, torch.Tensor]]] = {prefix: [] for prefix, _, _ in projector_specs} + + matched_vae: list[tuple[str, torch.Tensor]] = [] + matched_dit: list[tuple[str, torch.Tensor]] = [] + matched_patch: list[tuple[str, torch.Tensor]] = [] + matched_llm: list[tuple[str, torch.Tensor]] = [] + matched_speaker: list[tuple[str, torch.Tensor]] = [] + skipped_lm_head = 0 + + DIT_PREFIX = "velocity_field_predictor." + PATCH_PREFIX = "patch_encoder." + LLM_MODEL_PREFIX = "llm.model." + LLM_LM_HEAD_PREFIX = "llm.lm_head." + for name, tensor in weights: + if name.startswith(DIT_PREFIX): + candidate = name[len(DIT_PREFIX) :] + if candidate in dit_state_keys: + matched_dit.append((candidate, tensor)) + continue + if name.startswith(PATCH_PREFIX): + candidate = name[len(PATCH_PREFIX) :] + if candidate in patch_state_keys: + matched_patch.append((candidate, tensor)) + continue + if name.startswith(LLM_LM_HEAD_PREFIX): + # Qwen2.5-1.5B-Base has tie_word_embeddings=True, and our + # ``self.model`` is a Qwen2Model (no lm_head submodule). Drop. + skipped_lm_head += 1 + continue + if name.startswith(LLM_MODEL_PREFIX): + candidate = name[len(LLM_MODEL_PREFIX) :] + # Hand off untouched — Qwen2Model.load_weights does its own + # fusion (q/k/v_proj -> qkv_proj, gate/up_proj -> gate_up_proj) + # and tolerates unknown keys, so no local key set to gate. + matched_llm.append((candidate, tensor)) + continue + # Projectors: 5 candidate top-level prefixes, first hit wins. + hit_projector = False + for prefix, _attr, _mod in projector_specs: + if name.startswith(prefix): + candidate = name[len(prefix) :] + if candidate in projector_state_keys[prefix]: + projector_matched[prefix].append((candidate, tensor)) + hit_projector = True + break + if hit_projector: + continue + # Speaker encoder: set-membership match (no distinguishing prefix + # in checkpoint; placed before the AudioVAE fallback so its keys + # don't get silently dropped by the catch-all). + if name in speaker_state_keys: + matched_speaker.append((name, tensor)) + continue + candidate = name[len("vocoder.") :] if name.startswith("vocoder.") else name + if candidate in vae_state_keys: + matched_vae.append((candidate, tensor)) + + loaded: set[str] = set() + + if matched_vae: + vae_loader = AutoWeightsLoader(self._audio_vae) + loaded_vae = vae_loader.load_weights(iter(matched_vae)) + loaded.update(f"_audio_vae.{name}" for name in loaded_vae) + logger.info( + "DotsTTS load_weights: loaded %d/%d AudioVAE tensors.", + len(loaded_vae), + len(vae_state_keys), + ) + + if matched_dit: + dit_loader = AutoWeightsLoader(self._head) + loaded_dit = dit_loader.load_weights(iter(matched_dit)) + loaded.update(f"_head.{name}" for name in loaded_dit) + logger.info( + "DotsTTS load_weights: loaded %d/%d DiT tensors.", + len(loaded_dit), + len(dit_state_keys), + ) + + if matched_patch: + patch_loader = AutoWeightsLoader(self._patch_encoder) + loaded_patch = patch_loader.load_weights(iter(matched_patch)) + loaded.update(f"_patch_encoder.{name}" for name in loaded_patch) + logger.info( + "DotsTTS load_weights: loaded %d/%d patch_encoder tensors.", + len(loaded_patch), + len(patch_state_keys), + ) + + any_projector_matched = False + for prefix, attr_name, mod in projector_specs: + matched = projector_matched[prefix] + if not matched: + continue + any_projector_matched = True + proj_loader = AutoWeightsLoader(mod) + loaded_proj = proj_loader.load_weights(iter(matched)) + loaded.update(f"{attr_name}.{name}" for name in loaded_proj) + logger.info( + "DotsTTS load_weights: loaded %d/%d %s tensors.", + len(loaded_proj), + len(projector_state_keys[prefix]), + prefix[:-1], + ) + + if matched_llm: + # vLLM's Qwen2Model.load_weights does its own fusion + tolerates + # unknown keys. It returns the set of param names it actually + # accepted; re-prefix with ``model.`` for our return value (our + # local attr is ``self.model``). + loaded_llm = self.model.load_weights(iter(matched_llm)) + loaded.update(f"model.{name}" for name in loaded_llm) + logger.info( + "DotsTTS load_weights: loaded %d Qwen2 tensors (fed %d; %d llm.lm_head.* skipped — tied embeddings).", + len(loaded_llm), + len(matched_llm), + skipped_lm_head, + ) + + if matched_speaker: + speaker_loader = AutoWeightsLoader(self._speaker_encoder) + loaded_speaker = speaker_loader.load_weights(iter(matched_speaker)) + loaded.update(f"_speaker_encoder.{name}" for name in loaded_speaker) + logger.info( + "DotsTTS load_weights: loaded %d/%d CAM++ speaker_encoder tensors.", + len(loaded_speaker), + len(speaker_state_keys), + ) + + if ( + not matched_vae + and not matched_dit + and not matched_patch + and not any_projector_matched + and not matched_llm + and not matched_speaker + ): + logger.warning( + "DotsTTS load_weights: no AudioVAE / DiT / " + "patch_encoder / projector / Qwen2 / speaker_encoder keys matched." + ) + + return loaded diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py new file mode 100644 index 00000000000..fe17b5c0bc5 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py @@ -0,0 +1,1434 @@ +# 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 +# +# Merged from six upstream files (preserved as-is, inference behaviour +# unchanged; only internal imports were folded inline): +# - modules/vocoder/config.py -> AudioVAEConfig +# - modules/backbone/layers.py -> Conv1d, ConvTranspose1d +# (Mlp/MultiHeadAttention/RotaryEmbedding +# are not vendored here; they belong +# to the DiT side-path module added later) +# - modules/vocoder/alias_free_filter.py -> sinc, kaiser_sinc_filter1d, +# LowPassFilter1d +# - modules/vocoder/alias_free_resample.py-> UpSample1d, DownSample1d +# - modules/vocoder/alias_free_act.py -> Snake, SnakeBeta, Activation1d +# - modules/vocoder/bigvgan.py -> AudioVAE and all its sub-modules +# +# Per-section upstream attributions are preserved at each section header. +# AudioVAEConfig is the only piece that diverges from upstream: pydantic +# ConfigBase replaced with a stdlib @dataclass (extra fields not exercised by +# any shipped checkpoint). + +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass, field +from fractions import Fraction + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import pow, sin +from torch.nn import Parameter +from torch.nn.utils import remove_weight_norm, weight_norm + +# ============================================================================ +# Configuration (adapted from modules/vocoder/config.py) +# ---------------------------------------------------------------------------- +# Upstream subclasses a pydantic ``ConfigBase`` with ``extra="allow"``; both the +# dots.tts-soar and dots.tts-mf checkpoint ``config.json`` pass only the 17 +# declared fields below, so the escape hatch is never exercised. ``.get()`` +# helper preserves bigvgan.py's ``h.get(...)`` call sites without modification. +# ``num_decoder_lookahead`` / ``num_encoder_lookahead`` are looked up via +# ``h.get(...)`` but not present in any shipped config; they fall back to 2. +# ============================================================================ + + +@dataclass +class AudioVAEConfig: + sample_rate: int = 24000 + upsample_rates: list[int] = field(default_factory=list) + upsample_kernel_sizes: list[int] = field(default_factory=list) + upsample_initial_channel: int = 1536 + resblock: str = "1" + resblock_kernel_sizes: list[int] = field(default_factory=list) + resblock_dilation_sizes: list[list[int]] = field(default_factory=list) + downsample_rates: list[int] = field(default_factory=list) + downsample_channels: list[int] = field(default_factory=list) + activation: str = "snakebeta" + snake_logscale: bool = True + latent_dim: int = 128 + causal: bool = False + mi_num_layers: int = 4 + causal_encoder: bool = False + use_bias_at_final: bool = True + use_tanh_at_final: bool = True + + num_decoder_lookahead: int = 2 + num_encoder_lookahead: int = 2 + + def get(self, key: str, default=None): + return getattr(self, key, default) + + +# ============================================================================ +# Conv layers (vendored from modules/backbone/layers.py) +# ---------------------------------------------------------------------------- +# Only the Conv1d / ConvTranspose1d helpers are needed by the AudioVAE. +# Mlp / MultiHeadAttention / RotaryEmbedding from the same upstream file are +# reserved for the DiT flow-matching head, which gets its own vendored module +# in a later step. +# ============================================================================ + + +class Conv1d(nn.Conv1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int = 1, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + padding_mode: str = "zeros", + bias: bool = True, + padding=None, + causal: bool = False, + **_kwargs, + ): + self.causal = causal + if padding is None: + if causal: + padding = 0 + self.left_padding = dilation * (kernel_size - 1) + else: + padding = int((kernel_size * dilation - dilation) / 2) + + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + padding_mode=padding_mode, + bias=bias, + ) + + self.in_channels = in_channels + + def forward(self, x): + if self.causal: + x = F.pad(x.unsqueeze(2), (self.left_padding, 0, 0, 0)).squeeze(2) + return super().forward(x) + + +class ConvTranspose1d(nn.ConvTranspose1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + output_padding: int = 0, + groups: int = 1, + bias: bool = True, + dilation: int = 1, + padding=None, + padding_mode: str = "zeros", + causal: bool = False, + **_kwargs, + ): + if padding is None: + padding = 0 if causal else (kernel_size - stride) // 2 + if causal: + assert padding == 0, "padding is not allowed in causal ConvTranspose1d." + assert kernel_size == 2 * stride, "kernel_size must be equal to 2*stride in Causal ConvTranspose1d." + + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + output_padding=output_padding, + groups=groups, + bias=bias, + dilation=dilation, + padding_mode=padding_mode, + ) + + self.causal = causal + self.stride = stride + + def forward(self, x): + x = super().forward(x) + if self.causal: + x = x[:, :, : -self.stride] + return x + + +# ============================================================================ +# Alias-free low-pass filter (vendored from modules/vocoder/alias_free_filter.py) +# ---------------------------------------------------------------------------- +# Adapted from https://github.com/junjun3518/alias-free-torch under +# the Apache License 2.0. The sinc fallback is adopted from adefossez's +# julius.core.sinc (MIT). +# ============================================================================ + + +if "sinc" in dir(torch): + sinc = torch.sinc +else: + + def sinc(x: torch.Tensor): + """ + Implementation of sinc, i.e. sin(pi * x) / (pi * x) + __Warning__: Different to julius.sinc, the input is multiplied by `pi`! + """ + return torch.where( + x == 0, + torch.tensor(1.0, device=x.device, dtype=x.dtype), + torch.sin(math.pi * x) / math.pi / x, + ) + + +def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): # return filter [1,1,kernel_size] + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + + # For kaiser window + delta_f = 4 * half_width + A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + if A > 50.0: + beta = 0.1102 * (A - 8.7) + elif A >= 21.0: + beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + + # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio + if even: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + if cutoff == 0: + filter_ = torch.zeros_like(time) + else: + filter_ = 2 * cutoff * window * sinc(2 * cutoff * time) + # Normalize filter to have sum = 1, otherwise we will have a small leakage + # of the constant component in the input signal. + filter_ /= filter_.sum() + filter = filter_.view(1, 1, kernel_size) + + return filter + + +class LowPassFilter1d(nn.Module): + def __init__( + self, + cutoff=0.5, + half_width=0.6, + stride: int = 1, + padding: bool = True, + padding_mode: str = "replicate", + kernel_size: int = 12, + channels: int = 1, + causal: bool = True, + fixed_filter: bool = False, + ): + # kernel_size should be even number for stylegan3 setup, + # in this implementation, odd number is also possible. + super().__init__() + if cutoff < -0.0: + raise ValueError("Minimum cutoff must be larger than zero.") + if cutoff > 0.5: + raise ValueError("A cutoff above 0.5 does not make sense.") + self.kernel_size = kernel_size + if causal: + self.pad_left = kernel_size - 1 + self.pad_right = 0 + else: + self.even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(self.even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.padding = padding + self.padding_mode = padding_mode + self.fixed_filter = fixed_filter + filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) + if fixed_filter: + self.register_buffer("filter", filter) + else: + self.filter = nn.Parameter(filter.expand(channels, -1, -1).clone()) + + # input [B, C, T] + def forward(self, x): + _, C, _ = x.shape + if self.padding: + x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) + if self.fixed_filter: + out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + else: + out = F.conv1d(x, self.filter, stride=self.stride, groups=C) + return out + + +# ============================================================================ +# Alias-free up/down sampling (vendored from modules/vocoder/alias_free_resample.py) +# ---------------------------------------------------------------------------- +# Adapted from https://github.com/junjun3518/alias-free-torch (Apache 2.0). +# ============================================================================ + + +class UpSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None, channels=None, causal=True, fixed_filter=False): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.stride = ratio + self.channels = channels + self.causal = causal + self.fixed_filter = fixed_filter + if causal: + self.pad = 0 + else: + self.pad = self.kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 + filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) + if self.fixed_filter: + self.register_buffer("filter", filter) + else: + self.filter = nn.Parameter(filter.expand(channels, -1, -1).clone()) + + # x: [B, C, T] + def forward(self, x): + _, C, _ = x.shape + x = F.pad(x, (self.pad, self.pad), mode="replicate") + if self.fixed_filter: + x = self.ratio * F.conv_transpose1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + else: + x = self.ratio * F.conv_transpose1d(x, self.filter, stride=self.stride, groups=C) + if self.causal: + x = x[..., : -(self.kernel_size - self.stride)] + else: + x = x[..., self.pad_left : -self.pad_right] + + return x + + +class DownSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None, channels=None, causal=True, fixed_filter=False): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.lowpass = LowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=self.kernel_size, + channels=channels, + causal=causal, + fixed_filter=fixed_filter, + ) + + def forward(self, x): + return self.lowpass(x) + + +# ============================================================================ +# Snake activations + alias-free Activation1d (vendored from +# modules/vocoder/alias_free_act.py) +# ---------------------------------------------------------------------------- +# Adapted from https://github.com/junjun3518/alias-free-torch (Apache 2.0). +# Snake / SnakeBeta paper: https://arxiv.org/abs/2006.08195 +# ============================================================================ + + +class Activation1d(nn.Module): + def __init__( + self, + activation, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + causal=True, + fixed_filter=False, + ): + super().__init__() + # causal=False + self.up_ratio = up_ratio + self.down_ratio = down_ratio + self.act = activation + self.upsample = UpSample1d( + up_ratio, + up_kernel_size, + activation.in_features, + causal=causal, + fixed_filter=fixed_filter, + ) + self.downsample = DownSample1d( + down_ratio, + down_kernel_size, + activation.in_features, + causal=causal, + fixed_filter=fixed_filter, + ) + + # x: [B,C,T] + def forward(self, x): + x = self.upsample(x) + x = self.act(x) + return self.downsample(x) + + +class Snake(nn.Module): + """ + Implementation of a sine-based periodic activation function + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter + """ + + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + super().__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + # Snake := x + 1/a * sin^2 (xa) + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + if self.alpha_logscale: + alpha = torch.exp(alpha) + return x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + +class SnakeBeta(nn.Module): + """ + A modified Snake function which uses separate parameters for the magnitude + of the periodic components. + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + """ + + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + super().__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + self.beta = Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + self.beta = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + # SnakeBeta := x + 1/b * sin^2 (xa) + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + return x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + +# ============================================================================ +# BigVGAN AudioVAE (vendored from modules/vocoder/bigvgan.py) +# ---------------------------------------------------------------------------- +# Copyright (c) 2022 NVIDIA CORPORATION. +# Licensed under the MIT license. +# Upstream layout flattened: the only modification is that the three +# original ``from dots_tts.modules.{backbone,vocoder}.* import ...`` lines +# are dropped, since those names are already defined above in this file. +# ============================================================================ + + +@dataclass(slots=True) +class BigVGANStreamState: + lstm_hidden: tuple[torch.Tensor, torch.Tensor] + decoder: DecoderStreamState + + +@dataclass(slots=True) +class DecoderStreamState: + window: torch.Tensor + chunk_size: int + total_frames: int = 0 + emitted_frames: int = 0 + + +def _empty_chunk( + ref: torch.Tensor, + *, + channels: int | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + return ref.new_zeros( + (ref.size(0), channels or ref.size(1), 0), + dtype=dtype or ref.dtype, + ) + + +def _module_state_device_dtype(module: nn.Module) -> tuple[torch.device, torch.dtype]: + for name in ("weight", "bias", "filter"): + tensor = getattr(module, name, None) + if tensor is not None: + return tensor.device, tensor.dtype + for tensor in itertools.chain(module.parameters(), module.buffers()): + return tensor.device, tensor.dtype + raise RuntimeError(f"Unable to infer state dtype/device for {type(module).__name__}.") + + +def _stream_state_zeros( + batch_size: int, + channels: int, + length: int, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + return torch.zeros( + (batch_size, channels, max(0, int(length))), + device=device, + dtype=dtype, + ) + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +class Conv1d_S(nn.Module): + "Conv1d for spectral normalisation and orthogonal initialisation" + + def __init__( + self, + in_channels, + out_channels, + kernel_size=1, + stride=1, + dilation=1, + groups=1, + causal=False, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.kernel_size = kernel_size + self.stride = stride + self.dilation = dilation + self.groups = groups + self.causal = causal + pad = 0 if causal else dilation * (kernel_size - 1) // 2 + self.causal_pad = dilation * (kernel_size - 1) if causal else 0 + + self.layer = weight_norm( + nn.Conv1d( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + padding=pad, + dilation=dilation, + groups=groups, + ) + ) + + def forward(self, inputs): + if self.causal and self.causal_pad > 0: + inputs = F.pad(inputs, (self.causal_pad, 0)) + return self.layer(inputs) + + +class SLSTM(nn.Module): + """ + LSTM without worrying about the hidden state, nor the layout of the data. + Expects input as convolutional layout. + """ + + def __init__( + self, + dimension: int, + num_layers: int = 2, + skip: bool = True, + bidirectional: bool = False, + ): + super().__init__() + self.skip = skip + self.bidirectional = bidirectional + self.lstm = nn.LSTM( + input_size=dimension, + hidden_size=dimension, + num_layers=num_layers, + bidirectional=bidirectional, + batch_first=True, + ) + self._stream_num_layers = num_layers + self._stream_weight_ih = tuple(getattr(self.lstm, f"weight_ih_l{layer_idx}") for layer_idx in range(num_layers)) + self._stream_weight_hh = tuple(getattr(self.lstm, f"weight_hh_l{layer_idx}") for layer_idx in range(num_layers)) + self._stream_bias_ih = tuple(getattr(self.lstm, f"bias_ih_l{layer_idx}") for layer_idx in range(num_layers)) + self._stream_bias_hh = tuple(getattr(self.lstm, f"bias_hh_l{layer_idx}") for layer_idx in range(num_layers)) + if self.bidirectional: + self.proj_out = nn.Linear(dimension * 2, dimension) + + def forward(self, x): + y, _ = self.lstm(x) + if self.bidirectional: + y = self.proj_out(y) + if self.skip: + y = y + x + return y + + def stream_step( + self, + x: torch.Tensor, + hidden: tuple[torch.Tensor, torch.Tensor] | None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + if self.bidirectional: + raise RuntimeError("Streaming only supports unidirectional SLSTM.") + + residual = x + if hidden is None: + hidden = self.init_stream_state(x.size(0)) + + hidden_h, hidden_c = hidden + next_hidden_h = [] + next_hidden_c = [] + for layer_idx in range(self._stream_num_layers): + layer_input = x + hx = hidden_h[layer_idx] + cx = hidden_c[layer_idx] + outputs = [] + weight_ih = self._stream_weight_ih[layer_idx] + weight_hh = self._stream_weight_hh[layer_idx] + bias_ih = self._stream_bias_ih[layer_idx] + bias_hh = self._stream_bias_hh[layer_idx] + + for frame_idx in range(x.size(1)): + gates = F.linear(layer_input[:, frame_idx, :], weight_ih, bias_ih) + gates = gates + F.linear(hx, weight_hh, bias_hh) + input_gate, forget_gate, cell_gate, output_gate = gates.chunk(4, dim=-1) + input_gate = torch.sigmoid(input_gate) + forget_gate = torch.sigmoid(forget_gate) + cell_gate = torch.tanh(cell_gate) + output_gate = torch.sigmoid(output_gate) + cx = forget_gate * cx + input_gate * cell_gate + hx = output_gate * torch.tanh(cx) + outputs.append(hx) + + x = torch.stack(outputs, dim=1) + next_hidden_h.append(hx) + next_hidden_c.append(cx) + + y = x + if self.skip: + y = y + residual + return y, (torch.stack(next_hidden_h, dim=0), torch.stack(next_hidden_c, dim=0)) + + def init_stream_state( + self, + batch_size: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + num_directions = 2 if self.bidirectional else 1 + state_shape = ( + self.lstm.num_layers * num_directions, + batch_size, + self.lstm.hidden_size, + ) + weight = self.lstm.weight_hh_l0 + return ( + weight.new_zeros(state_shape), + weight.new_zeros(state_shape), + ) + + +class ResStack(nn.Module): + def __init__(self, channel, kernel_size=3, base=3, nums=4, causal=False): + super().__init__() + + self.layers = nn.ModuleList([]) + for i in range(nums): + dil = base**i + pad1 = dil * (kernel_size - 1) if causal else dil + pad2 = (kernel_size - 1) if causal else 1 + block = [ + nn.LeakyReLU(), + ] + if causal and pad1 > 0: + block.append(nn.ConstantPad1d((pad1, 0), 0.0)) + block.append( + nn.utils.weight_norm( + nn.Conv1d( + channel, + channel, + kernel_size=kernel_size, + dilation=dil, + padding=0 if causal else pad1, + ) + ) + ) + block.append(nn.LeakyReLU()) + if causal and pad2 > 0: + block.append(nn.ConstantPad1d((pad2, 0), 0.0)) + block.append( + nn.utils.weight_norm( + nn.Conv1d( + channel, + channel, + kernel_size=kernel_size, + dilation=1, + padding=0 if causal else pad2, + ) + ) + ) + self.layers.append(nn.Sequential(*block)) + + def forward(self, x): + for layer in self.layers: + x = x + layer(x) + return x + + +class Encoder(nn.Module): + def __init__( + self, + in_channels=1, + out_channels=100, + base_channels=12, + proj_kernel_size=3, + stack_kernel_size=3, + stack_dilation_base=2, + stacks=6, + channels=(12, 24, 48, 96, 192, 384, 768), + down_sample_factors=(2, 2, 2, 2, 4, 4), + causal=False, + lookahead=0, + ): + super().__init__() + + act_slope = 0.2 + layers = [] + # pre proj_layer + layers += [ + Conv1d_S( + in_channels, + base_channels, + kernel_size=proj_kernel_size, + stride=1, + causal=causal, + ), + nn.LeakyReLU(act_slope, True), + ] + + # channels: [512, 256, 128, 64], upsample_factors: [5, 2, 2] + for (in_c, out_c), down_f in zip(itertools.pairwise(channels), down_sample_factors, strict=True): + layers += [ + Conv1d_S(in_c, out_c, kernel_size=down_f * 2, stride=down_f, causal=causal), + ResStack(out_c, stack_kernel_size, stack_dilation_base, stacks, causal=causal), + nn.LeakyReLU(act_slope, True), + ] + + # post layers + if lookahead > 0: + layers += [ + Conv1d_S( + channels[-1], + out_channels, + kernel_size=lookahead * 2 + 1, + stride=1, + causal=False, + ), + ] + else: + layers += [ + Conv1d_S( + channels[-1], + out_channels, + kernel_size=proj_kernel_size, + stride=1, + causal=causal, + ), + ] + self.generator = nn.Sequential(*layers) + + def forward(self, conditions, _z_inputs=None): + return self.generator(conditions) + + +class AMPBlock1(torch.nn.Module): + def __init__( + self, + h, + channels, + kernel_size=3, + dilation=(1, 3, 5), + activation=None, + causal=True, + ): + super().__init__() + self.h = h + + self.convs1 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + causal=causal, + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + causal=causal, + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[2], + causal=causal, + ) + ), + ] + ) + self.convs1.apply(init_weights) + + self.convs2 = nn.ModuleList( + [ + weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)), + weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)), + weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)), + ] + ) + self.convs2.apply(init_weights) + + self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers + + if activation == "snake": # periodic nonlinearity with snake function and anti-aliasing + self.activations = nn.ModuleList( + [ + Activation1d( + activation=Snake(channels, alpha_logscale=h.snake_logscale), + causal=causal, + fixed_filter=True, + ) + for _ in range(self.num_layers) + ] + ) + elif activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing + self.activations = nn.ModuleList( + [ + Activation1d( + activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale), + causal=causal, + fixed_filter=True, + ) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + acts1, acts2 = self.activations[::2], self.activations[1::2] + for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2, strict=True): + xt = a1(x) + xt = c1(xt) + xt = a2(xt) + xt = c2(xt) + x = xt + x + + return x + + def remove_weight_norm(self): + for layer in self.convs1: + remove_weight_norm(layer) + for layer in self.convs2: + remove_weight_norm(layer) + + +class AMPBlock2(torch.nn.Module): + def __init__(self, h, channels, kernel_size=3, dilation=(1, 3), activation=None, causal=True): + super().__init__() + self.h = h + + self.convs = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + causal=causal, + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + causal=causal, + ) + ), + ] + ) + self.convs.apply(init_weights) + + self.num_layers = len(self.convs) # total number of conv layers + + if activation == "snake": # periodic nonlinearity with snake function and anti-aliasing + self.activations = nn.ModuleList( + [ + Activation1d( + activation=Snake(channels, alpha_logscale=h.snake_logscale), + causal=causal, + fixed_filter=True, + ) + for _ in range(self.num_layers) + ] + ) + elif activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing + self.activations = nn.ModuleList( + [ + Activation1d( + activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale), + causal=causal, + fixed_filter=True, + ) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + for c, a in zip(self.convs, self.activations, strict=True): + xt = a(x) + xt = c(xt) + x = xt + x + + return x + + def remove_weight_norm(self): + for layer in self.convs: + remove_weight_norm(layer) + + +class Decoder(nn.Module): + def __init__(self, h): + super().__init__() + self.h = h + causal = h.causal + self._stream_window_sizes: dict[int, int] = {} + + self.num_kernels = len(h.resblock_kernel_sizes) + self.num_upsamples = len(h.upsample_rates) + + num_decoder_lookahead = h.get("num_decoder_lookahead", 2) + # pre conv + self.conv_pre = weight_norm( + Conv1d( + h.latent_dim, + h.upsample_initial_channel, + kernel_size=2 * num_decoder_lookahead + 1, + stride=1, + causal=False, + ) + ) + + # define which AMPBlock to use. BigVGAN uses AMPBlock1 as default + resblock = AMPBlock1 if h.resblock == "1" else AMPBlock2 + + # transposed conv-based upsamplers. does not apply anti-aliasing + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes, strict=True)): + self.ups.append( + nn.ModuleList( + [ + weight_norm( + ConvTranspose1d( + h.upsample_initial_channel // (2**i), + h.upsample_initial_channel // (2 ** (i + 1)), + k, + u, + causal=causal, + ) + ) + ] + ) + ) + + # residual blocks using anti-aliased multi-periodicity composition modules (AMP) + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = h.upsample_initial_channel // (2 ** (i + 1)) + for k, d in zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes, strict=True): + self.resblocks.append(resblock(h, ch, k, d, activation=h.activation, causal=causal)) + + # post conv + if h.activation == "snake": # periodic nonlinearity with snake function and anti-aliasing + activation_post = Snake(ch, alpha_logscale=h.snake_logscale) + self.activation_post = Activation1d(activation=activation_post, causal=causal, fixed_filter=False) + elif h.activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing + activation_post = SnakeBeta(ch, alpha_logscale=h.snake_logscale) + self.activation_post = Activation1d(activation=activation_post, causal=causal, fixed_filter=False) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, causal=causal, bias=h.get("use_bias_at_final", True))) + + # weight initialization + for i in range(len(self.ups)): + self.ups[i].apply(init_weights) + self.conv_post.apply(init_weights) + + def forward(self, z): + # pre conv + x = self.conv_pre(z) + + for i in range(self.num_upsamples): + # upsampling + for i_up in range(len(self.ups[i])): + x = self.ups[i][i_up](x) + # AMP blocks + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + # post conv + x = self.activation_post(x) + x = self.conv_post(x) + if self.h.get("use_tanh_at_final", True): + x = torch.tanh(x) + else: + x = torch.clamp(x, min=-1.0, max=1.0) # Bound the output to [-1, 1] + + return x + + @property + def stream_lookahead(self) -> int: + return (self.conv_pre.kernel_size[0] - 1) // 2 + + @staticmethod + def _conv1d_left_context(layer) -> int: + dilation = layer.dilation[0] if isinstance(layer.dilation, tuple) else layer.dilation + kernel_size = layer.kernel_size[0] if isinstance(layer.kernel_size, tuple) else layer.kernel_size + if getattr(layer, "causal", False): + return dilation * (kernel_size - 1) + return layer.padding[0] if isinstance(layer.padding, tuple) else layer.padding + + @staticmethod + def _convtranspose1d_left_context(layer) -> int: + stride = layer.stride[0] if isinstance(layer.stride, tuple) else layer.stride + kernel_size = layer.kernel_size[0] if isinstance(layer.kernel_size, tuple) else layer.kernel_size + if not getattr(layer, "causal", False): + raise NotImplementedError("Streaming only supports causal ConvTranspose1d.") + if kernel_size != 2 * stride: + raise ValueError( + "Streaming ConvTranspose1d expects kernel_size == 2 * stride, got " + f"kernel_size={kernel_size} stride={stride}." + ) + return 1 + + @classmethod + def _activation_left_context(cls, activation: Activation1d) -> int: + upsample = activation.upsample + downsample = activation.downsample.lowpass + if not upsample.causal or not downsample.padding or downsample.pad_right != 0: + raise NotImplementedError("Streaming only supports causal alias-free activations.") + ratio = int(upsample.ratio) + if ratio != int(downsample.stride): + raise ValueError( + "Alias-free activation expects matched up/down ratios, got " + f"up_ratio={ratio} down_ratio={downsample.stride}." + ) + total_left = (upsample.kernel_size - 1) + (downsample.kernel_size - 1) + return (total_left + ratio - 1) // ratio + + @classmethod + def _ampblock_left_context(cls, block) -> int: + if isinstance(block, AMPBlock1): + left_context = 0 + acts1 = block.activations[::2] + acts2 = block.activations[1::2] + for conv1, conv2, act1, act2 in zip( + block.convs1, + block.convs2, + acts1, + acts2, + strict=True, + ): + left_context += ( + cls._activation_left_context(act1) + + cls._conv1d_left_context(conv1) + + cls._activation_left_context(act2) + + cls._conv1d_left_context(conv2) + ) + return left_context + if isinstance(block, AMPBlock2): + left_context = 0 + for conv, activation in zip(block.convs, block.activations, strict=True): + left_context += cls._activation_left_context(activation) + cls._conv1d_left_context(conv) + return left_context + raise TypeError(f"Unsupported resblock type: {type(block).__name__}.") + + def _stream_left_context(self) -> int: + left_context = Fraction(self._conv1d_left_context(self.conv_pre), 1) + current_scale = Fraction(1, 1) + for stage_idx, upsample_layers in enumerate(self.ups): + for upsample in upsample_layers: + left_context += current_scale * self._convtranspose1d_left_context(upsample) + stride = upsample.stride[0] if isinstance(upsample.stride, tuple) else upsample.stride + current_scale /= int(stride) + + stage_start = stage_idx * self.num_kernels + stage_end = stage_start + self.num_kernels + stage_context = max(self._ampblock_left_context(block) for block in self.resblocks[stage_start:stage_end]) + left_context += current_scale * stage_context + + left_context += current_scale * self._activation_left_context(self.activation_post) + left_context += current_scale * self._conv1d_left_context(self.conv_post) + return int(left_context.__ceil__()) + + def stream_window_size(self, chunk_size: int) -> int: + if chunk_size < 1: + raise ValueError(f"chunk_size must be >= 1, got {chunk_size}.") + cached = self._stream_window_sizes.get(chunk_size) + if cached is not None: + return cached + + window_size = chunk_size + self.stream_lookahead + self._stream_left_context() + self._stream_window_sizes[chunk_size] = window_size + return window_size + + def remove_weight_norm(self): + for upsample_layers in self.ups: + for upsample_layer in upsample_layers: + remove_weight_norm(upsample_layer) + for resblock in self.resblocks: + resblock.remove_weight_norm() + remove_weight_norm(self.conv_pre) + remove_weight_norm(self.conv_post) + + +class AudioVAE(nn.Module): + def __init__(self, h: AudioVAEConfig): + super().__init__() + self.config = h + + self.h = h + self.hop_size = int(np.prod(h.downsample_rates)) + self.sample_rate = h.sample_rate + self.decoder_lookahead = int(h.get("num_decoder_lookahead", 2)) + + self.audio_encoder = Encoder( + out_channels=h.latent_dim, + down_sample_factors=h.downsample_rates, + channels=h.downsample_channels, + causal=h.causal_encoder, + lookahead=h.get("num_encoder_lookahead", 2), + ) + + intermediate_size = h.latent_dim * 4 + self.enc_mi_layer = nn.Sequential( + nn.Linear(h.latent_dim, intermediate_size), + SLSTM(intermediate_size, num_layers=h.mi_num_layers), + nn.Linear(intermediate_size, h.latent_dim), + ) + self.dec_mi_layer = nn.Sequential( + nn.Linear(h.latent_dim, intermediate_size), + SLSTM(intermediate_size, num_layers=h.mi_num_layers), + nn.Linear(intermediate_size, h.latent_dim), + ) + self.pre_proj = Conv1d( + in_channels=h.latent_dim, + out_channels=h.latent_dim * 2, + kernel_size=1, + stride=1, + ) + self.post_proj = Conv1d(in_channels=h.latent_dim, out_channels=h.latent_dim, kernel_size=1, stride=1) + + self.decoder = Decoder(h) + + def inference(self, data): + latents = self.extract_latents(data["sample"]) + return {"sample": self.inference_from_latents(latents)} + + @torch.autocast(enabled=False, device_type="cuda") + def extract_latents(self, x, do_sample=False): + x = x.float() + x = self.audio_encoder(x) + x = x.permute(0, 2, 1) + x = self.enc_mi_layer(x) + x = x.permute(0, 2, 1) + x = self.pre_proj(x) + if do_sample: + m_q, logs_q = torch.split(x, self.h.latent_dim, dim=1) + x = m_q + torch.randn_like(m_q) * torch.exp(logs_q) + return x + + def inference_from_latents(self, x, do_sample=True, noise_scale=1.0): + if do_sample: + assert x.size(1) == self.h.latent_dim * 2, f"Input must be like [B, D, H], got {x.shape}" + m_q, logs_q = torch.split(x, self.h.latent_dim, dim=1) + x = m_q + torch.randn_like(m_q) * torch.exp(logs_q) * noise_scale + else: + assert x.size(1) == self.h.latent_dim, f"Input must be like [B, D, H], got {x.shape}" + x = self.post_proj(x) + x = x.permute(0, 2, 1) + x = self.dec_mi_layer(x) + x = x.permute(0, 2, 1) + return self.decoder(x) + + def _validate_stream_latents(self, latents: torch.Tensor) -> None: + if latents.ndim != 3: + raise ValueError( + f"Streaming latents must have shape [batch, latent_dim, frames], got {tuple(latents.shape)}." + ) + if latents.size(1) != self.h.latent_dim: + raise ValueError(f"Streaming latent_dim must be {self.h.latent_dim}, got {latents.size(1)}.") + + def init_stream_state( + self, + batch_size: int = 1, + chunk_size: int = 8, + ) -> BigVGANStreamState: + if not self.h.causal: + raise RuntimeError("Strict streaming requires a causal vocoder.") + if batch_size < 1: + raise ValueError(f"batch_size must be >= 1, got {batch_size}.") + window_size = self.decoder.stream_window_size(chunk_size) + device, dtype = _module_state_device_dtype(self.decoder.conv_pre) + return BigVGANStreamState( + lstm_hidden=self.dec_mi_layer[1].init_stream_state(batch_size), + decoder=DecoderStreamState( + window=_stream_state_zeros( + batch_size, + self.h.latent_dim, + window_size, + device=device, + dtype=dtype, + ), + chunk_size=int(chunk_size), + ), + ) + + def _decode_stream_latents( + self, + latents: torch.Tensor, + lstm_hidden: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + self._validate_stream_latents(latents) + latents = latents.float() + x = self.post_proj(latents) + x = x.permute(0, 2, 1) + x = self.dec_mi_layer[0](x) + x, lstm_hidden = self.dec_mi_layer[1].stream_step(x, lstm_hidden) + x = self.dec_mi_layer[2](x) + decoder_dtype = self.decoder.conv_pre.weight.dtype + return x.permute(0, 2, 1).to(dtype=decoder_dtype), lstm_hidden + + def _prepare_stream_decoder_input( + self, + latents: torch.Tensor, + state: BigVGANStreamState, + ) -> torch.Tensor: + decoder_input, state.lstm_hidden = self._decode_stream_latents( + latents, + state.lstm_hidden, + ) + return decoder_input + + def _append_stream_decoder_input_tensor( + self, + decoder_input: torch.Tensor, + window: torch.Tensor, + valid_frames: torch.Tensor, + ) -> torch.Tensor: + if window.dtype != decoder_input.dtype: + window = window.to(dtype=decoder_input.dtype) + chunk_size = int(decoder_input.size(-1)) + if chunk_size >= window.size(-1): + raise ValueError(f"decoder window size {window.size(-1)} must be larger than chunk_size {chunk_size}.") + positions = torch.arange( + window.size(-1), + device=window.device, + dtype=valid_frames.dtype, + ) + clipped_valid = valid_frames.clamp(min=0, max=window.size(-1)) + combined = torch.cat( + [window, decoder_input.new_zeros(window.size(0), window.size(1), chunk_size)], + dim=-1, + ) + insert_index = clipped_valid + torch.arange( + chunk_size, + device=window.device, + dtype=valid_frames.dtype, + ) + combined.scatter_( + -1, + insert_index.view(1, 1, -1).expand_as(decoder_input), + decoder_input, + ) + new_valid = (clipped_valid + chunk_size).clamp(max=window.size(-1)) + start = (clipped_valid + chunk_size - window.size(-1)).clamp(min=0) + gather_index = (start + positions).clamp(max=combined.size(-1) - 1) + gathered = combined.gather( + -1, + gather_index.view(1, 1, -1).expand_as(window), + ) + mask = (positions < new_valid).to(dtype=window.dtype).view(1, 1, -1) + return gathered * mask + + def _append_stream_decoder_input( + self, + decoder_input: torch.Tensor, + state: BigVGANStreamState, + ) -> torch.Tensor: + decoder_state = state.decoder + chunk_size = int(decoder_input.size(-1)) + if chunk_size != decoder_state.chunk_size: + raise ValueError(f"Streaming chunk_size must stay fixed at {decoder_state.chunk_size}, got {chunk_size}.") + window = decoder_state.window + valid_frames = min(decoder_state.total_frames, window.size(-1)) + valid_frames_tensor = window.new_tensor(valid_frames, dtype=torch.int64) + new_window = self._append_stream_decoder_input_tensor( + decoder_input, + window, + valid_frames_tensor, + ) + decoder_state.window = new_window + decoder_state.total_frames += chunk_size + return new_window + + def compiled_stream_step( + self, + latents: torch.Tensor, + hidden_h: torch.Tensor, + hidden_c: torch.Tensor, + window: torch.Tensor, + valid_frames: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + decoder_input, (hidden_h, hidden_c) = self._decode_stream_latents( + latents, + (hidden_h, hidden_c), + ) + new_window = self._append_stream_decoder_input_tensor( + decoder_input, + window, + valid_frames, + ) + audio_window = self.decode_stream_window(new_window) + return audio_window, hidden_h, hidden_c, new_window + + def decode_stream_window(self, window: torch.Tensor) -> torch.Tensor: + decoder_dtype = self.decoder.conv_pre.weight.dtype + if window.dtype != decoder_dtype: + window = window.to(dtype=decoder_dtype) + return self.decoder(window) + + def _slice_stream_audio_window( + self, + audio_window: torch.Tensor, + state: BigVGANStreamState, + *, + final: bool, + ) -> torch.Tensor: + decoder_state = state.decoder + stable_end = ( + decoder_state.total_frames if final else max(0, decoder_state.total_frames - self.decoder.stream_lookahead) + ) + if stable_end <= decoder_state.emitted_frames: + return _empty_chunk(audio_window, channels=1) + + window_size = decoder_state.window.size(-1) + valid_frames = min(decoder_state.total_frames, window_size) + window_start = decoder_state.total_frames - valid_frames + if decoder_state.emitted_frames < window_start: + raise RuntimeError("Decoder stream window is too short for fixed-graph decoding.") + + local_start = decoder_state.emitted_frames - window_start + local_end = stable_end - window_start + sample_start = local_start * self.hop_size + sample_end = local_end * self.hop_size + decoder_state.emitted_frames = stable_end + return audio_window[..., sample_start:sample_end] + + def stream_step( + self, + latents: torch.Tensor, + state: BigVGANStreamState, + ) -> torch.Tensor: + decoder_input = self._prepare_stream_decoder_input(latents, state) + window = self._append_stream_decoder_input(decoder_input, state) + audio_window = self.decode_stream_window(window) + return self._slice_stream_audio_window(audio_window, state, final=False) + + def stream_flush(self, state: BigVGANStreamState) -> torch.Tensor: + audio_window = self.decode_stream_window(state.decoder.window) + return self._slice_stream_audio_window(audio_window, state, final=True) + + def remove_weight_norm(self): + self.decoder.remove_weight_norm() + + +__all__ = ["AudioVAEConfig", "AudioVAE"] diff --git a/vllm_omni/model_executor/models/dots_tts/pipeline.py b/vllm_omni/model_executor/models/dots_tts/pipeline.py new file mode 100644 index 00000000000..22b50c62933 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/pipeline.py @@ -0,0 +1,46 @@ +# 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], + }, + ), + ), +) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 498f9eb85cb..530fb171216 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -301,6 +301,12 @@ "voxcpm2_talker", "VoxCPM2TalkerForConditionalGeneration", ), + ## dots.tts + "DotsTTSForConditionalGeneration": ( + "dots_tts", + "dots_tts_talker", + "DotsTTSForConditionalGeneration", + ), ## Voxtral TTS "VoxtralTTSForConditionalGeneration": ( "voxtral_tts", diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 88a23dc0c52..cdf105d693e 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -20,6 +20,7 @@ "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", "OmniVoiceConfig": "vllm_omni.transformers_utils.configs.omnivoice", @@ -42,6 +43,7 @@ "FishSpeechFastARConfig", "GLMTTSConfig", "VoxCPM2Config", + "DotsTTSConfig", "VoxtralTTSConfig", "CosyVoice3Config", "OmniVoiceConfig", @@ -70,6 +72,7 @@ def __dir__(): # Eagerly import all config modules so their AutoConfig.register() side-effects # run as soon as `vllm_omni.transformers_utils.configs` is imported. 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 diff --git a/vllm_omni/transformers_utils/configs/dots_tts.py b/vllm_omni/transformers_utils/configs/dots_tts.py new file mode 100644 index 00000000000..1950a2b3a17 --- /dev/null +++ b/vllm_omni/transformers_utils/configs/dots_tts.py @@ -0,0 +1,102 @@ +# 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/rednote-hilab/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"]