diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index acd3e87c1..98c746e09 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -35,6 +35,7 @@ th { |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | +|`CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | |`MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | |`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | diff --git a/pyproject.toml b/pyproject.toml index 95e43d871..99194c648 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -203,3 +203,5 @@ extend-ignore-identifiers-re = [ ".*[Oo]no_[Aa]nna.*", ".*cann.*", ] +[tool.typos.default.extend-words] +ue = "ue" diff --git a/requirements/common.txt b/requirements/common.txt index 7e5f5f9b2..548f8f0bf 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -12,6 +12,8 @@ torchsde>=0.2.6 openai-whisper>=20250625 imageio[ffmpeg]>=2.37.2 sox>=1.5.0 +x-transformers>=2.12.2 +einops>=0.8.1 prettytable>=3.8.0 aenum==3.1.16 pyzmq>=25.0.0 diff --git a/tests/model_executor/test_cosyvoice3_components.py b/tests/model_executor/test_cosyvoice3_components.py new file mode 100644 index 000000000..3b1471365 --- /dev/null +++ b/tests/model_executor/test_cosyvoice3_components.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for CosyVoice3 components.""" + +import pytest +import torch +import torch.nn as nn + + +class TestPreLookaheadLayer: + """Tests for PreLookaheadLayer.""" + + @pytest.fixture + def layer(self): + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer + + return PreLookaheadLayer(in_channels=512, channels=512, pre_lookahead_len=3) + + def test_forward_shape(self, layer): + """Test that output shape matches input shape.""" + batch, seq_len, channels = 2, 10, 512 + x = torch.randn(batch, seq_len, channels) + + out = layer(x) + + assert out.shape == x.shape + + def test_forward_with_context(self, layer): + """Test forward with context for streaming.""" + batch, seq_len, channels = 1, 10, 512 + x = torch.randn(batch, seq_len, channels) + context = torch.randn(batch, 3, channels) # pre_lookahead_len=3 + + layer.eval() + out = layer(x, context=context) + + assert out.shape == x.shape + + def test_residual_connection(self, layer): + """Test that residual connection is applied.""" + batch, seq_len, channels = 1, 5, 512 + x = torch.zeros(batch, seq_len, channels) + + # With zero input, output should also be close to zero due to residual + out = layer(x) + + # Output should be close to input (residual) plus conv output + assert out.shape == x.shape + + +class TestDiTAttention: + """Tests for DiTAttention with diffusion backend.""" + + @pytest.fixture + def attention(self): + from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTAttention + + return DiTAttention(dim=512, heads=8, dim_head=64, dropout=0.0) + + def test_forward_shape(self, attention): + """Test attention output shape.""" + batch, seq_len, dim = 2, 16, 512 + x = torch.randn(batch, seq_len, dim) + + out = attention(x) + + assert out.shape == x.shape + + def test_forward_with_mask(self, attention): + """Test attention with mask.""" + batch, seq_len, dim = 2, 16, 512 + x = torch.randn(batch, seq_len, dim) + mask = torch.ones(batch, seq_len, dtype=torch.bool) + mask[:, -3:] = False # Mask last 3 positions + + out = attention(x, mask=mask) + + assert out.shape == x.shape + # Masked positions should be zero + assert torch.allclose(out[:, -3:], torch.zeros_like(out[:, -3:])) + + def test_qkv_projections(self, attention): + """Test that Q/K/V projections exist and have correct dimensions.""" + assert hasattr(attention, "to_q") + assert hasattr(attention, "to_k") + assert hasattr(attention, "to_v") + assert attention.to_q.out_features == 512 # heads * dim_head + assert attention.to_k.out_features == 512 + assert attention.to_v.out_features == 512 + + +class TestDiTBlock: + """Tests for DiTBlock.""" + + @pytest.fixture + def block(self): + from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTBlock + + return DiTBlock(dim=512, heads=8, dim_head=64, ff_mult=4, dropout=0.0) + + def test_forward_shape(self, block): + """Test block output shape.""" + batch, seq_len, dim = 2, 16, 512 + x = torch.randn(batch, seq_len, dim) + t = torch.randn(batch, dim) # Timestep embedding + + out = block(x, t) + + assert out.shape == x.shape + + def test_adalayernorm_modulation(self, block): + """Test that AdaLayerNorm modulates based on timestep.""" + batch, seq_len, dim = 1, 8, 512 + x = torch.randn(batch, seq_len, dim) + t1 = torch.zeros(batch, dim) + t2 = torch.ones(batch, dim) + + out1 = block(x, t1) + out2 = block(x, t2) + + # Different timesteps should produce different outputs + assert not torch.allclose(out1, out2) + + +class TestDiT: + """Tests for the full DiT model.""" + + @pytest.fixture + def dit(self): + from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT + + return DiT( + dim=256, + depth=2, + heads=4, + dim_head=64, + dropout=0.0, + ff_mult=2, + mel_dim=80, + mu_dim=80, + spk_dim=80, + long_skip_connection=True, + ) + + def test_forward_shape(self, dit): + """Test DiT forward output shape.""" + batch, mel_dim, seq_len = 1, 80, 32 + x = torch.randn(batch, mel_dim, seq_len) + mask = torch.ones(batch, 1, seq_len) + mu = torch.randn(batch, mel_dim, seq_len) + t = torch.tensor([0.5]) + spks = torch.randn(batch, 80) + cond = torch.randn(batch, mel_dim, seq_len) + + out = dit(x, mask, mu, t, spks=spks, cond=cond) + + assert out.shape == (batch, mel_dim, seq_len) + + def test_timestep_embedding(self, dit): + """Test that different timesteps produce different outputs.""" + batch, mel_dim, seq_len = 1, 80, 16 + x = torch.randn(batch, mel_dim, seq_len) + mask = torch.ones(batch, 1, seq_len) + mu = torch.randn(batch, mel_dim, seq_len) + spks = torch.randn(batch, 80) + cond = torch.randn(batch, mel_dim, seq_len) + + out1 = dit(x, mask, mu, torch.tensor([0.0]), spks=spks, cond=cond) + out2 = dit(x, mask, mu, torch.tensor([1.0]), spks=spks, cond=cond) + + assert not torch.allclose(out1, out2) + + +class TestCFM: + """Tests for Conditional Flow Matching classes.""" + + @pytest.fixture + def dummy_estimator(self): + """Create a dummy estimator for testing.""" + + class DummyEstimator(nn.Module): + def __init__(self, mel_dim=80): + super().__init__() + self.mel_dim = mel_dim + + def forward(self, x, mask, mu, t, spks=None, cond=None): + return torch.zeros_like(x) + + return DummyEstimator() + + def test_causal_conditional_cfm_forward(self, dummy_estimator): + """Test CausalConditionalCFM forward pass.""" + from omegaconf import DictConfig + + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import CausalConditionalCFM + + cfm_params = DictConfig( + { + "sigma_min": 1e-6, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + } + ) + + cfm = CausalConditionalCFM( + in_channels=80, + cfm_params=cfm_params, + n_spks=1, + spk_emb_dim=80, + estimator=dummy_estimator, + ) + + batch, mel_dim, seq_len = 1, 80, 32 + mu = torch.randn(batch, mel_dim, seq_len) + mask = torch.ones(batch, 1, seq_len) + spks = torch.randn(batch, 80) + cond = torch.randn(batch, mel_dim, seq_len) + + out, _ = cfm(mu, mask, n_timesteps=2, spks=spks, cond=cond) + + assert out.shape == mu.shape + + +class TestSDPAFallback: + """Test SDPA fallback for float32 inputs.""" + + def test_float32_uses_sdpa(self): + """Test that float32 inputs use SDPA fallback.""" + from vllm_omni.diffusion.attention.layer import Attention + + attn = Attention( + num_heads=8, + head_size=64, + causal=False, + softmax_scale=1.0 / 8.0, + ) + + batch, seq_len, heads, dim = 1, 16, 8, 64 + q = torch.randn(batch, seq_len, heads, dim, dtype=torch.float32) + k = torch.randn(batch, seq_len, heads, dim, dtype=torch.float32) + v = torch.randn(batch, seq_len, heads, dim, dtype=torch.float32) + + # Should not raise error - SDPA fallback handles float32 + out = attn(q, k, v) + + assert out.shape == (batch, seq_len, heads, dim) + assert out.dtype == torch.float32 diff --git a/tests/model_executor/test_cosyvoice3_utils.py b/tests/model_executor/test_cosyvoice3_utils.py new file mode 100644 index 000000000..828bb2b14 --- /dev/null +++ b/tests/model_executor/test_cosyvoice3_utils.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for CosyVoice3 utility functions.""" + +import pytest +import torch + + +class TestMakePadMask: + """Tests for make_pad_mask utility.""" + + def test_basic_mask_creation(self): + """Test basic padding mask creation.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([5, 3, 2]) + mask = make_pad_mask(lengths) + + # Expected: [[F,F,F,F,F], [F,F,F,T,T], [F,F,T,T,T]] + expected = torch.tensor( + [ + [False, False, False, False, False], + [False, False, False, True, True], + [False, False, True, True, True], + ] + ) + assert torch.equal(mask, expected) + + def test_single_sequence(self): + """Test mask for single sequence.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([3]) + mask = make_pad_mask(lengths) + + expected = torch.tensor([[False, False, False]]) + assert torch.equal(mask, expected) + + def test_explicit_max_len(self): + """Test with explicit max_len parameter.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([2, 3]) + mask = make_pad_mask(lengths, max_len=5) + + expected = torch.tensor( + [ + [False, False, True, True, True], + [False, False, False, True, True], + ] + ) + assert torch.equal(mask, expected) + + def test_zero_max_len_uses_max_length(self): + """Test that max_len=0 falls back to max(lengths).""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([2, 4]) + mask = make_pad_mask(lengths, max_len=0) + + # Should use max(lengths) = 4 + assert mask.shape == (2, 4) + + def test_full_length_sequence(self): + """Test sequence that fills the entire length (no padding).""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([5]) + mask = make_pad_mask(lengths, max_len=5) + + # No positions should be masked + assert not mask.any() + + def test_empty_sequence(self): + """Test sequence with zero length (all padding).""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([0]) + mask = make_pad_mask(lengths, max_len=3) + + # All positions should be masked + assert mask.all() + + def test_device_preservation(self): + """Test that mask is created on the same device as input.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + + lengths = torch.tensor([3, 2]) + mask = make_pad_mask(lengths) + + assert mask.device == lengths.device + + +class TestDynamicRangeCompression: + """Tests for dynamic_range_compression_torch utility.""" + + def test_basic_compression(self): + """Test basic compression with default parameters.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import dynamic_range_compression_torch + + x = torch.tensor([1.0, 10.0, 100.0]) + result = dynamic_range_compression_torch(x) + + # Result should be log(x * 1) = log(x) + expected = torch.log(x) + assert torch.allclose(result, expected) + + def test_clipping_small_values(self): + """Test that small values are clipped.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import dynamic_range_compression_torch + + x = torch.tensor([0.0, 1e-10, 1e-6]) + clip_val = 1e-5 + result = dynamic_range_compression_torch(x, clip_val=clip_val) + + # All values below clip_val should be clipped + expected = torch.log(torch.tensor([clip_val, clip_val, clip_val])) + assert torch.allclose(result, expected) + + def test_scaling_factor(self): + """Test with different scaling factor c.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import dynamic_range_compression_torch + + x = torch.tensor([1.0, 10.0]) + c = 2.0 + result = dynamic_range_compression_torch(x, c=c) + + expected = torch.log(x * c) + assert torch.allclose(result, expected) + + +class TestSpectralNormalize: + """Tests for spectral_normalize_torch utility.""" + + def test_is_wrapper_for_compression(self): + """Test that spectral_normalize is a wrapper for dynamic_range_compression.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import ( + dynamic_range_compression_torch, + spectral_normalize_torch, + ) + + magnitudes = torch.tensor([1.0, 10.0, 100.0]) + result = spectral_normalize_torch(magnitudes) + expected = dynamic_range_compression_torch(magnitudes) + + assert torch.allclose(result, expected) + + +class TestExactDiv: + """Tests for exact_div utility.""" + + def test_exact_division(self): + """Test exact division with valid inputs.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import exact_div + + assert exact_div(10, 2) == 5 + assert exact_div(100, 25) == 4 + assert exact_div(0, 5) == 0 + + def test_non_exact_division_raises(self): + """Test that non-exact division raises AssertionError.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import exact_div + + with pytest.raises(AssertionError): + exact_div(10, 3) + + with pytest.raises(AssertionError): + exact_div(7, 2) + + +class TestConcatTextWithPromptIds: + """Tests for concat_text_with_prompt_ids utility.""" + + def test_basic_concat(self): + """Test basic concatenation.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import concat_text_with_prompt_ids + + text = torch.tensor([[5, 6, 7]], dtype=torch.int64) + text_len = torch.tensor([3], dtype=torch.int32) + prompt_text = torch.tensor([[1, 2]], dtype=torch.int64) + prompt_text_len = torch.tensor([2], dtype=torch.int32) + + result, result_len = concat_text_with_prompt_ids(text, text_len, prompt_text, prompt_text_len) + + assert result.tolist() == [[1, 2, 5, 6, 7]] + assert result_len.item() == 5 + + def test_empty_prompt(self): + """Test with empty prompt.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import concat_text_with_prompt_ids + + text = torch.tensor([[1, 2, 3]], dtype=torch.int64) + text_len = torch.tensor([3], dtype=torch.int32) + prompt_text = torch.tensor([[]], dtype=torch.int64).reshape(1, 0) + prompt_text_len = torch.tensor([0], dtype=torch.int32) + + result, result_len = concat_text_with_prompt_ids(text, text_len, prompt_text, prompt_text_len) + + assert result.tolist() == [[1, 2, 3]] + assert result_len.item() == 3 + + def test_batch_concat(self): + """Test batched concatenation.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import concat_text_with_prompt_ids + + text = torch.tensor([[5, 6], [7, 8]], dtype=torch.int64) + text_len = torch.tensor([2, 2], dtype=torch.int32) + prompt_text = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64) + prompt_text_len = torch.tensor([2, 2], dtype=torch.int32) + + result, result_len = concat_text_with_prompt_ids(text, text_len, prompt_text, prompt_text_len) + + assert result.tolist() == [[1, 2, 5, 6], [3, 4, 7, 8]] + assert result_len.tolist() == [4, 4] + + def test_prompt_comes_first(self): + """Test that prompt tokens come before text tokens.""" + from vllm_omni.model_executor.models.cosyvoice3.utils import concat_text_with_prompt_ids + + text = torch.tensor([[100, 200]], dtype=torch.int64) + text_len = torch.tensor([2], dtype=torch.int32) + prompt_text = torch.tensor([[1, 2, 3]], dtype=torch.int64) + prompt_text_len = torch.tensor([3], dtype=torch.int32) + + result, _ = concat_text_with_prompt_ids(text, text_len, prompt_text, prompt_text_len) + + # First 3 tokens should be prompt + assert result[0, :3].tolist() == [1, 2, 3] + # Last 2 tokens should be text + assert result[0, 3:].tolist() == [100, 200] diff --git a/vllm_omni/diffusion/models/cosyvoice3_audio/__init__.py b/vllm_omni/diffusion/models/cosyvoice3_audio/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/diffusion/models/cosyvoice3_audio/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py new file mode 100644 index 000000000..9c5af8f33 --- /dev/null +++ b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py @@ -0,0 +1,446 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/flow/DiT +# Refactored to use vllm_omni diffusion infrastructure for optimized attention backends + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F +from diffusers.models.normalization import AdaLayerNormZero +from einops import repeat +from torch import nn +from vllm.logger import init_logger +from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb + +from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention + +logger = init_logger(__name__) + +""" +in notation: +b - batch +n - sequence +nt - text sequence +nw - raw wave length +d - dimension +""" + + +def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0): + """Precompute rotary embedding frequencies.""" + theta *= theta_rescale_factor ** (dim / (dim - 2)) + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + t = torch.arange(end, device=freqs.device) + freqs = torch.outer(t, freqs).float() + freqs_cos = torch.cos(freqs) # real part + freqs_sin = torch.sin(freqs) # imaginary part + return torch.cat([freqs_cos, freqs_sin], dim=-1) + + +def get_pos_embed_indices(start, length, max_pos, scale=1.0): + """Get position embedding indices.""" + scale = scale * torch.ones_like(start, dtype=torch.float32) + pos = ( + start.unsqueeze(1) + + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long() + ) + pos = torch.where(pos < max_pos, pos, max_pos - 1) + return pos + + +class FeedForward(nn.Module): + """Feed-forward network with GELU activation.""" + + def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = "none"): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + activation = nn.GELU(approximate=approximate) + project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation) + self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)) + + def forward(self, x): + return self.ff(x) + + +class DiTAttention(nn.Module): + """Attention module using diffusion infrastructure for optimized backends. + + This replaces the original Attention class to leverage FlashAttention, + SageAttention, or SDPA backends automatically. + """ + + def __init__( + self, + dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + ): + super().__init__() + self.dim = dim + self.heads = heads + self.dim_head = dim_head + self.inner_dim = dim_head * heads + self.dropout = dropout + self.scale = 1.0 / math.sqrt(dim_head) + + # Q/K/V projections + self.to_q = nn.Linear(dim, self.inner_dim) + self.to_k = nn.Linear(dim, self.inner_dim) + self.to_v = nn.Linear(dim, self.inner_dim) + + # Output projection + self.to_out = nn.Sequential(nn.Linear(self.inner_dim, dim), nn.Dropout(dropout)) + + # Diffusion attention backend (Flash/Sage/SDPA) + self.attn = DiffusionAttention( + num_heads=heads, + head_size=dim_head, + softmax_scale=self.scale, + causal=False, + ) + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + rope=None, + ) -> torch.Tensor: + batch_size, seq_len = x.shape[0], x.shape[1] + + # Project to Q, K, V + query = self.to_q(x) + key = self.to_k(x) + value = self.to_v(x) + + # Apply rotary position embedding + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + # Reshape for attention: (batch, seq, heads, head_dim) + query = query.view(batch_size, seq_len, self.heads, self.dim_head) + key = key.view(batch_size, seq_len, self.heads, self.dim_head) + value = value.view(batch_size, seq_len, self.heads, self.dim_head) + + # Use diffusion attention backend + # The diffusion Attention layer expects (batch, seq, heads, head_dim) + out = self.attn(query, key, value, attn_metadata=None) + + # Reshape back: (batch, seq, dim) + out = out.view(batch_size, seq_len, self.inner_dim) + out = out.to(query.dtype) + + # Output projection + out = self.to_out(out) + + # Apply mask if provided + if mask is not None: + if mask.dim() == 2: + mask = mask.unsqueeze(-1) + elif mask.dim() == 4: + # (batch, heads, seq, seq) -> use last dim + mask = mask[:, 0, -1].unsqueeze(-1) + out = out.masked_fill(~mask.bool(), 0.0) + + return out + + +class DiTBlock(nn.Module): + """DiT block with AdaLayerNorm modulation.""" + + def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1): + super().__init__() + + self.attn_norm = AdaLayerNormZero(dim) + self.attn = DiTAttention( + dim=dim, + heads=heads, + dim_head=dim_head, + dropout=dropout, + ) + + self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") + + def forward(self, x, t, mask=None, rope=None): + # pre-norm & modulation for attention input + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) + + # attention + attn_output = self.attn(x=norm, mask=mask, rope=rope) + + # process attention output for input x + x = x + gate_msa.unsqueeze(1) * attn_output + + ff_norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(ff_norm) + x = x + gate_mlp.unsqueeze(1) * ff_output + + return x + + +class CausalConvPositionEmbedding(nn.Module): + """Causal convolutional position embedding.""" + + def __init__(self, dim, kernel_size=31, groups=16): + super().__init__() + assert kernel_size % 2 != 0 + self.kernel_size = kernel_size + self.conv1 = nn.Sequential( + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=0), + nn.Mish(), + ) + self.conv2 = nn.Sequential( + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=0), + nn.Mish(), + ) + + def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None): + if mask is not None: + mask = mask[..., None] + x = x.masked_fill(~mask, 0.0) + + x = x.permute(0, 2, 1) + x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) + x = self.conv1(x) + x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) + x = self.conv2(x) + out = x.permute(0, 2, 1) + + if mask is not None: + out = out.masked_fill(~mask, 0.0) + + return out + + +class AdaLayerNormZero_Final(nn.Module): + """AdaLayerNormZero for final layer - returns only modulated x.""" + + def __init__(self, dim): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 2) + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x, emb): + emb = self.linear(self.silu(emb)) + scale, shift = torch.chunk(emb, 2, dim=1) + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +class GRN(nn.Module): + """Global Response Normalization layer.""" + + def __init__(self, dim): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, dim)) + + def forward(self, x): + Gx = torch.norm(x, p=2, dim=1, keepdim=True) + Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) + return self.gamma * (x * Nx) + self.beta + x + + +class ConvNeXtV2Block(nn.Module): + """ConvNeXt-V2 Block.""" + + def __init__(self, dim: int, intermediate_dim: int, dilation: int = 1): + super().__init__() + padding = (dilation * (7 - 1)) // 2 + self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation) + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, intermediate_dim) + self.act = nn.GELU() + self.grn = GRN(intermediate_dim) + self.pwconv2 = nn.Linear(intermediate_dim, dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = x.transpose(1, 2) + x = self.dwconv(x) + x = x.transpose(1, 2) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + return residual + x + + +class SinusPositionEmbedding(nn.Module): + """Sinusoidal position embedding.""" + + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x, scale=1000): + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) + emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +class TimestepEmbedding(nn.Module): + """Timestep embedding with MLP.""" + + def __init__(self, dim, freq_embed_dim=256): + super().__init__() + self.time_embed = SinusPositionEmbedding(freq_embed_dim) + self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + + def forward(self, timestep: torch.Tensor): + time_hidden = self.time_embed(timestep) + time_hidden = time_hidden.to(timestep.dtype) + time = self.time_mlp(time_hidden) + return time + + +class TextEmbedding(nn.Module): + """Text embedding with optional ConvNeXt modeling.""" + + def __init__(self, text_num_embeds, text_dim, conv_layers=0, conv_mult=2): + super().__init__() + self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim) + + if conv_layers > 0: + self.extra_modeling = True + self.precompute_max_pos = 4096 + self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False) + self.text_blocks = nn.Sequential( + *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)] + ) + else: + self.extra_modeling = False + + def forward(self, text: torch.Tensor, seq_len, drop_text=False): + batch, text_len = text.shape[0], text.shape[1] + text = text + 1 + text = text[:, :seq_len] + text = F.pad(text, (0, seq_len - text_len), value=0) + + if drop_text: + text = torch.zeros_like(text) + + text = self.text_embed(text) + + if self.extra_modeling: + batch_start = torch.zeros((batch,), dtype=torch.long) + pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos) + text_pos_embed = self.freqs_cis[pos_idx] + text = text + text_pos_embed + text = self.text_blocks(text) + + return text + + +class InputEmbedding(nn.Module): + """Input embedding combining noised audio, condition, text, and speaker.""" + + def __init__(self, mel_dim, text_dim, out_dim, spk_dim=None): + super().__init__() + spk_dim = 0 if spk_dim is None else spk_dim + self.spk_dim = spk_dim + self.proj = nn.Linear(mel_dim * 2 + text_dim + spk_dim, out_dim) + self.conv_pos_embed = CausalConvPositionEmbedding(dim=out_dim) + + def forward(self, x, cond, text_embed, spks): + to_cat = [x, cond, text_embed] + if self.spk_dim > 0: + spks = repeat(spks, "b c -> b t c", t=x.shape[1]) + to_cat.append(spks) + + x = self.proj(torch.cat(to_cat, dim=-1)) + x = self.conv_pos_embed(x) + x + return x + + +class DiT(nn.Module): + """Diffusion Transformer backbone using optimized attention backends. + + This is a drop-in replacement for the original DiT that uses the + vllm_omni diffusion infrastructure for FlashAttention/SageAttention/SDPA. + """ + + def __init__( + self, + *, + dim, + depth=8, + heads=8, + dim_head=64, + dropout=0.1, + ff_mult=4, + mel_dim=80, + mu_dim=None, + long_skip_connection=False, + spk_dim=None, + out_channels=None, + static_chunk_size=50, + num_decoding_left_chunks=2, + ): + super().__init__() + + self.time_embed = TimestepEmbedding(dim) + if mu_dim is None: + mu_dim = mel_dim + self.input_embed = InputEmbedding(mel_dim, mu_dim, dim, spk_dim) + + self.rotary_embed = RotaryEmbedding(dim_head) + + self.dim = dim + self.depth = depth + + self.transformer_blocks = nn.ModuleList( + [DiTBlock(dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, dropout=dropout) for _ in range(depth)] + ) + self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None + + self.norm_out = AdaLayerNormZero_Final(dim) + self.proj_out = nn.Linear(dim, mel_dim) + self.out_channels = out_channels + self.static_chunk_size = static_chunk_size + self.num_decoding_left_chunks = num_decoding_left_chunks + + def forward(self, x, mask, mu, t, spks=None, cond=None): + x = x.transpose(1, 2) + mu = mu.transpose(1, 2) + cond = cond.transpose(1, 2) + spks = spks.unsqueeze(dim=1) + batch, seq_len = x.shape[0], x.shape[1] + if t.ndim == 0: + t = t.repeat(batch) + + # t: conditioning time, c: context (text + masked cond audio), x: noised input audio + t = self.time_embed(t) + x = self.input_embed(x, cond, mu, spks.squeeze(1)) + + rope = self.rotary_embed.forward_from_seq_len(seq_len) + + if self.long_skip_connection is not None: + residual = x + + attn_mask = mask.bool().repeat(1, x.size(1), 1).unsqueeze(dim=1) + + for block in self.transformer_blocks: + x = block(x, t, mask=attn_mask.bool(), rope=rope) + + if self.long_skip_connection is not None: + x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) + + x = self.norm_out(x, t) + output = self.proj_out(x).transpose(1, 2) + return output diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index dac096339..f0102cfff 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -16,6 +16,7 @@ def _register_omni_hf_configs() -> None: try: from transformers import AutoConfig + from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( Qwen3TTSConfig, ) @@ -25,6 +26,7 @@ def _register_omni_hf_configs() -> None: try: AutoConfig.register("qwen3_tts", Qwen3TTSConfig) + AutoConfig.register("cosyvoice3", CosyVoice3Config) except ValueError: # Already registered elsewhere; ignore. return diff --git a/vllm_omni/model_executor/models/cosyvoice3/__init__.py b/vllm_omni/model_executor/models/cosyvoice3/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/__init__.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py new file mode 100644 index 000000000..7281cd81f --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py @@ -0,0 +1,322 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/flow +"""Conditional Flow Matching (CFM) classes for audio generation.""" + +from abc import ABC + +import torch +import torch.nn as nn +from omegaconf import DictConfig +from torch.nn import functional as F +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + +logger = init_logger(__name__) + + +class BASECFM(torch.nn.Module, ABC): + def __init__( + self, + n_feats, + cfm_params, + n_spks=1, + spk_emb_dim=128, + ): + super().__init__() + self.n_feats = n_feats + self.n_spks = n_spks + self.spk_emb_dim = spk_emb_dim + self.solver = cfm_params.solver + if hasattr(cfm_params, "sigma_min"): + self.sigma_min = cfm_params.sigma_min + else: + self.sigma_min = 1e-4 + + self.estimator = None + + +class ConditionalCFM(BASECFM): + def __init__(self, in_channels, cfm_params, n_spks=1, spk_emb_dim=64, estimator: torch.nn.Module = None): + super().__init__( + n_feats=in_channels, + cfm_params=cfm_params, + n_spks=n_spks, + spk_emb_dim=spk_emb_dim, + ) + self.t_scheduler = cfm_params.t_scheduler + self.training_cfg_rate = cfm_params.training_cfg_rate + self.inference_cfg_rate = cfm_params.inference_cfg_rate + in_channels = in_channels + (spk_emb_dim if n_spks > 0 else 0) + # Just change the architecture of the estimator here + self.estimator = estimator + + @torch.inference_mode() + def forward( + self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, prompt_len=0, cache=torch.zeros(1, 80, 0, 2) + ): + """Forward diffusion + + Args: + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): output_mask + shape: (batch_size, 1, mel_timesteps) + n_timesteps (int): number of diffusion steps + temperature (float, optional): temperature for scaling noise. Defaults to 1.0. + spks (torch.Tensor, optional): speaker ids. Defaults to None. + shape: (batch_size, spk_emb_dim) + cond (Optional[Any], optional): Not used but kept for future purposes + + Returns: + sample (torch.Tensor): generated mel-spectrogram + shape: (batch_size, n_feats, mel_timesteps) + """ + + z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature + cache_size = cache.shape[2] + # fix prompt and overlap part mu and z + if cache_size != 0: + z[:, :, :cache_size] = cache[:, :, :, 0] + mu[:, :, :cache_size] = cache[:, :, :, 1] + z_cache = torch.concat([z[:, :, :prompt_len], z[:, :, -34:]], dim=2) + mu_cache = torch.concat([mu[:, :, :prompt_len], mu[:, :, -34:]], dim=2) + cache = torch.stack([z_cache, mu_cache], dim=-1) + + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), cache + + def solve_euler(self, x, t_span, mu, mask, spks, cond): + """ + Fixed euler solver for ODEs. + Args: + x (torch.Tensor): random noise + t_span (torch.Tensor): n_timesteps interpolated + shape: (n_timesteps + 1,) + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): output_mask + shape: (batch_size, 1, mel_timesteps) + spks (torch.Tensor, optional): speaker ids. Defaults to None. + shape: (batch_size, spk_emb_dim) + cond (Optional[Any], optional): Not used but kept for future purposes + """ + t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0] + t = t.unsqueeze(dim=0) + + sol = [] + + # Do not use concat, it may cause memory format changed and trt infer with wrong results! + # NOTE when flow run in amp mode, x.dtype is float32, which cause nan in trt fp16 + # inference, so set dtype=spks.dtype + x_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) + mask_in = torch.zeros([2, 1, x.size(2)], device=x.device, dtype=spks.dtype) + mu_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) + t_in = torch.zeros([2], device=x.device, dtype=spks.dtype) + spks_in = torch.zeros([2, 80], device=x.device, dtype=spks.dtype) + cond_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) + for step in range(1, len(t_span)): + # Classifier-Free Guidance inference introduced in VoiceBox + x_in[:] = x + mask_in[:] = mask + mu_in[0] = mu + t_in[:] = t.unsqueeze(0) + spks_in[0] = spks + cond_in[0] = cond + dphi_dt = self.forward_estimator(x_in, mask_in, mu_in, t_in, spks_in, cond_in) + dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [x.size(0), x.size(0)], dim=0) + dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt + x = x + dt * dphi_dt + t = t + dt + sol.append(x) + if step < len(t_span) - 1: + dt = t_span[step + 1] - t + + return sol[-1].float() + + def forward_estimator(self, x, mask, mu, t, spks, cond): + if isinstance(self.estimator, torch.nn.Module): + return self.estimator(x, mask, mu, t, spks, cond) + else: + [estimator, stream], trt_engine = self.estimator.acquire_estimator() + # NOTE need to synchronize when switching stream + torch.cuda.current_stream().synchronize() + with stream: + estimator.set_input_shape("x", (2, 80, x.size(2))) + estimator.set_input_shape("mask", (2, 1, x.size(2))) + estimator.set_input_shape("mu", (2, 80, x.size(2))) + estimator.set_input_shape("t", (2,)) + estimator.set_input_shape("spks", (2, 80)) + estimator.set_input_shape("cond", (2, 80, x.size(2))) + data_ptrs = [ + x.contiguous().data_ptr(), + mask.contiguous().data_ptr(), + mu.contiguous().data_ptr(), + t.contiguous().data_ptr(), + spks.contiguous().data_ptr(), + cond.contiguous().data_ptr(), + x.data_ptr(), + ] + for i, j in enumerate(data_ptrs): + estimator.set_tensor_address(trt_engine.get_tensor_name(i), j) + # run trt engine + assert estimator.execute_async_v3(torch.cuda.current_stream().cuda_stream) is True + torch.cuda.current_stream().synchronize() + self.estimator.release_estimator(estimator, stream) + return x + + +class CausalConditionalCFM(ConditionalCFM): + def __init__(self, in_channels, cfm_params, n_spks=1, spk_emb_dim=64, estimator: torch.nn.Module = None): + super().__init__(in_channels, cfm_params, n_spks, spk_emb_dim, estimator) + + @torch.inference_mode() + def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None): + """Forward diffusion + + Args: + mu (torch.Tensor): output of encoder + shape: (batch_size, n_feats, mel_timesteps) + mask (torch.Tensor): output_mask + shape: (batch_size, 1, mel_timesteps) + n_timesteps (int): number of diffusion steps + temperature (float, optional): temperature for scaling noise. Defaults to 1.0. + spks (torch.Tensor, optional): speaker ids. Defaults to None. + shape: (batch_size, spk_emb_dim) + cond (Optional[Any], optional): Not used but kept for future purposes + + Returns: + sample (torch.Tensor): generated mel-spectrogram + shape: (batch_size, n_feats, mel_timesteps) + """ + + z = ( + torch.randn( + (mu.size(0), mu.size(1), mu.size(2)), + device=mu.device, + dtype=mu.dtype, + ) + * temperature + ) + + # fix prompt and overlap part mu and z + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + + return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None + + +class CausalMaskedDiffWithDiT(torch.nn.Module): + def __init__( + self, + input_size: int = 512, + output_size: int = 80, + spk_embed_dim: int = 192, + output_type: str = "mel", + vocab_size: int = 4096, + input_frame_rate: int = 50, + only_mask_loss: bool = True, + token_mel_ratio: int = 2, + pre_lookahead_len: int = 3, + pre_lookahead_layer: torch.nn.Module = None, + decoder: torch.nn.Module = None, + decoder_conf: dict = { + "in_channels": 240, + "out_channel": 80, + "spk_emb_dim": 80, + "n_spks": 1, + "cfm_params": DictConfig( + { + "sigma_min": 1e-06, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + "reg_loss_type": "l1", + } + ), + "decoder_params": { + "channels": [256, 256], + "dropout": 0.0, + "attention_head_dim": 64, + "n_blocks": 4, + "num_mid_blocks": 12, + "num_heads": 8, + "act_fn": "gelu", + }, + }, + ): + super().__init__() + self.input_size = input_size + self.output_size = output_size + self.decoder_conf = decoder_conf + self.vocab_size = vocab_size + self.output_type = output_type + self.input_frame_rate = input_frame_rate + logger.info(f"input frame rate={self.input_frame_rate}") + self.input_embedding = nn.Embedding(vocab_size, input_size) + self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size) + self.pre_lookahead_len = pre_lookahead_len + self.pre_lookahead_layer = pre_lookahead_layer + self.decoder = decoder + self.only_mask_loss = only_mask_loss + self.token_mel_ratio = token_mel_ratio + + @torch.inference_mode() + def inference( + self, + token, + token_len, + prompt_token, + prompt_token_len, + prompt_feat, + prompt_feat_len, + embedding, + finalize, + ): + assert token.shape[0] == 1 + # xvec projection + + embedding = F.normalize(embedding, dim=1) + + embedding = self.spk_embed_affine_layer(embedding) + + # concat text and prompt_text + token, token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + token_len + mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding) + token = self.input_embedding(torch.clamp(token, min=0)) * mask + # text encode + if finalize is True: + h = self.pre_lookahead_layer(token) + else: + h = self.pre_lookahead_layer( + token[:, : -self.pre_lookahead_len], context=token[:, -self.pre_lookahead_len :] + ) + h = h.repeat_interleave(self.token_mel_ratio, dim=1) + + mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1] + + # get conditions + conds = torch.zeros([1, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype) + + conds[:, :mel_len1] = prompt_feat + + conds = conds.transpose(1, 2) + + mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h) + feat, _ = self.decoder( + mu=h.transpose(1, 2).contiguous(), + mask=mask.unsqueeze(1), + spks=embedding, + cond=conds, + n_timesteps=10, + ) + + feat = feat[:, :, mel_len1:] + assert feat.shape[2] == mel_len2 + return feat.float(), None diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py new file mode 100644 index 000000000..bd31443cc --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -0,0 +1,951 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/hifigan +# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu) + + +"""HIFI-GAN""" + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from scipy.signal import get_window +from torch import pow, sin +from torch.nn import Conv1d, ConvTranspose1d +from torch.nn.utils import remove_weight_norm + +try: + from torch.nn.utils.parametrizations import weight_norm +except ImportError: + from torch.nn.utils import weight_norm +from torch.distributions.uniform import Uniform +from torch.nn import Parameter + + +# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license. +# LICENSE is in incl_licenses directory. +class Snake(nn.Module): + """ + Implementation of a sine-based periodic activation function + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter + References: + - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snake(256) + >>> x = torch.randn(256) + >>> x = a1(x) + """ + + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + """ + Initialization. + INPUT: + - in_features: shape of the input + - alpha: trainable parameter + alpha is initialized to 1 by default, higher values = higher-frequency. + alpha will be trained along with the rest of your model. + """ + super().__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + """ + Forward pass of the function. + Applies the function to the input elementwise. + Snake ∶= x + 1/a * sin^2 (xa) + """ + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + if self.alpha_logscale: + alpha = torch.exp(alpha) + x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + return x + + +"""hifigan based generator implementation. + +This code is modified from https://github.com/jik876/hifi-gan + ,https://github.com/kan-bayashi/ParallelWaveGAN and + https://github.com/NVIDIA/BigVGAN + +""" + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +class ResBlock(torch.nn.Module): + """Residual block module in HiFiGAN/BigVGAN.""" + + def __init__( + self, + channels: int = 512, + kernel_size: int = 3, + dilations: list[int] = [1, 3, 5], + causal: bool = False, + ): + super().__init__() + self.causal = causal + self.convs1 = nn.ModuleList() + self.convs2 = nn.ModuleList() + + for dilation in dilations: + self.convs1.append( + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation, + padding=get_padding(kernel_size, dilation), + ) + if causal is False + else CausalConv1d(channels, channels, kernel_size, 1, dilation=dilation, causal_type="left") + ) + ) + self.convs2.append( + weight_norm( + Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)) + if causal is False + else CausalConv1d(channels, channels, kernel_size, 1, dilation=1, causal_type="left") + ) + ) + self.convs1.apply(init_weights) + self.convs2.apply(init_weights) + self.activations1 = nn.ModuleList([Snake(channels, alpha_logscale=False) for _ in range(len(self.convs1))]) + self.activations2 = nn.ModuleList([Snake(channels, alpha_logscale=False) for _ in range(len(self.convs2))]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for idx in range(len(self.convs1)): + xt = self.activations1[idx](x) + xt = self.convs1[idx](xt) + xt = self.activations2[idx](xt) + xt = self.convs2[idx](xt) + x = xt + x + return x + + def remove_weight_norm(self): + for idx in range(len(self.convs1)): + remove_weight_norm(self.convs1[idx]) + remove_weight_norm(self.convs2[idx]) + + +class SineGen(torch.nn.Module): + """Definition of sine generator + SineGen(samp_rate, harmonic_num = 0, + sine_amp = 0.1, noise_std = 0.003, + voiced_threshold = 0, + flag_for_pulse=False) + samp_rate: sampling rate in Hz + harmonic_num: number of harmonic overtones (default 0) + sine_amp: amplitude of sine-wavefrom (default 0.1) + noise_std: std of Gaussian noise (default 0.003) + voiced_thoreshold: F0 threshold for U/V classification (default 0) + flag_for_pulse: this SinGen is used inside PulseGen (default False) + Note: when flag_for_pulse is True, the first time step of a voiced + segment is always sin(np.pi) or cos(0) + """ + + def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0): + super().__init__() + self.sine_amp = sine_amp + self.noise_std = noise_std + self.harmonic_num = harmonic_num + self.sampling_rate = samp_rate + self.voiced_threshold = voiced_threshold + + def _f02uv(self, f0): + # generate uv signal + uv = (f0 > self.voiced_threshold).type(torch.float32) + return uv + + @torch.no_grad() + def forward(self, f0): + """sine_tensor, uv = forward(f0) + input F0: tensor(batchsize=1, dim=1, length) + f0 for unvoiced steps should be 0 + output sine_tensor: tensor(batchsize=1, length, dim) + output uv: tensor(batchsize=1, length, 1) + """ + f0 = f0.transpose(1, 2) + F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device) + for i in range(self.harmonic_num + 1): + F_mat[:, i : i + 1, :] = f0 * (i + 1) / self.sampling_rate + + theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1) + u_dist = Uniform(low=-np.pi, high=np.pi) + phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device) + phase_vec[:, 0, :] = 0 + + # generate sine waveforms + sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec) + + # generate uv signal + uv = self._f02uv(f0) + + # noise: for unvoiced should be similar to sine_amp + # std = self.sine_amp/3 -> max value ~ self.sine_amp + # . for voiced regions is self.noise_std + noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 + noise = noise_amp * torch.randn_like(sine_waves) + + # first: set the unvoiced part to 0 by uv + # then: additive noise + sine_waves = sine_waves * uv + noise + return sine_waves.transpose(1, 2), uv.transpose(1, 2), noise + + +class SineGen2(torch.nn.Module): + """Definition of sine generator + SineGen(samp_rate, harmonic_num = 0, + sine_amp = 0.1, noise_std = 0.003, + voiced_threshold = 0, + flag_for_pulse=False) + samp_rate: sampling rate in Hz + harmonic_num: number of harmonic overtones (default 0) + sine_amp: amplitude of sine-wavefrom (default 0.1) + noise_std: std of Gaussian noise (default 0.003) + voiced_thoreshold: F0 threshold for U/V classification (default 0) + flag_for_pulse: this SinGen is used inside PulseGen (default False) + Note: when flag_for_pulse is True, the first time step of a voiced + segment is always sin(np.pi) or cos(0) + """ + + def __init__( + self, + samp_rate, + upsample_scale, + harmonic_num=0, + sine_amp=0.1, + noise_std=0.003, + voiced_threshold=0, + flag_for_pulse=False, + causal=False, + ): + super().__init__() + self.sine_amp = sine_amp + self.noise_std = noise_std + self.harmonic_num = harmonic_num + self.dim = self.harmonic_num + 1 + self.sampling_rate = samp_rate + self.voiced_threshold = voiced_threshold + self.flag_for_pulse = flag_for_pulse + self.upsample_scale = upsample_scale + self.causal = causal + if causal is True: + self.rand_ini = torch.rand(1, 9) + self.rand_ini[:, 0] = 0 + self.sine_waves = torch.rand(1, 300 * 24000, 9) + + def _f02uv(self, f0): + # generate uv signal + uv = (f0 > self.voiced_threshold).type(torch.float32) + return uv + + def _f02sine(self, f0_values): + """f0_values: (batchsize, length, dim) + where dim indicates fundamental tone and overtones + """ + # convert to F0 in rad. The integer part n can be ignored + # because 2 * np.pi * n doesn't affect phase + rad_values = (f0_values / self.sampling_rate) % 1 + + # initial phase noise (no noise for fundamental component) + if self.training is False and self.causal is True: + rad_values[:, 0, :] = rad_values[:, 0, :] + self.rand_ini.to(rad_values.device) + else: + rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device) + rand_ini[:, 0] = 0 + rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini + + # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad) + if not self.flag_for_pulse: + rad_values = torch.nn.functional.interpolate( + rad_values.transpose(1, 2), scale_factor=1 / self.upsample_scale, mode="linear" + ).transpose(1, 2) + + phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi + phase = torch.nn.functional.interpolate( + phase.transpose(1, 2) * self.upsample_scale, + scale_factor=self.upsample_scale, + mode="nearest" if self.causal is True else "linear", + ).transpose(1, 2) + sines = torch.sin(phase) + else: + # If necessary, make sure that the first time step of every + # voiced segments is sin(pi) or cos(0) + # This is used for pulse-train generation + + # identify the last time step in unvoiced segments + uv = self._f02uv(f0_values) + uv_1 = torch.roll(uv, shifts=-1, dims=1) + uv_1[:, -1, :] = 1 + u_loc = (uv < 1) * (uv_1 > 0) + + # get the instantanouse phase + tmp_cumsum = torch.cumsum(rad_values, dim=1) + # different batch needs to be processed differently + for idx in range(f0_values.shape[0]): + temp_sum = tmp_cumsum[idx, u_loc[idx, :, 0], :] + temp_sum[1:, :] = temp_sum[1:, :] - temp_sum[0:-1, :] + # stores the accumulation of i.phase within + # each voiced segments + tmp_cumsum[idx, :, :] = 0 + tmp_cumsum[idx, u_loc[idx, :, 0], :] = temp_sum + + # rad_values - tmp_cumsum: remove the accumulation of i.phase + # within the previous voiced segment. + i_phase = torch.cumsum(rad_values - tmp_cumsum, dim=1) + + # get the sines + sines = torch.cos(i_phase * 2 * np.pi) + return sines + + def forward(self, f0): + """sine_tensor, uv = forward(f0) + input F0: tensor(batchsize=1, length, dim=1) + f0 for unvoiced steps should be 0 + output sine_tensor: tensor(batchsize=1, length, dim) + output uv: tensor(batchsize=1, length, 1) + """ + # fundamental component + fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device)) + + # generate sine waveforms + sine_waves = self._f02sine(fn) * self.sine_amp + + # generate uv signal + uv = self._f02uv(f0) + + # noise: for unvoiced should be similar to sine_amp + # std = self.sine_amp/3 -> max value ~ self.sine_amp + # . for voiced regions is self.noise_std + noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 + if self.training is False and self.causal is True: + noise = noise_amp * self.sine_waves[:, : sine_waves.shape[1]].to(sine_waves.device) + else: + noise = noise_amp * torch.randn_like(sine_waves) + + # first: set the unvoiced part to 0 by uv + # then: additive noise + sine_waves = sine_waves * uv + noise + return sine_waves, uv, noise + + +class SourceModuleHnNSF(torch.nn.Module): + """SourceModule for hn-nsf + SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1, + add_noise_std=0.003, voiced_threshold=0) + sampling_rate: sampling_rate in Hz + harmonic_num: number of harmonic above F0 (default: 0) + sine_amp: amplitude of sine source signal (default: 0.1) + add_noise_std: std of additive Gaussian noise (default: 0.003) + note that amplitude of noise in unvoiced is decided + by sine_amp + voiced_threshold: threshold to set U/V given F0 (default: 0) + Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) + F0_sampled (batchsize, length, 1) + Sine_source (batchsize, length, 1) + noise_source (batchsize, length 1) + uv (batchsize, length, 1) + """ + + def __init__( + self, + sampling_rate, + upsample_scale, + harmonic_num=0, + sine_amp=0.1, + add_noise_std=0.003, + voiced_threshold=0, + sinegen_type="1", + causal=False, + ): + super().__init__() + + self.sine_amp = sine_amp + self.noise_std = add_noise_std + + # to produce sine waveforms + if sinegen_type == "1": + self.l_sin_gen = SineGen(sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold) + else: + self.l_sin_gen = SineGen2( + sampling_rate, upsample_scale, harmonic_num, sine_amp, add_noise_std, voiced_threshold, causal=causal + ) + + # to merge source harmonics into a single excitation + self.l_linear = torch.nn.Linear(harmonic_num + 1, 1) + self.l_tanh = torch.nn.Tanh() + self.causal = causal + if causal is True: + self.uv = torch.rand(1, 300 * 24000, 1) + + def forward(self, x): + """ + Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) + F0_sampled (batchsize, length, 1) + Sine_source (batchsize, length, 1) + noise_source (batchsize, length 1) + """ + # source for harmonic branch + with torch.no_grad(): + sine_wavs, uv, _ = self.l_sin_gen(x) + sine_merge = self.l_tanh(self.l_linear(sine_wavs)) + + # source for noise branch, in the same shape as uv + if self.training is False and self.causal is True: + noise = self.uv[:, : uv.shape[1]] * self.sine_amp / 3 + else: + noise = torch.randn_like(uv) * self.sine_amp / 3 + return sine_merge, noise, uv + + +class HiFTGenerator(nn.Module): + """ + HiFTNet Generator: Neural Source Filter + ISTFTNet + https://arxiv.org/abs/2309.09493 + """ + + def __init__( + self, + in_channels: int = 80, + base_channels: int = 512, + nb_harmonics: int = 8, + sampling_rate: int = 22050, + nsf_alpha: float = 0.1, + nsf_sigma: float = 0.003, + nsf_voiced_threshold: float = 10, + upsample_rates: list[int] = [8, 8], + upsample_kernel_sizes: list[int] = [16, 16], + istft_params: dict[str, int] = {"n_fft": 16, "hop_len": 4}, + resblock_kernel_sizes: list[int] = [3, 7, 11], + resblock_dilation_sizes: list[list[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + source_resblock_kernel_sizes: list[int] = [7, 11], + source_resblock_dilation_sizes: list[list[int]] = [[1, 3, 5], [1, 3, 5]], + lrelu_slope: float = 0.1, + audio_limit: float = 0.99, + f0_predictor: torch.nn.Module = None, + ): + super().__init__() + + self.out_channels = 1 + self.nb_harmonics = nb_harmonics + self.sampling_rate = sampling_rate + self.istft_params = istft_params + self.lrelu_slope = lrelu_slope + self.audio_limit = audio_limit + + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + # NOTE in CosyVoice2, we use the original SineGen implementation + self.m_source = SourceModuleHnNSF( + sampling_rate=sampling_rate, + upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"], + harmonic_num=nb_harmonics, + sine_amp=nsf_alpha, + add_noise_std=nsf_sigma, + voiced_threshold=nsf_voiced_threshold, + sinegen_type="1" if self.sampling_rate == 22050 else "2", + causal=False, + ) + self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * istft_params["hop_len"]) + + self.conv_pre = weight_norm(Conv1d(in_channels, base_channels, 7, 1, padding=3)) + + # Up + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): + self.ups.append( + weight_norm( + ConvTranspose1d( + base_channels // (2**i), + base_channels // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + ) + + # Down + self.source_downs = nn.ModuleList() + self.source_resblocks = nn.ModuleList() + downsample_rates = [1] + upsample_rates[::-1][:-1] + downsample_cum_rates = np.cumprod(downsample_rates) + for i, (u, k, d) in enumerate( + zip(downsample_cum_rates[::-1], source_resblock_kernel_sizes, source_resblock_dilation_sizes) + ): + if u == 1: + self.source_downs.append(Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), 1, 1)) + else: + self.source_downs.append( + Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), u * 2, u, padding=(u // 2)) + ) + + self.source_resblocks.append(ResBlock(base_channels // (2 ** (i + 1)), k, d)) + + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = base_channels // (2 ** (i + 1)) + for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)): + self.resblocks.append(ResBlock(ch, k, d)) + + self.conv_post = weight_norm(Conv1d(ch, istft_params["n_fft"] + 2, 7, 1, padding=3)) + self.ups.apply(init_weights) + self.conv_post.apply(init_weights) + self.reflection_pad = nn.ReflectionPad1d((1, 0)) + self.stft_window = torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32)) + self.f0_predictor = f0_predictor + + def remove_weight_norm(self): + for layer in self.ups: + remove_weight_norm(layer) + for block in self.resblocks: + block.remove_weight_norm() + remove_weight_norm(self.conv_pre) + remove_weight_norm(self.conv_post) + self.m_source.remove_weight_norm() + for layer in self.source_downs: + remove_weight_norm(layer) + for block in self.source_resblocks: + block.remove_weight_norm() + + def _stft(self, x): + spec = torch.stft( + x, + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=self.stft_window.to(x.device), + return_complex=True, + ) + spec = torch.view_as_real(spec) # [B, F, TT, 2] + return spec[..., 0], spec[..., 1] + + def _istft(self, magnitude, phase): + magnitude = torch.clip(magnitude, max=1e2) + real = magnitude * torch.cos(phase) + img = magnitude * torch.sin(phase) + inverse_transform = torch.istft( + torch.complex(real, img), + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=self.stft_window.to(magnitude.device), + ) + return inverse_transform + + def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor: + s_stft_real, s_stft_imag = self._stft(s.squeeze(1)) + s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1) + + x = self.conv_pre(x) + for i in range(self.num_upsamples): + x = F.leaky_relu(x, self.lrelu_slope) + x = self.ups[i](x) + + if i == self.num_upsamples - 1: + x = self.reflection_pad(x) + + # fusion + si = self.source_downs[i](s_stft) + si = self.source_resblocks[i](si) + x = x + si + + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + x = F.leaky_relu(x) + x = self.conv_post(x) + magnitude = torch.exp(x[:, : self.istft_params["n_fft"] // 2 + 1, :]) + phase = torch.sin(x[:, self.istft_params["n_fft"] // 2 + 1 :, :]) # actually, sin is redundancy + + x = self._istft(magnitude, phase) + x = torch.clamp(x, -self.audio_limit, self.audio_limit) + return x + + def forward( + self, + batch: dict, + device: torch.device, + ) -> dict[str, torch.Tensor | None]: + speech_feat = batch["speech_feat"].transpose(1, 2).to(device) + # mel->f0 + f0 = self.f0_predictor(speech_feat) + # f0->source + s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t + s, _, _ = self.m_source(s) + s = s.transpose(1, 2) + # mel+source->speech + generated_speech = self.decode(x=speech_feat, s=s) + return generated_speech, f0 + + @torch.inference_mode() + def inference(self, speech_feat: torch.Tensor, cache_source: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor: + # mel->f0 + f0 = self.f0_predictor(speech_feat) + # f0->source + s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t + s, _, _ = self.m_source(s) + s = s.transpose(1, 2) + # use cache_source to avoid glitch + if cache_source.shape[2] != 0: + s[:, :, : cache_source.shape[2]] = cache_source + generated_speech = self.decode(x=speech_feat, s=s) + return generated_speech, s + + +class CausalHiFTGenerator(HiFTGenerator): + """ + HiFTNet Generator: Neural Source Filter + ISTFTNet + https://arxiv.org/abs/2309.09493 + """ + + def __init__( + self, + in_channels: int = 80, + base_channels: int = 512, + nb_harmonics: int = 8, + sampling_rate: int = 22050, + nsf_alpha: float = 0.1, + nsf_sigma: float = 0.003, + nsf_voiced_threshold: float = 10, + upsample_rates: list[int] = [8, 8], + upsample_kernel_sizes: list[int] = [16, 16], + istft_params: dict[str, int] = {"n_fft": 16, "hop_len": 4}, + resblock_kernel_sizes: list[int] = [3, 7, 11], + resblock_dilation_sizes: list[list[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + source_resblock_kernel_sizes: list[int] = [7, 11], + source_resblock_dilation_sizes: list[list[int]] = [[1, 3, 5], [1, 3, 5]], + lrelu_slope: float = 0.1, + audio_limit: float = 0.99, + conv_pre_look_right: int = 4, + f0_predictor: torch.nn.Module = None, + ): + torch.nn.Module.__init__(self) + + self.out_channels = 1 + self.nb_harmonics = nb_harmonics + self.sampling_rate = sampling_rate + self.istft_params = istft_params + self.lrelu_slope = lrelu_slope + self.audio_limit = audio_limit + + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + self.m_source = SourceModuleHnNSF( + sampling_rate=sampling_rate, + upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"], + harmonic_num=nb_harmonics, + sine_amp=nsf_alpha, + add_noise_std=nsf_sigma, + voiced_threshold=nsf_voiced_threshold, + sinegen_type="1" if self.sampling_rate == 22050 else "2", + causal=True, + ) + self.upsample_rates = upsample_rates + self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * istft_params["hop_len"]) + + self.conv_pre = weight_norm( + CausalConv1d(in_channels, base_channels, conv_pre_look_right + 1, 1, causal_type="right") + ) + + # Up + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): + self.ups.append( + weight_norm( + CausalConv1dUpsample( + base_channels // (2**i), + base_channels // (2 ** (i + 1)), + k, + u, + ) + ) + ) + + # Down + self.source_downs = nn.ModuleList() + self.source_resblocks = nn.ModuleList() + downsample_rates = [1] + upsample_rates[::-1][:-1] + downsample_cum_rates = np.cumprod(downsample_rates) + for i, (u, k, d) in enumerate( + zip(downsample_cum_rates[::-1], source_resblock_kernel_sizes, source_resblock_dilation_sizes) + ): + if u == 1: + self.source_downs.append( + CausalConv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), 1, 1, causal_type="left") + ) + else: + self.source_downs.append( + CausalConv1dDownSample(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), u * 2, u) + ) + + self.source_resblocks.append(ResBlock(base_channels // (2 ** (i + 1)), k, d, causal=True)) + + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = base_channels // (2 ** (i + 1)) + for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)): + self.resblocks.append(ResBlock(ch, k, d, causal=True)) + + self.conv_post = weight_norm(CausalConv1d(ch, istft_params["n_fft"] + 2, 7, 1, causal_type="left")) + self.ups.apply(init_weights) + self.conv_post.apply(init_weights) + self.reflection_pad = nn.ReflectionPad1d((1, 0)) + self.stft_window = torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32)) + self.conv_pre_look_right = conv_pre_look_right + self.f0_predictor = f0_predictor + + def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0), finalize: bool = True) -> torch.Tensor: + s_stft_real, s_stft_imag = self._stft(s.squeeze(1)) + if finalize is True: + x = self.conv_pre(x) + else: + x = self.conv_pre(x[:, :, : -self.conv_pre_look_right], x[:, :, -self.conv_pre_look_right :]) + s_stft_real = s_stft_real[:, :, : -int(np.prod(self.upsample_rates) * self.conv_pre_look_right)] + s_stft_imag = s_stft_imag[:, :, : -int(np.prod(self.upsample_rates) * self.conv_pre_look_right)] + s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1) + + for i in range(self.num_upsamples): + x = F.leaky_relu(x, self.lrelu_slope) + x = self.ups[i](x) + + if i == self.num_upsamples - 1: + x = self.reflection_pad(x) + + # fusion + si = self.source_downs[i](s_stft) + si = self.source_resblocks[i](si) + x = x + si + + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + x = F.leaky_relu(x) + x = self.conv_post(x) + magnitude = torch.exp(x[:, : self.istft_params["n_fft"] // 2 + 1, :]) + phase = torch.sin(x[:, self.istft_params["n_fft"] // 2 + 1 :, :]) # actually, sin is redundancy + + x = self._istft(magnitude, phase) + if finalize is False: + x = x[:, : -int(np.prod(self.upsample_rates) * self.istft_params["hop_len"])] + x = torch.clamp(x, -self.audio_limit, self.audio_limit) + return x + + @torch.inference_mode() + def inference(self, speech_feat: torch.Tensor, finalize: bool = True) -> torch.Tensor: + # mel->f0 NOTE f0_predictor precision is crucial for causal inference, move + # self.f0_predictor to cpu if necessary + self.f0_predictor.to("cpu") + f0 = self.f0_predictor(speech_feat.cpu(), finalize=finalize).to(speech_feat) + # f0->source + s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t + s, _, _ = self.m_source(s) + s = s.transpose(1, 2) + if finalize is True: + generated_speech = self.decode(x=speech_feat, s=s, finalize=finalize) + else: + generated_speech = self.decode( + x=speech_feat[:, :, : -self.f0_predictor.condnet[0].causal_padding], s=s, finalize=finalize + ) + return generated_speech, s + + +class CausalConv1dUpsample(torch.nn.Conv1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + device=None, + dtype=None, + ) -> None: + super().__init__( + in_channels, + out_channels, + kernel_size, + 1, + padding=0, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=padding_mode, + device=device, + dtype=dtype, + ) + assert dilation == 1 + self.causal_padding = kernel_size - 1 + self.upsample = torch.nn.Upsample(scale_factor=stride, mode="nearest") + + def forward(self, x: torch.Tensor, cache: torch.Tensor = torch.zeros(0, 0, 0)) -> tuple[torch.Tensor, torch.Tensor]: + x = self.upsample(x) + input_timestep = x.shape[2] + if cache.size(2) == 0: + x = F.pad(x, (self.causal_padding, 0), value=0.0) + else: + assert cache.size(2) == self.causal_padding + x = torch.concat([cache, x], dim=2) + x = super().forward(x) + assert input_timestep == x.shape[2] + return x + + +class CausalConv1dDownSample(torch.nn.Conv1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + device=None, + dtype=None, + ) -> None: + super().__init__( + in_channels, + out_channels, + kernel_size, + stride, + padding=0, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=padding_mode, + device=device, + dtype=dtype, + ) + assert stride != 1 and dilation == 1 + assert kernel_size % stride == 0 + self.causal_padding = stride - 1 + + def forward(self, x: torch.Tensor, cache: torch.Tensor = torch.zeros(0, 0, 0)) -> tuple[torch.Tensor, torch.Tensor]: + if cache.size(2) == 0: + x = F.pad(x, (self.causal_padding, 0), value=0.0) + else: + assert cache.size(2) == self.causal_padding + x = torch.concat([cache, x], dim=2) + x = super().forward(x) + return x + + +# NOTE(Xiang Lyu) causal conv module used in convolution-based vocoder +class CausalConv1d(torch.nn.Conv1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + causal_type: str = "left", + device=None, + dtype=None, + ) -> None: + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=padding_mode, + device=device, + dtype=dtype, + ) + assert stride == 1 + self.causal_padding = int((kernel_size * dilation - dilation) / 2) * 2 + (kernel_size + 1) % 2 + assert causal_type in ["left", "right"] + self.causal_type = causal_type + + def forward(self, x: torch.Tensor, cache: torch.Tensor = torch.zeros(0, 0, 0)) -> tuple[torch.Tensor]: + input_timestep = x.shape[2] + if cache.size(2) == 0: + cache = torch.zeros(x.shape[0], x.shape[1], self.causal_padding).to(x) + assert cache.size(2) == self.causal_padding + if self.causal_type == "left": + x = torch.concat([cache, x], dim=2) + else: + x = torch.concat([x, cache], dim=2) + x = super().forward(x) + assert x.shape[2] == input_timestep + return x + + +class CausalConvRNNF0Predictor(nn.Module): + def __init__(self, num_class: int = 1, in_channels: int = 80, cond_channels: int = 512): + super().__init__() + + self.num_class = num_class + self.condnet = nn.Sequential( + weight_norm(CausalConv1d(in_channels, cond_channels, kernel_size=4, causal_type="right")), + nn.ELU(), + weight_norm(CausalConv1d(cond_channels, cond_channels, kernel_size=3, causal_type="left")), + nn.ELU(), + weight_norm(CausalConv1d(cond_channels, cond_channels, kernel_size=3, causal_type="left")), + nn.ELU(), + weight_norm(CausalConv1d(cond_channels, cond_channels, kernel_size=3, causal_type="left")), + nn.ELU(), + weight_norm(CausalConv1d(cond_channels, cond_channels, kernel_size=3, causal_type="left")), + nn.ELU(), + ) + self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class) + + def forward(self, x: torch.Tensor, finalize: bool = True) -> torch.Tensor: + if finalize is True: + x = self.condnet[0](x) + else: + x = self.condnet[0](x[:, :, : -self.condnet[0].causal_padding], x[:, :, -self.condnet[0].causal_padding :]) + for i in range(1, len(self.condnet)): + x = self.condnet[i](x) + x = x.transpose(1, 2) + return torch.abs(self.classifier(x).squeeze(-1)) diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/layers.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/layers.py new file mode 100644 index 000000000..40d6bb54e --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/layers.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Neural network layers for CosyVoice3 Code2Wav.""" + +import torch +import torch.nn as nn +from torch.nn import functional as F + + +class PreLookaheadLayer(nn.Module): + """Pre-lookahead convolutional layer for causal processing.""" + + def __init__(self, in_channels: int, channels: int, pre_lookahead_len: int = 1): + super().__init__() + self.in_channels = in_channels + self.channels = channels + self.pre_lookahead_len = pre_lookahead_len + self.conv1 = nn.Conv1d( + in_channels, + channels, + kernel_size=pre_lookahead_len + 1, + stride=1, + padding=0, + ) + self.conv2 = nn.Conv1d( + channels, + in_channels, + kernel_size=3, + stride=1, + padding=0, + ) + + def forward(self, inputs: torch.Tensor, context: torch.Tensor = torch.zeros(0, 0, 0)) -> torch.Tensor: + """ + Args: + inputs: (batch_size, seq_len, channels) + context: Optional context tensor for streaming + """ + outputs = inputs.transpose(1, 2).contiguous() + context = context.transpose(1, 2).contiguous() + # look ahead + if context.size(2) == 0: + outputs = F.pad(outputs, (0, self.pre_lookahead_len), mode="constant", value=0.0) + else: + assert self.training is False, "you have passed context, make sure that you are running inference mode" + assert context.size(2) == self.pre_lookahead_len + outputs = F.pad( + torch.concat([outputs, context], dim=2), + (0, self.pre_lookahead_len - context.size(2)), + mode="constant", + value=0.0, + ) + outputs = F.leaky_relu(self.conv1(outputs)) + # outputs + outputs = F.pad(outputs, (self.conv2.kernel_size[0] - 1, 0), mode="constant", value=0.0) + outputs = self.conv2(outputs) + outputs = outputs.transpose(1, 2).contiguous() + + # residual connection + outputs = outputs + inputs + return outputs diff --git a/vllm_omni/model_executor/models/cosyvoice3/config.py b/vllm_omni/model_executor/models/cosyvoice3/config.py new file mode 100644 index 000000000..0c9a28997 --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/config.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from transformers.configuration_utils import PretrainedConfig + + +class CosyVoice3Config(PretrainedConfig): + model_type = "cosyvoice3" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.sample_rate = 24000 + self.llm_input_size = 896 + self.llm_output_size = 896 + self.hidden_size = self.llm_output_size + self.num_attention_heads = 14 + self.num_hidden_layers = 24 + self.spk_embed_dim = 192 + self.token_frame_rate = 25 + self.token_mel_ratio = 2 + self.vocab_size = 151923 + self.min_token_text_ratio = 2 + self.max_token_text_ratio = 20 + self.allowed_special = "all" + self.skip_special_tokens = True + self.target_sr = 24000 + self.feat_extractor = { + "n_fft": 1920, + "num_mels": 80, + "sampling_rate": self.sample_rate, + "hop_size": 480, + "win_size": 1920, + "fmin": 0, + "fmax": None, + "center": False, + } + self.speculative_config = None + self.qwen_pretrain_path = "CosyVoice-BlankEN" + self.campplus_onxx_path = "campplus.onnx" + self.speech_tokenizer_path = "speech_tokenizer_v3.onnx" + self.spk2info_path = "spk2info.pt" + self.version = "cosyvoice3" + self.llm = { + "llm_input_size": self.llm_input_size, + "llm_output_size": self.llm_output_size, + "speech_token_size": 6561, + "eos_token_id": 6561 + 1, + "length_normalized_loss": True, + "lsm_weight": 0, + "mix_ratio": [5, 15], + "llm": { + "pretrain_path": self.qwen_pretrain_path, + }, + "sampling": { + "top_p": 0.8, + "top_k": 25, + "win_size": 10, + "tau_r": 0.1, + }, + "spk_embed_dim": self.spk_embed_dim, + } + self.flow = { + "input_size": 80, + "output_size": 80, + "spk_embed_dim": self.spk_embed_dim, + "output_type": "mel", + "vocab_size": 6561, + "input_frame_rate": self.token_frame_rate, + "only_mask_loss": True, + "token_mel_ratio": self.token_mel_ratio, + "pre_lookahead_len": 3, + "pre_lookahead_layer": { + "in_channels": 80, + "channels": 1024, + "pre_lookahead_len": 3, + }, + "decoder": { + "in_channels": 240, + "n_spks": 1, + "spk_emb_dim": 80, + "cfm_params": { + "sigma_min": 1e-06, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + "reg_loss_type": "l1", + }, + "estimator": { + "dim": 1024, + "depth": 22, + "heads": 16, + "dim_head": 64, + "ff_mult": 2, + "mel_dim": 80, + "mu_dim": 80, + "spk_dim": 80, + "out_channels": 80, + "static_chunk_size": self.token_frame_rate * self.token_mel_ratio, + "num_decoding_left_chunks": -1, + }, + }, + } + self.hift = { + "in_channels": 80, + "base_channels": 512, + "nb_harmonics": 8, + "sampling_rate": self.sample_rate, + "nsf_alpha": 0.1, + "nsf_sigma": 0.003, + "nsf_voiced_threshold": 10, + "upsample_rates": [8, 5, 3], + "upsample_kernel_sizes": [16, 11, 7], + "istft_params": { + "n_fft": 16, + "hop_len": 4, + }, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "source_resblock_kernel_sizes": [7, 7, 11], + "source_resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "lrelu_slope": 0.1, + "audio_limit": 0.99, + "conv_pre_look_right": 4, + "f0_predictor": { + "num_class": 1, + "in_channels": 80, + "cond_channels": 512, + }, + } diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py new file mode 100644 index 000000000..87c5f323a --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -0,0 +1,468 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os +from collections.abc import Iterable, Mapping, Sequence +from functools import partial + +import numpy as np +import torch +import torch.nn as nn +from transformers.feature_extraction_utils import BatchFeature +from vllm.config import VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import SupportsMultiModal +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalDataDict, MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + ProcessorInputs, + PromptIndexTargets, + PromptInsertion, + PromptUpdate, +) +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config +from vllm_omni.model_executor.models.cosyvoice3.utils import ( + concat_text_with_prompt_ids, + extract_speech_feat, + extract_speech_token, + extract_spk_embedding, + extract_text_token, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + + +class CosyVoice3MultiModalProcessingInfo(BaseProcessingInfo): + def get_hf_config(self): + """If the config is not already present pass it + as a class and it will try to find it in your + model directory just copy the config class there also. + """ + return self.ctx.get_hf_config(CosyVoice3Config) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + """How many audio can you pass. I think I should keep it as 1 + For now I have kept it None. + """ + return {"audio": None} + + def get_data_parser(self): + return MultiModalDataParser( + target_sr=self.ctx.get_hf_config().target_sr, + expected_hidden_size=self._get_expected_hidden_size(), + ) + + +class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModalProcessingInfo]): + def _ensure_cached_runtime_components(self, model_dir: str, config: CosyVoice3Config) -> None: + cached_model_dir = getattr(self, "_cached_model_dir", None) + if cached_model_dir == model_dir: + return + + import onnxruntime + + from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer + from vllm_omni.model_executor.models.cosyvoice3.utils import mel_spectrogram + + option = onnxruntime.SessionOptions() + option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + option.intra_op_num_threads = 1 + + self.tokenizer = get_qwen_tokenizer( + token_path=os.path.join(model_dir, config.qwen_pretrain_path), + skip_special_tokens=config.skip_special_tokens, + version=config.version, + ) + self.speech_tokenizer = onnxruntime.InferenceSession( + os.path.join(model_dir, config.speech_tokenizer_path), + sess_options=option, + providers=["CUDAExecutionProvider" if torch.cuda.is_available() else "CPUExecutionProvider"], + ) + self.feat_extractor = partial(mel_spectrogram, **getattr(config, "feat_extractor", {})) + self.campplus_session = onnxruntime.InferenceSession( + os.path.join(model_dir, config.campplus_onxx_path), + sess_options=option, + providers=["CPUExecutionProvider"], + ) + self._cached_model_dir = model_dir + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + """ + apply-> cached_apply_hf_processor -> apply_hf_processor_mm -> + _call_hf_processor. + _call_hf_processor takes input prompt and mm_data and returns + token ids and tensors + """ + config = self.info.ctx.get_hf_config() + model_dir = self.info.ctx.model_config.model + self._ensure_cached_runtime_components(model_dir, config) + + audio = mm_data.get("audio", None) + + if audio is None: + audio = mm_data.get("audios") + if audio is not None: + audio = audio[0], config.target_sr + + text_token, text_token_len = extract_text_token(prompt, self.tokenizer, config.allowed_special) + if audio is None: + # Text-only path for profiling/cache + return BatchFeature({"input_ids": text_token, "input_len": [text_token_len]}) + + prompt_text = mm_kwargs.get("prompt_text") + + if not isinstance(prompt_text, str): + raise ValueError(f"prompt text is None : {prompt_text}") + + prompt_text_token, prompt_text_token_len = extract_text_token( + prompt_text, self.tokenizer, config.allowed_special + ) + + input_ids, input_len = concat_text_with_prompt_ids( + text_token, + text_token_len, + prompt_text_token, + prompt_text_token_len, + ) + logger.debug( + "cosyvoice _call_hf_processor: prompt_text_token=%s text_token=%s input_ids=%s " + "prompt_text_len=%s text_len=%s input_len=%s", + prompt_text_token.tolist(), + text_token.tolist(), + input_ids.tolist(), + int(prompt_text_token_len), + int(text_token_len), + int(input_len), + ) + device = "cpu" + + speech_token, speech_token_len = extract_speech_token(audio, self.speech_tokenizer, device) + speech_feat, speech_feat_len = extract_speech_feat(audio, self.feat_extractor, device) + + if config.sample_rate == 24000: + token_len = min(int(speech_feat.shape[1] / 2), speech_token.shape[1]) + speech_feat, speech_feat_len[:] = speech_feat[:, : 2 * token_len], 2 * token_len + speech_token, speech_token_len[:] = speech_token[:, :token_len], token_len + + embedding = extract_spk_embedding(audio, self.campplus_session, device) + + ft = BatchFeature( + { + "input_ids": input_ids, + "speech_feat": speech_feat, + "speech_token": speech_token, + "speech_token_len": [speech_token_len], + "embedding": embedding, + } + ) + + return ft + + def _get_mm_fields_config( + self, + hf_inputs: "BatchFeature", + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return { + "speech_feat": MultiModalFieldConfig.batched("audio"), + "speech_token": MultiModalFieldConfig.batched("audio"), + "speech_token_len": MultiModalFieldConfig.batched("audio"), + "embedding": MultiModalFieldConfig.batched("audio"), + } + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + def insertion_end(item_idx): + # TODO: Think if this can be done better + # sos + task + audio token ... ideally this needs to be split into + # two start and end but somehow I couldn't pass two of these + # wutg target .start() and .end() + token_len = out_mm_kwargs["audio"][0]["speech_token_len"].data[0].item() + return [1] * (1 + 1 + token_len) + + return [ + PromptInsertion( + modality="audio", + target=PromptIndexTargets.start(), + insertion=insertion_end, + ), + ] + + +class CosyVoice3DummyInputsBuilder(BaseDummyInputsBuilder[CosyVoice3MultiModalProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return "Hello, this is a test of the CosyVoice3 system capability." + + def get_dummy_mm_data( + self, seq_len: int, mm_counts: Mapping[str, int], mm_options: Mapping[str, BaseDummyOptions] | None = None + ) -> MultiModalDataDict: + num_audios = mm_counts.get("audio") + max_prompt_seconds = 30 + prompt_sample_rate = 24000 + target_audio_length = max_prompt_seconds * prompt_sample_rate + + audio_overrides = mm_options.get("audio") if mm_options else None + mm_data = { + "audio": ( + self._get_dummy_audios( + length=target_audio_length, + num_audios=num_audios, + overrides=audio_overrides, + )[0], + 24000, + ), + } + return mm_data + + def get_dummy_processor_inputs( + self, seq_len: int, mm_counts: Mapping[str, int], mm_options: Mapping[str, BaseDummyOptions] | None = None + ) -> ProcessorInputs: + inputs = super().get_dummy_processor_inputs(seq_len, mm_counts, mm_options) + inputs.hf_processor_mm_kwargs = {"prompt_text": "Testing my voices. Why should I not?"} + return inputs + + +@MULTIMODAL_REGISTRY.register_processor( + CosyVoice3MultiModalProcessor, + info=CosyVoice3MultiModalProcessingInfo, + dummy_inputs=CosyVoice3DummyInputsBuilder, +) +class CosyVoice3Model( + nn.Module, + SupportsMultiModal, +): + supports_multimodal_raw_input_only = True + supports_multimodal = True + requires_raw_input_tokens = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.config = vllm_config.model_config.hf_config + self.have_multimodal_outputs = True + self.model_stage = vllm_config.model_config.model_stage + self.model_dir = vllm_config.model_config.model + self.model = None + if self.model_stage == "talker": + # Initialize talker stage (text to speech tokens) + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_talker import CosyVoice3LM, VLLMQwen2Encoder + + llm_vllm_config = self._create_llm_vllm_config(vllm_config) + llm = VLLMQwen2Encoder(vllm_config=llm_vllm_config, prefix="model") + self.talker = CosyVoice3LM( + llm_input_size=self.config.llm["llm_input_size"], + llm_output_size=self.config.llm["llm_output_size"], + speech_token_size=self.config.llm["speech_token_size"], + llm=llm, + length_normalized_loss=self.config.llm["length_normalized_loss"], + lsm_weight=self.config.llm["lsm_weight"], + mix_ratio=self.config.llm["mix_ratio"], + ) + # KV cache is now managed externally by vLLM's PagedAttention + # No need for self.llm_cache + self.model = self.talker + elif self.model_stage == "code2wav": + # Initialize code2wav stage (flow matching + vocoder) + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + self.code2wav = CosyVoice3Code2Wav(self.config) + self.model = self.code2wav.flow_model + self.hift = self.code2wav.hift + + # 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 + else: + raise ValueError(f"Model stage not supported {self.model_stage}") + + def _create_llm_vllm_config(self, parent_config: VllmConfig) -> VllmConfig: + """Create VllmConfig for the inner Qwen2 LLM. + + This creates a modified VllmConfig with the Qwen2 HF config loaded from + the pretrained model directory. The cache config is inherited from the parent + to enable PagedAttention with the same memory configuration. + """ + from transformers import Qwen2Config + + qwen_config_path = os.path.join(self.model_dir, self.config.llm["llm"]["pretrain_path"]) + qwen_hf_config = Qwen2Config.from_pretrained(qwen_config_path) + + # Use parent's cache config - critical for PagedAttention to work correctly + return parent_config.with_hf_config(qwen_hf_config, architectures=["Qwen2Model"]) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput) -> torch.Tensor | None: + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if self.model_stage == "talker": + logits = self.model.llm_decoder(hidden_states) + vocab_size = self.config.vocab_size + pad_size = vocab_size - logits.size(-1) + pad_shape = logits.shape[:-1] + (pad_size,) + pad = logits.new_full(pad_shape, float("-inf")) + eos_token_val = logits[..., self.config.llm["eos_token_id"]].clone() + logits[..., -200:] = float("-inf") + logits[..., self.config.llm["eos_token_id"]] = eos_token_val + logits = torch.cat([logits, pad], dim=-1) + return logits + else: + raise RuntimeError(f"compute_logits is only valid for {self.model_stage}.") + + def embed_multimodal(self, **kwargs: object) -> torch.Tensor: + if self.model_stage == "talker": + speech_token = kwargs["speech_token"] + speech_token_emb = self.model.speech_embedding(speech_token) + return speech_token_emb + else: + raise RuntimeError(f"embed_multimodal is only valid for {self.model_stage}.") + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + if self.model_stage == "talker": + if is_multimodal is not None and any(is_multimodal): + embed_tokens = self.model.llm.model.embed_tokens(input_ids) + sos = self.model.speech_embedding.weight[self.model.sos].reshape(1, -1) + task_id = self.model.speech_embedding.weight[self.model.task_id].reshape(1, -1) + prompt_speech_token_emb = multimodal_embeddings[0] + pstoken_len = prompt_speech_token_emb.shape[0] # Get length from tensor shape + embed_tokens = torch.cat( + [sos, embed_tokens[2 + pstoken_len :], task_id, prompt_speech_token_emb], dim=0 + ) + else: + embed_tokens = self.model.speech_embedding.weight[input_ids] + return embed_tokens + elif self.model_stage == "code2wav": + assert input_ids.dim() == 1 + hidden = int(self.config.hidden_size) + return torch.zeros( + (input_ids.shape[0], hidden), + ) + else: + raise RuntimeError(f"embed_input_ids is not valid for {self.model_stage}.") + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + additional_information: dict[str, object] | None = None, + **kwargs: object, + ) -> OmniOutput: + if self.model_stage == "talker": + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + + # [total_tokens, hidden] + hidden_states = self.model.llm(inputs_embeds, positions) + + multimodal_outputs = {} + + if "speech_token" in kwargs: + # Wrap in lists to pass through gpu_ar_model_runner shape filtering + multimodal_outputs = { + "speech_token": [kwargs.get("speech_token")], + "speech_feat": [kwargs.get("speech_feat")], + "embedding": [kwargs.get("embedding")], + } + + return OmniOutput(text_hidden_states=hidden_states, multimodal_outputs=multimodal_outputs) + elif self.model_stage == "code2wav": + runtime_info = kwargs.get("runtime_additional_information", []) + if not runtime_info: + length = 30 * 24000 + audio = np.zeros((length,)) + return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audio}) + + # Remove the last eos token and add batch dimension + token = input_ids[..., :-1].unsqueeze(0) + + # Generate audio using code2wav + tts_speech = self.code2wav( + token=token, + prompt_token=runtime_info[0]["speech_token"][:1], + prompt_feat=runtime_info[0]["speech_feat"][:1], + embedding=runtime_info[0]["embedding"][:1], + n_timesteps=10, + ) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"audio": tts_speech}, + ) + else: + raise ValueError(f"Unsupported model_stage: {self.model_stage}") + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + if self.model_stage == "talker": + # Load weights for text to speech LM stage using vLLM's weight loading + llm_weight_path = os.path.join(self.model_dir, "llm.pt") + device = next(self.parameters()).device + checkpoint = torch.load(llm_weight_path, map_location=device) + + # 1. Load Qwen2 model weights into vLLM's Qwen2Model + # The checkpoint has prefix "llm.model.model." for the transformer weights + # vLLM's Qwen2Model expects just the model structure without extra prefixes + qwen_weights = [] + for name, weight in checkpoint.items(): + if name.startswith("llm.model.model."): + # Strip prefix: llm.model.model.X -> X (for vLLM's Qwen2Model) + vllm_name = name.replace("llm.model.model.", "") + qwen_weights.append((vllm_name, weight)) + + # Use vLLM's built-in load_weights which handles stacked params + # (q_proj+k_proj+v_proj -> qkv_proj, gate_proj+up_proj -> gate_up_proj) + self.model.llm.model.load_weights(iter(qwen_weights)) + + # 2. Load CosyVoice3LM-specific weights (speech_embedding, llm_decoder) + speech_emb_state = { + k.replace("speech_embedding.", ""): v + for k, v in checkpoint.items() + if k.startswith("speech_embedding.") + } + self.model.speech_embedding.load_state_dict(speech_emb_state) + + llm_decoder_state = { + k.replace("llm_decoder.", ""): v for k, v in checkpoint.items() if k.startswith("llm_decoder.") + } + self.model.llm_decoder.load_state_dict(llm_decoder_state) + + self.model.to(device).eval() + elif self.model_stage == "code2wav": + # Load weights for code2wav stage (flow + hift) + device = next(self.parameters()).device + self.code2wav.load_weights(self.model_dir, device) + else: + raise ValueError(f"{self.model_stage} not supported yet!") diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py new file mode 100644 index 000000000..f5e0d04a8 --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +CosyVoice3 Code2Wav Stage - Converts speech tokens to audio waveforms. + +This module contains the code2wav (token-to-waveform) stage which uses: +1. DiT (Diffusion Transformer) with optimized attention backends +2. CFM (Conditional Flow Matching) for mel spectrogram generation +3. HiFiGAN vocoder for waveform synthesis +""" + +from __future__ import annotations + +import numpy as np +import torch +import torch.nn as nn +from omegaconf import DictConfig +from torch.nn import functional as F +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import ( + CausalConditionalCFM, + CausalMaskedDiffWithDiT, +) +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( + CausalConvRNNF0Predictor, + CausalHiFTGenerator, +) +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer +from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config +from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask + +logger = init_logger(__name__) + + +class CosyVoice3Code2Wav(nn.Module): + """CosyVoice3 Code2Wav stage for token-to-waveform conversion. + + This class encapsulates: + - Flow matching decoder with DiT backbone (using diffusion attention) + - HiFiGAN vocoder for mel-to-waveform conversion + """ + + def __init__(self, config: CosyVoice3Config): + super().__init__() + self.config = config + + # Build flow matching components + pre_lookahead_layer = PreLookaheadLayer(**config.flow["pre_lookahead_layer"]) + + decoder_cfg = config.flow["decoder"] + cfm_params = DictConfig(decoder_cfg["cfm_params"]) + + # DiT estimator using diffusion attention (Flash/Sage/SDPA backends) + estimator = DiT(**decoder_cfg["estimator"]) + + decoder = CausalConditionalCFM( + in_channels=decoder_cfg["in_channels"], + estimator=estimator, + cfm_params=cfm_params, + n_spks=decoder_cfg["n_spks"], + spk_emb_dim=decoder_cfg["spk_emb_dim"], + ) + + self.flow_model = CausalMaskedDiffWithDiT( + input_size=config.flow["input_size"], + output_size=config.flow["output_size"], + spk_embed_dim=config.flow["spk_embed_dim"], + output_type=config.flow["output_type"], + vocab_size=config.flow["vocab_size"], + input_frame_rate=config.flow["input_frame_rate"], + only_mask_loss=config.flow["only_mask_loss"], + token_mel_ratio=config.flow["token_mel_ratio"], + pre_lookahead_len=config.flow["pre_lookahead_len"], + pre_lookahead_layer=pre_lookahead_layer, + decoder=decoder, + ) + + # Build HiFiGAN vocoder + f0_predictor = CausalConvRNNF0Predictor( + num_class=config.hift["f0_predictor"]["num_class"], + in_channels=config.hift["f0_predictor"]["in_channels"], + cond_channels=config.hift["f0_predictor"]["cond_channels"], + ) + + self.hift = CausalHiFTGenerator( + in_channels=config.hift["in_channels"], + base_channels=config.hift["base_channels"], + nb_harmonics=config.hift["nb_harmonics"], + sampling_rate=config.hift["sampling_rate"], + nsf_alpha=config.hift["nsf_alpha"], + nsf_sigma=config.hift["nsf_sigma"], + nsf_voiced_threshold=config.hift["nsf_voiced_threshold"], + upsample_rates=config.hift["upsample_rates"], + upsample_kernel_sizes=config.hift["upsample_kernel_sizes"], + istft_params=config.hift["istft_params"], + resblock_kernel_sizes=config.hift["resblock_kernel_sizes"], + resblock_dilation_sizes=config.hift["resblock_dilation_sizes"], + source_resblock_kernel_sizes=config.hift["source_resblock_kernel_sizes"], + source_resblock_dilation_sizes=config.hift["source_resblock_dilation_sizes"], + lrelu_slope=config.hift["lrelu_slope"], + audio_limit=config.hift["audio_limit"], + conv_pre_look_right=config.hift["conv_pre_look_right"], + f0_predictor=f0_predictor, + ) + # Run hift in float32 to avoid dtype mismatches in internal ops + self.hift = self.hift.float() + + # Streaming/chunking parameters + self.token_overlap_len = 20 + self.mel_overlap_len = int(self.token_overlap_len / self.flow_model.input_frame_rate * 22050 / 256) + self.mel_window = np.hamming(2 * self.mel_overlap_len) + self.mel_cache_len = 20 + self.source_cache_len = int(self.mel_cache_len * 256) + self.speech_window = np.hamming(2 * self.source_cache_len) + + @property + def input_frame_rate(self) -> int: + """Input frame rate from flow model.""" + return self.flow_model.input_frame_rate + + @property + def token_mel_ratio(self) -> int: + """Token to mel ratio.""" + return self.flow_model.token_mel_ratio + + @property + def output_size(self) -> int: + """Output mel dimension.""" + return self.flow_model.output_size + + @property + def input_embedding(self) -> nn.Embedding: + """Token embedding layer.""" + return self.flow_model.input_embedding + + @property + def pre_lookahead_layer(self) -> nn.Module: + """Pre-lookahead layer.""" + return self.flow_model.pre_lookahead_layer + + @property + def decoder(self) -> nn.Module: + """Flow matching decoder.""" + return self.flow_model.decoder + + @property + def spk_embed_affine_layer(self) -> nn.Linear: + """Speaker embedding affine layer.""" + return self.flow_model.spk_embed_affine_layer + + @torch.inference_mode() + def forward( + self, + token: torch.Tensor, + prompt_token: torch.Tensor, + prompt_feat: torch.Tensor, + embedding: torch.Tensor, + n_timesteps: int = 10, + ) -> torch.Tensor: + """Generate audio waveform from speech tokens. + + Args: + token: Speech tokens from talker stage [batch, seq_len] + prompt_token: Prompt speech tokens [batch, prompt_len] + prompt_feat: Prompt mel features [batch, feat_len, mel_dim] + embedding: Speaker embedding [batch, spk_dim] + n_timesteps: Number of diffusion steps + + Returns: + Audio waveform [batch, 1, audio_len] + """ + device = token.device + dtype = next(self.flow_model.parameters()).dtype + + # Normalize and project speaker embedding + embedding = embedding.to(device=device, dtype=dtype) + embedding = F.normalize(embedding, dim=1) + embedding = self.spk_embed_affine_layer(embedding) + + # Prepare tokens + prompt_token = prompt_token.to(device=device) + token_len1, token_len2 = prompt_token.shape[1], token.shape[1] + prompt_token_len = torch.tensor([token_len1], device=device, dtype=torch.int32) + token_len = torch.tensor([token_len2], device=device, dtype=torch.int32) + + # Concatenate prompt and target tokens + full_token = torch.cat([prompt_token, token], dim=1) + full_token_len = prompt_token_len + token_len + + # Create mask + mask = (~make_pad_mask(full_token_len)).unsqueeze(-1).to(embedding) + + # Token embedding + token_emb = self.input_embedding(torch.clamp(full_token, min=0)) * mask + + # Pre-lookahead processing + h = self.pre_lookahead_layer(token_emb) + h = h.repeat_interleave(self.token_mel_ratio, dim=1) + + # Calculate mel lengths + mel_len1 = prompt_feat.shape[1] + mel_len2 = h.shape[1] - mel_len1 + + # Build conditioning + conds = torch.zeros( + [1, mel_len1 + mel_len2, self.output_size], + device=device, + dtype=h.dtype, + ) + conds[:, :mel_len1] = prompt_feat + conds = conds.transpose(1, 2) + + # Create mel mask + mel_mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h) + + # Run flow matching decoder + feat, _ = self.decoder( + mu=h.transpose(1, 2).contiguous(), + mask=mel_mask.unsqueeze(1), + spks=embedding, + cond=conds, + n_timesteps=n_timesteps, + ) + + # Extract generated portion (after prompt) + feat = feat[:, :, mel_len1:] + + # Run vocoder + hift_weight = self.hift.m_source.l_linear.weight + tts_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype) + + if tts_mel.shape[-1] == 0: + tts_speech = torch.zeros( + (tts_mel.shape[0], 1, 0), + device=tts_mel.device, + dtype=tts_mel.dtype, + ) + else: + tts_speech, _ = self.hift.inference(speech_feat=tts_mel) + + return tts_speech + + def load_weights(self, model_dir: str, device: torch.device) -> None: + """Load flow.pt and hift.pt weights. + + Args: + model_dir: Model directory containing flow.pt and hift.pt + device: Device to load weights to + """ + import os + + # Load flow weights + flow_path = os.path.join(model_dir, "flow.pt") + self.flow_model.load_state_dict(torch.load(flow_path, map_location=device), strict=True) + self.flow_model.to(device).eval() + logger.info(f"Loaded flow weights from {flow_path}") + + # Load hift weights + hift_path = os.path.join(model_dir, "hift.pt") + hift_state_dict = { + k.replace("generator.", ""): v for k, v in torch.load(hift_path, map_location=device).items() + } + self.hift.load_state_dict(hift_state_dict, strict=True) + self.hift.to(device).eval() + logger.info(f"Loaded hift weights from {hift_path}") diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_talker.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_talker.py new file mode 100644 index 000000000..aeea2f8cd --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_talker.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/FunAudioLLM/CosyVoice/blob/main/cosyvoice/llm/llm.py +from collections.abc import Callable + +import torch +from torch import nn +from vllm.config import VllmConfig +from vllm.model_executor.models.qwen2 import Qwen2Model + + +class VLLMQwen2Encoder(torch.nn.Module): + """Qwen2 encoder using vLLM's Qwen2Model with external KV cache management. + + This replaces the HuggingFace Qwen2ForCausalLM with vLLM's optimized implementation + that uses PagedAttention and external KV cache via ForwardContext. + """ + + def __init__(self, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.model = Qwen2Model(vllm_config=vllm_config, prefix=prefix) + self.hidden_size = vllm_config.model_config.hf_config.hidden_size + + def forward(self, inputs_embeds: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Forward pass using vLLM's attention with external KV cache. + + Args: + inputs_embeds: Input embeddings [total_tokens, hidden_size] or [batch, seq, hidden] + positions: Position tensor for RoPE [total_tokens] + + Returns: + hidden_states: Output hidden states [total_tokens, hidden_size] + """ + # vLLM model expects flattened tensors [total_tokens, hidden_size] + if inputs_embeds.dim() == 3: + inputs_flat = inputs_embeds.view(-1, self.hidden_size) + else: + inputs_flat = inputs_embeds + positions_flat = positions.view(-1) + + # KV cache managed externally via ForwardContext (set by GPUARModelRunner) + # input_ids is required but ignored when inputs_embeds is provided + hidden_states = self.model( + input_ids=torch.zeros(inputs_flat.size(0), dtype=torch.long, device=inputs_flat.device), + positions=positions_flat, + intermediate_tensors=None, + inputs_embeds=inputs_flat, + ) + return hidden_states + + +class TransformerLM(torch.nn.Module): + def __init__( + self, + text_encoder_input_size: int, + llm_input_size: int, + llm_output_size: int, + text_token_size: int, + speech_token_size: int, + text_encoder: torch.nn.Module, + llm: torch.nn.Module, + length_normalized_loss: bool = True, + lsm_weight: float = 0.0, + spk_embed_dim: int = 192, + ): + super().__init__() + self.llm_input_size = llm_input_size + self.speech_token_size = speech_token_size + # 1. build text token inputs related modules + self.text_embedding = torch.nn.Embedding(text_token_size, text_encoder_input_size) + self.text_encoder = text_encoder + self.text_encoder_affine_layer = nn.Linear(self.text_encoder.output_size(), llm_input_size) + + # 2. build speech token language model related modules + self.sos = 0 + self.task_id = 1 + self.eos_token = self.speech_token_size + self.llm_embedding = torch.nn.Embedding(2, llm_input_size) + self.llm = llm + self.llm_decoder = nn.Linear(llm_output_size, speech_token_size + 1) + + # 3. [Optional] build speech token related modules + self.speech_embedding = torch.nn.Embedding(speech_token_size, llm_input_size) + self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, llm_input_size) + + +class Qwen2LM(TransformerLM): + def __init__( + self, + llm_input_size: int, + llm_output_size: int, + speech_token_size: int, + llm: torch.nn.Module, + sampling: Callable, + length_normalized_loss: bool = True, + lsm_weight: float = 0.0, + mix_ratio: list[int] = [5, 15], + ): + torch.nn.Module.__init__(self) + self.llm_input_size = llm_input_size + self.llm_output_size = llm_output_size + self.speech_token_size = speech_token_size + # 2. build speech token language model related modules + self.sos = 0 + self.task_id = 1 + self.eos_token = speech_token_size + self.fill_token = speech_token_size + 2 + + self.llm_embedding = torch.nn.Embedding(2, llm_input_size) + self.llm = llm + self.llm_decoder = nn.Linear(llm_output_size, speech_token_size + 3) + + # 3. [Optional] build speech token related modules + self.speech_embedding = torch.nn.Embedding(speech_token_size + 3, llm_input_size) + + # 4. sampling method + self.sampling = sampling + self.mix_ratio = mix_ratio + + # 5. vllm related + self.stop_token_ids = [speech_token_size + i for i in range(3)] + self.vllm_output_queue = {} + + +class CosyVoice3LM(Qwen2LM): + def __init__( + self, + llm_input_size: int, + llm_output_size: int, + speech_token_size: int, + llm: torch.nn.Module, + length_normalized_loss: bool = True, + lsm_weight: float = 0.0, + mix_ratio: list[int] = [5, 15], + ): + torch.nn.Module.__init__(self) + self.llm_input_size = llm_input_size + self.llm_output_size = llm_output_size + self.speech_token_size = speech_token_size + # 2. build speech token language model related modules + self.sos = speech_token_size + 0 + self.eos_token = speech_token_size + 1 + self.task_id = speech_token_size + 2 + self.fill_token = speech_token_size + 3 + + self.llm = llm + self.llm_decoder = nn.Linear(llm_output_size, speech_token_size + 200, bias=False) + + # 3. [Optional] build speech token related modules + self.speech_embedding = torch.nn.Embedding(speech_token_size + 200, llm_input_size) + + # 4. sampling method + # self.sampling = sampling + self.mix_ratio = mix_ratio + + # 5. vllm related + self.stop_token_ids = [speech_token_size + i for i in range(200)] + self.vllm_output_queue = {} diff --git a/vllm_omni/model_executor/models/cosyvoice3/tokenizer.py b/vllm_omni/model_executor/models/cosyvoice3/tokenizer.py new file mode 100644 index 000000000..6049df61e --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/tokenizer.py @@ -0,0 +1,356 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from functools import cache + +import torch +from transformers import AutoTokenizer + + +class CosyVoice2Tokenizer: + def __init__(self, token_path, skip_special_tokens=True): + super().__init__() + # NOTE: non-chat model, all these special tokens keep randomly initialized. + special_tokens = { + "eos_token": "<|endoftext|>", + "pad_token": "<|endoftext|>", + "additional_special_tokens": [ + "<|im_start|>", + "<|im_end|>", + "<|endofprompt|>", + "[breath]", + "", + "", + "[noise]", + "[laughter]", + "[cough]", + "[clucking]", + "[accent]", + "[quick_breath]", + "", + "", + "[hissing]", + "[sigh]", + "[vocalized-noise]", + "[lipsmack]", + "[mn]", + ], + } + self.special_tokens = special_tokens + self.tokenizer = AutoTokenizer.from_pretrained(token_path) + self.tokenizer.add_special_tokens(special_tokens) + self.skip_special_tokens = skip_special_tokens + + def encode(self, text, **kwargs): + tokens = self.tokenizer([text], return_tensors="pt") + tokens = tokens["input_ids"][0].cpu().tolist() + return tokens + + def decode(self, tokens): + tokens = torch.tensor(tokens, dtype=torch.int64) + text = self.tokenizer.batch_decode([tokens], skip_special_tokens=self.skip_special_tokens)[0] + return text + + +class CosyVoice3Tokenizer(CosyVoice2Tokenizer): + def __init__(self, token_path, skip_special_tokens=True): + # NOTE: non-chat model, all these special tokens keep randomly initialized. + special_tokens = { + "eos_token": "<|endoftext|>", + "pad_token": "<|endoftext|>", + "additional_special_tokens": [ + "<|im_start|>", + "<|im_end|>", + "<|endofprompt|>", + "[breath]", + "", + "", + "[noise]", + "[laughter]", + "[cough]", + "[clucking]", + "[accent]", + "[quick_breath]", + "", + "", + "[hissing]", + "[sigh]", + "[vocalized-noise]", + "[lipsmack]", + "[mn]", + "<|endofsystem|>", + "[AA]", + "[AA0]", + "[AA1]", + "[AA2]", + "[AE]", + "[AE0]", + "[AE1]", + "[AE2]", + "[AH]", + "[AH0]", + "[AH1]", + "[AH2]", + "[AO]", + "[AO0]", + "[AO1]", + "[AO2]", + "[AW]", + "[AW0]", + "[AW1]", + "[AW2]", + "[AY]", + "[AY0]", + "[AY1]", + "[AY2]", + "[B]", + "[CH]", + "[D]", + "[DH]", + "[EH]", + "[EH0]", + "[EH1]", + "[EH2]", + "[ER]", + "[ER0]", + "[ER1]", + "[ER2]", + "[EY]", + "[EY0]", + "[EY1]", + "[EY2]", + "[F]", + "[G]", + "[HH]", + "[IH]", + "[IH0]", + "[IH1]", + "[IH2]", + "[IT]", + "[IT0]", + "[IT1]", + "[IT2]", + "[JH]", + "[K]", + "[L]", + "[M]", + "[N]", + "[NG]", + "[OW]", + "[OW0]", + "[OW1]", + "[OW2]", + "[OY]", + "[OY0]", + "[OY1]", + "[OY2]", + "[P]", + "[R]", + "[S]", + "[SH]", + "[T]", + "[TH]", + "[UH]", + "[UH0]", + "[UH1]", + "[UH2]", + "[UW]", + "[UW0]", + "[UW1]", + "[UW2]", + "[V]", + "[W]", + "[Y]", + "[Z]", + "[ZH]", + "[a]", + "[ai]", + "[an]", + "[ang]", + "[ao]", + "[b]", + "[c]", + "[ch]", + "[d]", + "[e]", + "[ei]", + "[en]", + "[eng]", + "[f]", + "[g]", + "[h]", + "[i]", + "[ian]", + "[in]", + "[ing]", + "[iu]", + "[ià]", + "[iàn]", + "[iàng]", + "[iào]", + "[iá]", + "[ián]", + "[iáng]", + "[iáo]", + "[iè]", + "[ié]", + "[iòng]", + "[ióng]", + "[iù]", + "[iú]", + "[iā]", + "[iān]", + "[iāng]", + "[iāo]", + "[iē]", + "[iě]", + "[iōng]", + "[iū]", + "[iǎ]", + "[iǎn]", + "[iǎng]", + "[iǎo]", + "[iǒng]", + "[iǔ]", + "[j]", + "[k]", + "[l]", + "[m]", + "[n]", + "[o]", + "[ong]", + "[ou]", + "[p]", + "[q]", + "[r]", + "[s]", + "[sh]", + "[t]", + "[u]", + "[uang]", + "[ue]", + "[un]", + "[uo]", + "[uà]", + "[uài]", + "[uàn]", + "[uàng]", + "[uá]", + "[uái]", + "[uán]", + "[uáng]", + "[uè]", + "[ué]", + "[uì]", + "[uí]", + "[uò]", + "[uó]", + "[uā]", + "[uāi]", + "[uān]", + "[uāng]", + "[uē]", + "[uě]", + "[uī]", + "[uō]", + "[uǎ]", + "[uǎi]", + "[uǎn]", + "[uǎng]", + "[uǐ]", + "[uǒ]", + "[vè]", + "[w]", + "[x]", + "[y]", + "[z]", + "[zh]", + "[à]", + "[ài]", + "[àn]", + "[àng]", + "[ào]", + "[á]", + "[ái]", + "[án]", + "[áng]", + "[áo]", + "[è]", + "[èi]", + "[èn]", + "[èng]", + "[èr]", + "[é]", + "[éi]", + "[én]", + "[éng]", + "[ér]", + "[ì]", + "[ìn]", + "[ìng]", + "[í]", + "[ín]", + "[íng]", + "[ò]", + "[òng]", + "[òu]", + "[ó]", + "[óng]", + "[óu]", + "[ù]", + "[ùn]", + "[ú]", + "[ún]", + "[ā]", + "[āi]", + "[ān]", + "[āng]", + "[āo]", + "[ē]", + "[ēi]", + "[ēn]", + "[ēng]", + "[ě]", + "[ěi]", + "[ěn]", + "[ěng]", + "[ěr]", + "[ī]", + "[īn]", + "[īng]", + "[ō]", + "[ōng]", + "[ōu]", + "[ū]", + "[ūn]", + "[ǎ]", + "[ǎi]", + "[ǎn]", + "[ǎng]", + "[ǎo]", + "[ǐ]", + "[ǐn]", + "[ǐng]", + "[ǒ]", + "[ǒng]", + "[ǒu]", + "[ǔ]", + "[ǔn]", + "[ǘ]", + "[ǚ]", + "[ǜ]", + ], + } + self.special_tokens = special_tokens + self.tokenizer = AutoTokenizer.from_pretrained(token_path) + self.tokenizer.add_special_tokens(special_tokens) + self.skip_special_tokens = skip_special_tokens + + +@cache +def get_qwen_tokenizer(token_path: str, skip_special_tokens: bool, version: str = "cosyvoice3"): + if version == "cosyvoice2": + return CosyVoice2Tokenizer(token_path=token_path, skip_special_tokens=skip_special_tokens) + elif version == "cosyvoice3": + return CosyVoice3Tokenizer(token_path=token_path, skip_special_tokens=skip_special_tokens) + else: + raise ValueError diff --git a/vllm_omni/model_executor/models/cosyvoice3/utils.py b/vllm_omni/model_executor/models/cosyvoice3/utils.py new file mode 100644 index 000000000..e1310cd3b --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/utils.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os +from functools import cache, lru_cache + +import numpy as np +import torch +import torch.nn.functional as F +import torchaudio +import torchaudio.compliance.kaldi as kaldi +from librosa.filters import mel as librosa_mel_fn + +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) + + +def spectral_normalize_torch(magnitudes): + output = dynamic_range_compression_torch(magnitudes) + return output + + +@lru_cache +def _get_mel_basis( + sampling_rate: int, + n_fft: int, + num_mels: int, + fmin: float, + fmax: float | None, + device_str: str, +) -> torch.Tensor: + mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax) + return torch.from_numpy(mel).float().to(torch.device(device_str)) + + +@lru_cache +def _get_hann_window(win_size: int, device_str: str) -> torch.Tensor: + return torch.hann_window(win_size).to(torch.device(device_str)) + + +def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False): + device_str = str(y.device) + mel = _get_mel_basis( + int(sampling_rate), + int(n_fft), + int(num_mels), + float(fmin), + fmax, + device_str, + ) + window = _get_hann_window(int(win_size), device_str) + + y = torch.nn.functional.pad( + y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode="reflect" + ) + y = y.squeeze(1) + + spec = torch.view_as_real( + torch.stft( + y, + n_fft, + hop_length=hop_size, + win_length=win_size, + window=window, + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + ) + + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) + + spec = torch.matmul(mel, spec) + spec = spectral_normalize_torch(spec) + + return spec + + +def load_wav(wav, target_sr, min_sr=16000): + if not isinstance(wav, tuple): + speech, sample_rate = torchaudio.load(wav, backend="soundfile") + else: + speech, sample_rate = wav + if isinstance(speech, np.ndarray): + speech = torch.from_numpy(speech).float().unsqueeze(0) + + if sample_rate != target_sr: + if sample_rate < min_sr: + raise ValueError( + f"Audio sample rate {sample_rate} Hz is too low. " + f"Minimum required: {min_sr} Hz, target: {target_sr} Hz. " + f"Please provide audio with sample rate >= {min_sr} Hz." + ) + speech = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_sr)(speech) + + speech = speech.to(dtype=torch.float32) + return speech + + +def extract_speech_feat(prompt_wav, feat_extractor, device): + speech = load_wav(prompt_wav, 24000) + speech_feat = feat_extractor(speech).squeeze(dim=0).transpose(0, 1).to(device) + speech_feat = speech_feat.unsqueeze(dim=0) + speech_feat_len = torch.tensor([speech_feat.shape[1]], dtype=torch.int32).to(device) + return speech_feat, speech_feat_len + + +# Adopted from https://github.com/openai/whisper/blob/main/whisper/audio.py + + +def exact_div(x, y): + assert x % y == 0 + return x // y + + +@cache +def mel_filters(device, n_mels: int) -> torch.Tensor: + """ + load the mel filterbank matrix for projecting STFT into a Mel spectrogram. + Allows decoupling librosa dependency; saved using: + + np.savez_compressed( + "mel_filters.npz", + mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80), + mel_128=librosa.filters.mel(sr=16000, n_fft=400, n_mels=128), + ) + """ + assert n_mels in {80, 128}, f"Unsupported n_mels: {n_mels}" + + filters_path = os.path.join(os.path.dirname(__file__), "assets", "mel_filters.npz") + if not os.path.exists(filters_path): + source_url = "https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz" + raise FileNotFoundError( + "Missing CosyVoice3 mel filter asset:\n" + f" {filters_path}\n" + "Download it manually from:\n" + f" {source_url}\n" + "Example:\n" + f" mkdir -p {os.path.dirname(filters_path)} && " + f"curl -L {source_url} -o {filters_path}" + ) + + with np.load(filters_path, allow_pickle=False) as f: + return torch.from_numpy(f[f"mel_{n_mels}"]).to(device) + + +def log_mel_spectrogram( + audio: str | np.ndarray | torch.Tensor, + n_mels: int = 80, + padding: int = 0, + device: str | torch.device | None = None, +): + """ + Compute the log-Mel spectrogram of + + Parameters + ---------- + audio: Union[str, np.ndarray, torch.Tensor], shape = (*) + The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz + + n_mels: int + The number of Mel-frequency filters, only 80 and 128 are supported + + padding: int + Number of zero samples to pad to the right + + device: Optional[Union[str, torch.device]] + If given, the audio tensor is moved to this device before STFT + + Returns + ------- + torch.Tensor, shape = (n_mels, n_frames) + A Tensor that contains the Mel spectrogram + """ + N_FFT = 400 + HOP_LENGTH = 160 + + if not torch.is_tensor(audio): + raise Exception(f"audio is not tensor {type(audio)}") + + if device is not None: + audio = audio.to(device) + if padding > 0: + audio = F.pad(audio, (0, padding)) + window = torch.hann_window(N_FFT).to(audio.device) + stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True) + magnitudes = stft[..., :-1].abs() ** 2 + + filters = mel_filters(audio.device, n_mels) + mel_spec = filters @ magnitudes + + log_spec = torch.clamp(mel_spec, min=1e-10).log10() + log_spec = torch.maximum(log_spec, log_spec.max() - 8.0) + log_spec = (log_spec + 4.0) / 4.0 + 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) + feat = feat - feat.mean(dim=0, keepdim=True) + embedding = ( + campplus_session.run(None, {campplus_session.get_inputs()[0].name: feat.unsqueeze(dim=0).cpu().numpy()})[0] + .flatten() + .tolist() + ) + embedding = torch.tensor([embedding]).to(device) + return embedding + + +def extract_text_token(text, tokenizer, allowed_special): + text_token = tokenizer.encode(text, allowed_special=allowed_special) + text_token = torch.tensor([text_token], dtype=torch.int32) + text_token_len = text_token.shape[1] + return text_token, text_token_len + + +def concat_text_with_prompt_ids( + text: torch.Tensor, + text_len: torch.Tensor, + prompt_text: torch.Tensor, + prompt_text_len: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + text = torch.concat([prompt_text, text], dim=1) + text_len = text_len + prompt_text_len + return text, text_len + + +def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: + """Make mask tensor containing indices of padded part. + + See description of make_non_pad_mask. + + Args: + lengths (torch.Tensor): Batch of lengths (B,). + Returns: + torch.Tensor: Mask tensor containing indices of padded part. + + Examples: + >>> lengths = [5, 3, 2] + >>> make_pad_mask(lengths) + masks = [[0, 0, 0, 0 ,0], + [0, 0, 0, 1, 1], + [0, 0, 1, 1, 1]] + """ + batch_size = lengths.size(0) + max_len = max_len if max_len > 0 else lengths.max().item() + seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device) + seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len) + seq_length_expand = lengths.unsqueeze(-1) + mask = seq_range_expand >= seq_length_expand + return mask diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index ec879448f..496dfe305 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -52,6 +52,11 @@ _OMNI_MODELS = { "qwen3_omni_code2wav", "Qwen3OmniMoeCode2Wav", ), + "CosyVoice3Model": ( + "cosyvoice3", + "cosyvoice3", + "CosyVoice3Model", + ), "MammothModa2Qwen2ForCausalLM": ( "mammoth_moda2", "mammoth_moda2", diff --git a/vllm_omni/model_executor/stage_configs/cosyvoice3.yaml b/vllm_omni/model_executor/stage_configs/cosyvoice3.yaml new file mode 100644 index 000000000..cbf2193de --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/cosyvoice3.yaml @@ -0,0 +1,44 @@ +# Stage config for running CosyVoice3 with 2 stage architecture +# Stage 0: Talker (text prompt → speech tokens) +# Stage 1: Code2Wav (flow matching → acoustic features → waveform) +# Right now I have coded up stage 2 in stage 1 only; can be split in future + +stage_args: + - stage_id: 0 + runtime: + devices: 0 + engine_args: + model_stage: talker + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + model_arch: CosyVoice3Model + trust_remote_code: true + gpu_memory_utilization: 0.4 + engine_output_type: latent # Output speech tokens for chunk aware flow matching + disable_hybrid_kv_cache_manager: true + enable_prefix_caching: false + enforce_eager: false + mm_processor_cache_gb: 0 + skip_mm_profiling: true + dtype: "float32" + + - stage_id: 1 + runtime: + devices: 0 + engine_args: + model_stage: code2wav + model_arch: CosyVoice3Model + trust_remote_code: true + worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + engine_output_type: latent + gpu_memory_utilization: 0.2 + enforce_eager: false # CUDA graphs don't work with dynamic runtime_info access + disable_hybrid_kv_cache_manager: true + enable_prefix_caching: false + skip_mm_profiling: true + dtype: "float32" + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.cosyvoice3.text2flow + final_output: true + final_output_type: audio diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py new file mode 100644 index 000000000..b7f21eca8 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py @@ -0,0 +1,32 @@ +from typing import Any + +from vllm.inputs import TextPrompt + +from vllm_omni.inputs.data import OmniTokensPrompt + + +def text2flow( + stage_list: list[Any], + engine_input_source: list[int], + prompt: OmniTokensPrompt | TextPrompt = None, + requires_multimodal_data: bool = True, +): + """Build stage-1 inputs by prefixing stage-0 prompt ids to its outputs.""" + source_stage_id = engine_input_source[0] + source_outputs = stage_list[source_stage_id].engine_outputs + + if not isinstance(prompt, list): + prompt = [prompt] + + source_output = source_outputs[0] + output = source_output.outputs[0] + + multi_modal_data = output.multimodal_output + if multi_modal_data is None: + raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}") + + output_ids = output.token_ids + prefix_ids = source_output.prompt_token_ids + multi_modal_data["prefix_ids"] = prefix_ids + engine_input = OmniTokensPrompt(prompt_token_ids=output_ids, additional_information=multi_modal_data) + return [engine_input]