From fe786f1a1946c8617d96093ea3a3c091dd088e84 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Tue, 17 Mar 2026 18:54:12 -0400 Subject: [PATCH] Fix Base voice clone streaming quality and stop-token crash (#1945) Signed-off-by: linyueqian --- .../e2e/online_serving/test_qwen3_tts_base.py | 56 +++++++++++++++++++ .../test_qwen3_tts_async_chunk.py | 36 ++++++++++-- .../models/qwen3_tts/qwen3_tts_code2wav.py | 12 ++-- .../stage_input_processors/qwen3_tts.py | 24 +++++--- 4 files changed, 111 insertions(+), 17 deletions(-) diff --git a/tests/e2e/online_serving/test_qwen3_tts_base.py b/tests/e2e/online_serving/test_qwen3_tts_base.py index 9dd662f08..c09ee592c 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_base.py +++ b/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." + ) diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index 8d6d21430..54faa3b3a 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/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]} diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index eee751e16..79a2810bf 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/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) diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index c01517066..42bdb5a2c 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/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()