# Copyright 2026 The Alibaba Qwen team. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import base64 import io import urllib.request from urllib.parse import urlparse import numpy as np import soundfile as sf import torch from torch.nn.utils.rnn import pad_sequence from transformers import AutoConfig, AutoFeatureExtractor, AutoModel from vllm.multimodal.audio import AudioResampler from vllm.multimodal.media.audio import load_audio as _load_audio_file from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( Qwen3TTSTokenizerV2EncoderOutput, Qwen3TTSTokenizerV2Model, ) from .tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1 import Qwen3TTSTokenizerV1Config from .tokenizer_25hz.modeling_qwen3_tts_tokenizer_v1 import ( Qwen3TTSTokenizerV1EncoderOutput, Qwen3TTSTokenizerV1Model, ) AudioInput = ( str # wav path, or base64 string | np.ndarray # 1-D float array | list[str] | list[np.ndarray] ) class Qwen3TTSTokenizer: """ A wrapper for Qwen3 TTS Tokenizer 25Hz/12Hz with HuggingFace-style loading. - from_pretrained(): loads speech tokenizer model via AutoModel and feature_extractor via AutoFeatureExtractor. - encode(): supports wav path(s), base64 audio string(s), numpy array(s). - decode(): accepts either the raw model encode output, or a minimal dict/list-of-dicts. Notes: - For numpy array input, you must pass `sr` so the audio can be resampled to model sample rate. - Returned audio is float32 numpy arrays and the output sample rate. """ def __init__(self): self.model = None self.feature_extractor = None self.config = None self.device = None @classmethod def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs) -> "Qwen3TTSTokenizer": """ Initialize tokenizer with HuggingFace `from_pretrained` style. Args: pretrained_model_name_or_path (str): HuggingFace repo id or local directory. **kwargs (Any): Forwarded to `AutoModel.from_pretrained(...)` directly. Typical examples: device_map="cuda:0", dtype=torch.bfloat16, attn_implementation="eager". Returns: Qwen3TTSTokenizer: Initialized instance with `model`, `feature_extractor`, `config`. """ inst = cls() load_feature_extractor = bool(kwargs.pop("load_feature_extractor", True)) AutoConfig.register("qwen3_tts_tokenizer_12hz", Qwen3TTSTokenizerV2Config) AutoModel.register(Qwen3TTSTokenizerV2Config, Qwen3TTSTokenizerV2Model) AutoConfig.register("qwen3_tts_tokenizer_25hz", Qwen3TTSTokenizerV1Config) AutoModel.register(Qwen3TTSTokenizerV1Config, Qwen3TTSTokenizerV1Model) inst.model = AutoModel.from_pretrained(pretrained_model_name_or_path, **kwargs) inst.config = inst.model.config inst.feature_extractor = ( AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path) if load_feature_extractor else None ) inst.device = getattr(inst.model, "device", None) if inst.device is None: # fallback: infer from first parameter device try: inst.device = next(inst.model.parameters()).device except StopIteration: inst.device = torch.device("cpu") return inst def _is_probably_base64(self, s: str) -> bool: if s.startswith("data:audio"): return True # Heuristic: no filesystem path separators and long enough. if ("/" not in s and "\\" not in s) and len(s) > 256: return True return False def _is_url(self, s: str) -> bool: try: u = urlparse(s) return u.scheme in ("http", "https") and bool(u.netloc) except Exception: return False def _decode_base64_to_wav_bytes(self, b64: str) -> bytes: # Accept both "data:audio/wav;base64,...." and raw base64 if "," in b64 and b64.strip().startswith("data:"): b64 = b64.split(",", 1)[1] return base64.b64decode(b64) def load_audio( self, x: str, target_sr: int, ) -> np.ndarray: """ Load audio from wav path or base64 string, then resample to target_sr. Args: x (str): A wav file path, or a base64 audio string (raw or data URL). target_sr (int): Target sampling rate. Returns: np.ndarray: 1-D float32 waveform at target_sr. """ if self._is_url(x): with urllib.request.urlopen(x) as resp: audio_bytes = resp.read() with io.BytesIO(audio_bytes) as f: audio, sr = sf.read(f, dtype="float32", always_2d=False) elif self._is_probably_base64(x): wav_bytes = self._decode_base64_to_wav_bytes(x) with io.BytesIO(wav_bytes) as f: audio, sr = sf.read(f, dtype="float32", always_2d=False) else: audio, sr = _load_audio_file(x, sr=None, mono=True) if audio.ndim > 1: audio = np.mean(audio, axis=-1) if sr != target_sr: resampler = AudioResampler(target_sr=target_sr) audio = resampler.resample(audio, orig_sr=sr) return audio.astype(np.float32) def _normalize_audio_inputs( self, audios: AudioInput, sr: int | None, ) -> list[np.ndarray]: """ Normalize all supported input types into a list of 1-D numpy float32 waveforms at `self.feature_extractor.sampling_rate`. Args: audios (AudioInput): - str: wav path OR base64 audio string - np.ndarray: raw waveform (sr must be provided) - list[str] / list[np.ndarray] sr (Optional[int]): Sampling rate for raw numpy input. Required if input is np.ndarray or list[np.ndarray]. Returns: List[np.ndarray]: List of float32 waveforms resampled to model input SR. """ target_sr = int(self.feature_extractor.sampling_rate) if isinstance(audios, (str, np.ndarray)): audios = [audios] if len(audios) == 0: return [] if isinstance(audios[0], str): # wav path list or base64 list return [self.load_audio(x, target_sr=target_sr) for x in audios] # type: ignore[arg-type] # numpy list if sr is None: raise ValueError("For numpy waveform input, you must provide `sr` (original sampling rate).") out: list[np.ndarray] = [] for a in audios: # type: ignore[assignment] if not isinstance(a, np.ndarray): raise TypeError("Mixed input types are not supported. Use all paths/base64 or all numpy arrays.") if a.ndim > 1: a = np.mean(a, axis=-1) if int(sr) != target_sr: resampler = AudioResampler(target_sr=target_sr) a = resampler.resample(a.astype(np.float32), orig_sr=int(sr)) out.append(a.astype(np.float32)) return out def encode( self, audios: AudioInput, sr: int | None = None, return_dict: bool = True, ) -> Qwen3TTSTokenizerV1EncoderOutput | Qwen3TTSTokenizerV2EncoderOutput | tuple: """ Batch-encode audio into discrete codes (and optional conditioning, depending on 25Hz/12Hz). Args: audios (AudioInput): Supported forms: - np.ndarray: waveform (requires sr) - list[np.ndarray]: waveforms (requires sr) - str: wav path OR base64 audio string - list[str]: wav paths and/or base64 strings sr (Optional[int], default=None): Original sampling rate for numpy waveform input. return_dict (bool, default=True): Forwarded to model.encode(...). If True, returns ModelOutput. Returns: Qwen3TTSTokenizerV1EncoderOutput | Qwen3TTSTokenizerV2EncoderOutput | tuple: Encoder output or tuple returned by model.encode. If return_dict=True, returns a model-specific encoder output. For 25Hz models, this includes audio_codes/xvectors/ref_mels; for 12Hz models, this includes audio_codes. If return_dict=False, returns the raw tuple from model.encode. """ wavs = self._normalize_audio_inputs(audios, sr=sr) inputs = self.feature_extractor( raw_audio=wavs, sampling_rate=int(self.feature_extractor.sampling_rate), return_tensors="pt", ) inputs = inputs.to(self.device).to(self.model.dtype) with torch.inference_mode(): # model.encode expects (B, T) and (B, T) enc = self.model.encode( inputs["input_values"].squeeze(1), inputs["padding_mask"].squeeze(1), return_dict=return_dict, ) return enc def decode( self, encoded, ) -> tuple[list[np.ndarray], int]: """ Decode back to waveform. Usage: 1) Pass the raw output of `encode(...)` directly (recommended). - 25Hz: expects fields audio_codes, xvectors, ref_mels - 12Hz: expects field audio_codes 2) Pass a dict or list[dict] (minimal form) for custom pipelines: - 25Hz dict keys: {"audio_codes", "xvectors", "ref_mels"} - 12Hz dict keys: {"audio_codes"} Values can be torch tensors or numpy arrays. Args: encoded (Any): - ModelOutput returned by `encode()`, OR - dict, OR - list[dict] Returns: Tuple[List[np.ndarray], int]: - wavs: list of 1-D float32 numpy arrays - sample_rate: int, model output sampling rate """ model_type = self.model.get_model_type() def _to_tensor(x, dtype=None): if isinstance(x, torch.Tensor): return x x = np.asarray(x) t = torch.from_numpy(x) if dtype is not None: t = t.to(dtype) return t # Normalize `encoded` into the same shapes as the official demo uses. if hasattr(encoded, "audio_codes"): # ModelOutput from encode() audio_codes_list = encoded.audio_codes xvectors_list = getattr(encoded, "xvectors", None) ref_mels_list = getattr(encoded, "ref_mels", None) elif isinstance(encoded, dict): audio_codes_list = encoded["audio_codes"] xvectors_list = encoded.get("xvectors", None) ref_mels_list = encoded.get("ref_mels", None) elif isinstance(encoded, list): # list of dicts audio_codes_list = [e["audio_codes"] for e in encoded] xvectors_list = [e["xvectors"] for e in encoded] if ("xvectors" in encoded[0]) else None ref_mels_list = [e["ref_mels"] for e in encoded] if ("ref_mels" in encoded[0]) else None else: raise TypeError("`encoded` must be an encode output, a dict, or a list of dicts.") # Ensure list form for per-sample tensors if isinstance(audio_codes_list, torch.Tensor): # Could be a single sample tensor or an already padded batch tensor. t = audio_codes_list if t.dim() == 1: # 25Hz single sample: (C,) -> (1, C) t = t.unsqueeze(0) elif t.dim() == 2: # 12Hz single sample: (C, Q) -> (1, C, Q) t = t.unsqueeze(0) audio_codes_padded = t.to(self.device) else: # List[Tensor/np] audio_codes_list = [_to_tensor(c, dtype=torch.long) for c in audio_codes_list] audio_codes_padded = pad_sequence(audio_codes_list, batch_first=True, padding_value=-1).to(self.device) with torch.inference_mode(): if model_type == "qwen3_tts_tokenizer_25hz": if xvectors_list is None or ref_mels_list is None: raise ValueError("25Hz decode requires `xvectors` and `ref_mels`.") if isinstance(xvectors_list, torch.Tensor): xvectors_batch = xvectors_list if xvectors_batch.dim() == 1: # (D,) -> (1, D) xvectors_batch = xvectors_batch.unsqueeze(0) xvectors_batch = xvectors_batch.to(self.device).to(self.model.dtype) else: xvectors_list = [_to_tensor(x, dtype=torch.float32) for x in xvectors_list] xvectors_batch = torch.stack(xvectors_list, dim=0).to(self.device).to(self.model.dtype) if isinstance(ref_mels_list, torch.Tensor): ref_mels_padded = ref_mels_list if ref_mels_padded.dim() == 2: # (T, M) -> (1, T, M) ref_mels_padded = ref_mels_padded.unsqueeze(0) ref_mels_padded = ref_mels_padded.to(self.device).to(self.model.dtype) else: ref_mels_list = [_to_tensor(m, dtype=torch.float32) for m in ref_mels_list] ref_mels_padded = ( pad_sequence(ref_mels_list, batch_first=True, padding_value=0) .to(self.device) .to(self.model.dtype) ) dec = self.model.decode(audio_codes_padded, xvectors_batch, ref_mels_padded, return_dict=True) wav_tensors = dec.audio_values elif model_type == "qwen3_tts_tokenizer_12hz": dec = self.model.decode(audio_codes_padded, return_dict=True) wav_tensors = dec.audio_values else: raise ValueError(f"Unknown model type: {model_type}") wavs = [w.to(torch.float32).detach().cpu().numpy() for w in wav_tensors] return wavs, int(self.model.get_output_sample_rate()) def get_model_type(self) -> str: """ Get the underlying tokenizer model type. Returns: str: Model type string from `self.model.config.model_type` (e.g. "qwen3_tts_tokenizer_25hz" / "qwen3_tts_tokenizer_12hz"). """ return self.model.get_model_type() def get_input_sample_rate(self) -> int: """ Get the expected input sample rate for encoding. Returns: int: Input sample rate (Hz). """ return int(self.model.get_input_sample_rate()) def get_output_sample_rate(self) -> int: """ Get the output sample rate for decoded waveforms. Returns: int: Output sample rate (Hz). """ return int(self.model.get_output_sample_rate()) def get_encode_downsample_rate(self) -> int: """ Get the encoder downsample rate (waveform samples per code step). Returns: int: Encode downsample rate. """ return int(self.model.get_encode_downsample_rate()) def get_decode_upsample_rate(self) -> int: """ Get the decoder upsample rate (waveform samples per code step). Returns: int: Decode upsample rate. """ return int(self.model.get_decode_upsample_rate())