# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for CosyVoice3 components."""

import types
from contextlib import contextmanager
from types import SimpleNamespace

import pytest
import torch
import torch.nn as nn

from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.config import set_current_diffusion_config
from vllm_omni.diffusion.data import AttentionConfig
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
    CausalHiFTGenerator,
    HiFTGenerator,
)

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


@contextmanager
def _force_torch_sdpa():
    """Pin TORCH_SDPA so CPU shape tests do not pick CUDA-only backends (FA3)."""
    od_config = SimpleNamespace(
        diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"),
        parallel_config=SimpleNamespace(ring_degree=1),
    )
    with set_current_diffusion_config(od_config):
        yield


def test_batch_flow_runtime_flags_are_opt_in(monkeypatch):
    from vllm_omni.model_executor.models.cosyvoice3.runtime import (
        cosyvoice3_batch_flow_debug,
        cosyvoice3_batch_flow_enabled,
    )

    monkeypatch.delenv("COSYVOICE3_BATCH_FLOW", raising=False)
    monkeypatch.delenv("COSYVOICE3_BATCH_FLOW_DEBUG", raising=False)
    assert not cosyvoice3_batch_flow_enabled()
    assert not cosyvoice3_batch_flow_debug()

    monkeypatch.setenv("COSYVOICE3_BATCH_FLOW", "1")
    monkeypatch.setenv("COSYVOICE3_BATCH_FLOW_DEBUG", "1")
    assert cosyvoice3_batch_flow_enabled()
    assert cosyvoice3_batch_flow_debug()


@pytest.fixture
def causal_hift():
    return CausalHiFTGenerator(
        base_channels=32,
        upsample_rates=[2, 2],
        upsample_kernel_sizes=[4, 4],
        source_resblock_kernel_sizes=[3, 3],
        source_resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]],
    )


def test_causal_hift_moves_stft_window_with_model(causal_hift):
    assert causal_hift.get_buffer("stft_window") is causal_hift.stft_window
    assert "stft_window" not in causal_hift.state_dict()
    causal_hift.to(dtype=torch.float64)
    assert causal_hift.stft_window.dtype == torch.float64


def test_causal_hift_stft_moves_window_to_input_device(causal_hift):
    waveform = torch.empty((1, 64), device="meta")

    real, imag = causal_hift._stft(waveform)

    assert real.device == waveform.device
    assert imag.device == waveform.device
    assert causal_hift.stft_window.device == waveform.device


