[Model][TTS][WIP] Add dots.tts (rednote-hilab) — continuous-AR 48kHz TTS (#4765)
Signed-off-by: Moore-Z <zhumoore9661996@gmail.com>
This commit is contained in:
1 parent
dbc0dd6d85
commit
90a08c0eda
19 files changed
+5112
No files matched your search
@@ -59,6 +59,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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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()
|
||||
@@ -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 <path>` 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.
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -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 == {}
|
||||
@@ -47,6 +47,7 @@ from vllm_omni.model_executor.models.bagel.pipeline import (
|
||||
)
|
||||
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 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
|
||||
"hunyuan_video_15": HUNYUAN_VIDEO_15_PIPELINE,
|
||||
"wan2_2_ti2v": WAN2_2_TI2V_PIPELINE,
|
||||
"voxcpm2": VOXCPM2_PIPELINE,
|
||||
"dots_tts": DOTS_TTS_PIPELINE,
|
||||
"cosyvoice3": COSYVOICE3_PIPELINE,
|
||||
"audex_tts": AUDEX_TTS_PIPELINE,
|
||||
"audex_tta": AUDEX_TTA_PIPELINE,
|
||||
|
||||
@@ -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
|
||||
@@ -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.
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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 "[文本]<text>[文本对应语音]", audio_gen_start_id]
|
||||
|
||||
talker's preprocess() embeds these into ``hidden_states`` for the AR
|
||||
base LM. Decode then drives via ``state.curr_embed_for_next`` (DiT
|
||||
→ patch_encoder AR loopback).
|
||||
"""
|
||||
if not isinstance(text, str) or not text.strip():
|
||||
raise ValueError(f"text must be a non-empty string, got {text!r}")
|
||||
|
||||
prompt_text = f"{_TTS_TEXT_PREFIX}{text}{_TTS_AUDIO_PREFIX}"
|
||||
text_ids = tokenizer.encode(prompt_text, add_special_tokens=False)
|
||||
|
||||
audio_gen_start_id = tokenizer.convert_tokens_to_ids(_AUDIO_GEN_START_TOKEN)
|
||||
if not isinstance(audio_gen_start_id, int) or audio_gen_start_id == tokenizer.unk_token_id:
|
||||
raise ValueError(
|
||||
f"Tokenizer does not know {_AUDIO_GEN_START_TOKEN!r} "
|
||||
f"(got id={audio_gen_start_id!r}). Did you load the "
|
||||
f"dots.tts tokenizer (with added_tokens.json)?"
|
||||
)
|
||||
prompt_token_ids = list(text_ids) + [audio_gen_start_id]
|
||||
|
||||
prompt = tokens_input(prompt_token_ids=prompt_token_ids)
|
||||
prompt["additional_information"] = {
|
||||
# Reserved for future voice clone (ref_audio / speaker x-vector).
|
||||
# V1 zero-shot leaves this empty; preprocess() ignores it.
|
||||
}
|
||||
return prompt
|
||||
@@ -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)
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -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],
|
||||
},
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -301,6 +301,12 @@ _OMNI_MODELS = {
|
||||
"voxcpm2_talker",
|
||||
"VoxCPM2TalkerForConditionalGeneration",
|
||||
),
|
||||
## dots.tts
|
||||
"DotsTTSForConditionalGeneration": (
|
||||
"dots_tts",
|
||||
"dots_tts_talker",
|
||||
"DotsTTSForConditionalGeneration",
|
||||
),
|
||||
## Voxtral TTS
|
||||
"VoxtralTTSForConditionalGeneration": (
|
||||
"voxtral_tts",
|
||||
|
||||
@@ -20,6 +20,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
|
||||
"FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech",
|
||||
"GLMTTSConfig": "vllm_omni.transformers_utils.configs.glm_tts",
|
||||
"VoxCPM2Config": "vllm_omni.transformers_utils.configs.voxcpm2",
|
||||
"DotsTTSConfig": "vllm_omni.transformers_utils.configs.dots_tts",
|
||||
"VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts",
|
||||
"CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3",
|
||||
"OmniVoiceConfig": "vllm_omni.transformers_utils.configs.omnivoice",
|
||||
@@ -42,6 +43,7 @@ __all__ = [
|
||||
"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
|
||||
|
||||
@@ -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"]
|
||||
Reference in new issue
Block a user