Sy03andYueqian Lin
58007c0084
Signed-off-by: Sysy <1370724210@qq.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
85 lines
2.5 KiB
Python
85 lines
2.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import torch
|
|
|
|
from vllm_omni.model_executor.models.indextts2 import codec as codec_module
|
|
from vllm_omni.model_executor.models.indextts2 import preprocess_utils
|
|
from vllm_omni.model_executor.models.indextts2.codec import EnhancedCodec
|
|
|
|
|
|
def test_enhanced_codec_decode_restores_hidden_size_and_doubles_time():
|
|
codec = EnhancedCodec(
|
|
codebook_size=8,
|
|
hidden_size=4,
|
|
codebook_dim=2,
|
|
vocos_dim=4,
|
|
vocos_intermediate_dim=8,
|
|
vocos_num_layers=1,
|
|
downsample_scale=2,
|
|
).eval()
|
|
|
|
decoded = codec.decode(torch.tensor([[0, 1, 2]]))
|
|
|
|
assert decoded.shape == (1, 6, 4)
|
|
|
|
|
|
def test_enhanced_codec_inference_module_does_not_allocate_unused_encoder():
|
|
codec = EnhancedCodec(
|
|
codebook_size=8,
|
|
hidden_size=4,
|
|
codebook_dim=2,
|
|
vocos_dim=4,
|
|
vocos_intermediate_dim=8,
|
|
vocos_num_layers=1,
|
|
)
|
|
|
|
parameter_names = set(dict(codec.named_parameters()))
|
|
assert not any(name.startswith("encoder.") for name in parameter_names)
|
|
assert not any(name.startswith("down.") for name in parameter_names)
|
|
assert any(name.startswith("decoder.") for name in parameter_names)
|
|
assert any(name.startswith("quantizer.") for name in parameter_names)
|
|
|
|
|
|
def test_semantic_codec_loader_selects_enhanced_checkpoint(tmp_path, monkeypatch):
|
|
checkpoint = tmp_path / "codec.pth"
|
|
checkpoint.touch()
|
|
seen = {}
|
|
|
|
class FakeCodec:
|
|
def load_checkpoint(self, path):
|
|
seen["path"] = path
|
|
|
|
def to(self, **kwargs):
|
|
seen["to"] = kwargs
|
|
return self
|
|
|
|
def eval(self):
|
|
return self
|
|
|
|
def parameters(self):
|
|
return ()
|
|
|
|
monkeypatch.setattr(codec_module, "EnhancedCodec", lambda **kwargs: FakeCodec())
|
|
preprocess_utils._semantic_codec_cache.clear()
|
|
|
|
codec = preprocess_utils.load_semantic_codec(
|
|
str(tmp_path),
|
|
{"codebook_size": 8192},
|
|
torch.device("cpu"),
|
|
codec_type="enhanced",
|
|
checkpoint_name="codec.pth",
|
|
)
|
|
|
|
assert isinstance(codec, FakeCodec)
|
|
assert seen["path"] == str(checkpoint)
|
|
assert seen["to"] == {"device": torch.device("cpu"), "dtype": torch.float32}
|
|
|
|
|
|
def test_model_asset_resolver_supports_native_checkpoints_layout(tmp_path):
|
|
checkpoint = tmp_path / "checkpoints" / "codec.pth"
|
|
checkpoint.parent.mkdir()
|
|
checkpoint.touch()
|
|
|
|
assert preprocess_utils.resolve_model_file(str(tmp_path), "codec.pth") == str(checkpoint)
|