class _MinimalF0Predictor(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.proj = nn.Conv1d(80, 1, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.proj(x).squeeze(1).abs()


@pytest.mark.parametrize("sampling_rate", [22050, 24000])
def test_noncausal_hift_inference_runs_with_real_sinegen(sampling_rate):
    """Base (non-causal) HiFTGenerator with a real SineGen/SineGen2 (not a
    test double) must still run: SourceModuleHnNSF/SineGen/SineGen2/_f02sine
    all branch on `self.causal` to return the pre-streaming (shorter) tuple
    contract for non-causal callers, and any one of them regressing back to
    always returning the streaming-length tuple breaks this immediately.
    """
    hift = HiFTGenerator(
        base_channels=32,
        sampling_rate=sampling_rate,
        upsample_rates=[8, 5, 3],
        upsample_kernel_sizes=[16, 11, 7],
        source_resblock_kernel_sizes=[7, 7, 11],
        source_resblock_dilation_sizes=[[1, 3, 5]] * 3,
        f0_predictor=_MinimalF0Predictor(),
    ).eval()

    speech, _ = hift.inference(torch.randn(1, 80, 20))

    assert speech.shape[0] == 1
    assert torch.isfinite(speech).all()


def test_causal_hift_routes_npu_transforms_to_cpu_fallback(causal_hift, monkeypatch):
    class NpuTensor:
        device = SimpleNamespace(type="npu")

    tensor = NpuTensor()
    stft_calls = []
    istft_calls = []

    def fake_stft(value):
        stft_calls.append(value)
        return ("real", "imag")

    def fake_istft(magnitude, phase):
        istft_calls.append((magnitude, phase))
        return "waveform"

    monkeypatch.setattr(causal_hift, "_stft_on_cpu", fake_stft)
    monkeypatch.setattr(causal_hift, "_istft_on_cpu", fake_istft)

    assert causal_hift._stft(tensor) == ("real", "imag")
    assert causal_hift._istft(tensor, tensor) == "waveform"
    assert stft_calls == [tensor]
    assert istft_calls == [(tensor, tensor)]


def test_causal_hift_stft_cpu_fallback_uses_float32(causal_hift, monkeypatch):
    original_stft = torch.stft

    def assert_cpu_float32_stft(input, *args, window, **kwargs):
        assert input.device.type == "cpu"
        assert input.dtype == torch.float32
        assert window.device.type == "cpu"
        assert window.dtype == torch.float32
        return original_stft(input, *args, window=window, **kwargs)

    monkeypatch.setattr(torch, "stft", assert_cpu_float32_stft)

    waveform = torch.randn(1, 64, dtype=torch.float64)
    real, imag = causal_hift._stft_on_cpu(waveform)

    assert real.device == waveform.device
    assert imag.device == waveform.device
    assert real.dtype == waveform.dtype
    assert imag.dtype == waveform.dtype


def test_causal_hift_istft_cpu_fallback_uses_float32(causal_hift, monkeypatch):
    original_istft = torch.istft

    def assert_cpu_complex64_istft(input, *args, window, **kwargs):
        assert input.device.type == "cpu"
        assert input.dtype == torch.complex64
        assert window.device.type == "cpu"
        assert window.dtype == torch.float32
        return original_istft(input, *args, window=window, **kwargs)

    monkeypatch.setattr(torch, "istft", assert_cpu_complex64_istft)

    waveform = torch.randn(1, 64, dtype=torch.float32)
    spec = torch.stft(waveform, 16, 4, 16, window=causal_hift.stft_window, return_complex=True)
    magnitude = torch.abs(spec).to(torch.float64)
    phase = torch.angle(spec).to(torch.float64)
    reconstructed = causal_hift._istft_on_cpu(magnitude, phase)

    assert reconstructed.device == magnitude.device
    assert reconstructed.dtype == magnitude.dtype
    torch.testing.assert_close(reconstructed, waveform.to(torch.float64), rtol=1e-5, atol=1e-5)


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

        with _force_torch_sdpa():
            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

    @pytest.mark.core_model
    @pytest.mark.cpu
    def test_forward_passes_2d_mask_without_casting_qkv(self):
        """The float32 fallback receives the per-row validity mask."""
        from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTAttention

        class CaptureAttention(nn.Module):
            def __init__(self):
                super().__init__()
                self.last_metadata = None
                self.last_dtype = None

            def forward(self, q, k, v, attn_metadata=None):
                self.last_metadata = attn_metadata
                self.last_dtype = q.dtype
                return torch.zeros_like(q)

        attention = DiTAttention(dim=16, heads=2, dim_head=8, dropout=0.0)
        capture = CaptureAttention()
        attention.attn = capture

        x = torch.randn(2, 5, 16)
        mask = torch.tensor([[True, True, True, False, False], [True, True, True, True, False]])

        out = attention(x, mask=mask)

        assert out.shape == x.shape
        assert capture.last_metadata is not None
        assert torch.equal(capture.last_metadata.attn_mask, mask)
        assert capture.last_dtype == x.dtype
        assert out.dtype == x.dtype
        assert torch.allclose(out[~mask], torch.zeros_like(out[~mask]))


class TestDiTBlock:
    """Tests for DiTBlock."""

    @pytest.fixture
    def block(self):
        from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTBlock

        with _force_torch_sdpa():
            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

        with _force_torch_sdpa():
            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

    @pytest.mark.core_model
    @pytest.mark.cpu
    def test_trt_estimator_uses_stream_dependencies(self, monkeypatch):
        from omegaconf import DictConfig

        from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import ConditionalCFM

        class FakeStream:
            def __init__(self, name, state):
                self.name = name
                self.state = state
                self.cuda_stream = hash(name)
                self.synchronize_calls = 0
                self.waited_on = []

            def synchronize(self):
                self.synchronize_calls += 1

            def wait_stream(self, stream):
                self.waited_on.append(stream)

        class FakeStreamContext:
            def __init__(self, stream):
                self.stream = stream

            def __enter__(self):
                self.previous = self.stream.state.current
                self.stream.state.current = self.stream
                return self.stream

            def __exit__(self, exc_type, exc, traceback):
                self.stream.state.current = self.previous

        state = SimpleNamespace(current=None)
        caller_stream = FakeStream("caller", state)
        estimator_stream = FakeStream("estimator", state)
        state.current = caller_stream
        monkeypatch.setattr(torch.cuda, "current_stream", lambda *args, **kwargs: state.current)
        stream_contexts = []

        def stream_context(stream):
            stream_contexts.append(stream)
            return FakeStreamContext(stream)

        monkeypatch.setattr(torch.cuda, "stream", stream_context)

        class FakeContext:
            def __init__(self):
                self.execute_stream = None

            def set_input_shape(self, name, shape):
                pass

            def set_tensor_address(self, name, address):
                pass

            def execute_async_v3(self, stream):
                self.execute_stream = stream
                return True

        class FakeEngine:
            @staticmethod
            def get_tensor_name(index):
                return f"tensor_{index}"

        context = FakeContext()

        class FakeEstimatorPool:
            io_dtype = torch.float32

            def __init__(self):
                self.released = []

            def acquire_estimator(self):
                return [context, estimator_stream], FakeEngine()

            def release_estimator(self, released_context, released_stream):
                self.released.append((released_context, released_stream))

        estimator_pool = FakeEstimatorPool()
        cfm = ConditionalCFM(
            in_channels=80,
            cfm_params=DictConfig(
                {
                    "sigma_min": 1e-6,
                    "solver": "euler",
                    "t_scheduler": "cosine",
                    "training_cfg_rate": 0.2,
                    "inference_cfg_rate": 0.7,
                }
            ),
            n_spks=1,
            spk_emb_dim=80,
            estimator=estimator_pool,
        )

        x = torch.randn(2, 80, 4)
        mask = torch.ones(2, 1, 4)
        mu = torch.randn(2, 80, 4)
        timestep = torch.randn(2)
        speakers = torch.randn(2, 80)
        condition = torch.randn(2, 80, 4)

        output = cfm.forward_estimator(x, mask, mu, timestep, speakers, condition)

        assert output.shape == x.shape
        assert caller_stream.synchronize_calls == 0
        assert estimator_stream.synchronize_calls == 0
        assert estimator_stream.waited_on == [caller_stream]
        assert caller_stream.waited_on == [estimator_stream]
        assert stream_contexts == [estimator_stream]
        assert context.execute_stream == estimator_stream.cuda_stream
        assert estimator_pool.released == [(context, estimator_stream)]

    @pytest.mark.core_model
    @pytest.mark.cpu
    def test_trt_context_pool_stores_cuda_stream(self, monkeypatch):
        from vllm_omni.model_executor.models.cosyvoice3.flow_estimator_trt import TrtContextWrapper

        execution_context = object()

        class FakeEngine:
            @staticmethod
            def create_execution_context():
                return execution_context

        cuda_stream = object()
        monkeypatch.setattr(torch.cuda, "Stream", lambda device: cuda_stream)

        def reject_stream_context(stream):
            pytest.fail("TrtContextWrapper must store the CUDA stream, not a StreamContext")

        monkeypatch.setattr(torch.cuda, "stream", reject_stream_context)

        engine = FakeEngine()
        wrapper = TrtContextWrapper(engine, device="cuda:0")
        [context, stream], acquired_engine = wrapper.acquire_estimator()

        assert context is execution_context
        assert stream is cuda_stream
        assert acquired_engine is engine

        wrapper.release_estimator(context, stream)
        [reused_context, reused_stream], _ = wrapper.acquire_estimator()
        assert reused_context is context
        assert reused_stream is stream

    @pytest.mark.core_model
    @pytest.mark.cpu
    def test_causal_conditional_cfm_batches_cfg_estimator(self):
        """Batched flow calls should invoke the CFG estimator with 2B rows."""
        from omegaconf import DictConfig

        from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import CausalConditionalCFM

        class RecordingEstimator(nn.Module):
            def __init__(self):
                super().__init__()
                self.shapes: list[tuple[int, ...]] = []

            def forward(self, x, mask, mu, t, spks=None, cond=None):
                self.shapes.append(tuple(x.shape))
                return torch.zeros_like(x)

        estimator = RecordingEstimator()
        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=estimator,
        )

        batch, mel_dim, seq_len = 3, 80, 16
        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
        assert estimator.shapes == [(2 * batch, mel_dim, seq_len)] * 2


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

        with _force_torch_sdpa():
            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


def test_code2wav_forward_finalizes_hift_tail():
    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav

    class DummyHiFT(nn.Module):
        def __init__(self):
            super().__init__()
            self.m_source = SimpleNamespace(l_linear=SimpleNamespace(weight=torch.ones(1, dtype=torch.float32)))
            self.finalize_calls: list[bool] = []

        def inference(self, speech_feat, finalize=True):
            self.finalize_calls.append(bool(finalize))
            return (
                torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype),
                None,
                None,
            )

    model = object.__new__(CosyVoice3Code2Wav)
    nn.Module.__init__(model)
    model.hift = DummyHiFT()
    forward_mel_calls = []

    def fake_forward_mel(**kwargs):
        forward_mel_calls.append(kwargs)
        return torch.ones((1, 80, 8), dtype=torch.float32)

    model._forward_mel = fake_forward_mel

    out = model.forward(
        token=torch.tensor([[1, 2, 3]], dtype=torch.int32),
        prompt_token=torch.tensor([[4, 5]], dtype=torch.int32),
        prompt_feat=torch.ones((1, 4, 80), dtype=torch.float32),
        embedding=torch.ones((1, 192), dtype=torch.float32),
    )

    assert out.shape == (1, 1, 8)
    assert model.hift.finalize_calls == [True]
    assert forward_mel_calls[0]["token_offset_tokens"] == 0


