Fix Base voice clone streaming quality and stop-token crash (#1945)

Signed-off-by: linyueqian <linyueqian@outlook.com>

tests/e2e/online_serving/test_qwen3_tts_base.py

@@ -17,6 +17,7 @@ import tempfile
 from pathlib import Path
 
 import httpx
+import numpy as np
 import pytest
 
 from tests.conftest import (
@@ -99,6 +100,34 @@ def assert_not_silence(pcm_bytes: bytes):
 
 
 MIN_AUDIO_BYTES = 10000
+MIN_HNR_DB = 1.2  # Clean voice clone > 1.2 dB; distorted < 1.0 dB
+
+
+def compute_hnr_db(pcm_samples: np.ndarray, sr: int = 24000) -> float:
+    """Compute mean Harmonic-to-Noise Ratio (dB) for speech quality.
+
+    Clean cloned speech has HNR > 1.2 dB; distorted speech (e.g. lost
+    ref_code decoder context) drops below 1.0 dB.
+    """
+    frame_len = int(0.03 * sr)  # 30ms frames
+    hop = frame_len // 2
+    hnr_values: list[float] = []
+
+    for start in range(0, len(pcm_samples) - frame_len, hop):
+        frame = pcm_samples[start : start + frame_len]
+        if np.max(np.abs(frame)) < 0.01:
+            continue
+        ac = np.correlate(frame, frame, mode="full")[len(frame) - 1 :]
+        ac = ac / (ac[0] + 1e-10)
+        min_lag = int(sr / 400)
+        max_lag = min(int(sr / 80), len(ac))
+        if min_lag >= max_lag:
+            continue
+        peak = float(np.max(ac[min_lag:max_lag]))
+        if 0 < peak < 1:
+            hnr_values.append(10 * np.log10(peak / (1 - peak + 1e-10)))
+
+    return float(np.mean(hnr_values)) if hnr_values else 0.0
 
 
 class TestQwen3TTSBaseVoiceClone:
@@ -165,3 +194,30 @@ class TestQwen3TTSBaseVoiceClone:
         assert response.headers.get("content-type") == "audio/wav"
         assert verify_wav_audio(response.content), "Response is not valid WAV"
         assert len(response.content) > MIN_AUDIO_BYTES
+
+    @pytest.mark.core_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "L4"}, num_cards=4)
+    def test_base_voice_clone_no_distortion(self, omni_server) -> None:
+        """Audio must not be distorted (regression for ref_code context loss).
+
+        When the decoder loses ref_code speaker context on later streaming
+        chunks, HNR drops below 1.0 dB.  Clean voice clone should be > 1.2 dB.
+        See https://github.com/vllm-project/vllm-omni/issues/1944
+        """
+        response = make_base_speech_request(
+            host=omni_server.host,
+            port=omni_server.port,
+            response_format="pcm",
+        )
+
+        assert response.status_code == 200, f"Request failed: {response.text}"
+        assert len(response.content) > MIN_AUDIO_BYTES
+
+        pcm_samples = np.frombuffer(response.content, dtype=np.int16).astype(np.float32) / 32768.0
+        hnr = compute_hnr_db(pcm_samples)
+        print(f"Voice clone HNR: {hnr:.2f} dB (threshold: {MIN_HNR_DB} dB)")
+        assert hnr >= MIN_HNR_DB, (
+            f"Audio distortion detected: HNR={hnr:.2f} dB < {MIN_HNR_DB} dB. "
+            "Voice clone decoder may be losing ref_code speaker context on later chunks."
+        )

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -227,12 +227,14 @@ def test_first_streaming_chunk_prepends_ref_code_context():
     assert len(payload["code_predictor_codes"]) == _Q * 12
 
 
-def test_ref_code_context_only_applies_to_first_streaming_chunk():
+def test_ref_code_context_applies_to_all_streaming_chunks():
+    """ref_code is prepended as decoder context on every chunk, not just the first."""
     tm = _tm()
     rid = "r-ref2"
     tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(20)]
     tm.put_req_chunk[rid] = 1
     ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+    tm.request_payload[rid] = ref_code
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
@@ -242,8 +244,9 @@ def test_ref_code_context_only_applies_to_first_streaming_chunk():
     )
 
     assert payload is not None
-    assert payload["left_context_size"] == 10
-    assert len(payload["code_predictor_codes"]) == _Q * 20
+    # ref_code (2 frames) prepended as left context on second chunk too
+    assert payload["left_context_size"] == 10 + 2
+    assert len(payload["code_predictor_codes"]) == _Q * (20 + 2)
 
 
 def test_ref_code_context_can_be_buffered_before_first_emit():
@@ -276,9 +279,10 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
     )
 
     assert payload is not None
