109 lines
4.2 KiB
Python
109 lines
4.2 KiB
Python
import math
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.model_executor.models.fish_speech import fish_speech_slow_ar as slow_ar_module
|
|
from vllm_omni.model_executor.models.fish_speech.fish_speech_dac_decoder import FishSpeechDACDecoder
|
|
from vllm_omni.model_executor.models.fish_speech.fish_speech_slow_ar import (
|
|
FishSpeechSlowARForConditionalGeneration,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class _FakeCodec:
|
|
def decode(self, codes_bqf: torch.Tensor, feature_lengths: torch.Tensor):
|
|
del codes_bqf, feature_lengths
|
|
wav = torch.arange(100, dtype=torch.float32).view(1, 1, 100)
|
|
audio_lengths = torch.tensor([100], dtype=torch.long)
|
|
return wav, audio_lengths
|
|
|
|
|
|
class _FakeTokenizer:
|
|
def __init__(self, mapping, unk_token_id=-1):
|
|
self._mapping = mapping
|
|
self.unk_token_id = unk_token_id
|
|
|
|
def convert_tokens_to_ids(self, token: str) -> int:
|
|
return self._mapping.get(token, self.unk_token_id)
|
|
|
|
|
|
def test_dac_decoder_mixed_batch_empty_request_does_not_misalign_indices():
|
|
decoder = object.__new__(FishSpeechDACDecoder)
|
|
torch.nn.Module.__init__(decoder)
|
|
decoder._codec = _FakeCodec()
|
|
decoder._num_codebooks = 10
|
|
decoder._output_sample_rate = 44100
|
|
decoder._hop_length = 512
|
|
decoder._logged_codec_stats = False
|
|
decoder._ensure_codec_loaded = lambda: None
|
|
decoder._split_request_ids = lambda ids, seq_token_counts=None: [
|
|
torch.empty((0,), dtype=torch.long),
|
|
torch.arange(20, dtype=torch.long),
|
|
]
|
|
|
|
out = decoder.forward(
|
|
input_ids=torch.arange(20, dtype=torch.long),
|
|
runtime_additional_information=[{}, {"left_context_size": 1}],
|
|
)
|
|
|
|
audios = out.multimodal_outputs["model_outputs"]
|
|
assert len(audios) == 2
|
|
assert audios[0].numel() == 0
|
|
# 2 total frames with 1 frame of left context => proportional trim removes half the samples.
|
|
assert audios[1].shape[0] == 50
|
|
|
|
|
|
def test_structured_voice_clone_prefill_adds_full_codebooks_with_decode_scale(monkeypatch):
|
|
model = object.__new__(FishSpeechSlowARForConditionalGeneration)
|
|
torch.nn.Module.__init__(model)
|
|
model._num_codebooks = 2
|
|
model._codebook_size = 8
|
|
model._semantic_begin_id = 100
|
|
model.model_path = "unused"
|
|
|
|
hidden_size = 3
|
|
text_embed = torch.nn.Embedding(256, hidden_size)
|
|
codebook_embed = torch.nn.Embedding(model._num_codebooks * model._codebook_size, hidden_size)
|
|
with torch.no_grad():
|
|
text_embed.weight.zero_()
|
|
text_embed.weight[20] = torch.tensor([1.0, 2.0, 3.0])
|
|
text_embed.weight[21] = torch.tensor([4.0, 5.0, 6.0])
|
|
codebook_embed.weight.zero_()
|
|
codebook_embed.weight[1] = torch.tensor([10.0, 0.0, 0.0])
|
|
codebook_embed.weight[10] = torch.tensor([0.0, 20.0, 0.0])
|
|
codebook_embed.weight[3] = torch.tensor([30.0, 0.0, 0.0])
|
|
codebook_embed.weight[12] = torch.tensor([0.0, 40.0, 0.0])
|
|
|
|
model.embed_input_ids = lambda ids: text_embed(ids)
|
|
model.codebook_embeddings = codebook_embed
|
|
model._get_tokenizer = lambda: _FakeTokenizer({"<|audio_start|>": 10, "<|audio_end|>": 11})
|
|
|
|
monkeypatch.setattr(slow_ar_module.np, "load", lambda path: [0.0])
|
|
monkeypatch.setattr(slow_ar_module.os, "remove", lambda path: None)
|
|
monkeypatch.setattr(
|
|
slow_ar_module,
|
|
"encode_reference_audio_codes",
|
|
lambda *args, **kwargs: torch.tensor([[1, 2], [3, 4]], dtype=torch.long),
|
|
)
|
|
monkeypatch.setattr(
|
|
slow_ar_module,
|
|
"build_fish_voice_clone_prompt_ids",
|
|
lambda tokenizer, text, ref_text, semantic_token_ids: ([1, 10, 20, 21, 11, 2], None, None),
|
|
)
|
|
|
|
prefill = model._build_structured_voice_clone_prefill_embeds(
|
|
{
|
|
"ref_text": "ref",
|
|
"text": "target",
|
|
"ref_audio_path": "unused.npy",
|
|
"ref_audio_sr": 16000,
|
|
}
|
|
)
|
|
|
|
expected_0 = (torch.tensor([1.0, 2.0, 3.0]) + torch.tensor([10.0, 20.0, 0.0])) / math.sqrt(3.0)
|
|
expected_1 = (torch.tensor([4.0, 5.0, 6.0]) + torch.tensor([30.0, 40.0, 0.0])) / math.sqrt(3.0)
|
|
assert torch.allclose(prefill[2].to(dtype=torch.float32), expected_0, atol=2e-2, rtol=0)
|
|
assert torch.allclose(prefill[3].to(dtype=torch.float32), expected_1, atol=2e-2, rtol=0)
|