def test_code2wav_streaming_batch_pads_codec_tokens_and_preserves_lengths():
    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav

    model = object.__new__(CosyVoice3Code2Wav)
    model.flow_model = SimpleNamespace(pre_lookahead_len=1, token_mel_ratio=2)
    forward_mel_calls = []

    def fake_forward_mel(self, **kwargs):
        forward_mel_calls.append(kwargs)
        return torch.arange(2 * 80 * 8, dtype=torch.float32).reshape(2, 80, 8)

    def fake_stream_hift(self, feat, *, cache_state=None, finalize=False):
        return feat, None

    model._forward_mel = types.MethodType(fake_forward_mel, model)
    model._stream_hift_from_feat = types.MethodType(fake_stream_hift, model)
    items = [
        {
            "token": torch.ones(1, 3, dtype=torch.int32),
            "prompt_token": torch.ones(1, 4, dtype=torch.int32),
            "prompt_feat": torch.ones(1, 8, 80),
            "embedding": torch.ones(1, 192),
            "finalize": False,
        },
        {
            "token": torch.ones(1, 5, dtype=torch.int32),
            "prompt_token": torch.ones(1, 4, dtype=torch.int32),
            "prompt_feat": torch.ones(1, 8, 80),
            "embedding": torch.ones(1, 192),
            "finalize": False,
        },
    ]

    results = model.forward_streaming_batch(items)

    assert len(forward_mel_calls) == 1
    call = forward_mel_calls[0]
    assert call["token"].shape == (2, 5)
    assert torch.equal(call["token_lens"], torch.tensor([3, 5], dtype=torch.int32))
    assert results[0][0].shape[-1] == 4
    assert results[1][0].shape[-1] == 8


