[Model] Fun cosy voice3-0.5-b-2512 (#498)

Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Signed-off-by: dsinghvi <divyanshsinghvi@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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` |

pyproject.toml

@@ -203,3 +203,5 @@ extend-ignore-identifiers-re = [
     ".*[Oo]no_[Aa]nna.*",
     ".*cann.*",
 ]
+[tool.typos.default.extend-words]
+ue = "ue"

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

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

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]

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

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

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

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

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

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))

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

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,
+            },
+        }

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!")

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}")

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 = {}

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]",
+                "<strong>",
+                "</strong>",
+                "[noise]",
+                "[laughter]",
+                "[cough]",
+                "[clucking]",
+                "[accent]",
+                "[quick_breath]",
+                "<laughter>",
+                "</laughter>",
+                "[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]",
+                "<strong>",
+                "</strong>",
+                "[noise]",
+                "[laughter]",
+                "[cough]",
+                "[clucking]",
+                "[accent]",
+                "[quick_breath]",
+                "<laughter>",
+                "</laughter>",
+                "[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

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

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",

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

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]