diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index d23b80b31..9a012abaa 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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() 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 6be039df1..c75f6e52e 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 @@ -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) 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 934dfb20e..21815b09b 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/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]