diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py index d7559a72d..40e85da5b 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py @@ -2,18 +2,25 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from unittest.mock import patch import pytest import torch import torch.nn as nn -from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import Qwen3TTSCode2Wav +from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import ( + Qwen3TTSCode2Wav, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +_NUM_QUANTIZERS = 2 +_TOTAL_UPSAMPLE = 4 +_OUTPUT_SAMPLE_RATE = 24000 + class _FakeDecoder(nn.Module): - def __init__(self, total_upsample: int = 4): + def __init__(self, total_upsample: int = _TOTAL_UPSAMPLE): super().__init__() self.total_upsample = total_upsample @@ -24,18 +31,35 @@ class _FakeDecoder(nn.Module): return wav.view(1, 1, -1) -def _make_model() -> Qwen3TTSCode2Wav: - model = Qwen3TTSCode2Wav( - vllm_config=SimpleNamespace( - model_config=SimpleNamespace(model="unused"), - device_config=SimpleNamespace(device=torch.device("cpu")), - ) +def _fake_dec_config(): + return SimpleNamespace( + num_quantizers=_NUM_QUANTIZERS, + sliding_window=0, ) - model._decoder = _FakeDecoder() - model._num_quantizers = 2 - model._output_sample_rate = 24000 - model._total_upsample = 4 - model._ensure_speech_tokenizer_loaded = lambda: None + + +def _make_model() -> Qwen3TTSCode2Wav: + dec_config = _fake_dec_config() + tok_config = SimpleNamespace( + decoder_config=dec_config, + output_sample_rate=_OUTPUT_SAMPLE_RATE, + ) + with ( + patch( + "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav.Qwen3TTSTokenizerV2Config.from_pretrained", + return_value=tok_config, + ), + patch( + "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav.Qwen3TTSTokenizerV2Decoder._from_config", + return_value=_FakeDecoder(), + ), + ): + model = Qwen3TTSCode2Wav( + vllm_config=SimpleNamespace( + model_config=SimpleNamespace(model="unused"), + device_config=SimpleNamespace(device=torch.device("cpu")), + ) + ) return model 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 26adecce9..aec4f8eec 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 @@ -1,19 +1,24 @@ from __future__ import annotations -import os from collections.abc import Iterable from typing import Any import torch import torch.nn as nn -from transformers.utils.hub import cached_file from vllm.config import VllmConfig from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger +from vllm.model_executor.model_loader import DefaultModelLoader +from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.model_executor.models.output_templates import OmniOutput -from .qwen3_tts_tokenizer import Qwen3TTSTokenizer +from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2Config, +) +from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2Decoder, +) logger = init_logger(__name__) @@ -36,128 +41,24 @@ class Qwen3TTSCode2Wav(nn.Module): self.enable_update_additional_information = True self.requires_raw_input_tokens = True - self._speech_tokenizer: Qwen3TTSTokenizer | None = None - self._decoder: nn.Module | None = None - self._num_quantizers: int | None = None - self._output_sample_rate: int | None = None - self._total_upsample: int | None = None - self._decoder_sliding_window: int | None = None self._decode_chunk_frames = 300 self._decode_left_context_frames = 25 self._logged_codec_stats = False - @staticmethod - def _module_device(module: nn.Module) -> torch.device: - try: - return next(module.parameters()).device - except StopIteration: - for _, buf in module.named_buffers(recurse=True): - return buf.device - return torch.device("cpu") - - def _ensure_speech_tokenizer_loaded(self) -> None: - if self._decoder is not None: - return - - cfg_path = cached_file(self.model_path, "speech_tokenizer/config.json") - if cfg_path is None: - raise ValueError(f"{self.model_path}/speech_tokenizer/config.json not found") - speech_tokenizer_dir = os.path.dirname(cfg_path) - - prep_cfg = cached_file(self.model_path, "speech_tokenizer/preprocessor_config.json") - if prep_cfg is None: - raise ValueError( - f"{self.model_path}/speech_tokenizer/preprocessor_config.json not found. " - "Please make sure the checkpoint contains the required HF preprocessing files." - ) - - tok = Qwen3TTSTokenizer.from_pretrained( - speech_tokenizer_dir, - torch_dtype=torch.float32, - load_feature_extractor=False, + # Construct decoder from config so it is visible to vLLM's + # memory profiler at startup. Weights are loaded later in + # load_weights(). + tok_config = Qwen3TTSTokenizerV2Config.from_pretrained( + self.model_path, + subfolder="speech_tokenizer", ) - - if tok.model is not None: - tok.model.to(device=self.vllm_config.device_config.device) - tok.device = self._module_device(tok.model) - - dec_cfg = getattr(tok.model.config, "decoder_config", None) - num_q = getattr(dec_cfg, "num_quantizers", None) if dec_cfg is not None else None - if num_q is None: - raise ValueError("speech_tokenizer decoder_config.num_quantizers not found") - num_q = int(num_q) - if num_q <= 0: - raise ValueError(f"Invalid speech_tokenizer num_quantizers={num_q}") - - try: - upsample = int(tok.get_decode_upsample_rate()) - except Exception as e: - raise ValueError(f"Failed to get decode upsample rate: {e}") from e - if upsample <= 0: - raise ValueError(f"Invalid decode upsample rate: {upsample}") - - try: - out_sr = int(tok.get_output_sample_rate()) - except Exception as e: - raise ValueError(f"Failed to get output sample rate: {e}") from e - - decoder = tok.model.decoder - decoder.eval() - - self._speech_tokenizer = tok - self._decoder = decoder - self._num_quantizers = num_q - self._output_sample_rate = out_sr - self._total_upsample = int(decoder.total_upsample) - self._decoder_sliding_window = int(getattr(dec_cfg, "sliding_window", 0) or 0) - - # Precompute SnakeBeta exp caches (benefits both Triton and eager paths) - if hasattr(decoder, "precompute_snake_caches"): - decoder.precompute_snake_caches() - - chunk_frames = 0 - left_frames = 0 - model_cfg = getattr(self.vllm_config, "model_config", None) - connector_cfg = getattr(model_cfg, "stage_connector_config", None) - extra_cfg = ( - connector_cfg.get("extra", connector_cfg) - if isinstance(connector_cfg, dict) - else getattr(connector_cfg, "extra", None) - ) - if isinstance(extra_cfg, dict): - chunk_frames = int(extra_cfg.get("codec_chunk_frames") or 0) - left_frames = int(extra_cfg.get("codec_left_context_frames") or 0) - if getattr(model_cfg, "async_chunk", False) and chunk_frames > 0: - self._decode_chunk_frames = chunk_frames - self._decode_left_context_frames = left_frames if left_frames > 0 else 0 - - if hasattr(decoder, "enable_cudagraph"): - device = self._module_device(decoder) - if device.type == "cuda": - try: - if ( - chunk_frames > 0 - and left_frames > 0 - and self._decoder_sliding_window - and left_frames < self._decoder_sliding_window - ): - logger.warning( - "Qwen3-TTS streaming codec_left_context_frames=%d is smaller than " - "decoder sliding_window=%d; chunk-boundary distortion may occur. " - "Increase codec_left_context_frames to at least %d for streaming.", - left_frames, - self._decoder_sliding_window, - self._decoder_sliding_window, - ) - - decoder.enable_cudagraph( - device=device, - codec_chunk_frames=chunk_frames, - codec_left_context_frames=left_frames, - ) - logger.info("Code2Wav decoder CUDA Graph enabled") - except Exception: - logger.warning("Failed to enable CUDA Graph for Code2Wav decoder", exc_info=True) + dec_config = tok_config.decoder_config + self.decoder = Qwen3TTSTokenizerV2Decoder._from_config(dec_config) + self.decoder.eval() + self._num_quantizers = int(dec_config.num_quantizers) + self._output_sample_rate = int(tok_config.output_sample_rate) + self._total_upsample = int(self.decoder.total_upsample) + self._decoder_sliding_window = int(getattr(dec_config, "sliding_window", 0) or 0) def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: # This stage ignores token embeddings. Keep a stable dummy embedding for vLLM runner. @@ -210,12 +111,7 @@ class Qwen3TTSCode2Wav(nn.Module): Length management is done here instead of relying on HF's padding=-1 sentinel logic. """ - self._ensure_speech_tokenizer_loaded() - assert self._decoder is not None - assert self._num_quantizers is not None - assert self._total_upsample is not None - - decoder = self._decoder + decoder = self.decoder q = int(self._num_quantizers) upsample = int(self._total_upsample) sr_val = int(self._output_sample_rate) @@ -362,6 +258,80 @@ class Qwen3TTSCode2Wav(nn.Module): ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # SpeechTokenizer weights live under `speech_tokenizer/` and are loaded - # lazily from that directory. Ignore main checkpoint weights. - return set() + # The primary weights iterator contains no Code2Wav parameters. + # Drain it so callers don't hang on an unconsumed generator. + for _ in weights: + pass + + # Load decoder weights from the speech_tokenizer/ subfolder + # via vLLM's weight loader (handles sharded safetensors, index + # files, and all load formats). AutoWeightsLoader matches + # "decoder.*" weights to self.decoder and skips encoder weights. + model_loader = DefaultModelLoader(self.vllm_config.load_config) + source = DefaultModelLoader.Source( + model_or_path=self.model_path, + revision=self.vllm_config.model_config.revision, + subfolder="speech_tokenizer", + ) + subfolder_weights = model_loader._get_weights_iterator(source) + loaded = AutoWeightsLoader( + self, + skip_prefixes=["encoder."], + ).load_weights(subfolder_weights) + + device = self.vllm_config.device_config.device + self.decoder.to(device=device, dtype=torch.float32) + + # Precompute SnakeBeta exp caches (benefits both Triton and eager paths) + if hasattr(self.decoder, "precompute_snake_caches"): + self.decoder.precompute_snake_caches() + + # Read chunk config from stage connector and update decode params + chunk_frames = 0 + left_frames = 0 + model_cfg = getattr(self.vllm_config, "model_config", None) + connector_cfg = getattr(model_cfg, "stage_connector_config", None) + extra_cfg = ( + connector_cfg.get("extra", connector_cfg) + if isinstance(connector_cfg, dict) + else getattr(connector_cfg, "extra", None) + ) + if isinstance(extra_cfg, dict): + chunk_frames = int(extra_cfg.get("codec_chunk_frames") or 0) + left_frames = int(extra_cfg.get("codec_left_context_frames") or 0) + if getattr(model_cfg, "async_chunk", False) and chunk_frames > 0: + self._decode_chunk_frames = chunk_frames + self._decode_left_context_frames = left_frames if left_frames > 0 else 0 + + if hasattr(self.decoder, "enable_cudagraph") and device.type == "cuda": + try: + if ( + chunk_frames > 0 + and left_frames > 0 + and self._decoder_sliding_window + and left_frames < self._decoder_sliding_window + ): + logger.warning( + "Qwen3-TTS streaming codec_left_context_frames=%d " + "is smaller than decoder sliding_window=%d; " + "chunk-boundary distortion may occur. " + "Increase codec_left_context_frames to at least " + "%d for streaming.", + left_frames, + self._decoder_sliding_window, + self._decoder_sliding_window, + ) + + self.decoder.enable_cudagraph( + device=device, + codec_chunk_frames=chunk_frames, + codec_left_context_frames=left_frames, + ) + logger.info("Code2Wav decoder CUDA Graph enabled") + except Exception: + logger.warning( + "Failed to enable CUDA Graph for Code2Wav decoder", + exc_info=True, + ) + + return loaded