@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_code2wav_streaming_batch_matches_ragged_flow_numerics(monkeypatch):
    """A padded flow batch must match individual flow calls on valid mels."""
    from omegaconf import DictConfig

    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.layers import PreLookaheadLayer
    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav

    torch.manual_seed(0)
    estimator = DiT(
        dim=32,
        depth=1,
        heads=4,
        dim_head=8,
        dropout=0.0,
        ff_mult=2,
        mel_dim=80,
        mu_dim=80,
        spk_dim=80,
        out_channels=80,
    )
    decoder = CausalConditionalCFM(
        in_channels=80,
        cfm_params=DictConfig(
            {
                "sigma_min": 1e-6,
                "solver": "euler",
                "t_scheduler": "cosine",
                "training_cfg_rate": 0.2,
                "inference_cfg_rate": 0.7,
            }
        ),
        n_spks=1,
        spk_emb_dim=80,
        estimator=estimator,
    )
    flow_model = (
        CausalMaskedDiffWithDiT(
            input_size=80,
            output_size=80,
            spk_embed_dim=192,
            vocab_size=64,
            input_frame_rate=25,
            only_mask_loss=True,
            token_mel_ratio=2,
            pre_lookahead_len=1,
            pre_lookahead_layer=PreLookaheadLayer(in_channels=80, channels=80, pre_lookahead_len=1),
            decoder=decoder,
        )
        .to(device="cuda", dtype=torch.bfloat16)
        .eval()
    )

    model = object.__new__(CosyVoice3Code2Wav)
    nn.Module.__init__(model)
    model.flow_model = flow_model

    def return_mel(self, feat, *, cache_state=None, finalize=False):
        return feat, None

    model._stream_hift_from_feat = types.MethodType(return_mel, model)

    original_randn = torch.randn

    def length_consistent_randn(*size, **kwargs):
        shape = tuple(size[0]) if len(size) == 1 and isinstance(size[0], (tuple, list)) else tuple(size)
        if len(shape) == 3 and shape[1] == 80:
            device = kwargs.get("device")
            dtype = kwargs.get("dtype", torch.float32)
            channels = torch.arange(shape[1], device=device, dtype=torch.float32).view(1, -1, 1)
            positions = torch.arange(shape[2], device=device, dtype=torch.float32).view(1, 1, -1)
            noise = torch.sin(channels * 0.17 + positions * 0.31)
            return noise.expand(shape[0], -1, -1).to(dtype=dtype).clone()
        return original_randn(*size, **kwargs)

    monkeypatch.setattr(torch, "randn", length_consistent_randn)
    common = {
        "prompt_token": torch.tensor([[7, 8]], dtype=torch.int32),
        "prompt_feat": torch.linspace(-0.5, 0.5, 4 * 80).reshape(1, 4, 80),
        "embedding": torch.linspace(-1.0, 1.0, 192).reshape(1, 192),
        "finalize": False,
    }
    items = [
        {**common, "token": torch.tensor([[1, 2, 3]], dtype=torch.int32)},
        {**common, "token": torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.int32)},
    ]

    batched = model.forward_streaming_batch(items, n_timesteps=2)
    individual = [
        model.forward_streaming(
            token=item["token"],
            prompt_token=item["prompt_token"],
            prompt_feat=item["prompt_feat"],
            embedding=item["embedding"],
            n_timesteps=2,
        )
        for item in items
    ]

    for (batched_mel, _), (individual_mel, _) in zip(batched, individual):
        assert batched_mel.shape == individual_mel.shape
        rel_mean = (batched_mel - individual_mel).abs().mean() / individual_mel.abs().mean().clamp_min(1e-6)
        assert rel_mean.item() < 0.05