From 7fb86d51cd0ad6745734b367a28dc24370552f88 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Fri, 3 Apr 2026 14:56:33 +0800 Subject: [PATCH] [BugFix][FishSpeech] Fix structured voice clone prefill conditioning (#2446) --- .../models/test_fish_speech_regressions.py | 108 ++++++++++++++++++ .../models/fish_speech/dac_encoder.py | 90 ++++++++------- .../fish_speech/fish_speech_dac_decoder.py | 19 +-- .../models/fish_speech/fish_speech_slow_ar.py | 37 +++++- 4 files changed, 201 insertions(+), 53 deletions(-) create mode 100644 tests/model_executor/models/test_fish_speech_regressions.py diff --git a/tests/model_executor/models/test_fish_speech_regressions.py b/tests/model_executor/models/test_fish_speech_regressions.py new file mode 100644 index 000000000..1f8c3cf71 --- /dev/null +++ b/tests/model_executor/models/test_fish_speech_regressions.py @@ -0,0 +1,108 @@ +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) diff --git a/vllm_omni/model_executor/models/fish_speech/dac_encoder.py b/vllm_omni/model_executor/models/fish_speech/dac_encoder.py index e89815ab4..397530ca3 100644 --- a/vllm_omni/model_executor/models/fish_speech/dac_encoder.py +++ b/vllm_omni/model_executor/models/fish_speech/dac_encoder.py @@ -66,42 +66,22 @@ def _load_dac_codec( def _get_resample_kernel( source_sr: int, target_sr: int, - device_type: str, - device_index: int | None, - dtype_name: str, + device: torch.device, + dtype: torch.dtype, ): import torchaudio - device = torch.device(device_type, device_index) if device_index is not None else torch.device(device_type) - dtype = getattr(torch, dtype_name) + # lru_cache requires hashable key parts; torch.device and torch.dtype are. return torchaudio.transforms.Resample(source_sr, target_sr).to(device=device, dtype=dtype) -@torch.no_grad() -def encode_reference_audio( - model_path: str, +def _prepare_reference_audio_tensor( wav_samples: list[float] | np.ndarray | torch.Tensor, sample_rate: int, *, - device: torch.device | str | None = None, -) -> list[int]: - """Encode reference audio into semantic token IDs for prompt conditioning. - - Args: - model_path: HuggingFace model path (for locating codec.pth). - wav_samples: Audio waveform samples (mono, float). - sample_rate: Sample rate of the input audio. - - Returns: - List of semantic token IDs (151678 + code_value for each frame). - """ - if device is None: - device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") - else: - device = torch.device(device) - dtype = torch.float32 - codec = _load_dac_codec(model_path, device=device, dtype=dtype) - + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: if isinstance(wav_samples, torch.Tensor): wav_tensor = wav_samples.detach() else: @@ -124,28 +104,52 @@ def encode_reference_audio( resampler = _get_resample_kernel( int(sample_rate), DAC_SAMPLE_RATE, - device.type, - device.index, - "float32", + device, + dtype, ) wav_tensor = resampler(wav_tensor.unsqueeze(0)).squeeze(0) + return wav_tensor + + +@torch.no_grad() +def encode_reference_audio_codes( + model_path: str, + wav_samples: list[float] | np.ndarray | torch.Tensor, + sample_rate: int, + *, + device: torch.device | str | None = None, +) -> torch.Tensor: + """Encode reference audio into DAC codebook indices. + + Returns: + Tensor of shape [num_frames, num_codebooks] on the requested device + (dtype=torch.long). + """ + if device is None: + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + else: + device = torch.device(device) + dtype = torch.float32 + codec = _load_dac_codec(model_path, device=device, dtype=dtype) + wav_tensor = _prepare_reference_audio_tensor( + wav_samples, + sample_rate, + device=device, + dtype=dtype, + ) - # Encode: [1, 1, T] -> codes [1, num_codebooks, num_frames] wav_tensor = wav_tensor.unsqueeze(0).unsqueeze(0) feature_lengths = torch.tensor([wav_tensor.shape[-1]], device=device, dtype=torch.long) - codes, feature_lengths_out = codec.encode(wav_tensor, feature_lengths) - - # Extract semantic codebook (index 0) - shape [num_frames]. - semantic_codes = codes[0, 0, :].to(device="cpu", dtype=torch.long).tolist() - - # Convert to semantic token IDs: <|semantic:{i}|> = 151678 + i - SEMANTIC_TOKEN_OFFSET = 151678 - semantic_token_ids = [SEMANTIC_TOKEN_OFFSET + int(c) for c in semantic_codes] + codes, _ = codec.encode(wav_tensor, feature_lengths) + prepared_num_samples = int(wav_tensor.shape[-1]) + # [1, num_codebooks, num_frames] -> [num_frames, num_codebooks] + codes_fq = codes[0].transpose(0, 1).to(dtype=torch.long).contiguous() logger.info( - "Encoded reference audio: %d samples @ %dHz -> %d semantic tokens", - int(wav_tensor.shape[-1]), + "Encoded reference audio codes: %d samples @ %dHz -> frames=%d codebooks=%d", + prepared_num_samples, sample_rate, - len(semantic_token_ids), + int(codes_fq.shape[0]), + int(codes_fq.shape[1]), ) - return semantic_token_ids + return codes_fq diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py index 3a8042eb2..e121b0337 100644 --- a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py +++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py @@ -213,7 +213,9 @@ class FishSpeechDACDecoder(nn.Module): ids = input_ids.reshape(-1).to(dtype=torch.long) request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts")) - parsed_ctx_frames: list[int] = [] + num_req = len(request_ids_list) + parsed_ctx_frames = [0] * num_req + parsed_total_frames = [0] * num_req valid_codes_qf: list[torch.Tensor] = [] valid_indices: list[int] = [] left_context_size = [0] * len(request_ids_list) @@ -226,7 +228,6 @@ class FishSpeechDACDecoder(nn.Module): for i, req_ids in enumerate(request_ids_list): if req_ids.numel() < 1: - parsed_ctx_frames.append(0) continue ctx_frames = left_context_size[i] flat = req_ids @@ -238,15 +239,13 @@ class FishSpeechDACDecoder(nn.Module): n, q, ) - parsed_ctx_frames.append(0) continue frames = n // q codes_qf = flat.reshape(q, frames) - parsed_ctx_frames.append(ctx_frames) + parsed_ctx_frames[i] = ctx_frames + parsed_total_frames[i] = frames valid_codes_qf.append(codes_qf) valid_indices.append(i) - - num_req = len(request_ids_list) if not valid_codes_qf: return OmniOutput( text_hidden_states=None, @@ -297,11 +296,17 @@ class FishSpeechDACDecoder(nn.Module): for j, idx in enumerate(valid_indices): ctx_frames = parsed_ctx_frames[idx] + total_frames = parsed_total_frames[idx] audio_len = int(audio_lengths[j].item()) if audio_lengths.numel() > j else int(wav_batch.shape[-1]) wav = wav_batch[j, 0, :audio_len] # Trim context frames (left overlap for streaming). if ctx_frames > 0: - cut = ctx_frames * self._hop_length + # Decode length may deviate from (frames * hop_length) due to model + # internals (padding/rounding). Use proportional trimming to keep + # overlap removal aligned with the actual decoded length. + denom = max(int(total_frames), 1) + cut = int(ctx_frames / denom * wav.shape[0]) + cut = max(0, min(cut, int(wav.shape[0]))) if cut < wav.shape[0]: wav = wav[cut:] else: diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py index b2e8a9544..4ad2a1fa6 100644 --- a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py +++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py @@ -35,7 +35,7 @@ from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.models.output_templates import OmniOutput from .configuration_fish_speech import FishSpeechConfig, FishSpeechFastARConfig, FishSpeechSlowARConfig -from .dac_encoder import _load_dac_codec, encode_reference_audio +from .dac_encoder import _load_dac_codec, encode_reference_audio_codes from .fish_speech_fast_ar import FishSpeechFastAR from .prompt_utils import build_fish_voice_clone_prompt_ids @@ -530,12 +530,13 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module): ref_audio_wav = np.load(ref_audio_path) os.remove(ref_audio_path) - semantic_token_ids = encode_reference_audio( + ref_codes_fq = encode_reference_audio_codes( self.model_path, ref_audio_wav, ref_audio_sr, device=self.codebook_embeddings.weight.device, ) + semantic_token_ids = (ref_codes_fq[:, 0] + self._semantic_begin_id).tolist() prompt_ids, _, _ = build_fish_voice_clone_prompt_ids( tokenizer, text, @@ -547,7 +548,37 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module): dtype=torch.long, device=self.codebook_embeddings.weight.device, ) - return self.embed_input_ids(prompt_ids.unsqueeze(0)).squeeze(0).to(dtype=torch.bfloat16) + embeds = self.embed_input_ids(prompt_ids.unsqueeze(0)).squeeze(0).to(dtype=torch.bfloat16) + + audio_start_id = tokenizer.convert_tokens_to_ids("<|audio_start|>") + audio_end_id = tokenizer.convert_tokens_to_ids("<|audio_end|>") + start_pos = (prompt_ids == int(audio_start_id)).nonzero(as_tuple=False) + end_pos = (prompt_ids == int(audio_end_id)).nonzero(as_tuple=False) + if start_pos.numel() == 0 or end_pos.numel() == 0: + return embeds + s = int(start_pos[0].item()) + 1 + e = int(end_pos[0].item()) + if e <= s: + return embeds + + frames_in_prompt = e - s + if ref_codes_fq.device != embeds.device: + ref_codes_fq = ref_codes_fq.to(device=embeds.device, dtype=torch.long) + frames = min(int(ref_codes_fq.shape[0]), int(frames_in_prompt)) + if frames <= 0: + return embeds + + q = min(int(ref_codes_fq.shape[1]), self._num_codebooks) + offsets = (torch.arange(q, device=embeds.device, dtype=torch.long) * self._codebook_size).unsqueeze(0) + ref_codes_slice = ref_codes_fq[:frames, :q] + if bool((ref_codes_slice < 0).any().item()): + logger.warning("Fish Speech structured clone saw negative DAC codes; clamping them to zero") + code_with_offset = ref_codes_slice.clamp(min=0) + offsets + codebook_sum = self.codebook_embeddings(code_with_offset).sum(dim=1).to(dtype=embeds.dtype) + + result = embeds.clone() + result[s : s + frames] = (result[s : s + frames] + codebook_sum) / math.sqrt(self._num_codebooks + 1) + return result.to(dtype=torch.bfloat16) # -------------------- GPU-side MTP fast-path --------------------