[Fix] Qwen3 TTS audio handling for long ref_audio (#2104)

Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: Sy0307 <1370724210@qq.com>

vllm_omni/entrypoints/openai/serving_speech.py

@@ -10,6 +10,7 @@ from pathlib import Path
 from typing import Any
 
 import numpy as np
+import torch
 from fastapi import Request, UploadFile
 from fastapi.responses import Response, StreamingResponse
 from transformers.utils.hub import cached_file
@@ -1052,9 +1053,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
 
         if isinstance(audio_tensor, list):
-            import torch
-
-            audio_tensor = torch.cat(audio_tensor, dim=-1)
+            async_chunk = bool(getattr(self.engine_client.model_config, "async_chunk", False))
+            if async_chunk:
+                non_empty_chunks = [candidate for candidate in audio_tensor if candidate.numel() > 0]
+                audio_tensor = (
+                    torch.cat(non_empty_chunks, dim=-1) if non_empty_chunks else np.zeros((0,), dtype=np.float32)
+                )
+            else:
+                audio_history = audio_tensor
+                audio_tensor = np.zeros((0,), dtype=np.float32)
+                # Non-async Qwen3-TTS returns cumulative history snapshots, so keep the latest non-empty tensor.
+                for candidate in reversed(audio_history):
+                    if candidate.numel() > 0:
+                        audio_tensor = candidate
+                        break
         if hasattr(audio_tensor, "float"):
             audio_tensor = audio_tensor.float().detach().cpu().numpy()
 

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py

@@ -219,7 +219,13 @@ class Qwen3TTSCode2Wav(nn.Module):
                 if i >= len(left_context_size):
                     break
                 if "left_context_size" in info:
-                    left_context_size[i] = info["left_context_size"]
+                    # left_context_size may come through serialization as an int, [int], or tensor([int]).
+                    value = info["left_context_size"]
+                    if isinstance(value, list):
+                        value = value[0] if value else 0
+                    if isinstance(value, torch.Tensor):
+                        value = value.reshape(-1)[0].item() if value.numel() > 0 else 0
+                    left_context_size[i] = int(value)
         for i, req_ids in enumerate(request_ids_list):
             if req_ids.numel() < 1:
                 parsed.append((0, 0))
@@ -284,24 +290,21 @@ class Qwen3TTSCode2Wav(nn.Module):
         for j, idx in enumerate(valid_indices):
             ctx_frames, actual_frames = parsed[idx]
             wav = wav_tensors[j]
-            if ctx_frames > 0:
-                # 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.
+            # Drop the ref_code prefix from the decoded waveform, keeping only newly generated audio.
+            if ctx_frames <= 0:
+                expected_len = actual_frames * upsample
+                if wav.shape[0] > expected_len:
+                    wav = wav[:expected_len]
+            else:
                 cut = int(ctx_frames / max(actual_frames, 1) * wav.shape[0])
-                if cut < wav.shape[0]:
-                    wav = wav[cut:]
-                else:
+                if cut >= wav.shape[0]:
                     logger.warning(
                         "Context trim %d >= decoded length %d; returning empty audio.",
                         cut,
                         wav.shape[0],
                     )
                     continue
-            else:
-                expected_len = actual_frames * upsample
-                if wav.shape[0] > expected_len:
-                    wav = wav[:expected_len]
+                wav = wav[cut:]
             if wav.shape[0] > 0:
                 audios[idx] = wav.to(dtype=torch.float32).reshape(-1)
 

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -32,21 +32,63 @@ def talker2code2wav(
     talker_outputs = _validate_stage_inputs(stage_list, engine_input_source)
     code2wav_inputs: list[OmniTokensPrompt] = []
     for i, talker_output in enumerate(talker_outputs):
+        if not talker_output.finished:
+            # Non-async decode should only run once, after talker has
+            # accumulated the final code sequence.
+            continue
         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)
+        token_ids = output.token_ids
+        # token_ids provides an upper bound on the newly generated codec span.
+        # audio_codes may still contain zero-padded / invalid rows, so trim only
+        # after filtering valid frames instead of trying to align EOS indices.
+        seq_len = max(len(token_ids) - 1, 0)
         # 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]
+        if seq_len > 0 and audio_codes.ndim == 2 and int(audio_codes.shape[0]) > seq_len:
+            audio_codes = audio_codes[-seq_len:]
         ref_code = output.multimodal_output.get("ref_code")
+        ref_code_len = output.multimodal_output.get("ref_code_len")
+        if isinstance(ref_code_len, torch.Tensor):
+            ref_code_len = int(ref_code_len.reshape(-1)[-1].item()) if ref_code_len.numel() > 0 else 0
+        elif ref_code_len is None:
+            ref_code_len = 0
+        else:
+            ref_code_len = int(ref_code_len)
         if isinstance(ref_code, list):
             ref_code = ref_code[0] if ref_code else None
         if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
             ref_code = ref_code.to(torch.long).cpu().contiguous()
-            ref_code_len = int(ref_code.shape[0])
-            audio_codes = torch.cat([ref_code.to(audio_codes.device), audio_codes], dim=0)
+            if ref_code.ndim == 1:
+                num_quantizers = int(audio_codes.shape[1]) if audio_codes.ndim == 2 and audio_codes.shape[1] > 0 else 16
+                if ref_code.numel() % num_quantizers != 0:
+                    logger.warning(
+                        "Ignoring malformed ref_code with %d elements not divisible by num_quantizers=%d",
+                        ref_code.numel(),
+                        num_quantizers,
+                    )
+                    ref_code = None
+                else:
+                    ref_code = ref_code.reshape(-1, num_quantizers)
+            elif ref_code.ndim != 2:
+                logger.warning("Ignoring malformed ref_code shape %s", tuple(ref_code.shape))
+                ref_code = None
+            if isinstance(ref_code, torch.Tensor) and ref_code_len > 0 and int(ref_code.shape[0]) > ref_code_len:
+                logger.warning(
+                    "Trimming ref_code from %d frames to ref_code_len=%d before Code2Wav.",
+                    int(ref_code.shape[0]),
+                    ref_code_len,
+                )
+                ref_code = ref_code[:ref_code_len]
+            if not isinstance(ref_code, torch.Tensor):
+                ref_code_len = 0
+            else:
+                ref_code_len = int(ref_code.shape[0])
+                audio_codes = torch.cat([ref_code.to(audio_codes.device), audio_codes], dim=0)
         else:
             ref_code_len = 0
         # Code2Wav expects codebook-major flat: [Q*num_frames]