117 lines
3.6 KiB
Python
117 lines
3.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.model_executor.models.voxcpm2.runtime_config import _VoxCPM2RuntimeConfig
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@dataclass
|
|
class _ModelConfig:
|
|
voxcpm2_runtime_config: dict[str, object]
|
|
|
|
|
|
@dataclass
|
|
class _VllmConfig:
|
|
model_config: _ModelConfig
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("key", "value", "expected"),
|
|
[
|
|
("inference_timesteps", 4, 4),
|
|
("inference_timesteps", "6", 6),
|
|
("inference_timesteps", -3, 2),
|
|
("inference_timesteps", 0, 2),
|
|
("inference_timesteps", 1, 2),
|
|
("inference_timesteps", 2, 2),
|
|
("cfg_value", "1.5", 1.5),
|
|
("cfg_cutoff_ratio", "0.5", 0.5),
|
|
("cfg_cutoff_ratio", -1.0, 0.0),
|
|
("cfg_cutoff_ratio", 2.0, 1.0),
|
|
],
|
|
)
|
|
def test_generation_config(key: str, value: object, expected: int | float) -> None:
|
|
vllm_config = _VllmConfig(_ModelConfig({key: value}))
|
|
|
|
config = _VoxCPM2RuntimeConfig.from_vllm_config(vllm_config)
|
|
|
|
assert getattr(config, key) == expected
|
|
|
|
|
|
def test_generation_config_defaults() -> None:
|
|
config = _VoxCPM2RuntimeConfig()
|
|
|
|
assert config.inference_timesteps == 10
|
|
assert config.cfg_value == 2.0
|
|
assert config.cfg_cutoff_ratio == 1.0
|
|
|
|
|
|
class _CountingEstimator(torch.nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.calls = 0
|
|
|
|
def forward(self, x, mu, t, cond, dt):
|
|
self.calls += 1
|
|
return mu[:, :1, None].expand_as(x)
|
|
|
|
|
|
@pytest.mark.parametrize("value", ["nan", "inf", "-inf"])
|
|
def test_cfg_value_rejects_nonfinite(value: str) -> None:
|
|
with pytest.raises(ValueError, match="cfg_value must be finite"):
|
|
_VoxCPM2RuntimeConfig.from_vllm_config(_VllmConfig(_ModelConfig({"cfg_value": value})))
|
|
|
|
|
|
@pytest.mark.parametrize("timesteps", [-3, 0, 1, 2, 4, 10])
|
|
@pytest.mark.parametrize("batch_size", [1, 3])
|
|
@pytest.mark.parametrize("cfg_value", [1.5, 2.0])
|
|
@pytest.mark.parametrize("cutoff", [0.0, 1.0])
|
|
def test_generation_config_runs_estimator(timesteps: int, batch_size: int, cfg_value: float, cutoff: float) -> None:
|
|
from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
|
|
_CFMBufferManager,
|
|
_optimized_solve_euler_with_noise,
|
|
)
|
|
|
|
config = _VoxCPM2RuntimeConfig.from_vllm_config(
|
|
_VllmConfig(
|
|
_ModelConfig({"inference_timesteps": timesteps, "cfg_value": cfg_value, "cfg_cutoff_ratio": cutoff})
|
|
)
|
|
)
|
|
cfm = torch.nn.Module()
|
|
cfm.estimator = _CountingEstimator()
|
|
buffers = _CFMBufferManager(
|
|
device=torch.device("cpu"),
|
|
dtype=torch.float32,
|
|
feat_dim=2,
|
|
patch_size=2,
|
|
dit_hidden_size=2,
|
|
max_batch_size=batch_size,
|
|
)
|
|
noise = torch.ones(batch_size, 2, 2)
|
|
output = _optimized_solve_euler_with_noise(
|
|
cfm,
|
|
mu=torch.ones(batch_size, 2),
|
|
patch_size=2,
|
|
cond=torch.zeros_like(noise),
|
|
noise=noise,
|
|
n_timesteps=config.inference_timesteps,
|
|
cfg_value=config.cfg_value,
|
|
buffers=buffers,
|
|
cfg_cutoff_ratio=config.cfg_cutoff_ratio,
|
|
)
|
|
|
|
assert cfm.estimator.calls == max(2, timesteps) - 1
|
|
assert torch.isfinite(output).all()
|
|
assert not torch.equal(output, noise)
|
|
torch.testing.assert_close(noise, torch.ones_like(noise))
|
|
if cutoff == 1.0:
|
|
t_span = buffers.get_t_span(config.inference_timesteps)
|
|
expected = noise - cfg_value * (t_span[1] - t_span[-1])
|
|
torch.testing.assert_close(output, expected)
|