[Cleanup] Move cosyvoice3 tests to model subdirectory (#1666)
Signed-off-by: linyueqian <linyueqian@outlook.com>
This commit is contained in:
Yueqian Lin authored and GitHub committed
1 parent
c0b714a8ac
commit
8536dce759
2 files changed
No files matched your search
@@ -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
|
||||
@@ -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]
|
||||
Reference in new issue
Block a user