[Chore][TTS] Reference recipes.vllm.ai recipes, remove dead code, and fix stale references (#4567)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
This commit is contained in:
Yueqian Lin authored and GitHub committed 2026-06-19 15:05:26 +08:00
1 parent f05e03aea8
commit 208932210f
46 files changed
+141 -787

No files matched your search

+1
View File
@@ -68,6 +68,7 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
- [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
- [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
- [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
- [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving
## Contributing
+4 -1
View File
@@ -6,6 +6,8 @@ vLLM-Omni supports unified multimodal comprehension and generation models across
If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm_omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
For deployment recipes (TTS and multimodal), see [recipes.vllm.ai](https://recipes.vllm.ai) and the in-repo [`recipes/`](https://github.com/vllm-project/vllm-omni/tree/main/recipes) directory.
## List of Supported Models
<style>
@@ -63,13 +65,14 @@ th {
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | |
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
+7
View File
@@ -9,6 +9,13 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, and MOSS-TTS-Nano.
!!! tip "Deployment recipes"
TTS deployment recipes are published at
[recipes.vllm.ai](https://recipes.vllm.ai) (e.g.
[Qwen3-TTS](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice),
[Higgs-Audio v3](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b)).
The in-repo runbooks live under [`recipes/`](https://github.com/vllm-project/vllm-omni/tree/main/recipes).
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
## Quick Start
+10 -1
View File
@@ -12,6 +12,13 @@ that is aligned with
that external repository as a reference for structure rather than the place to
add files for this repo. Use one Markdown file per model family by default.
> **Note:** The vLLM-Omni TTS recipes (12 models across 8 providers) are also
> published on the rendered recipes site at
> [recipes.vllm.ai](https://recipes.vllm.ai) (see
> [vllm-project/recipes#554](https://github.com/vllm-project/recipes/pull/554)).
> Individual pages resolve at `https://recipes.vllm.ai/<provider>/<model-id>`,
> e.g. [`bosonai/higgs-audio-v3-tts-4b`](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b).
Example layout:
```text
@@ -30,11 +37,12 @@ recipes/
| [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
| [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
| [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs |
| [`BosonAI/Higgs-Audio-V3-TTS.md`](./BosonAI/Higgs-Audio-V3-TTS.md) | Online + offline multilingual TTS with voice cloning | 1x H100 80GB |
| [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
| [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
| [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
| [`inclusionAI/Ming-omni-tts-0.5B.md`](./inclusionAI/Ming-omni-tts-0.5B.md) | Offline + online dense Ming TTS/audio generation | 1x AMD `gfx942` (ROCm 7.2) |
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
@@ -42,6 +50,7 @@ recipes/
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
@@ -672,7 +672,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
if model_stage in _OMNIVOICE_TTS_MODEL_STAGES:
return "omnivoice"
if model_stage in _COVO_AUDIO_MODEL_STAGES:
model_arch = getattr(self._tts_stage.engine_args, "model_arch", None)
if model_arch and "CovoAudio" in model_arch:
return "covo_audio"
if model_stage in _VOXCPM2_TTS_MODEL_STAGES:
@@ -194,6 +194,3 @@ __all__ = [
"SpeechServingContext",
"TTSModelAdapter",
]
# Type alias kept for readability in adapter signatures.
NormalizeFn = Callable[["OpenAICreateSpeechRequest"], None]
@@ -479,14 +479,6 @@ class CosyVoice3Model(
# Keep additional information synchronized for async_chunk updates.
self.enable_update_additional_information = True
# Expose streaming parameters
self.token_overlap_len = self.code2wav.token_overlap_len
self.mel_overlap_len = self.code2wav.mel_overlap_len
self.mel_window = self.code2wav.mel_window
self.mel_cache_len = self.code2wav.mel_cache_len
self.source_cache_len = self.code2wav.source_cache_len
self.speech_window = self.code2wav.speech_window
self._stream_audio_cache_by_req: dict[str, torch.Tensor] = {}
self._stream_audio_cache_lock = Lock()
self._stream_vocoder_cache_by_req: dict[str, dict[str, torch.Tensor]] = {}
else:
@@ -511,27 +503,6 @@ class CosyVoice3Model(
# Use parent's cache config - critical for PagedAttention to work correctly
return parent_config.with_hf_config(qwen_hf_config, architectures=["Qwen2Model"])
@staticmethod
def _cross_fade_audio(audio: torch.Tensor, prev_tail: torch.Tensor) -> torch.Tensor:
"""Blend previous chunk tail into current chunk head using a Hamming window.
This mirrors upstream CosyVoice's `fade_in_out(...)` semantics:
update the current head in-place using a 2*overlap window, then
concatenate the unchanged remainder.
"""
if audio.numel() == 0 or prev_tail.numel() == 0:
return audio
overlap = min(int(audio.numel()), int(prev_tail.numel()))
if overlap <= 0:
return audio
window = torch.hamming_window(2 * overlap, periodic=False, dtype=audio.dtype, device=audio.device)
fade_in = window[:overlap]
fade_out = window[overlap:]
blended = audio[:overlap] * fade_in + prev_tail[-overlap:].to(device=audio.device, dtype=audio.dtype) * fade_out
if overlap == int(audio.numel()):
return blended
return torch.cat([blended, audio[overlap:]], dim=0)
def _stitch_stream_audio(self, req_id: str | None, audio: torch.Tensor, stream_finished: bool) -> torch.Tensor:
"""Pass-through stitching for async_chunk.
@@ -539,11 +510,9 @@ class CosyVoice3Model(
Applying an additional waveform-domain fade/cache step introduces either
duplicated overlap (if no tail trim) or duration shrink (if tail trim).
"""
if req_id is not None and stream_finished and hasattr(self, "_stream_audio_cache_by_req"):
if req_id is not None and stream_finished and hasattr(self, "_stream_vocoder_cache_by_req"):
with self._stream_audio_cache_lock:
self._stream_audio_cache_by_req.pop(req_id, None)
if hasattr(self, "_stream_vocoder_cache_by_req"):
self._stream_vocoder_cache_by_req.pop(req_id, None)
self._stream_vocoder_cache_by_req.pop(req_id, None)
return audio
@staticmethod
@@ -11,8 +11,6 @@ This module contains the code2wav (token-to-waveform) stage which uses:
from __future__ import annotations
from contextlib import nullcontext
import numpy as np
import torch
import torch.nn as nn
@@ -175,19 +173,18 @@ class CosyVoice3Code2Wav(nn.Module):
prompt_token_len = torch.tensor([prompt_token.shape[1]], device=device, dtype=torch.int32)
prompt_feat_len = torch.tensor([prompt_feat.shape[1]], device=device, dtype=torch.int32)
with nullcontext():
feat, _ = self.flow_model.inference(
token=token,
token_len=token_len,
prompt_token=prompt_token,
prompt_token_len=prompt_token_len,
prompt_feat=prompt_feat,
prompt_feat_len=prompt_feat_len,
embedding=embedding,
streaming=streaming,
finalize=finalize,
n_timesteps=n_timesteps,
)
feat, _ = self.flow_model.inference(
token=token,
token_len=token_len,
prompt_token=prompt_token,
prompt_token_len=prompt_token_len,
prompt_feat=prompt_feat,
prompt_feat_len=prompt_feat_len,
embedding=embedding,
streaming=streaming,
finalize=finalize,
n_timesteps=n_timesteps,
)
trim_mel = max(0, int(token_offset_tokens)) * int(self.token_mel_ratio)
if trim_mel > 0:
@@ -195,28 +192,6 @@ class CosyVoice3Code2Wav(nn.Module):
return feat
@staticmethod
def _fade_speech(
speech: torch.Tensor,
prev_speech: torch.Tensor,
) -> torch.Tensor:
"""Blend previous speech tail into current speech head."""
if speech.numel() == 0 or prev_speech.numel() == 0:
return speech
overlap = min(int(speech.shape[-1]), int(prev_speech.shape[-1]))
if overlap <= 0:
return speech
window = torch.hamming_window(2 * overlap, periodic=False, dtype=speech.dtype, device=speech.device)
fade_in = window[:overlap].view(1, -1)
fade_out = window[overlap:].view(1, -1)
blended_head = (
speech[:, :overlap] * fade_in
+ prev_speech[:, -overlap:].to(device=speech.device, dtype=speech.dtype) * fade_out
)
if overlap == int(speech.shape[-1]):
return blended_head
return torch.cat([blended_head, speech[:, overlap:]], dim=-1)
@torch.inference_mode()
def forward_streaming(
self,
@@ -238,17 +213,16 @@ class CosyVoice3Code2Wav(nn.Module):
suffix. That preserves causal look-right handling without double
trimming or duplicated overlap at chunk boundaries.
"""
with nullcontext():
feat = self._forward_mel(
token=token,
prompt_token=prompt_token,
prompt_feat=prompt_feat,
embedding=embedding,
n_timesteps=n_timesteps,
token_offset_tokens=token_offset_tokens,
streaming=True,
finalize=finalize,
)
feat = self._forward_mel(
token=token,
prompt_token=prompt_token,
prompt_feat=prompt_feat,
embedding=embedding,
n_timesteps=n_timesteps,
token_offset_tokens=token_offset_tokens,
streaming=True,
finalize=finalize,
)
hift_weight = self.hift.m_source.l_linear.weight
chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype)
@@ -268,8 +242,7 @@ class CosyVoice3Code2Wav(nn.Module):
if tts_mel.shape[-1] == 0:
tts_speech = torch.zeros((chunk_mel.shape[0], 1, 0), device=chunk_mel.device, dtype=chunk_mel.dtype)
else:
with nullcontext():
tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize)
tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize)
tts_speech = tts_speech.reshape(tts_speech.shape[0], -1)
speech_offset = max(0, min(speech_offset, int(tts_speech.shape[-1])))
@@ -4,8 +4,10 @@
Stage 0: Talker — text prompt → speech tokens (LLM autoregressive).
Stage 1: Code2Wav — flow-matching decoder → acoustic features → waveform.
* ``sync_process_input_func`` runs when ``deploy.async_chunk=false``:
stage 1 builds full-sequence flow input via ``text2flow``.
* ``sync_process_input_func`` (``text2flow_token_only``) runs when
``deploy.async_chunk=false``: stage 1 builds full-sequence flow input
from the emitted speech tokens. ``text2flow`` is the
``custom_process_input_func`` variant.
* ``async_chunk_process_next_stage_input_func`` runs when
``deploy.async_chunk=true``: stage 0 streams codec chunks to stage 1
through the shared-memory connector.
@@ -13,8 +13,6 @@ from vllm_omni.utils.audio import mel_filter_bank
logger = logging.getLogger(__name__)
IGNORE_ID = -1
def dynamic_range_compression_torch(x, c=1, clip_val=1e-5):
return torch.log(torch.clamp(x, min=clip_val) * c)
@@ -182,26 +180,6 @@ def log_mel_spectrogram(
return log_spec
def extract_speech_token(prompt_wav, speech_tokenizer_session, device):
speech = load_wav(prompt_wav, 16000)
assert speech.shape[1] / 16000 <= 30, "do not support extract speech token for audio longer than 30s"
feat = log_mel_spectrogram(speech, n_mels=128)
speech_token = (
speech_tokenizer_session.run(
None,
{
speech_tokenizer_session.get_inputs()[0].name: feat.detach().cpu().numpy(),
speech_tokenizer_session.get_inputs()[1].name: np.array([feat.shape[2]], dtype=np.int32),
},
)[0]
.flatten()
.tolist()
)
speech_token = torch.tensor([speech_token], dtype=torch.int32).to(device)
speech_token_len = torch.tensor([speech_token.shape[1]], dtype=torch.int32).to(device)
return speech_token, speech_token_len
def extract_spk_embedding(prompt_wav, campplus_session, device):
speech = load_wav(prompt_wav, 16000)
feat = kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000)
@@ -67,6 +67,5 @@ class CovoAudioCode2WavConfig(PretrainedConfig):
self.inference = {
"s_steps": 10,
"cfg_alpha": 1.0,
"dynamic_the": False,
"rescale_logits": False,
}
@@ -598,7 +598,7 @@ class Token2latentFlowMatching(nn.Module):
latent_cond = torch.zeros(b, tgt_lens, self.target_dim).to(device)
if prefix_target is not None:
latent_cond[:, : prefix_target.size(1), :] = prefix_target
sample, trajectory = self.sample(
sample, _ = self.sample(
tokens=token, audio=latent_cond, steps=s_steps, alpha=cfg_alpha, g_cond=ge, rescale_logits=rescale_logits
)
return sample
@@ -612,9 +612,8 @@ class Token2latentFlowMatching(nn.Module):
output = torch.cat([tokens, audio, z], dim=-1)
return self.vectorfield_forward(inputs=output, times=t.unsqueeze(0), self_attn_mask=None, g_cond=g_cond)
tokens_empty = torch.zeros(*audio.shape[:2], self.model_dim, device=tokens.device, dtype=tokens.dtype)
audio_empty = audio
tokens_t = torch.cat([tokens_empty, tokens], dim=0)
audio_t = torch.cat([audio_empty, audio], dim=0)
audio_t = torch.cat([audio, audio], dim=0)
audio_noizy_t = torch.cat([z, z], dim=0)
t_t = torch.stack([t, t], dim=0)
c = g_cond
@@ -348,18 +348,16 @@ class FishSpeechDACDecoder(nn.Module):
if req_ids.numel() < 1:
continue
flat = req_ids
n = flat.numel()
if n == 0 or n % q != 0:
if n > 0:
logger.warning(
"DAC decoder input_ids length %d not divisible by num_codebooks %d; returning empty audio.",
n,
q,
)
n = req_ids.numel()
if n % q != 0:
logger.warning(
"DAC decoder input_ids length %d not divisible by num_codebooks %d; returning empty audio.",
n,
q,
)
continue
frames = n // q
codes_qf = flat.reshape(q, frames)
codes_qf = req_ids.reshape(q, frames)
parsed_ctx_frames[i] = ctx_frames
parsed_total_frames[i] = frames
valid_codes_qf.append(codes_qf)
@@ -191,7 +191,11 @@ class _FastARMLP(nn.Module):
class _FastARDecoderLayer(nn.Module):
"""Transformer decoder layer for Fast AR (SDPA, no KV cache)."""
"""Transformer decoder layer for Fast AR.
The batched ``forward`` runs SDPA over the full sequence and is cache-free,
while ``forward_one`` decodes a single step using a per-call KV cache.
"""
def __init__(self, config: FishSpeechFastARConfig, *, prefix: str = "") -> None:
super().__init__()
@@ -519,7 +523,7 @@ class FishSpeechFastAR(nn.Module):
embed_buf[:bsz, 0, :] = projected
# Position 1: embedding of semantic code.
code_embed = self.fast_embeddings(semantic_code.clamp(min=0))
code_embed = self.fast_embeddings(semantic_code)
embed_buf[:bsz, 1, :] = code_embed
use_sampling = do_sample and temperature > 0
@@ -542,8 +546,7 @@ class FishSpeechFastAR(nn.Module):
logits = self.fast_output(self.fast_norm(hidden_out))
# Residual codebooks (step >= 1) only have 1024 entries.
if step >= 1:
logits = logits[:, :residual_codebook_size]
logits = logits[:, :residual_codebook_size]
if use_sampling:
scaled = logits * inv_temperature
@@ -132,13 +132,3 @@ def estimate_fish_voice_clone_prompt_len_from_normalized(
[0] * semantic_len,
)
return len(prompt_ids)
def estimate_fish_voice_clone_prompt_len(tokenizer: Any, text: str, ref_text: str, semantic_len: int) -> int:
normalized_text, normalized_ref_text = normalize_fish_voice_clone_texts(text, ref_text)
return estimate_fish_voice_clone_prompt_len_from_normalized(
tokenizer,
normalized_text,
normalized_ref_text,
semantic_len,
)
@@ -472,10 +472,6 @@ class GLMTTSDiT(nn.Module):
self.num_heads = heads
self.spkr_emb_adaln = spkr_emb_adaln
# Frame rates for block-causal mask token→mel conversion
self._input_frame_rate: float = 25.0
self._mel_framerate: float = 50.0
# Speaker embedding dimension for AdaLN
if spkr_emb_adaln:
additional_dim = spkr_dim + 256 if use_wavlm_emb else spkr_dim
@@ -155,15 +155,3 @@ class HiggsAudioV2Config(PretrainedConfig):
@property
def num_real_codes(self) -> int:
return self.audio_stream_bos_id # i.e. 1024 for the boson-ai checkpoint
# The full vocab size of each per-codebook output head.
@property
def codebook_output_size(self) -> int:
return self.codebook_size # i.e. 1026 for the boson-ai checkpoint
# Canonical MusicGen-style delay pattern used by
# `HiggsAudioV2DelayPatternLogitsProcessor`: codebook k starts emitting
# real codes only after k frames. Kept as a function so subclasses can
# override if upstream ever publishes a different convention.
def default_delay_pattern(self) -> list[int]:
return list(range(self.num_codebooks))
@@ -821,12 +821,12 @@ class HiggsAudioV2TalkerForConditionalGeneration(nn.Module):
except Exception:
self._last_step_query_start_loc = None
# Voice-clone ref-audio substitution: at the prompt-side audio
# placeholder span we overwrite the text-embed lookups with the
# multi-codebook embedding of the encoded reference clip carried in
# ``model_intermediate_buffer``. No-op for plain-text requests and
# for decode steps (which fall through to the audio-feedback path).
if input_ids is not None and inputs_embeds is None:
# Voice-clone ref-audio substitution: at the prompt-side audio
# placeholder span we overwrite the text-embed lookups with the
# multi-codebook embedding of the encoded reference clip carried in
# ``model_intermediate_buffer``. No-op for plain-text requests and
# for decode steps (which fall through to the audio-feedback path).
info_dicts = kwargs.get("model_intermediate_buffer")
if info_dicts is None:
info_dicts = kwargs.get("runtime_additional_information")
@@ -836,15 +836,14 @@ class HiggsAudioV2TalkerForConditionalGeneration(nn.Module):
info_dicts,
)
# Audio-feedback substitution: vLLM's AR runner calls forward()
# with ``inputs_embeds=None`` and lets the model embed its own
# input_ids inline. ``embed_input_ids`` is only invoked via the
# supports_mm_inputs/prompt_embeds branches, neither of which fires
# for the talker. So we apply the substitution directly here: at
# decode positions where input_ids == audio_token_id, swap the text
# embedding with embed_audio_codes(last audio frame stashed in
# ``self._audio_state[req_id]`` by :meth:`sample`).
if input_ids is not None and inputs_embeds is None:
# Audio-feedback substitution: vLLM's AR runner calls forward()
# with ``inputs_embeds=None`` and lets the model embed its own
# input_ids inline. ``embed_input_ids`` is only invoked via the
# supports_mm_inputs/prompt_embeds branches, neither of which fires
# for the talker. So we apply the substitution directly here: at
# decode positions where input_ids == audio_token_id, swap the text
# embedding with embed_audio_codes(last audio frame stashed in
# ``self._audio_state[req_id]`` by :meth:`sample`).
hidden_states = self._maybe_apply_audio_feedback(hidden_states, input_ids)
audio_mask = self.audio_token_mask(input_ids) if input_ids is not None else None
@@ -1506,7 +1505,7 @@ class HiggsAudioV2TalkerForConditionalGeneration(nn.Module):
sampled[~valid] = torch.argmax(logits[~valid].float(), dim=-1)
return sampled
def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> int:
def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> None:
"""Mask non-audio tokens at audio-mode positions, in-place on ``logits``.
Heuristic: per-request, find the last token seen so far (last of
@@ -1524,7 +1523,7 @@ class HiggsAudioV2TalkerForConditionalGeneration(nn.Module):
the stock sampler.
"""
if logits is None or logits.ndim != 2:
return 0
return
audio_bos = int(self.config.audio_bos_token_id)
audio_id = int(self.config.audio_token_id)
# (PyTorch align, corrected): upstream NEVER lets the model
@@ -1542,7 +1541,6 @@ class HiggsAudioV2TalkerForConditionalGeneration(nn.Module):
prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None)
output_ids = getattr(sampling_metadata, "output_token_ids", None)
num_rows = int(logits.shape[0])
biased = 0
# Fallback "previous token" source: read each request's LAST input
# position from ``self._last_step_input_ids`` using ``query_start_loc``
# from the forward context. The previous version of this code took
@@ -1619,7 +1617,6 @@ class HiggsAudioV2TalkerForConditionalGeneration(nn.Module):
mask[audio_eos_vocab] = row[audio_eos_vocab]
logits[i].copy_(mask)
slot_state["should_terminate"] = False
biased += 1
continue
# Always force ``audio_token_id`` — upstream mirror.
allowed: set[int] = {audio_id, *allowed_extra}
@@ -1629,9 +1626,7 @@ class HiggsAudioV2TalkerForConditionalGeneration(nn.Module):
if 0 <= tok < row.shape[-1]:
mask[tok] = row[tok]
logits[i].copy_(mask)
biased += 1
self._last_first_audio_after_bos = first_after_bos
return biased
has_postprocess: bool = True
@@ -1,111 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Bounded LRU cache for Higgs Audio v3 reference-audio codec encoding.
Voice-clone requests that re-send the same ref_audio are otherwise re-encoded
through the DAC encoder on every call. The cache keys by a stable content
hash (size + sample rate + head/mid/tail sentinel bytes) so the same audio
served via different URLs or repeated base64 payloads also hits warm.
Tensors are cloned on get/put so callers and the cache cannot alias storage.
"""
from __future__ import annotations
import hashlib
import threading
from collections import OrderedDict
from collections.abc import Callable
import numpy as np
import torch
__all__ = ["cached_encode_reference_audio", "cache_stats", "clear_cache"]
_DEFAULT_MAX_ENTRIES = 128
_SENTINEL_BYTES = 4096
_cache: OrderedDict[str, torch.Tensor] = OrderedDict()
_cache_lock = threading.Lock()
_hits = 0
_misses = 0
def _content_hash(wav: np.ndarray, sr: int) -> str:
"""Stable content hash. Head/mid/tail sentinel keeps cost O(1)."""
raw = wav.tobytes()
n = len(raw)
h = hashlib.sha256()
h.update(n.to_bytes(8, "little"))
h.update(int(sr).to_bytes(4, "little"))
if n == 0:
return h.hexdigest()
if n <= _SENTINEL_BYTES * 3:
h.update(raw)
return h.hexdigest()
h.update(raw[:_SENTINEL_BYTES])
mid = (n - _SENTINEL_BYTES) // 2
h.update(raw[mid : mid + _SENTINEL_BYTES])
h.update(raw[-_SENTINEL_BYTES:])
return h.hexdigest()
def cached_encode_reference_audio(
wav: np.ndarray,
sr: int,
encode_fn: Callable[[np.ndarray, int], torch.Tensor],
max_entries: int = _DEFAULT_MAX_ENTRIES,
) -> torch.Tensor:
"""LRU-cached wrapper around ``encode_reference_audio``.
``encode_fn`` is the actual codec encoder; passed in so this module
stays decoupled from the tokenizer import (avoids circular imports
during stage-input-processor discovery).
Returns a CPU tensor clone — never the cache's interior storage.
"""
global _hits, _misses
key = _content_hash(wav, sr)
with _cache_lock:
cached = _cache.get(key)
if cached is not None:
_hits += 1
_cache.move_to_end(key)
return cached.clone()
# Miss: encode outside the lock so concurrent misses on different
# refs don't serialize on this single mutex.
codes = encode_fn(wav, sr)
codes_cpu = codes.detach().to("cpu").contiguous()
with _cache_lock:
if key in _cache:
# Concurrent miss winner already stored it; refresh LRU position
# and return that copy to avoid divergent tensors per request.
_cache.move_to_end(key)
_misses += 1 # we still paid the encode cost
return _cache[key].clone()
_cache[key] = codes_cpu
if len(_cache) > max_entries:
_cache.popitem(last=False)
_misses += 1
return codes_cpu.clone()
def cache_stats() -> dict:
"""Return (hits, misses, size, hit_rate) — for /metrics-style reporting."""
with _cache_lock:
total = _hits + _misses
return {
"hits": _hits,
"misses": _misses,
"size": len(_cache),
"hit_rate": (_hits / total) if total > 0 else 0.0,
}
def clear_cache() -> None:
"""Drop all cached entries. Used by tests."""
global _hits, _misses
with _cache_lock:
_cache.clear()
_hits = 0
_misses = 0
@@ -108,20 +108,6 @@ class AudioProjector(nn.Module):
layers.append(nn.Linear(llm_dim, llm_dim))
self.proj = nn.Sequential(*layers)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Project audio features with temporal downsampling.
Args:
x: [B, T, audio_dim] audio encoder output (channel-last).
Returns:
[B, T', llm_dim] projected features (channel-last),
where T' = (T - ds_kernel_size + 2*(ds_kernel_size//2)) // ds_stride + 1.
"""
# Conv1d expects [B, C, T], so transpose input
x = x.transpose(-1, -2) # [B, audio_dim, T]
return self.proj(x)
def forward_packed(
self,
packed: torch.Tensor,
@@ -155,19 +141,6 @@ class AudioProjector(nn.Module):
packed_proj = mlp(packed_proj)
return packed_proj, proj_lens
def compute_output_length(self, input_length: torch.Tensor) -> torch.Tensor:
"""Compute output sequence length after Conv1d downsampling.
Args:
input_length: Original mel spectrogram lengths.
Returns:
Output lengths after both convolutions.
"""
length = (input_length - 3 + 2 * 1) // 2 + 1
length = (length - self.ds_kernel_size + 2 * (self.ds_kernel_size // 2)) // self.ds_stride + 1
return length
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
@@ -15,19 +15,6 @@ from vllm.logger import init_logger
logger = init_logger(__name__)
# Ported from
# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/front/toolkit.py
_TOKENIZE_PATTERN = re.compile(r"(?:[a-zA-Z]\.)+|[a-zA-Z]+(?:['\-][a-zA-Z]+)*|\d+(?:\.\d+)?|[\u4e00-\u9fff]|\s+|\S")
def tokenize_mixed_text(text: str) -> list[str]:
return re.findall(_TOKENIZE_PATTERN, text)
def tokenize_mixed_text_iterator(text: str):
for match in _TOKENIZE_PATTERN.finditer(text):
yield match.group(0)
# Ported from
# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/front/text_segment_cut.py
@@ -218,84 +205,6 @@ def cut_text_by_semantic_length(
return [f.replace(_DOT_PLACEHOLDER, ".") for f in result_fragments]
# Streaming sentence boundary detection
_RE_CJK = re.compile(r"[\u4e00-\u9fff]")
_RE_DIGIT_LAST = re.compile(r"[0-9]")
# Left for reference for now
# Alternative to `cut_text_by_semantic_length`
def detect_sentence_boundaries(
text: str,
max_length: int = 50,
) -> list[str]:
"""Accumulate tokens and flush at sentence boundaries.
Ported from the streaming sentence detection loop from the Ming repo
TTS branch, but operates on the full text since we have it available upfront.
"""
sentences: list[str] = []
streaming_text: list[str] = []
count = 0
for ele in tokenize_mixed_text_iterator(text):
if len(ele) == 0:
continue
should_process = False
min_tokens = 12 if count == 0 else 8
if ele[-1] in "!?。,!?":
if len(streaming_text) >= min_tokens:
should_process = True
streaming_text.append(ele)
elif ele[-1] == ".":
if (
len(streaming_text) >= min_tokens
and streaming_text
and not _RE_DIGIT_LAST.search(streaming_text[-1][-1])
):
should_process = True
streaming_text.append(ele)
elif ele[-1] == "\n":
if streaming_text:
joined = "".join(streaming_text)
if _RE_CJK.search(joined):
if _RE_CJK.search(streaming_text[-1][-1]):
ele = ","
streaming_text.append(ele)
else:
if len(ele) > 1 and re.search(r"[a-zA-Z]", ele[-2]):
ele = ele[:-1] + "."
else:
ele = ele[:-1]
streaming_text.append(ele)
if len(streaming_text) >= min_tokens:
should_process = True
else:
streaming_text.append(ele)
continue
if should_process:
joined = "".join(streaming_text)
fragments = cut_text_by_semantic_length(joined, max_length)
sentences.extend(fragments)
streaming_text = []
count += 1
# Flush remaining
if streaming_text and re.search(r"[a-zA-Z\u4e00-\u9fff1-9]", "".join(streaming_text)):
joined = "".join(streaming_text)
fragments = cut_text_by_semantic_length(joined, max_length)
sentences.extend(fragments)
return sentences
# number normalization for English. Ported from
# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/front/number_en.py
@@ -30,7 +30,6 @@ from .constants import (
KEY_MAX_DECODE_STEPS,
KEY_MIN_DECODE_STEPS,
KEY_NEXT_EMBEDS,
KEY_PROMPT_LATENT_TAIL,
KEY_PROMPT_LATENTS,
KEY_REQUEST_ID,
KEY_SIGMA,
@@ -184,7 +183,6 @@ __all__ = [
"KEY_MAX_DECODE_STEPS",
"KEY_MIN_DECODE_STEPS",
"KEY_NEXT_EMBEDS",
"KEY_PROMPT_LATENT_TAIL",
"KEY_PROMPT_LATENTS",
"KEY_REQUEST_ID",
"KEY_SIGMA",
@@ -55,7 +55,6 @@ KEY_DECODE_STEP = "ming_decode_step"
KEY_LAST_STOP_PROB = "ming_last_stop_prob"
KEY_NEXT_EMBEDS = "ming_next_embeds"
KEY_PROMPT_LATENTS = "ming_prompt_latents"
KEY_PROMPT_LATENT_TAIL = "ming_prompt_latent_tail"
KEY_SPEAKER_EMBEDDING = "ming_speaker_embedding"
KEY_REQUEST_ID = "ming_request_id"
KEY_CHUNK_ID = "ming_chunk_id"
@@ -16,7 +16,6 @@ from vllm.v1.sample.sampler import Sampler
from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
from .audio_prep import (
_coerce_prompt_latents,
_find_audio_placeholder_positions,
_find_speaker_placeholder_positions,
_initial_history,
@@ -32,7 +31,6 @@ from .config_ming_tts import (
KEY_MAX_DECODE_STEPS,
KEY_MIN_DECODE_STEPS,
KEY_NEXT_EMBEDS,
KEY_PROMPT_LATENT_TAIL,
KEY_REQUEST_ID,
KEY_SIGMA,
KEY_SPEAKER_EMBEDDING,
@@ -175,7 +173,6 @@ class MingTTSForConditionalGeneration(nn.Module, SupportsPP, CustomProcessMixin)
dtype=torch.float32,
)
update[KEY_LATENT_HISTORY] = history.detach().to("cpu").contiguous()
update[KEY_PROMPT_LATENT_TAIL] = update[KEY_LATENT_HISTORY]
speaker_embedding = info_dict.get(KEY_SPEAKER_EMBEDDING, info_dict.get("speaker_embedding"))
speaker_embeddings = None
@@ -262,7 +259,6 @@ def _copy_runtime_controls(update: dict[str, Any], info_dict: dict[str, Any]) ->
__all__ = [
"MingTTSForConditionalGeneration",
"_ModelSampleAdapter",
"_coerce_prompt_latents",
"_find_audio_placeholder_positions",
"_find_speaker_placeholder_positions",
"_initial_history",
@@ -3,7 +3,6 @@
# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/modeling_bailingmm.py
from __future__ import annotations
import warnings
from collections.abc import Iterable
from typing import Any
@@ -346,10 +345,10 @@ class MingLLMModel(nn.Module):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
params_dict = dict(self.named_parameters(remove_duplicate=False))
loaded_params = AutoWeightsLoader(self).load_weights(weights, mapper=self.hf_to_vllm_mapper)
_warn_missing_prefix("flowloss", params_dict, loaded_params, prefix="flowloss.", fatal=True)
_warn_missing_prefix("linear_proj_audio", params_dict, loaded_params, prefix="linear_proj_audio.", fatal=True)
_warn_missing_prefix("stop_head", params_dict, loaded_params, prefix="stop_head.", fatal=True)
_warn_missing_prefix("spk_head", params_dict, loaded_params, prefix="spk_head.", fatal=True)
_warn_missing_prefix("flowloss", params_dict, loaded_params, prefix="flowloss.")
_warn_missing_prefix("linear_proj_audio", params_dict, loaded_params, prefix="linear_proj_audio.")
_warn_missing_prefix("stop_head", params_dict, loaded_params, prefix="stop_head.")
_warn_missing_prefix("spk_head", params_dict, loaded_params, prefix="spk_head.")
return loaded_params
@@ -383,7 +382,6 @@ def _warn_missing_prefix(
params_dict: dict[str, nn.Parameter],
loaded_params: set[str],
prefix: str,
fatal: bool = False,
) -> None:
missing = {key for key in params_dict if key.startswith(prefix)} - loaded_params
if not missing:
@@ -392,9 +390,7 @@ def _warn_missing_prefix(
f"MingLLMModel: {len(missing)} {module_name} params not loaded "
f"(prefix={prefix}). First few: {sorted(missing)[:5]}"
)
if fatal:
raise RuntimeError(msg)
warnings.warn(msg, stacklevel=3)
raise RuntimeError(msg)
__all__ = [
@@ -78,7 +78,7 @@ class MossAudioTokenizerConfig(PretrainedConfig):
@property
def codebook_size(self) -> int:
return int(self.quantizer_kwargs.get("codebook_size", 4096))
return int(self.quantizer_kwargs.get("codebook_size", 1024))
@property
def frame_rate(self) -> float:
@@ -31,8 +31,6 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
logger = init_logger(__name__)
_AUDIO_PAD_ID = 0 # padding id used before delay slot fires
def _maybe_prefix(prefix: str, name: str) -> str:
return f"{prefix}.{name}" if prefix else name
@@ -1058,7 +1056,7 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
history_per_codebook=hist_per_cb,
).squeeze(0) # (n_vq,)
if int(state.get("step", 0)) < 5 or int(state.get("step", 0)) % 50 == 0:
logger.info(
logger.debug(
"[MossTTSRealtime make_omni] step=%d ch0=%d cursor=%d/%d",
int(state.get("step", 0)),
int(new_codes[0].item()),
@@ -1,60 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Configuration for MOSS-TTS-Nano in vLLM-Omni single-stage pipeline."""
from transformers.configuration_utils import PretrainedConfig
class MossTTSNanoConfig(PretrainedConfig):
"""Config for MOSS-TTS-Nano (OpenMOSS-Team/MOSS-TTS-Nano).
The model is a Global-Local GPT2-based AR LM (0.1B) paired with the
MOSS-Audio-Tokenizer-Nano codec. Both components run in a single
vLLM-Omni generation stage.
Relevant fields from config.json:
gpt2_config – backbone config (n_layer=12, n_embd=768, n_head=12, ...)
n_vq – number of RVQ codebooks (16)
audio_vocab_size – per-codebook vocabulary size (1024)
audio_tokenizer_pretrained_name_or_path – HF hub path for codec model
"""
model_type = "moss_tts_nano"
def __init__(self, **kwargs):
gpt2_cfg = kwargs.pop("gpt2_config", None) or {}
if hasattr(gpt2_cfg, "to_dict"):
gpt2_cfg = gpt2_cfg.to_dict()
super().__init__(**kwargs)
# --- GPT2 backbone parameters (exposed at top level for vLLM) ---
self.hidden_size = gpt2_cfg.get("n_embd", 768)
self.num_hidden_layers = gpt2_cfg.get("n_layer", 12)
self.num_attention_heads = gpt2_cfg.get("n_head", 12)
self.num_key_value_heads = self.num_attention_heads # no GQA in GPT2
self.head_dim = self.hidden_size // self.num_attention_heads
self.vocab_size = gpt2_cfg.get("vocab_size", 16384)
self.max_position_embeddings = gpt2_cfg.get("n_positions", 32768)
self.intermediate_size = gpt2_cfg.get("n_inner", self.hidden_size * 4)
# --- Audio codec parameters ---
self.n_vq: int = getattr(self, "n_vq", 16)
self.audio_vocab_size: int = getattr(self, "audio_vocab_size", 1024)
self.audio_start_token_id: int = getattr(self, "audio_start_token_id", 6)
self.audio_end_token_id: int = getattr(self, "audio_end_token_id", 7)
self.audio_user_slot_token_id: int = getattr(self, "audio_user_slot_token_id", 8)
self.audio_assistant_slot_token_id: int = getattr(self, "audio_assistant_slot_token_id", 9)
self.audio_tokenizer_pretrained_name_or_path: str = getattr(
self,
"audio_tokenizer_pretrained_name_or_path",
"OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano",
)
self.audio_tokenizer_sample_rate: int = getattr(self, "audio_tokenizer_sample_rate", 48000)
# vLLM requires speculative_config to be absent or None
self.speculative_config = None
def get_text_config(self, **kwargs):
"""Return self so vLLM uses our top-level config."""
return self
@@ -172,7 +172,6 @@ class MossTTSNanoForGeneration(nn.Module):
has_preprocess = False
has_postprocess = False
enable_update_additional_information = True
inject_omni_request_id_into_runtime_info = True
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
super().__init__()
@@ -379,13 +379,12 @@ class OmniVoiceModel(
target_ids = torch.full((num_codebooks, target_len), mask_id, dtype=torch.long, device=device)
# Conditional: [text] [ref_audio?] [target_mask]
cond_audio_start = text_ids.shape[1]
if ref_audio_tokens is not None:
ref_tokens = ref_audio_tokens.to(device) # [8, T_ref]
cond_ids = torch.cat([text_ids, ref_tokens, target_ids], dim=1)
cond_audio_start = text_ids.shape[1]
else:
cond_ids = torch.cat([text_ids, target_ids], dim=1)
cond_audio_start = text_ids.shape[1]
cond_len = cond_ids.shape[1]
@@ -6,8 +6,9 @@ OmniVoice Generator (Stage 0) - Iterative unmasking with Qwen3 backbone.
Generates 8-codebook audio tokens from text via 32-step non-autoregressive
iterative masked prediction with classifier-free guidance.
Uses vLLM-Omni's DiffusionAttention for optimized full (bidirectional) attention
via FlashAttention/SageAttention/SDPA backends.
Uses full bidirectional attention computed directly with PyTorch SDPA
(torch.nn.functional.scaled_dot_product_attention); no auto-selected
FlashAttention/SageAttention/DiffusionAttention backend is used.
"""
from __future__ import annotations
@@ -237,7 +238,7 @@ def _gumbel_sample(logits: torch.Tensor, temperature: float, generator: torch.Ge
# ---------------------------------------------------------------------------
# Qwen3-style transformer blocks using DiffusionAttention
# Qwen3-style transformer blocks using PyTorch SDPA
# ---------------------------------------------------------------------------
@@ -256,7 +257,7 @@ class OmniVoiceRMSNorm(nn.Module):
class OmniVoiceAttention(nn.Module):
"""Qwen3-style GQA attention using DiffusionAttention backend."""
"""Qwen3-style GQA attention using PyTorch SDPA (full bidirectional)."""
def __init__(self, config: OmniVoiceConfig):
super().__init__()
@@ -315,18 +316,17 @@ class OmniVoiceAttention(nn.Module):
v = v.permute(0, 2, 1, 3)
# Convert [B, 1, S, S] bool mask to float mask for SDPA
# (0.0 where attend, -inf where masked)
sdpa_mask = None
if attention_mask is not None:
sdpa_mask = attention_mask.to(dtype=q.dtype)
sdpa_mask = sdpa_mask.masked_fill(~attention_mask, float("-inf"))
sdpa_mask = sdpa_mask.masked_fill(attention_mask, 0.0)
sdpa_mask = torch.zeros_like(attention_mask, dtype=q.dtype).masked_fill_(~attention_mask, float("-inf"))
out = F.scaled_dot_product_attention(
q,
k,
v,
attn_mask=sdpa_mask,
scale=1.0 / math.sqrt(self.head_dim),
scale=self.scale,
)
# Back to (batch, seq, heads * head_dim)
@@ -351,7 +351,7 @@ class OmniVoiceMLP(nn.Module):
class OmniVoiceTransformerBlock(nn.Module):
"""Single Qwen3 transformer block with DiffusionAttention."""
"""Single Qwen3 transformer block with PyTorch SDPA attention."""
def __init__(self, config: OmniVoiceConfig):
super().__init__()
@@ -643,10 +643,9 @@ class OmniVoiceGenerator(nn.Module):
- 32-step iterative unmasking with classifier-free guidance
Optimizations:
- DiffusionAttention (FlashAttn/SageAttn/SDPA auto-selected)
- TeaCache / Cache-DiT compatible (hook-based, non-intrusive)
- Full bidirectional attention via PyTorch SDPA (no auto-selected
FlashAttn/SageAttn/DiffusionAttention backend)
- regionally_compile() compatible for torch.compile on repeated blocks
- Sequence parallelism via SP hooks for multi-GPU
"""
# For regionally_compile() support
@@ -27,7 +27,7 @@ class Qwen3TTSSpeakerEncoderConfig(PretrainedConfig):
Args:
mel_dim (`int`, *optional*, defaults to 128):
The dimension of the input mel-spectrogram.
enc_dim (`int`, *optional*, defaults to 192):
enc_dim (`int`, *optional*, defaults to 1024):
The dimension of the final speaker embedding.
enc_channels (`list[int]`, *optional*, defaults to `[512, 512, 512, 512, 1536]`):
A list of output channels for each TDNN/SERes2Net layer in the encoder.
@@ -80,11 +80,11 @@ class Qwen3TTSTalkerCodePredictorConfig(PretrainedConfig):
Args:
vocab_size (`int`, *optional*, defaults to 151936):
vocab_size (`int`, *optional*, defaults to 2048):
Vocabulary size of the Qwen3TTSTalkerCodePredictor model.
Defines the number of different tokens that can be represented by the
`inputs_ids` passed when calling [`Qwen3TTSTalkerCodePredictorModel`]
hidden_size (`int`, *optional*, defaults to 4096):
hidden_size (`int`, *optional*, defaults to 1024):
Dimension of the hidden representations.
intermediate_size (`int`, *optional*, defaults to 22016):
Dimension of the MLP representations.
@@ -271,11 +271,11 @@ class Qwen3TTSTalkerConfig(PretrainedConfig):
Args:
vocab_size (`int`, *optional*, defaults to 151936):
vocab_size (`int`, *optional*, defaults to 3072):
Vocabulary size of the Qwen3TTSTalker model.
Defines the number of different tokens that can be represented by the
`inputs_ids` passed when calling [`Qwen3TTSTalkerModel`]
hidden_size (`int`, *optional*, defaults to 2048):
hidden_size (`int`, *optional*, defaults to 1024):
Dimension of the hidden representations.
intermediate_size (`int`, *optional*, defaults to 6144):
Dimension of the MLP representations.
@@ -436,7 +436,6 @@ class Qwen3TTSTalkerConfig(PretrainedConfig):
self.rope_scaling["rope_type"] = self.rope_scaling["type"]
if code_predictor_config is None:
code_predictor_config = {}
self.code_predictor_config = Qwen3TTSTalkerCodePredictorConfig()
logger.info("code_predictor_config is None. Initializing code_predictor model with default values")
elif isinstance(code_predictor_config, Qwen3TTSTalkerCodePredictorConfig):
@@ -1320,9 +1320,7 @@ class Qwen3TTSPromptEmbedsBuilder:
# Keep it at least 1D; embedding on a 0-d tensor can return 1D.
spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long)
spk_embed = codec_embed(spk_tensor)
if spk_embed.ndim == 1:
spk_embed = spk_embed.view(1, 1, -1)
elif spk_embed.ndim == 2:
if spk_embed.ndim in (1, 2):
spk_embed = spk_embed.view(1, 1, -1)
speaker_embed = spk_embed
codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1)
@@ -1523,9 +1521,7 @@ class Qwen3TTSPromptEmbedsBuilder:
ref_code_len: int | None = None
if isinstance(ref_code, list):
if ref_code and isinstance(ref_code[0], list):
ref_code_len = len(ref_code)
elif ref_code:
if ref_code:
ref_code_len = len(ref_code)
elif hasattr(ref_code, "shape"):
try:
@@ -286,8 +286,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
)
self._codec_eos_token_id = int(getattr(self.talker_config, "codec_eos_token_id", -1))
self._eos_logit_bias: float = 0.0
self.have_multimodal_outputs = True
self.has_preprocess = True
self.has_postprocess = True
@@ -517,11 +515,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
# Mask out invalid codec ids using the pre-built constant buffer.
logits = logits.masked_fill(self._codec_disallowed_mask, float("-inf"))
if self._eos_logit_bias != 0.0:
eos_id = self._codec_eos_token_id
if 0 <= eos_id < logits.shape[-1]:
logits[:, eos_id] = logits[:, eos_id] + self._eos_logit_bias
return logits
# -------------------- Omni multimodal output plumbing --------------------
@@ -912,11 +912,6 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
self._cudagraph_wrapper.compile_shapes,
)
def disable_cudagraph(self):
self._cudagraph_enabled = False
self._cudagraph_wrapper = None
logger.info("CUDA Graph disabled for decoder")
def forward(self, codes):
if codes.shape[1] != self.config.num_quantizers:
raise ValueError(f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}")
@@ -33,7 +33,6 @@
import random
import typing as tp
from math import ceil
import numpy as np
import torch
@@ -45,10 +44,6 @@ from vllm.logger import init_logger
logger = init_logger(__name__)
def round_up_multiple(num, mult):
return ceil(num / mult) * mult
def default(val: tp.Any, d: tp.Any) -> tp.Any:
return val if val is not None else d
@@ -26,17 +26,17 @@ class StepAudio2ForConditionalGeneration(nn.Module, SupportsMultiModal, Supports
Step-Audio2 Main Controller
Manages two-stage inference pipeline:
- Stage 1 (Thinker): Audio understanding and token generation
- Stage 2 (Token2Wav): Audio token to waveform synthesis
- Stage 0 (Thinker): Audio understanding and token generation
- Stage 1 (Token2Wav): Audio token to waveform synthesis
Usage:
# Stage 1: Thinker
# Stage 0: Thinker
model = StepAudio2ForConditionalGeneration(
vllm_config=config,
model_stage="thinker"
)
# Stage 2: Token2Wav
# Stage 1: Token2Wav
model = StepAudio2ForConditionalGeneration(
vllm_config=config,
model_stage="token2wav"
@@ -42,39 +42,6 @@ class StepAudio2TokenConfig:
return self.audio_start + self.audio_vocab_size - 1
@dataclass(frozen=True)
class StepAudio2EncoderConfig:
"""
Step-Audio2 audio encoder architecture configuration.
These values define the encoder structure and should match your model weights.
"""
# Mel spectrogram parameters
n_mels: int = 128
"""Number of mel frequency bins"""
# Encoder architecture
n_audio_ctx: int = 1500
"""Audio context length (max sequence length)"""
n_audio_state: int = 512
"""Audio encoder hidden state dimension"""
n_audio_head: int = 8
"""Number of attention heads in encoder"""
n_audio_layer: int = 6
"""Number of encoder layers"""
# Adapter parameters
kernel_size: int = 3
"""Convolution kernel size for adapter"""
adapter_stride: int = 2
"""Stride for adapter convolution (downsampling factor)"""
@dataclass(frozen=True)
class StepAudio2Token2WavConfig:
"""Token2Wav synthesis configuration."""
@@ -141,7 +108,6 @@ class StepAudio2ModelConfig:
# Default configuration instances
DEFAULT_TOKEN_CONFIG = StepAudio2TokenConfig()
DEFAULT_ENCODER_CONFIG = StepAudio2EncoderConfig()
DEFAULT_TOKEN2WAV_CONFIG = StepAudio2Token2WavConfig()
DEFAULT_STREAM_CONFIG = StepAudio2StreamConfig()
DEFAULT_MODEL_CONFIG = StepAudio2ModelConfig()
@@ -162,12 +128,10 @@ STEP_AUDIO2_DEFAULT_PROMPT_WAV = "default_female.wav"
__all__ = [
"StepAudio2TokenConfig",
"StepAudio2EncoderConfig",
"StepAudio2Token2WavConfig",
"StepAudio2StreamConfig",
"StepAudio2ModelConfig",
"DEFAULT_TOKEN_CONFIG",
"DEFAULT_ENCODER_CONFIG",
"DEFAULT_TOKEN2WAV_CONFIG",
"DEFAULT_STREAM_CONFIG",
"DEFAULT_MODEL_CONFIG",
@@ -609,7 +609,7 @@ class StepAudio2ThinkerForConditionalGeneration(nn.Module, SupportsMultiModal, S
Architecture:
AudioEncoder (6 layers, 512 hidden) →
Adaptor (512 → LLM dim) →
Qwen2 LLM (vocab=64012)
Qwen2 LLM (vocab=158720)
"""
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -868,31 +868,3 @@ class StepAudio2ThinkerForConditionalGeneration(nn.Module, SupportsMultiModal, S
def has_audio_output(token_ids: list[int], audio_start: int = DEFAULT_TOKEN_CONFIG.audio_start) -> bool:
"""Check if generated tokens contain audio tokens"""
return any(tid >= audio_start for tid in token_ids)
class StepAudio2OutputProcessor:
"""Helper class to process Step-Audio2 outputs"""
@staticmethod
def process_output(output_ids: torch.Tensor, tokenizer, remove_audio_padding: bool = True) -> dict:
"""Process model output and separate text/audio tokens"""
if isinstance(output_ids, torch.Tensor):
output_ids = output_ids.squeeze().tolist()
# Separate tokens
text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(output_ids)
# Remove audio padding if requested
if remove_audio_padding and audio_tokens:
audio_tokens = [t for t in audio_tokens if t < DEFAULT_TOKEN_CONFIG.audio_eos]
# Decode text
text = tokenizer.decode(text_tokens, skip_special_tokens=False)
return {
"text": text,
"text_tokens": text_tokens,
"audio_tokens": audio_tokens,
"has_audio": len(audio_tokens) > 0,
"all_tokens": output_ids,
}
@@ -94,8 +94,6 @@ class StepAudio2Token2WavCore(nn.Module):
self.mel_cache_len = DEFAULT_STREAM_CONFIG.mel_cache_len
self.source_cache_len = STREAM_SOURCE_CACHE_LEN
self.speech_window: torch.Tensor | None = None # created lazily on device
self.stream_cache: dict | None = None # flow model causal cache
self.hift_cache_dict: dict = {} # HiFT vocoder overlap cache
def _ensure_models_loaded(self):
"""Lazy load models on first use"""
@@ -253,125 +251,6 @@ class StepAudio2Token2WavCore(nn.Module):
else:
return wav
# ------------------------------------------------------------------
# Streaming (chunked) inference
# ------------------------------------------------------------------
def setup_stream(self, prompt_wav: str) -> None:
"""Initialize flow + HiFT caches for chunked streaming inference.
Must be called once before the first ``stream_chunk()`` call for a
given utterance.
"""
if prompt_wav not in self.cache:
self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
(
prompt_speech_tokens,
_,
spk_emb,
prompt_mels,
_,
) = self.cache[prompt_wav]
# Lazily create the Hamming window on the correct device
if self.speech_window is None or self.speech_window.device != self.device:
self.speech_window = torch.from_numpy(np.hamming(2 * self.source_cache_len)).to(
device=self.device, dtype=torch.float32
)
# Flow model: setup causal cache with prompt tokens + lookahead
pre_lookahead = DEFAULT_STREAM_CONFIG.pre_lookahead_len
self.stream_cache = self.flow.setup_cache(
torch.cat([prompt_speech_tokens, prompt_speech_tokens[:, :pre_lookahead]], dim=1),
prompt_mels,
spk_emb,
n_timesteps=self.n_timesteps,
)
# HiFT vocoder: empty overlap buffers
self.hift_cache_dict = {
"mel": torch.zeros(1, prompt_mels.shape[2], 0, device=self.device),
"source": torch.zeros(1, 1, 0, device=self.device),
"speech": torch.zeros(1, 0, device=self.device),
}
def stream_chunk(
self,
audio_tokens: list[int],
prompt_wav: str,
last_chunk: bool = False,
) -> torch.Tensor:
"""Process one chunk of audio tokens and return a waveform segment.
Args:
audio_tokens: Token IDs for this chunk (includes lookahead for
non-last chunks).
prompt_wav: Path to the prompt wav (for cached speaker embeddings).
last_chunk: Whether this is the final chunk of the utterance.
Returns:
1-D float tensor of audio samples (24 kHz).
"""
if self.stream_cache is None:
raise ValueError("stream_cache not initialised – call setup_stream() first")
if prompt_wav not in self.cache:
self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
_, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav]
token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device)
with torch.amp.autocast(
str(self.device.type),
dtype=torch.float16 if self.float16 else torch.float32,
):
chunk_mel, self.stream_cache = self.flow.inference_chunk(
token=token_tensor,
spk=spk_emb,
cache=self.stream_cache,
last_chunk=last_chunk,
n_timesteps=self.n_timesteps,
)
# Trim estimator attention cache to avoid unbounded growth
keep = DEFAULT_STREAM_CONFIG.estimator_cache_keep
est_att = self.stream_cache.get("estimator_att_cache")
if est_att is not None and est_att.shape[4] > (prompt_mels.shape[1] + keep):
self.stream_cache["estimator_att_cache"] = torch.cat(
[est_att[:, :, :, :, : prompt_mels.shape[1]], est_att[:, :, :, :, -keep:]],
dim=4,
)
# ---- HiFT vocoder with overlap-add ----
hift_cache_mel = self.hift_cache_dict["mel"]
hift_cache_source = self.hift_cache_dict["source"]
hift_cache_speech = self.hift_cache_dict["speech"]
mel = torch.cat([hift_cache_mel, chunk_mel], dim=2)
speech, source = self.hift(mel, hift_cache_source)
# Cross-fade overlap region for smooth concatenation
if hift_cache_speech.shape[-1] > 0:
speech = fade_in_out(speech, hift_cache_speech, self.speech_window)
# Update vocoder cache for the next chunk
self.hift_cache_dict = {
"mel": mel[..., -self.mel_cache_len :].clone().detach(),
"source": source[:, :, -self.source_cache_len :].clone().detach(),
"speech": speech[:, -self.source_cache_len :].clone().detach(),
}
# For non-last chunks, trim the tail that will be re-synthesised
if not last_chunk:
speech = speech[:, : -self.source_cache_len]
return speech.squeeze(0) # [samples]
def reset_stream(self) -> None:
"""Clear all streaming state after an utterance is complete."""
self.stream_cache = None
self.hift_cache_dict = {}
# ------------------------------------------------------------------
# Per-request streaming (does NOT mutate self – uses external state)
# ------------------------------------------------------------------
@@ -65,7 +65,6 @@ class _PagedMiniCPM4Attention(nn.Module):
prefix: str = "",
) -> None:
super().__init__()
self.layer_idx = layer_idx
self.hidden_size = hidden_size
self.num_heads = num_attention_heads
self.head_dim = kv_channels if kv_channels else hidden_size // num_attention_heads
@@ -318,18 +317,12 @@ class MiniCPM4PagedForVoxCPM2(nn.Module):
).detach()
def compile_selective(self) -> list[str]:
"""Compile the full model forward as one graph.
"""Compile the full model ``forward`` as one graph.
Earlier versions compiled ``layer.mlp`` + ``layer.self_attn.o_proj``
(PR #2690) and then the whole ``layer`` (perf/voxcpm2-streaming-vae).
Both still paid one Dynamo dispatch per layer per decode step.
V3 profiling showed 1,332 per-layer dispatches (~28 layers × ~47
decode steps) costing ~726 ms of CPU self-time for a long prompt.
Compiling ``forward`` at the model level lets Dynamo unroll the
28-layer Python loop inside the graph. Graph breaks at
PagedAttention produce sub-graphs but Dynamo memoises the whole
trace once, so the per-step dispatch drops from 28 to just a few.
``torch.compile`` is applied at the model level so Dynamo unrolls the
per-layer Python loop inside the graph. Graph breaks at PagedAttention
produce sub-graphs, but Dynamo memoises the whole trace once, so the
per-decode-step dispatch drops to just a few instead of one per layer.
"""
if self._compiled_layers:
return []
@@ -312,7 +312,6 @@ class _RequestState:
decode_step_count: int = 0
request_start_time: float = 0.0
prefill_completed: bool = False
prefill_text: str = ""
prompt_cache: dict | None = None
prefill_masks: tuple | None = None
is_stopping: bool = False
@@ -2510,7 +2509,6 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
token_ids = token_ids[1:]
state = self._get_or_create_state(req_id)
state.prefill_text = ""
state.decode_pad = None
state.prefill_completed = False
state.decode_step_count = 0
@@ -2572,29 +2570,26 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
)
if state.prompt_cache is None and (ref_audio or (prompt_audio and prompt_text)):
if state.prompt_cache is None:
try:
state.prompt_cache = self._build_prompt_cache(
ref_audio=ref_audio,
prompt_audio=prompt_audio,
prompt_text=prompt_text,
try:
state.prompt_cache = self._build_prompt_cache(
ref_audio=ref_audio,
prompt_audio=prompt_audio,
prompt_text=prompt_text,
)
if (
voice_name
and state.prompt_cache is not None
and state.prompt_cache.get("mode") == "reference"
and "ref_audio_feat" in state.prompt_cache
):
_key = self._speaker_cache.make_cache_key(
voice_name, model_type="voxcpm2", created_at=_created_at
)
if (
voice_name
and state.prompt_cache is not None
and state.prompt_cache.get("mode") == "reference"
and "ref_audio_feat" in state.prompt_cache
):
_key = self._speaker_cache.make_cache_key(
voice_name, model_type="voxcpm2", created_at=_created_at
)
self._speaker_cache.put(
_key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()}
)
logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name)
except Exception as e:
logger.warning("build_prompt_cache failed: %s; falling back to zero-shot", e)
state.prompt_cache = None
self._speaker_cache.put(_key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()})
logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name)
except Exception as e:
logger.warning("build_prompt_cache failed: %s; falling back to zero-shot", e)
state.prompt_cache = None
inputs = self._build_prefill_inputs(token_ids, dev, req_id)
tts = self.tts
@@ -244,9 +244,6 @@ class VoxtralTTSForConditionalGeneration(
# Delegate to generation model for multimodal processing
return self.model.embed_multimodal(**kwargs)
def last_index_of(self, list, value):
return len(list) - 1 - list[::-1].index(value)
def forward(
self,
input_ids: torch.Tensor,
@@ -246,7 +246,6 @@ class BidirectionalAttention(nn.Module):
self.n_local_heads: int = args.n_heads
self.n_local_kv_heads: int = args.n_kv_heads
self.repeats = self.n_local_heads
self.layer_id = layer_id
self.head_dim = args.head_dim
@@ -276,7 +275,6 @@ class BidirectionalAttention(nn.Module):
bias=args.use_biases,
)
self.softmax_scale: float = self.args.head_dim**-0.5
self.repeats = self.n_local_heads // self.n_local_kv_heads
def _native_attention(
@@ -420,13 +418,6 @@ class FlowMatchingAudioTransformer(nn.Module):
self.acoustic_transformer_args = args
assert isinstance(self.acoustic_transformer_args, AcousticTransformerArgs)
# currently assuming always 1 semantic codebook + N acoustic codebooks
self.num_non_acoustic_embeddings = 1
self.num_acoustic_codebooks = len(self.model_args.get_codebook_sizes()) - self.num_non_acoustic_embeddings
# flow matching utils
self.sigma = args.sigma
# codebook sizes
acoustic_codebook_sizes = self.model_args.get_codebook_sizes(
pad_to_multiple=None, include_special_tokens=False
@@ -715,7 +706,6 @@ class VoxtralTTSProcessorAdapter:
)
text_tokens_for_audio = list[torch.Tensor]()
assert audios is not None
audio_tokens_pt = list[torch.Tensor]()
for audio_token_array in audio_tokens:
assert isinstance(audio_token_array, np.ndarray)
@@ -1008,7 +998,7 @@ class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMulti
if audio_arrays is not None:
if not isinstance(audio_arrays, (torch.Tensor, list)):
raise ValueError(f"Incorrect type of images. Got type: {type(audio_arrays)}")
raise ValueError(f"Incorrect type of audio_arrays. Got type: {type(audio_arrays)}")
if isinstance(audio_arrays, torch.Tensor) and audio_arrays.dim() == 3:
audio_arrays = flatten_bn(audio_arrays)
if isinstance(audio_arrays, torch.Tensor):
@@ -1053,24 +1043,11 @@ class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMulti
fake_logits[~is_eos, self.audio_tok_id] = 1.0
return fake_logits
# TODO(chenyo): Remove this
def compute_logits(
self,
hidden_states: torch.Tensor,
) -> torch.Tensor | None:
text_logits = self.language_model.compute_logits(
hidden_states,
)
assert text_logits is not None
return text_logits
def compute_mm_logits(
self,
hidden_states: torch.Tensor,
cfg_alpha: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor | None]:
audio_codes = None
mm_tokens = None
audio_codes = self.acoustic_transformer(
llm_hidden=hidden_states,
cfg_alpha=cfg_alpha,
@@ -1088,7 +1065,6 @@ class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMulti
return fake_eos, mm_tokens
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
# fmt: on
remapping_rules = [
(r"^acoustic_transformer\.(.*)$", r"\1"), # noqa: E501
(r"^audio_tokenizer\.(.*)$", r"\1"), # noqa: E501
@@ -652,7 +652,7 @@ class TransformerBlock(nn.Module):
def layer_id(self) -> int:
return self._layer_id
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
def forward(self, x: torch.Tensor) -> torch.Tensor:
r = self.attention.forward(self.attention_norm(x))
if self.post_attention_norm is not None:
r = self.post_attention_norm(r)
@@ -757,7 +757,7 @@ class VoxtralTTSAudioTokenizer(nn.Module):
# transformer
layer_args = deepcopy(args)
layer_args.attn_sliding_window_size = cur_window_size
assert layer_args.qk_norm, "qk_norm must be True for decoder"
assert layer_args.qk_norm, "qk_norm must be True for encoder"
encoder_transformer = Transformer(
args=layer_args,
n_layers=n_layers,
@@ -143,7 +143,6 @@ class MingWhisperFeatureExtractor(FeatureExtractionMixin):
# (conv1 has stride=1 and does not change T)
# 2. AudioProjector Conv1d: kernel=3, stride=2, padding=1
# Combined: T → ((T-1)//2 + 1 - 1)//2 + 1
# See also: AudioProjector.compute_output_length()
encoder_feats_lengths = ((audio_feats_lengths - 3 + 2 * 1) // 2 + 1 - 3 + 2 * 1) // 2 + 1
audio_feats = torch.cat(audio_feat_list, dim=0).unsqueeze(0) # [1, T_total, n_mels]