+    # ref_code (2 frames) is kept (not popped) for subsequent chunks
     assert payload["left_context_size"] == 2
     assert len(payload["code_predictor_codes"]) == _Q * 12
-    assert rid not in tm.request_payload
+    assert rid in tm.request_payload
 
 
 def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
@@ -317,3 +321,27 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
         4,
         8,
     ]
+
+
+def test_non_async_processor_filters_out_of_range_codec_values():
+    """Frames with values >= codebook_size (e.g. stop_token_id=2150) are filtered."""
+    ref_code = torch.tensor([[9, 9, 9, 9]], dtype=torch.long)
+    audio_codes = torch.tensor(
+        [
+            [0, 0, 0, 0],  # zero-padded (filtered)
+            [1, 2, 3, 4],  # valid
+            [2150, 0, 0, 0],  # stop token (filtered)
+            [5, 6, 7, 8],  # valid
+        ],
+        dtype=torch.long,
+    )
+    output = SimpleNamespace(multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code})
+    stage = SimpleNamespace(engine_outputs=[SimpleNamespace(outputs=[output])])
+
+    prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0])
+
+    assert len(prompts) == 1
+    prompt = prompts[0]
+    # Only ref_code (1 frame) + 2 valid frames = 3 frames * 4 quantizers = 12 codes
+    assert len(prompt["prompt_token_ids"]) == 4 * 3
+    assert prompt["additional_information"] == {"left_context_size": [1]}

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py

@@ -284,11 +284,11 @@ class Qwen3TTSCode2Wav(nn.Module):
         for j, idx in enumerate(valid_indices):
             ctx_frames, actual_frames = parsed[idx]
             wav = wav_tensors[j]
-            expected_len = actual_frames * upsample
-            if wav.shape[0] > expected_len:
-                wav = wav[:expected_len]
             if ctx_frames > 0:
-                cut = ctx_frames * upsample
+                # Proportional trim matching the official HF implementation:
+                # the decoder output length may not be exactly frames * upsample
+                # so compute cut as a proportion of total decoded length.
+                cut = int(ctx_frames / max(actual_frames, 1) * wav.shape[0])
                 if cut < wav.shape[0]:
                     wav = wav[cut:]
                 else:
@@ -298,6 +298,10 @@ class Qwen3TTSCode2Wav(nn.Module):
                         wav.shape[0],
                     )
                     continue
+            else:
+                expected_len = actual_frames * upsample
+                if wav.shape[0] > expected_len:
+                    wav = wav[:expected_len]
             if wav.shape[0] > 0:
                 audios[idx] = wav.to(dtype=torch.float32).reshape(-1)
 

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -29,8 +29,10 @@ def talker2code2wav(
         output = talker_output.outputs[0]
         # audio_codes shape: [num_frames, Q] where Q=num_quantizers (16)
         audio_codes = output.multimodal_output["audio_codes"].to(torch.long)
-        # Filter zero-padded frames (EOS/invalid steps), matching _extract_last_frame behavior
-        valid_mask = audio_codes.any(dim=1)
+        # Filter invalid frames: zero-padded (EOS) and frames containing
+        # out-of-range values (e.g. stop_token_id=2150 exceeds codebook_size=2048).
+        _CODEBOOK_SIZE = 2048
+        valid_mask = audio_codes.any(dim=1) & (audio_codes.max(dim=1).values < _CODEBOOK_SIZE)
         audio_codes = audio_codes[valid_mask]
         ref_code = output.multimodal_output.get("ref_code")
         if isinstance(ref_code, list):
@@ -177,13 +179,17 @@ def talker2code2wav_async_chunk(
     left_context_size = max(0, end_index - context_length)
     window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
 
-    # Prepend ref_code to the first decoder window only.
-    if transfer_manager.put_req_chunk[request_id] == 0:
-        ref_code = request_payload.pop(request_id, None)
-        if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
-            ref_frames = ref_code.tolist()
-            window_frames = ref_frames + window_frames
-            left_context_size += len(ref_frames)
+    # Prepend ref_code as decoder context for every chunk so the vocoder
+    # maintains voice-clone speaker identity throughout the stream.  The HF
+    # reference decodes ref_code + all_codes in one pass; without ref_code
+    # context on later chunks the decoder loses speaker identity and produces
+    # distorted audio.  Use `.get()` (not `.pop()`) to keep ref_code for
+    # subsequent chunks.
+    ref_code = request_payload.get(request_id)
+    if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
+        ref_frames = ref_code.tolist()
+        window_frames = ref_frames + window_frames
+        left_context_size += len(ref_frames)
 
     code_predictor_codes = torch.tensor(window_frames).transpose(0, 1).reshape(-1).tolist()