[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:
1 parent
f05e03aea8
commit
208932210f
46 files changed
+141
-787
No files matched your search
@@ -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
|
||||
|
||||
|
||||
@@ -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` | ✅︎ | ✅︎ | | |
|
||||
|
||||
@@ -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
@@ -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 --------------------
|
||||
|
||||
-5
@@ -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]
|
||||
|
||||
|
||||
Reference in new issue
Block a user