[BugFix][FishSpeech] Fix structured voice clone prefill conditioning (#2446)

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)

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

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:

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 --------------------