masterr-main

Dockerfile

@@ -1,38 +0,0 @@
-# Use Python 3.10 as the base image
-FROM python:3.10-slim
-
-# Set up a new user named "user" with user ID 1000
-RUN useradd -m -u 1000 user
-
-# Switch to the "user" user
-USER user
-
-# Set home and path environment variables
-ENV HOME=/home/user \
-    PATH=/home/user/.local/bin:$PATH
-
-# Set the working directory
-WORKDIR $HOME/app
-
-# Copy the requirements file first to leverage Docker cache
-COPY --chown=user requirements.txt $HOME/app/requirements.txt
-
-# Install dependencies
-RUN pip install --no-cache-dir --upgrade pip && \
-    pip install --no-cache-dir -r requirements.txt
-
-# Copy the rest of the application code
-# Copy the rest of the application code
-#COPY --chown=user download_models.py $HOME/app/download_models.py
-
-# Download models during build
-#RUN python download_models.py
-
-COPY --chown=user . $HOME/app
-
-# Expose the default port for Hugging Face Spaces
-EXPOSE 7860
-
-# Run the application
-# We use python app.py since it calls graph.launch()
-CMD ["python", "app.py"]

README.md

@@ -1,16 +1,12 @@
 ---
 title: Ohm Audio Studio
-emoji: 🔥
-colorFrom: red
-colorTo: yellow
+emoji: 🏆
+colorFrom: yellow
+colorTo: purple
 sdk: docker
-app_port: 7860
-pinned: true
+pinned: false
 license: apache-2.0
-short_description: Professional Voice Design, Cloning & ASR Studio by Ohm.
-tags:
-- daggr
-- tts
-- asr
-- qwen
----
+short_description: Daggr+Qwen3 ASR + Qwen3 TTS
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,471 +0,0 @@
-"""
-Ohm Audio Studio
-======================
-
-A professional interface for Qwen2-Audio ASR and TTS models.
-This application uses Daggr and Gradio to provide a seamless user experience
-for Voice Design, Voice Cloning, Custom Voice Synthesis, and Automatic Speech Recognition.
-
-Author: Ohm
-Date: 2026
-"""
-
-import os
-import gc
-import base64
-import io
-import logging
-import numpy as np
-import torch
-import torchaudio
-import soundfile as sf
-import gradio as gr
-from typing import Any, Dict, List, Optional, Tuple, Union
-from dataclasses import dataclass
-
-from huggingface_hub import snapshot_download, login
-from daggr import FnNode, Graph
-
-# Configure Logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-# --- Configuration ---
-
-@dataclass
-class AppConfig:
-    HF_TOKEN: Optional[str] = os.environ.get('HF_TOKEN')
-    OMP_NUM_THREADS: str = "1"
-    
-    MODEL_SIZES = ["0.6B", "1.7B"]
-    SPEAKERS = [
-        "Aiden", "Dylan", "Eric", "Ono_anna", "Ryan", "Serena", "Sohee", "Uncle_fu", "Vivian"
-    ]
-    
-    TTS_LANGUAGES = [
-        "Auto", "English", "Japanese", "Korean", "French", "German", 
-        "Spanish", "Portuguese", "Russian"
-    ]
-    
-    ASR_SUPPORTED_LANGUAGES = [
-        "English", "Arabic", "German", "French", "Spanish", "Portuguese", 
-        "Indonesian", "Italian", "Korean", "Russian", "Thai", "Vietnamese", 
-        "Japanese", "Turkish", "Hindi", "Malay", "Dutch", "Swedish", "Danish", 
-        "Finnish", "Polish", "Czech", "Filipino", "Persian", "Greek", 
-        "Romanian", "Hungarian", "Macedonian"
-    ]
-
-# Apply Environment Variables
-os.environ["OMP_NUM_THREADS"] = AppConfig.OMP_NUM_THREADS
-if AppConfig.HF_TOKEN:
-    login(token=AppConfig.HF_TOKEN)
-
-# --- Utilities ---
-
-class AudioUtils:
-    """Utilities for audio processing and normalization."""
-
-    @staticmethod
-    def title_case_display(s: str) -> str:
-        s = (s or "").strip()
-        s = s.replace("_", " ")
-        return " ".join([w[:1].upper() + w[1:] if w else "" for w in s.split()])
-
-    @staticmethod
-    def build_choices_and_map(items: Optional[List[str]]) -> Tuple[List[str], Dict[str, str]]:
-        if not items:
-            return [], {}
-        display = [AudioUtils.title_case_display(x) for x in items]
-        mapping = {d: r for d, r in zip(display, items)}
-        return display, mapping
-
-    @staticmethod
-    def normalize_audio(wav: np.ndarray, eps: float = 1e-12, clip: bool = True) -> np.float32:
-        """Normalize audio to float32 in [-1, 1] range."""
-        x = np.asarray(wav)
-        
-        if np.issubdtype(x.dtype, np.integer):
-            info = np.iinfo(x.dtype)
-            if info.min < 0:
-                y = x.astype(np.float32) / max(abs(info.min), info.max)
-            else:
-                mid = (info.max + 1) / 2.0
-                y = (x.astype(np.float32) - mid) / mid
-        elif np.issubdtype(x.dtype, np.floating):
-            y = x.astype(np.float32)
-            m = np.max(np.abs(y)) if y.size else 0.0
-            if m > 1.0 + 1e-6:
-                y = y / (m + eps)
-        else:
-            y = x.astype(np.float32)
-
-        if clip:
-            y = np.clip(y, -1.0, 1.0)
-            
-        if y.ndim > 1:
-            y = np.mean(y, axis=-1).astype(np.float32)
-            
-        return y
-
-    @staticmethod
-    def process_input(audio_input: Any) -> Optional[Tuple[np.float32, int]]:
-        """
-        Handles Filepaths, Data URIs (base64), and Numpy arrays.
-        Returns (numpy_float32, sample_rate_int)
-        """
-        if audio_input is None:
-            return None
-
-        try:
-            # Handle Path or Base64
-            if isinstance(audio_input, str):
-                if audio_input.startswith("data:"):
-                    try:
-                        header, encoded = audio_input.split(",", 1)
-                        data = base64.b64decode(encoded)
-                        wav, sr = sf.read(io.BytesIO(data))
-                        return AudioUtils.normalize_audio(wav), int(sr)
-                    except Exception as e:
-                        logger.error(f"Failed to decode base64 audio: {e}")
-                        return None
-                
-                if os.path.exists(audio_input):
-                    wav_tensor, sr = torchaudio.load(audio_input)
-                    wav = wav_tensor.mean(dim=0).numpy()
-                    return AudioUtils.normalize_audio(wav), int(sr)
-                else:
-                    logger.error(f"Input string is not a file or valid data URI: {audio_input[:50]}...")
-                    return None
-
-            # Handle Tuple (sample_rate, data) or (data, sample_rate)
-            if isinstance(audio_input, tuple) and len(audio_input) == 2:
-                a0, a1 = audio_input
-                if isinstance(a0, int):
-                    return AudioUtils.normalize_audio(a1), int(a0)
-                else:
-                    return AudioUtils.normalize_audio(a0), int(a1)
-
-            # Handle Dictionary
-            if isinstance(audio_input, dict):
-                if "name" in audio_input:
-                    return AudioUtils.process_input(audio_input["name"])
-                if "path" in audio_input:
-                    return AudioUtils.process_input(audio_input["path"])
-                if "sampling_rate" in audio_input and "data" in audio_input:
-                    return AudioUtils.normalize_audio(audio_input["data"]), int(audio_input["sampling_rate"])
-
-            return None
-
-        except Exception as e:
-            logger.error(f"Audio Processing Error: {e}")
-            return None
-
-# --- Model Management ---
-
-class ModelManager:
-    """Manages loading and unloading of AI models."""
-    
-    def __init__(self):
-        self._loaded_models = {}
-
-    def _get_model_path(self, model_type: str, model_size: str) -> str:
-        """Download/Get model path based on type and size."""
-        if model_type == "ASR":
-            return "Qwen/Qwen3-ASR-1.7B"
-        return snapshot_download(f"Qwen/Qwen3-TTS-12Hz-{model_size}-{model_type}")
-
-    def get_model(self, model_type: str, model_size: str):
-        """
-        Lazy load models. Unloads previous models if VRAM is tight.
-        """
-        key = (model_type, model_size)
-        
-        if key not in self._loaded_models:
-            logger.info(f"Clearing Cache before loading {model_type}...")
-            self._loaded_models.clear()
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-
-            logger.info(f"Loading Model: {model_type} {model_size}...")
-            
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-
-            if model_type == "ASR":
-                from qwen_asr import Qwen3ASRModel
-                self._loaded_models[key] = Qwen3ASRModel.from_pretrained(
-                    "Qwen/Qwen3-ASR-1.7B",
-                    dtype=dtype,
-                    device_map=device,
-                    forced_aligner="Qwen/Qwen3-ForcedAligner-0.6B",
-                    forced_aligner_kwargs=dict(dtype=dtype, device_map=device),
-                    max_inference_batch_size=4, 
-                    attn_implementation="sdpa",
-                )
-            else:
-                from qwen_tts import Qwen3TTSModel
-                model_path = self._get_model_path(model_type, model_size)
-                self._loaded_models[key] = Qwen3TTSModel.from_pretrained(
-                    model_path,
-                    device_map=device,
-                    dtype=dtype,
-                    token=AppConfig.HF_TOKEN,
-                )
-            
-        return self._loaded_models[key]
-
-# --- Core Service ---
-
-class QwenService:
-    """Core service logic connecting the ModelManager and AudioUtils."""
-    
-    def __init__(self):
-        self.models = ModelManager()
-
-    def _cleanup_resources(self):
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-    def voice_design(self, text, language, voice_description):
-        """Voice Design (Prompt-to-Speech)"""
-        self._cleanup_resources()
-
-        if not text: return None, "Text required"
-        if not voice_description: return None, "Description required"
-
-        try:
-            tts = self.models.get_model("VoiceDesign", "1.7B") 
-            
-            wavs, sr = tts.generate_voice_design(
-                text=text.strip(),
-                language=language,
-                instruct=voice_description.strip(),
-                non_streaming_mode=True,
-                max_new_tokens=2048,
-            )
-            return (sr, wavs[0]), "Success"
-        except Exception as e:
-            logger.exception("Voice Design Error")
-            return None, f"Error: {str(e)}"
-
-    def voice_clone(self, ref_audio, ref_text, target_text, language, use_xvector_only, model_size):
-        """Voice Cloning (Zero-Shot)"""
-        self._cleanup_resources()
-
-        if not target_text: return None, "Target text required"
-        
-        audio_tuple = AudioUtils.process_input(ref_audio)
-        if audio_tuple is None:
-            return None, "Error: Could not process reference audio. Please upload a valid WAV/MP3."
-
-        if not use_xvector_only and not ref_text:
-            return None, "Error: Reference text required (or check 'Use x-vector only')"
-
-        try:
-            tts = self.models.get_model("Base", model_size)
-            
-            wavs, sr = tts.generate_voice_clone(
-                text=target_text.strip(),
-                language=language,
-                ref_audio=audio_tuple,
-                ref_text=ref_text.strip() if ref_text else None,
-                x_vector_only_mode=use_xvector_only,
-                max_new_tokens=2048,
-            )
-            return (sr, wavs[0]), "Success"
-        except Exception as e:
-            logger.exception("Voice Clone Error")
-            return None, f"Error: {str(e)}"
-
-    def custom_voice(self, text, language, speaker, instruct, model_size):
-        """Standard TTS"""
-        self._cleanup_resources()
-
-        if not text: return None, "Text required"
-
-        try:
-            tts = self.models.get_model("CustomVoice", model_size)
-            
-            wavs, sr = tts.generate_custom_voice(
-                text=text.strip(),
-                language=language,
-                speaker=speaker.lower().replace(" ", "_"),
-                instruct=instruct.strip() if instruct else None,
-                non_streaming_mode=True,
-                max_new_tokens=2048,
-            )
-            return (sr, wavs[0]), "Success"
-        except Exception as e:
-            logger.exception("Custom Voice Error")
-            return None, f"Error: {str(e)}"
-
-    def asr(self, audio_upload, lang_disp):
-        """Automatic Speech Recognition"""
-        self._cleanup_resources()
-
-        if audio_upload is None:
-            return "", "", "No Audio"
-
-        processed_audio = AudioUtils.process_input(audio_upload)
-        if processed_audio is None:
-            return "", "", "Error processing audio"
-        
-        language = None
-        if lang_disp and lang_disp != "Auto":
-            # Assuming ASR_LANG_MAP is globally available or we rebuild it
-            # For efficiency let's reuse if possible, or rebuild locally
-            _, mapping = AudioUtils.build_choices_and_map(AppConfig.ASR_SUPPORTED_LANGUAGES)
-            language = mapping.get(lang_disp, lang_disp)
-
-        try:
-            asr_model = self.models.get_model("ASR", "1.7B")
-
-            results = asr_model.transcribe(
-                audio=processed_audio,
-                language=language,
-                return_time_stamps=False,
-            )
-
-            if not isinstance(results, list) or len(results) != 1:
-                return "", "", "Unexpected result format"
-
-            r = results[0]
-            detected_lang = getattr(r, "language", "") or ""
-            transcribed_text = getattr(r, "text", "") or ""
-
-            return detected_lang, transcribed_text, "Success"
-
-        except Exception as e:
-            logger.exception("ASR Error")
-            return "", "", f"Error: {str(e)}"
-
-
-# --- Graph Construction ---
-
-# Initialize Service
-service = QwenService()
-ASR_LANG_DISPLAY, _ = AudioUtils.build_choices_and_map(AppConfig.ASR_SUPPORTED_LANGUAGES)
-ASR_LANG_CHOICES = ["Auto"] + ASR_LANG_DISPLAY
-
-# Define Nodes
-voice_design_node = FnNode(
-    fn=service.voice_design,
-    inputs={
-        "text": gr.Textbox(
-            label="Text to Synthesize (Voice Design)", 
-            lines=4, 
-            value="Welcome to Ohm Audio Studio. Experience the future of voice design."
-        ),
-        "language": gr.Dropdown(
-            label="Language (Voice Design)", 
-            choices=AppConfig.TTS_LANGUAGES, 
-            value="Auto"
-        ),
-        "voice_description": gr.Textbox(
-            label="Voice Description (Voice Design)", 
-            lines=3, 
-            value="A professional, warm and inviting voice with a clear, confident tone."
-        ),
-    },
-    outputs={
-        "generated_audio": gr.Audio(label="Generated Audio", type="numpy"),
-        "status": gr.Textbox(label="Status", interactive=False),
-    },
-    name="Voice Design"
-)
-
-custom_voice_node = FnNode(
-    fn=service.custom_voice,
-    inputs={
-        "text": gr.Textbox(
-            label="Text to Synthesize (Custom Voice)", 
-            lines=4, 
-            value="Welcome to Ohm Audio Studio coverage of the latest in AI audio technology."
-        ),
-        "language": gr.Dropdown(
-            label="Language (Custom Voice)", 
-            choices=AppConfig.TTS_LANGUAGES, 
-            value="English"
-        ),
-        "speaker": gr.Dropdown(
-            label="Speaker (Custom Voice)", 
-            choices=AppConfig.SPEAKERS, 
-            value="Ryan"
-        ),
-        "instruct": gr.Textbox(
-            label="Style Instruction (Custom Voice)", 
-            lines=2, 
-            placeholder="e.g. Happy, Sad", 
-            value="Neutral"
-        ),
-        "model_size": gr.Dropdown(
-            label="Model Size (Custom Voice)", 
-            choices=AppConfig.MODEL_SIZES, 
-            value="1.7B"
-        ),
-    },
-    outputs={
-        "tts_audio": gr.Audio(label="Generated Audio", type="numpy"),
-        "status": gr.Textbox(label="Status", interactive=False),
-    },
-    name="Custom Voice"
-)
-
-voice_clone_node = FnNode(
-    fn=service.voice_clone,
-    inputs={
-        "ref_audio": gr.Audio(label="Reference Audio (Voice Clone)", type="filepath"),
-        "ref_text": gr.Textbox(label="Reference Transcript (Voice Clone)", lines=2),
-        "target_text": gr.Textbox(label="Target Text (Voice Clone)", lines=4),
-        "language": gr.Dropdown(
-            label="Language (Voice Clone)", 
-            choices=AppConfig.TTS_LANGUAGES, 
-            value="Auto"
-        ),
-        "use_xvector_only": gr.Checkbox(label="Use x-vector only (Voice Clone)", value=False),
-        "model_size": gr.Dropdown(
-            label="Model Size (Voice Clone)", 
-            choices=AppConfig.MODEL_SIZES, 
-            value="1.7B"
-        ),
-    },
-    outputs={
-        "cloned_audio": gr.Audio(label="Cloned Audio", type="numpy"),
-        "status": gr.Textbox(label="Status", interactive=False),
-    },
-    name="Voice Clone"
-)
-
-asr_node = FnNode(
-    fn=service.asr,
-    inputs={
-        "audio_upload": gr.Audio(
-            label="Upload Audio (Qwen3 ASR)", 
-            type="numpy", 
-            sources=["upload", "microphone"]
-        ),
-        "lang_disp": gr.Dropdown(
-            label="Language (Qwen3 ASR)", 
-            choices=ASR_LANG_CHOICES, 
-            value="Auto"
-        ),
-    },
-    outputs={
-        "detected_lang": gr.Textbox(label="Detected Language", interactive=False),
-        "transcription": gr.Textbox(label="Transcription Result", lines=6, interactive=True),
-        "status": gr.Textbox(label="Status", interactive=False),
-    },
-    name="Qwen3 ASR"
-)
-
-# Create and Launch Graph
-graph = Graph(
-    name="Ohm-Audio-Studio",
-    nodes=[voice_design_node, custom_voice_node, voice_clone_node, asr_node]
-)
-
-if __name__ == "__main__":
-    port = int(os.environ.get("PORT", 7860))
-    graph.launch(host="0.0.0.0", port=port)

download_models.py

@@ -1,37 +0,0 @@
-
-import os
-from huggingface_hub import snapshot_download
-
-# List of models used in the app
-models_to_download = [
-    # ASR Models
-    "Qwen/Qwen3-ASR-1.7B",
-    "Qwen/Qwen3-ForcedAligner-0.6B",
-    
-    # TTS Models (1.7B versions as they are the default/main ones used)
-    "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
-    "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
-    "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
-    
-    # Uncomment if you plan to use 0.6B TTS models
-    # "Qwen/Qwen3-TTS-12Hz-0.6B-VoiceDesign",
-    # "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
-    # "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice",
-]
-
-print(f"Starting download of {len(models_to_download)} models...")
-
-for model_id in models_to_download:
-    print(f"\nDownloading: {model_id}")
-    try:
-        # If HF_TOKEN is empty/None, treat as anonymous (public models only)
-        token = os.environ.get('HF_TOKEN')
-        if token and not token.strip():
-            token = None
-            
-        path = snapshot_download(repo_id=model_id, token=token)
-        print(f"Successfully downloaded to: {path}")
-    except Exception as e:
-        print(f"Failed to download {model_id}: {e}")
-
-print("\nAll downloads complete!")

pre-requirements.txt

@@ -1 +0,0 @@
-pip>=23.0.0

qwen_asr/__init__.py

@@ -1,25 +0,0 @@
-# coding=utf-8
-# 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.
-"""
-qwen_asr: Qwen3-ASR package.
-"""
-
-from .inference.qwen3_asr import Qwen3ASRModel
-from .inference.qwen3_forced_aligner import Qwen3ForcedAligner
-
-from .inference.utils import parse_asr_output
-
-__all__ = ["__version__"]

qwen_asr/__main__.py

@@ -1,25 +0,0 @@
-# coding=utf-8
-# 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.
-def main():
-    print(
-        "qwen_asr package.\n"
-        "Use CLI entrypoints:\n"
-        "  - qwen-asr-demo\n"
-        "  - qwen-asr-serve\n"
-    )
-
-if __name__ == "__main__":
-    main()

qwen_asr/cli/demo.py

@@ -1,523 +0,0 @@
-# coding=utf-8
-# Copyright 2026 The Alibaba Qwen team.
-# SPDX-License-Identifier: Apache-2.0
-"""
-A gradio demo for Qwen3 ASR models.
-"""
-
-import argparse
-import base64
-import io
-import json
-import os
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import gradio as gr
-import numpy as np
-import torch
-from qwen_asr import Qwen3ASRModel
-from scipy.io.wavfile import write as wav_write
-
-
-def _title_case_display(s: str) -> str:
-    s = (s or "").strip()
-    s = s.replace("_", " ")
-    return " ".join([w[:1].upper() + w[1:] if w else "" for w in s.split()])
-
-
-def _build_choices_and_map(items: Optional[List[str]]) -> Tuple[List[str], Dict[str, str]]:
-    if not items:
-        return [], {}
-    display = [_title_case_display(x) for x in items]
-    mapping = {d: r for d, r in zip(display, items)}
-    return display, mapping
-
-
-def _dtype_from_str(s: str) -> torch.dtype:
-    s = (s or "").strip().lower()
-    if s in ("bf16", "bfloat16"):
-        return torch.bfloat16
-    if s in ("fp16", "float16", "half"):
-        return torch.float16
-    if s in ("fp32", "float32"):
-        return torch.float32
-    raise ValueError(f"Unsupported torch dtype: {s}. Use bfloat16/float16/float32.")
-
-
-def _normalize_audio(wav, eps=1e-12, clip=True):
-    x = np.asarray(wav)
-
-    if np.issubdtype(x.dtype, np.integer):
-        info = np.iinfo(x.dtype)
-        if info.min < 0:
-            y = x.astype(np.float32) / max(abs(info.min), info.max)
-        else:
-            mid = (info.max + 1) / 2.0
-            y = (x.astype(np.float32) - mid) / mid
-    elif np.issubdtype(x.dtype, np.floating):
-        y = x.astype(np.float32)
-        m = np.max(np.abs(y)) if y.size else 0.0
-        if m > 1.0 + 1e-6:
-            y = y / (m + eps)
-    else:
-        raise TypeError(f"Unsupported dtype: {x.dtype}")
-
-    if clip:
-        y = np.clip(y, -1.0, 1.0)
-
-    if y.ndim > 1:
-        y = np.mean(y, axis=-1).astype(np.float32)
-
-    return y
-
-
-def _audio_to_tuple(audio: Any) -> Optional[Tuple[np.ndarray, int]]:
-    """
-    Accept gradio audio:
-      - {"sampling_rate": int, "data": np.ndarray}
-      - (sr, np.ndarray)  [some gradio versions]
-    Return: (wav_float32_mono, sr)
-    """
-    if audio is None:
-        return None
-
-    if isinstance(audio, dict) and "sampling_rate" in audio and "data" in audio:
-        sr = int(audio["sampling_rate"])
-        wav = _normalize_audio(audio["data"])
-        return wav, sr
-
-    if isinstance(audio, tuple) and len(audio) == 2:
-        a0, a1 = audio
-        if isinstance(a0, int):
-            sr = int(a0)
-            wav = _normalize_audio(a1)
-            return wav, sr
-        if isinstance(a1, int):
-            wav = _normalize_audio(a0)
-            sr = int(a1)
-            return wav, sr
-
-    return None
-
-
-def _parse_audio_any(audio: Any) -> Union[str, Tuple[np.ndarray, int]]:
-    if audio is None:
-        raise ValueError("Audio is required.")
-    at = _audio_to_tuple(audio)
-    if at is not None:
-        return at
-    raise ValueError("Unsupported audio input format.")
-
-
-def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(
-        prog="qwen-asr-demo",
-        description=(
-            "Launch a Gradio demo for Qwen3 ASR models (Transformers / vLLM).\n\n"
-            "Examples:\n"
-            "  qwen-asr-demo --asr-checkpoint Qwen/Qwen3-ASR-1.7B\n"
-            "  qwen-asr-demo --asr-checkpoint Qwen/Qwen3-ASR-1.7B --aligner-checkpoint Qwen/Qwen3-ForcedAligner-0.6B\n"
-            "  qwen-asr-demo --backend vllm --cuda-visible-devices 0\n"
-            "  qwen-asr-demo --backend transformers --backend-kwargs '{\"device_map\":\"cuda:0\",\"dtype\":\"bfloat16\",\"attn_implementation\":\"flash_attention_2\"}'\n"
-            "  qwen-asr-demo --backend vllm --backend-kwargs '{\"gpu_memory_utilization\":0.85}'\n"
-        ),
-        formatter_class=argparse.RawTextHelpFormatter,
-        add_help=True,
-    )
-
-    parser.add_argument("--asr-checkpoint", required=True, help="Qwen3-ASR model checkpoint path or HF repo id.")
-    parser.add_argument(
-        "--aligner-checkpoint",
-        default=None,
-        help="Qwen3-ForcedAligner checkpoint path or HF repo id (optional; enables timestamps when provided).",
-    )
-
-    parser.add_argument(
-        "--backend",
-        default="transformers",
-        choices=["transformers", "vllm"],
-        help="Backend for ASR model loading (default: transformers).",
-    )
-
-    parser.add_argument(
-        "--cuda-visible-devices",
-        default="0",
-        help=(
-            "Set CUDA_VISIBLE_DEVICES for the demo process (default: 0). "
-            "Use e.g. '0' or '1'"
-        ),
-    )
-
-    parser.add_argument(
-        "--backend-kwargs",
-        default=None,
-        help=(
-            "JSON dict for backend-specific kwargs excluding checkpoints.\n"
-            "Examples:\n"
-            "  transformers: '{\"device_map\":\"cuda:0\",\"dtype\":\"bfloat16\",\"attn_implementation\":\"flash_attention_2\",\"max_inference_batch_size\":32}'\n"
-            "  vllm        : '{\"gpu_memory_utilization\":0.8,\"max_inference_batch_size\":32}'\n"
-        ),
-    )
-    parser.add_argument(
-        "--aligner-kwargs",
-        default=None,
-        help=(
-            "JSON dict for forced aligner kwargs (only used when --aligner-checkpoint is set).\n"
-            "Example: '{\"dtype\":\"bfloat16\",\"device_map\":\"cuda:0\"}'\n"
-        ),
-    )
-
-    # Gradio server args
-    parser.add_argument("--ip", default="0.0.0.0", help="Server bind IP for Gradio (default: 0.0.0.0).")
-    parser.add_argument("--port", type=int, default=8000, help="Server port for Gradio (default: 8000).")
-    parser.add_argument(
-        "--share/--no-share",
-        dest="share",
-        default=False,
-        action=argparse.BooleanOptionalAction,
-        help="Whether to create a public Gradio link (default: disabled).",
-    )
-    parser.add_argument("--concurrency", type=int, default=16, help="Gradio queue concurrency (default: 16).")
-
-    # HTTPS args
-    parser.add_argument("--ssl-certfile", default=None, help="Path to SSL certificate file for HTTPS (optional).")
-    parser.add_argument("--ssl-keyfile", default=None, help="Path to SSL key file for HTTPS (optional).")
-    parser.add_argument(
-        "--ssl-verify/--no-ssl-verify",
-        dest="ssl_verify",
-        default=True,
-        action=argparse.BooleanOptionalAction,
-        help="Whether to verify SSL certificate (default: enabled).",
-    )
-
-    return parser
-
-
-def _parse_json_dict(s: Optional[str], *, name: str) -> Dict[str, Any]:
-    if s is None or not str(s).strip():
-        return {}
-    try:
-        obj = json.loads(s)
-    except Exception as e:
-        raise ValueError(f"Invalid JSON for {name}: {e}")
-    if not isinstance(obj, dict):
-        raise ValueError(f"{name} must be a JSON object (dict).")
-    return obj
-
-
-def _apply_cuda_visible_devices(cuda_visible_devices: str) -> None:
-    v = (cuda_visible_devices or "").strip()
-    if not v:
-        return
-    os.environ["CUDA_VISIBLE_DEVICES"] = v
-
-
-def _default_backend_kwargs(backend: str) -> Dict[str, Any]:
-    if backend == "transformers":
-        return dict(
-            dtype=torch.bfloat16,
-            device_map="cuda:0",
-            attn_implementation="flash_attention_2",
-            max_inference_batch_size=32,
-        )
-    else:
-        return dict(
-            gpu_memory_utilization=0.8,
-            max_inference_batch_size=32,
-        )
-
-
-def _default_aligner_kwargs() -> Dict[str, Any]:
-    return dict(
-        dtype=torch.bfloat16,
-        device_map="cuda:0",
-    )
-
-
-def _merge_dicts(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
-    out = dict(base)
-    out.update(override)
-    return out
-
-
-def _coerce_special_types(d: Dict[str, Any]) -> Dict[str, Any]:
-    out: Dict[str, Any] = {}
-    for k, v in d.items():
-        if k == "dtype" and isinstance(v, str):
-            out[k] = _dtype_from_str(v)
-        else:
-            out[k] = v
-    return out
-
-
-def _make_timestamp_html(audio_upload: Any, timestamps: Any) -> str:
-    """
-    Build HTML with per-token audio slices, using base64 data URLs (no filesystem caching).
-    Expect timestamps as list[dict] with keys: text, start_time, end_time (ms).
-    """
-    at = _audio_to_tuple(audio_upload)
-    if at is None:
-        raise ValueError("Audio input is required for visualization.")
-    audio, sr = at
-
-    if not timestamps:
-        return "<div style='color:#666'>No timestamps to visualize.</div>"
-    if not isinstance(timestamps, list):
-        raise ValueError("Timestamps must be a list (JSON array).")
-
-    html_content = """
-    <style>
-        .word-alignment-container { display: flex; flex-wrap: wrap; gap: 10px; }
-        .word-box {
-            border: 1px solid #ddd; border-radius: 8px; padding: 10px;
-            background-color: #f9f9f9; box-shadow: 0 2px 4px rgba(0,0,0,0.06);
-            text-align: center;
-        }
-        .word-text { font-size: 18px; font-weight: 700; margin-bottom: 5px; }
-        .word-time { font-size: 12px; color: #666; margin-bottom: 8px; }
-        .word-audio audio { width: 140px; height: 30px; }
-        details { border: 1px solid #ddd; border-radius: 6px; padding: 10px; background-color: #f7f7f7; }
-        summary { font-weight: 700; cursor: pointer; }
-    </style>
-    """
-
-    html_content += """
-    <details open>
-        <summary>Timestamps Visualization (时间戳可视化结果）</summary>
-        <div class="word-alignment-container" style="margin-top: 14px;">
-    """
-
-    for item in timestamps:
-        if not isinstance(item, dict):
-            continue
-        word = str(item.get("text", "") or "")
-        start = item.get("start_time", None)
-        end = item.get("end_time", None)
-        if start is None or end is None:
-            continue
-
-        start = float(start)
-        end = float(end)
-        if end <= start:
-            continue
-
-        start_sample = max(0, int(start * sr))
-        end_sample = min(len(audio), int(end * sr))
-        if end_sample <= start_sample:
-            continue
-
-        seg = audio[start_sample:end_sample]
-        seg_i16 = (np.clip(seg, -1.0, 1.0) * 32767.0).astype(np.int16)
-
-        mem = io.BytesIO()
-        wav_write(mem, sr, seg_i16)
-        mem.seek(0)
-        b64 = base64.b64encode(mem.read()).decode("utf-8")
-        audio_src = f"data:audio/wav;base64,{b64}"
-
-        html_content += f"""
-        <div class="word-box">
-            <div class="word-text">{word}</div>
-            <div class="word-time">{start} - {end} s</div>
-            <div class="word-audio">
-                <audio controls preload="none" src="{audio_src}"></audio>
-            </div>
-        </div>
-        """
-
-    html_content += "</div></details>"
-    return html_content
-
-
-def build_demo(
-    asr: Qwen3ASRModel,
-    asr_ckpt: str,
-    backend: str,
-    aligner_ckpt: Optional[str] = None,
-) -> gr.Blocks:
-    supported_langs_raw = asr.get_supported_languages()
-    lang_choices_disp, lang_map = _build_choices_and_map([x for x in supported_langs_raw])
-    lang_choices = ["Auto"] + lang_choices_disp
-
-    has_aligner = bool(aligner_ckpt)
-
-    theme = gr.themes.Soft(
-        font=[gr.themes.GoogleFont("Source Sans Pro"), "Arial", "sans-serif"],
-    )
-    css = ".gradio-container {max-width: none !important;}"
-
-    with gr.Blocks(theme=theme, css=css) as demo:
-        gr.Markdown(
-            f"""
-# Qwen3 ASR Demo
-**Backend:** `{backend}`  
-**ASR Checkpoint:** `{asr_ckpt}`  
-**Forced Aligner:** `{aligner_ckpt if aligner_ckpt else "(none)"}`  
-"""
-        )
-
-        with gr.Row():
-            with gr.Column(scale=2):
-                audio_in = gr.Audio(label="Audio Input (上传音频)", type="numpy")
-                lang_in = gr.Dropdown(
-                    label="Language (语种)",
-                    choices=lang_choices,
-                    value="Auto",
-                    interactive=True,
-                )
-                if has_aligner:
-                    ts_in = gr.Checkbox(
-                        label="Return Timestamps (是否返回时间戳)",
-                        value=True,
-                    )
-                else:
-                    ts_in = gr.State(False)
-
-                btn = gr.Button("Transcribe (识别)", variant="primary")
-
-            with gr.Column(scale=2):
-                out_lang = gr.Textbox(label="Detected Language", lines=1)
-                out_text = gr.Textbox(label="Result Text", lines=12)
-
-            if has_aligner:
-                with gr.Column(scale=3):
-                    out_ts = gr.JSON(label="Timestamps（时间戳结果）")
-                    viz_btn = gr.Button("Visualize Timestamps (可视化时间戳)", variant="secondary")
-            else:
-                with gr.Column(scale=3):
-                    out_ts = gr.State(None)
-                    viz_btn = gr.State(None)
-
-        # Put the visualization panel below the three columns
-        if has_aligner:
-            with gr.Row():
-                out_ts_html = gr.HTML(label="Timestamps Visualization (时间戳可视化结果)")
-        else:
-            out_ts_html = gr.State("")
-
-        def run(audio_upload: Any, lang_disp: str, return_ts: bool):
-            audio_obj = _parse_audio_any(audio_upload)
-
-            language = None
-            if lang_disp and lang_disp != "Auto":
-                language = lang_map.get(lang_disp, lang_disp)
-
-            return_ts = bool(return_ts) and has_aligner
-
-            results = asr.transcribe(
-                audio=audio_obj,
-                language=language,
-                return_time_stamps=return_ts,
-            )
-            if not isinstance(results, list) or len(results) != 1:
-                raise RuntimeError(
-                    f"Unexpected result size: {type(results)} "
-                    f"len={len(results) if isinstance(results, list) else 'N/A'}"
-                )
-
-            r = results[0]
-
-            if has_aligner:
-                ts_payload = None
-                if return_ts:
-                    ts_payload = [
-                        dict(
-                            text=getattr(t, "text", None),
-                            start_time=getattr(t, "start_time", None),
-                            end_time=getattr(t, "end_time", None),
-                        )
-                        for t in (getattr(r, "time_stamps", None) or [])
-                    ]
-                return (
-                    getattr(r, "language", "") or "",
-                    getattr(r, "text", "") or "",
-                    gr.update(value=ts_payload) if return_ts else gr.update(value=None),
-                    gr.update(value=""),  # clear html on each transcribe
-                )
-            else:
-                return (
-                    getattr(r, "language", "") or "",
-                    getattr(r, "text", "") or "",
-                )
-
-        def visualize(audio_upload: Any, timestamps_json: Any):
-            return _make_timestamp_html(audio_upload, timestamps_json)
-
-        if has_aligner:
-            btn.click(
-                run,
-                inputs=[audio_in, lang_in, ts_in],
-                outputs=[out_lang, out_text, out_ts, out_ts_html],
-            )
-            viz_btn.click(
-                visualize,
-                inputs=[audio_in, out_ts],
-                outputs=[out_ts_html],
-            )
-        else:
-            btn.click(
-                run,
-                inputs=[audio_in, lang_in, ts_in],
-                outputs=[out_lang, out_text],
-            )
-
-    return demo
-
-
-def main(argv=None) -> int:
-    parser = build_parser()
-    args = parser.parse_args(argv)
-
-    _apply_cuda_visible_devices(args.cuda_visible_devices)
-
-    backend = args.backend
-    asr_ckpt = args.asr_checkpoint
-    aligner_ckpt = args.aligner_checkpoint
-
-    user_backend_kwargs = _parse_json_dict(args.backend_kwargs, name="--backend-kwargs")
-    user_aligner_kwargs = _parse_json_dict(args.aligner_kwargs, name="--aligner-kwargs")
-
-    backend_kwargs = _merge_dicts(_default_backend_kwargs(backend), user_backend_kwargs)
-    backend_kwargs = _coerce_special_types(backend_kwargs)
-
-    forced_aligner = None
-    forced_aligner_kwargs = None
-    if aligner_ckpt:
-        forced_aligner = aligner_ckpt
-        aligner_kwargs = _merge_dicts(_default_aligner_kwargs(), user_aligner_kwargs)
-        forced_aligner_kwargs = _coerce_special_types(aligner_kwargs)
-
-    if backend == "transformers":
-        asr = Qwen3ASRModel.from_pretrained(
-            asr_ckpt,
-            forced_aligner=forced_aligner,
-            forced_aligner_kwargs=forced_aligner_kwargs,
-            **backend_kwargs,
-        )
-    else:
-        asr = Qwen3ASRModel.LLM(
-            model=asr_ckpt,
-            forced_aligner=forced_aligner,
-            forced_aligner_kwargs=forced_aligner_kwargs,
-            **backend_kwargs,
-        )
-
-    demo = build_demo(asr, asr_ckpt, backend, aligner_ckpt=aligner_ckpt)
-
-    launch_kwargs: Dict[str, Any] = dict(
-        server_name=args.ip,
-        server_port=args.port,
-        share=args.share,
-        ssl_verify=True if args.ssl_verify else False,
-    )
-    if args.ssl_certfile is not None:
-        launch_kwargs["ssl_certfile"] = args.ssl_certfile
-    if args.ssl_keyfile is not None:
-        launch_kwargs["ssl_keyfile"] = args.ssl_keyfile
-
-    demo.queue(default_concurrency_limit=int(args.concurrency)).launch(**launch_kwargs)
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

qwen_asr/cli/serve.py

@@ -1,46 +0,0 @@
-# coding=utf-8
-# 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 sys
-
-from qwen_asr.core.transformers_backend import (
-    Qwen3ASRConfig,
-    Qwen3ASRForConditionalGeneration,
-    Qwen3ASRProcessor,
-)
-from transformers import AutoConfig, AutoModel, AutoProcessor
-
-AutoConfig.register("qwen3_asr", Qwen3ASRConfig)
-AutoModel.register(Qwen3ASRConfig, Qwen3ASRForConditionalGeneration)
-AutoProcessor.register(Qwen3ASRConfig, Qwen3ASRProcessor)
-
-try:
-    from qwen_asr.core.vllm_backend import Qwen3ASRForConditionalGeneration
-    from vllm import ModelRegistry
-    ModelRegistry.register_model("Qwen3ASRForConditionalGeneration", Qwen3ASRForConditionalGeneration)
-except Exception as e:
-    raise ImportError(
-        "vLLM is not available, to use qwen-asr-serve, please install with: pip install qwen-asr[vllm]"
-    ) from e
-
-from vllm.entrypoints.cli.main import main as vllm_main
-
-def main():
-    sys.argv.insert(1, "serve")
-    vllm_main()
-
-
-if __name__ == "__main__":
-    main()

qwen_asr/core/transformers_backend/__init__.py

@@ -1,18 +0,0 @@
-# coding=utf-8
-# 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.
-from .configuration_qwen3_asr import Qwen3ASRConfig
-from .modeling_qwen3_asr import Qwen3ASRForConditionalGeneration
-from .processing_qwen3_asr import Qwen3ASRProcessor

qwen_asr/core/transformers_backend/configuration_qwen3_asr.py

@@ -1,425 +0,0 @@
-# coding=utf-8
-# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# 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.
-from transformers.configuration_utils import PretrainedConfig
-from transformers.utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class Qwen3ASRAudioEncoderConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`Qwen3ASRAudioEncoder`]. It is used to instantiate a
-    Qwen3-ASR audio encoder according to the specified arguments, defining the model architecture. Instantiating a
-    configuration with the defaults will yield a similar configuration to that of the audio encoder of the Qwen2-Audio
-    architecture.
-
-    e.g. [Qwen/Qwen3-ASR-1.7B](https://huggingface.co/Qwen/Qwen3-ASR-1.7B)
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-    Args:
-        num_mel_bins (`int`, *optional*, defaults to 128):
-            Number of mel features used per input features. Should correspond to the value used in the
-            `Qwen3ASRProcessor` class.
-        encoder_layers (`int`, *optional*, defaults to 32):
-            Number of encoder layers.
-        encoder_attention_heads (`int`, *optional*, defaults to 20):
-            Number of attention heads for each attention layer in the Transformer encoder.
-        encoder_ffn_dim (`int`, *optional*, defaults to 5120):
-            Dimensionality of the "intermediate" (often named feed-forward) layer in encoder.
-        d_model (`int`, *optional*, defaults to 1280):
-            Dimensionality of the layers.
-        dropout (`float`, *optional*, defaults to 0.0):
-            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
-        attention_dropout (`float`, *optional*, defaults to 0.0):
-            The dropout ratio for the attention probabilities.
-        activation_function (`str`, *optional*, defaults to `"gelu"`):
-            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
-            `"relu"`, `"silu"` and `"gelu_new"` are supported.
-        activation_dropout (`float`, *optional*, defaults to 0.0):
-            The dropout ratio for activations inside the fully connected layer.
-        scale_embedding (`bool`, *optional*, defaults to `False`):
-            Scale embeddings by diving by sqrt(d_model).
-        initializer_range (`float`, *optional*, defaults to 0.02):
-            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-        max_source_positions (`int`, *optional*, defaults to 1500):
-            The maximum sequence length of log-mel filter-bank features that this model might ever be used with.
-        n_window (`int`, *optional*, defaults to 100):
-            The chunk for conv and flash attn in AudioEncoder.
-        output_dim (`int`, *optional*, defaults to 3584):
-            The output dimension of AudioEncoder.
-
-    Example:
-
-    ```python
-    >>> from transformers import Qwen3ASRAudioEncoderConfig, Qwen3ASRAudioEncoder
-
-    >>> # Initializing a Qwen3ASRAudioEncoderConfig
-    >>> configuration = Qwen3ASRAudioEncoderConfig()
-
-    >>> # Initializing a Qwen3ASRAudioEncoder (with random weights)
-    >>> model = Qwen3ASRAudioEncoder(configuration)
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config
-    ```"""
-
-    model_type = "qwen3_asr_audio_encoder"
-
-    def __init__(
-        self,
-        num_mel_bins=128,
-        encoder_layers=32,
-        encoder_attention_heads=20,
-        encoder_ffn_dim=5120,
-        d_model=1280,
-        dropout=0,
-        attention_dropout=0,
-        activation_function="gelu",
-        activation_dropout=0,
-        scale_embedding=False,
-        initializer_range=0.02,
-        max_source_positions=1500,
-        n_window=100,
-        output_dim=3584,
-        n_window_infer=400,
-        conv_chunksize=500,
-        downsample_hidden_size=480,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-
-        self.num_mel_bins = num_mel_bins
-        self.d_model = d_model
-        self.encoder_layers = encoder_layers
-        self.encoder_attention_heads = encoder_attention_heads
-        self.encoder_ffn_dim = encoder_ffn_dim
-        self.dropout = dropout
-        self.attention_dropout = attention_dropout
-        self.activation_function = activation_function
-        self.activation_dropout = activation_dropout
-        self.num_hidden_layers = encoder_layers
-        self.initializer_range = initializer_range
-        self.scale_embedding = scale_embedding  # scale factor will be sqrt(d_model) if True
-        self.max_source_positions = max_source_positions
-        self.n_window = n_window
-        self.output_dim = output_dim
-        self.n_window_infer = n_window_infer
-        self.conv_chunksize = conv_chunksize
-        self.downsample_hidden_size = downsample_hidden_size
-
-
-class Qwen3ASRTextConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`Qwen3ASRTextModel`]. It is used to instantiate a
-    Qwen3-ASR model according to the specified arguments, defining the model architecture. Instantiating a configuration
-    with the defaults will yield a similar configuration to that of
-    Qwen3-ASR-1.7B [Qwen/Qwen3-ASR-1.7B](https://huggingface.co/Qwen/Qwen3-ASR-1.7B)
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-    Args:
-        vocab_size (`int`, *optional*, defaults to 151936):
-            Vocabulary size of the Qwen3ASR model. Defines the number of different tokens that can be represented by the
-            `inputs_ids` passed when calling [`Qwen3ASRModel`]
-        hidden_size (`int`, *optional*, defaults to 4096):
-            Dimension of the hidden representations.
-        intermediate_size (`int`, *optional*, defaults to 22016):
-            Dimension of the MLP representations.
-        num_hidden_layers (`int`, *optional*, defaults to 32):
-            Number of hidden layers in the Transformer encoder.
-        num_attention_heads (`int`, *optional*, defaults to 32):
-            Number of attention heads for each attention layer in the Transformer encoder.
-        num_key_value_heads (`int`, *optional*, defaults to 32):
-            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
-            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
-            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
-            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
-            by meanpooling all the original heads within that group. For more details, check out [this
-            paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`.
-        head_dim (`int`, *optional*, defaults to 128):
-            The dimension of the head. If not specified, will default to `hidden_size // num_attention_heads`.
-        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
-            The non-linear activation function (function or string) in the decoder.
-        max_position_embeddings (`int`, *optional*, defaults to 128000):
-            The maximum sequence length that this model might ever be used with.
-        initializer_range (`float`, *optional*, defaults to 0.02):
-            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
-            The epsilon used by the rms normalization layers.
-        use_cache (`bool`, *optional*, defaults to `True`):
-            Whether or not the model should return the last key/values attentions (not used by all models). Only
-            relevant if `config.is_decoder=True`.
-        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
-            Whether the model's input and output word embeddings should be tied.
-        rope_theta (`float`, *optional*, defaults to 5000000.0):
-            The base period of the RoPE embeddings.
-        rope_scaling (`Dict`, *optional*):
-            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
-            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
-            accordingly.
-            Expected contents:
-                `rope_type` (`str`):
-                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
-                    'llama3'], with 'default' being the original RoPE implementation.
-                `factor` (`float`, *optional*):
-                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
-                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
-                    original maximum pre-trained length.
-                `original_max_position_embeddings` (`int`, *optional*):
-                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
-                    pretraining.
-                `attention_factor` (`float`, *optional*):
-                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
-                    computation. If unspecified, it defaults to value recommended by the implementation, using the
-                    `factor` field to infer the suggested value.
-                `beta_fast` (`float`, *optional*):
-                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
-                    ramp function. If unspecified, it defaults to 32.
-                `beta_slow` (`float`, *optional*):
-                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
-                    ramp function. If unspecified, it defaults to 1.
-                `short_factor` (`list[float]`, *optional*):
-                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
-                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
-                    size divided by the number of attention heads divided by 2
-                `long_factor` (`list[float]`, *optional*):
-                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
-                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
-                    size divided by the number of attention heads divided by 2
-                `low_freq_factor` (`float`, *optional*):
-                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
-                `high_freq_factor` (`float`, *optional*):
-                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
-        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
-            Whether to use a bias in the query, key, value and output projection layers during self-attention.
-        attention_dropout (`float`, *optional*, defaults to 0.0):
-            The dropout ratio for the attention probabilities.
-
-    ```python
-    >>> from transformers import Qwen3ASRTextModel, Qwen3ASRTextConfig
-
-    >>> # Initializing a Qwen3ASR style configuration
-    >>> configuration = Qwen3ASRTextConfig()
-
-    >>> # Initializing a model from the Qwen3-VL-7B style configuration
-    >>> model = Qwen3ASRTextModel(configuration)
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config
-    ```"""
-
-    model_type = "qwen3_asr_text"
-    base_config_key = "text_config"
-
-    def __init__(
-        self,
-        vocab_size=151936,
-        hidden_size=4096,
-        intermediate_size=22016,
-        num_hidden_layers=32,
-        num_attention_heads=32,
-        num_key_value_heads=32,
-        head_dim=128,
-        hidden_act="silu",
-        max_position_embeddings=128000,
-        initializer_range=0.02,
-        rms_norm_eps=1e-6,
-        use_cache=True,
-        tie_word_embeddings=False,
-        rope_theta=5000000.0,
-        rope_scaling=None,
-        attention_bias=False,
-        attention_dropout=0.0,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-
-        # for backward compatibility
-        if num_key_value_heads is None:
-            num_key_value_heads = num_attention_heads
-
-        self.num_key_value_heads = num_key_value_heads
-        self.head_dim = head_dim
-        self.hidden_act = hidden_act
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling
-        self.attention_bias = attention_bias
-        self.attention_dropout = attention_dropout
-        # Validate the correctness of rotary position embeddings parameters
-        # BC: if there is a 'type' field, move it to 'rope_type'.
-        if self.rope_scaling is not None and "type" in self.rope_scaling:
-            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
-
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
-
-
-class Qwen3ASRThinkerConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`Qwen3ASRThinker`]. It is used to instantiate a
-    Qwen3-ASR-Thinker model according to the specified arguments, defining the model architecture. Instantiating a
-    configuration with the defaults will yield a similar configuration to that of the thinker component of the Qwen3-Omni
-    architecture.
-
-    e.g. [Qwen/Qwen3-ASR-1.7B](https://huggingface.co/Qwen/Qwen3-ASR-1.7B)
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-    Args:
-        audio_config (`dict`, *optional*):
-            The config dictionary of the audio backbone.
-        text_config (`dict`, *optional*):
-            The config dictionary of the text backbone.
-        audio_token_id (`int`, *optional*, defaults to 151646):
-            The audio token id to encode the audio prompt.
-        audio_start_token_id (`int`, *optional*, defaults to 151647):
-            The audio start token id to encode the audio prompt.
-        user_token_id (`int`, *optional*, defaults to 872):
-            The user token id to encode the user token.
-        initializer_range (`float`, *optional*, defaults to 0.02):
-            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-
-    Example:
-
-    ```python
-    >>> from transformers import Qwen3ASRThinkerModel, Qwen3ASRThinkerConfig
-
-    >>> # Initializing a default Qwen3ASRThinkerConfig
-    >>> configuration = Qwen3ASRThinkerConfig()
-
-    >>> # Initializing a model (with random weights) from the default configuration
-    >>> model = Qwen3ASRThinkerModel(configuration)
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config
-    ```"""
-
-    model_type = "qwen3_asr_thinker"
-
-    attribute_map = {}
-    sub_configs = {
-        "audio_config": Qwen3ASRAudioEncoderConfig,
-        "text_config": Qwen3ASRTextConfig,
-    }
-
-    def __init__(
-        self,
-        audio_config=None,
-        text_config=None,
-        audio_token_id=151646,
-        audio_start_token_id=151647,
-        user_token_id=872,
-        initializer_range=0.02,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.user_token_id = user_token_id
-        self.audio_start_token_id = audio_start_token_id
-        self.initializer_range = initializer_range
-
-        if isinstance(audio_config, dict):
-            audio_config = Qwen3ASRAudioEncoderConfig(**audio_config)
-        elif audio_config is None:
-            audio_config = Qwen3ASRAudioEncoderConfig()
-        self.audio_config = audio_config
-
-        if isinstance(text_config, dict):
-            text_config = Qwen3ASRTextConfig(**text_config)
-        elif text_config is None:
-            text_config = Qwen3ASRTextConfig()
-        self.text_config = text_config
-        self.audio_token_id = audio_token_id
-
-
-class Qwen3ASRConfig(PretrainedConfig):
-    """
-    This is the configuration class to store the configuration of a [`Qwen3ASRForConditionalGeneration`]. It is used to instantiate a Qwen3ASR
-    model according to the specified sub-models configurations, defining the model architecture.
-
-    Instantiating a configuration with the defaults will yield a similar configuration to that of the
-    [Qwen/Qwen3-ASR-1.7B](https://huggingface.co/Qwen/Qwen3-ASR-1.7B) architecture.
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-    Args:
-        thinker_config (`dict`, *optional*): Configuration of the underlying thinker sub-model.
-        support_languages (`List[str]`, *optional*): The languages supported by the model.
-
-    Example:
-
-    ```python
-    >>> from transformers import (
-    ...     Qwen3ASRThinkerConfig,
-    ...     Qwen3ASRForConditionalGeneration,
-    ...     Qwen3ASRConfig,
-    ... )
-
-    >>> # Initializing a Qwen3ASR style configuration
-    >>> configuration = Qwen3ASRConfig()
-
-    >>> # Initializing a model from the configuration
-    >>> model = Qwen3ASRForConditionalGeneration(configuration)
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config
-    ```"""
-
-    model_type = "qwen3_asr"
-    sub_configs = {
-        "thinker_config": Qwen3ASRThinkerConfig,
-    }
-
-    def __init__(
-        self,
-        thinker_config=None,
-        support_languages=None,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        if thinker_config is None:
-            thinker_config = {}
-
-        self.thinker_config = Qwen3ASRThinkerConfig(**thinker_config)
-        self.support_languages = support_languages
-
-    def get_text_config(self, decoder=False) -> "PretrainedConfig":
-        """
-        Returns the config that is meant to be used with text IO. On most models, it is the original config instance
-        itself. On specific composite models, it is under a set of valid names.
-
-        Args:
-            decoder (`Optional[bool]`, *optional*, defaults to `False`):
-                If set to `True`, then only search for decoder config names.
-        """
-        # Overridden for deeply nested config like Qwen2.5-Omni. We don't have any omni model
-        # except for Qwen yet. This has to be generalized if more deeply nested configs are
-        # added. NOTE: currently method used only by vLLM
-        return self.thinker_config.get_text_config()
-
-
-__all__ = ["Qwen3ASRConfig", "Qwen3ASRThinkerConfig", "Qwen3ASRAudioEncoderConfig"]

qwen_asr/core/transformers_backend/modeling_qwen3_asr.py

@@ -1,1361 +0,0 @@
-# coding=utf-8
-# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# 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 math
-from dataclasses import dataclass
-from typing import Callable, Optional, Union
-
-import numpy as np
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache, DynamicCache
-from transformers.generation import GenerationMixin
-from transformers.integrations import use_kernel_forward_from_hub
-from transformers.masking_utils import create_causal_mask
-from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
-from transformers.modeling_layers import GradientCheckpointingLayer
-from transformers.modeling_outputs import (
-    BaseModelOutput,
-    BaseModelOutputWithPast,
-    MoeCausalLMOutputWithPast,
-)
-from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
-from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
-from transformers.processing_utils import Unpack
-from transformers.utils import auto_docstring, can_return_tuple
-from transformers.utils.deprecation import deprecate_kwarg
-from transformers.utils.generic import TransformersKwargs, check_model_inputs
-
-from .configuration_qwen3_asr import (
-    Qwen3ASRAudioEncoderConfig,
-    Qwen3ASRConfig,
-    Qwen3ASRThinkerConfig,
-)
-
-
-@use_kernel_forward_from_hub("RMSNorm")
-class Qwen3ASRTextRMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
-        """
-        Qwen3ASRTextRMSNorm is equivalent to T5LayerNorm
-        """
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states.to(input_dtype)
-
-    def extra_repr(self):
-        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
-
-
-def rotate_half(x):
-    """Rotates half the hidden dims of the input."""
-    x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2 :]
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """
-    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
-    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
-    """
-    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
-    if n_rep == 1:
-        return hidden_states
-    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
-    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
-
-
-def eager_attention_forward(
-    module: nn.Module,
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    attention_mask: Optional[torch.Tensor],
-    scaling: float,
-    dropout: float = 0.0,
-    **kwargs: Unpack[TransformersKwargs],
-):
-    key_states = repeat_kv(key, module.num_key_value_groups)
-    value_states = repeat_kv(value, module.num_key_value_groups)
-
-    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
-    if attention_mask is not None:
-        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
-        attn_weights = attn_weights + causal_mask
-
-    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
-    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
-    attn_output = torch.matmul(attn_weights, value_states)
-    attn_output = attn_output.transpose(1, 2).contiguous()
-
-    return attn_output, attn_weights
-
-
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
-    """Applies Rotary Position Embedding to the query and key tensors.
-
-    Args:
-        q (`torch.Tensor`): The query tensor.
-        k (`torch.Tensor`): The key tensor.
-        cos (`torch.Tensor`): The cosine part of the rotary embedding.
-        sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`, *optional*):
-            Deprecated and unused.
-        unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
-    Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
-    """
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
-
-
-class Qwen3ASRTextAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
-
-    def __init__(self, config: Qwen3ASRConfig, layer_idx: int):
-        super().__init__()
-        self.config = config
-        self.layer_idx = layer_idx
-        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
-        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
-        self.scaling = self.head_dim**-0.5
-        self.attention_dropout = config.attention_dropout
-        self.is_causal = True
-
-        self.q_proj = nn.Linear(
-            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.k_proj = nn.Linear(
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.v_proj = nn.Linear(
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.o_proj = nn.Linear(
-            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
-        )
-        self.q_norm = Qwen3ASRTextRMSNorm(
-            self.head_dim, eps=config.rms_norm_eps
-        )  # unlike olmo, only on the head dim!
-        self.k_norm = Qwen3ASRTextRMSNorm(
-            self.head_dim, eps=config.rms_norm_eps
-        )  # thus post q_norm does not need reshape
-
-    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        position_embeddings: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: Optional[torch.Tensor],
-        past_key_values: Optional[Cache] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        **kwargs: Unpack[FlashAttentionKwargs],
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
-        input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, self.head_dim)
-
-        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
-        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
-        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-
-        cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
-        if past_key_values is not None:
-            # sin and cos are specific to RoPE models; cache_position needed for the static cache
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
-        attention_interface: Callable = eager_attention_forward
-        if self.config._attn_implementation != "eager":
-            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
-
-        attn_output, attn_weights = attention_interface(
-            self,
-            query_states,
-            key_states,
-            value_states,
-            attention_mask,
-            dropout=0.0 if not self.training else self.attention_dropout,
-            scaling=self.scaling,
-            **kwargs,
-        )
-
-        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
-        attn_output = self.o_proj(attn_output)
-        return attn_output, attn_weights
-
-
-class Qwen3ASRTextMLP(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = config.intermediate_size
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_act]
-
-    def forward(self, x):
-        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
-        return down_proj
-
-
-class Qwen3ASRThinkerTextDecoderLayer(GradientCheckpointingLayer):
-    def __init__(self, config: Qwen3ASRConfig, layer_idx: int):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-
-        self.self_attn = Qwen3ASRTextAttention(config=config, layer_idx=layer_idx)
-
-        self.mlp = Qwen3ASRTextMLP(config)
-        self.input_layernorm = Qwen3ASRTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = Qwen3ASRTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        position_embeddings: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Cache] = None,
-        use_cache: Optional[bool] = False,
-        cache_position: Optional[torch.LongTensor] = None,
-        **kwargs: Unpack[TransformersKwargs],
-    ) -> torch.Tensor:
-        residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
-        # Self Attention
-        hidden_states, _ = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            use_cache=use_cache,
-            cache_position=cache_position,
-            position_embeddings=position_embeddings,
-            **kwargs,
-        )
-        hidden_states = residual + hidden_states
-
-        # Fully Connected
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
-        return hidden_states
-
-
-@auto_docstring
-class Qwen3ASRPreTrainedModel(PreTrainedModel):
-    config: Qwen3ASRConfig
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = True
-    _skip_keys_device_placement = "past_key_values"
-    _supports_flash_attn = True
-    _supports_sdpa = True
-
-    _can_compile_fullgraph = True
-    _supports_attention_backend = True
-    _can_record_outputs = {
-        "attentions": Qwen3ASRTextAttention,
-    }
-
-
-@dataclass
-class Qwen3ASRThinkerCausalLMOutputWithPast(MoeCausalLMOutputWithPast):
-    r"""
-    Args:
-        rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
-            The rope index difference between sequence length and multimodal rope.
-    """
-
-    rope_deltas: Optional[torch.LongTensor] = None
-
-
-def _get_feat_extract_output_lengths(input_lengths):
-    """
-    Computes the output length of the convolutional layers and the output length of the audio encoder
-    """
-
-    input_lengths_leave = input_lengths % 100
-    feat_lengths = (input_lengths_leave - 1) // 2 + 1
-    output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
-    return output_lengths
-
-
-class Qwen3ASRPreTrainedModelForConditionalGeneration(Qwen3ASRPreTrainedModel):
-    def _prepare_4d_causal_attention_mask_with_cache_position(
-        self,
-        attention_mask: torch.Tensor,
-        sequence_length: int,
-        target_length: int,
-        dtype: torch.dtype,
-        device: torch.device,
-        min_dtype: float,
-        cache_position: torch.Tensor,
-        batch_size: int,
-    ):
-        """
-        Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
-        `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
-
-        Args:
-            attention_mask (`torch.Tensor`):
-                A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
-            sequence_length (`int`):
-                The sequence length being processed.
-            target_length (`int`):
-                The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
-            dtype (`torch.dtype`):
-                The dtype to use for the 4D attention mask.
-            device (`torch.device`):
-                The device to place the 4D attention mask on.
-            min_dtype (`float`):
-                The minimum value representable with the dtype `dtype`.
-            cache_position (`torch.Tensor`):
-                Indices depicting the position of the input sequence tokens in the sequence.
-            batch_size (`torch.Tensor`):
-                Batch size.
-        """
-        if attention_mask is not None and attention_mask.dim() == 4:
-            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
-            causal_mask = attention_mask
-        else:
-            causal_mask = torch.full(
-                (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
-            )
-            if sequence_length != 1:
-                causal_mask = torch.triu(causal_mask, diagonal=1)
-            causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
-            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
-            if attention_mask is not None:
-                causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
-                mask_length = attention_mask.shape[-1]
-                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
-                padding_mask = padding_mask == 0
-                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
-                    padding_mask, min_dtype
-                )
-
-        return causal_mask
-
-
-    def get_chunked_index(
-        self, token_indices: torch.Tensor, tokens_per_chunk: int, remove_index: int
-    ) -> list[tuple[int, int]]:
-        """
-        Splits token index list into chunks based on token value ranges.
-
-        Given a list of token indices, returns a list of (start, end) index tuples representing
-        slices of the list where the token values fall within successive ranges of `t_ntoken_per_chunk`.
-
-        For example, if `t_ntoken_per_chunk` is 1000, the function will create chunks such that:
-        - the first chunk contains token values < 1000,
-        - the second chunk contains values >= 1000 and < 2000, and so on.
-
-        Parameters:
-            token_indices (`torch.Tensor` of shape `(seq_len, )`): A monotonically increasing list of
-                                token index values.
-            t_ntoken_per_chunk (`int`): Number of tokens per chunk (used as the chunk size threshold).
-            remove_index (`int`) An index id to subtract from `token_indices` before chunking
-
-        Returns:
-            `list[tuple[int, int]]`: A list of tuples, each representing the start (inclusive)
-                                and end (exclusive) indices of a chunk in `token_indices`.
-        """
-
-        def _iter():
-            i, start_idx = 0, 0  # skip bos token
-            current_chunk = 1
-            while i < len(token_indices):  # skip eos token
-                if token_indices[i] - remove_index >= current_chunk * tokens_per_chunk:
-                    yield (start_idx, i)
-                    start_idx = i
-                    current_chunk += 1
-                i += 1
-            yield (start_idx, len(token_indices))
-
-        return list(_iter())
-
-    def get_rope_index(
-        self,
-        attention_mask: Optional[torch.Tensor] = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Calculate the rope index in LLM.
-
-        Explanation:
-            Each embedding sequence contains text embedding.
-
-        Args:
-            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
-                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
-                it.
-            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
-                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
-                - 1 for tokens that are **not masked**,
-                - 0 for tokens that are **masked**.
-            audio_seqlens (`torch.LongTensor` of shape `(num_audios)`, *optional*):
-                The length of feature shape of each audio in LLM.
-
-        Returns:
-            position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
-            mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
-        """
-        mrope_position_deltas = []
-
-        position_ids = attention_mask.float().cumsum(-1) - 1
-        position_ids.masked_fill_(attention_mask == 0, 1)
-        position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
-        max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
-        mrope_position_deltas = max_position_ids + 1 - torch.sum(attention_mask, dim=-1, keepdim=True)
-
-        return position_ids, mrope_position_deltas
-
-
-class Qwen3ASRAudioAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
-
-    def __init__(self, config):
-        super().__init__()
-        self.embed_dim = config.d_model
-        self.num_heads = config.encoder_attention_heads
-        self.dropout = config.attention_dropout
-        self.head_dim = self.embed_dim // self.num_heads
-        self.num_key_value_groups = 1  # needed for eager attention
-        self.config = config
-
-        if (self.head_dim * self.num_heads) != self.embed_dim:
-            raise ValueError(
-                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
-                f" and `num_heads`: {self.num_heads})."
-            )
-        self.scaling = self.head_dim**-0.5
-        self.attention_dropout = 0.0
-        self.is_decoder = False
-        self.is_causal = False
-        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
-        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
-        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
-        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        cu_seqlens: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        **kwargs,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
-        """Input shape: Batch x Time x Channel"""
-
-        seq_length, _ = hidden_states.size()
-
-        query_states = self.q_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
-        key_states = self.k_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
-        value_states = self.v_proj(hidden_states).reshape(seq_length, self.num_heads, -1)
-
-        query_states = query_states.transpose(0, 1).unsqueeze(0)
-        key_states = key_states.transpose(0, 1).unsqueeze(0)
-        value_states = value_states.transpose(0, 1).unsqueeze(0)
-        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
-
-        attention_interface: Callable = eager_attention_forward
-        if self.config._attn_implementation != "eager":
-            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
-
-        attn_output, _ = attention_interface(
-            self,
-            query_states,
-            key_states,
-            value_states,
-            attention_mask=attention_mask,
-            dropout=0.0 if not self.training else self.attention_dropout,
-            scaling=self.scaling,
-            cu_seq_lens_q=cu_seqlens,  # pass cu seq lens for FA2
-            cu_seq_lens_k=cu_seqlens,
-            max_length_q=max_seqlen,
-            max_length_k=max_seqlen,
-            is_causal=False,
-            **kwargs,
-        )
-
-        attn_output = attn_output.reshape(seq_length, -1).contiguous()
-        attn_output = self.out_proj(attn_output)
-
-        return attn_output
-
-
-class Qwen3ASRAudioEncoderLayer(GradientCheckpointingLayer):
-    def __init__(self, config: Qwen3ASRAudioEncoderConfig):
-        super().__init__()
-        self.embed_dim = config.d_model
-        self.self_attn = Qwen3ASRAudioAttention(config)
-        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
-        self.dropout = config.dropout
-        self.activation_fn = ACT2FN[config.activation_function]
-        self.activation_dropout = config.activation_dropout
-        self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
-        self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
-        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        cu_seqlens: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        **kwargs,
-    ) -> torch.Tensor:
-        """
-        Args:
-            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
-            attention_mask (`torch.FloatTensor`): attention mask of size
-                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
-            layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
-                `(encoder_attention_heads,)`.
-            output_attentions (`bool`, *optional*):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
-        """
-        residual = hidden_states
-        hidden_states = self.self_attn_layer_norm(hidden_states)
-        hidden_states = self.self_attn(
-            hidden_states=hidden_states,
-            cu_seqlens=cu_seqlens,
-            attention_mask=attention_mask,
-            **kwargs,
-        )
-        hidden_states = residual + hidden_states
-        residual = hidden_states
-        hidden_states = self.final_layer_norm(hidden_states)
-        hidden_states = self.fc1(hidden_states)
-        hidden_states = self.activation_fn(hidden_states)
-        hidden_states = self.fc2(hidden_states)
-        hidden_states = residual + hidden_states
-
-        if hidden_states.dtype == torch.float16:
-            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
-            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
-
-        outputs = (hidden_states,)
-
-        return outputs
-
-
-class SinusoidsPositionEmbedding(nn.Module):
-    def __init__(self, length, channels, max_timescale=10000):
-        super().__init__()
-        if channels % 2 != 0:
-            raise ValueError("SinusoidsPositionEmbedding needs even channels input")
-        log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
-        inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2).float())
-        scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
-        self.register_buffer(
-            "positional_embedding",
-            torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1),
-            persistent=False,
-        )
-
-    def forward(self, seqlen: int):
-        return self.positional_embedding[:seqlen, :]
-
-
-@auto_docstring(
-    custom_intro="""
-    Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
-    [`Qwen3ASRAudioEncoderLayer`].
-    """
-)
-class Qwen3ASRAudioEncoder(Qwen3ASRPreTrainedModel):
-    config: Qwen3ASRAudioEncoderConfig
-    main_input_name = "input_features"
-    _no_split_modules = ["Qwen3ASRAudioEncoderLayer"]
-    _supports_sdpa = True
-
-    def __init__(self, config: Qwen3ASRAudioEncoderConfig):
-        super().__init__(config)
-        self.dropout = config.dropout
-
-        embed_dim = config.d_model
-        self.num_mel_bins = config.num_mel_bins
-        self.max_source_positions = config.max_source_positions
-        self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
-        self.n_window = config.n_window
-        self.positional_embedding = SinusoidsPositionEmbedding(self.max_source_positions, embed_dim)
-        self.layers = nn.ModuleList([Qwen3ASRAudioEncoderLayer(config) for _ in range(config.encoder_layers)])
-        self.ln_post = nn.LayerNorm(config.d_model)
-        self.gradient_checkpointing = False
-        self.conv2d1 = nn.Conv2d(1, config.downsample_hidden_size, 3, 2, padding=1)
-        self.conv2d2 = nn.Conv2d(config.downsample_hidden_size, config.downsample_hidden_size, 3, 2, padding=1)
-        self.conv2d3 = nn.Conv2d(config.downsample_hidden_size, config.downsample_hidden_size, 3, 2, padding=1)
-        self.conv_out = nn.Linear(
-            config.downsample_hidden_size * ((((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2),
-            config.d_model,
-            bias=False,
-        )
-        self.proj1 = nn.Linear(config.d_model, config.d_model)
-        self.act = ACT2FN[config.activation_function]
-        self.proj2 = nn.Linear(config.d_model, config.output_dim)
-        self.n_window_infer = self.config.n_window_infer
-        self.conv_chunksize = self.config.conv_chunksize
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def _freeze_parameters(self):
-        for param in self.parameters():
-            param.requires_grad = False
-        self._requires_grad = False
-
-    def get_input_embeddings(self) -> nn.Module:
-        return self.conv1
-
-    def set_input_embeddings(self, value: nn.Module):
-        self.conv1 = value
-
-    def _prepare_attention_mask(self, inputs_tensor: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
-        # Flash Attention 2 doesn't need a 4D mask and relies on `cu_seqlens/max_seqlen`
-        # NOTE: the created attention masl only approximates the ragged FA2 attention by
-        # allowing bidirectional attention within `cu_seqlens` blocks, and not attending between
-        # blocks. Though it will not be a 100% match for FA2's `varlen` path
-        if self.config._attn_implementation == "flash_attention_2":
-            return None
-
-        seq_length = inputs_tensor.shape[0]
-        attention_mask = torch.full(
-            [1, 1, seq_length, seq_length],
-            torch.finfo(inputs_tensor.dtype).min,
-            device=inputs_tensor.device,
-            dtype=inputs_tensor.dtype,
-        )
-        for i in range(1, len(cu_seqlens)):
-            attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0
-        return attention_mask
-
-    @auto_docstring
-    def forward(
-        self,
-        input_features,
-        feature_lens=None,
-        aftercnn_lens=None,
-    ):
-        r"""
-        feature_lens (`torch.LongTensor` of shape `(batch_size,)`):
-            mel length
-        aftercnn_lens (`torch.LongTensor` of shape `(batch_size,)`):
-            mel length after cnn
-        """
-        aftercnn_lens = _get_feat_extract_output_lengths(feature_lens)
-        chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()
-
-        chunk_lengths = torch.tensor(
-            [self.n_window * 2] * chunk_num.sum(),
-            dtype=torch.long,
-            device=feature_lens.device,
-        )
-        tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
-        chunk_lengths[tail_chunk_index] = feature_lens % (self.n_window * 2)
-        chunk_lengths[chunk_lengths == 0] = self.n_window * 2
-
-        chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)
-        padded_feature = nn.utils.rnn.pad_sequence(chunk_list, batch_first=True).transpose(1, 2)
-        feature_lens_after_cnn = _get_feat_extract_output_lengths(chunk_lengths)
-        padded_mask_after_cnn = nn.utils.rnn.pad_sequence(
-            [torch.ones(length, dtype=torch.bool, device=padded_feature.device) for length in feature_lens_after_cnn],
-            batch_first=True,
-        )
-        padded_feature = padded_feature.unsqueeze(1)
-        # Split to chunk to avoid OOM during convolution
-        padded_embeds = []
-        for chunk in padded_feature.split(self.conv_chunksize, dim=0):
-            padded_embed = F.gelu(self.conv2d1(chunk))
-            padded_embed = F.gelu(self.conv2d2(padded_embed))
-            padded_embed = F.gelu(self.conv2d3(padded_embed))
-            padded_embeds.append(padded_embed)
-        padded_embed = torch.cat(padded_embeds, dim=0)
-        b, c, f, t = padded_embed.size()
-        padded_embed = self.conv_out(padded_embed.permute(0, 3, 1, 2).contiguous().view(b, t, c * f))
-
-        positional_embedding = (
-            self.positional_embedding.positional_embedding[: padded_embed.shape[1], :]
-            .unsqueeze(0)
-            .to(padded_embed.dtype)
-        )
-        padded_embed = padded_embed + positional_embedding
-        hidden_states = padded_embed[padded_mask_after_cnn]
-        cu_chunk_lens = [0]
-        window_aftercnn = padded_mask_after_cnn.shape[-1] * (self.n_window_infer // (self.n_window * 2))
-        for cnn_len in aftercnn_lens:
-            cu_chunk_lens += [window_aftercnn] * (cnn_len // window_aftercnn)
-            remainder = cnn_len % window_aftercnn
-            if remainder != 0:
-                cu_chunk_lens += [remainder]
-        cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum(-1, dtype=torch.int32)
-
-        for encoder_layer in self.layers:
-            layer_outputs = encoder_layer(
-                hidden_states,
-                cu_seqlens,
-            )
-
-            hidden_states = layer_outputs[0]
-
-        hidden_states = self.ln_post(hidden_states)
-        hidden_states = self.proj1(hidden_states)
-        hidden_states = self.act(hidden_states)
-        hidden_states = self.proj2(hidden_states)
-        return BaseModelOutput(last_hidden_state=hidden_states)
-
-    def padded_and_mask_function(self, tensor_list, tensor_len, padding_value=0, padding_side="right"):
-        """
-        Pads a sequence of tensors to their maximum length on indicated `padding_side`.
-        Then prepares a mask so that pad tokens are not attended to.
-        """
-        max_len = tensor_len.max()
-        dim = tensor_list[0].shape[0]
-        padded_tensor = torch.full(
-            size=(len(tensor_list), dim, max_len),
-            fill_value=padding_value,
-            dtype=self.dtype,
-            device=tensor_list[0].device,
-        )
-
-        batch_mask = torch.zeros(
-            (len(tensor_len), max_len),
-            dtype=torch.long,
-            device=padded_tensor.device,
-        )
-        for i, length in enumerate(tensor_len):
-            batch_mask[i, :length] = 1
-            padded_tensor[i, :, :length] = tensor_list[i]
-
-        feature_lens_after_cnn = (tensor_len - 1) // 2 + 1
-        max_len_after_cnn = feature_lens_after_cnn.max()
-        batch_mask_after_cnn = torch.zeros(
-            (len(tensor_len), max_len_after_cnn),
-            dtype=torch.long,
-            device=padded_tensor.device,
-        )
-        for i, length in enumerate(feature_lens_after_cnn):
-            batch_mask_after_cnn[i, :length] = 1
-        return (
-            padded_tensor,
-            batch_mask.unsqueeze(1),
-            batch_mask_after_cnn.bool(),
-        )
-
-
-class Qwen3ASRThinkerTextRotaryEmbedding(nn.Module):
-    inv_freq: torch.Tensor  # fix linting for `register_buffer`
-
-    def __init__(self, config: Qwen3ASRConfig, device=None):
-        super().__init__()
-        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
-            self.rope_type = config.rope_scaling.get("rope_type", "default")
-        else:
-            self.rope_type = "default"
-        self.max_seq_len_cached = config.max_position_embeddings
-        self.original_max_seq_len = config.max_position_embeddings
-
-        self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
-
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        self.mrope_section = config.rope_scaling.get("mrope_section", [24, 20, 20])
-
-    def apply_interleaved_mrope(self, freqs, mrope_section):
-        """Apply interleaved MRoPE to 3D rotary embeddings.
-        Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
-        interleaved [THTHWHTHW...TT], preserving frequency continuity.
-        args:
-            x: (3, bs, seq_len, head_dim // 2)
-            mrope_section: (3,)
-        returns:
-            x_t: (bs, seq_len, head_dim // 2)
-        """
-        freqs_t = freqs[0]  # just overwrite the first dimension T
-        for dim, offset in enumerate((1, 2), start=1):  # H, W
-            length = mrope_section[dim] * 3
-            idx = slice(offset, length, 3)
-            freqs_t[..., idx] = freqs[dim, ..., idx]
-        return freqs_t
-
-    @torch.no_grad()
-    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
-    def forward(self, x, position_ids):
-        # In contrast to other models, Qwen3ASRThinker has different position ids for the grids
-        # So we expand the inv_freq to shape (3, ...)
-        if position_ids.ndim == 2:
-            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
-        inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
-        position_ids_expanded = position_ids[:, :, None, :].float()  # shape (3, bs, 1, positions)
-
-        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
-            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
-            freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
-            emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos() * self.attention_scaling
-            sin = emb.sin() * self.attention_scaling
-
-        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
-
-
-class Qwen3ASRThinkerTextMLP(nn.Module):
-    def __init__(self, config, intermediate_size=None):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_act]
-
-    def forward(self, x):
-        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
-        return down_proj
-
-
-@use_kernel_forward_from_hub("RMSNorm")
-class Qwen3ASRThinkerTextRMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
-        """
-        Qwen3ASRThinkerTextRMSNorm is equivalent to T5LayerNorm
-        """
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states):
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states.to(input_dtype)
-
-    def extra_repr(self):
-        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
-
-
-class Qwen3ASRThinkerTextAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
-
-    def __init__(self, config, layer_idx):
-        super().__init__()
-        self.config = config
-        self.layer_idx = layer_idx
-        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
-        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
-        self.scaling = self.head_dim**-0.5
-        self.attention_dropout = config.attention_dropout
-        self.is_causal = True
-
-        self.q_proj = nn.Linear(
-            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.k_proj = nn.Linear(
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.v_proj = nn.Linear(
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.o_proj = nn.Linear(
-            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
-        )
-        self.q_norm = Qwen3ASRThinkerTextRMSNorm(
-            self.head_dim, eps=config.rms_norm_eps
-        )  # unlike olmo, only on the head dim!
-        self.k_norm = Qwen3ASRThinkerTextRMSNorm(
-            self.head_dim, eps=config.rms_norm_eps
-        )  # thus post q_norm does not need reshape
-        self.sliding_window = None
-
-    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        position_embeddings: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: Optional[torch.Tensor],
-        past_key_values: Optional[Cache] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        **kwargs: Unpack[FlashAttentionKwargs],
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
-        input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, self.head_dim)
-
-        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
-        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
-        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-
-        cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
-        if past_key_values is not None:
-            # sin and cos are specific to RoPE models; cache_position needed for the static cache
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
-        attention_interface: Callable = eager_attention_forward
-        if self.config._attn_implementation != "eager":
-            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
-
-        attn_output, attn_weights = attention_interface(
-            self,
-            query_states,
-            key_states,
-            value_states,
-            attention_mask,
-            dropout=0.0 if not self.training else self.attention_dropout,
-            scaling=self.scaling,
-            sliding_window=self.sliding_window,  # diff with Llama
-            **kwargs,
-        )
-
-        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
-        attn_output = self.o_proj(attn_output)
-        return attn_output, attn_weights
-
-
-@auto_docstring(
-    custom_intro=(
-        "Text part of Qwen3ASRThinker, "
-    )
-)
-class Qwen3ASRThinkerTextModel(Qwen3ASRPreTrainedModel):
-    config: Qwen3ASRConfig
-    _no_split_modules = ["Qwen3ASRThinkerTextDecoderLayer"]
-    config_class = Qwen3ASRConfig
-    _can_record_outputs = {
-        "hidden_states": Qwen3ASRThinkerTextDecoderLayer,
-        "attentions": Qwen3ASRThinkerTextAttention,
-    }
-
-    def __init__(self, config: Qwen3ASRConfig):
-        super().__init__(config)
-        self.padding_idx = config.pad_token_id
-        self.vocab_size = config.vocab_size
-
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
-        self.layers = nn.ModuleList(
-            [Qwen3ASRThinkerTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
-        )
-        self.norm = Qwen3ASRTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rotary_emb = Qwen3ASRThinkerTextRotaryEmbedding(config)
-        self.gradient_checkpointing = False
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    @check_model_inputs()
-    @auto_docstring
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Cache] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        **kwargs: Unpack[FlashAttentionKwargs],
-    ) -> Union[tuple, BaseModelOutputWithPast]:
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-
-        # torch.jit.trace() doesn't support cache objects in the output
-        if use_cache and past_key_values is None and not torch.jit.is_tracing():
-            past_key_values = DynamicCache(config=self.config)
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
-
-        if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(
-                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
-            )
-
-        # the hard coded `3` is for temporal, height and width.
-        if position_ids is None:
-            position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
-        elif position_ids.ndim == 2:
-            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
-
-        if position_ids.ndim == 3 and position_ids.shape[0] == 4:
-            text_position_ids = position_ids[0]
-            position_ids = position_ids[1:]
-        else:
-            text_position_ids = position_ids[0]
-
-        attention_mask = create_causal_mask(
-            config=self.config,
-            input_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            cache_position=cache_position,
-            past_key_values=past_key_values,
-            position_ids=text_position_ids,
-        )
-
-        hidden_states = inputs_embeds
-
-        # create position embeddings to be shared across the decoder layers
-        position_embeddings = self.rotary_emb(hidden_states, position_ids)
-
-        # decoder layers
-        for layer_idx, decoder_layer in enumerate(self.layers):
-            layer_outputs = decoder_layer(
-                hidden_states,
-                attention_mask=attention_mask,
-                position_ids=text_position_ids,
-                past_key_values=past_key_values,
-                cache_position=cache_position,
-                position_embeddings=position_embeddings,
-                **kwargs,
-            )
-            hidden_states = layer_outputs
-
-        hidden_states = self.norm(hidden_states)
-
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=past_key_values,
-        )
-
-
-@auto_docstring(
-    custom_intro="""
-    The Qwen3ASRThinker model which consists of a audio backbone and a language model.
-    """
-)
-class Qwen3ASRThinkerForConditionalGeneration(Qwen3ASRPreTrainedModelForConditionalGeneration, GenerationMixin):
-    config: Qwen3ASRThinkerConfig
-    base_model_prefix = "thinker"
-    _tied_weights_keys = ["model.embed_tokens.weight", "lm_head.weight"]
-    _no_split_modules = [
-        "Qwen3ASRAudioEncoderLayer",
-        "Qwen3ASRThinkerTextDecoderLayer",
-    ]
-    _can_record_outputs = {
-        "hidden_states": Qwen3ASRThinkerTextDecoderLayer,
-        "attentions": Qwen3ASRThinkerTextAttention,
-    }
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.audio_tower = Qwen3ASRAudioEncoder._from_config(config.audio_config)
-        self.vocab_size = config.text_config.vocab_size
-        self.model = Qwen3ASRThinkerTextModel._from_config(config.text_config)
-        if "forced_aligner" in config.model_type:
-            self.lm_head = nn.Linear(config.text_config.hidden_size, config.classify_num, bias=False)
-        else:
-            self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
-        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
-        self.rope_deltas = None
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.model.get_input_embeddings()
-
-    def set_input_embeddings(self, value):
-        self.model.set_input_embeddings(value)
-
-    def get_audio_features(
-        self,
-        input_features: torch.FloatTensor,
-        feature_attention_mask: Optional[torch.LongTensor] = None,
-        audio_feature_lengths: Optional[torch.LongTensor] = None,
-    ):
-        """
-        Encodes audios into continuous embeddings that can be forwarded to the language model.
-
-        Args:
-            input_features (`torch.FloatTensor`):
-                The tensors corresponding to the input audios.
-            feature_attention_mask (`torch.LongTensor`, *optional*):
-                Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`:
-            audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
-                The length of feature shape of each audio in LLM.
-        """
-        if feature_attention_mask is not None:
-            audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)
-            input_features = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
-        else:
-            audio_feature_lengths = None
-
-        feature_lens = audio_feature_lengths if audio_feature_lengths is not None else feature_attention_mask.sum(-1)
-        audio_outputs = self.audio_tower(
-            input_features,
-            feature_lens=feature_lens,
-        )
-        audio_features = audio_outputs.last_hidden_state
-
-        return audio_features
-
-    def get_placeholder_mask(
-        self,
-        input_ids: torch.LongTensor,
-        inputs_embeds: torch.FloatTensor,
-    ):
-        """
-        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
-        equal to the length of multimodal features. If the lengths are different, an error is raised.
-        """
-        if input_ids is None:
-            special_audio_mask = (
-                inputs_embeds
-                == self.get_input_embeddings()(
-                    torch.tensor(self.config.audio_token_id, dtype=torch.long, device=inputs_embeds.device)
-                )
-            ).all(-1)
-        else:
-            special_audio_mask = input_ids == self.config.audio_token_id
-
-        special_audio_mask = special_audio_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
-        return special_audio_mask
-
-    @can_return_tuple
-    @auto_docstring
-    def forward(
-        self,
-        input_ids=None,
-        input_features=None,
-        attention_mask=None,
-        feature_attention_mask=None,
-        audio_feature_lengths=None,
-        position_ids=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        rope_deltas=None,
-        labels=None,
-        use_cache=None,
-        cache_position=None,
-        **kwargs,
-    ) -> Union[tuple, Qwen3ASRThinkerCausalLMOutputWithPast]:
-        r"""
-        feature_attention_mask (`torch.Tensor` of shape `(batch_size, feature_sequence_length)`, *optional*):
-            Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`:
-            - 1 for tokens that are **not masked**,
-            - 0 for tokens that are **masked**.
-        audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*):
-            The length of feature shape of each audio in LLM.
-        rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
-            The rope index difference between sequence length and multimodal rope.
-        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
-            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
-            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-        """
-
-        if inputs_embeds is None:
-            # 1. Extract the input embeddings
-            inputs_embeds = self.get_input_embeddings()(input_ids)
-
-        # 2. Merge text, audios
-        if input_features is not None:
-            audio_features = self.get_audio_features(
-                input_features,
-                feature_attention_mask=feature_attention_mask,
-                audio_feature_lengths=audio_feature_lengths,
-            )
-            audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype)
-            audio_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds)
-            inputs_embeds = inputs_embeds.masked_scatter(audio_mask, audio_features)
-
-        if feature_attention_mask is not None:
-            audio_feature_lengths = torch.sum(feature_attention_mask, dim=1)
-        else:
-            audio_feature_lengths = None
-
-        if attention_mask is not None and position_ids is None:
-            if (
-                cache_position is None
-                or (cache_position is not None and cache_position[0] == 0)
-                or self.rope_deltas is None
-            ):
-                delta0 = (1 - attention_mask).sum(dim=-1).unsqueeze(1)
-                position_ids, rope_deltas = self.get_rope_index(
-                    attention_mask,
-                )
-                rope_deltas = rope_deltas - delta0
-                self.rope_deltas = rope_deltas
-            else:
-                batch_size, seq_length = input_ids.shape
-                delta = cache_position[0] + self.rope_deltas if cache_position is not None else 0
-                position_ids = torch.arange(seq_length, device=input_ids.device)
-                position_ids = position_ids.view(1, -1).expand(batch_size, -1)
-                position_ids = position_ids.add(delta)
-                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
-
-        outputs = self.model(
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            cache_position=cache_position,
-            **kwargs,
-        )
-
-        hidden_states = outputs[0]
-        logits = self.lm_head(hidden_states)
-
-        loss = None
-        if labels is not None:
-            loss = self.loss_function(
-                logits=logits, labels=labels, vocab_size=self.config.get_text_config().vocab_size
-            )
-
-        return Qwen3ASRThinkerCausalLMOutputWithPast(
-            loss=loss,
-            logits=logits,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-            past_key_values=outputs.past_key_values,
-            rope_deltas=self.rope_deltas,
-        )
-
-    def prepare_inputs_for_generation(
-        self,
-        input_ids,
-        past_key_values=None,
-        attention_mask=None,
-        inputs_embeds=None,
-        cache_position=None,
-        position_ids=None,
-        use_cache=True,
-        input_features=None,
-        feature_attention_mask=None,
-        **kwargs,
-    ):
-        model_inputs = super().prepare_inputs_for_generation(
-            input_ids,
-            past_key_values=past_key_values,
-            attention_mask=attention_mask,
-            inputs_embeds=inputs_embeds,
-            cache_position=cache_position,
-            position_ids=position_ids,
-            use_cache=use_cache,
-            input_features=input_features,
-            feature_attention_mask=feature_attention_mask,
-            **kwargs,
-        )
-
-        model_inputs["position_ids"] = None
-
-        if cache_position[0] != 0:
-            model_inputs["input_features"] = None
-
-        return model_inputs
-
-
-@auto_docstring
-class Qwen3ASRThinkerTextPreTrainedModel(PreTrainedModel):
-    config = Qwen3ASRConfig
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = True
-    _no_split_modules = ["Qwen3ASRThinkerTextDecoderLayer"]
-    _skip_keys_device_placement = ["past_key_values"]
-    _supports_flash_attn = True
-    _supports_sdpa = True
-    _supports_flex_attn = True
-    _can_compile_fullgraph = False  # MoE models don't work with torch.compile (`torch.where(condition)` not supported)
-    _supports_attention_backend = True
-    _can_record_outputs = {
-        "hidden_states": Qwen3ASRThinkerTextDecoderLayer,
-        "attentions": Qwen3ASRThinkerTextAttention,
-    }
-    config_class = Qwen3ASRConfig
-
-
-class Qwen3ASRForConditionalGeneration(Qwen3ASRPreTrainedModel, GenerationMixin):
-    config_class = Qwen3ASRConfig
-
-    def __init__(self, config: Qwen3ASRConfig):
-        super().__init__(config)
-        self.config = config
-
-        self.thinker = Qwen3ASRThinkerForConditionalGeneration._from_config(config.thinker_config)
-        self.post_init()
-    
-    def get_support_languages(self):
-        return self.config.support_languages
-
-    @torch.no_grad()
-    def generate(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        max_new_tokens: int = 8192,
-        eos_token_id: int | list[int] = [151645, 151643],
-        **kwargs,
-    ):
-        shared_kwargs = {}
-        thinker_kwargs = {
-            "max_new_tokens": max_new_tokens,
-            "eos_token_id": eos_token_id,
-        }
-
-        for key, value in kwargs.items():
-            # Process special input values
-            if key == "feature_attention_mask":
-                thinker_kwargs[key] = value
-            elif key in ("input_features", "attention_mask"):
-                thinker_kwargs[key] = value
-            # Put other key to shared kwargs
-            else:
-                shared_kwargs[key] = value
-
-        # Merge kwargs
-        for key, value in shared_kwargs.items():
-            if key not in thinker_kwargs:
-                thinker_kwargs[key] = value
-
-        thinker_result = self.thinker.generate(input_ids=input_ids, return_dict_in_generate=True, **thinker_kwargs)
-
-        return thinker_result
-
-
-__all__ = [
-    "Qwen3ASRForConditionalGeneration",
-    "Qwen3ASRThinkerTextModel",
-    "Qwen3ASRThinkerForConditionalGeneration",
-    "Qwen3ASRPreTrainedModel",
-    "Qwen3ASRPreTrainedModelForConditionalGeneration",
-    "Qwen3ASRThinkerTextPreTrainedModel",
-]

qwen_asr/core/transformers_backend/processing_qwen3_asr.py

@@ -1,209 +0,0 @@
-# coding=utf-8
-# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# 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 re
-
-import numpy as np
-
-from transformers.audio_utils import AudioInput
-from transformers.feature_extraction_utils import BatchFeature
-from transformers.processing_utils import ProcessingKwargs, ProcessorMixin
-from transformers.tokenization_utils_base import TextInput
-
-
-class Qwen3ASRProcessorKwargs(ProcessingKwargs, total=False):
-    _defaults = {
-        "text_kwargs": {
-            "padding": False,
-            "padding_side": "left",
-        },
-        "audio_kwargs": {
-            "sampling_rate": 16000,
-            "padding": True,
-            "return_attention_mask": True,
-        },
-    }
-
-
-def _get_feat_extract_output_lengths(input_lengths):
-    """
-    Computes the output length of the convolutional layers and the output length of the audio encoder
-    """
-
-    input_lengths_leave = input_lengths % 100
-    feat_lengths = (input_lengths_leave - 1) // 2 + 1
-    output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
-    return output_lengths
-
-
-class Qwen3ASRProcessor(ProcessorMixin):
-    r"""
-    Constructs a Qwen3ASR processor.
-    [`Qwen3ASRProcessor`] offers all the functionalities of [`WhisperFeatureExtractor`], and [`Qwen2TokenizerFast`]. See the
-    [`~Qwen3ASRProcessor.__call__`] and [`~Qwen3ASRProcessor.decode`] for more information.
-
-    Args:
-        feature_extractor ([`WhisperFeatureExtractor`], *optional*):
-            The audio feature extractor.
-        tokenizer ([`Qwen2TokenizerFast`], *optional*):
-            The text tokenizer.
-        chat_template (`Optional[str]`, *optional*):
-            The Jinja template to use for formatting the conversation. If not provided, the default chat template is used.
-    """
-
-    attributes = ["feature_extractor", "tokenizer"]
-    feature_extractor_class = "WhisperFeatureExtractor"
-    tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
-
-    def __init__(
-        self, feature_extractor=None, tokenizer=None, chat_template=None
-    ):
-        super().__init__(feature_extractor, tokenizer, chat_template=chat_template)
-        self.audio_token = self.tokenizer.audio_token
-        self.audio_bos_token = self.tokenizer.audio_bos_token
-        self.audio_eos_token = self.tokenizer.audio_eos_token
-
-    def __call__(
-        self,
-        text: TextInput = None,
-        audio: AudioInput = None,
-        **kwargs,
-    ) -> BatchFeature:
-        """
-        Main method to prepare for the model one or several sequences(s) and audio(s). This method forwards the `text`
-        and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
-        the text. To prepare the audio(s), this method forwards the `audio` and `kwargs` arguments to
-        WhisperFeatureExtractor's [`~WhisperFeatureExtractor.__call__`] if `audio` is not `None`. Please refer to the doctsring
-        of the above two methods for more information.
-
-        Args:
-            text (`str`, `List[str]`, `List[List[str]]`):
-                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
-                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
-                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
-            audio (`np.ndarray`, `List[np.ndarray]`):
-                The audio or batch of audio to be prepared. Each audio can be a NumPy array.
-        """
-
-        if text is None:
-            raise ValueError("You need to specify either a `text` input to process.")
-
-        output_kwargs = self._merge_kwargs(
-            Qwen3ASRProcessorKwargs,
-            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
-            **kwargs,
-        )
-
-        if audio is not None:
-            output_kwargs["audio_kwargs"]["padding"] = True
-            output_kwargs["audio_kwargs"]["truncation"] = False
-            audio_inputs = self.feature_extractor(audio, **output_kwargs["audio_kwargs"])
-            audio_inputs["feature_attention_mask"] = audio_inputs.pop(
-                "attention_mask"
-            )  # rename feature_attention_mask to prevent conflicts later on
-            audio_inputs["input_features"] = audio_inputs.pop(
-                "input_features"
-            )  # rename input_features to prevent conflicts later on
-            audio_lengths = iter(_get_feat_extract_output_lengths(audio_inputs["feature_attention_mask"].sum(-1)))
-        else:
-            audio_inputs = {}
-            audio_lengths = iter([])
-
-        if not isinstance(text, list):
-            text = [text]
-
-        text = self.replace_multimodal_special_tokens(
-            text,
-            audio_lengths,
-        )
-
-        texts_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
-
-        return BatchFeature(
-            data={**texts_inputs, **audio_inputs},
-            tensor_type=kwargs.get("return_tensors"),
-        )
-
-    def replace_multimodal_special_tokens(
-        self,
-        text,
-        audio_lengths,
-    ):
-
-        processed_text = []
-        for sample in text:
-            positions = []
-            special_tokens = [re.escape(tok) for tok in [self.audio_token]]
-            pattern = "|".join(special_tokens)
-            positions = sorted([(match.start(), match.group()) for match in re.finditer(pattern, sample)])
-            positions.sort(key=lambda x: x[0])
-
-            for _, special_token in positions:
-                if special_token == self.audio_token:
-                    sample = sample.replace(self.audio_token, "<|audio_placeholder|>" * next(audio_lengths), 1)
-
-            sample = sample.replace("<|audio_placeholder|>", self.audio_token)
-            processed_text.append(sample)
-        return processed_text
-
-    def get_chunked_index(self, token_indices: np.ndarray, tokens_per_chunk: int) -> list[tuple[int, int]]:
-        """
-        Splits token index list into chunks based on token value ranges.
-
-        Given a list of token indices, returns a list of (start, end) index tuples representing
-        slices of the list where the token values fall within successive ranges of `t_ntoken_per_chunk`.
-
-        For example, if `t_ntoken_per_chunk` is 1000, the function will create chunks such that:
-        - the first chunk contains token values < 1000,
-        - the second chunk contains values >= 1000 and < 2000, and so on.
-
-        Parameters:
-            token_indices (`np.ndarray`): A monotonically increasing list of token index values.
-            t_ntoken_per_chunk (`int`): Number of tokens per chunk (used as the chunk size threshold).
-
-        Returns:
-            `list[tuple[int, int]]`: A list of tuples, each representing the start (inclusive)
-                                and end (exclusive) indices of a chunk in `token_indices`.
-        """
-
-        def _iter():
-            i, start_idx = 0, 0  # skip bos token
-            current_chunk = 1
-            while i < len(token_indices):  # skip eos token
-                if token_indices[i] >= current_chunk * tokens_per_chunk:
-                    yield (start_idx, i)
-                    start_idx = i
-                    current_chunk += 1
-                i += 1
-            yield (start_idx, len(token_indices))
-
-        return list(_iter())
-
-    def apply_chat_template(self, conversations, chat_template=None, **kwargs):
-        return super().apply_chat_template(conversations, chat_template, **kwargs)
-
-    @property
-    def model_input_names(self):
-        tokenizer_input_names = self.tokenizer.model_input_names
-        feature_extractor_input_names = self.feature_extractor.model_input_names
-        return list(
-            dict.fromkeys(
-                tokenizer_input_names
-                + feature_extractor_input_names
-                + ["feature_attention_mask"]
-            )
-        )
-
-
-__all__ = ["Qwen3ASRProcessor"]

qwen_asr/core/vllm_backend/__init__.py

@@ -1,16 +0,0 @@
-# coding=utf-8
-# 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.
-from .qwen3_asr import Qwen3ASRForConditionalGeneration

qwen_asr/core/vllm_backend/qwen3_asr.py

@@ -1,997 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# Copyright 2026 The Qwen team.
-# Copyright 2023 The vLLM team.
-# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
-#
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
-# and OPT implementations in this library. It has been modified from its
-# original forms to accommodate minor architectural differences compared
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
-#
-# 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.
-"""Inference-only Qwen3-ASR model."""
-
-from collections.abc import Iterable, Mapping, Sequence
-from typing import Any, Literal, cast
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from transformers.feature_extraction_utils import BatchFeature
-from transformers.models.whisper import WhisperFeatureExtractor
-
-from vllm.config import MultiModalConfig, ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
-from vllm.distributed import get_tensor_model_parallel_world_size
-from vllm.inputs.data import PromptType
-from vllm.logger import init_logger
-from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
-from vllm.model_executor.layers.attention.mm_encoder_attention import (
-    MMEncoderAttention,
-)
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    QKVParallelLinear,
-    RowParallelLinear,
-)
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-from vllm.model_executor.models.interfaces import (
-    MultiModalEmbeddings,
-    SupportsMRoPE,
-    SupportsMultiModal,
-    SupportsPP,
-    SupportsTranscription,
-)
-from vllm.model_executor.models.module_mapping import MultiModelKeys
-from vllm.model_executor.models.qwen3 import Qwen3ForCausalLM
-from vllm.model_executor.models.qwen3_omni_moe_thinker import (
-    Qwen2_5OmniAudioFeatureInputs,
-    Qwen3OmniMoeThinkerMultiModalProcessor,
-)
-from vllm.model_executor.models.utils import (
-    AutoWeightsLoader,
-    WeightsMapper,
-    _merge_multimodal_embeddings,
-    maybe_prefix,
-)
-from vllm.model_executor.models.whisper import ISO639_1_SUPPORTED_LANGS
-from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.inputs import (
-    AudioItem,
-    ModalityData,
-    MultiModalDataDict,
-    MultiModalFeatureSpec,
-    MultiModalFieldConfig,
-    MultiModalKwargsItems,
-)
-from vllm.multimodal.parse import (
-    AudioProcessorItems,
-    DictEmbeddingItems,
-    ModalityDataItems,
-    MultiModalDataItems,
-    MultiModalDataParser,
-)
-from vllm.multimodal.processing import (
-    BaseProcessingInfo,
-    PromptReplacement,
-    PromptUpdate,
-)
-from vllm.sequence import IntermediateTensors
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.tokenizers import cached_tokenizer_from_config
-from vllm.transformers_utils.processor import cached_processor_from_config
-from vllm.model_executor.models.vision import (
-    get_vit_attn_backend,
-)
-from ..transformers_backend.configuration_qwen3_asr import (
-    Qwen3ASRConfig,
-    Qwen3ASRThinkerConfig,
-    Qwen3ASRAudioEncoderConfig
-)
-from ..transformers_backend.processing_qwen3_asr import (
-    Qwen3ASRProcessor,
-)
-
-try:
-    from vllm.multimodal.profiling import BaseDummyInputsBuilder
-except:
-    from vllm.multimodal.processing import BaseDummyInputsBuilder
-
-logger = init_logger(__name__)
-
-
-def _get_feat_extract_output_lengths(input_lengths: torch.Tensor):
-    input_lengths_leave = input_lengths % 100
-    feat_lengths = (input_lengths_leave - 1) // 2 + 1
-    output_lengths = (
-        ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
-    )
-    return output_lengths
-
-
-# ============= Audio Encoder Components =============
-
-
-class SinusoidsPositionEmbedding(nn.Module):
-    """Sinusoidal position embedding for audio encoder."""
-
-    def __init__(self, length: int, channels: int, max_timescale: int = 10000):
-        super().__init__()
-        self.length = length
-        self.channels = channels
-        self.max_timescale = max_timescale
-
-        if channels % 2 != 0:
-            raise ValueError("SinusoidsPositionEmbedding needs even channels input")
-
-        log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
-        inv_timescales = torch.exp(
-            -log_timescale_increment * torch.arange(channels // 2).float()
-        )
-        scaled_time = (
-            torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
-        )
-        positional_embedding = torch.cat(
-            [torch.sin(scaled_time), torch.cos(scaled_time)], dim=1
-        )
-        self.register_buffer(
-            "positional_embedding", positional_embedding, persistent=False
-        )
-
-    def forward(self, seqlen: int) -> torch.Tensor:
-        return self.positional_embedding[:seqlen, :]
-
-
-class Qwen3ASRAudioAttention(nn.Module):
-    """Multi-headed attention for Qwen3-Omni Audio Encoder using MMEncoderAttention."""
-
-    def __init__(
-        self,
-        config: Qwen3ASRAudioEncoderConfig,
-        multimodal_config: MultiModalConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.embed_dim = config.d_model
-        self.num_heads = config.encoder_attention_heads
-        self.head_dim = self.embed_dim // self.num_heads
-        tp_size = get_tensor_model_parallel_world_size()
-        self.num_local_heads = self.num_heads // tp_size
-
-        if (self.head_dim * self.num_heads) != self.embed_dim:
-            raise ValueError(
-                f"embed_dim must be divisible by num_heads (got `embed_dim`: "
-                f"{self.embed_dim} and `num_heads`: {self.num_heads})."
-            )
-
-        self.scaling = self.head_dim**-0.5
-
-        self.qkv = QKVParallelLinear(
-            hidden_size=self.embed_dim,
-            head_size=self.head_dim,
-            total_num_heads=self.num_heads,
-            total_num_kv_heads=self.num_heads,
-            bias=True,
-            prefix=f"{prefix}.qkv",
-        )
-
-        self.out_proj = RowParallelLinear(
-            input_size=self.embed_dim,
-            output_size=self.embed_dim,
-            bias=True,
-            prefix=f"{prefix}.out_proj",
-        )
-
-        self.attn = MMEncoderAttention(
-            num_heads=self.num_local_heads,
-            head_size=self.head_dim,
-            scale=self.scaling,
-            multimodal_config=multimodal_config,
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        cu_seqlens: torch.Tensor,
-        max_seqlen: torch.Tensor | None,
-    ) -> torch.Tensor:
-        seq_length, _ = hidden_states.size()
-        qkv, _ = self.qkv(hidden_states)
-        q, k, v = qkv.chunk(3, dim=-1)
-        q = q.view(1, seq_length, -1, self.head_dim)
-        k = k.view(1, seq_length, -1, self.head_dim)
-        v = v.view(1, seq_length, -1, self.head_dim)
-
-        attn_output = self.attn(
-            query=q,
-            key=k,
-            value=v,
-            cu_seqlens=cu_seqlens,
-            max_seqlen=max_seqlen,
-        )
-
-        attn_output = attn_output.view(seq_length, -1)
-        output, _ = self.out_proj(attn_output)
-        return output
-
-
-class Qwen3ASRAudioEncoderLayer(nn.Module):
-    """Transformer encoder layer for Qwen3-Omni Audio Encoder."""
-
-    def __init__(
-        self,
-        config: Qwen3ASRAudioEncoderConfig,
-        multimodal_config: MultiModalConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.embed_dim = config.d_model
-        self.self_attn = Qwen3ASRAudioAttention(
-            config, multimodal_config=multimodal_config, prefix=f"{prefix}.self_attn"
-        )
-        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
-        self.activation_fn = _ACTIVATION_REGISTRY[config.activation_function]
-        self.fc1 = ColumnParallelLinear(
-            self.embed_dim,
-            config.encoder_ffn_dim,
-            bias=True,
-            prefix=f"{prefix}.fc1",
-        )
-        self.fc2 = RowParallelLinear(
-            config.encoder_ffn_dim,
-            self.embed_dim,
-            bias=True,
-            prefix=f"{prefix}.fc2",
-        )
-        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        cu_seqlens: torch.Tensor,
-        max_seqlen: torch.Tensor | None,
-    ) -> torch.Tensor:
-        """
-        Args:
-            hidden_states: Input tensor of shape (seq_len, hidden_size)
-            cu_seqlens: Cumulative sequence lengths
-            max_seqlen: Maximum sequence length in the batch
-        """
-        residual = hidden_states
-        hidden_states = self.self_attn_layer_norm(hidden_states)
-        hidden_states = self.self_attn(
-            hidden_states=hidden_states,
-            cu_seqlens=cu_seqlens,
-            max_seqlen=max_seqlen,
-        )
-        hidden_states = residual + hidden_states
-
-        residual = hidden_states
-        hidden_states = self.final_layer_norm(hidden_states)
-        hidden_states, _ = self.fc1(hidden_states)
-        hidden_states = self.activation_fn(hidden_states)
-        hidden_states, _ = self.fc2(hidden_states)
-        hidden_states = residual + hidden_states
-
-        # Clamp for numerical stability with fp16
-        if hidden_states.dtype == torch.float16:
-            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
-            hidden_states = torch.clamp(
-                hidden_states, min=-clamp_value, max=clamp_value
-            )
-
-        return hidden_states
-
-
-class Qwen3ASRAudioEncoder(nn.Module):
-    """vLLM-native Qwen3-ASR Audio Encoder."""
-
-    def __init__(
-        self,
-        config: Qwen3ASRAudioEncoderConfig,
-        multimodal_config: MultiModalConfig | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-
-        embed_dim = config.d_model
-        self.num_mel_bins = config.num_mel_bins
-        self.max_source_positions = config.max_source_positions
-        self.n_window = config.n_window
-        self.n_window_infer = config.n_window_infer
-        self.conv_chunksize = config.conv_chunksize
-
-        # Position embedding
-        self.positional_embedding = SinusoidsPositionEmbedding(
-            self.max_source_positions, embed_dim
-        )
-
-        # Convolutional layers for mel-spectrogram processing
-        self.conv2d1 = nn.Conv2d(1, config.downsample_hidden_size, 3, 2, padding=1)
-        self.conv2d2 = nn.Conv2d(
-            config.downsample_hidden_size,
-            config.downsample_hidden_size,
-            3,
-            2,
-            padding=1,
-        )
-        self.conv2d3 = nn.Conv2d(
-            config.downsample_hidden_size,
-            config.downsample_hidden_size,
-            3,
-            2,
-            padding=1,
-        )
-
-        conv_out_dim = config.downsample_hidden_size * (
-            (((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2
-        )
-        self.conv_out = nn.Linear(conv_out_dim, config.d_model, bias=False)
-
-        # Transformer encoder layers
-        self.layers = nn.ModuleList(
-            [
-                Qwen3ASRAudioEncoderLayer(
-                    config,
-                    multimodal_config=multimodal_config,
-                    prefix=f"{prefix}.layers.{i}",
-                )
-                for i in range(config.encoder_layers)
-            ]
-        )
-
-        # Output layers
-        self.ln_post = nn.LayerNorm(config.d_model)
-        self.proj1 = nn.Linear(config.d_model, config.d_model)
-        self.act = _ACTIVATION_REGISTRY[config.activation_function]
-        self.proj2 = nn.Linear(config.d_model, config.output_dim)
-
-        # Get attention backend
-        attn_backend_override = (
-            multimodal_config.mm_encoder_attn_backend
-            if multimodal_config is not None
-            else None
-        )
-        self.attn_backend = get_vit_attn_backend(
-            head_size=config.d_model // config.encoder_attention_heads,
-            dtype=torch.get_default_dtype(),
-            attn_backend_override=attn_backend_override,
-        )
-
-    def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> torch.Tensor | None:
-        """Compute max_seqlen only for flash attention backends."""
-        max_seqlen = None
-        if self.attn_backend in {
-            AttentionBackendEnum.FLASH_ATTN,
-            AttentionBackendEnum.ROCM_AITER_FA,
-        }:
-            max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
-        return max_seqlen
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return self.conv2d1.weight.dtype
-
-    @property
-    def device(self) -> torch.device:
-        return self.conv2d1.weight.device
-
-    def forward(
-        self,
-        input_features: torch.Tensor,
-        feature_lens: torch.Tensor,
-        aftercnn_lens: torch.Tensor,
-    ):
-        # Compute chunk information
-        chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()
-
-        chunk_lengths = torch.tensor(
-            [self.n_window * 2] * chunk_num.sum(),
-            dtype=torch.long,
-            device=feature_lens.device,
-        )
-        tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
-        chunk_lengths[tail_chunk_index] = feature_lens % (self.n_window * 2)
-        chunk_lengths[chunk_lengths == 0] = self.n_window * 2
-
-        # Split input features into chunks and pad
-        chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)
-        padded_feature = nn.utils.rnn.pad_sequence(
-            chunk_list, batch_first=True
-        ).transpose(1, 2)
-
-        # Compute feature lengths after CNN
-        feature_lens_after_cnn = self._get_cnn_output_lengths(chunk_lengths)
-        # Vectorized mask creation: avoid creating many small tensors
-        max_len_after_cnn = feature_lens_after_cnn.max().item()
-        indices = torch.arange(max_len_after_cnn, device=padded_feature.device)
-        padded_mask_after_cnn = indices.unsqueeze(0) < feature_lens_after_cnn.unsqueeze(
-            1
-        )
-
-        # Add channel dimension for conv2d
-        padded_feature = padded_feature.unsqueeze(1)
-
-        # Apply convolutional layers (chunk if needed to avoid OOM)
-        if padded_feature.size(0) <= self.conv_chunksize:
-            # Fast path: no chunking needed
-            padded_embed = F.gelu(self.conv2d1(padded_feature))
-            padded_embed = F.gelu(self.conv2d2(padded_embed))
-            padded_embed = F.gelu(self.conv2d3(padded_embed))
-        else:
-            # Chunked processing to avoid OOM
-            padded_embeds = []
-            for chunk in padded_feature.split(self.conv_chunksize, dim=0):
-                padded_embed = F.gelu(self.conv2d1(chunk))
-                padded_embed = F.gelu(self.conv2d2(padded_embed))
-                padded_embed = F.gelu(self.conv2d3(padded_embed))
-                padded_embeds.append(padded_embed)
-            padded_embed = torch.cat(padded_embeds, dim=0)
-
-        # (batch, channels, freq, time) -> (batch, time, channels*freq)
-        b, c, f, t = padded_embed.size()
-        padded_embed = self.conv_out(
-            padded_embed.permute(0, 3, 1, 2).contiguous().view(b, t, c * f)
-        )
-
-        # Add positional embedding
-        positional_embedding = (
-            self.positional_embedding.positional_embedding[: padded_embed.shape[1], :]
-            .unsqueeze(0)
-            .to(padded_embed.dtype)
-        )
-        padded_embed = padded_embed + positional_embedding
-
-        # Extract valid hidden states and compute cu_seqlens
-        hidden_states = padded_embed[padded_mask_after_cnn]
-
-        # Compute cumulative sequence lengths for chunked attention
-        cu_chunk_lens = [0]
-        window_aftercnn = padded_mask_after_cnn.shape[-1] * (
-            self.n_window_infer // (self.n_window * 2)
-        )
-        # Use tolist() for efficient batch conversion from tensor to Python
-        for cnn_len in aftercnn_lens.tolist():
-            num_full_chunks = cnn_len // window_aftercnn
-            remainder = cnn_len % window_aftercnn
-            cu_chunk_lens.extend([window_aftercnn] * num_full_chunks)
-            if remainder:
-                cu_chunk_lens.append(remainder)
-        cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum(
-            -1, dtype=torch.int32
-        )
-
-        max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
-
-        # Apply transformer layers
-        for encoder_layer in self.layers:
-            hidden_states = encoder_layer(
-                hidden_states,
-                cu_seqlens,
-                max_seqlen,
-            )
-
-        # Apply output layers
-        hidden_states = self.ln_post(hidden_states)
-        hidden_states = self.proj1(hidden_states)
-        hidden_states = self.act(hidden_states)
-        hidden_states = self.proj2(hidden_states)
-
-        return hidden_states
-
-    def _get_cnn_output_lengths(self, input_lengths: torch.Tensor) -> torch.Tensor:
-        """Compute output lengths after the three conv2d layers."""
-        lengths = input_lengths
-        for _ in range(3):
-            lengths = (lengths - 1) // 2 + 1
-        return lengths
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        """Load weights with mapping from HuggingFace format."""
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("self_attn.qkv.", "self_attn.q_proj.", "q"),
-            ("self_attn.qkv.", "self_attn.k_proj.", "k"),
-            ("self_attn.qkv.", "self_attn.v_proj.", "v"),
-        ]
-        params_dict = dict(self.named_parameters(remove_duplicate=False))
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                param = params_dict.get(name)
-                if param is not None:
-                    weight_loader = getattr(
-                        param, "weight_loader", default_weight_loader
-                    )
-                    weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
-
-
-class Qwen3ASRProcessingInfo(BaseProcessingInfo):
-    def get_hf_config(self):
-        return self.ctx.get_hf_config(Qwen3ASRConfig).thinker_config
-
-    def get_hf_processor(self, **kwargs: object) -> Qwen3ASRProcessor:
-        processor = self.ctx.get_hf_processor(
-            Qwen3ASRProcessor,
-            use_fast=kwargs.pop("use_fast", True),
-            **kwargs,
-        )
-        if not hasattr(processor, "audio_token"):
-            processor.audio_token = "<|audio_pad|>"
-        return processor
-
-    def get_feature_extractor(self, **kwargs: object) -> WhisperFeatureExtractor:
-        hf_processor = self.get_hf_processor(**kwargs)
-        feature_extractor = hf_processor.feature_extractor
-        assert isinstance(feature_extractor, WhisperFeatureExtractor)
-        return feature_extractor
-
-    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
-        return {"audio": None}
-
-
-class Qwen3ASRDummyInputsBuilder(BaseDummyInputsBuilder[Qwen3ASRProcessingInfo]):
-    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
-        num_audios = mm_counts.get("audio", 0)
-
-        hf_processor = self.info.get_hf_processor()
-        audio_token = hf_processor.audio_token
-
-        return audio_token * num_audios
-
-    def get_dummy_mm_data(
-        self,
-        seq_len: int,
-        mm_counts: Mapping[str, int],
-        mm_options: Mapping[str, BaseDummyOptions] | None = None,
-    ) -> MultiModalDataDict:
-        num_audios = mm_counts.get("audio", 0)
-
-        feature_extractor = self.info.get_feature_extractor()
-
-        target_audio_length = (
-            min(
-                feature_extractor.chunk_length,
-                30,
-            )
-            * feature_extractor.sampling_rate
-        )
-
-        audio_overrides = mm_options.get("audio") if mm_options else None
-
-        return {
-            "audio": self._get_dummy_audios(
-                length=target_audio_length,
-                num_audios=num_audios,
-                overrides=audio_overrides,
-            ),
-        }
-
-
-def _qwen3asr_field_config(hf_inputs: Mapping[str, torch.Tensor]):
-    audio_feature_lengths = hf_inputs.get("audio_feature_lengths", torch.empty((0,)))
-    return dict(
-        input_audio_features=MultiModalFieldConfig.flat_from_sizes(
-            "audio", audio_feature_lengths, dim=1
-        ),
-        feature_attention_mask=MultiModalFieldConfig.batched("audio"),
-        audio_feature_lengths=MultiModalFieldConfig.batched("audio"),
-    )
-
-
-class Qwen3ASRMultiModalDataParser(MultiModalDataParser):
-    def _parse_audio_data(
-        self,
-        data: dict[str, torch.Tensor] | ModalityData[AudioItem],
-    ) -> ModalityDataItems[Any, Any] | None:
-        if isinstance(data, dict):
-            return DictEmbeddingItems(
-                data,
-                modality="audio",
-                required_fields={"input_audio_features", "audio_feature_lengths"},
-                fields_factory=_qwen3asr_field_config,
-            )
-
-        return super()._parse_audio_data(data)
-
-
-class Qwen3ASRMultiModalProcessor(
-    Qwen3OmniMoeThinkerMultiModalProcessor,
-):
-    def _get_data_parser(self) -> MultiModalDataParser:
-        feature_extractor = self.info.get_feature_extractor()
-        return Qwen3ASRMultiModalDataParser(
-            target_sr=feature_extractor.sampling_rate,
-        )
-
-    def _get_mm_fields_config(
-        self,
-        hf_inputs: BatchFeature,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> Mapping[str, MultiModalFieldConfig]:
-        return _qwen3asr_field_config(hf_inputs)
-
-    def _get_prompt_updates(
-        self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, Any],
-        out_mm_kwargs: MultiModalKwargsItems,
-    ) -> Sequence[PromptUpdate]:
-        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
-        tokenizer = self.info.get_tokenizer()
-        vocab = tokenizer.get_vocab()
-
-        audio_token = processor.audio_token
-        audio_token_id = vocab[audio_token]
-
-        out_mm_data = out_mm_kwargs.get_data()
-        audio_feature_lengths = out_mm_data.get("audio_feature_lengths")
-        feature_attention_mask = out_mm_data.get("feature_attention_mask")
-        if audio_feature_lengths is None and feature_attention_mask is None:
-            audio_output_lengths = []
-        elif audio_feature_lengths is not None:
-            audio_output_lens = _get_feat_extract_output_lengths(audio_feature_lengths)
-            audio_output_lengths = audio_output_lens.tolist()
-        elif feature_attention_mask is not None:
-            assert isinstance(feature_attention_mask, torch.Tensor)
-            audio_output_lens = _get_feat_extract_output_lengths(
-                feature_attention_mask.sum(-1)
-            )
-            audio_output_lengths = audio_output_lens.tolist()
-
-        def get_replacement_qwen2_audio(item_idx: int):
-            num_features = audio_output_lengths[item_idx]
-            if num_features == 0:
-                audios = mm_items.get_items("audio", AudioProcessorItems)
-                audio = audios.get(item_idx)
-                raise ValueError(
-                    f"The audio {audio} (len={len(audio)}) is too short "
-                    "to be represented inside the model"
-                )
-
-            return [audio_token_id] * num_features
-
-        return [
-            PromptReplacement(
-                modality="audio",
-                target=audio_token,
-                replacement=get_replacement_qwen2_audio,
-            ),
-        ]
-
-
-@MULTIMODAL_REGISTRY.register_processor(
-    Qwen3ASRMultiModalProcessor,
-    info=Qwen3ASRProcessingInfo,
-    dummy_inputs=Qwen3ASRDummyInputsBuilder,
-)
-class Qwen3ASRForConditionalGeneration(
-    nn.Module,
-    SupportsMultiModal,
-    SupportsPP,
-    SupportsMRoPE,
-    SupportsTranscription,
-):
-    supported_languages = ISO639_1_SUPPORTED_LANGS
-
-    hf_to_vllm_mapper = WeightsMapper(
-        orig_to_new_prefix={
-            "thinker.lm_head.": "language_model.lm_head.",
-            "thinker.model.": "language_model.model.",
-            "thinker.": "",
-        }
-    )
-
-    @classmethod
-    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
-        if modality.startswith("audio"):
-            return "<|audio_start|><|audio_pad|><|audio_end|>"
-
-        raise ValueError("Only audio modality is supported")
-
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        super().__init__()
-        self.vllm_config = vllm_config  # needed for torch compile forward context
-        thinker_config: Qwen3ASRThinkerConfig = (
-            vllm_config.model_config.hf_config.thinker_config
-        )
-        quant_config = vllm_config.quant_config
-        multimodal_config = vllm_config.model_config.multimodal_config
-        self.config = thinker_config
-        self.multimodal_config = multimodal_config
-
-        self.audio_tower = Qwen3ASRAudioEncoder(
-            thinker_config.audio_config,
-            multimodal_config=multimodal_config,
-            prefix=maybe_prefix(prefix, "audio_tower"),
-        )
-        self.quant_config = quant_config
-
-        self.language_model = Qwen3ForCausalLM(
-            vllm_config=vllm_config.with_hf_config(
-                thinker_config.text_config, architectures=["Qwen3ForCausalLM"]
-            ),
-            prefix=maybe_prefix(prefix, "language_model"),
-        )
-
-        self.make_empty_intermediate_tensors = (
-            self.language_model.make_empty_intermediate_tensors
-        )
-
-    def _parse_and_validate_audio_input(
-        self, **kwargs: object
-    ) -> Qwen2_5OmniAudioFeatureInputs | None:
-        input_audio_features = kwargs.pop("input_audio_features", None)
-        audio_feature_lengths = kwargs.pop("audio_feature_lengths", None)
-        feature_attention_mask = kwargs.pop("feature_attention_mask", None)
-        if input_audio_features is None:
-            return None
-
-        return Qwen2_5OmniAudioFeatureInputs(
-            type="audio_features",
-            input_features=input_audio_features,
-            audio_feature_lengths=audio_feature_lengths,
-            feature_attention_mask=feature_attention_mask,
-        )
-
-    def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
-        mm_input_by_modality = {}
-
-        # Preserve the order of modalities if there are multiple of them
-        # from the order of kwargs.
-        for input_key in kwargs:
-            if (
-                input_key in ("input_audio_features")
-                and "audio" not in mm_input_by_modality
-            ):
-                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
-                    **kwargs
-                )
-        return mm_input_by_modality
-
-    def _process_audio_input(
-        self,
-        audio_input: Qwen2_5OmniAudioFeatureInputs,
-        audio_hashes: list[str] | None = None,
-        cached_audio_features: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        input_features = audio_input["input_features"]
-        audio_feature_lengths = audio_input["audio_feature_lengths"]
-
-        audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths)
-
-        audio_features = self.audio_tower(
-            input_features.to(self.audio_tower.dtype),
-            feature_lens=audio_feature_lengths,
-            aftercnn_lens=audio_output_lengths,
-        )
-        return audio_features.split(audio_output_lengths.tolist())
-
-    def get_language_model(self) -> torch.nn.Module:
-        return self.language_model
-
-    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:
-        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
-        if not mm_input_by_modality:
-            return []
-
-        # The result multimodal_embeddings is tuple of tensors, with each
-        # tensor correspoending to a multimodal data item (image or video).
-        multimodal_embeddings: tuple[torch.Tensor, ...] = ()
-
-        # NOTE: It is important to iterate over the keys in this dictionary
-        # to preserve the order of the modalities.
-        for modality in mm_input_by_modality:
-            multimodal_input = mm_input_by_modality[modality]
-            if modality == "audio":
-                audio_embeddings = self._process_audio_input(multimodal_input)
-                multimodal_embeddings += tuple(audio_embeddings)
-        return multimodal_embeddings
-
-    def embed_input_ids(
-        self,
-        input_ids: torch.Tensor,
-        multimodal_embeddings: MultiModalEmbeddings | None = None,
-        *,
-        is_multimodal: torch.Tensor | None = None,
-        handle_oov_mm_token: bool = False,
-    ) -> torch.Tensor:
-        inputs_embeds = self._embed_text_input_ids(
-            input_ids,
-            self.language_model.embed_input_ids,
-            is_multimodal=is_multimodal,
-            handle_oov_mm_token=handle_oov_mm_token,
-        )
-
-        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
-            return inputs_embeds
-
-        inputs_embeds = _merge_multimodal_embeddings(
-            inputs_embeds=inputs_embeds,
-            multimodal_embeddings=multimodal_embeddings,
-            is_multimodal=is_multimodal,
-        )
-
-        return inputs_embeds
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        intermediate_tensors: IntermediateTensors | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-        **kwargs: object,
-    ) -> torch.Tensor | IntermediateTensors:
-        if intermediate_tensors is not None:
-            inputs_embeds = None
-
-        hidden_states = self.language_model.model(
-            input_ids,
-            positions,
-            intermediate_tensors,
-            inputs_embeds=inputs_embeds,
-        )
-
-        return hidden_states
-
-    def compute_logits(
-        self,
-        hidden_states: torch.Tensor,
-    ) -> torch.Tensor | None:
-        return self.language_model.compute_logits(hidden_states)
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(
-            self,
-            skip_prefixes=["talker.", "code2wav."],
-        )
-        loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
-
-        return loaded_weights
-
-    def get_mrope_input_positions(
-        self,
-        input_tokens: list[int],
-        mm_features: list[MultiModalFeatureSpec],
-    ) -> tuple[torch.Tensor, int]:
-        seq_len = len(input_tokens)
-
-        if not mm_features:
-            # No audio features, just return linear positions
-            llm_positions = (
-                torch.arange(seq_len, dtype=torch.long).view(1, -1).expand(3, -1)
-            )
-            return llm_positions.clone(), 0
-
-        llm_pos_ids_list: list[torch.Tensor] = []
-        st = 0
-
-        for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
-            offset = mm_feature.mm_position.offset
-
-            # Get audio feature length from mm_feature data
-            audio_feature_length = mm_feature.data["audio_feature_lengths"].data
-            if isinstance(audio_feature_length, torch.Tensor):
-                audio_feature_length = audio_feature_length.item()
-            audio_len = _get_feat_extract_output_lengths(
-                torch.tensor(audio_feature_length)
-            ).item()
-
-            # Text segment before audio (includes audio_start token)
-            text_len = offset - st
-            st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-            text_positions = (
-                torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1)
-                + st_idx
-            )
-            llm_pos_ids_list.append(text_positions)
-            st_idx = st_idx + text_len
-
-            # Audio token segment
-            audio_positions = (
-                torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1)
-                + st_idx
-            )
-            llm_pos_ids_list.append(audio_positions)
-
-            st = offset + audio_len
-
-        # Handle remaining text (includes audio_end and any trailing text)
-        if st < seq_len:
-            st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
-            text_len = seq_len - st
-            final_text_positions = (
-                torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1)
-                + st_idx
-            )
-            llm_pos_ids_list.append(final_text_positions)
-
-        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
-        if llm_positions.shape[1] != seq_len:
-            raise RuntimeError("Position ids length mismatch with input ids length")
-
-        mrope_position_delta = (llm_positions.max() + 1 - seq_len).item()
-        return llm_positions, mrope_position_delta
-
-    def get_mm_mapping(self) -> MultiModelKeys:
-        """
-        Get the module prefix in multimodal models
-        """
-        return MultiModelKeys.from_string_field(
-            language_model="language_model",
-            tower_model=["audio_tower."],
-        )
-
-    @classmethod
-    def get_speech_to_text_config(
-        cls, model_config: ModelConfig, task_type: str
-    ) -> SpeechToTextConfig:
-        processor = cached_processor_from_config(model_config)
-        feature_extractor: WhisperFeatureExtractor = processor.feature_extractor
-        return SpeechToTextConfig(
-            max_audio_clip_s=feature_extractor.chunk_length,
-            sample_rate=feature_extractor.sampling_rate,
-        )
-
-    @classmethod
-    def get_generation_prompt(
-        cls,
-        audio: np.ndarray,
-        model_config: ModelConfig,
-        stt_config: SpeechToTextConfig,
-        language: str | None,
-        task_type: Literal["transcribe", "translate"],
-        request_prompt: str,
-        to_language: str | None,
-    ) -> PromptType:
-        """Get the generation prompt to be used for transcription requests."""
-        tokenizer = cached_tokenizer_from_config(model_config)
-        audio_placeholder = cls.get_placeholder_str("audio", 0)
-
-        if task_type not in ("transcribe", "translate"):
-            raise ValueError(
-                f"Unsupported task_type '{task_type}'. "
-                "Supported task types are 'transcribe' and 'translate'."
-            )
-        full_lang_name_to = cls.supported_languages.get(to_language, to_language)
-        if to_language is None:
-            prompt = (
-                f"<|im_start|>user\n{audio_placeholder}<|im_end|>\n"
-                f"<|im_start|>assistant\n"
-            )
-        else:
-            prompt = (
-                f"<|im_start|>user\n{audio_placeholder}<|im_end|>\n"
-                f"<|im_start|>assistant\nlanguage {full_lang_name_to}<asr_text>"
-            )
-
-        prompt_token_ids = tokenizer.encode(prompt)
-        prompt_dict = {
-            "prompt_token_ids": prompt_token_ids,
-            "multi_modal_data": {"audio": audio},
-        }
-        return cast(PromptType, prompt_dict)

qwen_asr/inference/qwen3_asr.py

@@ -1,519 +0,0 @@
-# coding=utf-8
-# 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.
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
-
-import numpy as np
-import torch
-from qwen_asr.core.transformers_backend import (
-    Qwen3ASRConfig,
-    Qwen3ASRForConditionalGeneration,
-    Qwen3ASRProcessor,
-)
-from transformers import AutoConfig, AutoModel, AutoProcessor
-
-AutoConfig.register("qwen3_asr", Qwen3ASRConfig)
-AutoModel.register(Qwen3ASRConfig, Qwen3ASRForConditionalGeneration)
-AutoProcessor.register(Qwen3ASRConfig, Qwen3ASRProcessor)
-
-from .qwen3_forced_aligner import Qwen3ForcedAligner
-from .utils import (
-    MAX_ASR_INPUT_SECONDS,
-    MAX_FORCE_ALIGN_INPUT_SECONDS,
-    SAMPLE_RATE,
-    SUPPORTED_LANGUAGES,
-    AudioChunk,
-    AudioLike,
-    chunk_list,
-    merge_languages,
-    normalize_audios,
-    normalize_language_name,
-    parse_asr_output,
-    split_audio_into_chunks,
-    validate_language,
-)
-
-try:
-    from qwen_asr.core.vllm_backend import Qwen3ASRForConditionalGeneration
-    from vllm import ModelRegistry
-    ModelRegistry.register_model("Qwen3ASRForConditionalGeneration", Qwen3ASRForConditionalGeneration)
-except:
-    pass
-
-
-@dataclass
-class ASRTranscription:
-    """
-    One transcription result.
-
-    Attributes:
-        language (str):
-            Merged language string for the sample, e.g. "Chinese" or "Chinese,English".
-            Empty string if unknown or silent audio.
-        text (str):
-            Transcribed text.
-        time_stamps (Optional[Any]):
-            Forced aligner output (ForcedAlignResult).
-            Present only when return_time_stamps=True.
-    """
-    language: str
-    text: str
-    time_stamps: Optional[Any] = None
-
-
-class Qwen3ASRModel:
-    """
-    Unified inference wrapper for Qwen3-ASR with two backends:
-      - Transformers backend 
-      - vLLM backend
-
-    It optionally supports time stamp output via Qwen3-ForcedAligner.
-
-    Notes:
-      - Each request uses a context text and exactly one audio.
-      - If language is provided, the prompt will force the output to be text-only by appending
-        "language {Language}<asr_text>" to the assistant prompt.
-    """
-
-    def __init__(
-        self,
-        backend: str,
-        model: Any,
-        processor: Any,
-        sampling_params: Optional[Any] = None,
-        forced_aligner: Optional[Qwen3ForcedAligner] = None,
-        max_inference_batch_size: int = -1,
-    ):
-        self.backend = backend  # "transformers" | "vllm"
-        self.model = model
-        self.processor = processor
-        self.sampling_params = sampling_params
-        self.forced_aligner = forced_aligner
-        self.max_inference_batch_size = int(max_inference_batch_size)
-
-        if backend == "transformers":
-            self.device = getattr(model, "device", None)
-            if self.device is None:
-                try:
-                    self.device = next(model.parameters()).device
-                except StopIteration:
-                    self.device = torch.device("cpu")
-            self.dtype = getattr(model, "dtype", torch.float32)
-        else:
-            self.device = None
-            self.dtype = None
-
-    @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path: str,
-        forced_aligner: Optional[str] = None,
-        forced_aligner_kwargs: Optional[Dict[str, Any]] = None,
-        max_inference_batch_size: int = -1,
-        **kwargs,
-    ) -> "Qwen3ASRModel":
-        """
-        Initialize using Transformers backend.
-
-        Args:
-            pretrained_model_name_or_path:
-                HuggingFace repo id or local directory.
-            forced_aligner:
-                Optional forced aligner model path/repo id.
-            forced_aligner_kwargs:
-                Optional kwargs forwarded to Qwen3ForcedAligner.from_pretrained(...).
-            max_inference_batch_size:
-                Batch size limit for inference. -1 means no chunking. Small values can avoid OOM.
-            **kwargs:
-                Forwarded to AutoModel.from_pretrained(...).
-
-        Returns:
-            Qwen3ASRModel
-        """
-
-        model = AutoModel.from_pretrained(pretrained_model_name_or_path, **kwargs)
-
-        processor = AutoProcessor.from_pretrained(pretrained_model_name_or_path, fix_mistral_regex=True)
-
-        if forced_aligner is not None:
-            forced_aligner_model = Qwen3ForcedAligner.from_pretrained(
-                forced_aligner, **(forced_aligner_kwargs or {})
-            )
-
-        return cls(
-            backend="transformers",
-            model=model,
-            processor=processor,
-            sampling_params=None,
-            forced_aligner=forced_aligner_model,
-            max_inference_batch_size=max_inference_batch_size,
-        )
-
-    @classmethod
-    def LLM(
-        cls,
-        model: str,
-        forced_aligner: Optional[str] = None,
-        forced_aligner_kwargs: Optional[Dict[str, Any]] = None,
-        max_inference_batch_size: int = -1,
-        max_new_tokens: Optional[int] = 8192,
-        **kwargs,
-    ) -> "Qwen3ASRModel":
-        """
-        Initialize using vLLM backend.
-
-        Import is isolated to keep vLLM optional.
-
-        Args:
-            model:
-                Model path/repo for vLLM.
-            forced_aligner:
-                Optional forced aligner model path/repo id.
-            forced_aligner_kwargs:
-                Optional kwargs forwarded to Qwen3ForcedAligner.from_pretrained(...).
-            max_inference_batch_size:
-                Batch size limit for inference. -1 means no chunking. Small values can avoid OOM.
-            max_new_tokens:
-                Maximum number of tokens to generate.
-            **kwargs:
-                Forwarded to vllm.LLM(...).
-
-        Returns:
-            Qwen3ASRModel
-
-        Raises:
-            ImportError: If vLLM is not installed.
-        """
-        try:
-            from vllm import LLM as vLLM
-            from vllm import SamplingParams
-        except Exception as e:
-            raise ImportError(
-                "vLLM is not available. Install with: pip install qwen-asr[vllm]"
-            ) from e
-
-        llm = vLLM(model=model, **kwargs)
-
-        processor = Qwen3ASRProcessor.from_pretrained(model, fix_mistral_regex=True)
-        sampling_params = SamplingParams(**({"temperature": 0.0, "max_tokens": max_new_tokens}))
-
-        if forced_aligner is not None:
-            forced_aligner_model = Qwen3ForcedAligner.from_pretrained(
-                forced_aligner, **(forced_aligner_kwargs or {})
-            )
-
-        return cls(
-            backend="vllm",
-            model=llm,
-            processor=processor,
-            sampling_params=sampling_params,
-            forced_aligner=forced_aligner_model,
-            max_inference_batch_size=max_inference_batch_size,
-        )
-
-    def get_supported_languages(self) -> List[str]:
-        """
-        Returns the supported language list.
-
-        Returns:
-            List[str]: Canonical language names.
-        """
-        return list(SUPPORTED_LANGUAGES)
-
-    @torch.no_grad()
-    def transcribe(
-        self,
-        audio: Union[AudioLike, List[AudioLike]],
-        context: Union[str, List[str]] = "",
-        language: Optional[Union[str, List[Optional[str]]]] = None,
-        return_time_stamps: bool = False,
-    ) -> List[ASRTranscription]:
-        """
-        Transcribe audio with optional context and optional forced alignment timestamps.
-
-        Args:
-            audio:
-                Audio input(s). Supported:
-                  - str: local path / URL / base64 data url
-                  - (np.ndarray, sr)
-                  - list of above
-            context:
-                Context string(s). If scalar, it will be broadcast to batch size.
-            language:
-                Optional language(s). If provided, it must be in supported languages.
-                If scalar, it will be broadcast to batch size.
-                If provided, the prompt will force output to be transcription text only.
-            return_time_stamps:
-                If True, timestamps are produced via forced aligner and merged across chunks.
-                This requires forced_aligner initialized.
-
-        Returns:
-            List[ASRTranscription]: One result per input audio.
-
-        Raises:
-            ValueError:
-                - If return_time_stamps=True but forced_aligner is not provided.
-                - If language is unsupported.
-                - If batch sizes mismatch for context/language.
-        """
-        if return_time_stamps and self.forced_aligner is None:
-            raise ValueError("return_time_stamps=True requires `forced_aligner` to be provided at initialization.")
-
-        wavs = normalize_audios(audio)
-        n = len(wavs)
-
-        ctxs = context if isinstance(context, list) else [context]
-        if len(ctxs) == 1 and n > 1:
-            ctxs = ctxs * n
-        if len(ctxs) != n:
-            raise ValueError(f"Batch size mismatch: audio={n}, context={len(ctxs)}")
-
-        langs_in: List[Optional[str]]
-        if language is None:
-            langs_in = [None] * n
-        else:
-            langs_in = language if isinstance(language, list) else [language]
-            if len(langs_in) == 1 and n > 1:
-                langs_in = langs_in * n
-            if len(langs_in) != n:
-                raise ValueError(f"Batch size mismatch: audio={n}, language={len(langs_in)}")
-
-        langs_norm: List[Optional[str]] = []
-        for l in langs_in:
-            if l is None or str(l).strip() == "":
-                langs_norm.append(None)
-            else:
-                ln = normalize_language_name(str(l))
-                validate_language(ln)
-                langs_norm.append(ln)
-
-        max_chunk_sec = MAX_FORCE_ALIGN_INPUT_SECONDS if return_time_stamps else MAX_ASR_INPUT_SECONDS
-
-        # chunk audios and record mapping
-        chunks: List[AudioChunk] = []
-        for i, wav in enumerate(wavs):
-            parts = split_audio_into_chunks(
-                wav=wav,
-                sr=SAMPLE_RATE,
-                max_chunk_sec=max_chunk_sec,
-            )
-            for j, (cwav, offset_sec) in enumerate(parts):
-                chunks.append(AudioChunk(orig_index=i, chunk_index=j, wav=cwav, sr=SAMPLE_RATE, offset_sec=offset_sec))
-
-        # run ASR on chunks
-        chunk_ctx: List[str] = [ctxs[c.orig_index] for c in chunks]
-        chunk_lang: List[Optional[str]] = [langs_norm[c.orig_index] for c in chunks]
-        chunk_wavs: List[np.ndarray] = [c.wav for c in chunks]
-        raw_outputs = self._infer_asr(chunk_ctx, chunk_wavs, chunk_lang)
-
-        # parse outputs, prepare for optional alignment
-        per_chunk_lang: List[str] = []
-        per_chunk_text: List[str] = []
-        for out, forced_lang in zip(raw_outputs, chunk_lang):
-            lang, txt = parse_asr_output(out, user_language=forced_lang)
-            per_chunk_lang.append(lang)
-            per_chunk_text.append(txt)
-
-        # forced alignment (optional)
-        per_chunk_align: List[Optional[Any]] = [None] * len(chunks)
-        if return_time_stamps:
-            to_align_audio = []
-            to_align_text = []
-            to_align_lang = []
-            to_align_idx = []
-
-            for idx, (c, txt, lang_pred) in enumerate(zip(chunks, per_chunk_text, per_chunk_lang)):
-                if txt.strip() == "":
-                    continue
-                to_align_audio.append((c.wav, c.sr))
-                to_align_text.append(txt)
-                to_align_lang.append(lang_pred)
-                to_align_idx.append(idx)
-
-            # batch align with max_inference_batch_size
-            aligned_results: List[Any] = []
-            for a_chunk, t_chunk, l_chunk in zip(
-                chunk_list(to_align_audio, self.max_inference_batch_size),
-                chunk_list(to_align_text, self.max_inference_batch_size),
-                chunk_list(to_align_lang, self.max_inference_batch_size),
-            ):
-                aligned_results.extend(
-                    self.forced_aligner.align(audio=a_chunk, text=t_chunk, language=l_chunk)
-                )
-
-            # offset fix
-            for k, idx in enumerate(to_align_idx):
-                c = chunks[idx]
-                r = aligned_results[k]
-                per_chunk_align[idx] = self._offset_align_result(r, c.offset_sec)
-
-        # merge chunks back to original samples
-        out_langs: List[List[str]] = [[] for _ in range(n)]
-        out_texts: List[List[str]] = [[] for _ in range(n)]
-        out_aligns: List[List[Any]] = [[] for _ in range(n)]
-
-        for c, lang, txt, al in zip(chunks, per_chunk_lang, per_chunk_text, per_chunk_align):
-            out_langs[c.orig_index].append(lang)
-            out_texts[c.orig_index].append(txt)
-            if return_time_stamps and al is not None:
-                out_aligns[c.orig_index].append(al)
-
-        results: List[ASRTranscription] = []
-        for i in range(n):
-            merged_text = "".join([t for t in out_texts[i] if t is not None])
-            merged_language = merge_languages(out_langs[i])
-            merged_align = None
-            if return_time_stamps:
-                merged_align = self._merge_align_results(out_aligns[i])
-            results.append(ASRTranscription(language=merged_language, text=merged_text, time_stamps=merged_align))
-
-        return results
-
-    def _build_messages(self, context: str, audio_payload: Any) -> List[Dict[str, Any]]:
-        return [
-            {"role": "system", "content": context or ""},
-            {"role": "user", "content": [{"type": "audio", "audio": audio_payload}]},
-        ]
-
-    def _build_text_prompt(self, context: str, force_language: Optional[str]) -> str:
-        """
-        Build the string prompt for one request.
-
-        If force_language is provided, "language X<asr_text>" is appended after the generation prompt
-        to request text-only output.
-        """
-        msgs = self._build_messages(context=context, audio_payload="")
-        base = self.processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
-        if force_language:
-            base = base + f"language {force_language}{'<asr_text>'}"
-        return base
-
-    def _infer_asr(
-        self,
-        contexts: List[str],
-        wavs: List[np.ndarray],
-        languages: List[Optional[str]],
-    ) -> List[str]:
-        """
-        Run backend inference for chunk-level items.
-
-        Args:
-            contexts: List of system context strings.
-            wavs: List of mono waveforms (np.ndarray).
-            languages: List of forced languages or None.
-
-        Returns:
-            List[str]: Raw decoded strings (one per chunk).
-        """
-        if self.backend == "transformers":
-            return self._infer_asr_transformers(contexts, wavs, languages)
-        if self.backend == "vllm":
-            return self._infer_asr_vllm(contexts, wavs, languages)
-        raise RuntimeError(f"Unknown backend: {self.backend}")
-
-    def _infer_asr_transformers(
-        self,
-        contexts: List[str],
-        wavs: List[np.ndarray],
-        languages: List[Optional[str]],
-    ) -> List[str]:
-        outs: List[str] = []
-
-        texts = [self._build_text_prompt(context=c, force_language=fl) for c, fl in zip(contexts, languages)]
-
-        batch_size = self.max_inference_batch_size
-        if batch_size is None or batch_size < 0:
-            batch_size = len(texts)
-
-        for i in range(0, len(texts), batch_size):
-            sub_text = texts[i : i + batch_size]
-            sub_wavs = wavs[i : i + batch_size]
-            inputs = self.processor(text=sub_text, audio=sub_wavs, return_tensors="pt", padding=True)
-            inputs = inputs.to(self.model.device).to(self.model.dtype)
-
-            text_ids = self.model.generate(**inputs)
-
-            decoded = self.processor.batch_decode(
-                text_ids.sequences[:, inputs["input_ids"].shape[1]:],
-                skip_special_tokens=True,
-                clean_up_tokenization_spaces=False,
-            )
-            outs.extend(list(decoded))
-
-        return outs
-
-    def _infer_asr_vllm(
-        self,
-        contexts: List[str],
-        wavs: List[np.ndarray],
-        languages: List[Optional[str]],
-    ) -> List[str]:
-        inputs: List[Dict[str, Any]] = []
-        for c, w, fl in zip(contexts, wavs, languages):
-            prompt = self._build_text_prompt(context=c, force_language=fl)
-            inputs.append({"prompt": prompt, "multi_modal_data": {"audio": [w]}})
-
-        outs: List[str] = []
-        for batch in chunk_list(inputs, self.max_inference_batch_size):
-            outputs = self.model.generate(batch, sampling_params=self.sampling_params, use_tqdm=False)
-            for o in outputs:
-                outs.append(o.outputs[0].text)
-        return outs
-
-    def _offset_align_result(self, result: Any, offset_sec: float) -> Any:
-        """
-        Apply time offset to a ForcedAlignResult-like object.
-
-        This function assumes:
-          - result has attribute `.items` which is a list of items with start_time/end_time in seconds.
-          - dataclasses are frozen in upstream implementation, so we reconstruct by type.
-
-        Args:
-            result: ForcedAlignResult
-            offset_sec: Offset in seconds
-
-        Returns:
-            ForcedAlignResult: New object with shifted timestamps.
-        """
-        if result is None:
-            return None
-        items = []
-        for it in result.items:
-            items.append(type(it)(text=it.text, 
-                                  start_time=round(it.start_time + offset_sec, 3), 
-                                  end_time=round(it.end_time + offset_sec, 3)))
-        return type(result)(items=items)
-
-    def _merge_align_results(self, results: List[Any]) -> Optional[Any]:
-        """
-        Merge multiple ForcedAlignResult objects into a single one by concatenating items.
-
-        Args:
-            results: List of ForcedAlignResult
-
-        Returns:
-            ForcedAlignResult or None
-        """
-        if not results:
-            return None
-        all_items = []
-        for r in results:
-            if r is None:
-                continue
-            all_items.extend(list(r.items))
-        if not all_items:
-            return None
-        return type(results[0])(items=all_items)

qwen_asr/inference/qwen3_forced_aligner.py

@@ -1,484 +0,0 @@
-# coding=utf-8
-# 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 os
-import unicodedata
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import nagisa
-import numpy as np
-import torch
-from qwen_asr.core.transformers_backend import (
-    Qwen3ASRConfig,
-    Qwen3ASRForConditionalGeneration,
-    Qwen3ASRProcessor,
-)
-from transformers import AutoConfig, AutoModel, AutoProcessor
-
-from .utils import (
-    AudioLike,
-    ensure_list,
-    normalize_audios,
-)
-
-
-class Qwen3ForceAlignProcessor():
-    def __init__(self):
-        ko_dict_path = os.path.join(os.path.dirname(__file__), "assets", "korean_dict_jieba.dict")
-        ko_scores = {}
-        with open(ko_dict_path, "r", encoding="utf-8") as f:
-            for line in f:
-                line = line.strip()
-                if not line:
-                    continue
-                word = line.split()[0]
-                ko_scores[word] = 1.0
-        self.ko_score = ko_scores
-        self.ko_tokenizer = None
-
-    def is_kept_char(self, ch: str) -> bool:
-        if ch == "'":
-            return True
-        cat = unicodedata.category(ch)
-        if cat.startswith("L") or cat.startswith("N"):
-            return True
-        return False
-
-    def clean_token(self, token: str) -> str:
-        return "".join(ch for ch in token if self.is_kept_char(ch))
-
-    def is_cjk_char(self, ch: str) -> bool:
-        code = ord(ch)
-        return (
-            0x4E00 <= code <= 0x9FFF   # CJK Unified Ideographs
-            or 0x3400 <= code <= 0x4DBF  # Extension A
-            or 0x20000 <= code <= 0x2A6DF  # Extension B
-            or 0x2A700 <= code <= 0x2B73F  # Extension C
-            or 0x2B740 <= code <= 0x2B81F  # Extension D
-            or 0x2B820 <= code <= 0x2CEAF  # Extension E
-            or 0xF900 <= code <= 0xFAFF    # Compatibility Ideographs
-        )
-
-    def tokenize_chinese_mixed(self, text: str) -> List[str]:
-        tokens: List[str] = []
-        current_latin: List[str] = []
-
-        def flush_latin():
-            nonlocal current_latin
-            if current_latin:
-                token = "".join(current_latin)
-                cleaned = self.clean_token(token)
-                if cleaned:
-                    tokens.append(cleaned)
-                current_latin = []
-
-        for ch in text:
-            if self.is_cjk_char(ch):
-                flush_latin()
-                tokens.append(ch)
-            else:
-                if self.is_kept_char(ch):
-                    current_latin.append(ch)
-                else:
-                    flush_latin()
-
-        flush_latin()
-
-        return tokens
-
-    def tokenize_japanese(self, text: str) -> List[str]:
-        words = nagisa.tagging(text).words
-        tokens: List[str] = []
-        for w in words:
-            cleaned = self.clean_token(w)
-            if cleaned:
-                tokens.append(cleaned)
-        return tokens
-
-    def tokenize_korean(self, ko_tokenizer, text: str) -> List[str]:
-        raw_tokens = ko_tokenizer.tokenize(text)
-        tokens: List[str] = []
-        for w in raw_tokens:
-            w_clean = self.clean_token(w)
-            if w_clean:
-                tokens.append(w_clean)
-        return tokens
-
-    def split_segment_with_chinese(self, seg: str) -> List[str]:
-        tokens: List[str] = []
-        buf: List[str] = []
-
-        def flush_buf():
-            nonlocal buf
-            if buf:
-                tokens.append("".join(buf))
-                buf = []
-
-        for ch in seg:
-            if self.is_cjk_char(ch):
-                flush_buf()
-                tokens.append(ch)
-            else:
-                buf.append(ch)
-
-        flush_buf()
-        return tokens
-
-    def tokenize_space_lang(self, text: str) -> List[str]:
-        tokens: List[str] = []
-        for seg in text.split():
-            cleaned = self.clean_token(seg)
-            if cleaned:
-                tokens.extend(self.split_segment_with_chinese(cleaned))
-        return tokens
-
-    def fix_timestamp(self, data) -> List[int]:
-        data = data.tolist()
-        n = len(data)
-
-        dp = [1] * n
-        parent = [-1] * n
-        
-        for i in range(1, n):
-            for j in range(i):
-                if data[j] <= data[i] and dp[j] + 1 > dp[i]:
-                    dp[i] = dp[j] + 1
-                    parent[i] = j
-        
-        max_length = max(dp)
-        max_idx = dp.index(max_length)
-        
-        lis_indices = []
-        idx = max_idx
-        while idx != -1:
-            lis_indices.append(idx)
-            idx = parent[idx]
-        lis_indices.reverse()
-
-        is_normal = [False] * n
-        for idx in lis_indices:
-            is_normal[idx] = True
-        
-        result = data.copy()
-        i = 0
-        
-        while i < n:
-            if not is_normal[i]:
-                j = i
-                while j < n and not is_normal[j]:
-                    j += 1
-                
-                anomaly_count = j - i
-                
-                if anomaly_count <= 2:
-                    left_val = None
-                    for k in range(i - 1, -1, -1):
-                        if is_normal[k]:
-                            left_val = result[k]
-                            break
-                    
-                    right_val = None
-                    for k in range(j, n):
-                        if is_normal[k]:
-                            right_val = result[k]
-                            break
-                    
-                    for k in range(i, j):
-                        if left_val is None:
-                            result[k] = right_val
-                        elif right_val is None:
-                            result[k] = left_val
-                        else:
-                            result[k] = left_val if (k - (i - 1)) <= ((j) - k) else right_val
-                
-                else:
-                    left_val = None
-                    for k in range(i - 1, -1, -1):
-                        if is_normal[k]:
-                            left_val = result[k]
-                            break
-
-                    right_val = None
-                    for k in range(j, n):
-                        if is_normal[k]:
-                            right_val = result[k]
-                            break
-                    
-                    if left_val is not None and right_val is not None:
-                        step = (right_val - left_val) / (anomaly_count + 1)
-                        for k in range(i, j):
-                            result[k] = left_val + step * (k - i + 1)
-                    elif left_val is not None:
-                        for k in range(i, j):
-                            result[k] = left_val
-                    elif right_val is not None:
-                        for k in range(i, j):
-                            result[k] = right_val
-                
-                i = j
-            else:
-                i += 1
-
-        return [int(res) for res in result]
-
-    def encode_timestamp(self, text: str, language: str) -> List[str]:
-        language = language.lower()
-
-        if language.lower() == "japanese":
-            word_list = self.tokenize_japanese(text)
-        elif language.lower() == "korean":
-            if self.ko_tokenizer is None:
-                from soynlp.tokenizer import LTokenizer
-                self.ko_tokenizer = LTokenizer(scores=self.ko_score)
-            word_list = self.tokenize_korean(self.ko_tokenizer, text)
-        else:
-            word_list = self.tokenize_space_lang(text)
-        
-        input_text = "<timestamp><timestamp>".join(word_list) + "<timestamp><timestamp>"
-        input_text = "<|audio_start|><|audio_pad|><|audio_end|>" + input_text
-
-        return word_list, input_text
-
-    def parse_timestamp(self, word_list, timestamp):
-        timestamp_output = []
-
-        timestamp_fixed = self.fix_timestamp(timestamp)
-        for i, word in enumerate(word_list):
-            start_time = timestamp_fixed[i * 2]
-            end_time = timestamp_fixed[i * 2 + 1]
-            timestamp_output.append({
-                "text": word,
-                "start_time": start_time,
-                "end_time": end_time
-            })
-        
-        return timestamp_output
-
-
-@dataclass(frozen=True)
-class ForcedAlignItem:
-    """
-    One aligned item span.
-
-    Attributes:
-        text (str):
-            The aligned unit (cjk character or word) produced by the forced aligner processor.
-        start_time (float):
-            Start time in seconds.
-        end_time (float):
-            End time in seconds.
-    """
-    text: str
-    start_time: int
-    end_time: int
-
-
-@dataclass(frozen=True)
-class ForcedAlignResult:
-    """
-    Forced alignment output for one sample.
-
-    Attributes:
-        items (List[ForcedAlignItem]):
-            Aligned token spans.
-    """
-    items: List[ForcedAlignItem]
-
-    def __iter__(self):
-        return iter(self.items)
-
-    def __len__(self):
-        return len(self.items)
-
-    def __getitem__(self, idx: int) -> ForcedAlignItem:
-        return self.items[idx]
-
-
-class Qwen3ForcedAligner:
-    """
-    A HuggingFace-style wrapper for Qwen3-ForcedAligner model inference.
-
-    This wrapper provides:
-      - `from_pretrained()` initialization via HuggingFace AutoModel/AutoProcessor
-      - audio input normalization (path/URL/base64/(np.ndarray, sr))
-      - batch and single-sample forced alignment
-      - structured output with attribute access (`.text`, `.start_time`, `.end_time`)
-    """
-
-    def __init__(
-        self,
-        model: Qwen3ASRForConditionalGeneration,
-        processor: Qwen3ASRProcessor,
-        aligner_processor: Qwen3ForceAlignProcessor,
-    ):
-        self.model = model
-        self.processor = processor
-        self.aligner_processor = aligner_processor
-
-        self.device = getattr(model, "device", None)
-        if self.device is None:
-            try:
-                self.device = next(model.parameters()).device
-            except StopIteration:
-                self.device = torch.device("cpu")
-
-        self.timestamp_token_id = int(model.config.timestamp_token_id)
-        self.timestamp_segment_time = float(model.config.timestamp_segment_time)
-
-    @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path: str,
-        **kwargs,
-    ) -> "Qwen3ForcedAligner":
-        """
-        Load Qwen3-ForcedAligner model and initialize processors.
-
-        This method:
-          1) Registers config/model/processor for HF auto classes.
-          2) Loads the model using `AutoModel.from_pretrained(...)`.
-          3) Initializes:
-             - HF processor (`AutoProcessor.from_pretrained(...)`)
-             - forced alignment text processor (`Qwen3ForceAlignProcessor()`)
-
-        Args:
-            pretrained_model_name_or_path (str):
-                HuggingFace repo id or local directory.
-            **kwargs:
-                Forwarded to `AutoModel.from_pretrained(...)`.
-                Typical examples: device_map="cuda:0", dtype=torch.bfloat16.
-
-        Returns:
-            Qwen3ForcedAligner:
-                Initialized wrapper instance.
-        """
-        AutoConfig.register("qwen3_asr", Qwen3ASRConfig)
-        AutoModel.register(Qwen3ASRConfig, Qwen3ASRForConditionalGeneration)
-        AutoProcessor.register(Qwen3ASRConfig, Qwen3ASRProcessor)
-
-        model = AutoModel.from_pretrained(pretrained_model_name_or_path, **kwargs)
-        if not isinstance(model, Qwen3ASRForConditionalGeneration):
-            raise TypeError(
-                f"AutoModel returned {type(model)}, expected Qwen3ASRForConditionalGeneration."
-            )
-
-        processor = AutoProcessor.from_pretrained(pretrained_model_name_or_path, fix_mistral_regex=True)
-        aligner_processor = Qwen3ForceAlignProcessor()
-
-        return cls(model=model, processor=processor, aligner_processor=aligner_processor)
-
-    def _to_structured_items(self, timestamp_output: List[Dict[str, Any]]) -> ForcedAlignResult:
-        items: List[ForcedAlignItem] = []
-        for it in timestamp_output:
-            items.append(
-                ForcedAlignItem(
-                    text=str(it.get("text", "")),
-                    start_time=float(it.get("start_time", 0)),
-                    end_time=float(it.get("end_time", 0)),
-                )
-            )
-        return ForcedAlignResult(items=items)
-
-    @torch.inference_mode()
-    def align(
-        self,
-        audio: Union[AudioLike, List[AudioLike]],
-        text: Union[str, List[str]],
-        language: Union[str, List[str]],
-    ) -> List[ForcedAlignResult]:
-        """
-        Run forced alignment for batch or single sample.
-
-        Args:
-            audio:
-                Audio input(s). Each item supports:
-                  - local path / https URL / base64 string
-                  - (np.ndarray, sr)
-                All audios will be converted into mono 16k float32 arrays in [-1, 1].
-            text:
-                Transcript(s) for alignment.
-            language:
-                Language(s) for each sample (e.g., "Chinese", "English").
-
-        Returns:
-            List[ForcedAlignResult]:
-                One result per sample. Each result contains `items`, and each token can be accessed via
-                `.text`, `.start_time`, `.end_time`.
-        """
-        texts = ensure_list(text)
-        languages = ensure_list(language)
-        audios = normalize_audios(audio)
-
-        if len(languages) == 1 and len(audios) > 1:
-            languages = languages * len(audios)
-
-        if not (len(audios) == len(texts) == len(languages)):
-            raise ValueError(
-                f"Batch size mismatch: audio={len(audios)}, text={len(texts)}, language={len(languages)}"
-            )
-
-        word_lists = []
-        aligner_input_texts = []
-        for t, lang in zip(texts, languages):
-            word_list, aligner_input_text = self.aligner_processor.encode_timestamp(t, lang)
-            word_lists.append(word_list)
-            aligner_input_texts.append(aligner_input_text)
-
-        inputs = self.processor(
-            text=aligner_input_texts,
-            audio=audios,
-            return_tensors="pt",
-            padding=True,
-        )
-        inputs = inputs.to(self.model.device).to(self.model.dtype)
-
-        logits = self.model.thinker(**inputs).logits
-        output_ids = logits.argmax(dim=-1)
-
-        results: List[ForcedAlignResult] = []
-        for input_id, output_id, word_list in zip(inputs["input_ids"], output_ids, word_lists):
-            masked_output_id = output_id[input_id == self.timestamp_token_id]
-            timestamp_ms = (masked_output_id * self.timestamp_segment_time).to("cpu").numpy()
-            timestamp_output = self.aligner_processor.parse_timestamp(word_list, timestamp_ms)
-            for it in timestamp_output:
-                it['start_time'] = round(it['start_time'] / 1000.0, 3)
-                it['end_time'] = round(it['end_time'] / 1000.0, 3)
-            results.append(self._to_structured_items(timestamp_output))
-
-        return results
-    
-    def get_supported_languages(self) -> Optional[List[str]]:
-        """
-        List supported language names for the current model.
-
-        This is a thin wrapper around `self.model.get_support_languages()`.
-        If the underlying model does not expose language constraints (returns None),
-        this method also returns None.
-
-        Returns:
-            Optional[List[str]]:
-                - A sorted list of supported language names (lowercased), if available.
-                - None if the model does not provide supported languages.
-        """
-        fn = getattr(self.model, "get_support_languages", None)
-        if not callable(fn):
-            return None
-
-        langs = fn()
-        if langs is None:
-            return None
-
-        return sorted({str(x).lower() for x in langs})

qwen_asr/inference/utils.py

@@ -1,497 +0,0 @@
-# coding=utf-8
-# 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 dataclasses import dataclass
-from typing import Any, Iterable, List, Optional, Tuple, Union
-from urllib.parse import urlparse
-
-import librosa
-import numpy as np
-import soundfile as sf
-
-AudioLike = Union[
-    str,                      # wav path / URL / base64
-    Tuple[np.ndarray, int],   # (waveform, sr)
-]
-MaybeList = Union[Any, List[Any]]
-
-SAMPLE_RATE = 16000
-MAX_ASR_INPUT_SECONDS = 1200
-MAX_FORCE_ALIGN_INPUT_SECONDS = 180
-MIN_ASR_INPUT_SECONDS = 0.5
-SUPPORTED_LANGUAGES: List[str] = [
-    "Chinese",
-    "English",
-    "Cantonese",
-    "Arabic",
-    "German",
-    "French",
-    "Spanish",
-    "Portuguese",
-    "Indonesian",
-    "Italian",
-    "Korean",
-    "Russian",
-    "Thai",
-    "Vietnamese",
-    "Japanese",
-    "Turkish",
-    "Hindi",
-    "Malay",
-    "Dutch",
-    "Swedish",
-    "Danish",
-    "Finnish",
-    "Polish",
-    "Czech",
-    "Filipino",
-    "Persian",
-    "Greek",
-    "Romanian",
-    "Hungarian",
-    "Macedonian"
-]
-_ASR_TEXT_TAG = "<asr_text>"
-_LANG_PREFIX = "language "
-
-
-def normalize_language_name(language: str) -> str:
-    """
-    Normalize language name to the canonical format used by Qwen3-ASR:
-    first letter uppercase, the rest lowercase (e.g., 'cHINese' -> 'Chinese').
-
-    Args:
-        language (str): Input language name.
-
-    Returns:
-        str: Normalized language name.
-
-    Raises:
-        ValueError: If language is empty.
-    """
-    if language is None:
-        raise ValueError("language is None")
-    s = str(language).strip()
-    if not s:
-        raise ValueError("language is empty")
-    return s[:1].upper() + s[1:].lower()
-
-
-def validate_language(language: str) -> None:
-    """
-    Validate the language is supported.
-
-    Args:
-        language (str): Canonical language name.
-
-    Raises:
-        ValueError: If unsupported.
-    """
-    if language not in SUPPORTED_LANGUAGES:
-        raise ValueError(f"Unsupported language: {language}. Supported: {SUPPORTED_LANGUAGES}")
-
-
-def ensure_list(x: MaybeList) -> List[Any]:
-        return x if isinstance(x, list) else [x]
-
-
-def is_url(s: str) -> bool:
-    try:
-        u = urlparse(s)
-        return u.scheme in ("http", "https") and bool(u.netloc)
-    except Exception:
-        return False
-
-
-def is_probably_base64(s: str) -> bool:
-    if s.startswith("data:audio"):
-        return True
-    if ("/" not in s and "\\" not in s) and len(s) > 256:
-        return True
-    return False
-
-
-def decode_base64_bytes(b64: str) -> bytes:
-    if "," in b64 and b64.strip().startswith("data:"):
-        b64 = b64.split(",", 1)[1]
-    return base64.b64decode(b64)
-
-
-def load_audio_any(x: str) -> Tuple[np.ndarray, int]:
-    if 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 is_probably_base64(x):
-        audio_bytes = decode_base64_bytes(x)
-        with io.BytesIO(audio_bytes) as f:
-            audio, sr = sf.read(f, dtype="float32", always_2d=False)
-    else:
-        audio, sr = librosa.load(x, sr=None, mono=False)
-
-    audio = np.asarray(audio, dtype=np.float32)
-    sr = int(sr)
-    return audio, sr
-
-
-def to_mono(audio: np.ndarray) -> np.ndarray:
-    if audio.ndim == 1:
-        return audio
-    # soundfile can return shape (T, C); some pipelines use (C, T)
-    if audio.ndim == 2:
-        if audio.shape[0] <= 8 and audio.shape[1] > audio.shape[0]:
-            audio = audio.T
-        return np.mean(audio, axis=-1).astype(np.float32)
-    raise ValueError(f"Unsupported audio ndim={audio.ndim}")
-
-
-def float_range_normalize(audio: np.ndarray) -> np.ndarray:
-    audio = audio.astype(np.float32)
-    if audio.size == 0:
-        return audio
-    peak = float(np.max(np.abs(audio)))
-    if peak == 0.0:
-        return audio
-    # If decoded audio is int-like scaled or out-of-range, normalize conservatively.
-    if peak > 1.0:
-        audio = audio / peak
-    audio = np.clip(audio, -1.0, 1.0)
-    return audio
-
-
-def normalize_audio_input(a: AudioLike) -> np.ndarray:
-    """
-    Normalize one audio input to mono 16k float32 waveform in [-1, 1].
-
-    Supported inputs:
-        - str: local file path / https URL / base64 audio string
-        - (np.ndarray, sr): waveform and sampling rate
-
-    Returns:
-        np.ndarray:
-            Mono 16k float32 waveform in [-1, 1].
-    """
-    if isinstance(a, str):
-        audio, sr = load_audio_any(a)
-    elif isinstance(a, tuple) and len(a) == 2 and isinstance(a[0], np.ndarray):
-        audio, sr = a[0], int(a[1])
-    else:
-        raise TypeError(f"Unsupported audio input type: {type(a)}")
-
-    audio = to_mono(np.asarray(audio))
-    if sr != SAMPLE_RATE:
-        audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE).astype(np.float32)
-    audio = float_range_normalize(audio)
-    return audio
-
-
-def normalize_audios(audios: Union[AudioLike, List[AudioLike]]) -> List[np.ndarray]:
-    items = ensure_list(audios)
-    return [normalize_audio_input(a) for a in items]
-
-
-def chunk_list(xs: List[Any], chunk_size: int) -> Iterable[List[Any]]:
-    """
-    Yield chunks of a list.
-
-    Args:
-        xs (List[Any]): Input list.
-        chunk_size (int): Chunk size.
-
-    Yields:
-        List[Any]: Slices of xs.
-    """
-    if chunk_size <= 0:
-        yield xs
-        return
-    for i in range(0, len(xs), chunk_size):
-        yield xs[i : i + chunk_size]
-
-
-@dataclass(frozen=True)
-class AudioChunk:
-    """
-    One chunk cut from an original audio.
-
-    Attributes:
-        orig_index: Index of the original sample in the input batch.
-        chunk_index: Index of this chunk within the original sample.
-        wav: Mono float32 waveform.
-        sr: Sampling rate.
-        offset_sec: Start offset of this chunk in the original audio, in seconds.
-    """
-    orig_index: int
-    chunk_index: int
-    wav: np.ndarray
-    sr: int
-    offset_sec: float
-
-
-def split_audio_into_chunks(
-    wav: np.ndarray,
-    sr: int,
-    max_chunk_sec: float,
-    search_expand_sec: float = 5.0,
-    min_window_ms: float = 100.0,
-) -> List[Tuple[np.ndarray, float]]:
-    """
-    Split a long audio into chunks close to max_chunk_sec, using a low-energy boundary.
-
-    This implementation guarantees:
-      - Concatenating all returned chunks reproduces the original audio exactly
-        (total number of samples is identical, no overlaps, no gaps).
-
-    Args:
-        wav: Mono waveform float32.
-        sr: Sampling rate.
-        max_chunk_sec: Target max chunk duration in seconds.
-        search_expand_sec: Boundary search half-window in seconds.
-        min_window_ms: Sliding window in milliseconds for energy estimation.
-
-    Returns:
-        List[Tuple[np.ndarray, float]]: List of (chunk_wav, offset_sec).
-    """
-    wav = np.asarray(wav, dtype=np.float32)
-    if wav.ndim > 1:
-        wav = np.mean(wav, axis=-1).astype(np.float32)
-
-    total_len = int(wav.shape[0])
-    total_sec = total_len / float(sr)
-    if total_sec <= max_chunk_sec:
-        return [(wav, 0.0)]
-
-    max_len = int(max_chunk_sec * sr)
-    expand = int(search_expand_sec * sr)
-    win = max(4, int((min_window_ms / 1000.0) * sr))
-
-    chunks: List[Tuple[np.ndarray, float]] = []
-
-    start = 0
-    offset_sec = 0.0
-
-    while (total_len - start) > max_len:
-        cut = start + max_len
-
-        left = max(start, cut - expand)
-        right = min(total_len, cut + expand)
-
-        if right - left <= win:
-            boundary = cut
-        else:
-            seg = wav[left:right]
-            seg_abs = np.abs(seg)
-
-            window_sums = np.convolve(seg_abs, np.ones(win, dtype=np.float32), mode="valid")
-
-            min_pos = int(np.argmin(window_sums))
-
-            wstart = min_pos
-            wend = min_pos + win
-            local = seg_abs[wstart:wend]
-            inner = int(np.argmin(local))
-            boundary = left + wstart + inner
-
-        boundary = int(max(boundary, start + 1))
-        boundary = int(min(boundary, total_len))
-
-        chunk = wav[start:boundary]
-        chunks.append((chunk, offset_sec))
-
-        offset_sec += (boundary - start) / float(sr)
-        start = boundary
-
-    tail = wav[start:total_len]
-    chunks.append((tail, offset_sec))
-
-    # Pad too-short chunks to at least MIN_ASR_INPUT_SECONDS (zero-padding at tail)
-    min_len = int(MIN_ASR_INPUT_SECONDS * sr)
-    padded: List[Tuple[np.ndarray, float]] = []
-    for c, off in chunks:
-        if c.shape[0] < min_len:
-            pad = min_len - int(c.shape[0])
-            c = np.pad(c, (0, pad), mode="constant", constant_values=0.0).astype(np.float32)
-        padded.append((c, off))
-    chunks = padded
-
-    return chunks
-
-
-def detect_and_fix_repetitions(text, threshold=20):
-    def fix_char_repeats(s, thresh):
-        res = []
-        i = 0
-        n = len(s)
-        while i < n:
-            count = 1
-            while i + count < n and s[i + count] == s[i]:
-                count += 1
-
-            if count > thresh:
-                res.append(s[i])
-                i += count
-            else:
-                res.append(s[i:i+count])
-                i += count
-        return ''.join(res)
-
-    def fix_pattern_repeats(s, thresh, max_len=20):
-        n = len(s)
-        min_repeat_chars = thresh * 2
-        if n < min_repeat_chars:
-            return s
-            
-        i = 0
-        result = []
-        while i <= n - min_repeat_chars:
-            found = False
-            for k in range(1, max_len + 1):
-                if i + k * thresh > n:
-                    break
-                    
-                pattern = s[i:i+k]
-                valid = True
-                for rep in range(1, thresh):
-                    start_idx = i + rep * k
-                    if s[start_idx:start_idx+k] != pattern:
-                        valid = False
-                        break
-                
-                if valid:
-                    total_rep = thresh
-                    end_index = i + thresh * k
-                    while end_index + k <= n and s[end_index:end_index+k] == pattern:
-                        total_rep += 1
-                        end_index += k
-                    result.append(pattern)
-                    result.append(fix_pattern_repeats(s[end_index:], thresh, max_len))
-                    i = n
-                    found = True
-                    break
-            
-            if found:
-                break
-            else:
-                result.append(s[i])
-                i += 1
-
-        if not found:
-            result.append(s[i:])
-        return ''.join(result)
-    
-    text_raw = text
-    text = fix_char_repeats(text_raw, threshold)
-    text = fix_pattern_repeats(text, threshold)
-    return text
-
-
-def parse_asr_output(
-    raw: str,
-    user_language: Optional[str] = None,
-) -> Tuple[str, str]:
-    """
-    Parse Qwen3-ASR raw output into (language, text).
-
-    Cases:
-      - With tag: "language Chinese<asr_text>...."
-      - With newlines: "language Chinese\\n...\\n<asr_text>...."
-      - No tag: treat whole string as text.
-      - "language None<asr_text>": treat as empty audio -> ("", "")
-
-    If user_language is provided, language is forced to user_language and raw is treated as text-only
-    (the model is expected to output plain transcription without metadata).
-
-    Args:
-        raw: Raw decoded string.
-        user_language: Canonical language name if user forced language.
-
-    Returns:
-        Tuple[str, str]: (language, text)
-    """
-    if raw is None:
-        return "", ""
-    s = str(raw).strip()
-    if not s:
-        return "", ""
-
-    s = detect_and_fix_repetitions(s)
-
-    if user_language:
-        # user explicitly forced language => model output is treated as pure text
-        return user_language, s
-
-    meta_part = s
-    text_part = ""
-    has_tag = _ASR_TEXT_TAG in s
-    if has_tag:
-        meta_part, text_part = s.split(_ASR_TEXT_TAG, 1)
-    else:
-        # no tag => pure text
-        return "", s.strip()
-
-    meta_lower = meta_part.lower()
-
-    # empty audio heuristic
-    if "language none" in meta_lower:
-        t = text_part.strip()
-        if not t:
-            return "", ""
-        # if model still returned something, keep it but language unknown
-        return "", t
-
-    # extract "language xxx" from meta
-    lang = ""
-    for line in meta_part.splitlines():
-        line = line.strip()
-        if not line:
-            continue
-        low = line.lower()
-        if low.startswith(_LANG_PREFIX):
-            val = line[len(_LANG_PREFIX):].strip()
-            if val:
-                lang = normalize_language_name(val)
-            break
-
-    return lang, text_part.strip()
-
-
-def merge_languages(langs: List[str]) -> str:
-    """
-    Merge per-chunk languages into a compact comma-separated string,
-    keeping order and removing consecutive duplicates and empty entries.
-
-    Example:
-      ["Chinese", "English", "English"] -> "Chinese,English"
-
-    Args:
-        langs: List of canonical language names.
-
-    Returns:
-        str: Merged language string.
-    """
-    out: List[str] = []
-    prev = None
-    for x in langs:
-        x = (x or "").strip()
-        if not x:
-            continue
-        if x == prev:
-            continue
-        out.append(x)
-        prev = x
-    return ",".join(out)

qwen_tts/__init__.py

@@ -1,25 +0,0 @@
-# coding=utf-8
-# 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.
-
-"""
-qwen_tts: Qwen-TTS package.
-"""
-
-from .inference.qwen3_tts_model import Qwen3TTSModel, VoiceClonePromptItem
-from .inference.qwen3_tts_tokenizer import Qwen3TTSTokenizer
-
-__all__ = ["__version__"]
-__version__ = "0.0.1"

qwen_tts/__main__.py

@@ -1,24 +0,0 @@
-# coding=utf-8
-# 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.
-def main():
-    print(
-        "qwen_tts package.\n"
-        "Use CLI entrypoints:\n"
-        "  - qwen-tts-demo\n"
-    )
-
-if __name__ == "__main__":
-    main()

qwen_tts/cli/demo.py

@@ -1,633 +0,0 @@
-# coding=utf-8
-# 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.
-"""
-A gradio demo for Qwen3 TTS models.
-"""
-
-import argparse
-import os
-import tempfile
-from dataclasses import asdict
-from typing import Any, Dict, List, Optional, Tuple
-
-import gradio as gr
-import numpy as np
-import torch
-
-from .. import Qwen3TTSModel, VoiceClonePromptItem
-
-
-def _title_case_display(s: str) -> str:
-    s = (s or "").strip()
-    s = s.replace("_", " ")
-    return " ".join([w[:1].upper() + w[1:] if w else "" for w in s.split()])
-
-
-def _build_choices_and_map(items: Optional[List[str]]) -> Tuple[List[str], Dict[str, str]]:
-    if not items:
-        return [], {}
-    display = [_title_case_display(x) for x in items]
-    mapping = {d: r for d, r in zip(display, items)}
-    return display, mapping
-
-
-def _dtype_from_str(s: str) -> torch.dtype:
-    s = (s or "").strip().lower()
-    if s in ("bf16", "bfloat16"):
-        return torch.bfloat16
-    if s in ("fp16", "float16", "half"):
-        return torch.float16
-    if s in ("fp32", "float32"):
-        return torch.float32
-    raise ValueError(f"Unsupported torch dtype: {s}. Use bfloat16/float16/float32.")
-
-
-def _maybe(v):
-    return v if v is not None else gr.update()
-
-
-def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(
-        prog="qwen-tts-demo",
-        description=(
-            "Launch a Gradio demo for Qwen3 TTS models (CustomVoice / VoiceDesign / Base).\n\n"
-            "Examples:\n"
-            "  qwen-tts-demo Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice\n"
-            "  qwen-tts-demo Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign --port 8000 --ip 127.0.0.01\n"
-            "  qwen-tts-demo Qwen/Qwen3-TTS-12Hz-1.7B-Base --device cuda:0\n"
-            "  qwen-tts-demo Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --dtype bfloat16 --no-flash-attn\n"
-        ),
-        formatter_class=argparse.RawTextHelpFormatter,
-        add_help=True,
-    )
-
-    # Positional checkpoint (also supports -c/--checkpoint)
-    parser.add_argument(
-        "checkpoint_pos",
-        nargs="?",
-        default=None,
-        help="Model checkpoint path or HuggingFace repo id (positional).",
-    )
-    parser.add_argument(
-        "-c",
-        "--checkpoint",
-        default=None,
-        help="Model checkpoint path or HuggingFace repo id (optional if positional is provided).",
-    )
-
-    # Model loading / from_pretrained args
-    parser.add_argument(
-        "--device",
-        default="cuda:0",
-        help="Device for device_map, e.g. cpu, cuda, cuda:0 (default: cuda:0).",
-    )
-    parser.add_argument(
-        "--dtype",
-        default="bfloat16",
-        choices=["bfloat16", "bf16", "float16", "fp16", "float32", "fp32"],
-        help="Torch dtype for loading the model (default: bfloat16).",
-    )
-    parser.add_argument(
-        "--flash-attn/--no-flash-attn",
-        dest="flash_attn",
-        default=True,
-        action=argparse.BooleanOptionalAction,
-        help="Enable FlashAttention-2 (default: enabled).",
-    )
-
-    # Gradio server args
-    parser.add_argument(
-        "--ip",
-        default="0.0.0.0",
-        help="Server bind IP for Gradio (default: 0.0.0.0).",
-    )
-    parser.add_argument(
-        "--port",
-        type=int,
-        default=8000,
-        help="Server port for Gradio (default: 8000).",
-    )
-    parser.add_argument(
-        "--share/--no-share",
-        dest="share",
-        default=False,
-        action=argparse.BooleanOptionalAction,
-        help="Whether to create a public Gradio link (default: disabled).",
-    )
-    parser.add_argument(
-        "--concurrency",
-        type=int,
-        default=16,
-        help="Gradio queue concurrency (default: 16).",
-    )
-
-    # HTTPS args
-    parser.add_argument(
-        "--ssl-certfile",
-        default=None,
-        help="Path to SSL certificate file for HTTPS (optional).",
-    )
-    parser.add_argument(
-        "--ssl-keyfile",
-        default=None,
-        help="Path to SSL key file for HTTPS (optional).",
-    )
-    parser.add_argument(
-        "--ssl-verify",
-        default=None,
-        help="SSL verify setting for Gradio (optional).",
-    )
-
-    # Optional generation args
-    parser.add_argument("--max-new-tokens", type=int, default=None, help="Max new tokens for generation (optional).")
-    parser.add_argument("--temperature", type=float, default=None, help="Sampling temperature (optional).")
-    parser.add_argument("--top-k", type=int, default=None, help="Top-k sampling (optional).")
-    parser.add_argument("--top-p", type=float, default=None, help="Top-p sampling (optional).")
-    parser.add_argument("--repetition-penalty", type=float, default=None, help="Repetition penalty (optional).")
-    parser.add_argument("--subtalker-top-k", type=int, default=None, help="Subtalker top-k (optional, only for tokenizer v2).")
-    parser.add_argument("--subtalker-top-p", type=float, default=None, help="Subtalker top-p (optional, only for tokenizer v2).")
-    parser.add_argument(
-        "--subtalker-temperature", type=float, default=None, help="Subtalker temperature (optional, only for tokenizer v2)."
-    )
-
-    return parser
-
-
-def _resolve_checkpoint(args: argparse.Namespace) -> str:
-    ckpt = args.checkpoint or args.checkpoint_pos
-    if not ckpt:
-        raise SystemExit(0)  # main() prints help
-    return ckpt
-
-
-def _collect_gen_kwargs(args: argparse.Namespace) -> Dict[str, Any]:
-    mapping = {
-        "max_new_tokens": args.max_new_tokens,
-        "temperature": args.temperature,
-        "top_k": args.top_k,
-        "top_p": args.top_p,
-        "repetition_penalty": args.repetition_penalty,
-        "subtalker_top_k": args.subtalker_top_k,
-        "subtalker_top_p": args.subtalker_top_p,
-        "subtalker_temperature": args.subtalker_temperature,
-    }
-    return {k: v for k, v in mapping.items() if v is not None}
-
-
-def _normalize_audio(wav, eps=1e-12, clip=True):
-    x = np.asarray(wav)
-
-    if np.issubdtype(x.dtype, np.integer):
-        info = np.iinfo(x.dtype)
-
-        if info.min < 0:
-            y = x.astype(np.float32) / max(abs(info.min), info.max)
-        else:
-            mid = (info.max + 1) / 2.0
-            y = (x.astype(np.float32) - mid) / mid
-
-    elif np.issubdtype(x.dtype, np.floating):
-        y = x.astype(np.float32)
-        m = np.max(np.abs(y)) if y.size else 0.0
-
-        if m <= 1.0 + 1e-6:
-            pass
-        else:
-            y = y / (m + eps)
-    else:
-        raise TypeError(f"Unsupported dtype: {x.dtype}")
-
-    if clip:
-        y = np.clip(y, -1.0, 1.0)
-    
-    if y.ndim > 1:
-        y = np.mean(y, axis=-1).astype(np.float32)
-
-    return y
-
-
-def _audio_to_tuple(audio: Any) -> Optional[Tuple[np.ndarray, int]]:
-    if audio is None:
-        return None
-
-    if isinstance(audio, tuple) and len(audio) == 2 and isinstance(audio[0], int):
-        sr, wav = audio
-        wav = _normalize_audio(wav)
-        return wav, int(sr)
-
-    if isinstance(audio, dict) and "sampling_rate" in audio and "data" in audio:
-        sr = int(audio["sampling_rate"])
-        wav = _normalize_audio(audio["data"])
-        return wav, sr
-
-    return None
-
-
-def _wav_to_gradio_audio(wav: np.ndarray, sr: int) -> Tuple[int, np.ndarray]:
-    wav = np.asarray(wav, dtype=np.float32)
-    return sr, wav
-
-
-def _detect_model_kind(ckpt: str, tts: Qwen3TTSModel) -> str:
-    mt = getattr(tts.model, "tts_model_type", None)
-    if mt in ("custom_voice", "voice_design", "base"):
-        return mt
-    else:
-        raise ValueError(f"Unknown Qwen-TTS model type: {mt}")
-
-
-def build_demo(tts: Qwen3TTSModel, ckpt: str, gen_kwargs_default: Dict[str, Any]) -> gr.Blocks:
-    model_kind = _detect_model_kind(ckpt, tts)
-
-    supported_langs_raw = None
-    if callable(getattr(tts.model, "get_supported_languages", None)):
-        supported_langs_raw = tts.model.get_supported_languages()
-
-    supported_spks_raw = None
-    if callable(getattr(tts.model, "get_supported_speakers", None)):
-        supported_spks_raw = tts.model.get_supported_speakers()
-
-    lang_choices_disp, lang_map = _build_choices_and_map([x for x in (supported_langs_raw or [])])
-    spk_choices_disp, spk_map = _build_choices_and_map([x for x in (supported_spks_raw or [])])
-
-    def _gen_common_kwargs() -> Dict[str, Any]:
-        return dict(gen_kwargs_default)
-
-    theme = gr.themes.Soft(
-        font=[gr.themes.GoogleFont("Source Sans Pro"), "Arial", "sans-serif"],
-    )
-
-    css = ".gradio-container {max-width: none !important;}"
-
-    with gr.Blocks(theme=theme, css=css) as demo:
-        gr.Markdown(
-            f"""
-# Qwen3 TTS Demo
-**Checkpoint:** `{ckpt}`  
-**Model Type:** `{model_kind}`  
-"""
-        )
-
-        if model_kind == "custom_voice":
-            with gr.Row():
-                with gr.Column(scale=2):
-                    text_in = gr.Textbox(
-                        label="Text (待合成文本)",
-                        lines=4,
-                        placeholder="Enter text to synthesize (输入要合成的文本).",
-                    )
-                    with gr.Row():
-                        lang_in = gr.Dropdown(
-                            label="Language (语种)",
-                            choices=lang_choices_disp,
-                            value="Auto",
-                            interactive=True,
-                        )
-                        spk_in = gr.Dropdown(
-                            label="Speaker (说话人)",
-                            choices=spk_choices_disp,
-                            value="Vivian",
-                            interactive=True,
-                        )
-                    instruct_in = gr.Textbox(
-                        label="Instruction (Optional) (控制指令，可不输入)",
-                        lines=2,
-                        placeholder="e.g. Say it in a very angry tone (例如：用特别伤心的语气说).",
-                    )
-                    btn = gr.Button("Generate (生成)", variant="primary")
-                with gr.Column(scale=3):
-                    audio_out = gr.Audio(label="Output Audio (合成结果)", type="numpy")
-                    err = gr.Textbox(label="Status (状态)", lines=2)
-
-            def run_instruct(text: str, lang_disp: str, spk_disp: str, instruct: str):
-                try:
-                    if not text or not text.strip():
-                        return None, "Text is required (必须填写文本)."
-                    if not spk_disp:
-                        return None, "Speaker is required (必须选择说话人)."
-                    language = lang_map.get(lang_disp, "Auto")
-                    speaker = spk_map.get(spk_disp, spk_disp)
-                    kwargs = _gen_common_kwargs()
-                    wavs, sr = tts.generate_custom_voice(
-                        text=text.strip(),
-                        language=language,
-                        speaker=speaker,
-                        instruct=(instruct or "").strip() or None,
-                        **kwargs,
-                    )
-                    return _wav_to_gradio_audio(wavs[0], sr), "Finished. (生成完成)"
-                except Exception as e:
-                    return None, f"{type(e).__name__}: {e}"
-
-            btn.click(run_instruct, inputs=[text_in, lang_in, spk_in, instruct_in], outputs=[audio_out, err])
-
-        elif model_kind == "voice_design":
-            with gr.Row():
-                with gr.Column(scale=2):
-                    text_in = gr.Textbox(
-                        label="Text (待合成文本)",
-                        lines=4,
-                        value="It's in the top drawer... wait, it's empty? No way, that's impossible! I'm sure I put it there!"
-                    )
-                    with gr.Row():
-                        lang_in = gr.Dropdown(
-                            label="Language (语种)",
-                            choices=lang_choices_disp,
-                            value="Auto",
-                            interactive=True,
-                        )
-                    design_in = gr.Textbox(
-                        label="Voice Design Instruction (音色描述)",
-                        lines=3,
-                        value="Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
-                    )
-                    btn = gr.Button("Generate (生成)", variant="primary")
-                with gr.Column(scale=3):
-                    audio_out = gr.Audio(label="Output Audio (合成结果)", type="numpy")
-                    err = gr.Textbox(label="Status (状态)", lines=2)
-
-            def run_voice_design(text: str, lang_disp: str, design: str):
-                try:
-                    if not text or not text.strip():
-                        return None, "Text is required (必须填写文本)."
-                    if not design or not design.strip():
-                        return None, "Voice design instruction is required (必须填写音色描述)."
-                    language = lang_map.get(lang_disp, "Auto")
-                    kwargs = _gen_common_kwargs()
-                    wavs, sr = tts.generate_voice_design(
-                        text=text.strip(),
-                        language=language,
-                        instruct=design.strip(),
-                        **kwargs,
-                    )
-                    return _wav_to_gradio_audio(wavs[0], sr), "Finished. (生成完成)"
-                except Exception as e:
-                    return None, f"{type(e).__name__}: {e}"
-
-            btn.click(run_voice_design, inputs=[text_in, lang_in, design_in], outputs=[audio_out, err])
-
-        else:  # voice_clone for base
-            with gr.Tabs():
-                with gr.Tab("Clone & Generate (克隆并合成)"):
-                    with gr.Row():
-                        with gr.Column(scale=2):
-                            ref_audio = gr.Audio(
-                                label="Reference Audio (参考音频)",
-                            )
-                            ref_text = gr.Textbox(
-                                label="Reference Text (参考音频文本)",
-                                lines=2,
-                                placeholder="Required if not set use x-vector only (不勾选use x-vector only时必填).",
-                            )
-                            xvec_only = gr.Checkbox(
-                                label="Use x-vector only (仅用说话人向量，效果有限，但不用传入参考音频文本)",
-                                value=False,
-                            )
-
-                        with gr.Column(scale=2):
-                            text_in = gr.Textbox(
-                                label="Target Text (待合成文本)",
-                                lines=4,
-                                placeholder="Enter text to synthesize (输入要合成的文本).",
-                            )
-                            lang_in = gr.Dropdown(
-                                label="Language (语种)",
-                                choices=lang_choices_disp,
-                                value="Auto",
-                                interactive=True,
-                            )
-                            btn = gr.Button("Generate (生成)", variant="primary")
-
-                        with gr.Column(scale=3):
-                            audio_out = gr.Audio(label="Output Audio (合成结果)", type="numpy")
-                            err = gr.Textbox(label="Status (状态)", lines=2)
-
-                    def run_voice_clone(ref_aud, ref_txt: str, use_xvec: bool, text: str, lang_disp: str):
-                        try:
-                            if not text or not text.strip():
-                                return None, "Target text is required (必须填写待合成文本)."
-                            at = _audio_to_tuple(ref_aud)
-                            if at is None:
-                                return None, "Reference audio is required (必须上传参考音频)."
-                            if (not use_xvec) and (not ref_txt or not ref_txt.strip()):
-                                return None, (
-                                    "Reference text is required when use x-vector only is NOT enabled.\n"
-                                    "(未勾选 use x-vector only 时，必须提供参考音频文本；否则请勾选 use x-vector only，但效果会变差.)"
-                                )
-                            language = lang_map.get(lang_disp, "Auto")
-                            kwargs = _gen_common_kwargs()
-                            wavs, sr = tts.generate_voice_clone(
-                                text=text.strip(),
-                                language=language,
-                                ref_audio=at,
-                                ref_text=(ref_txt.strip() if ref_txt else None),
-                                x_vector_only_mode=bool(use_xvec),
-                                **kwargs,
-                            )
-                            return _wav_to_gradio_audio(wavs[0], sr), "Finished. (生成完成)"
-                        except Exception as e:
-                            return None, f"{type(e).__name__}: {e}"
-
-                    btn.click(
-                        run_voice_clone,
-                        inputs=[ref_audio, ref_text, xvec_only, text_in, lang_in],
-                        outputs=[audio_out, err],
-                    )
-
-                with gr.Tab("Save / Load Voice (保存/加载克隆音色)"):
-                    with gr.Row():
-                        with gr.Column(scale=2):
-                            gr.Markdown(
-                                """
-### Save Voice (保存音色)
-Upload reference audio and text, choose use x-vector only or not, then save a reusable voice prompt file.  
-(上传参考音频和参考文本，选择是否使用 use x-vector only 模式后保存为可复用的音色文件)
-"""
-                            )
-                            ref_audio_s = gr.Audio(label="Reference Audio (参考音频)", type="numpy")
-                            ref_text_s = gr.Textbox(
-                                label="Reference Text (参考音频文本)",
-                                lines=2,
-                                placeholder="Required if not set use x-vector only (不勾选use x-vector only时必填).",
-                            )
-                            xvec_only_s = gr.Checkbox(
-                                label="Use x-vector only (仅用说话人向量，效果有限，但不用传入参考音频文本)",
-                                value=False,
-                            )
-                            save_btn = gr.Button("Save Voice File (保存音色文件)", variant="primary")
-                            prompt_file_out = gr.File(label="Voice File (音色文件)")
-
-                        with gr.Column(scale=2):
-                            gr.Markdown(
-                                """
-### Load Voice & Generate (加载音色并合成)
-Upload a previously saved voice file, then synthesize new text.  
-(上传已保存提示文件后，输入新文本进行合成)
-"""
-                            )
-                            prompt_file_in = gr.File(label="Upload Prompt File (上传提示文件)")
-                            text_in2 = gr.Textbox(
-                                label="Target Text (待合成文本)",
-                                lines=4,
-                                placeholder="Enter text to synthesize (输入要合成的文本).",
-                            )
-                            lang_in2 = gr.Dropdown(
-                                label="Language (语种)",
-                                choices=lang_choices_disp,
-                                value="Auto",
-                                interactive=True,
-                            )
-                            gen_btn2 = gr.Button("Generate (生成)", variant="primary")
-
-                        with gr.Column(scale=3):
-                            audio_out2 = gr.Audio(label="Output Audio (合成结果)", type="numpy")
-                            err2 = gr.Textbox(label="Status (状态)", lines=2)
-
-                    def save_prompt(ref_aud, ref_txt: str, use_xvec: bool):
-                        try:
-                            at = _audio_to_tuple(ref_aud)
-                            if at is None:
-                                return None, "Reference audio is required (必须上传参考音频)."
-                            if (not use_xvec) and (not ref_txt or not ref_txt.strip()):
-                                return None, (
-                                    "Reference text is required when use x-vector only is NOT enabled.\n"
-                                    "(未勾选 use x-vector only 时，必须提供参考音频文本；否则请勾选 use x-vector only，但效果会变差.)"
-                                )
-                            items = tts.create_voice_clone_prompt(
-                                ref_audio=at,
-                                ref_text=(ref_txt.strip() if ref_txt else None),
-                                x_vector_only_mode=bool(use_xvec),
-                            )
-                            payload = {
-                                "items": [asdict(it) for it in items],
-                            }
-                            fd, out_path = tempfile.mkstemp(prefix="voice_clone_prompt_", suffix=".pt")
-                            os.close(fd)
-                            torch.save(payload, out_path)
-                            return out_path, "Finished. (生成完成)"
-                        except Exception as e:
-                            return None, f"{type(e).__name__}: {e}"
-
-                    def load_prompt_and_gen(file_obj, text: str, lang_disp: str):
-                        try:
-                            if file_obj is None:
-                                return None, "Voice file is required (必须上传音色文件)."
-                            if not text or not text.strip():
-                                return None, "Target text is required (必须填写待合成文本)."
-
-                            path = getattr(file_obj, "name", None) or getattr(file_obj, "path", None) or str(file_obj)
-                            payload = torch.load(path, map_location="cpu", weights_only=True)
-                            if not isinstance(payload, dict) or "items" not in payload:
-                                return None, "Invalid file format (文件格式不正确)."
-
-                            items_raw = payload["items"]
-                            if not isinstance(items_raw, list) or len(items_raw) == 0:
-                                return None, "Empty voice items (音色为空)."
-
-                            items: List[VoiceClonePromptItem] = []
-                            for d in items_raw:
-                                if not isinstance(d, dict):
-                                    return None, "Invalid item format in file (文件内部格式错误)."
-                                ref_code = d.get("ref_code", None)
-                                if ref_code is not None and not torch.is_tensor(ref_code):
-                                    ref_code = torch.tensor(ref_code)
-                                ref_spk = d.get("ref_spk_embedding", None)
-                                if ref_spk is None:
-                                    return None, "Missing ref_spk_embedding (缺少说话人向量)."
-                                if not torch.is_tensor(ref_spk):
-                                    ref_spk = torch.tensor(ref_spk)
-
-                                items.append(
-                                    VoiceClonePromptItem(
-                                        ref_code=ref_code,
-                                        ref_spk_embedding=ref_spk,
-                                        x_vector_only_mode=bool(d.get("x_vector_only_mode", False)),
-                                        icl_mode=bool(d.get("icl_mode", not bool(d.get("x_vector_only_mode", False)))),
-                                        ref_text=d.get("ref_text", None),
-                                    )
-                                )
-
-                            language = lang_map.get(lang_disp, "Auto")
-                            kwargs = _gen_common_kwargs()
-                            wavs, sr = tts.generate_voice_clone(
-                                text=text.strip(),
-                                language=language,
-                                voice_clone_prompt=items,
-                                **kwargs,
-                            )
-                            return _wav_to_gradio_audio(wavs[0], sr), "Finished. (生成完成)"
-                        except Exception as e:
-                            return None, (
-                                f"Failed to read or use voice file. Check file format/content.\n"
-                                f"(读取或使用音色文件失败，请检查文件格式或内容)\n"
-                                f"{type(e).__name__}: {e}"
-                            )
-
-                    save_btn.click(save_prompt, inputs=[ref_audio_s, ref_text_s, xvec_only_s], outputs=[prompt_file_out, err2])
-                    gen_btn2.click(load_prompt_and_gen, inputs=[prompt_file_in, text_in2, lang_in2], outputs=[audio_out2, err2])
-
-        gr.Markdown(
-            """
-**Disclaimer (免责声明)**  
-- The audio is automatically generated/synthesized by an AI model solely to demonstrate the model’s capabilities; it may be inaccurate or inappropriate, does not represent the views of the developer/operator, and does not constitute professional advice. You are solely responsible for evaluating, using, distributing, or relying on this audio; to the maximum extent permitted by applicable law, the developer/operator disclaims liability for any direct, indirect, incidental, or consequential damages arising from the use of or inability to use the audio, except where liability cannot be excluded by law. Do not use this service to intentionally generate or replicate unlawful, harmful, defamatory, fraudulent, deepfake, or privacy/publicity/copyright/trademark‑infringing content; if a user prompts, supplies materials, or otherwise facilitates any illegal or infringing conduct, the user bears all legal consequences and the developer/operator is not responsible.
-- 音频由人工智能模型自动生成/合成，仅用于体验与展示模型效果，可能存在不准确或不当之处；其内容不代表开发者/运营方立场，亦不构成任何专业建议。用户应自行评估并承担使用、传播或依赖该音频所产生的一切风险与责任；在适用法律允许的最大范围内，开发者/运营方不对因使用或无法使用本音频造成的任何直接、间接、附带或后果性损失承担责任（法律另有强制规定的除外）。严禁利用本服务故意引导生成或复制违法、有害、诽谤、欺诈、深度伪造、侵犯隐私/肖像/著作权/商标等内容；如用户通过提示词、素材或其他方式实施或促成任何违法或侵权行为，相关法律后果由用户自行承担，与开发者/运营方无关。
-"""
-        )
-
-    return demo
-
-
-def main(argv=None) -> int:
-    parser = build_parser()
-    args = parser.parse_args(argv)
-
-    if not args.checkpoint and not args.checkpoint_pos:
-        parser.print_help()
-        return 0
-
-    ckpt = _resolve_checkpoint(args)
-
-    dtype = _dtype_from_str(args.dtype)
-    attn_impl = "flash_attention_2" if args.flash_attn else None
-
-    tts = Qwen3TTSModel.from_pretrained(
-        ckpt,
-        device_map=args.device,
-        dtype=dtype,
-        attn_implementation=attn_impl,
-    )
-
-    gen_kwargs_default = _collect_gen_kwargs(args)
-    demo = build_demo(tts, ckpt, gen_kwargs_default)
-
-    launch_kwargs: Dict[str, Any] = dict(
-        server_name=args.ip,
-        server_port=args.port,
-        share=args.share,
-    )
-    if args.ssl_certfile is not None:
-        launch_kwargs["ssl_certfile"] = args.ssl_certfile
-    if args.ssl_keyfile is not None:
-        launch_kwargs["ssl_keyfile"] = args.ssl_keyfile
-    if args.ssl_verify is not None:
-        launch_kwargs["ssl_verify"] = args.ssl_verify
-
-    demo.queue(default_concurrency_limit=int(args.concurrency)).launch(**launch_kwargs)
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

qwen_tts/core/__init__.py

@@ -1,19 +0,0 @@
-# coding=utf-8
-# 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.
-from .tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1 import Qwen3TTSTokenizerV1Config
-from .tokenizer_25hz.modeling_qwen3_tts_tokenizer_v1 import Qwen3TTSTokenizerV1Model
-from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config
-from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Model

qwen_tts/core/models/__init__.py

@@ -1,18 +0,0 @@
-# coding=utf-8
-# 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.
-from .configuration_qwen3_tts import Qwen3TTSConfig
-from .modeling_qwen3_tts import Qwen3TTSForConditionalGeneration
-from .processing_qwen3_tts import Qwen3TTSProcessor

qwen_tts/core/models/configuration_qwen3_tts.py

@@ -1,502 +0,0 @@
-# coding=utf-8
-# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# 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.
-from transformers.configuration_utils import PretrainedConfig, layer_type_validation
-from transformers.modeling_rope_utils import rope_config_validation
-from transformers.utils import logging
-
-logger = logging.get_logger(__name__)
-
-
-class Qwen3TTSSpeakerEncoderConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`Qwen3TTSSpeakerEncoder`].
-    It is used to instantiate a Qwen3TTS speaker encoder model according to the specified arguments, defining the model
-    architecture. The architecture is based on the ECAPA-TDNN model.
-
-    Args:
-        mel_dim (`int`, *optional*, defaults to 128):
-            The dimension of the input mel-spectrogram.
-        enc_dim (`int`, *optional*, defaults to 192):
-            The dimension of the final speaker embedding.
-        enc_channels (`list[int]`, *optional*, defaults to `[512, 512, 512, 512, 1536]`):
-            A list of output channels for each TDNN/SERes2Net layer in the encoder. The first channel size is for the initial TDNN layer,
-            the intermediate ones for the `SqueezeExcitationRes2NetBlock` layers, and the last one for the multi-layer feature aggregation.
-        enc_kernel_sizes (`list[int]`, *optional*, defaults to `[5, 3, 3, 3, 1]`):
-            A list of kernel sizes for each layer in the encoder, corresponding to `enc_channels`.
-        enc_dilations (`list[int]`, *optional*, defaults to `[1, 2, 3, 4, 1]`):
-            A list of dilations for each layer in the encoder, corresponding to `enc_channels`.
-        enc_attention_channels (`int`, *optional*, defaults to 128):
-            The number of attention channels in the `AttentiveStatisticsPooling` layer.
-        enc_res2net_scale (`int`, *optional*,defaults to 8):
-            The scale of the `Res2NetBlock` in the encoder.
-        enc_se_channels (`int`, *optional*, defaults to 128):
-            The number of channels in the squeeze part of the `SqueezeExcitationBlock`.
-    """
-    def __init__(
-        self,
-        mel_dim=128,
-        enc_dim=1024,
-        enc_channels=[512, 512, 512, 512, 1536],
-        enc_kernel_sizes=[5, 3, 3, 3, 1],
-        enc_dilations=[1, 2, 3, 4, 1],
-        enc_attention_channels=128,
-        enc_res2net_scale=8,
-        enc_se_channels=128,
-        sample_rate=24000,
-    ):
-        self.mel_dim = mel_dim
-        self.enc_dim = enc_dim
-        self.enc_channels = enc_channels
-        self.enc_kernel_sizes = enc_kernel_sizes
-        self.enc_dilations = enc_dilations
-        self.enc_attention_channels = enc_attention_channels
-        self.enc_res2net_scale = enc_res2net_scale
-        self.enc_se_channels = enc_se_channels
-        self.sample_rate = sample_rate
-
-
-class Qwen3TTSTalkerCodePredictorConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`Qwen3TTSTalkerCodePredictorModel`]. It is used to instantiate a
-    Qwen3TTSTalkerCodePredictor model according to the specified arguments, defining the model architecture.
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-
-    Args:
-        vocab_size (`int`, *optional*, defaults to 151936):
-            Vocabulary size of the Qwen3TTSTalkerCodePredictor model. Defines the number of different tokens that can be represented by the
-            `inputs_ids` passed when calling [`Qwen3TTSTalkerCodePredictorModel`]
-        hidden_size (`int`, *optional*, defaults to 4096):
-            Dimension of the hidden representations.
-        intermediate_size (`int`, *optional*, defaults to 22016):
-            Dimension of the MLP representations.
-        num_hidden_layers (`int`, *optional*, defaults to 32):
-            Number of hidden layers in the Transformer encoder.
-        num_attention_heads (`int`, *optional*, defaults to 32):
-            Number of attention heads for each attention layer in the Transformer encoder.
-        num_key_value_heads (`int`, *optional*, defaults to 32):
-            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
-            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
-            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
-            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
-            by meanpooling all the original heads within that group. For more details, check out [this
-            paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`.
-        head_dim (`int`, *optional*, defaults to 128):
-            The attention head dimension.
-        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
-            The non-linear activation function (function or string) in the decoder.
-        max_position_embeddings (`int`, *optional*, defaults to 32768):
-            The maximum sequence length that this model might ever be used with.
-        initializer_range (`float`, *optional*, defaults to 0.02):
-            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
-            The epsilon used by the rms normalization layers.
-        use_cache (`bool`, *optional*, defaults to `True`):
-            Whether or not the model should return the last key/values attentions (not used by all models). Only
-            relevant if `config.is_decoder=True`.
-        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
-            Whether the model's input and output word embeddings should be tied.
-        rope_theta (`float`, *optional*, defaults to 10000.0):
-            The base period of the RoPE embeddings.
-        rope_scaling (`Dict`, *optional*):
-            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
-            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
-            accordingly.
-            Expected contents:
-                `rope_type` (`str`):
-                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
-                    'llama3'], with 'default' being the original RoPE implementation.
-                `factor` (`float`, *optional*):
-                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
-                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
-                    original maximum pre-trained length.
-                `original_max_position_embeddings` (`int`, *optional*):
-                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
-                    pretraining.
-                `attention_factor` (`float`, *optional*):
-                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
-                    computation. If unspecified, it defaults to value recommended by the implementation, using the
-                    `factor` field to infer the suggested value.
-                `beta_fast` (`float`, *optional*):
-                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
-                    ramp function. If unspecified, it defaults to 32.
-                `beta_slow` (`float`, *optional*):
-                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
-                    ramp function. If unspecified, it defaults to 1.
-                `short_factor` (`list[float]`, *optional*):
-                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
-                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
-                    size divided by the number of attention heads divided by 2
-                `long_factor` (`list[float]`, *optional*):
-                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
-                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
-                    size divided by the number of attention heads divided by 2
-                `low_freq_factor` (`float`, *optional*):
-                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
-                `high_freq_factor` (`float`, *optional*):
-                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
-        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
-            Whether to use a bias in the query, key, value and output projection layers during self-attention.
-        use_sliding_window (`bool`, *optional*, defaults to `False`):
-            Whether to use sliding window attention.
-        sliding_window (`int`, *optional*, defaults to 4096):
-            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
-        max_window_layers (`int`, *optional*, defaults to 28):
-            The number of layers using full attention. The first `max_window_layers` layers will use full attention, while any
-            additional layer afterwards will use SWA (Sliding Window Attention).
-        layer_types (`list`, *optional*):
-            Attention pattern for each layer.
-        attention_dropout (`float`, *optional*, defaults to 0.0):
-            The dropout ratio for the attention probabilities.
-
-    """
-
-    model_type = "qwen3_tts_talker_code_predictor"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    # Default tensor parallel plan for base model `Qwen3TTSTalkerCodePredictor`
-    base_model_tp_plan = {
-        "layers.*.self_attn.q_proj": "colwise",
-        "layers.*.self_attn.k_proj": "colwise",
-        "layers.*.self_attn.v_proj": "colwise",
-        "layers.*.self_attn.o_proj": "rowwise",
-        "layers.*.mlp.gate_proj": "colwise",
-        "layers.*.mlp.up_proj": "colwise",
-        "layers.*.mlp.down_proj": "rowwise",
-    }
-    base_model_pp_plan = {
-        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
-        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
-        "norm": (["hidden_states"], ["hidden_states"]),
-    }
-
-    def __init__(
-        self,
-        vocab_size=2048,
-        hidden_size=1024,
-        intermediate_size=3072,
-        num_hidden_layers=5,
-        num_attention_heads=16,
-        num_key_value_heads=8,
-        head_dim=128,
-        hidden_act="silu",
-        max_position_embeddings=32768,
-        initializer_range=0.02,
-        rms_norm_eps=0.000001,
-        use_cache=True,
-        tie_word_embeddings=False,
-        rope_theta=10000,
-        rope_scaling=None,
-        attention_bias=False,
-        use_sliding_window=False,
-        sliding_window=4096,
-        max_window_layers=28,
-        layer_types=None,
-        attention_dropout=0,
-        num_code_groups=32,
-        **kwargs,
-    ):
-        super().__init__(
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs,
-        )
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.use_sliding_window = use_sliding_window
-        self.sliding_window = sliding_window if self.use_sliding_window else None
-        self.max_window_layers = max_window_layers
-
-        # for backward compatibility
-        if num_key_value_heads is None:
-            num_key_value_heads = num_attention_heads
-
-        self.num_key_value_heads = num_key_value_heads
-        self.head_dim = head_dim
-        self.hidden_act = hidden_act
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling
-        self.attention_bias = attention_bias
-        self.attention_dropout = attention_dropout
-        # Validate the correctness of rotary position embeddings parameters
-        # BC: if there is a 'type' field, move it to 'rope_type'.
-        if self.rope_scaling is not None and "type" in self.rope_scaling:
-            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
-        rope_config_validation(self)
-
-        self.layer_types = layer_types
-        if self.layer_types is None:
-            self.layer_types = [
-                "sliding_attention"
-                if self.sliding_window is not None and i >= self.max_window_layers
-                else "full_attention"
-                for i in range(self.num_hidden_layers)
-            ]
-        layer_type_validation(self.layer_types)
-        self.num_code_groups = num_code_groups
-
-
-class Qwen3TTSTalkerConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`Qwen3TTSTalkerModel`]. It is used to instantiate a
-    Qwen3TTSTalker model according to the specified arguments, defining the model architecture.
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-
-    Args:
-        vocab_size (`int`, *optional*, defaults to 151936):
-            Vocabulary size of the Qwen3TTSTalker model. Defines the number of different tokens that can be represented by the
-            `inputs_ids` passed when calling [`Qwen3TTSTalkerModel`]
-        hidden_size (`int`, *optional*, defaults to 2048):
-            Dimension of the hidden representations.
-        intermediate_size (`int`, *optional*, defaults to 6144):
-            Dimension of the MLP representations.
-        num_hidden_layers (`int`, *optional*, defaults to 24):
-            Number of hidden layers in the Transformer encoder.
-        num_attention_heads (`int`, *optional*, defaults to 32):
-            Number of attention heads for each attention layer in the Transformer encoder.
-        num_key_value_heads (`int`, *optional*, defaults to 4):
-            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
-            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
-            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
-            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
-            by meanpooling all the original heads within that group. For more details, check out [this
-            paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`.
-
-        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
-            The non-linear activation function (function or string) in the decoder.
-        max_position_embeddings (`int`, *optional*, defaults to 32768):
-            The maximum sequence length that this model might ever be used with.
-        initializer_range (`float`, *optional*, defaults to 0.02):
-            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
-            The epsilon used by the rms normalization layers.
-        use_cache (`bool`, *optional*, defaults to `True`):
-            Whether or not the model should return the last key/values attentions (not used by all models). Only
-            relevant if `config.is_decoder=True`.
-        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
-            Whether the model's input and output word embeddings should be tied.
-        rope_theta (`float`, *optional*, defaults to 10000.0):
-            The base period of the RoPE embeddings.
-        rope_scaling (`Dict`, *optional*):
-            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
-            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
-            accordingly.
-            Expected contents:
-                `rope_type` (`str`):
-                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
-                    'llama3'], with 'default' being the original RoPE implementation.
-                `factor` (`float`, *optional*):
-                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
-                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
-                    original maximum pre-trained length.
-                `original_max_position_embeddings` (`int`, *optional*):
-                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
-                    pretraining.
-                `attention_factor` (`float`, *optional*):
-                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
-                    computation. If unspecified, it defaults to value recommended by the implementation, using the
-                    `factor` field to infer the suggested value.
-                `beta_fast` (`float`, *optional*):
-                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
-                    ramp function. If unspecified, it defaults to 32.
-                `beta_slow` (`float`, *optional*):
-                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
-                    ramp function. If unspecified, it defaults to 1.
-                `short_factor` (`list[float]`, *optional*):
-                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
-                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
-                    size divided by the number of attention heads divided by 2
-                `long_factor` (`list[float]`, *optional*):
-                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
-                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
-                    size divided by the number of attention heads divided by 2
-                `low_freq_factor` (`float`, *optional*):
-                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
-                `high_freq_factor` (`float`, *optional*):
-                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
-        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
-            Whether to use a bias in the query, key, value and output projection layers during self-attention.
-        use_sliding_window (`bool`, *optional*, defaults to `False`):
-            Whether to use sliding window attention.
-        sliding_window (`int`, *optional*, defaults to 4096):
-            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
-        attention_dropout (`float`, *optional*, defaults to 0.0):
-            The dropout ratio for the attention probabilities.
-    """
-
-    model_type = "qwen3_tts_talker"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    # Default tensor parallel plan for base model `Qwen3TTSTalker`
-    base_model_tp_plan = {
-        "layers.*.self_attn.q_proj": "colwise",
-        "layers.*.self_attn.k_proj": "colwise",
-        "layers.*.self_attn.v_proj": "colwise",
-        "layers.*.self_attn.o_proj": "rowwise",
-        "layers.*.mlp.gate_proj": "colwise",
-        "layers.*.mlp.up_proj": "colwise",
-        "layers.*.mlp.down_proj": "rowwise",
-    }
-    base_model_pp_plan = {
-        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
-        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
-        "norm": (["hidden_states"], ["hidden_states"]),
-    }
-    sub_configs = {"code_predictor_config": Qwen3TTSTalkerCodePredictorConfig}
-
-    def __init__(
-        self,
-        code_predictor_config=None,
-        vocab_size=3072,
-        hidden_size=1024,
-        intermediate_size=2048,
-        num_hidden_layers=20,
-        num_attention_heads=16,
-        num_key_value_heads=2,
-        hidden_act="silu",
-        max_position_embeddings=32768,
-        initializer_range=0.02,
-        rms_norm_eps=0.000001,
-        use_cache=True,
-        tie_word_embeddings=False,
-        rope_theta=10000,
-        rope_scaling=None,
-        attention_bias=False,
-        use_sliding_window=False,
-        sliding_window=4096,
-        attention_dropout=0,
-        num_code_groups=32,
-        text_hidden_size=2048,
-        codec_eos_token_id=4198,
-        codec_think_id=4202,
-        codec_nothink_id=4203,
-        codec_think_bos_id=4204,
-        codec_think_eos_id=4205,
-        codec_pad_id=4196,
-        codec_bos_id=4197,
-        spk_id=None,
-        spk_is_dialect=None,
-        codec_language_id=None,
-        **kwargs,
-    ):
-        super().__init__(
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs,
-        )
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.use_sliding_window = use_sliding_window
-        self.sliding_window = sliding_window if use_sliding_window else None
-
-        self.num_key_value_heads = num_key_value_heads
-        self.hidden_act = hidden_act
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling
-        self.attention_bias = attention_bias
-        self.attention_dropout = attention_dropout
-        # Validate the correctness of rotary position embeddings parameters
-        # BC: if there is a 'type' field, move it to 'rope_type'.
-        if self.rope_scaling is not None and "type" in self.rope_scaling:
-            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
-
-        if code_predictor_config is None:
-            code_predictor_config = {}
-            self.code_predictor_config = Qwen3TTSTalkerCodePredictorConfig()
-            logger.info("code_predictor_config is None. Initializing code_predictor model with default values")
-        elif isinstance(code_predictor_config, Qwen3TTSTalkerCodePredictorConfig):
-            self.code_predictor_config = code_predictor_config
-        else:
-            self.code_predictor_config = Qwen3TTSTalkerCodePredictorConfig(**code_predictor_config)
-        self.num_code_groups = num_code_groups
-        self.text_hidden_size = text_hidden_size
-        self.codec_eos_token_id = codec_eos_token_id
-        self.codec_think_id = codec_think_id
-        self.codec_language_id = codec_language_id
-        self.codec_nothink_id = codec_nothink_id
-        self.codec_think_bos_id = codec_think_bos_id
-        self.codec_think_eos_id = codec_think_eos_id
-        self.codec_pad_id = codec_pad_id
-        self.codec_bos_id = codec_bos_id
-        self.spk_id = spk_id
-        self.spk_is_dialect = spk_is_dialect
-
-
-class Qwen3TTSConfig(PretrainedConfig):
-    """
-    This is the configuration class to store the configuration of a [`Qwen3TTSForConditionalGeneration`]. 
-    """
-
-    model_type = "qwen3_tts"
-    sub_configs = {
-        "talker_config": Qwen3TTSTalkerConfig,
-        "speaker_encoder_config": Qwen3TTSSpeakerEncoderConfig,
-    }
-
-    def __init__(
-        self,
-        talker_config=None,
-        speaker_encoder_config=None,
-        tokenizer_type=None,
-        tts_model_size=None,
-        tts_model_type=None,
-        im_start_token_id=151644,
-        im_end_token_id=151645,
-        tts_pad_token_id=151671,
-        tts_bos_token_id=151672,
-        tts_eos_token_id=151673,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-
-        if talker_config is None:
-            talker_config = {}
-            logger.info("talker_config is None. Initializing talker model with default values")
-        if speaker_encoder_config is None:
-            speaker_encoder_config = {}
-            logger.info("speaker_encoder_config is None. Initializing talker model with default values")
-
-        self.talker_config = Qwen3TTSTalkerConfig(**talker_config)
-        self.speaker_encoder_config = Qwen3TTSSpeakerEncoderConfig(**speaker_encoder_config)
-
-        self.tokenizer_type = tokenizer_type
-        self.tts_model_size = tts_model_size
-        self.tts_model_type = tts_model_type
-
-        self.im_start_token_id = im_start_token_id
-        self.im_end_token_id = im_end_token_id
-        self.tts_pad_token_id = tts_pad_token_id
-        self.tts_bos_token_id = tts_bos_token_id
-        self.tts_eos_token_id = tts_eos_token_id
-
-
-__all__ = ["Qwen3TTSConfig", "Qwen3TTSTalkerConfig", "Qwen3TTSSpeakerEncoderConfig"]

qwen_tts/core/models/modeling_qwen3_tts.py

@@ -1,2246 +0,0 @@
-# coding=utf-8
-# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# 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.
-"""PyTorch Qwen3TTS model."""
-
-import json
-import os
-from dataclasses import dataclass
-from typing import Callable, Optional
-
-import torch
-from librosa.filters import mel as librosa_mel_fn
-from torch import nn
-from torch.nn import functional as F
-from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache, DynamicCache
-from transformers.generation import GenerationMixin
-from transformers.integrations import use_kernel_forward_from_hub
-from transformers.masking_utils import (
-    create_causal_mask,
-    create_sliding_window_causal_mask,
-)
-from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
-from transformers.modeling_layers import GradientCheckpointingLayer
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-    ModelOutput,
-)
-from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
-from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
-from transformers.processing_utils import Unpack
-from transformers.utils import can_return_tuple, logging
-from transformers.utils.hub import cached_file
-
-from ...inference.qwen3_tts_tokenizer import Qwen3TTSTokenizer
-from .configuration_qwen3_tts import (
-    Qwen3TTSConfig,
-    Qwen3TTSSpeakerEncoderConfig,
-    Qwen3TTSTalkerCodePredictorConfig,
-    Qwen3TTSTalkerConfig,
-)
-
-logger = logging.get_logger(__name__)
-
-
-class Res2NetBlock(torch.nn.Module):
-    def __init__(self, in_channels, out_channels, scale=8, kernel_size=3, dilation=1):
-        super().__init__()
-
-        in_channel = in_channels // scale
-        hidden_channel = out_channels // scale
-
-        self.blocks = nn.ModuleList(
-            [
-                TimeDelayNetBlock(
-                    in_channel,
-                    hidden_channel,
-                    kernel_size=kernel_size,
-                    dilation=dilation,
-                )
-                for i in range(scale - 1)
-            ]
-        )
-        self.scale = scale
-
-    def forward(self, hidden_states):
-        outputs = []
-        for i, hidden_part in enumerate(torch.chunk(hidden_states, self.scale, dim=1)):
-            if i == 0:
-                output_part = hidden_part
-            elif i == 1:
-                output_part = self.blocks[i - 1](hidden_part)
-            else:
-                output_part = self.blocks[i - 1](hidden_part + output_part)
-            outputs.append(output_part)
-        output = torch.cat(outputs, dim=1)
-        return output
-
-
-class SqueezeExcitationBlock(nn.Module):
-    def __init__(self, in_channels, se_channels, out_channels):
-        super().__init__()
-
-        self.conv1 = nn.Conv1d(
-            in_channels=in_channels,
-            out_channels=se_channels,
-            kernel_size=1,
-            padding="same",
-            padding_mode="reflect",
-        )
-        self.relu = nn.ReLU(inplace=True)
-        self.conv2 = nn.Conv1d(
-            in_channels=se_channels,
-            out_channels=out_channels,
-            kernel_size=1,
-            padding="same",
-            padding_mode="reflect",
-        )
-        self.sigmoid = nn.Sigmoid()
-
-    def forward(self, hidden_states):
-        hidden_states_mean = hidden_states.mean(dim=2, keepdim=True)
-
-        hidden_states_mean = self.relu(self.conv1(hidden_states_mean))
-        hidden_states_mean = self.sigmoid(self.conv2(hidden_states_mean))
-
-        return hidden_states * hidden_states_mean
-
-
-class AttentiveStatisticsPooling(nn.Module):
-    """This class implements an attentive statistic pooling layer for each channel.
-    It returns the concatenated mean and std of the input tensor.
-    """
-
-    def __init__(self, channels, attention_channels=128):
-        super().__init__()
-
-        self.eps = 1e-12
-        self.tdnn = TimeDelayNetBlock(channels * 3, attention_channels, 1, 1)
-        self.tanh = nn.Tanh()
-        self.conv = nn.Conv1d(
-            in_channels=attention_channels,
-            out_channels=channels,
-            kernel_size=1,
-            padding="same",
-            padding_mode="reflect",
-        )
-
-    def _length_to_mask(self, length, max_len=None, dtype=None, device=None):
-        """Creates a binary mask for each sequence.
-
-        Reference: https://discuss.pytorch.org/t/how-to-generate-variable-length-mask/23397/3
-
-        Arguments
-        ---------
-        length : torch.LongTensor
-            Containing the length of each sequence in the batch. Must be 1D.
-        max_len : int
-            Max length for the mask, also the size of the second dimension.
-        dtype : torch.dtype, default: None
-            The dtype of the generated mask.
-        device: torch.device, default: None
-            The device to put the mask variable.
-
-        Returns
-        -------
-        mask : tensor
-            The binary mask.
-        """
-
-        if max_len is None:
-            max_len = length.max().long().item()  # using arange to generate mask
-        mask = torch.arange(max_len, device=length.device, dtype=length.dtype).expand(
-            len(length), max_len
-        ) < length.unsqueeze(1)
-
-        mask = torch.as_tensor(mask, dtype=dtype, device=device)
-        return mask
-
-    def _compute_statistics(self, x, m, dim=2):
-        mean = (m * x).sum(dim)
-        std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps))
-        return mean, std
-
-    def forward(self, hidden_states):
-        seq_length = hidden_states.shape[-1]
-        lengths = torch.ones(hidden_states.shape[0], device=hidden_states.device)
-
-        # Make binary mask of shape [N, 1, L]
-        mask = self._length_to_mask(
-            lengths * seq_length, max_len=seq_length, dtype=hidden_states.dtype, device=hidden_states.device
-        )
-        mask = mask.unsqueeze(1)
-
-        # Expand the temporal context of the pooling layer by allowing the
-        # self-attention to look at global properties of the utterance.
-        total = mask.sum(dim=2, keepdim=True)
-
-        mean, std = self._compute_statistics(hidden_states, mask / total)
-        mean = mean.unsqueeze(2).repeat(1, 1, seq_length)
-        std = std.unsqueeze(2).repeat(1, 1, seq_length)
-        attention = torch.cat([hidden_states, mean, std], dim=1)
-
-        # Apply layers
-        attention = self.conv(self.tanh(self.tdnn(attention)))
-
-        # Filter out zero-paddings
-        attention = attention.masked_fill(mask == 0, float("-inf"))
-
-        attention = F.softmax(attention, dim=2)
-        mean, std = self._compute_statistics(hidden_states, attention)
-        # Append mean and std of the batch
-        pooled_stats = torch.cat((mean, std), dim=1)
-        pooled_stats = pooled_stats.unsqueeze(2)
-
-        return pooled_stats
-
-class TimeDelayNetBlock(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        dilation,
-    ):
-        super().__init__()
-        self.conv = nn.Conv1d(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            kernel_size=kernel_size,
-            dilation=dilation,
-            padding="same",
-            padding_mode="reflect",
-        )
-        self.activation = nn.ReLU()
-
-    def forward(self, hidden_states: torch.Tensor):
-        return self.activation(self.conv(hidden_states))
-
-class SqueezeExcitationRes2NetBlock(nn.Module):
-    """An implementation of building block in ECAPA-TDNN, i.e.,
-    TDNN-Res2Net-TDNN-SqueezeExcitationBlock.
-    """
-
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        res2net_scale=8,
-        se_channels=128,
-        kernel_size=1,
-        dilation=1,
-    ):
-        super().__init__()
-        self.out_channels = out_channels
-        self.tdnn1 = TimeDelayNetBlock(
-            in_channels,
-            out_channels,
-            kernel_size=1,
-            dilation=1,
-        )
-        self.res2net_block = Res2NetBlock(out_channels, out_channels, res2net_scale, kernel_size, dilation)
-        self.tdnn2 = TimeDelayNetBlock(
-            out_channels,
-            out_channels,
-            kernel_size=1,
-            dilation=1,
-        )
-        self.se_block = SqueezeExcitationBlock(out_channels, se_channels, out_channels)
-
-    def forward(self, hidden_state):
-        residual = hidden_state
-
-        hidden_state = self.tdnn1(hidden_state)
-        hidden_state = self.res2net_block(hidden_state)
-        hidden_state = self.tdnn2(hidden_state)
-        hidden_state = self.se_block(hidden_state)
-
-        return hidden_state + residual
-
-
-class Qwen3TTSSpeakerEncoder(torch.nn.Module):
-    """An implementation of the speaker embedding model in a paper.
-    "ECAPA-TDNN: Emphasized Channel Attention, Propagation and Aggregation in
-    TDNN Based Speaker Verification" (https://huggingface.co/papers/2005.07143).
-    Use for Qwen3TTS extract speaker embedding.
-    """
-
-    def __init__(self, config: Qwen3TTSSpeakerEncoderConfig):
-        super().__init__()
-        if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(config.enc_channels) != len(
-            config.enc_dilations
-        ):
-            raise ValueError("enc_channels, enc_kernel_sizes and enc_dilations should have same length")
-        self.channels = config.enc_channels
-        self.blocks = nn.ModuleList()
-
-        # The initial TDNN layer
-        self.blocks.append(
-            TimeDelayNetBlock(
-                config.mel_dim,
-                config.enc_channels[0],
-                config.enc_kernel_sizes[0],
-                config.enc_dilations[0],
-            )
-        )
-
-        # SE-Res2Net layers
-        for i in range(1, len(config.enc_channels) - 1):
-            self.blocks.append(
-                SqueezeExcitationRes2NetBlock(
-                    config.enc_channels[i - 1],
-                    config.enc_channels[i],
-                    res2net_scale=config.enc_res2net_scale,
-                    se_channels=config.enc_se_channels,
-                    kernel_size=config.enc_kernel_sizes[i],
-                    dilation=config.enc_dilations[i],
-                )
-            )
-
-        # Multi-layer feature aggregation
-        self.mfa = TimeDelayNetBlock(
-            config.enc_channels[-1],
-            config.enc_channels[-1],
-            config.enc_kernel_sizes[-1],
-            config.enc_dilations[-1],
-        )
-
-        # Attentive Statistical Pooling
-        self.asp = AttentiveStatisticsPooling(
-            config.enc_channels[-1],
-            attention_channels=config.enc_attention_channels,
-        )
-
-        # Final linear transformation
-        self.fc = nn.Conv1d(
-            in_channels=config.enc_channels[-1] * 2,
-            out_channels=config.enc_dim,
-            kernel_size=1,
-            padding="same",
-            padding_mode="reflect",
-        )
-
-    def forward(self, hidden_states):
-        # Minimize transpose for efficiency
-        hidden_states = hidden_states.transpose(1, 2)
-
-        hidden_states_list = []
-        for layer in self.blocks:
-            hidden_states = layer(hidden_states)
-            hidden_states_list.append(hidden_states)
-
-        # Multi-layer feature aggregation
-        hidden_states = torch.cat(hidden_states_list[1:], dim=1)
-        hidden_states = self.mfa(hidden_states)
-
-        # Attentive Statistical Pooling
-        hidden_states = self.asp(hidden_states)
-
-        # Final linear transformation
-        hidden_states = self.fc(hidden_states)
-
-        hidden_states = hidden_states.squeeze(-1)
-        return hidden_states
-
-
-def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
-    return torch.log(torch.clamp(x, min=clip_val) * C)
-
-def mel_spectrogram(
-    y: torch.Tensor,
-    n_fft: int,
-    num_mels: int,
-    sampling_rate: int,
-    hop_size: int,
-    win_size: int,
-    fmin: int,
-    fmax: int = None,
-    center: bool = False,
-) -> torch.Tensor:
-    """
-    Calculate the mel spectrogram of an input signal.
-    This function uses slaney norm for the librosa mel filterbank (using librosa.filters.mel) and uses Hann window for STFT (using torch.stft).
-
-    Args:
-        y (torch.Tensor): Input signal.
-        n_fft (int): FFT size.
-        num_mels (int): Number of mel bins.
-        sampling_rate (int): Sampling rate of the input signal.
-        hop_size (int): Hop size for STFT.
-        win_size (int): Window size for STFT.
-        fmin (int): Minimum frequency for mel filterbank.
-        fmax (int): Maximum frequency for mel filterbank. If None, defaults to half the sampling rate (fmax = sr / 2.0) inside librosa_mel_fn
-        center (bool): Whether to pad the input to center the frames. Default is False.
-
-    Returns:
-        torch.Tensor: Mel spectrogram.
-    """
-    if torch.min(y) < -1.0:
-        print(f"[WARNING] Min value of input waveform signal is {torch.min(y)}")
-    if torch.max(y) > 1.0:
-        print(f"[WARNING] Max value of input waveform signal is {torch.max(y)}")
-
-    device = y.device
-
-    mel = librosa_mel_fn(
-        sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
-    )
-
-    mel_basis = torch.from_numpy(mel).float().to(device)
-    hann_window = torch.hann_window(win_size).to(device)
-
-    padding = (n_fft - hop_size) // 2
-    y = torch.nn.functional.pad(
-        y.unsqueeze(1), (padding, padding), mode="reflect"
-    ).squeeze(1)
-
-    spec = torch.stft(
-        y,
-        n_fft,
-        hop_length=hop_size,
-        win_length=win_size,
-        window=hann_window,
-        center=center,
-        pad_mode="reflect",
-        normalized=False,
-        onesided=True,
-        return_complex=True,
-    )
-    spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
-
-    mel_spec = torch.matmul(mel_basis, spec)
-    mel_spec = dynamic_range_compression_torch(mel_spec)
-
-    return mel_spec
-
-
-class Qwen3TTSPreTrainedModel(PreTrainedModel):
-    config_class = Qwen3TTSConfig
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = True
-    _no_split_modules = ["Qwen3TTSDecoderLayer"]
-    _skip_keys_device_placement = "past_key_values"
-    _supports_flash_attn_2 = True
-    _supports_sdpa = True
-    _supports_cache_class = True
-    _supports_static_cache = False
-    _supports_attention_backend = True
-
-    def _init_weights(self, module):
-        # important: this ported version of Qwen2.5OmniThinker isn't meant for training from scratch - only
-        # inference and fine-tuning - so the proper init weights code has been removed
-        std = self.config.initializer_range if hasattr(self.config, "initializer_range") else 0.02
-
-        if isinstance(module, (nn.Linear, nn.Conv1d, nn.Conv3d, nn.ConvTranspose1d)):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, nn.LayerNorm):
-            if module.weight is not None:
-                module.weight.data.fill_(1.0)
-            if module.bias is not None:
-                module.bias.data.zero_()
-
-
-class Qwen3TTSTalkerTextPreTrainedModel(PreTrainedModel):
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = True
-    _no_split_modules = []
-    _skip_keys_device_placement = ["past_key_values"]
-    _supports_flash_attn_3 = True
-    _supports_flash_attn_2 = True
-    _supports_sdpa = True
-    _supports_flex_attn = True
-    _supports_cache_class = True
-    _supports_quantized_cache = True
-    _supports_static_cache = False
-    _supports_attention_backend = True
-
-    def _init_weights(self, module):
-        std = self.config.initializer_range
-        if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, Qwen3TTSRMSNorm):
-            module.weight.data.fill_(1.0)
-
-
-class Qwen3TTSTalkerRotaryEmbedding(nn.Module):
-    def __init__(self, config: Qwen3TTSTalkerConfig, device=None):
-        super().__init__()
-        # BC: "rope_type" was originally "type"
-        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
-            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
-        else:
-            self.rope_type = "default"
-        self.max_seq_len_cached = config.max_position_embeddings
-        self.original_max_seq_len = config.max_position_embeddings
-
-        self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
-
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-    @torch.no_grad()
-    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
-    def forward(self, x, position_ids):
-        # In contrast to other models, Qwen3TTSThinkerText has different position ids for the grids
-        # So we expand the inv_freq to shape (3, ...)
-        inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
-        position_ids_expanded = position_ids[:, :, None, :].float()  # shape (3, bs, 1, positions)
-
-        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
-            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
-            emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos() * self.attention_scaling
-            sin = emb.sin() * self.attention_scaling
-
-        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
-
-class Qwen3TTSRotaryEmbedding(nn.Module):
-    def __init__(self, config: Qwen3TTSConfig, device=None):
-        super().__init__()
-        # BC: "rope_type" was originally "type"
-        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
-            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
-        else:
-            self.rope_type = "default"
-        self.max_seq_len_cached = config.max_position_embeddings
-        self.original_max_seq_len = config.max_position_embeddings
-
-        self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
-
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-    @torch.no_grad()
-    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
-    def forward(self, x, position_ids):
-        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
-        position_ids_expanded = position_ids[:, None, :].float()
-
-        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
-            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
-            emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos() * self.attention_scaling
-            sin = emb.sin() * self.attention_scaling
-
-        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
-
-
-@use_kernel_forward_from_hub("RMSNorm")
-class Qwen3TTSRMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
-        """
-        Qwen3TTSRMSNorm is equivalent to T5LayerNorm
-        """
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states):
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states.to(input_dtype)
-
-    def extra_repr(self):
-        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
-
-def rotate_half(x):
-    """Rotates half the hidden dims of the input."""
-    x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2 :]
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """
-    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
-    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
-    """
-    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
-    if n_rep == 1:
-        return hidden_states
-    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
-    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
-
-
-def eager_attention_forward(
-    module: nn.Module,
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    attention_mask: Optional[torch.Tensor],
-    scaling: float,
-    dropout: float = 0.0,
-    **kwargs,
-):
-    key_states = repeat_kv(key, module.num_key_value_groups)
-    value_states = repeat_kv(value, module.num_key_value_groups)
-
-    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
-    if attention_mask is not None:
-        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
-        attn_weights = attn_weights + causal_mask
-
-    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
-    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
-    attn_output = torch.matmul(attn_weights, value_states)
-    attn_output = attn_output.transpose(1, 2).contiguous()
-
-    return attn_output, attn_weights
-
-
-def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, mrope_interleaved=False, unsqueeze_dim=1):
-    """Applies Rotary Position Embedding with Multimodal Sections to the query and key tensors (https://qwenlm.github.io/blog/qwen2-vl/).
-
-    Explanation:
-        Multimodal 3D rotary position embedding is an extension to 1D rotary position embedding. The input embedding
-        sequence contains vision (images / videos) embedding and text embedding or just contains text embedding. For
-        vision embedding part, we apply rotary position embedding on temporal, height and width dimension separately.
-        Here we split the channel dimension to 3 chunks for the temporal, height and width rotary position embedding.
-        For text embedding part, we just apply 1D rotary position embedding. The three rotary position index (temporal,
-        height and width) of text embedding is always the same, so the text embedding rotary position embedding has no
-        difference with modern LLMs.
-
-    Args:
-        q (`torch.Tensor`): The query tensor.
-        k (`torch.Tensor`): The key tensor.
-        cos (`torch.Tensor`): The cosine part of the rotary embedding.
-        sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`):
-            The position indices of the tokens corresponding to the query and key tensors. For example, this can be
-            used to pass offsetted position ids when working with a KV-cache.
-        mrope_section(`List(int)`):
-            Multimodal rope section is for channel dimension of temporal, height and width in rope calculation.
-        unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
-    Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
-    """
-    if mrope_interleaved:
-
-        def apply_interleaved_rope(x, modality_num):
-            x_t = x[0].clone()
-            index_ranges = []
-            for i, n in enumerate(mrope_section[1:], 1):
-                beg_idx = i
-                end_idx = n * modality_num
-                index_ranges.append((beg_idx, end_idx))
-            for beg_idx, end_idx in index_ranges:
-                x_t[..., beg_idx:end_idx:modality_num] = x[beg_idx, ..., beg_idx:end_idx:modality_num]
-            return x_t
-
-        dim = cos.shape[-1]
-        modality_num = len(mrope_section)
-        cos = torch.cat([apply_interleaved_rope(cos[..., : dim // 2], modality_num)] * 2, dim=-1).unsqueeze(
-            unsqueeze_dim
-        )
-        sin = torch.cat([apply_interleaved_rope(sin[..., : dim // 2], modality_num)] * 2, dim=-1).unsqueeze(
-            unsqueeze_dim
-        )
-    else:
-        mrope_section = mrope_section * 2
-        cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(
-            unsqueeze_dim
-        )
-        sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(
-            unsqueeze_dim
-        )
-
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
-
-
-class Qwen3TTSTalkerAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
-
-    def __init__(self, config, layer_idx):
-        super().__init__()
-        self.config = config
-        self.layer_idx = layer_idx
-        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
-        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
-        self.scaling = self.head_dim**-0.5
-        self.attention_dropout = config.attention_dropout
-        self.is_causal = True
-
-        self.q_proj = nn.Linear(
-            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.k_proj = nn.Linear(
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.v_proj = nn.Linear(
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.o_proj = nn.Linear(
-            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
-        )
-        self.q_norm = Qwen3TTSRMSNorm(
-            self.head_dim, eps=config.rms_norm_eps
-        )  # unlike olmo, only on the head dim!
-        self.k_norm = Qwen3TTSRMSNorm(
-            self.head_dim, eps=config.rms_norm_eps
-        )  # thus post q_norm does not need reshape
-        self.sliding_window = getattr(config, "sliding_window", None)
-        self.rope_scaling = config.rope_scaling
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        position_embeddings: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: Optional[torch.Tensor],
-        past_key_values: Optional[Cache] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        **kwargs: Unpack[FlashAttentionKwargs],
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
-        input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, self.head_dim)
-
-        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
-        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
-        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-
-        cos, sin = position_embeddings
-        query_states, key_states = apply_multimodal_rotary_pos_emb(
-            query_states, key_states, cos, sin, self.rope_scaling["mrope_section"], self.rope_scaling["interleaved"]
-        )
-
-        if past_key_values is not None:
-            # sin and cos are specific to RoPE models; cache_position needed for the static cache
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
-        attention_interface: Callable = eager_attention_forward
-        if self.config._attn_implementation != "eager":
-            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
-
-        attn_output, attn_weights = attention_interface(
-            self,
-            query_states,
-            key_states,
-            value_states,
-            attention_mask,
-            dropout=0.0 if not self.training else self.attention_dropout,
-            scaling=self.scaling,
-            sliding_window=self.sliding_window,  # diff with Llama
-            **kwargs,
-        )
-
-        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
-        attn_output = self.o_proj(attn_output)
-        return attn_output, attn_weights
-
-
-class Qwen3TTSTalkerResizeMLP(nn.Module):
-    def __init__(self, input_size: int, intermediate_size: int, output_size: int, act: str, bias=False):
-        super().__init__()
-        self.linear_fc1 = nn.Linear(input_size, intermediate_size, bias=bias)
-        self.linear_fc2 = nn.Linear(intermediate_size, output_size, bias=bias)
-        self.act_fn = ACT2FN[act]
-
-    def forward(self, hidden_state):
-        return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))
-
-
-@dataclass
-class Qwen3TTSTalkerCodePredictorOutputWithPast(ModelOutput):
-    r"""
-    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
-        Language modeling loss (for next-token prediction).
-    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
-        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
-    past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
-        Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
-        `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
-
-        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
-        `past_key_values` input) to speed up sequential decoding.
-    """
-
-    loss: Optional[torch.FloatTensor] = None
-    logits: torch.FloatTensor = None
-    past_key_values: Optional[list[torch.FloatTensor]] = None
-    hidden_states: Optional[tuple[torch.FloatTensor]] = None
-    attentions: Optional[tuple[torch.FloatTensor]] = None
-    generation_steps: Optional[int] = None
-
-
-class Qwen3TTSTalkerTextMLP(nn.Module):
-    def __init__(self, config, intermediate_size=None):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_act]
-
-    def forward(self, x):
-        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
-        return down_proj
-
-
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
-    """Applies Rotary Position Embedding to the query and key tensors.
-
-    Args:
-        q (`torch.Tensor`): The query tensor.
-        k (`torch.Tensor`): The key tensor.
-        cos (`torch.Tensor`): The cosine part of the rotary embedding.
-        sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`, *optional*):
-            Deprecated and unused.
-        unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
-    Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
-    """
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
-
-
-class Qwen3TTSAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
-
-    def __init__(self, config: Qwen3TTSConfig, layer_idx: int):
-        super().__init__()
-        self.config = config
-        self.layer_idx = layer_idx
-        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
-        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
-        self.scaling = self.head_dim**-0.5
-        self.attention_dropout = config.attention_dropout
-        self.is_causal = True
-
-        self.q_proj = nn.Linear(
-            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.k_proj = nn.Linear(
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.v_proj = nn.Linear(
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.o_proj = nn.Linear(
-            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
-        )
-        self.q_norm = Qwen3TTSRMSNorm(self.head_dim, eps=config.rms_norm_eps)  # unlike olmo, only on the head dim!
-        self.k_norm = Qwen3TTSRMSNorm(
-            self.head_dim, eps=config.rms_norm_eps
-        )  # thus post q_norm does not need reshape
-        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        position_embeddings: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: Optional[torch.Tensor],
-        past_key_values: Optional[Cache] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        **kwargs: Unpack[FlashAttentionKwargs],
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
-        input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, self.head_dim)
-
-        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
-        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
-        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-
-        cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
-        if past_key_values is not None:
-            # sin and cos are specific to RoPE models; cache_position needed for the static cache
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
-        attention_interface: Callable = eager_attention_forward
-        if self.config._attn_implementation != "eager":
-            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
-
-        attn_output, attn_weights = attention_interface(
-            self,
-            query_states,
-            key_states,
-            value_states,
-            attention_mask,
-            dropout=0.0 if not self.training else self.attention_dropout,
-            scaling=self.scaling,
-            sliding_window=self.sliding_window,  # diff with Llama
-            **kwargs,
-        )
-
-        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
-        attn_output = self.o_proj(attn_output)
-        return attn_output, attn_weights
-
-
-class Qwen3TTSDecoderLayer(GradientCheckpointingLayer):
-    def __init__(self, config: Qwen3TTSConfig, layer_idx: int):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-
-        self.self_attn = Qwen3TTSAttention(config=config, layer_idx=layer_idx)
-
-        self.mlp = Qwen3TTSTalkerTextMLP(config)
-        self.input_layernorm = Qwen3TTSRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = Qwen3TTSRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.attention_type = config.layer_types[layer_idx]
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Cache] = None,
-        output_attentions: Optional[bool] = False,
-        use_cache: Optional[bool] = False,
-        cache_position: Optional[torch.LongTensor] = None,
-        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
-        **kwargs: Unpack[FlashAttentionKwargs],
-    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
-        residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
-
-        # Self Attention
-        hidden_states, self_attn_weights = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            output_attentions=output_attentions,
-            use_cache=use_cache,
-            cache_position=cache_position,
-            position_embeddings=position_embeddings,
-            **kwargs,
-        )
-        hidden_states = residual + hidden_states
-
-        # Fully Connected
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
-
-        outputs = (hidden_states,)
-        if output_attentions:
-            outputs += (self_attn_weights,)
-
-        return outputs
-
-
-class Qwen3TTSTalkerCodePredictorModel(Qwen3TTSPreTrainedModel):
-    config_class = Qwen3TTSTalkerCodePredictorConfig
-    base_model_prefix = "talker.code_predictor.model"
-
-    def __init__(self, config: Qwen3TTSTalkerCodePredictorConfig, embedding_dim: int):
-        super().__init__(config)
-        self.padding_idx = config.pad_token_id
-        self.vocab_size = config.vocab_size
-        self.layers = nn.ModuleList(
-            [Qwen3TTSDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
-        )
-        self.norm = Qwen3TTSRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rotary_emb = Qwen3TTSRotaryEmbedding(config=config)
-        self.gradient_checkpointing = False
-        self.has_sliding_layers = "sliding_attention" in self.config.layer_types
-        self.codec_embedding = nn.ModuleList(
-            [nn.Embedding(config.vocab_size, embedding_dim) for _ in range(config.num_code_groups - 1)]
-        )
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.codec_embedding
-
-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
-
-    @can_return_tuple
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        position_ids=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        use_cache=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        cache_position=None,
-        generation_steps=None,
-        **flash_attn_kwargs,
-    ) -> BaseModelOutputWithPast:
-        if input_ids is not None:
-            raise ValueError("`input_ids` is expected to be `None`")
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-
-        if self.gradient_checkpointing and self.training and use_cache:
-            logger.warning_once(
-                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
-            )
-            use_cache = False
-
-        # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache
-        if not isinstance(past_key_values, (type(None), Cache)):
-            raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
-
-        if use_cache and past_key_values is None:
-            past_key_values = DynamicCache()
-
-        if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(
-                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
-            )
-
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0)
-
-        # It may already have been prepared by e.g. `generate`
-        if not isinstance(causal_mask_mapping := attention_mask, dict):
-            # Prepare mask arguments
-            mask_kwargs = {
-                "config": self.config,
-                "input_embeds": inputs_embeds,
-                "attention_mask": attention_mask,
-                "cache_position": cache_position,
-                "past_key_values": past_key_values,
-            }
-            # Create the masks
-            causal_mask_mapping = {
-                "full_attention": create_causal_mask(**mask_kwargs),
-            }
-            # The sliding window alternating layers are not always activated depending on the config
-            if self.has_sliding_layers:
-                causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
-
-        hidden_states = inputs_embeds
-
-        # create position embeddings to be shared across the decoder layers
-        position_embeddings = self.rotary_emb(hidden_states, position_ids)
-
-        # decoder layers
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attns = () if output_attentions else None
-
-        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
-            if output_hidden_states:
-                all_hidden_states += (hidden_states,)
-
-            layer_outputs = decoder_layer(
-                hidden_states,
-                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
-                position_ids=position_ids,
-                past_key_values=past_key_values,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                cache_position=cache_position,
-                position_embeddings=position_embeddings,
-                **flash_attn_kwargs,
-            )
-
-            hidden_states = layer_outputs[0]
-
-            if output_attentions:
-                all_self_attns += (layer_outputs[1],)
-
-        hidden_states = self.norm(hidden_states)
-
-        # add hidden states from the last decoder layer
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=past_key_values if use_cache else None,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
-        )
-
-
-class Qwen3TTSTalkerCodePredictorModelForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin):
-    _tied_weights_keys = ["lm_head.weight"]
-    _tp_plan = {"lm_head": "colwise_rep"}
-    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
-    config_class = Qwen3TTSTalkerCodePredictorConfig
-    base_model_prefix = "talker.code_predictor"
-
-    def __init__(self, config: Qwen3TTSTalkerCodePredictorConfig, talker_config: Qwen3TTSTalkerConfig):
-        super().__init__(config)
-        self.model = Qwen3TTSTalkerCodePredictorModel(config, talker_config.hidden_size)
-        self.vocab_size = config.vocab_size
-        self.lm_head = nn.ModuleList(
-            [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(config.num_code_groups - 1)]
-        )
-
-        if config.hidden_size != talker_config.hidden_size:
-            self.small_to_mtp_projection = torch.nn.Linear(talker_config.hidden_size, config.hidden_size, bias=True)
-        else:
-            self.small_to_mtp_projection = torch.nn.Identity()
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.model.get_input_embeddings()
-
-    def set_input_embeddings(self, value):
-        self.model.embed_tokens = value
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, new_embeddings):
-        self.lm_head = new_embeddings
-
-    def set_decoder(self, decoder):
-        self.model = decoder
-
-    def get_decoder(self):
-        return self.model
-    
-    def forward_finetune(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        position_ids=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        labels=None,
-        use_cache=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        cache_position=None,
-        generation_steps=None,
-        **kwargs,
-    ) -> CausalLMOutputWithPast:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-
-        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-        outputs: BaseModelOutputWithPast = self.model(
-            input_ids=None,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            cache_position=cache_position,
-            **kwargs,
-        )
-
-        hidden_states = outputs.last_hidden_state
-        
-        logits = []
-        for i in range(1, self.config.num_code_groups):
-            logits.append(self.lm_head[i-1](hidden_states[:, i]))
-        logits = torch.stack(logits, dim=1)
-
-        loss = None
-        if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
-
-        return Qwen3TTSTalkerCodePredictorOutputWithPast(
-            loss=loss,
-            logits=logits
-        )
-
-    @can_return_tuple
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        position_ids=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        labels=None,
-        use_cache=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        cache_position=None,
-        generation_steps=None,
-        **kwargs,
-    ) -> CausalLMOutputWithPast:
-        r"""
-        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
-            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
-            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-        """
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-
-        # Prefill stage
-        if inputs_embeds is not None and inputs_embeds.shape[1] > 1:
-            generation_steps = inputs_embeds.shape[1] - 2  # hidden & layer 0
-        # Generation stage
-        else:
-            inputs_embeds = self.model.get_input_embeddings()[generation_steps - 1](input_ids)
-        inputs_embeds = self.small_to_mtp_projection(inputs_embeds)
-
-        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-        outputs: BaseModelOutputWithPast = self.model(
-            input_ids=None,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            cache_position=cache_position,
-            **kwargs,
-        )
-
-        hidden_states = outputs.last_hidden_state
-        logits = self.lm_head[generation_steps](hidden_states)
-
-        loss = None
-        if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
-
-        return Qwen3TTSTalkerCodePredictorOutputWithPast(
-            loss=loss,
-            logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-            generation_steps=generation_steps + 1,
-        )
-
-    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False, num_new_tokens=1):
-        model_kwargs = super()._update_model_kwargs_for_generation(
-            outputs, model_kwargs, is_encoder_decoder, num_new_tokens
-        )
-        model_kwargs["generation_steps"] = outputs.generation_steps
-        return model_kwargs
-
-
-@dataclass
-class Qwen3TTSTalkerOutputWithPast(ModelOutput):
-    r"""
-    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
-        Language modeling loss (for next-token prediction).
-    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
-        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
-    past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
-        Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
-        `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
-
-        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
-        `past_key_values` input) to speed up sequential decoding.
-    """
-
-    loss: Optional[torch.FloatTensor] = None
-    logits: Optional[torch.FloatTensor] = None
-    past_key_values: Optional[list[torch.FloatTensor]] = None
-    hidden_states: Optional[tuple[torch.FloatTensor]] = None
-    attentions: Optional[tuple[torch.FloatTensor]] = None
-    past_hidden: Optional[torch.FloatTensor] = None
-    generation_step: Optional[int] = None
-    trailing_text_hidden: Optional[torch.FloatTensor] = None
-    tts_pad_embed: Optional[torch.FloatTensor] = None
-
-
-class Qwen3TTSTalkerDecoderLayer(GradientCheckpointingLayer):
-    def __init__(self, config, layer_idx):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.self_attn = Qwen3TTSTalkerAttention(config, layer_idx)
-
-        self.mlp = Qwen3TTSTalkerTextMLP(config, intermediate_size=config.intermediate_size)
-
-        self.input_layernorm = Qwen3TTSRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = Qwen3TTSRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[tuple[torch.Tensor]] = None,
-        output_attentions: Optional[bool] = False,
-        use_cache: Optional[bool] = False,
-        cache_position: Optional[torch.LongTensor] = None,
-        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
-        **kwargs: Unpack[FlashAttentionKwargs],
-    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
-        """
-        Args:
-            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
-            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
-                `(batch, sequence_length)` where padding elements are indicated by 0.
-            output_attentions (`bool`, *optional*):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
-            use_cache (`bool`, *optional*):
-                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
-                (see `past_key_values`).
-            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
-            cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
-                Indices depicting the position of the input sequence tokens in the sequence.
-            position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
-                Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
-                with `head_dim` being the embedding dimension of each attention head.
-            kwargs (`dict`, *optional*):
-                Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
-                into the model
-        """
-
-        residual = hidden_states
-
-        hidden_states = self.input_layernorm(hidden_states)
-
-        # Self Attention
-        hidden_states, self_attn_weights = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            output_attentions=output_attentions,
-            use_cache=use_cache,
-            cache_position=cache_position,
-            position_embeddings=position_embeddings,
-            **kwargs,
-        )
-        hidden_states = residual + hidden_states
-
-        # Fully Connected
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-
-        hidden_states = self.mlp(hidden_states)
-
-        hidden_states = residual + hidden_states
-
-        outputs = (hidden_states,)
-
-        if output_attentions:
-            outputs += (self_attn_weights,)
-
-        return outputs
-
-
-class Qwen3TTSTalkerModel(Qwen3TTSTalkerTextPreTrainedModel):
-    config_class = Qwen3TTSTalkerConfig
-    base_model_prefix = "talker.model"
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.padding_idx = config.pad_token_id
-        self.vocab_size = config.vocab_size
-        self.layers = nn.ModuleList(
-            [Qwen3TTSTalkerDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
-        )
-        self.norm = Qwen3TTSRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rotary_emb = Qwen3TTSTalkerRotaryEmbedding(config)
-        self.gradient_checkpointing = False
-        self.codec_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.text_embedding = nn.Embedding(config.text_vocab_size, config.text_hidden_size)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.codec_embedding
-    
-    def get_text_embeddings(self):
-        return self.text_embedding
-
-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
-
-    @can_return_tuple
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[list[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
-    ) -> BaseModelOutputWithPast:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-
-        if self.gradient_checkpointing and self.training:
-            if use_cache:
-                logger.warning_once(
-                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
-                )
-                use_cache = False
-
-        if use_cache and past_key_values is None:
-            past_key_values = DynamicCache()
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
-
-        if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(
-                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
-            )
-
-        # the hard coded `3` is for temporal, height and width.
-        if position_ids is None:
-            position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
-        elif position_ids.ndim == 2:
-            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
-
-        if position_ids.ndim == 3 and position_ids.shape[0] == 4:
-            text_position_ids = position_ids[0]
-            position_ids = position_ids[1:]
-        else:
-            text_position_ids = position_ids[0]
-        
-        mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask
-        causal_mask = mask_function(
-            config=self.config,
-            input_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            cache_position=cache_position,
-            past_key_values=past_key_values,
-            position_ids=text_position_ids,
-        )
-
-        hidden_states = inputs_embeds
-
-        # create position embeddings to be shared across the decoder layers
-        position_embeddings = self.rotary_emb(hidden_states, position_ids)
-
-        # decoder layers
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attns = () if output_attentions else None
-
-        for decoder_layer in self.layers:
-            if output_hidden_states:
-                all_hidden_states += (hidden_states,)
-
-            layer_outputs = decoder_layer(
-                hidden_states,
-                attention_mask=causal_mask,
-                position_ids=text_position_ids,
-                past_key_values=past_key_values,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                cache_position=cache_position,
-                position_embeddings=position_embeddings,
-                **flash_attn_kwargs,
-            )
-
-            hidden_states = layer_outputs[0]
-
-            if output_attentions:
-                all_self_attns += (layer_outputs[1],)
-
-        hidden_states = self.norm(hidden_states)
-
-        # add hidden states from the last decoder layer
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=past_key_values,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
-        )
-
-
-class Qwen3TTSTalkerForConditionalGeneration(Qwen3TTSTalkerTextPreTrainedModel, GenerationMixin):
-    _tied_weights_keys = ["lm_head.weight"]
-    _tp_plan = {"lm_head": "colwise_rep"}
-    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
-    config_class = Qwen3TTSTalkerConfig
-    base_model_prefix = "talker"
-
-    def __init__(self, config: Qwen3TTSTalkerConfig):
-        super().__init__(config)
-        self.model = Qwen3TTSTalkerModel(config)
-        self.vocab_size = config.vocab_size
-        self.text_projection = Qwen3TTSTalkerResizeMLP(
-            config.text_hidden_size, config.text_hidden_size, config.hidden_size, config.hidden_act, bias=True
-        )
-
-        self.codec_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.code_predictor = Qwen3TTSTalkerCodePredictorModelForConditionalGeneration(
-            config=config.code_predictor_config,
-            talker_config=config
-        )
-        self.rope_deltas = None
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-        # TODO: hack, modular cannot inherit multiple classes
-
-    def get_input_embeddings(self):
-        return self.model.get_input_embeddings()
-
-    def get_text_embeddings(self):
-        return self.model.get_text_embeddings()
-    
-    def set_input_embeddings(self, value):
-        self.model.embed_tokens = value
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, new_embeddings):
-        self.lm_head = new_embeddings
-
-    def set_decoder(self, decoder):
-        self.model = decoder
-
-    def get_decoder(self):
-        return self.model
-    
-    def forward_sub_talker_finetune(self, codec_ids, talker_hidden_states):
-        assert len(codec_ids.shape) == 2
-        assert len(talker_hidden_states.shape) == 2
-        assert codec_ids.shape[0] == talker_hidden_states.shape[0]
-        assert talker_hidden_states.shape[1] == self.config.hidden_size
-        assert codec_ids.shape[1] == self.config.num_code_groups
-
-        sub_talker_inputs_embeds = [talker_hidden_states.unsqueeze(1)]
-
-        for i in range(self.config.num_code_groups - 1):
-            if i == 0:
-                sub_talker_inputs_embeds.append(self.get_input_embeddings()(codec_ids[:, :1]))
-            else:
-                sub_talker_inputs_embeds.append(self.code_predictor.get_input_embeddings()[i-1](codec_ids[:, i:i+1]))
-        sub_talker_inputs_embeds = torch.cat(sub_talker_inputs_embeds, dim=1)
-        
-        sub_talker_outputs = self.code_predictor.forward_finetune(inputs_embeds=sub_talker_inputs_embeds,
-                                                                 labels=codec_ids[:, 1:])
-        
-        sub_talker_logits = sub_talker_outputs.logits
-        sub_talker_loss = sub_talker_outputs.loss
-        return sub_talker_logits, sub_talker_loss
-
-    @can_return_tuple
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        position_ids=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        labels=None,
-        use_cache=None,
-        output_attentions=None,
-        output_hidden_states=None,
-        cache_position=None,
-        past_hidden=None,
-        trailing_text_hidden=None,
-        tts_pad_embed=None,
-        generation_step=None,
-        subtalker_dosample=None,
-        subtalker_top_p=None,
-        subtalker_top_k=None,
-        subtalker_temperature=None,
-        **kwargs,
-    ) -> CausalLMOutputWithPast:
-        r"""
-        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
-            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
-            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-        ```"""
-        # Prefill
-        if inputs_embeds is not None and inputs_embeds.shape[1] > 1:
-            generation_step = -1
-            codec_ids = None
-        # Generate
-        else:
-            last_id_hidden = self.get_input_embeddings()(input_ids)
-            predictor_result = self.code_predictor.generate(
-                inputs_embeds=torch.cat((past_hidden, last_id_hidden), dim=1),
-                max_new_tokens=self.config.num_code_groups - 1,
-                do_sample=subtalker_dosample,
-                top_p=subtalker_top_p,
-                top_k=subtalker_top_k,
-                temperature=subtalker_temperature,
-                output_hidden_states=True,
-                return_dict_in_generate=True,
-            )
-            codec_ids = torch.cat((input_ids, predictor_result.sequences), dim=-1)
-            codec_hiddens = torch.cat(
-                [last_id_hidden]
-                + [self.code_predictor.get_input_embeddings()[i](predictor_result.sequences[..., i:i+1]) for i in range(self.config.num_code_groups - 1)],
-                dim=1,
-            )
-            inputs_embeds = codec_hiddens.sum(1, keepdim=True)
-
-            if generation_step < trailing_text_hidden.shape[1]:
-                inputs_embeds = inputs_embeds + trailing_text_hidden[:, generation_step].unsqueeze(1)
-            else:
-                inputs_embeds = inputs_embeds + tts_pad_embed
-        if attention_mask is not None:
-            if (
-                cache_position is None
-                or (cache_position is not None and cache_position[0] == 0)
-                or self.rope_deltas is None
-            ):
-                delta0 = (1 - attention_mask).sum(dim=-1).unsqueeze(1)
-                position_ids, rope_deltas = self.get_rope_index(
-                    attention_mask,
-                )
-                rope_deltas = rope_deltas - delta0
-                self.rope_deltas = rope_deltas
-            else:
-                batch_size, seq_length = input_ids.shape
-                delta = cache_position[0] + self.rope_deltas if cache_position is not None else 0
-                position_ids = torch.arange(seq_length, device=input_ids.device)
-                position_ids = position_ids.view(1, -1).expand(batch_size, -1)
-                position_ids = position_ids.add(delta)
-                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
-
-        outputs: BaseModelOutputWithPast = self.model(
-            input_ids=None,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            cache_position=cache_position,
-            **kwargs,
-        )
-
-        hidden_states = outputs.last_hidden_state
-        logits = self.codec_head(hidden_states)
-
-        loss = None
-        if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
-
-
-        return Qwen3TTSTalkerOutputWithPast(
-            loss=loss,
-            logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=(outputs.hidden_states, codec_ids),
-            attentions=outputs.attentions,
-            past_hidden=hidden_states[:, -1:, :],
-            generation_step=generation_step + 1,
-            trailing_text_hidden=trailing_text_hidden,
-            tts_pad_embed=tts_pad_embed,
-        )
-
-    def get_rope_index(
-        self,
-        attention_mask: Optional[torch.Tensor] = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Calculate the 3D rope index based on image and video's temporal, height and width in LLM.
-
-        Explanation:
-            Each embedding sequence contains vision embedding and text embedding or just contains text embedding.
-
-            For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.
-            Examples:
-                input_ids: [T T T T T], here T is for text.
-                temporal position_ids: [0, 1, 2, 3, 4]
-                height position_ids: [0, 1, 2, 3, 4]
-                width position_ids: [0, 1, 2, 3, 4]
-
-            For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part
-            and 1D rotary position embedding for text part.
-            Examples:
-                Temporal (Time): 3 patches, representing different segments of the video in time.
-                Height: 2 patches, dividing each frame vertically.
-                Width: 2 patches, dividing each frame horizontally.
-                We also have some important parameters:
-                fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second.
-                interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs.
-                input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.
-                text temporal position_ids: [101, 102, 103, 104, 105]
-                text height position_ids: [101, 102, 103, 104, 105]
-                text width position_ids: [101, 102, 103, 104, 105]
-                Here we calculate the text start position_ids as the max vision position_ids plus 1.
-
-        Args:
-            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
-                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
-                it.
-            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
-                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
-                - 1 for tokens that are **not masked**,
-                - 0 for tokens that are **masked**.
-
-        Returns:
-            position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
-            mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
-        """
-        mrope_position_deltas = []
-
-        position_ids = attention_mask.float().cumsum(-1) - 1
-        position_ids.masked_fill_(attention_mask == 0, 1)
-        position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
-        max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
-        mrope_position_deltas = max_position_ids + 1 - torch.sum(attention_mask, dim=-1, keepdim=True)
-
-        return position_ids, mrope_position_deltas
-
-    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False, num_new_tokens=1):
-        model_kwargs = super()._update_model_kwargs_for_generation(
-            outputs, model_kwargs, is_encoder_decoder, num_new_tokens
-        )
-        model_kwargs["past_hidden"] = outputs.past_hidden
-        model_kwargs["generation_step"] = outputs.generation_step
-        model_kwargs["trailing_text_hidden"] = outputs.trailing_text_hidden
-        model_kwargs["tts_pad_embed"] = outputs.tts_pad_embed
-        return model_kwargs
-
-
-class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin):
-    config_class = Qwen3TTSConfig
-
-    def __init__(self, config: Qwen3TTSConfig):
-        super().__init__(config)
-        self.config = config
-
-        self.talker = Qwen3TTSTalkerForConditionalGeneration(self.config.talker_config)
-
-        if config.tts_model_type == "base":
-            self.speaker_encoder = Qwen3TTSSpeakerEncoder(self.config.speaker_encoder_config)
-        else:
-            self.speaker_encoder = None
-
-        self.speech_tokenizer = None
-        self.generate_config = None
-
-        self.supported_speakers = self.config.talker_config.spk_id.keys()
-        self.supported_languages = ["auto"]
-        for language_id in self.config.talker_config.codec_language_id.keys():
-            if "dialect" not in language_id:
-                self.supported_languages.append(language_id)
-        
-        self.speaker_encoder_sample_rate = self.config.speaker_encoder_config.sample_rate
-        self.tokenizer_type = self.config.tokenizer_type
-        self.tts_model_size = self.config.tts_model_size
-        self.tts_model_type = self.config.tts_model_type
-
-        self.post_init()
-    
-    def load_speech_tokenizer(self, speech_tokenizer):
-        self.speech_tokenizer = speech_tokenizer
-    
-    def load_generate_config(self, generate_config):
-        self.generate_config = generate_config
-    
-    def get_supported_speakers(self):
-        return self.supported_speakers
-    
-    def get_supported_languages(self):
-        return self.supported_languages
-    
-    @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path,
-        *model_args,
-        config=None,
-        cache_dir=None,
-        ignore_mismatched_sizes=False,
-        force_download=False,
-        local_files_only=False,
-        token=None,
-        revision="main",
-        use_safetensors=None,
-        weights_only=True,
-        **kwargs,
-    ):
-        model = super().from_pretrained(
-            pretrained_model_name_or_path,
-            *model_args,
-            config=config,
-            cache_dir=cache_dir,
-            ignore_mismatched_sizes=ignore_mismatched_sizes,
-            force_download=force_download,
-            local_files_only=local_files_only,
-            token=token,
-            revision=revision,
-            use_safetensors=use_safetensors,
-            weights_only=weights_only,
-            **kwargs,
-        )
-        speech_tokenizer_path = cached_file(
-            pretrained_model_name_or_path,
-            "speech_tokenizer/config.json",
-            subfolder=kwargs.pop("subfolder", None),
-            cache_dir=kwargs.pop("cache_dir", None),
-            force_download=kwargs.pop("force_download", False),
-            proxies=kwargs.pop("proxies", None),
-            resume_download=kwargs.pop("resume_download", None),
-            local_files_only=kwargs.pop("local_files_only", False),
-            token=kwargs.pop("use_auth_token", None),
-            revision=kwargs.pop("revision", None),
-        )
-        if speech_tokenizer_path is None:
-            raise ValueError(f"""{pretrained_model_name_or_path}/{speech_tokenizer_path} not exists""")
-        speech_tokenizer_dir = os.path.dirname(speech_tokenizer_path)
-        speech_tokenizer = Qwen3TTSTokenizer.from_pretrained(
-            speech_tokenizer_dir,
-            *model_args,
-            **kwargs,
-        )
-        model.load_speech_tokenizer(speech_tokenizer)
-
-        generate_config_path = cached_file(
-            pretrained_model_name_or_path,
-            "generation_config.json",
-            subfolder=kwargs.pop("subfolder", None),
-            cache_dir=kwargs.pop("cache_dir", None),
-            force_download=kwargs.pop("force_download", False),
-            proxies=kwargs.pop("proxies", None),
-            resume_download=kwargs.pop("resume_download", None),
-            local_files_only=kwargs.pop("local_files_only", False),
-            token=kwargs.pop("use_auth_token", None),
-            revision=kwargs.pop("revision", None),
-        )
-        with open(generate_config_path, "r", encoding="utf-8") as f:
-            generate_config = json.load(f)
-        model.load_generate_config(generate_config)
-
-        return model
-    
-    @torch.inference_mode()
-    def extract_speaker_embedding(self, audio, sr):
-        assert sr == 24000, "Only support 24kHz audio"
-        mels = mel_spectrogram(
-            torch.from_numpy(audio).unsqueeze(0), 
-            n_fft=1024, 
-            num_mels=128, 
-            sampling_rate=24000,
-            hop_size=256, 
-            win_size=1024, 
-            fmin=0, 
-            fmax=12000
-        ).transpose(1, 2)
-        speaker_embedding = self.speaker_encoder(mels.to(self.device).to(self.dtype))[0]
-        return speaker_embedding
-    
-    @torch.inference_mode()
-    def generate_speaker_prompt(
-        self,
-        voice_clone_prompt: list[dict]
-    ):
-        voice_clone_spk_embeds = []
-        for index in range(len(voice_clone_prompt['ref_spk_embedding'])):
-            ref_spk_embedding = voice_clone_prompt["ref_spk_embedding"][index].to(self.talker.device).to(self.talker.dtype)            
-            voice_clone_spk_embeds.append(ref_spk_embedding)
-        
-        return voice_clone_spk_embeds
-
-    def generate_icl_prompt(
-        self,
-        text_id: torch.Tensor,
-        ref_id: torch.Tensor,
-        ref_code: torch.Tensor,
-        tts_pad_embed: torch.Tensor,
-        tts_eos_embed: torch.Tensor,
-        non_streaming_mode: bool,
-    ):
-        # text embed (ref id + text id + eos) 1 T1 D
-        text_embed = self.talker.text_projection(
-            self.talker.get_text_embeddings()(torch.cat([ref_id, text_id], 
-                                                            dim=-1)))
-        text_embed = torch.cat([text_embed, tts_eos_embed], dim=1)
-        # codec embed (codec bos + codec) 1 T2 D
-        codec_embed = []
-        for i in range(self.talker.config.num_code_groups):
-            if i == 0:
-                codec_embed.append(self.talker.get_input_embeddings()(ref_code[:, :1]))
-            else:
-                codec_embed.append(self.talker.code_predictor.get_input_embeddings()[i-1](ref_code[:, i:i+1]))
-        codec_embed = torch.cat(codec_embed, dim=1).sum(1).unsqueeze(0)
-        codec_embed = torch.cat([self.talker.get_input_embeddings()(
-                                    torch.tensor(
-                                        [[
-                                            self.config.talker_config.codec_bos_id,
-                                        ]],
-                                        device=self.talker.device,
-                                        dtype=text_id.dtype,
-                                    )
-                                ), codec_embed], dim=1)
-        # compute lens
-        text_lens = text_embed.shape[1]
-        codec_lens = codec_embed.shape[1]
-        if non_streaming_mode:
-            icl_input_embed = text_embed + self.talker.get_input_embeddings()(
-                                                torch.tensor(
-                                                    [[
-                                                        self.config.talker_config.codec_pad_id,
-                                                    ] * text_lens],
-                                                    device=self.talker.device,
-                                                    dtype=text_id.dtype,
-                                                )
-                                            )
-            icl_input_embed = torch.cat([icl_input_embed, codec_embed + tts_pad_embed], dim=1)
-            return icl_input_embed, tts_pad_embed
-        else:
-            if text_lens > codec_lens:
-                return text_embed[:, :codec_lens] + codec_embed, text_embed[:, codec_lens:]
-            else:
-                text_embed = torch.cat([text_embed] + [tts_pad_embed] * (codec_lens - text_lens), dim=1)
-                return text_embed + codec_embed, tts_pad_embed
-
-    @torch.no_grad()
-    def generate(
-        self,
-        input_ids: Optional[list[torch.Tensor]] = None,
-        instruct_ids: Optional[list[torch.Tensor]] = None,
-        ref_ids: Optional[list[torch.Tensor]] = None,
-        voice_clone_prompt: list[dict] = None,
-        languages: list[str] = None,
-        speakers: list[str] = None,
-        non_streaming_mode = False,
-        max_new_tokens: int = 4096,
-        do_sample: bool = True,
-        top_k: int = 50,
-        top_p: float = 1.0,
-        temperature: float = 0.9,
-        subtalker_dosample: bool = True,
-        subtalker_top_k: int = 50,
-        subtalker_top_p: float = 1.0,
-        subtalker_temperature: float = 0.9,
-        eos_token_id: Optional[int] = None,
-        repetition_penalty: float = 1.05,
-        **kwargs,
-    ):
-        talker_kwargs = {
-            "max_new_tokens": max_new_tokens,
-            "min_new_tokens": 2,
-            "do_sample": do_sample,
-            "top_k": top_k,
-            "top_p": top_p,
-            "temperature": temperature,
-            "subtalker_dosample": subtalker_dosample, 
-            "subtalker_top_k": subtalker_top_k,
-            "subtalker_top_p": subtalker_top_p,
-            "subtalker_temperature": subtalker_temperature,
-            "eos_token_id": eos_token_id
-            if eos_token_id is not None
-            else self.config.talker_config.codec_eos_token_id,
-            "repetition_penalty": repetition_penalty,
-            "suppress_tokens": [
-                i
-                for i in range(self.config.talker_config.vocab_size - 1024, self.config.talker_config.vocab_size)
-                if i not in (self.config.talker_config.codec_eos_token_id,)
-            ],
-            "output_hidden_states": getattr(kwargs, "output_hidden_states", True),
-            "return_dict_in_generate": getattr(kwargs, "return_dict_in_generate", True)
-        }
-        
-        talker_input_embeds = [[] for _ in range(len(input_ids))]
-
-        voice_clone_spk_embeds = None
-        # voice clone speaker prompt generate
-        if voice_clone_prompt is not None:
-            voice_clone_spk_embeds = self.generate_speaker_prompt(voice_clone_prompt)
-        
-        # instruct text prompt generate
-        if instruct_ids is not None:
-            for index, instruct_id in enumerate(instruct_ids):
-                if instruct_id is not None:
-                    talker_input_embeds[index].append(self.talker.text_projection(
-                                                  self.talker.get_text_embeddings()(instruct_id)))
-
-        # tts text prompt generate
-        trailing_text_hiddens = []
-        if speakers is None:
-            speakers = [None] * len(input_ids)
-        for index, (input_id, language, speaker) in enumerate(zip(input_ids, languages, speakers)):
-            if voice_clone_spk_embeds is None:
-                if speaker == "" or speaker == None: # Instruct create speaker
-                    speaker_embed = None
-                else:
-                    if speaker.lower() not in self.config.talker_config.spk_id:
-                        raise NotImplementedError(f"Speaker {speaker} not implemented")
-                    else:
-                        spk_id = self.config.talker_config.spk_id[speaker.lower()]
-                        speaker_embed = self.talker.get_input_embeddings()(
-                                            torch.tensor(
-                                                spk_id,
-                                                device=self.talker.device,
-                                                dtype=input_id.dtype,
-                                            )
-                                        )
-            else:
-                if voice_clone_prompt["x_vector_only_mode"][index] or voice_clone_prompt["icl_mode"][index]:
-                    speaker_embed = voice_clone_spk_embeds[index]
-                else:
-                    speaker_embed = None
-
-            assert language is not None
-
-            if language.lower() == "auto":
-                language_id = None
-            else:
-                if language.lower() not in self.config.talker_config.codec_language_id:
-                    raise NotImplementedError(f"Language {language} not implemented")
-                else:
-                    language_id = self.config.talker_config.codec_language_id[language.lower()]
-            
-            if (language.lower() in ["chinese", "auto"] and \
-                   speaker != "" and speaker is not None and \
-                     self.config.talker_config.spk_is_dialect[speaker.lower()] != False):
-                dialect = self.config.talker_config.spk_is_dialect[speaker.lower()]
-                language_id = self.config.talker_config.codec_language_id[dialect]
-            
-            tts_bos_embed, tts_eos_embed, tts_pad_embed = self.talker.text_projection(
-                self.talker.get_text_embeddings()(
-                    torch.tensor(
-                        [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
-                        device=self.talker.device,
-                        dtype=input_id.dtype,
-                    )
-                )
-            ).chunk(3, dim=1)  # 3 * [1 1 d]
-            
-            # codec: tag and speaker
-            if language_id is None:
-                codec_prefill_list = [[
-                                        self.config.talker_config.codec_nothink_id,
-                                        self.config.talker_config.codec_think_bos_id,
-                                        self.config.talker_config.codec_think_eos_id,
-                                    ]]
-            else:
-                codec_prefill_list = [[
-                                        self.config.talker_config.codec_think_id,
-                                        self.config.talker_config.codec_think_bos_id,
-                                        language_id,
-                                        self.config.talker_config.codec_think_eos_id,
-                                    ]]
-
-            codec_input_emebdding_0 = self.talker.get_input_embeddings()(
-                                                    torch.tensor(
-                                                        codec_prefill_list,
-                                                        device=self.talker.device,
-                                                        dtype=input_id.dtype,
-                                                    )
-                                                )
-            codec_input_emebdding_1 = self.talker.get_input_embeddings()(
-                                                    torch.tensor(
-                                                        [[
-                                                            self.config.talker_config.codec_pad_id,
-                                                            self.config.talker_config.codec_bos_id,
-                                                        ]],
-                                                        device=self.talker.device,
-                                                        dtype=input_id.dtype,
-                                                    )
-                                                )
-            if speaker_embed is None:
-                codec_input_emebdding = torch.cat([codec_input_emebdding_0,
-                                                   codec_input_emebdding_1], dim=1)
-            else:
-                codec_input_emebdding = torch.cat([codec_input_emebdding_0,
-                                                   speaker_embed.view(1, 1, -1),
-                                                   codec_input_emebdding_1], dim=1)
-
-            # '<|im_start|>assistant\n我叫通义千问，是阿里云的开源大模型。<|im_end|>\n<|im_start|>assistant\n'
-
-            # <|im_start|>assistant\n
-            _talker_input_embed_role = self.talker.text_projection(
-                                        self.talker.get_text_embeddings()(input_id[:, :3])
-                                        )
-
-            # tts_pad * 4 + tts_bos
-            _talker_input_embed = torch.cat((tts_pad_embed.expand(-1, codec_input_emebdding.shape[1] - 2, -1),
-                                            tts_bos_embed,
-                                            ), dim=1) + codec_input_emebdding[:, :-1]
-
-            talker_input_embed = torch.cat((_talker_input_embed_role, _talker_input_embed), dim=1)
-
-            if voice_clone_prompt is not None and voice_clone_prompt["ref_code"] is not None and voice_clone_prompt["icl_mode"][index]:
-                icl_input_embed, trailing_text_hidden = self.generate_icl_prompt(
-                    text_id=input_id[:, 3:-5],
-                    ref_id=ref_ids[index][:, 3:-2],
-                    ref_code=voice_clone_prompt["ref_code"][index].to(self.talker.device),
-                    tts_pad_embed=tts_pad_embed,
-                    tts_eos_embed=tts_eos_embed,
-                    non_streaming_mode=non_streaming_mode,
-                )
-                talker_input_embed = torch.cat([talker_input_embed, icl_input_embed], dim=1)
-            else:
-                #  tts_text_first_token
-                talker_input_embed = torch.cat([talker_input_embed, 
-                                                self.talker.text_projection(self.talker.get_text_embeddings()(input_id[:, 3:4])) + codec_input_emebdding[:, -1:]], 
-                                                dim=1)
-                if non_streaming_mode:
-                    talker_input_embed = talker_input_embed[:, :-1] # 去掉原本放进去的text
-                    talker_input_embed = torch.cat([talker_input_embed,
-                                                    torch.cat((self.talker.text_projection(
-                                                        self.talker.get_text_embeddings()(input_id[:, 3:-5])
-                                                    ), tts_eos_embed), dim=1) + self.talker.get_input_embeddings()(
-                                                        torch.tensor(
-                                                            [[
-                                                                self.config.talker_config.codec_pad_id,
-                                                            ] * (input_id[:, 3:-5].shape[1] + 1)],
-                                                            device=self.talker.device,
-                                                            dtype=input_id.dtype,
-                                                        )
-                                                    ), 
-                                                    tts_pad_embed + self.talker.get_input_embeddings()(
-                                                        torch.tensor(
-                                                            [[
-                                                                self.config.talker_config.codec_bos_id,
-                                                            ]],
-                                                            device=self.talker.device,
-                                                            dtype=input_id.dtype,
-                                                        )
-                                                    ) 
-                                                    ], dim=1)
-                    trailing_text_hidden = tts_pad_embed
-                else:
-                    # 叫通义千问，是阿里云的开源大模型。
-                    trailing_text_hidden = torch.cat((self.talker.text_projection(
-                                                        self.talker.get_text_embeddings()(input_id[:, 4:-5])
-                                                    ), tts_eos_embed), dim=1)
-            talker_input_embeds[index].append(talker_input_embed)
-            trailing_text_hiddens.append(trailing_text_hidden)
-        
-        for index, talker_input_embed in enumerate(talker_input_embeds):
-            talker_input_embeds[index] = torch.cat([item for item in talker_input_embed if item is not None], dim=1)
-
-        # for batch inferquence
-        original_lengths = torch.tensor([t.shape[1] for t in talker_input_embeds])
-        # left padding for talker input embeds
-        sequences = [t.squeeze(0) for t in talker_input_embeds]
-        sequences_reversed = [t.flip(dims=[0]) for t in sequences]
-        padded_reversed = torch.nn.utils.rnn.pad_sequence(
-            sequences_reversed,
-            batch_first=True,
-            padding_value=0.0
-        )
-        talker_input_embeds = padded_reversed.flip(dims=[1])
-        # generate mask
-        batch_size, max_len = talker_input_embeds.shape[0], talker_input_embeds.shape[1]
-        indices = torch.arange(max_len).expand(batch_size, -1)
-        num_pads = max_len - original_lengths
-        talker_attention_mask = (indices >= num_pads.unsqueeze(1)).long().to(talker_input_embeds.device)
-        # padding trailing text hiddens
-        pad_embedding_vector = tts_pad_embed.squeeze()
-        sequences_to_pad = [t.squeeze(0) for t in trailing_text_hiddens]
-        trailing_text_original_lengths = [s.shape[0] for s in sequences_to_pad]
-        padded_hiddens = torch.nn.utils.rnn.pad_sequence(
-            sequences_to_pad,
-            batch_first=True,
-            padding_value=0.0
-        )
-        arange_tensor = torch.arange(max(trailing_text_original_lengths), 
-                                     device=padded_hiddens.device).expand(len(trailing_text_original_lengths), -1)
-        lengths_tensor = torch.tensor(trailing_text_original_lengths, device=padded_hiddens.device).unsqueeze(1)
-        padding_mask = arange_tensor >= lengths_tensor
-        padded_hiddens[padding_mask] = pad_embedding_vector
-        trailing_text_hiddens = padded_hiddens
-
-        # forward
-        talker_result = self.talker.generate(
-            inputs_embeds=talker_input_embeds,
-            attention_mask=talker_attention_mask,
-            trailing_text_hidden=trailing_text_hiddens,
-            tts_pad_embed=tts_pad_embed,
-            **talker_kwargs,
-        )
-
-        talker_codes = torch.stack([hid[-1] for hid in talker_result.hidden_states if hid[-1] is not None], dim=1)
-        talker_hidden_states = torch.cat([hid[0][-1][:, -1:] for hid in talker_result.hidden_states], dim=1)[:, :-1]
-        
-        first_codebook = talker_codes[:, :, 0]
-        is_stop_token = (first_codebook ==  self.config.talker_config.codec_eos_token_id)
-        stop_indices = torch.argmax(is_stop_token.int(), dim=1)
-        has_stop_token = is_stop_token.any(dim=1)
-        effective_lengths = torch.where(has_stop_token, stop_indices, talker_codes.shape[1])
-        
-        talker_codes_list = [talker_codes[i, :length, ] for i, length in enumerate(effective_lengths)]
-        talker_hidden_states_list = [talker_hidden_states[i, :length, :] for i, length in enumerate(effective_lengths)]
-        
-        return talker_codes_list, talker_hidden_states_list
-
-__all__ = [
-    "Qwen3TTSForConditionalGeneration",
-    "Qwen3TTSTalkerForConditionalGeneration",
-    "Qwen3TTSPreTrainedModel",
-    "Qwen3TTSTalkerModel",
-]

qwen_tts/core/models/processing_qwen3_tts.py

@@ -1,106 +0,0 @@
-# coding=utf-8
-# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# 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.
-from transformers.feature_extraction_utils import BatchFeature
-from transformers.processing_utils import ProcessingKwargs, ProcessorMixin
-
-
-class Qwen3TTSProcessorKwargs(ProcessingKwargs, total=False):
-    _defaults = {
-        "text_kwargs": {
-            "padding": False,
-            "padding_side": "left",
-        }
-    }
-
-class Qwen3TTSProcessor(ProcessorMixin):
-    r"""
-    Constructs a Qwen3TTS processor.
-
-    Args:
-        tokenizer ([`Qwen2TokenizerFast`], *optional*):
-            The text tokenizer.
-        chat_template (`Optional[str]`, *optional*):
-            The Jinja template to use for formatting the conversation. If not provided, the default chat template is used.
-    """
-
-    attributes = ["tokenizer"]
-    tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
-
-    def __init__(
-        self, tokenizer=None, chat_template=None
-    ):
-        super().__init__(tokenizer, chat_template=chat_template)
-
-    def __call__(self, text=None, **kwargs) -> BatchFeature:
-        """
-        Main method to prepare for the model one or several sequences(s) and audio(s). This method forwards the `text`
-        and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
-        the text. 
-
-        Args:
-            text (`str`, `List[str]`, `List[List[str]]`):
-                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
-                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
-                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
-        """
-
-        if text is None:
-            raise ValueError("You need to specify either a `text` input to process.")
-
-        output_kwargs = self._merge_kwargs(
-            Qwen3TTSProcessorKwargs,
-            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
-            **kwargs,
-        )
-        if not isinstance(text, list):
-            text = [text]
-
-        texts_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
-
-        return BatchFeature(
-            data={**texts_inputs},
-            tensor_type=kwargs.get("return_tensors"),
-        )
-
-    def batch_decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
-        refer to the docstring of this method for more information.
-        """
-        return self.tokenizer.batch_decode(*args, **kwargs)
-
-    def decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
-        the docstring of this method for more information.
-        """
-        return self.tokenizer.decode(*args, **kwargs)
-
-    def apply_chat_template(self, conversations, chat_template=None, **kwargs):
-        if isinstance(conversations[0], dict):
-            conversations = [conversations]
-        return super().apply_chat_template(conversations, chat_template, **kwargs)
-
-    @property
-    def model_input_names(self):
-        tokenizer_input_names = self.tokenizer.model_input_names
-        return list(
-            dict.fromkeys(
-                tokenizer_input_names
-            )
-        )
-
-
-__all__ = ["Qwen3TTSProcessor"]

qwen_tts/core/tokenizer_12hz/configuration_qwen3_tts_tokenizer_v2.py

@@ -1,172 +0,0 @@
-# coding=utf-8
-# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# 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.
-"""Qwen3TTSTokenizerV2 model configuration"""
-
-from transformers.configuration_utils import PretrainedConfig
-from transformers.utils import logging
-
-from transformers import MimiConfig
-
-
-logger = logging.get_logger(__name__)
-
-
-class Qwen3TTSTokenizerV2DecoderConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`Qwen3TTSTokenizerV2DecoderConfig`].
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-    Args:
-        codebook_size (`int`, *optional*, defaults to 2048):
-            Number of entries in each residual codebook used for acoustic token quantization.
-        hidden_size (`int`, *optional*, defaults to 1024):
-            Dimensionality of the hidden states and embeddings in the autoregressive transformer decoder.
-        max_position_embeddings (`int`, *optional*, defaults to 8000):
-            Maximum sequence length that the autoregressive decoder can handle. Determines positional embedding size.
-        rope_theta (`float`, *optional*, defaults to 10000.0):
-            The base period for rotary position embeddings (RoPE) applied to attention layers.
-        num_attention_heads (`int`, *optional*, defaults to 16):
-            Number of attention heads for each attention layer in the decoder.
-        num_key_value_heads (`int`, *optional*, defaults to 16):
-            Number of key and value attention heads used in grouped-query attention (if applicable).
-        attention_bias (`bool`, *optional*, defaults to `False`):
-            Whether to use bias in the attention projection layers.
-        sliding_window (`int`, *optional*, defaults to 72):
-            Window size for local attention mechanism, limiting attention context to improve efficiency.
-        intermediate_size (`int`, *optional*, defaults to 3072):
-            Dimensionality of the feed-forward (intermediate) layer in each transformer block.
-        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
-            The non-linear activation function used in the feed-forward layers. Supports `"silu"`, `"relu"`, `"gelu"`, etc.
-        layer_scale_initial_scale (`float`, *optional*, defaults to 0.01):
-            Initial value for LayerScale applied in transformer blocks, helping stabilize training.
-        rms_norm_eps (`float`, *optional*, defaults to 1e-5):
-            Epsilon value for RMS normalization layers to prevent division by zero.
-        num_hidden_layers (`int`, *optional*, defaults to 8):
-            Number of transformer blocks in the autoregressive decoder.
-        num_quantizers (`int`, *optional*, defaults to 16):
-            Number of residual vector quantizers used in the vocoder for fine-grained audio reconstruction.
-        upsample_rates (`Tuple[int]`, *optional*, defaults to `(8, 5, 4, 3)`):
-            Rate at which features are upsampled in the final waveform synthesis stage.
-        upsampling_ratios (`Tuple[int]`, *optional*, defaults to `(2, 2)`):
-            Ratios used in transposed convolutional layers to progressively upsample feature maps to waveform.
-        decoder_dim (`int`, *optional*, defaults to 1536):
-            Final dimensionality of the decoder's output before waveform generation.
-        attention_dropout (`float`, *optional*, defaults to 0.0):
-            Dropout probability applied to attention weights in the decoder.
-    """
-
-    def __init__(
-        self,
-        codebook_size=2048,
-        hidden_size=1024,
-        latent_dim=1024,
-        max_position_embeddings=8000,
-        rope_theta=10000,
-        num_attention_heads=16,
-        num_key_value_heads=16,
-        attention_bias=False,
-        sliding_window=72,
-        intermediate_size=3072,
-        hidden_act="silu",
-        layer_scale_initial_scale=0.01,
-        rms_norm_eps=1e-5,
-        num_hidden_layers=8,
-        num_quantizers=16,
-        upsample_rates=(8, 5, 4, 3),
-        upsampling_ratios=(2, 2),
-        decoder_dim=1536,
-        attention_dropout=0.0,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.codebook_size = codebook_size
-        self.hidden_size = hidden_size
-        self.latent_dim = latent_dim
-        self.max_position_embeddings = max_position_embeddings
-        self.rope_theta = rope_theta
-        self.num_attention_heads = num_attention_heads
-        self.num_key_value_heads = num_key_value_heads
-        self.attention_bias = attention_bias
-        self.sliding_window = sliding_window
-        self.intermediate_size = intermediate_size
-        self.hidden_act = hidden_act
-        self.layer_scale_initial_scale = layer_scale_initial_scale
-        self.rms_norm_eps = rms_norm_eps
-        self.num_hidden_layers = num_hidden_layers
-        self.num_quantizers = num_quantizers
-        self.upsample_rates = upsample_rates
-        self.upsampling_ratios = upsampling_ratios
-        self.decoder_dim = decoder_dim
-        self.attention_dropout = attention_dropout
-
-    @property
-    def layer_types(self):
-        """
-        All layer in code2wav should be sliding attention
-        """
-        return ["sliding_attention"] * self.num_hidden_layers
-
-
-class Qwen3TTSTokenizerV2Config(PretrainedConfig):
-    """
-    This is the configuration class to store the configuration of a [`Qwen3TTSTokenizerV2Config`]. It is used to instantiate a Qwen3TTSTokenizerV2Model
-    model according to the specified sub-models configurations, defining the model architecture.
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-    Args:
-        encoder_config (`dict`, *optional*): Configuration of the underlying encoder sub-model.
-        decoder_config (`dict`, *optional*): Configuration of the underlying decoder sub-model.
-    """
-
-    model_type = "qwen3_tts_tokenizer_12hz"
-    sub_configs = {
-        "encoder_config": MimiConfig,
-        "decoder_config": Qwen3TTSTokenizerV2DecoderConfig,
-    }
-
-    def __init__(
-        self,
-        encoder_config=None,
-        decoder_config=None,
-        encoder_valid_num_quantizers=16,
-        input_sample_rate=24000,
-        output_sample_rate=24000,
-        decode_upsample_rate=1920,
-        encode_downsample_rate=1920,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        if encoder_config is None:
-            encoder_config = {}
-            logger.info("encoder_config is None. Initializing encoder with default values")
-        if decoder_config is None:
-            decoder_config = {}
-            logger.info("decoder_config is None. Initializing decoder with default values")
-
-        self.encoder_config = MimiConfig(**encoder_config)
-        self.decoder_config = Qwen3TTSTokenizerV2DecoderConfig(**decoder_config)
-
-        self.encoder_valid_num_quantizers = encoder_valid_num_quantizers
-        self.input_sample_rate = input_sample_rate
-        self.output_sample_rate = output_sample_rate
-        self.decode_upsample_rate = decode_upsample_rate
-        self.encode_downsample_rate = encode_downsample_rate
-
-
-__all__ = ["Qwen3TTSTokenizerV2Config", "Qwen3TTSTokenizerV2DecoderConfig"]

qwen_tts/core/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py

@@ -1,1025 +0,0 @@
-# coding=utf-8
-# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# 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.
-"""PyTorch Qwen3TTSTokenizerV2 model."""
-
-import math
-from dataclasses import dataclass
-from typing import Callable, Optional, Union, List
-
-import numpy as np
-import torch
-from torch import nn
-from torch.nn import Parameter
-from torch.nn import functional as F
-from transformers import MimiConfig, MimiModel
-from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache, DynamicCache
-from transformers.integrations import use_kernel_forward_from_hub
-from transformers.masking_utils import (
-    create_causal_mask,
-    create_sliding_window_causal_mask,
-)
-from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
-from transformers.modeling_layers import GradientCheckpointingLayer
-from transformers.modeling_outputs import BaseModelOutputWithPast
-from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
-from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
-from transformers.processing_utils import Unpack
-from transformers.utils import ModelOutput, auto_docstring, logging
-from transformers.utils.deprecation import deprecate_kwarg
-from transformers.utils.generic import check_model_inputs
-
-from .configuration_qwen3_tts_tokenizer_v2 import (
-    Qwen3TTSTokenizerV2Config,
-    Qwen3TTSTokenizerV2DecoderConfig,
-)
-
-logger = logging.get_logger(__name__)
-
-
-@dataclass
-@auto_docstring
-class Qwen3TTSTokenizerV2EncoderOutput(ModelOutput):
-    r"""
-    audio_codes (`List[torch.LongTensor]`):
-        Discret code embeddings computed using `model.encode`, each tensor has shape (codes_length_i, num_quantizers).
-    """
-
-    audio_codes: List[torch.LongTensor] = None
-
-
-@dataclass
-@auto_docstring
-class Qwen3TTSTokenizerV2DecoderOutput(ModelOutput):
-    r"""
-    audio_values (`List[torch.FloatTensor]`):
-        Decoded audio values, obtained using the decoder part of Qwen3TTSTokenizerV1.
-        Each tensor has shape (segment_length_i).
-    """
-
-    audio_values: List[torch.FloatTensor] = None
-
-
-def rotate_half(x):
-    """Rotates half the hidden dims of the input."""
-    x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2 :]
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
-    """Applies Rotary Position Embedding to the query and key tensors.
-
-    Args:
-        q (`torch.Tensor`): The query tensor.
-        k (`torch.Tensor`): The key tensor.
-        cos (`torch.Tensor`): The cosine part of the rotary embedding.
-        sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`, *optional*):
-            Deprecated and unused.
-        unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
-    Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
-    """
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
-
-
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """
-    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
-    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
-    """
-    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
-    if n_rep == 1:
-        return hidden_states
-    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
-    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
-
-
-def eager_attention_forward(
-    module: nn.Module,
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    attention_mask: Optional[torch.Tensor],
-    scaling: float,
-    dropout: float = 0.0,
-    **kwargs,
-):
-    key_states = repeat_kv(key, module.num_key_value_groups)
-    value_states = repeat_kv(value, module.num_key_value_groups)
-
-    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
-    if attention_mask is not None:
-        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
-        attn_weights = attn_weights + causal_mask
-
-    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
-    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
-    attn_output = torch.matmul(attn_weights, value_states)
-    attn_output = attn_output.transpose(1, 2).contiguous()
-
-    return attn_output, attn_weights
-
-
-@auto_docstring
-class Qwen3TTSTokenizerV2DecoderPreTrainedModel(PreTrainedModel):
-    config: Qwen3TTSTokenizerV2DecoderConfig
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = True
-    _skip_keys_device_placement = "past_key_values"
-    _supports_flash_attn = True
-    _supports_sdpa = True
-    _can_compile_fullgraph = False
-    _supports_attention_backend = True
-
-
-class Qwen3TTSTokenizerV2CausalConvNet(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        dilation=1,
-        stride=1,
-        groups=1,
-    ):
-        super().__init__()
-        self.conv = nn.Conv1d(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride=stride,
-            dilation=dilation,
-            groups=groups,
-        )
-        self.stride = stride
-        self.kernel_size = (kernel_size - 1) * dilation + 1
-        self.dilation = dilation
-        self.padding = self.kernel_size - self.stride
-
-    def _get_extra_padding_for_conv1d(self, hidden_state: torch.Tensor) -> int:
-        length = hidden_state.shape[-1]
-        n_frames = (length - self.kernel_size + self.padding) / self.stride + 1
-        ideal_length = (math.ceil(n_frames) - 1) * self.stride + (self.kernel_size - self.padding)
-        return ideal_length - length
-
-    def forward(self, hidden_state):
-        extra_padding = self._get_extra_padding_for_conv1d(hidden_state)
-        hidden_state = F.pad(hidden_state, (self.padding, extra_padding), mode="constant", value=0)
-        return self.conv(hidden_state).contiguous()
-
-
-class Qwen3TTSTokenizerV2CausalTransConvNet(nn.Module):
-    def __init__(self, in_channels, out_channels, kernel_size, stride=1):
-        super().__init__()
-        self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride)
-
-        pad = kernel_size - stride
-        self.left_pad = math.ceil(pad)
-        self.right_pad = pad = self.left_pad
-
-    def forward(self, hidden_state):
-        hidden_state = self.conv(hidden_state)
-        hidden_state = hidden_state[..., self.left_pad : hidden_state.shape[-1] - self.right_pad]
-        return hidden_state.contiguous()
-
-
-class Qwen3TTSTokenizerV2ConvNeXtBlock(nn.Module):
-    def __init__(self, dim: int):
-        super().__init__()
-        self.dwconv = Qwen3TTSTokenizerV2CausalConvNet(
-            dim,
-            dim,
-            kernel_size=7,
-            groups=dim,
-            dilation=1,
-        )
-        self.norm = nn.LayerNorm(dim, eps=1e-6)
-        self.pwconv1 = nn.Linear(dim, 4 * dim)
-        self.act = nn.GELU()
-        self.pwconv2 = nn.Linear(4 * dim, dim)
-        self.gamma = nn.Parameter(1e-6 * torch.ones(dim))
-
-    def forward(self, hidden_states):
-        input = hidden_states
-
-        hidden_states = self.dwconv(hidden_states)
-        hidden_states = hidden_states.permute(0, 2, 1)
-        hidden_states = self.norm(hidden_states)
-        hidden_states = self.pwconv1(hidden_states)
-        hidden_states = self.act(hidden_states)
-        hidden_states = self.pwconv2(hidden_states)
-
-        hidden_states = self.gamma * hidden_states
-
-        hidden_states = hidden_states.permute(0, 2, 1)
-
-        hidden_states = input + hidden_states
-
-        return hidden_states
-
-
-class Qwen3TTSTokenizerV2DecoderRotatoryEmbedding(nn.Module):
-    inv_freq: torch.Tensor  # fix linting for `register_buffer`
-
-    def __init__(self, config: Qwen3TTSTokenizerV2DecoderConfig, device=None):
-        super().__init__()
-        # BC: "rope_type" was originally "type"
-        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
-            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
-        else:
-            self.rope_type = "default"
-        self.max_seq_len_cached = config.max_position_embeddings
-        self.original_max_seq_len = config.max_position_embeddings
-
-        self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
-
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-    @torch.no_grad()
-    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
-    def forward(self, x, position_ids):
-        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
-        position_ids_expanded = position_ids[:, None, :].float()
-
-        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
-            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
-            emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos() * self.attention_scaling
-            sin = emb.sin() * self.attention_scaling
-
-        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
-
-
-class Qwen3TTSTokenizerV2DecoderAttention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
-
-    def __init__(self, config: Qwen3TTSTokenizerV2DecoderConfig, layer_idx):
-        super().__init__()
-        self.config = config
-        self.layer_idx = layer_idx
-        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
-        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
-        self.scaling = self.head_dim**-0.5
-        self.attention_dropout = config.attention_dropout
-        self.is_causal = True
-
-        self.q_proj = nn.Linear(
-            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.k_proj = nn.Linear(
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.v_proj = nn.Linear(
-            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.o_proj = nn.Linear(
-            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
-        )
-        self.q_norm = nn.Identity()
-        self.k_norm = nn.Identity()
-        self.sliding_window = config.sliding_window
-
-    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        position_embeddings: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: Optional[torch.Tensor],
-        past_key_values: Optional[Cache] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        **kwargs: Unpack[FlashAttentionKwargs],
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
-        input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, self.head_dim)
-
-        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
-        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
-        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-
-        cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
-        if past_key_values is not None:
-            # sin and cos are specific to RoPE models; cache_position needed for the static cache
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
-        attention_interface: Callable = eager_attention_forward
-        if self.config._attn_implementation != "eager":
-            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
-
-        attn_output, attn_weights = attention_interface(
-            self,
-            query_states,
-            key_states,
-            value_states,
-            attention_mask,
-            dropout=0.0 if not self.training else self.attention_dropout,
-            scaling=self.scaling,
-            sliding_window=self.sliding_window,  # diff with Llama
-            **kwargs,
-        )
-
-        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
-        attn_output = self.o_proj(attn_output)
-        return attn_output, attn_weights
-
-
-class Qwen3TTSTokenizerV2DecoderMlp(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = config.intermediate_size
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_act]
-
-    def forward(self, x):
-        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
-        return down_proj
-
-
-@use_kernel_forward_from_hub("RMSNorm")
-class Qwen3TTSTokenizerV2DecoderRMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
-        """
-        Qwen3TTSTokenizerV2DecoderRMSNorm is equivalent to T5LayerNorm
-        """
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states.to(input_dtype)
-
-    def extra_repr(self):
-        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
-
-
-class Qwen3TTSTokenizerV2DecoderLayerScale(nn.Module):
-    """Layer scale from [Touvron et al 2021] (https://huggingface.co/papers/2103.17239).
-    This rescales diagonally the residual outputs close to 0, with a learnt scale.
-    """
-
-    def __init__(self, config):
-        super().__init__()
-        channels = config.hidden_size
-        initial_scale = config.layer_scale_initial_scale
-        self.scale = nn.Parameter(torch.full((channels,), initial_scale, requires_grad=True))
-
-    def forward(self, x: torch.Tensor):
-        return self.scale * x
-
-
-class Qwen3TTSTokenizerV2DecoderTransformerLayer(GradientCheckpointingLayer):
-    def __init__(self, config: Qwen3TTSTokenizerV2DecoderConfig, layer_idx):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.self_attn = Qwen3TTSTokenizerV2DecoderAttention(config, layer_idx)
-        self.mlp = Qwen3TTSTokenizerV2DecoderMlp(config)
-        self.input_layernorm = Qwen3TTSTokenizerV2DecoderRMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.post_attention_layernorm = Qwen3TTSTokenizerV2DecoderRMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.self_attn_layer_scale = Qwen3TTSTokenizerV2DecoderLayerScale(config)
-        self.mlp_layer_scale = Qwen3TTSTokenizerV2DecoderLayerScale(config)
-        self.attention_type = "sliding_attention"
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Cache] = None,
-        use_cache: Optional[bool] = False,
-        cache_position: Optional[torch.LongTensor] = None,
-        **kwargs,
-    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
-        """
-        Args:
-            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
-            attention_mask (`torch.FloatTensor`, *optional*):
-                attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
-                query_sequence_length, key_sequence_length)` if default attention is used.
-            output_attentions (`bool`, *optional*):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
-            use_cache (`bool`, *optional*):
-                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
-                (see `past_key_values`).
-            past_key_values (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
-            cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
-                Indices depicting the position of the input sequence tokens in the sequence
-            kwargs (`dict`, *optional*):
-                Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
-                into the model
-        """
-        residual = hidden_states
-
-        hidden_states = self.input_layernorm(hidden_states)
-
-        # Self Attention
-        hidden_states, _ = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            use_cache=use_cache,
-            cache_position=cache_position,
-            **kwargs,
-        )
-        hidden_states = residual + self.self_attn_layer_scale(hidden_states)
-
-        # Fully Connected
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + self.mlp_layer_scale(hidden_states)
-
-        return hidden_states
-
-
-@auto_docstring
-class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
-    _can_record_outputs = {
-        "hidden_states": Qwen3TTSTokenizerV2DecoderTransformerLayer,
-        "attentions": Qwen3TTSTokenizerV2DecoderAttention,
-    }
-
-    def __init__(self, config: Qwen3TTSTokenizerV2DecoderConfig):
-        super().__init__(config)
-        self.layers = nn.ModuleList(
-            [Qwen3TTSTokenizerV2DecoderTransformerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
-        )
-        self.norm = Qwen3TTSTokenizerV2DecoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rotary_emb = Qwen3TTSTokenizerV2DecoderRotatoryEmbedding(config=config)
-        self.gradient_checkpointing = False
-        self.has_sliding_layers = "sliding_attention" in self.config.layer_types
-        self.window_size = config.sliding_window
-
-        self.input_proj = nn.Linear(config.latent_dim, config.hidden_size)
-        self.output_proj = nn.Linear(config.hidden_size, config.latent_dim)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    @check_model_inputs()
-    @auto_docstring
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        position_ids=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        use_cache=None,
-        cache_position=None,
-        **kwargs,
-    ) -> BaseModelOutputWithPast:
-        if input_ids is not None:
-            raise ValueError("input_ids is not expected")
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
-        
-        inputs_embeds = self.input_proj(inputs_embeds)
-
-        if use_cache and past_key_values is None:
-            past_key_values = DynamicCache(config=self.config)
-
-        if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(
-                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
-            )
-
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0)
-
-        # It may already have been prepared by e.g. `generate`
-        if not isinstance(causal_mask_mapping := attention_mask, dict):
-            # Prepare mask arguments
-            mask_kwargs = {
-                "config": self.config,
-                "input_embeds": inputs_embeds,
-                "attention_mask": attention_mask,
-                "cache_position": cache_position,
-                "past_key_values": past_key_values,
-                "position_ids": position_ids,
-            }
-            # Create the masks
-            causal_mask_mapping = {
-                "full_attention": create_causal_mask(**mask_kwargs),
-            }
-            # The sliding window alternating layers are not always activated depending on the config
-            if self.has_sliding_layers:
-                causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
-
-        hidden_states = inputs_embeds
-
-        # create position embeddings to be shared across the decoder layers
-        position_embeddings = self.rotary_emb(hidden_states, position_ids)
-
-        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
-            hidden_states = decoder_layer(
-                hidden_states,
-                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
-                position_ids=position_ids,
-                past_key_values=past_key_values,
-                use_cache=use_cache,
-                cache_position=cache_position,
-                position_embeddings=position_embeddings,
-                **kwargs,
-            )
-
-        hidden_states = self.norm(hidden_states)
-        hidden_states = self.output_proj(hidden_states)
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=past_key_values if use_cache else None,
-        )
-
-
-class SnakeBeta(nn.Module):
-    """
-    A modified Snake function which uses separate parameters for the magnitude of the periodic components
-    Shape:
-        - Input: (B, C, T)
-        - Output: (B, C, T), same shape as the input
-    Parameters:
-        - alpha - trainable parameter that controls frequency
-        - beta - trainable parameter that controls magnitude
-    References:
-        - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
-        https://huggingface.co/papers/2006.08195
-    """
-
-    def __init__(self, in_features, alpha=1.0):
-        super().__init__()
-        self.in_features = in_features
-
-        # initialize alpha
-        self.alpha = Parameter(torch.zeros(in_features) * alpha)
-        self.beta = Parameter(torch.zeros(in_features) * alpha)
-
-        self.no_div_by_zero = 0.000000001
-
-    def forward(self, hidden_states):
-        """
-        Forward pass of the function.
-        Applies the function to the input elementwise.
-        SnakeBeta ∶= x + 1/b * sin^2 (xa)
-        """
-        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
-        beta = self.beta.unsqueeze(0).unsqueeze(-1)
-        alpha = torch.exp(alpha)
-        beta = torch.exp(beta)
-        hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(
-            torch.sin(hidden_states * alpha), 2
-        )
-
-        return hidden_states
-
-
-class Qwen3TTSTokenizerV2DecoderDecoderResidualUnit(nn.Module):
-    def __init__(self, dim: int = 16, dilation: int = 1):
-        super().__init__()
-
-        self.act1 = SnakeBeta(dim)
-        self.conv1 = Qwen3TTSTokenizerV2CausalConvNet(dim, dim, kernel_size=7, dilation=dilation)
-        self.act2 = SnakeBeta(dim)
-        self.conv2 = Qwen3TTSTokenizerV2CausalConvNet(dim, dim, kernel_size=1)
-
-    def forward(self, hidden_state):
-        residual = hidden_state
-
-        hidden_state = self.act1(hidden_state)
-        hidden_state = self.conv1(hidden_state)
-        hidden_state = self.act2(hidden_state)
-        hidden_state = self.conv2(hidden_state)
-        return hidden_state + residual
-
-
-class Qwen3TTSTokenizerV2DecoderDecoderBlock(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
-    def __init__(self, config: Qwen3TTSTokenizerV2DecoderConfig, layer_idx):
-        super().__init__(config)
-        in_dim = config.decoder_dim // 2**layer_idx
-        out_dim = config.decoder_dim // 2 ** (layer_idx + 1)
-        upsample_rate = config.upsample_rates[layer_idx]
-
-        block = [
-            SnakeBeta(in_dim),
-            Qwen3TTSTokenizerV2CausalTransConvNet(in_dim, out_dim, 2 * upsample_rate, upsample_rate),
-        ]
-
-        for dilation in (1, 3, 9):
-            block.append(Qwen3TTSTokenizerV2DecoderDecoderResidualUnit(out_dim, dilation))
-
-        self.block = nn.ModuleList(block)
-
-    def forward(self, hidden):
-        for block in self.block:
-            hidden = block(hidden)
-        return hidden
-
-
-class EuclideanCodebook(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        codebook_size: int,
-        epsilon: float = 1e-5,
-    ):
-        super().__init__()
-        self.dim = dim
-        self.codebook_size = codebook_size
-        self.epsilon = epsilon
-
-        self.cluster_usage = nn.Parameter(torch.ones(codebook_size))
-        self.embedding_sum = nn.Parameter(torch.zeros(codebook_size, dim))
-
-    def decode(self, codes: torch.Tensor) -> torch.Tensor:
-        embedding = self.embedding_sum / self.cluster_usage.clamp(min=self.epsilon)[:, None]
-        quantized = F.embedding(codes, embedding)
-        return quantized
-
-
-class VectorQuantization(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        codebook_size: int,
-        codebook_dim: Optional[int] = None,
-        epsilon: float = 1e-5,
-    ):
-        super().__init__()
-        if codebook_dim is None:
-            codebook_dim = dim
-
-        requires_projection = codebook_dim != dim
-
-        self.project_out = (
-            nn.Linear(codebook_dim, dim) if requires_projection else nn.Identity()
-        )
-        self.epsilon = epsilon
-        self._codebook = EuclideanCodebook(
-            dim=codebook_dim,
-            codebook_size=codebook_size,
-            epsilon=epsilon
-        )
-        self.codebook_size = codebook_size
-
-    def decode(self, codes: torch.Tensor) -> torch.Tensor:
-        quantized = self._codebook.decode(codes)
-        quantized = self.project_out(quantized)
-        quantized = quantized.transpose(1, 2)
-        return quantized
-
-
-class ResidualVectorQuantization(nn.Module):
-    def __init__(self, *, num_quantizers: int, **kwargs):
-        super().__init__()
-        self.layers = nn.ModuleList(
-            [VectorQuantization(**kwargs) for _ in range(num_quantizers)]
-        )
-
-    def decode(self, codes: torch.Tensor) -> torch.Tensor:
-        quantized = torch.zeros([1], device=codes.device)[0]
-        for idx, layer_codes in enumerate(codes):
-            layer = self.layers[idx]
-            assert isinstance(layer, VectorQuantization)
-            quantized = quantized + layer.decode(layer_codes)
-        return quantized
-
-
-class ResidualVectorQuantizer(nn.Module):
-    def __init__(
-        self,
-        dimension: int = 128,
-        input_dimension: Optional[int] = None,
-        output_dimension: Optional[int] = None,
-        n_q: int = 8,
-        q_dropout: bool = False,
-        no_quantization_rate: float = 0.0,
-        bins: int = 1024,
-        decay: float = 0.99,
-        force_projection: bool = False,
-    ):
-        super().__init__()
-        self.max_n_q = n_q
-        self.n_q = n_q
-        self.q_dropout = q_dropout
-        self.no_quantization_rate = no_quantization_rate
-        self.dimension = dimension
-        self.input_dimension = input_dimension or dimension
-        self.output_dimension = output_dimension or dimension
-        self.bins = bins
-        self.decay = decay
-        self.input_proj: torch.nn.Module
-        self.output_proj: torch.nn.Module
-        if self.input_dimension == self.dimension and not force_projection:
-            self.input_proj = torch.nn.Identity()
-        else:
-            self.input_proj = torch.nn.Conv1d(
-                self.input_dimension, self.dimension, 1, bias=False
-            )
-        if self.output_dimension == self.dimension and not force_projection:
-            self.output_proj = torch.nn.Identity()
-        else:
-            self.output_proj = torch.nn.Conv1d(
-                self.dimension, self.output_dimension, 1, bias=False
-            )
-        self.vq = ResidualVectorQuantization(
-            dim=self.dimension,
-            codebook_size=self.bins,
-            num_quantizers=self.n_q
-        )
-
-    def decode(self, codes: torch.Tensor) -> torch.Tensor:
-        codes = codes.transpose(0, 1)
-        quantized = self.vq.decode(codes)
-        quantized = self.output_proj(quantized)
-        return quantized
-
-
-class SplitResidualVectorQuantizer(nn.Module):
-    """Residual Vector Quantizer with separate projections for the first quantizer and the rest.
-
-    Args:
-        n_q (int): Number of residual vector quantizers used.
-        n_semantic_q (int): Number of residual vector quantizers used for the semantic quantizer.
-        **kwargs: Arguments to the constructor of `ResidualVectorQuantizer` that are shared between both.
-    """
-
-    def __init__(
-        self,
-        *,
-        n_q: int = 8,
-        n_q_semantic: int = 1,
-        **kwargs,
-    ):
-        super().__init__()
-        assert n_q > n_q_semantic, (
-            f"Number of quantizers {n_q} must be larger "
-            f"than the number of semantic quantizers {n_q_semantic}."
-        )
-        self.max_n_q = n_q
-        self.n_q_semantic = n_q_semantic
-        self.n_q_acoustic = n_q - n_q_semantic
-        q_dropout = kwargs.pop("q_dropout", False)
-        self.rvq_first = ResidualVectorQuantizer(
-            n_q=n_q_semantic, force_projection=True, q_dropout=False, **kwargs
-        )
-        self.rvq_rest = ResidualVectorQuantizer(
-            n_q=n_q - n_q_semantic,
-            force_projection=True,
-            q_dropout=q_dropout,
-            **kwargs,
-        )
-
-    def decode(self, codes: torch.Tensor) -> torch.Tensor:
-        """Decode the given codes to the quantized representation."""
-        # codes is [B, K, T], with T frames, K nb of codebooks.
-        quantized = self.rvq_first.decode(codes[:, : self.n_q_semantic])
-        if codes.shape[1] > self.n_q_semantic:
-            quantized += self.rvq_rest.decode(codes[:, self.n_q_semantic :])
-        return quantized
-
-
-class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
-    def __init__(self, config: Qwen3TTSTokenizerV2DecoderConfig):
-        super().__init__(config)
-        self.total_upsample = np.prod(config.upsample_rates + config.upsampling_ratios)
-        self.pre_transformer = Qwen3TTSTokenizerV2DecoderTransformerModel._from_config(config)
-        
-        self.quantizer = SplitResidualVectorQuantizer(
-            dimension=config.codebook_dim // 2,
-            n_q=config.num_quantizers,
-            n_q_semantic=1,
-            bins=config.codebook_size,
-            input_dimension=config.codebook_dim,
-            output_dimension=config.codebook_dim,
-        )
-
-        self.pre_conv = Qwen3TTSTokenizerV2CausalConvNet(
-            config.codebook_dim,
-            config.latent_dim,
-            kernel_size=3,
-        )
-
-        upsample = []
-        for factor in config.upsampling_ratios:
-            upsample.append(
-                nn.ModuleList(
-                    [
-                        Qwen3TTSTokenizerV2CausalTransConvNet(config.latent_dim, config.latent_dim, factor, factor),
-                        Qwen3TTSTokenizerV2ConvNeXtBlock(config.latent_dim),
-                    ]
-                )
-            )
-        self.upsample = nn.ModuleList(upsample)
-
-        decoder = [Qwen3TTSTokenizerV2CausalConvNet(config.latent_dim, config.decoder_dim, 7)]
-        for i in range(len(config.upsample_rates)):
-            decoder.append(Qwen3TTSTokenizerV2DecoderDecoderBlock(config, i))
-        output_dim = config.decoder_dim // 2 ** len(config.upsample_rates)
-        decoder += [
-            SnakeBeta(output_dim),
-            Qwen3TTSTokenizerV2CausalConvNet(output_dim, 1, 7),
-        ]
-        self.decoder = nn.ModuleList(decoder)
-
-        self.post_init()
-
-    def forward(self, codes):
-        if codes.shape[1] != self.config.num_quantizers:
-            raise ValueError(f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}")
-
-        hidden = self.quantizer.decode(codes)
-        hidden = self.pre_conv(hidden).transpose(1, 2)
-
-        hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state
-        hidden = hidden.permute(0, 2, 1)
-        for blocks in self.upsample:
-            for block in blocks:
-                hidden = block(hidden)
-        wav = hidden
-        for block in self.decoder:
-            wav = block(wav)
-        return wav.clamp(min=-1, max=1)
-
-    def chunked_decode(self, codes, chunk_size=300, left_context_size=25):
-        wavs = []
-        start_index = 0
-        while start_index < codes.shape[-1]:
-            end_index = min(start_index + chunk_size, codes.shape[-1])
-            context_size = left_context_size if start_index - left_context_size > 0 else start_index
-            codes_chunk = codes[..., start_index - context_size : end_index]
-            wav_chunk = self(codes_chunk)
-            wavs.append(wav_chunk[..., context_size * self.total_upsample :])
-            start_index = end_index
-        return torch.cat(wavs, dim=-1)
-
-
-class Qwen3TTSTokenizerV2Encoder(MimiModel):
-    def __init__(self, config: MimiConfig):
-        super().__init__(config)
-        self.config = config
-
-        self.upsample = None
-        self.decoder_transformer = None
-        self.decoder = None
-
-        self.post_init()
-
-
-@auto_docstring
-class Qwen3TTSTokenizerV2PreTrainedModel(PreTrainedModel):
-    config: Qwen3TTSTokenizerV2Config
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = True
-    _skip_keys_device_placement = "past_key_values"
-    _supports_flash_attn = True
-    _supports_sdpa = True
-    _can_compile_fullgraph = False
-    _supports_attention_backend = True
-
-
-@auto_docstring(
-    custom_intro="""
-    The Qwen3TTSTokenizerV2 model.
-    """
-)
-class Qwen3TTSTokenizerV2Model(Qwen3TTSTokenizerV2PreTrainedModel):
-    def __init__(self, config: Qwen3TTSTokenizerV2Config):
-        super().__init__(config)
-        self.config = config
-
-        self.encoder_valid_num_quantizers = config.encoder_valid_num_quantizers
-
-        self.input_sample_rate = config.input_sample_rate
-        self.output_sample_rate = config.output_sample_rate
-
-        self.decode_upsample_rate = config.decode_upsample_rate
-        self.encode_downsample_rate = config.encode_downsample_rate
-
-        self.encoder = Qwen3TTSTokenizerV2Encoder._from_config(self.config.encoder_config)
-        self.decoder = Qwen3TTSTokenizerV2Decoder._from_config(self.config.decoder_config)
-
-        self.post_init()
-    
-    def get_model_type(self):
-        return self.config.model_type
-    
-    def get_input_sample_rate(self):
-        return self.input_sample_rate
-    
-    def get_output_sample_rate(self):
-        return self.output_sample_rate
-    
-    def get_encode_downsample_rate(self):
-        return self.encode_downsample_rate
-    
-    def get_decode_upsample_rate(self):
-        return self.decode_upsample_rate
-    
-    def encode(
-        self,
-        input_values: torch.Tensor,
-        padding_mask: Optional[torch.Tensor] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[tuple[torch.Tensor, Optional[torch.Tensor]], Qwen3TTSTokenizerV2EncoderOutput]:
-        """
-        Encodes the input audio waveform into discrete codes.
-
-        Args:
-            input_values (`torch.Tensor` of shape `(batch_size, sequence_length)`):
-                Float values of the input audio waveform.
-            padding_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`):
-                Indicates which inputs are to be ignored due to padding, where elements are either 1 for *not masked* or 0
-                for *masked*.
-            return_dict (`bool`, *optional*):
-                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-        """
-        return_dict = return_dict if return_dict is not None else self.config.return_dict
-
-        encoded_frames = self.encoder.encode(input_values=input_values.unsqueeze(1),
-                                             return_dict=True)
-        audio_codes = encoded_frames.audio_codes[:, :self.encoder_valid_num_quantizers]
-        audio_codes = [code[..., :-(-mask.sum() // self.encode_downsample_rate)].transpose(0, 1) for code, mask in zip(audio_codes, padding_mask)]
-
-        if not return_dict:
-            return (
-                audio_codes,
-            )
-
-        return Qwen3TTSTokenizerV2EncoderOutput(audio_codes)
-
-    def decode(
-        self,
-        audio_codes: torch.Tensor,
-        return_dict: Optional[bool] = None,
-    ) -> Union[tuple[torch.Tensor, torch.Tensor], Qwen3TTSTokenizerV2DecoderOutput]:
-        """
-        Decodes the given frames into an output audio waveform.
-
-        Note that the output might be a bit bigger than the input. In that case, any extra steps at the end can be
-        trimmed.
-
-        Args:
-            audio_codes (`torch.LongTensor`  of shape `(batch_size, codes_length, num_quantizers)`, *optional*):
-                Discret code embeddings computed using `model.encode`.
-            return_dict (`bool`, *optional*):
-                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-
-        """
-        return_dict = return_dict if return_dict is not None else self.config.return_dict
-
-        audio_values = self.decoder.chunked_decode(audio_codes.transpose(1, 2)).squeeze(1)
-
-        audio_lengths = (audio_codes[..., 0] > 0).sum(1) * self.decode_upsample_rate
-        audio_values = [a[:l] for a, l in zip(audio_values, audio_lengths)]
-
-        if not return_dict:
-            return (
-                audio_values,
-            )
-
-        return Qwen3TTSTokenizerV2DecoderOutput(audio_values)
-
-
-__all__ = ["Qwen3TTSTokenizerV2Model", "Qwen3TTSTokenizerV2PreTrainedModel"]

qwen_tts/core/tokenizer_25hz/configuration_qwen3_tts_tokenizer_v1.py

@@ -1,332 +0,0 @@
-# coding=utf-8
-# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# 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.
-"""Qwen3TTSTokenizerV1 model configuration"""
-
-from transformers.configuration_utils import PretrainedConfig
-from transformers.utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class Qwen3TTSTokenizerV1DecoderDiTConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of the Qwen3TTSTokenizerV1DecoderToken2WavDiT.
-    It defines the architecture of the DiT model, which is used for generating mel-spectrograms from tokens.
-
-    Args:
-        hidden_size (`int`, *optional*, defaults to 1024):
-            The dimension of the model.
-        num_hidden_layers (`int`, *optional*, defaults to 22):
-            The number of transformer blocks in the DiT model.
-        num_attention_heads (`int`, *optional*, defaults to 16):
-            The number of attention heads in each transformer block.
-        ff_mult (`int`, *optional*, defaults to 2):
-            The multiplier for the feedforward layer in each transformer block.
-        emb_dim (`int`, *optional*, defaults to 512):
-            The dimension of the embedding layer.
-        head_dim (`int`, *optional*, defaults to 64):
-            The dimension of each attention head.
-        repeats (`int`, *optional*, defaults to 2):
-            The number of times the codec embeddings are repeated.
-        num_embeds (`int`, *optional*, defaults to 8193):
-            The number of unique embeddings in the codec.
-        mel_dim (`int`, *optional*, defaults to 80):
-            The dimension of the mel-spectrogram.
-        dropout (`float`, *optional*, defaults to 0.1):
-            The dropout rate for the transformer blocks.
-
-        enc_emb_dim (`int`, *optional*, defaults to 192):
-            The dimension of the pre-trained speaker embedding.
-        enc_dim (`int`, *optional*, defaults to 128):
-            The dimension of the encoder output.
-        enc_channels (`list[int]`, *optional*, defaults to `[256, 256, 256, 256, 768]`):
-            A list of output channels for each TDNN/SERes2Net layer in the encoder.
-        enc_kernel_sizes (`list[int]`, *optional*, defaults to `[5, 3, 3, 3, 1]`):
-            A list of kernel sizes for each layer in the encoder.
-        enc_dilations (`list[int]`, *optional*, defaults to `[1, 2, 3, 4, 1]`):
-            A list of dilations for each layer in the encoder.
-        enc_attention_channels (`int`, *optional*, defaults to 64):
-            The number of attention channels in the SqueezeExcitationBlock.
-        enc_res2net_scale (`int`, *optional*, defaults to 2):
-            The scale of the Res2Net block in the encoder.
-        enc_se_channels (`int`, *optional*, defaults to 64):
-            The number of output channels after squeeze in the SqueezeExcitationBlock.
-    """
-
-    model_type = "qwen3_tts_tokenizer_v1_decoder_dit"
-
-    def __init__(
-        self,
-        hidden_size=1024,
-        num_hidden_layers=22,
-        num_attention_heads=16,
-        ff_mult=2,
-        emb_dim=512,
-        head_dim=64,
-        rope_theta=10000.0,
-        max_position_embeddings=32768,
-        block_size=24,
-        look_ahead_layers=[10],
-        look_backward_layers=[0, 20],
-        repeats=2,
-        num_embeds=8193,
-        mel_dim=80,
-        dropout=0.1,
-        enc_emb_dim=192,
-        enc_dim=128,
-        enc_channels=[256, 256, 256, 256, 768],
-        enc_kernel_sizes=[5, 3, 3, 3, 1],
-        enc_dilations=[1, 2, 3, 4, 1],
-        enc_attention_channels=64,
-        enc_res2net_scale=2,
-        enc_se_channels=64,
-        **kwargs,
-    ):
-        self.hidden_size = hidden_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.ff_mult = ff_mult
-        self.emb_dim = emb_dim
-        self.head_dim = head_dim
-        self.rope_theta = rope_theta
-        self.max_position_embeddings = max_position_embeddings
-        self.block_size = block_size
-        self.look_ahead_layers = look_ahead_layers
-        self.look_backward_layers = look_backward_layers
-        self.repeats = repeats
-        self.num_embeds = num_embeds
-        self.mel_dim = mel_dim
-        self.dropout = dropout
-        self.enc_emb_dim = enc_emb_dim
-        self.enc_dim = enc_dim
-        self.enc_channels = enc_channels
-        self.enc_kernel_sizes = enc_kernel_sizes
-        self.enc_dilations = enc_dilations
-        self.enc_attention_channels = enc_attention_channels
-        self.enc_res2net_scale = enc_res2net_scale
-        self.enc_se_channels = enc_se_channels
-        super().__init__(**kwargs)
-
-
-class Qwen3TTSTokenizerV1DecoderBigVGANConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of the Qwen3TTSTokenizerV1DecoderToken2WavBigVGAN module.
-    It defines the architecture of the BigVGAN model, which is used for converting mel-spectrograms to waveforms.
-
-    Args:
-        mel_dim (`int`, *optional*, defaults to 80):
-            The dimension of the mel-spectrogram.
-        upsample_initial_channel (`int`, *optional*, defaults to 1536):
-            The number of channels in the initial upsampling layer.
-        resblock_kernel_sizes (`list[int]`, *optional*, defaults to `[3, 7, 11]`):
-            A list of kernel sizes for each residual block.
-        resblock_dilation_sizes (`list[list[int]]`, *optional*, defaults to `[[1, 3, 5], [1, 3, 5], [1, 3, 5]]`):
-            A list of dilation sizes for each residual block.
-        upsample_rates (`list[int]`, *optional*, defaults to `[5, 3, 2, 2, 2, 2]`):
-            A list of upsampling rates for each upsampling layer.
-        upsample_kernel_sizes (`list[int]`, *optional*, defaults to `[11, 7, 4, 4, 4, 4]`):
-            A list of kernel sizes for each upsampling layer.
-    """
-
-    model_type = "qwen3_tts_tokenizer_v1_decoder_bigvgan"
-
-    def __init__(
-        self,
-        mel_dim=80,
-        upsample_initial_channel=1536,
-        resblock_kernel_sizes=[3, 7, 11],
-        resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
-        upsample_rates=[5, 3, 2, 2, 2, 2],
-        upsample_kernel_sizes=[11, 7, 4, 4, 4, 4],
-        **kwargs,
-    ):
-        self.mel_dim = mel_dim
-        self.upsample_initial_channel = upsample_initial_channel
-        self.resblock_kernel_sizes = resblock_kernel_sizes
-        self.resblock_dilation_sizes = resblock_dilation_sizes
-        self.upsample_rates = upsample_rates
-        self.upsample_kernel_sizes = upsample_kernel_sizes
-        super().__init__(**kwargs)
-
-
-class Qwen3TTSTokenizerV1DecoderConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`Qwen3TTSTokenizerV1DecoderConfig`].
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-    Args:
-        dit_config ([`DiT_Args`], *optional*):
-            Configuration class for the Diffusion Transformer (DiT) module responsible for generating mel-spectrograms.
-        bigvgan_config ([`BigVGAN_Args`], *optional*):
-            Configuration class for the BigVGAN module responsible for converting mel-spectrograms to waveforms.
-    """
-
-    model_type = "qwen3_tts_tokenizer_v1_decoder"
-    sub_configs = {
-        "dit_config": Qwen3TTSTokenizerV1DecoderDiTConfig,
-        "bigvgan_config": Qwen3TTSTokenizerV1DecoderBigVGANConfig,
-    }
-
-    def __init__(self, dit_config=None, bigvgan_config=None, **kwargs):
-        if dit_config is None:
-            dit_config = {}
-        if bigvgan_config is None:
-            bigvgan_config = {}
-        self.dit_config = Qwen3TTSTokenizerV1DecoderDiTConfig(**dit_config)
-        self.bigvgan_config = Qwen3TTSTokenizerV1DecoderBigVGANConfig(**bigvgan_config)
-        super().__init__(**kwargs)
-
-
-class Qwen3TTSTokenizerV1EncoderConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of the Qwen3TTSTokenizerV1 Encoder.
-
-    The encoder typically takes mel-spectrogram features and produces high-level audio representations, then (optionally)
-    applies an Audio-VQ module (e.g., GRVQ) to discretize continuous representations into codes.
-
-    Args:
-        n_mels (`int`, *optional*, defaults to 128):
-            Number of mel bins in the input mel-spectrogram.
-        n_ctx (`int`, *optional*, defaults to 1500):
-            Maximum input sequence length (in frames/tokens) for the encoder.
-        n_state (`int`, *optional*, defaults to 1280):
-            Hidden size (model dimension) of the encoder transformer.
-        n_head (`int`, *optional*, defaults to 20):
-            Number of attention heads in each transformer layer.
-        n_layer (`int`, *optional*, defaults to 32):
-            Number of transformer layers.
-        n_window (`int`, *optional*, defaults to 100):
-            Window size used by the model for local attention / chunking (implementation-dependent).
-        output_dim (`int`, *optional*, defaults to 3584):
-            Output feature dimension produced by the encoder head (before/after projection, implementation-dependent).
-
-        grad_checkpointing (`bool`, *optional*, defaults to `False`):
-            Whether to enable gradient checkpointing to reduce memory usage during training.
-        enable_mp (`bool`, *optional*, defaults to `False`):
-            Whether to enable model parallel features (implementation-dependent).
-        audio_sequence_parallel (`bool`, *optional*, defaults to `False`):
-            Whether to enable sequence parallelism for audio branch (implementation-dependent).
-
-        audio_vq_type (`str`, *optional*, defaults to `"GRVQ"`):
-            Type of audio vector-quantization module. Common choices: `"GRVQ"`, `"RVQ"`, etc.
-        audio_vq_layers (`int`, *optional*, defaults to 6):
-            Number of VQ layers / quantizers (e.g., number of residual quantizers for RVQ/GRVQ-like designs).
-        audio_vq_codebook_size (`int`, *optional*, defaults to 32768):
-            Size of each codebook (number of entries).
-        audio_vq_codebook_dim (`int`, *optional*, defaults to 1280):
-            Dimension of codebook vectors (often equals encoder hidden size).
-        audio_vq_pe (`bool`, *optional*, defaults to `True`):
-            Whether to use positional encoding (or position embeddings) inside the VQ module.
-        audio_vq_ds_rate (`int`, *optional*, defaults to 2):
-            Downsampling rate applied before VQ (e.g., temporal downsample factor).
-    """
-
-    model_type = "qwen3_tts_tokenizer_v1_encoder"
-
-    def __init__(
-        self,
-        n_mels=128,
-        n_ctx=1500,
-        n_state=1280,
-        n_head=20,
-        n_layer=32,
-        n_window=100,
-        output_dim=3584,
-        grad_checkpointing=False,
-        enable_mp=False,
-        audio_sequence_parallel=False,
-        audio_vq_type="GRVQ",
-        audio_vq_layers=6,
-        audio_vq_codebook_size=32768,
-        audio_vq_codebook_dim=1280,
-        audio_vq_pe=True,
-        audio_vq_ds_rate=2,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.n_mels = n_mels
-        self.n_ctx = n_ctx
-        self.n_state = n_state
-        self.n_head = n_head
-        self.n_layer = n_layer
-        self.n_window = n_window
-        self.output_dim = output_dim
-        self.grad_checkpointing = grad_checkpointing
-        self.enable_mp = enable_mp
-        self.audio_sequence_parallel = audio_sequence_parallel
-        self.audio_vq_type = audio_vq_type
-        self.audio_vq_layers = audio_vq_layers
-        self.audio_vq_codebook_size = audio_vq_codebook_size
-        self.audio_vq_codebook_dim = audio_vq_codebook_dim
-        self.audio_vq_pe = audio_vq_pe
-        self.audio_vq_ds_rate = audio_vq_ds_rate
-
-
-class Qwen3TTSTokenizerV1Config(PretrainedConfig):
-    """
-    This is the configuration class to store the configuration of a [`Qwen3TTSTokenizerV1Config`]. It is used to instantiate a Qwen3TTSTokenizerV1Model
-    model according to the specified sub-models configurations, defining the model architecture.
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
-
-    Args:
-        encoder_config (`dict`, *optional*): Configuration of the underlying encoder sub-model.
-        decoder_config (`dict`, *optional*): Configuration of the underlying decoder sub-model.
-    """
-
-    model_type = "qwen3_tts_tokenizer_25hz"
-    sub_configs = {
-        "encoder_config": Qwen3TTSTokenizerV1EncoderConfig,
-        "decoder_config": Qwen3TTSTokenizerV1DecoderConfig,
-    }
-
-    def __init__(
-        self,
-        encoder_config=None,
-        decoder_config=None,
-        input_sample_rate=24000,
-        output_sample_rate=24000,
-        decode_upsample_rate=1920,
-        encode_downsample_rate=1920,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        if encoder_config is None:
-            encoder_config = {}
-            logger.info("encoder_config is None. Initializing encoder with default values")
-        if decoder_config is None:
-            decoder_config = {}
-            logger.info("decoder_config is None. Initializing decoder with default values")
-
-        self.encoder_config = Qwen3TTSTokenizerV1EncoderConfig(**encoder_config)
-        self.decoder_config = Qwen3TTSTokenizerV1DecoderConfig(**decoder_config)
-
-        self.input_sample_rate = input_sample_rate
-        self.output_sample_rate = output_sample_rate
-        self.decode_upsample_rate = decode_upsample_rate
-        self.encode_downsample_rate = encode_downsample_rate
-
-
-__all__ = [
-    "Qwen3TTSTokenizerV1Config", 
-    "Qwen3TTSTokenizerV1EncoderConfig",
-    "Qwen3TTSTokenizerV1DecoderConfig", 
-    "Qwen3TTSTokenizerV1DecoderBigVGANConfig",
-    "Qwen3TTSTokenizerV1DecoderDiTConfig"
-]

qwen_tts/core/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py

@@ -1,1528 +0,0 @@
-# coding=utf-8
-# Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# 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.
-"""PyTorch Qwen3TTSTokenizerV1 model."""
-
-import math
-from dataclasses import dataclass
-from typing import Optional, Union, List
-
-import numpy as np
-import torch
-from torch import nn
-from torch.nn import Parameter
-from torch.nn import functional as F
-from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
-from transformers.utils import ModelOutput, auto_docstring, logging
-from transformers.utils.hub import cached_file
-
-from torch.nn.utils.rnn import pad_sequence
-
-from .vq.whisper_encoder import get_mel_audio, get_T_after_cnn
-from .vq.speech_vq import WhisperEncoderVQ, XVectorExtractor
-
-from .configuration_qwen3_tts_tokenizer_v1 import (
-    Qwen3TTSTokenizerV1Config,
-    Qwen3TTSTokenizerV1EncoderConfig,
-    Qwen3TTSTokenizerV1DecoderConfig,
-    Qwen3TTSTokenizerV1DecoderBigVGANConfig,
-    Qwen3TTSTokenizerV1DecoderDiTConfig
-)
-
-logger = logging.get_logger(__name__)
-
-
-@dataclass
-@auto_docstring
-class Qwen3TTSTokenizerV1EncoderOutput(ModelOutput):
-    r"""
-    audio_codes (`List[torch.LongTensor]`):
-        Discret code embeddings computed using `model.encode`, each tensor has shape (codes_length_i,).
-    xvectors (`List[torch.FloatTensor]`):
-        X-vector embeddings computed using `model.encode`, each tensor has shape (xvector_dim,).
-    ref_mels (`List[torch.FloatTensor]`):
-        Reference mel spectrogram computed using `model.encode`, each tensor has shape (mel_length_i, mel_dim,).
-    """
-
-    audio_codes: List[torch.LongTensor] = None
-    xvectors: List[torch.FloatTensor] = None
-    ref_mels: List[torch.FloatTensor] = None
-
-
-@dataclass
-@auto_docstring
-class Qwen3TTSTokenizerV1DecoderOutput(ModelOutput):
-    r"""
-    audio_values (`List[torch.FloatTensor]`):
-        Decoded audio values, obtained using the decoder part of Qwen3TTSTokenizerV1.
-        Each tensor has shape (segment_length_i).
-    """
-
-    audio_values: List[torch.FloatTensor] = None
-
-
-@auto_docstring
-class Qwen3TTSTokenizerV1DecoderPreTrainedModel(PreTrainedModel):
-    config: Qwen3TTSTokenizerV1DecoderConfig
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = True
-    _skip_keys_device_placement = "past_key_values"
-    _supports_flash_attn = True
-    _supports_sdpa = True
-    _can_compile_fullgraph = False
-    _supports_attention_backend = True
-
-
-@auto_docstring
-class Qwen3TTSTokenizerV1EncoderPreTrainedModel(PreTrainedModel):
-    config: Qwen3TTSTokenizerV1EncoderConfig
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = True
-    _skip_keys_device_placement = "past_key_values"
-    _supports_flash_attn = True
-    _supports_sdpa = True
-    _can_compile_fullgraph = False
-    _supports_attention_backend = True
-
-
-class Qwen3TTSTokenizerV1DecoderDiTRotaryEmbedding(nn.Module):
-    inv_freq: torch.Tensor  # fix linting for `register_buffer`
-
-    def __init__(self, dim, base=10000):
-        super().__init__()
-
-        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
-        self.register_buffer("inv_freq", inv_freq)
-
-    def forward(self, x):
-        batch_size, seq_len = x.shape[0], x.shape[1]
-        t = torch.arange(seq_len, device=x.device)
-        device_type = x.device.type
-        device_type = device_type if device_type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):
-            freqs = t.unsqueeze(1).float() @ self.inv_freq.unsqueeze(0).float()
-            freqs = torch.stack((freqs, freqs), dim=-1)
-            freqs = freqs.reshape(*freqs.shape[:-2], -1)
-            freqs = freqs.repeat(batch_size, *([1] * freqs.dim()))
-            cos = freqs.cos()
-            sin = freqs.sin()
-
-        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
-
-
-class TimeDelayNetBlock(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        dilation,
-    ):
-        super().__init__()
-        self.conv = nn.Conv1d(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            kernel_size=kernel_size,
-            dilation=dilation,
-            padding="same",
-            padding_mode="reflect",
-        )
-        self.activation = nn.ReLU()
-
-    def forward(self, hidden_states: torch.Tensor):
-        return self.activation(self.conv(hidden_states))
-
-
-class Res2NetBlock(torch.nn.Module):
-    def __init__(self, in_channels, out_channels, scale=8, kernel_size=3, dilation=1):
-        super().__init__()
-
-        in_channel = in_channels // scale
-        hidden_channel = out_channels // scale
-
-        self.blocks = nn.ModuleList(
-            [
-                TimeDelayNetBlock(
-                    in_channel,
-                    hidden_channel,
-                    kernel_size=kernel_size,
-                    dilation=dilation,
-                )
-                for i in range(scale - 1)
-            ]
-        )
-        self.scale = scale
-
-    def forward(self, hidden_states):
-        outputs = []
-        for i, hidden_part in enumerate(torch.chunk(hidden_states, self.scale, dim=1)):
-            if i == 0:
-                output_part = hidden_part
-            elif i == 1:
-                output_part = self.blocks[i - 1](hidden_part)
-            else:
-                output_part = self.blocks[i - 1](hidden_part + output_part)
-            outputs.append(output_part)
-        output = torch.cat(outputs, dim=1)
-        return output
-
-
-class SqueezeExcitationBlock(nn.Module):
-    def __init__(self, in_channels, se_channels, out_channels):
-        super().__init__()
-
-        self.conv1 = nn.Conv1d(
-            in_channels=in_channels,
-            out_channels=se_channels,
-            kernel_size=1,
-            padding="same",
-            padding_mode="reflect",
-        )
-        self.relu = nn.ReLU(inplace=True)
-        self.conv2 = nn.Conv1d(
-            in_channels=se_channels,
-            out_channels=out_channels,
-            kernel_size=1,
-            padding="same",
-            padding_mode="reflect",
-        )
-        self.sigmoid = nn.Sigmoid()
-
-    def forward(self, hidden_states):
-        hidden_states_mean = hidden_states.mean(dim=2, keepdim=True)
-
-        hidden_states_mean = self.relu(self.conv1(hidden_states_mean))
-        hidden_states_mean = self.sigmoid(self.conv2(hidden_states_mean))
-
-        return hidden_states * hidden_states_mean
-
-
-class AttentiveStatisticsPooling(nn.Module):
-    """This class implements an attentive statistic pooling layer for each channel.
-    It returns the concatenated mean and std of the input tensor.
-    """
-
-    def __init__(self, channels, attention_channels=128):
-        super().__init__()
-
-        self.eps = 1e-12
-        self.tdnn = TimeDelayNetBlock(channels * 3, attention_channels, 1, 1)
-        self.tanh = nn.Tanh()
-        self.conv = nn.Conv1d(
-            in_channels=attention_channels,
-            out_channels=channels,
-            kernel_size=1,
-            padding="same",
-            padding_mode="reflect",
-        )
-
-    def _length_to_mask(self, length, max_len=None, dtype=None, device=None):
-        """Creates a binary mask for each sequence.
-
-        Reference: https://discuss.pytorch.org/t/how-to-generate-variable-length-mask/23397/3
-
-        Arguments
-        ---------
-        length : torch.LongTensor
-            Containing the length of each sequence in the batch. Must be 1D.
-        max_len : int
-            Max length for the mask, also the size of the second dimension.
-        dtype : torch.dtype, default: None
-            The dtype of the generated mask.
-        device: torch.device, default: None
-            The device to put the mask variable.
-
-        Returns
-        -------
-        mask : tensor
-            The binary mask.
-        """
-
-        if max_len is None:
-            max_len = length.max().long().item()  # using arange to generate mask
-        mask = torch.arange(max_len, device=length.device, dtype=length.dtype).expand(
-            len(length), max_len
-        ) < length.unsqueeze(1)
-
-        mask = torch.as_tensor(mask, dtype=dtype, device=device)
-        return mask
-
-    def _compute_statistics(self, x, m, dim=2):
-        mean = (m * x).sum(dim)
-        std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps))
-        return mean, std
-
-    def forward(self, hidden_states):
-        seq_length = hidden_states.shape[-1]
-        lengths = torch.ones(hidden_states.shape[0], device=hidden_states.device)
-
-        # Make binary mask of shape [N, 1, L]
-        mask = self._length_to_mask(
-            lengths * seq_length, max_len=seq_length, dtype=hidden_states.dtype, device=hidden_states.device
-        )
-        mask = mask.unsqueeze(1)
-
-        # Expand the temporal context of the pooling layer by allowing the
-        # self-attention to look at global properties of the utterance.
-        total = mask.sum(dim=2, keepdim=True)
-
-        mean, std = self._compute_statistics(hidden_states, mask / total)
-        mean = mean.unsqueeze(2).repeat(1, 1, seq_length)
-        std = std.unsqueeze(2).repeat(1, 1, seq_length)
-        attention = torch.cat([hidden_states, mean, std], dim=1)
-
-        # Apply layers
-        attention = self.conv(self.tanh(self.tdnn(attention)))
-
-        # Filter out zero-paddings
-        attention = attention.masked_fill(mask == 0, float("-inf"))
-
-        attention = F.softmax(attention, dim=2)
-        mean, std = self._compute_statistics(hidden_states, attention)
-        # Append mean and std of the batch
-        pooled_stats = torch.cat((mean, std), dim=1)
-        pooled_stats = pooled_stats.unsqueeze(2)
-
-        return pooled_stats
-
-
-class SqueezeExcitationRes2NetBlock(nn.Module):
-    """An implementation of building block in ECAPA-TDNN, i.e.,
-    TDNN-Res2Net-TDNN-SqueezeExcitationBlock.
-    """
-
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        res2net_scale=8,
-        se_channels=128,
-        kernel_size=1,
-        dilation=1,
-    ):
-        super().__init__()
-        self.out_channels = out_channels
-        self.tdnn1 = TimeDelayNetBlock(
-            in_channels,
-            out_channels,
-            kernel_size=1,
-            dilation=1,
-        )
-        self.res2net_block = Res2NetBlock(out_channels, out_channels, res2net_scale, kernel_size, dilation)
-        self.tdnn2 = TimeDelayNetBlock(
-            out_channels,
-            out_channels,
-            kernel_size=1,
-            dilation=1,
-        )
-        self.se_block = SqueezeExcitationBlock(out_channels, se_channels, out_channels)
-
-    def forward(self, hidden_state):
-        residual = hidden_state
-
-        hidden_state = self.tdnn1(hidden_state)
-        hidden_state = self.res2net_block(hidden_state)
-        hidden_state = self.tdnn2(hidden_state)
-        hidden_state = self.se_block(hidden_state)
-
-        return hidden_state + residual
-
-
-class ECAPA_TimeDelayNet(torch.nn.Module):
-    """An implementation of the speaker embedding model in a paper.
-    "ECAPA-TDNN: Emphasized Channel Attention, Propagation and Aggregation in
-    TDNN Based Speaker Verification" (https://huggingface.co/papers/2005.07143).
-    """
-
-    def __init__(self, config: Qwen3TTSTokenizerV1DecoderBigVGANConfig):
-        super().__init__()
-        if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(config.enc_channels) != len(
-            config.enc_dilations
-        ):
-            raise ValueError("enc_channels, enc_kernel_sizes and enc_dilations should have same length")
-        self.channels = config.enc_channels
-        self.blocks = nn.ModuleList()
-
-        # The initial TDNN layer
-        self.blocks.append(
-            TimeDelayNetBlock(
-                config.mel_dim,
-                config.enc_channels[0],
-                config.enc_kernel_sizes[0],
-                config.enc_dilations[0],
-            )
-        )
-
-        # SE-Res2Net layers
-        for i in range(1, len(config.enc_channels) - 1):
-            self.blocks.append(
-                SqueezeExcitationRes2NetBlock(
-                    config.enc_channels[i - 1],
-                    config.enc_channels[i],
-                    res2net_scale=config.enc_res2net_scale,
-                    se_channels=config.enc_se_channels,
-                    kernel_size=config.enc_kernel_sizes[i],
-                    dilation=config.enc_dilations[i],
-                )
-            )
-
-        # Multi-layer feature aggregation
-        self.mfa = TimeDelayNetBlock(
-            config.enc_channels[-1],
-            config.enc_channels[-1],
-            config.enc_kernel_sizes[-1],
-            config.enc_dilations[-1],
-        )
-
-        # Attentive Statistical Pooling
-        self.asp = AttentiveStatisticsPooling(
-            config.enc_channels[-1],
-            attention_channels=config.enc_attention_channels,
-        )
-
-        # Final linear transformation
-        self.fc = nn.Conv1d(
-            in_channels=config.enc_channels[-1] * 2,
-            out_channels=config.enc_dim,
-            kernel_size=1,
-            padding="same",
-            padding_mode="reflect",
-        )
-
-    def forward(self, hidden_states):
-        # Minimize transpose for efficiency
-        hidden_states = hidden_states.transpose(1, 2)
-
-        hidden_states_list = []
-        for layer in self.blocks:
-            hidden_states = layer(hidden_states)
-            hidden_states_list.append(hidden_states)
-
-        # Multi-layer feature aggregation
-        hidden_states = torch.cat(hidden_states_list[1:], dim=1)
-        hidden_states = self.mfa(hidden_states)
-
-        # Attentive Statistical Pooling
-        hidden_states = self.asp(hidden_states)
-
-        # Final linear transformation
-        hidden_states = self.fc(hidden_states)
-
-        hidden_states = hidden_states.squeeze(-1)
-        return hidden_states
-
-
-class DiTInputEmbedding(nn.Module):
-    def __init__(self, config: Qwen3TTSTokenizerV1DecoderBigVGANConfig):
-        super().__init__()
-        self.proj = nn.Linear(
-            config.mel_dim + config.enc_dim + config.enc_emb_dim + config.emb_dim,
-            config.hidden_size,
-        )
-        self.spk_encoder = ECAPA_TimeDelayNet(config)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        speaker_embedding: torch.Tensor,
-        condition_vector: torch.Tensor,
-        code_embed: torch.Tensor,
-        drop_audio_cond: Optional[bool] = False,
-        code_embed_uncond: Optional[bool] = None,
-        apply_cfg: Optional[bool] = True,
-    ):
-        if apply_cfg:
-            hidden_states = torch.cat([hidden_states, hidden_states], dim=0)
-            speaker_embedding = torch.cat([speaker_embedding, torch.zeros_like(speaker_embedding)], dim=0)
-            condition_vector = torch.cat([condition_vector, torch.zeros_like(condition_vector)], dim=0)
-            code_embed = torch.cat([code_embed, code_embed_uncond], dim=0)
-        elif drop_audio_cond:  # cfg for cond audio
-            condition_vector = torch.zeros_like(condition_vector)
-            speaker_embedding = torch.zeros_like(speaker_embedding)
-        condition_vector = self.spk_encoder(condition_vector).unsqueeze(1).repeat(1, hidden_states.size(1), 1)
-        hidden_states = self.proj(torch.cat((hidden_states, condition_vector, code_embed, speaker_embedding), dim=-1))
-
-        return hidden_states
-
-
-# Transformer backbone using DiT blocks
-class DiTCodecEmbedding(nn.Module):
-    def __init__(self, codec_num_embeds, codec_dim, repeats):
-        super().__init__()
-        self.repeats = repeats
-        self.codec_embed = nn.Embedding(codec_num_embeds + 1, codec_dim)
-
-    def forward(self, code, drop_code=False):
-        if drop_code:
-            code = torch.zeros_like(code)
-        code_embed = self.codec_embed(code)
-
-        code_embed = torch.repeat_interleave(code_embed, repeats=self.repeats, dim=1)
-        return code_embed
-
-
-# AdaLayerNormZero
-# return with modulated x for attn input, and params for later mlp modulation
-class AdaLayerNormZero(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-
-        self.silu = nn.SiLU()
-        self.linear = nn.Linear(dim, dim * 6)
-
-        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-
-    def forward(self, hidden_states, emb=None):
-        emb = self.linear(self.silu(emb))
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
-
-        hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
-        return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp
-
-
-# AdaLayerNormZero for final layer
-# return only with modulated x for attn input, cuz no more mlp modulation
-class AdaLayerNormZero_Final(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-
-        self.silu = nn.SiLU()
-        self.linear = nn.Linear(dim, dim * 2)
-
-        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-
-    def forward(self, hidden_states, emb):
-        emb = self.linear(self.silu(emb))
-        scale, shift = torch.chunk(emb, 2, dim=1)
-
-        hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
-        return hidden_states
-
-
-# FeedForward
-class DiTMLP(nn.Module):
-    def __init__(self, dim, mult=4, dropout=0.0):
-        super().__init__()
-        inner_dim = int(dim * mult)
-
-        self.ff = nn.ModuleList(
-            [
-                nn.Linear(dim, inner_dim),
-                nn.GELU(approximate="tanh"),
-                nn.Dropout(dropout),
-                nn.Linear(inner_dim, dim),
-            ]
-        )
-
-    def forward(self, hidden_states):
-        for layer in self.ff:
-            hidden_states = layer(hidden_states)
-        return hidden_states
-
-
-# Modified from Llama with a different rotate function, will fixed in next release
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
-    """Applies Rotary Position Embedding to the query and key tensors.
-
-    Args:
-        q (`torch.Tensor`): The query tensor.
-        k (`torch.Tensor`): The key tensor.
-        cos (`torch.Tensor`): The cosine part of the rotary embedding.
-        sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`, *optional*):
-            Deprecated and unused.
-        unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
-    Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
-    """
-
-    def rotate_half_codec(x):
-        # x = rearrange(x, "... (d r) -> ... d r", r=2)
-        x = x.reshape(*x.shape[:-1], -1, 2)
-        x1, x2 = x.unbind(dim=-1)
-        x = torch.stack((-x2, x1), dim=-1)
-        return x.reshape(*x.shape[:-2], -1)
-
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
-    q_embed = (q * cos) + (rotate_half_codec(q) * sin)
-    k_embed = (k * cos) + (rotate_half_codec(k) * sin)
-    return q_embed, k_embed
-
-
-class DiTAttention(nn.Module):
-    def __init__(self, config: Qwen3TTSTokenizerV1DecoderBigVGANConfig):
-        super().__init__()
-
-        self.config = config
-        self.dim = config.hidden_size
-        self.heads = config.num_attention_heads
-        self.inner_dim = config.head_dim * config.num_attention_heads
-        self.dropout = config.dropout
-        self.is_causal = False
-
-        self.to_q = nn.Linear(config.hidden_size, self.inner_dim)
-        self.to_k = nn.Linear(config.hidden_size, self.inner_dim)
-        self.to_v = nn.Linear(config.hidden_size, self.inner_dim)
-
-        self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, config.hidden_size), nn.Dropout(config.dropout)])
-
-    def forward(
-        self,
-        hidden_states,  # noised input x
-        position_embeddings=None,  # rotary position embedding for x
-        attention_mask=None,
-    ) -> torch.Tensor:
-        batch_size = hidden_states.shape[0]
-
-        # `sample` projections.
-        query = self.to_q(hidden_states)
-        key = self.to_k(hidden_states)
-        value = self.to_v(hidden_states)
-
-        # attention
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // self.heads
-        query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
-        key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
-        value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
-
-        # apply rotary position embedding
-        # Due to training process, only first head is applied with RoPE, will be fixed at next release
-        cos, sin = position_embeddings
-        query, key = apply_rotary_pos_emb(query, key, cos, sin)
-
-        attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
-        attention_weights, _ = attention_interface(
-            self,
-            query,
-            key,
-            value,
-            attention_mask=attention_mask,
-            is_causal=False,
-        )
-
-        # mask. e.g. inference got a batch with different target durations, mask out the padding
-        attention_weights = attention_weights.reshape(batch_size, -1, self.heads * head_dim)
-        attention_weights = attention_weights.to(query.dtype)
-
-        # linear proj
-        attention_output = self.to_out[0](attention_weights)
-        attention_output = self.to_out[1](attention_output)
-
-        return attention_output
-
-
-# time step conditioning embedding
-class SinusPositionEmbedding(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        self.dim = dim
-
-    def forward(self, hidden_states, scale=1000):
-        device = hidden_states.device
-        half_dim = self.dim // 2
-        emb = math.log(10000) / (half_dim - 1)
-        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
-        emb = scale * hidden_states.unsqueeze(1) * emb.unsqueeze(0)
-        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
-        return emb.type_as(hidden_states)
-
-
-class DiTTimestepEmbedding(nn.Module):
-    def __init__(self, dim, freq_embed_dim=256):
-        super().__init__()
-        self.time_embed = SinusPositionEmbedding(freq_embed_dim)
-        self.time_mlp = nn.ModuleList([nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)])
-
-    def forward(self, timestep):
-        time_hidden = self.time_embed(timestep)
-        time_hidden = time_hidden.to(timestep.dtype)
-        for layer in self.time_mlp:
-            time_hidden = layer(time_hidden)  # b d
-        return time_hidden
-
-
-class DiTDecoderLayer(nn.Module):
-    def __init__(self, config: Qwen3TTSTokenizerV1DecoderBigVGANConfig, look_ahead_block=0, look_backward_block=0):
-        super().__init__()
-        self.attn_norm = AdaLayerNormZero(config.hidden_size)
-
-        self.attn = DiTAttention(config)
-        self.look_ahead_block = look_ahead_block
-        self.look_backward_block = look_backward_block
-        self.ff_norm = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=1e-6)
-        self.ff = DiTMLP(dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout)
-
-    def forward(
-        self, hidden_states, timestep, position_embeddings=None, block_diff=None
-    ):  # x: noised input, t: time embedding
-        # pre-norm & modulation for attention input
-        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(hidden_states, emb=timestep)
-
-        # attention
-        attn_output = self.attn(
-            hidden_states=norm,
-            position_embeddings=position_embeddings,
-            attention_mask=(block_diff >= -float(self.look_backward_block))
-            & (block_diff <= float(self.look_ahead_block)),
-        )
-
-        # process attention output for input x
-        hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_output
-
-        norm = self.ff_norm(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
-        ff_output = self.ff(norm)
-        hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
-
-        return hidden_states
-
-
-class SnakeBeta(nn.Module):
-    """
-    A modified Snake function which uses separate parameters for the magnitude of the periodic components
-    Shape:
-        - Input: (B, C, T)
-        - Output: (B, C, T), same shape as the input
-    Parameters:
-        - alpha - trainable parameter that controls frequency
-        - beta - trainable parameter that controls magnitude
-    References:
-        - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
-        https://huggingface.co/papers/2006.08195
-    """
-
-    def __init__(self, in_features, alpha=1.0):
-        super().__init__()
-        self.in_features = in_features
-
-        # initialize alpha
-        self.alpha = Parameter(torch.zeros(in_features) * alpha)
-        self.beta = Parameter(torch.zeros(in_features) * alpha)
-
-        self.no_div_by_zero = 0.000000001
-
-    def forward(self, hidden_states):
-        """
-        Forward pass of the function.
-        Applies the function to the input elementwise.
-        SnakeBeta ∶= x + 1/b * sin^2 (xa)
-        """
-        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
-        beta = self.beta.unsqueeze(0).unsqueeze(-1)
-        alpha = torch.exp(alpha)
-        beta = torch.exp(beta)
-        hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(
-            torch.sin(hidden_states * alpha), 2
-        )
-
-        return hidden_states
-
-
-def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
-    """Generates a 1D Kaiser-windowed sinc filter.
-
-    Args:
-        cutoff (float): Normalized cutoff frequency (0 to 0.5).
-        half_width (float): Transition bandwidth.
-        kernel_size (int): Number of filter taps.
-
-    Returns:
-        torch.Tensor: A tensor of shape (1, 1, kernel_size) representing the filter.
-    """
-    is_even = kernel_size % 2 == 0
-    half_size = kernel_size // 2
-
-    # Compute Kaiser window parameters
-    delta_f = 4 * half_width
-    attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
-
-    if attenuation > 50.0:
-        beta = 0.1102 * (attenuation - 8.7)
-    elif attenuation >= 21.0:
-        beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0)
-    else:
-        beta = 0.0
-
-    kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
-
-    # Compute time indices
-    if is_even:
-        time_indices = torch.arange(-half_size, half_size) + 0.5
-    else:
-        time_indices = torch.arange(kernel_size) - half_size
-
-    # Compute sinc filter
-    if cutoff == 0:
-        return torch.zeros((1, 1, kernel_size), dtype=torch.float32)  # Ensures correct shape
-
-    sinc_filter = torch.sinc(2 * cutoff * time_indices)
-    normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
-
-    # Normalize to ensure sum = 1 (avoid leakage of constant component)
-    normalized_filter /= normalized_filter.sum()
-
-    return normalized_filter.view(1, 1, kernel_size)
-
-
-class UpSample1d(nn.Module):
-    def __init__(self, ratio=2, kernel_size=None):
-        super().__init__()
-        self.ratio = ratio
-        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
-        self.stride = ratio
-        self.pad = self.kernel_size // ratio - 1
-        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
-        self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
-
-        filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
-        self.register_buffer("filter", filter, persistent=False)
-
-    def forward(self, hidden_states):
-        channels = hidden_states.shape[1]
-
-        hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate")
-        hidden_states = self.ratio * F.conv_transpose1d(
-            hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels
-        )
-        hidden_states = hidden_states[..., self.pad_left : -self.pad_right]
-
-        return hidden_states
-
-
-class DownSample1d(nn.Module):
-    def __init__(self, ratio=2, kernel_size=None):
-        super().__init__()
-        cutoff = 0.5 / ratio
-        half_width = 0.6 / ratio
-
-        if cutoff < 0.0:
-            raise ValueError("Minimum cutoff must be larger than zero.")
-        if cutoff > 0.5:
-            raise ValueError("A cutoff above 0.5 does not make sense.")
-
-        self.even = kernel_size % 2 == 0
-        self.pad_left = kernel_size // 2 - int(self.even)
-        self.pad_right = kernel_size // 2
-        self.stride = ratio
-        filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
-        self.register_buffer("filter", filter, persistent=False)
-
-    def forward(self, hidden_states):
-        channels = hidden_states.shape[1]
-        hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate")
-        out = F.conv1d(hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels)
-        return out
-
-
-class TorchActivation1d(nn.Module):
-    def __init__(
-        self,
-        activation,
-        up_ratio: int = 2,
-        down_ratio: int = 2,
-        up_kernel_size: int = 12,
-        down_kernel_size: int = 12,
-    ):
-        super().__init__()
-        if not callable(activation):
-            raise TypeError("Activation function must be callable")
-        self.act = activation
-        self.upsample = UpSample1d(up_ratio, up_kernel_size)
-        self.downsample = DownSample1d(down_ratio, down_kernel_size)
-
-    def forward(self, hidden_states):
-        hidden_states = self.upsample(hidden_states)
-        hidden_states = self.act(hidden_states)
-        hidden_states = self.downsample(hidden_states)
-
-        return hidden_states
-
-
-class CausalConv1d(nn.Conv1d):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.causal_padding = self.dilation[0] * (self.kernel_size[0] - 1)
-
-    def forward(self, x):
-        return self._conv_forward(F.pad(x, [self.causal_padding, 0]), self.weight, self.bias)
-
-
-class AMPBlock(torch.nn.Module):
-    def __init__(
-        self,
-        channels,
-        kernel_size=3,
-        dilation=(1, 3, 5),
-        causal_type='1',
-    ):
-        super().__init__()
-
-        self.convs1 = nn.ModuleList(
-            [
-                CausalConv1d(
-                    channels,
-                    channels,
-                    kernel_size,
-                    1,
-                    dilation=dilation[0],
-                ),
-                CausalConv1d(
-                    channels,
-                    channels,
-                    kernel_size,
-                    1,
-                    dilation=dilation[1],
-                ),
-                CausalConv1d(
-                    channels,
-                    channels,
-                    kernel_size,
-                    1,
-                    dilation=dilation[2],
-                ),
-            ]
-        )
-
-        if causal_type == '1':
-            self.convs2 = nn.ModuleList(
-                [
-                    nn.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=1,
-                        padding=self._get_padding(kernel_size, 1),
-                    ),
-                    nn.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=1,
-                        padding=self._get_padding(kernel_size, 1),
-                    ),
-                    nn.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=1,
-                        padding=self._get_padding(kernel_size, 1),
-                    ),
-                ]
-            )
-        else:
-            self.convs2 = nn.ModuleList(
-                [
-                    CausalConv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=1,
-                    ),
-                    CausalConv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=1,
-                    ),
-                    CausalConv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=1,
-                    ),
-                ]
-            )
-
-        self.num_layers = len(self.convs1) + len(self.convs2)  # total number of conv layers
-
-        self.activations = nn.ModuleList(
-            [TorchActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
-        )
-
-        if causal_type == '2':
-            self.pre_conv = nn.Conv1d(
-                                channels,
-                                channels,
-                                kernel_size,
-                                stride=1,
-                                padding=self._get_padding(kernel_size, 1),
-                            )
-            self.pre_act = TorchActivation1d(activation=SnakeBeta(channels))
-        else:
-            self.pre_conv = nn.Identity()
-            self.pre_act = nn.Identity()
-
-    def _get_padding(self, kernel_size, dilation=1):
-        return int((kernel_size * dilation - dilation) / 2)
-
-    def forward(self, x):
-        hidden_states = self.pre_conv(x)
-        hidden_states = self.pre_act(hidden_states)
-        acts1, acts2 = self.activations[::2], self.activations[1::2]
-        for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2):
-            hidden_states = act1(hidden_states)
-            hidden_states = conv1(hidden_states)
-            hidden_states = act2(hidden_states)
-            hidden_states = conv2(hidden_states)
-            x = x + hidden_states
-        return x
-
-
-@auto_docstring
-class Qwen3TTSTokenizerV1DecoderBigVGANModel(Qwen3TTSTokenizerV1DecoderPreTrainedModel):
-    config: Qwen3TTSTokenizerV1DecoderBigVGANConfig
-
-    def __init__(self, config: Qwen3TTSTokenizerV1DecoderBigVGANConfig):
-        super().__init__(config)
-        self.num_residual_blocks = len(config.resblock_kernel_sizes)
-        self.num_upsample_layers = len(config.upsample_rates)
-
-        self.conv_pre = nn.Conv1d(config.mel_dim, config.upsample_initial_channel, 5, 1, padding=2)
-
-        # Removing extra ModuleList breaks official state dict
-        ups = [
-            nn.ModuleList(
-                [
-                    nn.ConvTranspose1d(
-                        config.upsample_initial_channel // (2**layer_idx),
-                        config.upsample_initial_channel // (2 ** (layer_idx + 1)),
-                        kernel_size,
-                        stride,
-                        padding=(kernel_size - stride) // 2,
-                    )
-                ]
-            )
-            for layer_idx, (stride, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes))
-        ]
-        self.ups = nn.ModuleList(ups)
-
-        self.resblocks = nn.ModuleList(
-            [
-                AMPBlock(config.upsample_initial_channel // (2 ** (layer_idx + 1)), kernel_size, dilation, '1' if layer_idx > 1 else '2')
-                for layer_idx in range(self.num_upsample_layers)
-                for kernel_size, dilation in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes)
-            ]
-        )
-
-        self.activation_post = TorchActivation1d(
-            activation=SnakeBeta(config.upsample_initial_channel // (2**self.num_upsample_layers))
-        )
-        self.conv_post = nn.Conv1d(
-            config.upsample_initial_channel // (2**self.num_upsample_layers), 1, 7, 1, padding=3, bias=False
-        )
-
-    def normalize_spectrogram(self, spectrogram, max_value, min_db):
-        return torch.clamp((2 * max_value) * ((spectrogram - min_db) / (-min_db)) - max_value, -max_value, max_value)
-
-    def amplitude_to_db(self, amplitude, min_db_level):
-        min_level = torch.exp(
-            torch.tensor(min_db_level / 20.0 * np.log(10), device=amplitude.device, dtype=amplitude.dtype)
-        )
-        return 20 * torch.log10(torch.clamp(amplitude, min=min_level))
-
-    def process_mel_spectrogram(self, mel_spectrogram):
-        amplitude_spectrum = torch.exp(mel_spectrogram)
-        decibel_spectrum = self.amplitude_to_db(amplitude_spectrum, -115) - 20
-        return self.normalize_spectrogram(decibel_spectrum, 1, -115)
-
-    def forward(self, mel_spectrogram):
-        processed_spectrogram = self.process_mel_spectrogram(mel_spectrogram)
-        hidden_representation = self.conv_pre(processed_spectrogram)
-
-        for layer_index in range(self.num_upsample_layers):
-            hidden_representation = self.ups[layer_index][0](hidden_representation)
-            residual_output = sum(
-                self.resblocks[layer_index * self.num_residual_blocks + block_index](hidden_representation)
-                for block_index in range(self.num_residual_blocks)
-            )
-            residual_output = residual_output / self.num_residual_blocks
-            hidden_representation = residual_output
-
-        hidden_representation = self.activation_post(hidden_representation)
-        output_waveform = self.conv_post(hidden_representation)
-        return torch.clamp(output_waveform, min=-1.0, max=1.0).squeeze(1)
-
-
-@auto_docstring
-class Qwen3TTSTokenizerV1DecoderDiTModel(Qwen3TTSTokenizerV1DecoderPreTrainedModel):
-    config: Qwen3TTSTokenizerV1DecoderDiTConfig
-    _no_split_modules = ["DiTDecoderLayer"]
-
-    def __init__(self, config: Qwen3TTSTokenizerV1DecoderDiTConfig):
-        super().__init__(config)
-        self.mel_dim = config.mel_dim
-        self.repeats = config.repeats
-        self.time_embed = DiTTimestepEmbedding(config.hidden_size)
-
-        self.text_embed = DiTCodecEmbedding(config.num_embeds, config.emb_dim, config.repeats)
-        self.input_embed = DiTInputEmbedding(config)
-
-        self.rotary_embed = Qwen3TTSTokenizerV1DecoderDiTRotaryEmbedding(config.head_dim)
-
-        self.hidden_size = config.hidden_size
-        self.layers = config.num_hidden_layers
-        self.block_size = config.block_size
-        self.num_attention_heads = config.num_attention_heads
-
-        self.transformer_blocks = nn.ModuleList()
-        for i in range(config.num_hidden_layers):
-            self.transformer_blocks.append(
-                DiTDecoderLayer(
-                    config,
-                    look_ahead_block=1 if i in config.look_ahead_layers else 0,
-                    look_backward_block=1 if i in config.look_backward_layers else 0,
-                )
-            )
-
-        self.norm_out = AdaLayerNormZero_Final(config.hidden_size)  # final modulation
-        self.proj_out = nn.Linear(config.hidden_size, config.mel_dim)
-
-    def _create_block_diff(self, hidden_states):
-        batch, seq_len = hidden_states.shape[0], hidden_states.shape[1]
-        block_indices = torch.arange(seq_len, device=hidden_states.device) // self.block_size  # [seq_length]
-
-        block_i = block_indices.unsqueeze(1)  # [seq_length, 1]
-        block_j = block_indices.unsqueeze(0)  # [1, seq_length]
-        block_diff = block_j - block_i  # (n, n)
-
-        return block_diff.expand(batch, self.num_attention_heads, seq_len, seq_len)
-
-    def forward(
-        self,
-        hidden_states,
-        condition_vector,
-        speaker_embedding,
-        quantized_code,
-        time_step,
-        drop_audio_conditioning=False,
-        drop_code=False,
-        apply_cfg=True,
-    ):
-        batch_size = hidden_states.shape[0] * 2
-        if time_step.ndim == 0:
-            time_step = time_step.repeat(batch_size)
-
-        # Compute embeddings
-        time_embedding = self.time_embed(time_step)
-        text_embedding = self.text_embed(quantized_code, drop_code=False if apply_cfg else drop_code)
-        text_embedding_unconditioned = self.text_embed(quantized_code, drop_code=True) if apply_cfg else None
-
-        hidden_states = self.input_embed(
-            hidden_states,
-            speaker_embedding,
-            condition_vector,
-            text_embedding,
-            drop_audio_cond=drop_audio_conditioning,
-            code_embed_uncond=text_embedding_unconditioned,
-            apply_cfg=apply_cfg,
-        )
-
-        # Compute positional encodings
-        position_embeddings = self.rotary_embed(hidden_states)
-        blockwise_difference = self._create_block_diff(hidden_states)
-
-        # Transformer blocks
-        for transformer_block in self.transformer_blocks:
-            hidden_states = transformer_block(
-                hidden_states,
-                time_embedding,
-                position_embeddings=position_embeddings,
-                block_diff=blockwise_difference,
-            )
-
-        hidden_states = self.norm_out(hidden_states, time_embedding)
-        output = self.proj_out(hidden_states)
-
-        return output
-
-    def optimized_scale(self, positive_flat, negative_flat):
-        # Calculate dot production
-        dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
-        # Squared norm of uncondition
-        squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8
-        # st_star = v_cond^T * v_uncond / ||v_uncond||^2
-        st_star = dot_product / squared_norm
-        return st_star
-
-    @torch.no_grad()
-    def sample(
-        self,
-        conditioning_vector,
-        reference_mel_spectrogram,
-        quantized_code,
-        num_steps=10,
-        guidance_scale=0.5,
-        sway_coefficient=-1.0,
-    ):
-        noise_initialization = torch.randn([quantized_code.shape[0], 30000, self.mel_dim], dtype=reference_mel_spectrogram.dtype)
-        maximum_duration = quantized_code.shape[1] * self.repeats
-        initial_state = noise_initialization[:, :maximum_duration].to(quantized_code.device)
-        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, maximum_duration, 1)
-
-        def ode_function(time_step, hidden_states):
-            if guidance_scale < 1e-5:
-                prediction = self(
-                    hidden_states=hidden_states,
-                    speaker_embedding=conditioning_vector,
-                    condition_vector=reference_mel_spectrogram,
-                    quantized_code=quantized_code,
-                    time_step=time_step,
-                    drop_audio_conditioning=False,
-                    drop_code=False,
-                )
-                return prediction
-
-            model_output = self(
-                hidden_states=hidden_states,
-                quantized_code=quantized_code,
-                speaker_embedding=conditioning_vector,
-                condition_vector=reference_mel_spectrogram,
-                time_step=time_step,
-                apply_cfg=True,
-            )
-            guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0)
-
-            return guided_prediction + (guided_prediction - null_prediction) * guidance_scale
-
-        initial_time = 0
-        time_embedding = torch.linspace(
-            initial_time, 1, num_steps, device=quantized_code.device, dtype=conditioning_vector.dtype
-        )
-
-        if sway_coefficient is not None:
-            time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
-
-        values = initial_state.clone()
-        for t0, t1 in zip(time_embedding[:-1], time_embedding[1:]):
-            dt = t1 - t0
-            vt = ode_function(t0, values)
-            values = values + vt * dt
-
-        generated_mel_spectrogram = values.permute(0, 2, 1)
-        return generated_mel_spectrogram
-
-
-@auto_docstring
-class Qwen3TTSTokenizerV1Decoder(Qwen3TTSTokenizerV1DecoderPreTrainedModel):
-    config: Qwen3TTSTokenizerV1DecoderConfig
-    base_model_prefix = "model"
-    _no_split_modules = ["Qwen3TTSTokenizerV1DecoderDiTModel", "Qwen3TTSTokenizerV1DecoderBigVGANModel"]
-
-    def __init__(self, config: Qwen3TTSTokenizerV1DecoderConfig):
-        super().__init__(config)
-        attn_impl = config._attn_implementation
-        if config._attn_implementation == "flash_attention_2":
-            logger.warning_once(
-                "Qwen3TTSTokenizerV1Decoder must inference with fp32, but flash_attention_2 only supports fp16 and bf16, "
-                "attention implementation of Qwen3TTSTokenizerV1Decoder will fallback to sdpa."
-            )
-            attn_impl = "sdpa"
-        elif config._attn_implementation == "eager":
-            logger.warning_once(
-                "Qwen3TTSTokenizerV1Decoder does not support eager attention implementation, fall back to sdpa"
-            )
-            attn_impl = "sdpa"
-        self.dit = Qwen3TTSTokenizerV1DecoderDiTModel._from_config(
-            config.dit_config, attn_implementation=attn_impl
-        )
-        self.bigvgan = Qwen3TTSTokenizerV1DecoderBigVGANModel._from_config(
-            config.bigvgan_config, attn_implementation=attn_impl
-        )
-
-    def forward(
-        self,
-        code,
-        conditioning,
-        reference_mel,
-        num_steps=10,
-        guidance_scale=0.5,
-        sway_coefficient=-1.0,
-        **kwargs,
-    ):
-        """Generates a waveform from input code and conditioning parameters."""
-
-        mel_spectrogram = self.dit.sample(
-            conditioning,
-            reference_mel,
-            code,
-            num_steps=num_steps,
-            guidance_scale=guidance_scale,
-            sway_coefficient=sway_coefficient,
-        )
-
-        waveform = self.bigvgan(mel_spectrogram)
-
-        return waveform
-
-
-class Qwen3TTSTokenizerV1Encoder(Qwen3TTSTokenizerV1EncoderPreTrainedModel):
-    config: Qwen3TTSTokenizerV1EncoderConfig
-    def __init__(self, config: Qwen3TTSTokenizerV1EncoderConfig):
-        super().__init__(config)
-
-        self.tokenizer = WhisperEncoderVQ(
-            n_mels=config.n_mels,
-            n_ctx=config.n_ctx,
-            n_state=config.n_state,
-            n_head=config.n_head,
-            n_layer=config.n_layer,
-            n_window=config.n_window,
-            output_dim=config.output_dim,
-            grad_checkpointing=config.grad_checkpointing,
-            enable_mp=config.enable_mp,
-            audio_sequence_parallel=config.audio_sequence_parallel,
-            audio_vq_type=config.audio_vq_type,
-            audio_vq_layers=config.audio_vq_layers,
-            audio_vq_codebook_size=config.audio_vq_codebook_size,
-            audio_vq_codebook_dim=config.audio_vq_codebook_dim,
-            audio_vq_pe=config.audio_vq_pe,
-            audio_vq_ds_rate=config.audio_vq_ds_rate,
-        )
-
-        self.padding = True
-        self.audio_vq_ds_rate = self.tokenizer.audio_vq_ds_rate
-
-    def speech2mel(self, speechs):
-        mels = [
-            get_mel_audio(
-                speech, padding = self.padding, audio_vq_ds_rate = self.audio_vq_ds_rate
-            ).to(speech.dtype).to(self.tokenizer.conv1.weight.device)
-            for speech in speechs
-        ]
-        return mels
-
-    def mel2code(self, mels):
-        audio_mellens = [mel.size(-1) for mel in mels]
-        audio_aftercnnlens = [get_T_after_cnn(T) for T in audio_mellens]
-        audio_seqlens = [T + 2 for T in audio_aftercnnlens]
-
-        with torch.no_grad():
-            _, indices = self.tokenizer(
-                x_list = mels, 
-                audio_mellens = audio_mellens, 
-                audio_aftercnnlens = audio_aftercnnlens, 
-                audio_seqlens = audio_seqlens, 
-                return_indices=True,
-            )
-        
-        indice_lens = [T // self.tokenizer.audio_vq_ds_rate for T in audio_aftercnnlens]
-        indices  = pad_sequence(torch.split(indices, indice_lens), batch_first=True, padding_value=0)
-
-        return indices, indice_lens
-
-    def quantize_speech(self, speechs):
-        mels = self.speech2mel(speechs)
-        indices, indice_lens = self.mel2code(mels)
-        return indices, indice_lens
-
-
-@auto_docstring
-class Qwen3TTSTokenizerV1PreTrainedModel(PreTrainedModel):
-    config: Qwen3TTSTokenizerV1Config
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = True
-    _skip_keys_device_placement = "past_key_values"
-    _supports_flash_attn = True
-    _supports_sdpa = True
-    _can_compile_fullgraph = False
-    _supports_attention_backend = True
-
-
-@auto_docstring(
-    custom_intro="""
-    The Qwen3TTSTokenizerV1 model.
-    """
-)
-class Qwen3TTSTokenizerV1Model(Qwen3TTSTokenizerV1PreTrainedModel):
-    def __init__(self, config: Qwen3TTSTokenizerV1Config):
-        super().__init__(config)
-        self.config = config
-
-        self.input_sample_rate = config.input_sample_rate
-        self.output_sample_rate = config.output_sample_rate
-
-        self.decode_upsample_rate = config.decode_upsample_rate
-        self.encode_downsample_rate = config.encode_downsample_rate
-
-        self.encoder = Qwen3TTSTokenizerV1Encoder._from_config(self.config.encoder_config)
-        self.decoder = Qwen3TTSTokenizerV1Decoder._from_config(self.config.decoder_config)
-
-        self.encoder_xvector_extractor = None
-
-        self.post_init()
-    
-    def load_encoder_xvector_extractor(self, model_path):
-        self.encoder_xvector_extractor = XVectorExtractor(model_path)
-    
-    def get_model_type(self):
-        return self.config.model_type
-    
-    def get_input_sample_rate(self):
-        return self.input_sample_rate
-    
-    def get_output_sample_rate(self):
-        return self.output_sample_rate
-    
-    def get_encode_downsample_rate(self):
-        return self.encode_downsample_rate
-    
-    def get_decode_upsample_rate(self):
-        return self.decode_upsample_rate
-    
-    @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path,
-        *model_args,
-        config=None,
-        cache_dir=None,
-        ignore_mismatched_sizes=False,
-        force_download=False,
-        local_files_only=False,
-        token=None,
-        revision="main",
-        use_safetensors=None,
-        weights_only=True,
-        **kwargs,
-    ):
-        model = super().from_pretrained(
-            pretrained_model_name_or_path,
-            *model_args,
-            config=config,
-            cache_dir=cache_dir,
-            ignore_mismatched_sizes=ignore_mismatched_sizes,
-            force_download=force_download,
-            local_files_only=local_files_only,
-            token=token,
-            revision=revision,
-            use_safetensors=use_safetensors,
-            weights_only=weights_only,
-            **kwargs,
-        )
-        encoder_xvector_extractor_path = cached_file(
-            pretrained_model_name_or_path,
-            "campplus.onnx",
-            subfolder=kwargs.pop("subfolder", None),
-            cache_dir=kwargs.pop("cache_dir", None),
-            force_download=kwargs.pop("force_download", False),
-            proxies=kwargs.pop("proxies", None),
-            resume_download=kwargs.pop("resume_download", None),
-            local_files_only=kwargs.pop("local_files_only", False),
-            token=kwargs.pop("use_auth_token", None),
-            revision=kwargs.pop("revision", None),
-        )
-        if encoder_xvector_extractor_path is None:
-            raise ValueError(f"""{pretrained_model_name_or_path}/{encoder_xvector_extractor_path} not exists""")
-        model.load_encoder_xvector_extractor(encoder_xvector_extractor_path)
-
-        return model
-
-    def encode(     
-        self,
-        input_values: torch.Tensor,
-        padding_mask: Optional[torch.Tensor] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[tuple[torch.Tensor, Optional[torch.Tensor]], Qwen3TTSTokenizerV1EncoderOutput]:
-        """
-        Encodes the input audio waveform into discrete codes.
-
-        Args:
-            input_values (`torch.Tensor` of shape `(batch_size, sequence_length)`):
-                Float values of the input audio waveform.
-            padding_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`):
-                Indicates which inputs are to be ignored due to padding, where elements are either 1 for *not masked* or 0
-                for *masked*.
-            return_dict (`bool`, *optional*):
-                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-        """
-        return_dict = return_dict if return_dict is not None else self.config.return_dict
-
-        wavs = [value[:mask.sum()] for value, mask in zip(input_values, padding_mask)]
-
-        codes, codes_lens = self.encoder.quantize_speech(wavs)
-        codes = [c[:l] for c, l in zip(codes, codes_lens)]
-
-        xvectors = []
-        ref_mels = []
-        for wav in wavs:
-            xvector, ref_mel = self.encoder_xvector_extractor.extract_code(wav.cpu().numpy())
-            xvector = torch.tensor(xvector).to(wav.dtype).to(wav.device)
-            ref_mel = torch.tensor(ref_mel).to(wav.dtype).to(wav.device)
-            xvectors.append(xvector)
-            ref_mels.append(ref_mel)
-
-        if not return_dict:
-            return (
-                codes,
-                xvectors,
-                ref_mels
-            )
-
-        return Qwen3TTSTokenizerV1EncoderOutput(codes, xvectors, ref_mels)
-
-    def decode(
-        self,
-        audio_codes: torch.Tensor,
-        xvectors: torch.Tensor,
-        ref_mels: torch.Tensor,
-        return_dict: Optional[bool] = None,
-    ) -> Union[tuple[torch.Tensor, torch.Tensor], Qwen3TTSTokenizerV1DecoderOutput]:
-        """
-        Decodes the given frames into an output audio waveform.
-
-        Note that the output might be a bit bigger than the input. In that case, any extra steps at the end can be
-        trimmed.
-
-        Args:
-            audio_codes (`torch.LongTensor`  of shape `(batch_size, codes_length)`, *optional*):
-                Discret code embeddings computed using `model.encode`.
-            xvectors (`torch.FloatTensor` of shape `(batch_size, xvector_dim)`, *optional*):
-                X-vector embeddings computed using `model.encode`.
-            ref_mels (`torch.FloatTensor` of shape `(batch_size, mel_length, mel_dim)`, *optional*):
-                Reference mel spectrogram computed using `model.encode`.
-            return_dict (`bool`, *optional*):
-                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
-
-        """
-        return_dict = return_dict if return_dict is not None else self.config.return_dict
-
-        audio_values = self.decoder(code=audio_codes,
-                                    reference_mel=ref_mels,
-                                    conditioning=xvectors)
-        
-        audio_lengths = (audio_codes > 0).sum(1) * self.decode_upsample_rate
-        audio_values = [a[:l] for a, l in zip(audio_values, audio_lengths)]
-
-        if not return_dict:
-            return (
-                audio_values,
-            )
-
-        return Qwen3TTSTokenizerV1DecoderOutput(audio_values)
-
-
-__all__ = ["Qwen3TTSTokenizerV1Model", "Qwen3TTSTokenizerV1PreTrainedModel"]

qwen_tts/core/tokenizer_25hz/vq/assets/mel_filters.npz

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7450ae70723a5ef9d341e3cee628c7cb0177f36ce42c44b7ed2bf3325f0f6d4c
-size 4271

qwen_tts/core/tokenizer_25hz/vq/core_vq.py

@@ -1,523 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-#
-# This implementation is inspired from
-# https://github.com/lucidrains/vector-quantize-pytorch
-# which is released under MIT License. Hereafter, the original license:
-# MIT License
-#
-# Copyright (c) 2020 Phil Wang
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-"""Core vector quantization implementation."""
-import random
-import typing as tp
-from random import randrange
-
-import numpy as np
-from einops import rearrange, repeat
-from math import ceil
-import torch
-from torch import nn
-import torch.nn.functional as F
-
-
-def round_up_multiple(num, mult):
-    return ceil(num / mult) * mult
-
-def default(val: tp.Any, d: tp.Any) -> tp.Any:
-    return val if val is not None else d
-
-
-def ema_inplace(moving_avg, new, decay: float):
-    moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
-
-
-def laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):
-    return (x + epsilon) / (x.sum() + n_categories * epsilon)
-
-
-def uniform_init(*shape: int):
-    t = torch.empty(shape)
-    nn.init.kaiming_uniform_(t)
-    return t
-
-
-def sample_vectors(samples, num: int):
-    num_samples, device = samples.shape[0], samples.device
-
-    if num_samples >= num:
-        indices = torch.randperm(num_samples, device=device)[:num]
-    else:
-        indices = torch.randint(0, num_samples, (num,), device=device)
-
-    return samples[indices]
-
-
-@torch.no_grad()
-def kmeans(samples, num_clusters: int, num_iters: int = 10):
-    dim, dtype = samples.shape[-1], samples.dtype
-
-    means = sample_vectors(samples, num_clusters)
-
-    for _ in range(num_iters):
-        dists = -(
-                samples.pow(2).sum(1, keepdim=True)
-                - 2 * torch.matmul(samples, means.t())
-                + means.t().pow(2).sum(0, keepdim=True)
-        )
-
-        buckets = dists.max(dim=-1).indices
-        del dists
-        bins = torch.bincount(buckets, minlength=num_clusters)
-        zero_mask = bins == 0
-        bins_min_clamped = bins.masked_fill(zero_mask, 1)
-
-        new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
-        new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples)
-        new_means = new_means / bins_min_clamped[..., None]
-
-        means = torch.where(zero_mask[..., None], means, new_means)
-    return means, bins
-
-
-def preprocess(x):
-    x = rearrange(x, "... d -> (...) d")
-    return x
-
-
-def postprocess_emb(embed_ind, shape):
-    return embed_ind.view(*shape[:-1])
-
-
-class EuclideanCodebook(nn.Module):
-    """Codebook with Euclidean distance.
-    Args:
-        dim (int): Dimension.
-        codebook_size (int): Codebook size.
-        kmeans_init (bool): Whether to use k-means to initialize the codebooks.
-            If set to true, run the k-means algorithm on the first training batch and use
-            the learned centroids as initialization.
-        kmeans_iters (int): Number of iterations used for k-means algorithm at initialization.
-        decay (float): Decay for exponential moving average over the codebooks.
-        epsilon (float): Epsilon value for numerical stability.
-        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
-            that have an exponential moving average cluster size less than the specified threshold with
-            randomly selected vector from the current batch.
-    """
-
-    def __init__(
-            self,
-            dim: int,
-            codebook_size: int,
-            kmeans_init: int = False,
-            kmeans_iters: int = 10,
-            decay: float = 0.99,
-            epsilon: float = 1e-5,
-            threshold_ema_dead_code: float = 2.0,
-    ):
-        super().__init__()
-        self.decay = decay
-        self.codebook_size = codebook_size
-        self.kmeans_iters = kmeans_iters
-        self.epsilon = epsilon
-        self.threshold_ema_dead_code = threshold_ema_dead_code
-
-        self.inited = None
-        self.cluster_size = None
-        self.embed = None
-        self.embed_avg = None
-        self.training = True
-
-    def init_embed_(self, data):
-        if self.inited:
-            return
-
-        embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
-        self.embed.data.copy_(embed)
-        self.embed_avg.data.copy_(embed.clone())
-        self.cluster_size.data.copy_(cluster_size)
-        self.inited.data.copy_(torch.Tensor([True]))
-        # Make sure all buffers across workers are in sync after initialization
-        # distrib.broadcast_tensors([self.embed, self.embed_avg, self.cluster_size, self.inited])
-
-    def replace_(self, samples, mask):
-        modified_codebook = torch.where(
-            mask[..., None], sample_vectors(samples, self.codebook_size), self.embed
-        )
-        self.embed.data.copy_(modified_codebook)
-
-    def expire_codes_(self, batch_samples):
-        if self.threshold_ema_dead_code == 0:
-            return
-
-        cluster_size = self.cluster_size / sum(self.cluster_size) * self.codebook_size
-        expired_codes = cluster_size < self.threshold_ema_dead_code
-        if not torch.any(expired_codes):
-            return
-        else:
-            print(f"VQ expire infos: num_expire={sum(expired_codes)}, cluster_size[:5]={cluster_size[:5]}")
-
-        batch_samples = rearrange(batch_samples, "... d -> (...) d")
-        self.replace_(batch_samples, mask=expired_codes)
-        # sync buffers outside for efficiency
-        # distrib.broadcast_tensors(self.buffers())
-
-    def quantize(self, x):
-        embed = self.embed.t()
-        dist = -(
-            x.pow(2).sum(1, keepdim=True)
-            - 2 * x @ embed
-            + embed.pow(2).sum(0, keepdim=True)
-        )
-        embed_ind = dist.max(dim=-1).indices
-        return embed_ind
-
-    def dequantize(self, embed_ind):
-        quantize = F.embedding(embed_ind, self.embed)
-        return quantize
-
-    def encode(self, x, buffers):
-        self.inited, self.cluster_size, self.embed, self.embed_avg = buffers
-
-        shape = x.shape
-        # pre-process
-        x = preprocess(x)
-        # quantize
-        embed_ind = self.quantize(x)
-        # post-process
-        embed_ind = postprocess_emb(embed_ind, shape)
-        return embed_ind
-
-    def decode(self, embed_ind, buffers):
-        self.inited, self.cluster_size, self.embed, self.embed_avg = buffers
-
-        quantize = self.dequantize(embed_ind)
-        return quantize
-
-    def forward(self, x, buffers):
-        self.inited, self.cluster_size, self.embed, self.embed_avg = buffers
-
-        shape, dtype = x.shape, x.dtype
-        x = preprocess(x)
-
-        self.init_embed_(x)
-        if self.training:
-            # We do the expiry of code at that point as buffers are in sync
-            # and all the workers will take the same decision.
-            self.expire_codes_(x)
-
-        embed_ind = self.quantize(x)
-        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
-        embed_ind = postprocess_emb(embed_ind, shape)
-        quantize = self.dequantize(embed_ind)
-
-        if self.training:
-            ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
-            embed_sum = x.t() @ embed_onehot
-            ema_inplace(self.embed_avg, embed_sum.t(), self.decay)
-            cluster_size = (
-                laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon)
-                * self.cluster_size.sum()
-            )
-            embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
-            self.embed.data.copy_(embed_normalized)
-            # Note: after ema update, there is a very small difference between codebooks on GPUs.
-            # The impact can be very small, ignore it.
-
-        return quantize, embed_ind
-
-
-class VectorQuantization(nn.Module):
-    """Vector quantization implementation.
-    Currently, supports only euclidean distance.
-    Args:
-        dim (int): Dimension
-        codebook_size (int): Codebook size
-        codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim.
-        decay (float): Decay for exponential moving average over the codebooks.
-        epsilon (float): Epsilon value for numerical stability.
-        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
-        kmeans_iters (int): Number of iterations used for kmeans initialization.
-        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
-            that have an exponential moving average cluster size less than the specified threshold with
-            randomly selected vector from the current batch.
-        commitment_weight (float): Weight for commitment loss.
-    """
-    def __init__(
-            self,
-            dim: int,
-            codebook_size: int,
-            codebook_dim: tp.Optional[int] = None,
-            decay: float = 0.99,
-            epsilon: float = 1e-5,
-            kmeans_init: bool = True,
-            kmeans_iters: int = 50,
-            threshold_ema_dead_code: float = 2.0,
-            commitment_weight: float = 1.,
-    ):
-        super().__init__()
-        _codebook_dim: int = default(codebook_dim, dim)
-
-        requires_projection = _codebook_dim != dim
-        self.project_in = (nn.Linear(dim, _codebook_dim)) if requires_projection else (nn.Identity())
-        self.project_out = (nn.Linear(_codebook_dim, dim)) if requires_projection else (nn.Identity())
-
-        self.epsilon = epsilon
-        self.commitment_weight = commitment_weight
-
-        self._codebook = EuclideanCodebook(dim=_codebook_dim, codebook_size=codebook_size,
-                                           kmeans_init=kmeans_init, kmeans_iters=kmeans_iters,
-                                           decay=decay, epsilon=epsilon,
-                                           threshold_ema_dead_code=threshold_ema_dead_code)
-        self.codebook_size = codebook_size
-        self.training = True
-
-    @property
-    def codebook(self):
-        return self._codebook.embed
-
-    def encode(self, x, buffers):
-        # x = rearrange(x, "b d n -> b n d")
-        x = self.project_in(x)
-        embed_in = self._codebook.encode(x, buffers)
-        return embed_in
-
-    def decode(self, embed_ind, buffers):
-        quantize = self._codebook.decode(embed_ind, buffers)
-        quantize = self.project_out(quantize)
-        # quantize = rearrange(quantize, "b n d -> b d n")
-        return quantize
-
-    def forward(self, x, buffers):
-        device = x.device
-        # x = rearrange(x, "b d n -> b n d")
-        x = self.project_in(x)
-
-        quantize, embed_ind = self._codebook(x, buffers)
-
-        if self.training:
-            quantize = x + (quantize - x).detach()
-
-        loss = torch.tensor([0.0], device=device, requires_grad=self.training)
-
-        if self.training:
-            if self.commitment_weight > 0:
-                commit_loss = F.mse_loss(quantize.detach(), x)
-                loss = loss + commit_loss * self.commitment_weight
-
-        quantize = self.project_out(quantize)
-        # quantize = rearrange(quantize, "b n d -> b d n")
-        return quantize, embed_ind, loss
-
-
-class DistributedResidualVectorQuantization(nn.Module):
-    """Efficient distributed residual vector quantization implementation.
-    Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf
-    """
-    def __init__(self, *,
-                 num_quantizers,
-                 quantize_dropout: bool = False,
-                 rand_num_quant: tp.Optional[tp.List] = None,
-                 **kwargs):
-        super().__init__()
-        """
-        dim: int,
-        codebook_size: int,
-        codebook_dim: tp.Optional[int] = None,
-        """
-        codebook_size, codebook_dim = kwargs["codebook_size"], kwargs["codebook_dim"] if kwargs["codebook_dim"] else kwargs["dim"]
-        kmeans_init = kwargs["kmeans_init"]
-        if isinstance(kmeans_init, bool):
-            if not kwargs["kmeans_init"]:
-                # use uniform init
-                embed = uniform_init(num_quantizers, codebook_size, codebook_dim)
-                inited = True
-            else:
-                # to perform kmeans init on first batch
-                embed = torch.zeros(num_quantizers, codebook_size, codebook_dim)
-                inited = False
-        elif isinstance(kmeans_init, str):
-            # use prepared kmeans init
-            embed = np.load(kmeans_init)
-            embed = torch.from_numpy(embed)
-            if embed.dim() == 2:
-                embed = embed.unsqueeze(0)
-            inited = True
-        else:
-            raise TypeError("kmeans_init should be either a bool or string path to init weights.")
-
-        self.register_buffer("inited", torch.Tensor([[inited] for _ in range(num_quantizers)]))
-        self.register_buffer("cluster_size", torch.zeros(num_quantizers, codebook_size))
-        self.register_buffer("embed", embed)
-        self.register_buffer("embed_avg", embed.clone())
-
-        self.q0_ds_ratio = 1
-        if "q0_ds_ratio" in kwargs:
-            self.q0_ds_ratio = kwargs.pop("q0_ds_ratio")
-
-        self.layers = nn.ModuleList()
-        for i in range(num_quantizers):
-            vq_args = dict(**kwargs)
-            vq = VectorQuantization(**vq_args)
-            self.layers.append(vq)
-
-        self.quantize_dropout = quantize_dropout
-        self.rand_num_quant = rand_num_quant
-
-    def forward(self, x, n_q: tp.Optional[int] = None):
-        quantized_out = torch.zeros_like(x)
-        residual = x
-        bb, cc, tt = x.shape
-        device = x.device
-
-        all_losses = []
-        all_indices = []
-        all_sub_quants = []
-        n_q = n_q or len(self.layers)
-
-        should_quantize_dropout = self.training and self.quantize_dropout and self.rand_num_quant is not None
-        if should_quantize_dropout:
-            rand_quantize_dropout_index = random.choice(self.rand_num_quant)
-
-            null_indices_shape = (x.shape[0], x.shape[2])
-            null_indices = torch.full(null_indices_shape, -1., device=device, dtype=torch.long)
-            null_loss = torch.full((1,), 0., device=device, dtype=x.dtype)
-            null_sub_quant = torch.full(x.shape, -1, device=device, dtype=x.dtype)
-
-        for quantizer_index, layer in enumerate(self.layers[:n_q]):
-            # dropout except the first quantizer
-            if should_quantize_dropout and quantizer_index >= rand_quantize_dropout_index:
-                all_indices.append(null_indices)
-                all_losses.append(null_loss)
-                all_sub_quants.append(null_sub_quant)
-                continue
-
-            quant_in = residual
-            if self.q0_ds_ratio > 1 and quantizer_index == 0:
-                quant_in = F.interpolate(quant_in, size=[tt//2])
-            quantized, indices, loss = layer(quant_in, [
-                self.inited[quantizer_index],
-                self.cluster_size[quantizer_index],
-                self.embed[quantizer_index],
-                self.embed_avg[quantizer_index]
-            ])
-            if self.q0_ds_ratio > 1 and quantizer_index == 0:
-                quantized = F.interpolate(quantized, size=[tt])
-                indices = F.interpolate(indices.unsqueeze(1).float(), size=[tt]).squeeze(1).long()
-            residual = residual - quantized
-            quantized_out = quantized_out + quantized
-
-            all_indices.append(indices)
-            all_losses.append(loss)
-            all_sub_quants.append(quantized)
-
-        # sync buffers after one forward step
-        # distrib.broadcast_tensors(self.buffers())
-        out_losses, out_indices, out_sub_quants = map(torch.stack, (all_losses, all_indices, all_sub_quants))
-
-        return quantized_out, out_indices, out_losses
-
-    def encode(self, x: torch.Tensor, n_q: tp.Optional[int] = None) -> torch.Tensor:
-        residual = x
-        all_indices = []
-        n_q = n_q or len(self.layers)
-        for i, layer in enumerate(self.layers[:n_q]):
-            indices = layer.encode(residual, [
-                self.inited[i],
-                self.cluster_size[i],
-                self.embed[i],
-                self.embed_avg[i]
-            ])
-            quantized = layer.decode(indices, [
-                self.inited[i],
-                self.cluster_size[i],
-                self.embed[i],
-                self.embed_avg[i]
-            ])
-            residual = residual - quantized
-            all_indices.append(indices)
-        out_indices = torch.stack(all_indices)
-        return out_indices
-
-    def decode(self, q_indices: torch.Tensor) -> torch.Tensor:
-        quantized_out = torch.tensor(0.0, device=q_indices.device)
-        for i, indices in enumerate(q_indices):
-            layer = self.layers[i]
-            quantized = layer.decode(indices, [
-                self.inited[i],
-                self.cluster_size[i],
-                self.embed[i],
-                self.embed_avg[i]
-            ])
-            quantized_out = quantized_out + quantized
-        return quantized_out
-
-
-class DistributedGroupResidualVectorQuantization(nn.Module):
-    """Efficient distributed group residual vector quantization implementation.
-    Follows Algorithm 1. in https://arxiv.org/abs/2305.02765
-    Group Then rvq
-    """
-    def __init__(self, *,
-                 num_groups,
-                 num_quantizers,
-                 quantize_dropout: bool = False,
-                 rand_num_quant: tp.Optional[tp.List] = None,
-                 **kwargs):
-        super().__init__()
-        self.rvqs = nn.ModuleList(
-            [
-                DistributedResidualVectorQuantization(
-                    num_quantizers=num_quantizers,
-                    quantize_dropout=quantize_dropout,
-                    rand_num_quant=rand_num_quant,
-                    **kwargs
-                )
-                for _ in range(num_groups)
-            ]
-        )
-        self.num_groups = num_groups
-
-    def forward(self, x, n_q: tp.Optional[int] = None):
-        x_lst = torch.chunk(x, chunks=self.num_groups, dim=1)
-        all_quantized_out = []
-        all_indices = []
-        all_losses = []
-        for mod, item in zip(self.rvqs, x_lst):
-            quantized_out, out_indices, out_losses = mod(item, n_q)
-            all_quantized_out.append(quantized_out)
-            all_indices.append(out_indices)
-            all_losses.append(out_losses)
-
-        out_losses = torch.stack(all_losses, dim=1).mean(dim=1)
-
-        return torch.cat(all_quantized_out, dim=1), torch.stack(all_indices, dim=1), out_losses
-
-    def encode(self, x: torch.Tensor, n_q: tp.Optional[int] = None) -> torch.Tensor:
-        x_lst = torch.chunk(x, chunks=self.num_groups, dim=1)
-        return torch.stack([mod.encode(item, n_q) for mod, item in zip(self.rvqs, x_lst)], dim=1)
-
-    def decode(self, q_indices: torch.Tensor) -> torch.Tensor:
-        q_indices_lst = torch.chunk(q_indices, chunks=self.num_groups, dim=1)
-        return torch.cat([mod.decode(item.squeeze(1)) for mod, item in zip(self.rvqs, q_indices_lst)], dim=1)

qwen_tts/core/tokenizer_25hz/vq/speech_vq.py

@@ -1,357 +0,0 @@
-# coding=utf-8
-# 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 sox
-import copy
-import torch
-import operator
-import onnxruntime
-
-import torch.nn as nn
-import torch.nn.functional as F
-import torchaudio.compliance.kaldi as kaldi
-
-from librosa.filters import mel as librosa_mel_fn
-from itertools import accumulate
-from typing import List
-from torch import Tensor
-
-from .core_vq import DistributedGroupResidualVectorQuantization
-from .whisper_encoder import WhisperEncoder, Conv1d, ConvTranspose1d
-
-
-def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
-    return torch.log(torch.clamp(x, min=clip_val) * C)
-
-def spectral_normalize_torch(magnitudes):
-    output = dynamic_range_compression_torch(magnitudes)
-    return output
-
-class MelSpectrogramFeatures(nn.Module):
-    """
-    Calculate the BigVGAN style mel spectrogram of an input signal.
-    Args:
-        filter_length (int): The number of samples in the filter window, used for the Fourier Transform. Default is 1024.
-        hop_length (int): The number of samples between successive frames (stride of the STFT). Default is 160.
-        win_length (int): The length of the window function applied to each frame, usually less than or equal to the filter length. Default is 640.
-        n_mel_channels (int): The number of Mel-frequency channels to output from the Mel-scale spectrogram. Default is 80.
-        mel_fmin (int): The minimum frequency (in Hz) of the Mel-scale spectrogram. Default is 0.
-        mel_fmax (int): The maximum frequency (in Hz) of the Mel-scale spectrogram. Default is 8000.
-        sampling_rate (int): The sampling rate of the audio data (in Hz). Default is 16000.
-        sampling_rate_org (int, optional): The original sampling rate of the audio data before any resampling (in Hz), if applicable. Default is None.
-        padding (str): The padding mode for the input signal. 'center' pads the signal symmetrically around its center. Default is 'center'.
- 
-    Returns:
-        torch.Tensor: Mel spectrogram.
-    """
-    def __init__(self, 
-                 filter_length=1024,
-                 hop_length=160,
-                 win_length=640,
-                 n_mel_channels=80,
-                 mel_fmin=0,
-                 mel_fmax=8000,
-                 sampling_rate=16000,
-                 sampling_rate_org=None,
-                 padding='center',
-                 use_db = False,
-                 ):
-        super().__init__()
-        if padding not in ["center", "same"]:
-            raise ValueError("Padding must be 'center' or 'same'.")
-        self.padding = padding
-    
-        self.filter_length = filter_length
-        self.hop_length = hop_length
-        self.win_length = win_length
-        self.n_mel_channels = n_mel_channels
-        self.mel_fmin = mel_fmin
-        self.mel_fmax = mel_fmax
-        self.sampling_rate = sampling_rate
-        self.sampling_rate_org = sampling_rate_org if sampling_rate_org is not None else sampling_rate
-        self.mel_basis = {}
-        self.hann_window = {}
-
-    def forward(self, audio: torch.Tensor, **kwargs) -> torch.Tensor:
-        with torch.no_grad():
-            feats = self.extract(audio, **kwargs) 
-        return feats
-    
-    def extract(self, audio, **kwargs):
-
-        if len(audio.shape) == 3:
-            audio = audio.squeeze(1) if audio.shape[1] == 1 else audio.squeeze(2)
-        assert len(audio.shape) == 2
-
-        y = audio
-        if len(list(self.mel_basis.keys())) == 0:
-            mel = librosa_mel_fn(sr=self.sampling_rate, n_fft=self.filter_length, n_mels=self.n_mel_channels, fmin=self.mel_fmin, fmax=self.mel_fmax)
-            self.mel_basis[str(self.mel_fmax)+'_'+str(y.device)] = torch.from_numpy(mel).float().to(y.device)
-            self.hann_window[str(y.device)] = torch.hann_window(self.win_length).to(y.device)
-
-        y = torch.nn.functional.pad(y.unsqueeze(1), (int((self.filter_length-self.hop_length)/2), int((self.filter_length-self.hop_length)/2)), mode='reflect')
-        y = y.squeeze(1)
-
-        spec = torch.stft(y, self.filter_length, hop_length=self.hop_length, win_length=self.win_length, window=self.hann_window[str(y.device)],
-                          center=False, pad_mode='reflect', normalized=False, onesided=True, return_complex=True)
-        spec = torch.view_as_real(spec)
-        spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9))
-
-        spec = torch.matmul(self.mel_basis[str(self.mel_fmax)+'_'+str(y.device)], spec)
-        spec = spectral_normalize_torch(spec)
-    
-        return spec
-        
-
-class XVectorExtractor(nn.Module):
-    def __init__(self, audio_codec_with_xvector):
-        super().__init__()
-        option = onnxruntime.SessionOptions()
-        option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
-        option.intra_op_num_threads = 1
-        providers = ["CPUExecutionProvider"]
-        self.ort_session = onnxruntime.InferenceSession(audio_codec_with_xvector, sess_options=option, providers=providers)
-
-        self.tfm = sox.Transformer()
-        self.tfm.norm(db_level=-6)
-
-        self.mel_ext = MelSpectrogramFeatures(
-            filter_length=1024,
-            hop_length=160,
-            win_length=640,
-            n_mel_channels=80,
-            mel_fmin=0,
-            mel_fmax=8000,
-            sampling_rate=16000
-        )
-
-    def extract_code(self, audio):
-        with torch.no_grad():
-            norm_audio = self.sox_norm(audio)
-
-            norm_audio = torch.from_numpy(copy.deepcopy(norm_audio)).unsqueeze(0)
-            feat = kaldi.fbank(norm_audio,
-                            num_mel_bins=80,
-                            dither=0,
-                            sample_frequency=16000)
-            feat = feat - feat.mean(dim=0, keepdim=True)
-            norm_embedding = self.ort_session.run(None, {self.ort_session.get_inputs()[0].name: feat.unsqueeze(dim=0).cpu().numpy()})[0].flatten()
-            norm_embedding = F.normalize(torch.from_numpy(norm_embedding), dim=0)
-            
-            ref_mel = self.mel_ext.extract(audio=norm_audio)
-        
-        return norm_embedding.numpy(), ref_mel.permute(0,2,1).squeeze(0).numpy()
-    
-    def sox_norm(self, audio):
-        wav_norm = self.tfm.build_array(input_array=audio, sample_rate_in=16000)
-        return wav_norm
-
-
-class WhisperEncoderVQ(WhisperEncoder):
-    def __init__(
-            self,
-            n_mels: int,
-            n_ctx: int,
-            n_state: int,
-            n_head: int,
-            n_layer: int,
-            n_window: int = 1500,
-            output_dim: int = 512,
-            grad_checkpointing: bool = False,
-            enable_mp: bool = False,
-            audio_sequence_parallel: bool = False,
-            audio_vq_layers: int = -1,
-            audio_vq_type: str = "NULL",
-            audio_vq_codebook_size: int = 4096,
-            audio_vq_pe: bool = False,
-            audio_vq_commit_loss: float = 0.0,
-            audio_vq_out_commit_loss: float = 0.0,
-            audio_vq_no_quantize: bool = False,
-            audio_vq_ff_layer: int = 0,
-            audio_vq_threshold_ema_dead_code: float = 0.1,
-            audio_vq_codebook_dim: int = None,
-            audio_vq_ds_rate: int = None,
-    ):
-        super().__init__(n_mels, n_ctx, n_state, n_head, n_layer, n_window, output_dim, grad_checkpointing, enable_mp, audio_sequence_parallel)
-
-        self.audio_vq_layers = audio_vq_layers
-        self.audio_vq_type = audio_vq_type
-        self.audio_vq_codebook_size = audio_vq_codebook_size
-        self.audio_vq_pe = audio_vq_pe
-        self.audio_vq_commit_loss = audio_vq_commit_loss
-        self.audio_vq_out_commit_loss = audio_vq_out_commit_loss
-        self.audio_vq_no_quantize = audio_vq_no_quantize
-        self.audio_vq_ff_layer = audio_vq_ff_layer
-
-        if audio_vq_layers > 0:
-            self.vq_feature_dim = self.n_state
-            self.audio_vq_ds_rate = 1
-        else:
-            raise NotImplementedError(f"Unsupported audio_vq_layers: {audio_vq_layers}")
-        
-        if self.audio_vq_ds_rate == audio_vq_ds_rate:
-            self.audio_vq_downsample = nn.Identity()
-            self.audio_vq_upsample   = nn.Identity()
-        else:
-            assert audio_vq_ds_rate % self.audio_vq_ds_rate == 0
-            stride = audio_vq_ds_rate // self.audio_vq_ds_rate
-            self.audio_vq_downsample = Conv1d(self.vq_feature_dim, self.vq_feature_dim, kernel_size=stride, stride=stride)
-            self.audio_vq_upsample = ConvTranspose1d(self.vq_feature_dim, self.vq_feature_dim, kernel_size=stride, stride=stride)
-            self.audio_vq_ds_rate = audio_vq_ds_rate
-        
-        if audio_vq_type == "GRVQ":
-            self.audio_quantizer = DistributedGroupResidualVectorQuantization(
-                codebook_size = audio_vq_codebook_size,
-                dim = self.vq_feature_dim, 
-                codebook_dim = self.vq_codebook_dim if audio_vq_codebook_dim is None else audio_vq_codebook_dim, 
-                num_groups=1,
-                num_quantizers=1,
-                kmeans_init=False,
-                threshold_ema_dead_code = audio_vq_threshold_ema_dead_code
-            )
-        else:
-            raise NotImplementedError(f"Unsupported audio_vq_type: {audio_vq_type}")
-        
-        if self.audio_vq_pe:
-            self.project_after_vq_pe = nn.Linear(self.n_state, self.n_state)
-
-    def _calc_quantize_activities(self, indices):
-        indices_onehot = F.one_hot(indices.long().flatten(), self.audio_vq_codebook_size).sum(dim=0)
-        vq_num_activities = sum(indices_onehot>0)
-        vq_num_tokens = sum(indices_onehot)
-        return {
-            "vq_num_activities": vq_num_activities,
-            "vq_num_tokens": vq_num_tokens,
-        }
-    
-    def _do_quantize(self, x, pe=None, y=None):
-        """
-            x: torch.Tensor, shape = (T, D)
-            q: torch.Tensor, shape = (T, D)
-            i: torch.Tensor, shape = (T)
-        """
-        if self.audio_vq_out_commit_loss > 0:
-            x_teacher = x.clone()
-        x = x.unsqueeze(0)
-
-        x = self.audio_vq_downsample(x.transpose(1, 2))
-        x = x.transpose(1, 2)
-
-        vq_stats = {}
-
-        if self.audio_vq_type == "GRVQ":
-            if self.training:
-                raise NotImplementedError
-            else:
-                indices = self.audio_quantizer.encode(x)
-                x = self.audio_quantizer.decode(indices)
-                indices = indices.squeeze(2).squeeze(1)
-
-        vq_stats.update(self._calc_quantize_activities(indices))
-
-        x, indices = x.squeeze(0), indices.squeeze(0)
-        if self.audio_vq_pe:
-            x = x + pe
-            x = self.project_after_vq_pe(x)
-        
-        x = self.audio_vq_upsample(x.unsqueeze(0).transpose(1, 2))
-        x = x.transpose(1, 2).squeeze(0)
-
-        if self.audio_vq_out_commit_loss > 0:
-            vq_out_commit_loss = F.mse_loss(x_teacher.detach(), x)
-            vq_stats["vq_out_commit_loss"] = vq_out_commit_loss * self.audio_vq_out_commit_loss
-
-        return x, indices, vq_stats
-
-    def forward(self, x_list: List[Tensor], audio_mellens:List[int], audio_aftercnnlens:List[int], audio_seqlens:List[int], return_indices=False, audio_pitchs=None):
-        """
-        x : torch.Tensor, shape = (n_mels, n_ctx)
-            the mel spectrogram of the audio
-        """
-
-        aftercnn_x_list = []
-        pe_for_vq_list = []
-        for each_x in x_list:
-            each_x_split_list = each_x.split(self.n_window * 2, dim=1)
-            for each_x_split in each_x_split_list:
-                each_x_split = F.gelu(self.conv1(each_x_split))
-                each_x_split = F.gelu(self.conv2(each_x_split))
-                each_x_split = each_x_split.permute(1, 0) # L,D
-
-                each_positional_embedding_split = self.positional_embedding[:each_x_split.shape[0]]
-                aftercnn_x_list.append(each_x_split+each_positional_embedding_split.to(each_x_split.dtype))
-
-                pe_for_vq_split = self.positional_embedding[:each_x_split.shape[0] // self.audio_vq_ds_rate]
-                pe_for_vq_list.append(pe_for_vq_split.to(each_x_split.dtype))
-
-        pe_for_vq = torch.cat(pe_for_vq_list, dim=0)
-        x = torch.cat(aftercnn_x_list, dim=0)
-        src_len = x.size(0)
-
-        output_list = []
-        for item in audio_aftercnnlens:
-            while item > self.n_window:
-                output_list.append(self.n_window)
-                item -= self.n_window
-            output_list.append(item)
-
-        cu_seqlens = list(accumulate(output_list, func=operator.add,initial=0))
-        cu_seqlens = torch.Tensor(cu_seqlens).to(device=x.device, dtype=torch.int32)
-
-        layer_id = 0
-
-        for block in self.blocks:
-            layer_id+=1
-
-            x = block(x, cu_seqlens=cu_seqlens)
-            
-            if self.audio_vq_layers == layer_id: # vq inside encoder
-                x, indices, vq_stats = self._do_quantize(x, pe_for_vq)
-                if return_indices:
-                    return x, indices
-
-        if self.avg_pooler:
-            x_list = x.split(audio_aftercnnlens, dim=0)
-            token_x_list = []
-            for x in x_list:
-                x = x.permute(1, 0)
-                x = self.avg_pooler(x)
-                x = x.permute(1, 0)
-                token_x_list.append(x)
-            x = torch.cat(token_x_list, dim=0)
-
-        x = self.ln_post(x)
-
-        x = self.proj(x)
-
-        output = torch.zeros(
-            (x.size(0) + len(audio_seqlens) * 2, x.size(1)),
-            device=x.device, dtype=x.dtype
-        )
-
-        audio_seqlens_acc = list(accumulate(audio_seqlens, func=operator.add, initial=0))
-        start_ids = torch.tensor(audio_seqlens_acc[:-1], device=x.device, dtype=torch.int32)
-        end_ids = torch.tensor(audio_seqlens_acc[1:], device=x.device, dtype=torch.int32) - 1
-
-        audio_tokens_mask = torch.ones(output.size(0), device=x.device, dtype=torch.bool)
-        audio_tokens_mask[start_ids] = False
-        audio_tokens_mask[end_ids] = False
-        output[start_ids] = self.audio_bos_eos_token.weight[0].to(x.dtype)
-        output[end_ids] = self.audio_bos_eos_token.weight[1].to(x.dtype)
-        output[audio_tokens_mask] = x
-
-        if self.audio_vq_type != "NULL":
-            return output, vq_stats
-        return output

qwen_tts/core/tokenizer_25hz/vq/whisper_encoder.py

@@ -1,406 +0,0 @@
-# coding=utf-8
-# 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 os
-import math
-import torch
-import operator
-
-import numpy as np
-import torch.nn.functional as F
-
-from functools import lru_cache
-from typing import Optional, Union, List
-from torch import nn, Tensor
-from itertools import accumulate
-
-try:
-    from flash_attn.flash_attn_interface import flash_attn_varlen_func as flash_attn_varlen_func
-except ImportError:
-    try:
-        from flash_attn.flash_attn_interface import flash_attn_unpadded_func as flash_attn_varlen_func
-    except ImportError:
-        print("\n********\nWarning: flash-attn is not installed. Will only run the manual PyTorch version. Please install flash-attn for faster inference.\n********\n ")
-        flash_attn_varlen_func = None
-
-
-N_FFT = 400
-HOP_LENGTH = 160
-
-
-@lru_cache(maxsize=None)
-def mel_filters(device, n_mels: int) -> torch.Tensor:
-    """
-    load the mel filterbank matrix for projecting STFT into a Mel spectrogram.
-    Allows decoupling librosa dependency; saved using:
-
-        np.savez_compressed(
-            "mel_filters.npz",
-            mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80),
-            mel_128=librosa.filters.mel(sr=16000, n_fft=400, n_mels=128),
-        )
-    """
-    assert n_mels in {80, 128}, f"Unsupported n_mels: {n_mels}"
-
-    filters_path = os.path.join(os.path.dirname(__file__), "assets", "mel_filters.npz")
-    with np.load(filters_path, allow_pickle=False) as f:
-        return torch.from_numpy(f[f"mel_{n_mels}"]).to(device)
-
-
-def log_mel_spectrogram(
-    audio: Union[str, np.ndarray, torch.Tensor],
-    n_mels: int = 80,
-    padding: int = 0,
-    device: Optional[Union[str, torch.device]] = None,
-):
-    """
-    Compute the log-Mel spectrogram of
-
-    Parameters
-    ----------
-    audio: Union[str, np.ndarray, torch.Tensor], shape = (*)
-        The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz
-
-    n_mels: int
-        The number of Mel-frequency filters, only 80 is supported
-
-    padding: int
-        Number of zero samples to pad to the right
-
-    device: Optional[Union[str, torch.device]]
-        If given, the audio tensor is moved to this device before STFT
-
-    Returns
-    -------
-    torch.Tensor, shape = (80, n_frames)
-        A Tensor that contains the Mel spectrogram
-    """
-    if not torch.is_tensor(audio):
-        audio = torch.from_numpy(audio)
-
-    if device is not None:
-        audio = audio.to(device)
-    if padding > 0:
-        audio = F.pad(audio, (0, padding))
-    window = torch.hann_window(N_FFT).to(audio.device)
-    stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
-    magnitudes = stft[..., :-1].abs() ** 2
-
-    filters = mel_filters(audio.device, n_mels)
-    mel_spec = filters @ magnitudes
-
-    log_spec = torch.clamp(mel_spec, min=1e-10).log10()
-    log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
-    log_spec = (log_spec + 4.0) / 4.0
-    return log_spec
-
-
-def get_T_after_cnn(L_in, dilation=1):
-    for (padding, kernel_size, stride) in eval("[(1,3,1)] + [(1,3,2)] "):
-        L_out = L_in + 2 * padding - dilation * (kernel_size - 1) - 1
-        L_out = 1 + L_out // stride
-        L_in = L_out
-    return L_out
-
-
-def get_mel_audio(audio, padding=False, audio_vq_ds_rate = 1, n_mels = 128):
-    audio_len = len(audio)
-    if padding:
-        reduction = 160 * 2 * audio_vq_ds_rate
-        audio_pad = math.ceil(audio_len / reduction) * reduction - audio_len
-        mel = log_mel_spectrogram(audio, n_mels=n_mels, padding=audio_pad)
-    else:
-        mel = log_mel_spectrogram(audio, n_mels=n_mels)  # [F,T]
-    return mel
-
-
-def sinusoids(length, channels, max_timescale=10000):
-    """Returns sinusoids for positional embedding"""
-    assert channels % 2 == 0
-    log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
-    inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
-    scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
-    return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
-
-
-class Conv1d(nn.Conv1d):
-    def _conv_forward(
-        self, x: Tensor, weight: Tensor, bias: Optional[Tensor]
-    ) -> Tensor:
-        return super()._conv_forward(
-            x, weight.to(x.dtype), None if bias is None else bias.to(x.dtype)
-        )
-
-
-class ConvTranspose1d(nn.ConvTranspose1d):
-    def _conv_forward(
-        self, x: Tensor, weight: Tensor, bias: Optional[Tensor]
-    ) -> Tensor:
-        return super()._conv_forward(
-            x, weight.to(x.dtype), None if bias is None else bias.to(x.dtype)
-        )
-
-
-class Linear(nn.Linear):
-    def forward(self, x: Tensor) -> Tensor:
-        return F.linear(x, self.weight.to(x.dtype), None if self.bias is None else self.bias.to(x.dtype) )
-
-
-class MultiHeadAttention(nn.Module):
-    def __init__(self, n_state: int, n_head: int):
-        super().__init__()
-        self.n_head = n_head
-        self.query = Linear(n_state, n_state)
-        self.key = Linear(n_state, n_state, bias=False)
-        self.value = Linear(n_state, n_state)
-        self.out = Linear(n_state, n_state)
-
-        self.use_flash_attention = True
-
-    def forward(
-        self,
-        x: Tensor,
-        cu_seqlens = None,
-    ):
-        q = self.query(x)
-        k = self.key(x)
-        v = self.value(x)
-        
-        if self.use_flash_attention:
-            if flash_attn_varlen_func is None:
-                x = self.qkv_attention_manual(q, k, v, cu_seqlens=cu_seqlens)
-            else:
-                if q.dtype not in [torch.float16, torch.bfloat16]:
-                    x = self.qkv_attention_manual(q, k, v, cu_seqlens=cu_seqlens)
-                    self.use_flash_attention = False
-                else:
-                    x = self.qkv_flash_attention(q, k, v, cu_seqlens=cu_seqlens)
-        else:
-            x = self.qkv_attention_manual(q, k, v, cu_seqlens=cu_seqlens)
-
-        output = self.out(x)
-        return output
-
-    def qkv_flash_attention(
-        self, q: Tensor, k: Tensor, v: Tensor, cu_seqlens=None
-    ):
-        n_ctx, n_state = q.shape
-        # scale = (n_state // self.n_head) ** -0.25
-        q = q.view(n_ctx, self.n_head, -1)# (batch_size, seqlen, nheads, headdim)
-        k = k.view(n_ctx, self.n_head, -1)
-        v = v.view(n_ctx, self.n_head, -1)
-
-        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
-
-
-        x = flash_attn_varlen_func(
-            q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, dropout_p=0.0
-        )
-        x = x.reshape(n_ctx, n_state)
-        return x
-
-    def qkv_attention_manual(
-        self, q: Tensor, k: Tensor, v: Tensor, cu_seqlens: Tensor
-    ):
-        n_ctx, n_state = q.shape
-        head_dim = n_state // self.n_head
-        scale = head_dim ** -0.5
-
-        q = q.view(n_ctx, self.n_head, head_dim)
-        k = k.view(n_ctx, self.n_head, head_dim)
-        v = v.view(n_ctx, self.n_head, head_dim)
-
-        seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
-        batch_size = len(seqlens)
-        max_seqlen = max(seqlens)
-
-        q_padded = torch.zeros(batch_size, max_seqlen, self.n_head, head_dim, dtype=q.dtype, device=q.device)
-        k_padded = torch.zeros_like(q_padded)
-        v_padded = torch.zeros_like(q_padded)
-
-        for i in range(batch_size):
-            start_idx = cu_seqlens[i]
-            end_idx = cu_seqlens[i+1]
-            seq_len = seqlens[i]
-            q_padded[i, :seq_len] = q[start_idx:end_idx]
-            k_padded[i, :seq_len] = k[start_idx:end_idx]
-            v_padded[i, :seq_len] = v[start_idx:end_idx]
-        
-        q_padded = q_padded.transpose(1, 2)
-        k_padded = k_padded.transpose(1, 2)
-        v_padded = v_padded.transpose(1, 2)
-
-        attn_mask = torch.arange(max_seqlen, device=q.device)[None, :] < torch.tensor(seqlens, device=q.device)[:, None]
-        attn_mask = attn_mask.unsqueeze(1).unsqueeze(2)
-
-        attn_mask = attn_mask.masked_fill(attn_mask == 0, -torch.finfo(q.dtype).max)
-
-        attn_scores = torch.matmul(q_padded, k_padded.transpose(-2, -1)) * scale
-        attn_scores = attn_scores + attn_mask
-        attn_weights = F.softmax(attn_scores, dim=-1)
-
-        context = torch.matmul(attn_weights, v_padded)
-        
-        context = context.transpose(1, 2).contiguous().view(batch_size, max_seqlen, n_state)
-
-        output_packed = torch.cat([context[i, :seqlens[i]] for i in range(batch_size)], dim=0)
-
-        assert output_packed.shape == (n_ctx, n_state)
-        
-        return output_packed
-
-
-class ResidualAttentionBlock(nn.Module):
-    def __init__(self, n_state: int, n_head: int,
-                 enable_mp: bool = False, sequence_parallel: bool = False):
-        super().__init__()
-        n_mlp = n_state * 4
-        self.attn_ln = nn.LayerNorm(n_state)
-        self.mlp_ln = nn.LayerNorm(n_state)
-
-        self.attn = MultiHeadAttention(n_state, n_head)
-        self.mlp = nn.Sequential(
-                Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state)
-            )
-
-    def forward(
-        self,
-        x: Tensor,
-        cu_seqlens = None
-    ):
-        x = x + self.attn(self.attn_ln(x), cu_seqlens=cu_seqlens)
-        x = x + self.mlp(self.mlp_ln(x))
-        return x
-
-
-class WhisperEncoder(nn.Module):
-    def __init__(
-            self,
-            n_mels: int,
-            n_ctx: int,
-            n_state: int,
-            n_head: int,
-            n_layer: int,
-            n_window: int = 1500,
-            output_dim: int = 512,
-            grad_checkpointing: bool = False,
-            enable_mp: bool = False,
-            audio_sequence_parallel: bool = False,
-    ):
-        super().__init__()
-        self.conv1 = Conv1d(n_mels, n_state, kernel_size=3, padding=1)
-        self.conv2 = Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1)
-        self.register_buffer("positional_embedding", sinusoids(n_ctx, n_state))
-        self.n_layer = n_layer
-        self.n_mels = n_mels
-
-        self.blocks = nn.ModuleList(
-            [ResidualAttentionBlock(n_state, n_head, enable_mp=enable_mp, sequence_parallel=audio_sequence_parallel)
-             for _ in range(n_layer)]
-        )
-        self.ln_post = nn.LayerNorm(n_state)
-        self.avg_pooler = nn.AvgPool1d(2, stride=2)
-
-        self.proj = torch.nn.Linear(n_state, output_dim)
-
-        self.audio_bos_eos_token = nn.Embedding(2, output_dim)
-
-        self.output_dim = output_dim
-        self.grad_checkpointing = grad_checkpointing
-        self.enable_mp = enable_mp
-        self.n_head = n_head
-        self.n_state = n_state
-        self.n_window = n_window
-
-        self.audio_sequence_parallel = audio_sequence_parallel
-
-        self.tp_world_size = 1
-
-        self.set_audio_sync()
-
-    def set_audio_sync(self):
-        for name, param in self.named_parameters():
-            if not name.startswith("blocks"):
-                setattr(param, "audio_sync", True)
-
-    def forward(self, x_list: List[Tensor], audio_mellens:List[int], audio_aftercnnlens:List[int], audio_seqlens:List[int]):
-        """
-        x : torch.Tensor, shape = (n_mels, n_ctx)
-            the mel spectrogram of the audio
-        """
-
-        aftercnn_x_list = []
-        for each_x in x_list:
-            each_x_split_list = each_x.split(self.n_window * 2, dim=1)
-            for each_x_split in each_x_split_list:
-                each_x_split = F.gelu(self.conv1(each_x_split))
-                each_x_split = F.gelu(self.conv2(each_x_split))
-                each_x_split = each_x_split.permute(1, 0) # L,D
-                each_positional_embedding_split = self.positional_embedding[:each_x_split.shape[0]]
-                aftercnn_x_list.append(each_x_split+each_positional_embedding_split.to(each_x_split.dtype))
-
-        x = torch.cat(aftercnn_x_list, dim=0)
-        src_len = x.size(0)
-
-        output_list = []
-        for item in audio_aftercnnlens:
-            while item > self.n_window:
-                output_list.append(self.n_window)
-                item -= self.n_window
-            output_list.append(item)
-
-        cu_seqlens = list(accumulate(output_list, func=operator.add,initial=0))
-        cu_seqlens = torch.Tensor(cu_seqlens).to(device=x.device, dtype=torch.int32)
-
-        layer_id = 0
-        for block in self.blocks:
-            layer_id+=1
-            x = block(x, cu_seqlens=cu_seqlens)
-
-        if self.avg_pooler:
-            x_list = x.split(audio_aftercnnlens, dim=0)
-            token_x_list = []
-            for x in x_list:
-                x = x.permute(1, 0)
-                x = self.avg_pooler(x)
-                x = x.permute(1, 0)
-                token_x_list.append(x)
-            x = torch.cat(token_x_list, dim=0)
-
-        x = self.ln_post(x)
-        x = self.proj(x)
-
-        output = torch.zeros(
-            (x.size(0) + len(audio_seqlens) * 2, x.size(1)),
-            device=x.device, dtype=x.dtype
-        )
-
-        audio_seqlens_acc = list(accumulate(audio_seqlens, func=operator.add, initial=0))
-        start_ids = torch.tensor(audio_seqlens_acc[:-1], device=x.device, dtype=torch.int32)
-        end_ids = torch.tensor(audio_seqlens_acc[1:], device=x.device, dtype=torch.int32) - 1
-
-        audio_tokens_mask = torch.ones(output.size(0), device=x.device, dtype=torch.bool)
-        audio_tokens_mask[start_ids] = False
-        audio_tokens_mask[end_ids] = False
-        output[start_ids] = self.audio_bos_eos_token.weight[0].to(x.dtype)
-        output[end_ids] = self.audio_bos_eos_token.weight[1].to(x.dtype)
-        output[audio_tokens_mask] = x
-        return output
-
-    def lock(self, layers: int):
-        self.conv1.requires_grad_(False)
-        self.conv2.requires_grad_(False)
-        for i in range(min(layers, len(self.blocks))):
-            self.blocks[i].requires_grad_(False)

qwen_tts/inference/qwen3_tts_model.py

@@ -1,874 +0,0 @@
-# coding=utf-8
-# 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 dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
-from urllib.parse import urlparse
-
-import librosa
-import numpy as np
-import soundfile as sf
-import torch
-from transformers import AutoConfig, AutoModel, AutoProcessor
-
-from ..core.models import Qwen3TTSConfig, Qwen3TTSForConditionalGeneration, Qwen3TTSProcessor
-
-AudioLike = Union[
-    str,                     # wav path, URL, base64
-    np.ndarray,              # waveform (requires sr)
-    Tuple[np.ndarray, int],  # (waveform, sr)
-]
-
-MaybeList = Union[Any, List[Any]]
-
-
-@dataclass
-class VoiceClonePromptItem:
-    """
-    Container for one sample's voice-clone prompt information that can be fed to the model.
-
-    Fields are aligned with `Qwen3TTSForConditionalGeneration.generate(..., voice_clone_prompt=...)`.
-    """
-    ref_code: Optional[torch.Tensor]                 # (T, Q) or (T,) depending on tokenizer 25Hz/12Hz
-    ref_spk_embedding: torch.Tensor                  # (D,)
-    x_vector_only_mode: bool
-    icl_mode: bool
-    ref_text: Optional[str] = None
-
-
-class Qwen3TTSModel:
-    """
-    A HuggingFace-style wrapper for Qwen3 TTS models (CustomVoice/VoiceDesign/Base) that provides:
-      - from_pretrained() initialization via AutoModel/AutoProcessor
-      - generation APIs for:
-          * CustomVoice: generate_custom_voice()
-          * VoiceDesign: generate_voice_design()
-          * Base: generate_voice_clone() + create_voice_clone_prompt()
-      - consistent output: (wavs: List[np.ndarray], sample_rate: int)
-
-    Notes:
-      - This wrapper expects the underlying model class to be `Qwen3TTSForConditionalGeneration`
-      - Language / speaker validation is done via model methods:
-          model.get_supported_languages(), model.get_supported_speakers()
-    """
-
-    def __init__(self, model: Qwen3TTSForConditionalGeneration, processor, generate_defaults: Optional[Dict[str, Any]] = None):
-        self.model = model
-        self.processor = processor
-        self.generate_defaults = generate_defaults or {}
-
-        self.device = getattr(model, "device", None)
-        if self.device is None:
-            try:
-                self.device = next(model.parameters()).device
-            except StopIteration:
-                self.device = torch.device("cpu")
-
-    @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path: str,
-        **kwargs,
-    ) -> "Qwen3TTSModel":
-        """
-        Load a Qwen3 TTS model and its processor in HuggingFace `from_pretrained` style.
-
-        This method:
-          1) Loads config via AutoConfig (so your side can register model_type -> config/model).
-          2) Loads the model via AutoModel.from_pretrained(...), forwarding `kwargs` unchanged.
-          3) Loads the processor via AutoProcessor.from_pretrained(model_path).
-          4) Loads optional `generate_config.json` from the model directory/repo snapshot if present.
-
-        Args:
-            pretrained_model_name_or_path (str):
-                HuggingFace repo id or local directory of the model.
-            **kwargs:
-                Forwarded as-is into `AutoModel.from_pretrained(...)`.
-                Typical examples: device_map="cuda:0", dtype=torch.bfloat16, attn_implementation="flash_attention_2".
-
-        Returns:
-            Qwen3TTSModel:
-                Wrapper instance containing `model`, `processor`, and generation defaults.
-        """
-        AutoConfig.register("qwen3_tts", Qwen3TTSConfig)
-        AutoModel.register(Qwen3TTSConfig, Qwen3TTSForConditionalGeneration)
-        AutoProcessor.register(Qwen3TTSConfig, Qwen3TTSProcessor)
-
-        model = AutoModel.from_pretrained(pretrained_model_name_or_path, **kwargs)
-        if not isinstance(model, Qwen3TTSForConditionalGeneration):
-            raise TypeError(
-                f"AutoModel returned {type(model)}, expected Qwen3TTSForConditionalGeneration. "
-            )
-
-        processor = AutoProcessor.from_pretrained(pretrained_model_name_or_path, fix_mistral_regex=True,)
-
-        generate_defaults = model.generate_config
-        return cls(model=model, processor=processor, generate_defaults=generate_defaults)
-
-    def _supported_languages_set(self) -> Optional[set]:
-        langs = getattr(self.model, "get_supported_languages", None)
-        if callable(langs):
-            v = langs()
-            if v is None:
-                return None
-            return set([str(x).lower() for x in v])
-        return None
-
-    def _supported_speakers_set(self) -> Optional[set]:
-        spks = getattr(self.model, "get_supported_speakers", None)
-        if callable(spks):
-            v = spks()
-            if v is None:
-                return None
-            return set([str(x).lower() for x in v])
-        return None
-
-    def _validate_languages(self, languages: List[str]) -> None:
-        """
-        Validate that requested languages are supported by the model.
-
-        Args:
-            languages (List[str]): Language names for each sample.
-
-        Raises:
-            ValueError: If any language is not supported.
-        """
-        supported = self._supported_languages_set()
-        if supported is None:
-            return
-
-        bad = []
-        for lang in languages:
-            if lang is None:
-                bad.append(lang)
-                continue
-            if str(lang).lower() not in supported:
-                bad.append(lang)
-        if bad:
-            raise ValueError(f"Unsupported languages: {bad}. Supported: {sorted(supported)}")
-
-    def _validate_speakers(self, speakers: List[Optional[str]]) -> None:
-        """
-        Validate that requested speakers are supported by the Instruct model.
-
-        Args:
-            speakers (List[Optional[str]]): Speaker names for each sample.
-
-        Raises:
-            ValueError: If any speaker is not supported.
-        """
-        supported = self._supported_speakers_set()
-        if supported is None:
-            return
-
-        bad = []
-        for spk in speakers:
-            if spk is None or spk == "":
-                continue
-            if str(spk).lower() not in supported:
-                bad.append(spk)
-        if bad:
-            raise ValueError(f"Unsupported speakers: {bad}. Supported: {sorted(supported)}")
-
-    def _is_probably_base64(self, s: str) -> bool:
-        if s.startswith("data:audio"):
-            return True
-        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:
-        if "," in b64 and b64.strip().startswith("data:"):
-            b64 = b64.split(",", 1)[1]
-        return base64.b64decode(b64)
-
-    def _load_audio_to_np(self, x: str) -> Tuple[np.ndarray, int]:
-        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 = librosa.load(x, sr=None, mono=True)
-
-        if audio.ndim > 1:
-            audio = np.mean(audio, axis=-1)
-
-        return audio.astype(np.float32), int(sr)
-
-    def _normalize_audio_inputs(self, audios: Union[AudioLike, List[AudioLike]]) -> List[Tuple[np.ndarray, int]]:
-        """
-        Normalize audio inputs into a list of (waveform, sr).
-
-        Supported forms:
-          - str: wav path / URL / base64 audio string
-          - (np.ndarray, sr): waveform + sampling rate
-          - list of the above
-
-        Args:
-            audios:
-                Audio input(s).
-
-        Returns:
-            List[Tuple[np.ndarray, int]]:
-                List of (float32 waveform, original sr).
-
-        Raises:
-            ValueError: If a numpy waveform is provided without sr.
-        """
-        if isinstance(audios, list):
-            items = audios
-        else:
-            items = [audios]
-
-        out: List[Tuple[np.ndarray, int]] = []
-        for a in items:
-            if isinstance(a, str):
-                out.append(self._load_audio_to_np(a))
-            elif isinstance(a, tuple) and len(a) == 2 and isinstance(a[0], np.ndarray):
-                out.append((a[0].astype(np.float32), int(a[1])))
-            elif isinstance(a, np.ndarray):
-                raise ValueError("For numpy waveform input, pass a tuple (audio, sr).")
-            else:
-                raise TypeError(f"Unsupported audio input type: {type(a)}")
-        for i, a in enumerate(out):
-            if a[0].ndim > 1:
-                a[0] = np.mean(a[0], axis=-1).astype(np.float32)
-                out[i] = (a[0], a[1])
-        return out
-
-    def _ensure_list(self, x: MaybeList) -> List[Any]:
-        return x if isinstance(x, list) else [x]
-
-    def _build_assistant_text(self, text: str) -> str:
-        return f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
-
-    def _build_ref_text(self, text: str) -> str:
-        return f"<|im_start|>assistant\n{text}<|im_end|>\n"
-
-    def _build_instruct_text(self, instruct: str) -> str:
-        return f"<|im_start|>user\n{instruct}<|im_end|>\n"
-
-    def _tokenize_texts(self, texts: List[str]) -> List[torch.Tensor]:
-        input_ids = []
-        for text in texts:
-            input = self.processor(text=text, return_tensors="pt", padding=True)
-            input_id = input["input_ids"].to(self.device)
-            input_id = input_id.unsqueeze(0) if input_id.dim() == 1 else input_id
-            input_ids.append(input_id)
-        return input_ids
-
-    def _merge_generate_kwargs(
-        self,
-        non_streaming_mode: Optional[bool] = None,
-        do_sample: Optional[bool] = None,
-        top_k: Optional[int] = None,
-        top_p: Optional[float] = None,
-        temperature: Optional[float] = None,
-        repetition_penalty: Optional[float] = None,
-        subtalker_dosample: Optional[bool] = None,
-        subtalker_top_k: Optional[int] = None,
-        subtalker_top_p: Optional[float] = None,
-        subtalker_temperature: Optional[float] = None,
-        max_new_tokens: Optional[int] = None,
-        **kwargs,
-    ) -> Dict[str, Any]:
-        """
-        Merge user-provided generation arguments with defaults from `generate_config.json`.
-
-        Rule:
-          - If the user explicitly passes a value (not None), use it.
-          - Otherwise, use the value from generate_config.json if present.
-          - Otherwise, fall back to the hard defaults.
-
-        Args:
-            non_streaming_mode, do_sample, top_k, top_p, temperature, repetition_penalty,
-            subtalker_dosample, subtalker_top_k, subtalker_top_p, subtalker_temperature, max_new_tokens:
-                Common generation parameters.
-            **kwargs:
-                Other arguments forwarded to model.generate().
-
-        Returns:
-            Dict[str, Any]: Final kwargs to pass into model.generate().
-        """
-        hard_defaults = dict(
-            non_streaming_mode=False,
-            do_sample=True,
-            top_k=50,
-            top_p=1.0,
-            temperature=0.9,
-            repetition_penalty=1.05,
-            subtalker_dosample=True,
-            subtalker_top_k=50,
-            subtalker_top_p=1.0,
-            subtalker_temperature=0.9,
-            max_new_tokens=2048,
-        )
-
-        def pick(name: str, user_val: Any) -> Any:
-            if user_val is not None:
-                return user_val
-            if name in self.generate_defaults:
-                return self.generate_defaults[name]
-            return hard_defaults[name]
-
-        merged = dict(kwargs)
-        merged.update(
-            non_streaming_mode=pick("non_streaming_mode", non_streaming_mode),
-            do_sample=pick("do_sample", do_sample),
-            top_k=pick("top_k", top_k),
-            top_p=pick("top_p", top_p),
-            temperature=pick("temperature", temperature),
-            repetition_penalty=pick("repetition_penalty", repetition_penalty),
-            subtalker_dosample=pick("subtalker_dosample", subtalker_dosample),
-            subtalker_top_k=pick("subtalker_top_k", subtalker_top_k),
-            subtalker_top_p=pick("subtalker_top_p", subtalker_top_p),
-            subtalker_temperature=pick("subtalker_temperature", subtalker_temperature),
-            max_new_tokens=pick("max_new_tokens", max_new_tokens),
-        )
-        return merged
-
-    # voice clone model
-    @torch.inference_mode()
-    def create_voice_clone_prompt(
-        self,
-        ref_audio: Union[AudioLike, List[AudioLike]],
-        ref_text: Optional[Union[str, List[Optional[str]]]] = None,
-        x_vector_only_mode: Union[bool, List[bool]] = False,
-    ) -> List[VoiceClonePromptItem]:
-        """
-        Build voice-clone prompt items from reference audio (and optionally reference text) using Base model.
-
-        Modes:
-          - x_vector_only_mode=True:
-              Only speaker embedding is used to clone voice; ref_text/ref_code are ignored.
-              This is mutually exclusive with ICL.
-          - x_vector_only_mode=False:
-              ICL mode is enabled automatically (icl_mode=True). In this case ref_text is required,
-              because the model continues/conditions on the reference text + reference speech codes.
-
-        Batch behavior:
-          - ref_audio can be a single item or a list.
-          - ref_text and x_vector_only_mode can be scalars or lists.
-          - If any of them are lists with length > 1, lengths must match.
-
-        Audio input:
-          - str: local wav path / URL / base64
-          - (np.ndarray, sr): waveform + sampling rate
-
-        Args:
-            ref_audio:
-                Reference audio(s) used to extract:
-                  - ref_code via `model.speech_tokenizer.encode(...)`
-                  - ref_spk_embedding via `model.extract_speaker_embedding(...)` (resampled to 24k)
-            ref_text:
-                Reference transcript(s). Required when x_vector_only_mode=False (ICL mode).
-            x_vector_only_mode:
-                Whether to use speaker embedding only. If False, ICL mode will be used.
-
-        Returns:
-            List[VoiceClonePromptItem]:
-                List of prompt items that can be converted into `voice_clone_prompt` dict.
-
-        Raises:
-            ValueError:
-                - If x_vector_only_mode=False but ref_text is missing.
-                - If batch lengths mismatch.
-        """
-        if self.model.tts_model_type != "base":
-            raise ValueError(
-                f"model with \ntokenizer_type: {self.model.tokenizer_type}\n"
-                f"tts_model_size: {self.model.tts_model_size}\n"
-                f"tts_model_type: {self.model.tts_model_type}\n"
-                "does not support create_voice_clone_prompt, Please check Model Card or Readme for more details."
-            )
-        
-        ref_audio_list = self._ensure_list(ref_audio)
-        ref_text_list = self._ensure_list(ref_text) if isinstance(ref_text, list) else ([ref_text] * len(ref_audio_list))
-        xvec_list = self._ensure_list(x_vector_only_mode) if isinstance(x_vector_only_mode, list) else ([x_vector_only_mode] * len(ref_audio_list))
-
-        if len(ref_text_list) != len(ref_audio_list) or len(xvec_list) != len(ref_audio_list):
-            raise ValueError(
-                f"Batch size mismatch: ref_audio={len(ref_audio_list)}, ref_text={len(ref_text_list)}, x_vector_only_mode={len(xvec_list)}"
-            )
-
-        normalized = self._normalize_audio_inputs(ref_audio_list)
-
-        ref_wavs_for_code: List[np.ndarray] = []
-        ref_sr_for_code: List[int] = []
-        for wav, sr in normalized:
-            ref_wavs_for_code.append(wav)
-            ref_sr_for_code.append(sr)
-
-        if len(set(ref_sr_for_code)) == 1:
-            enc = self.model.speech_tokenizer.encode(ref_wavs_for_code, sr=ref_sr_for_code[0])
-            ref_codes = enc.audio_codes
-        else:
-            ref_codes = []
-            for wav, sr in normalized:
-                ref_codes.append(self.model.speech_tokenizer.encode(wav, sr=sr).audio_codes[0])
-
-        items: List[VoiceClonePromptItem] = []
-        for i, ((wav, sr), code, rtext, xvec_only) in enumerate(zip(normalized, ref_codes, ref_text_list, xvec_list)):
-            if not xvec_only:
-                if rtext is None or rtext == "":
-                    raise ValueError(f"ref_text is required when x_vector_only_mode=False (ICL mode). Bad index={i}")
-
-            wav_resample = wav
-            if sr != self.model.speaker_encoder_sample_rate:
-                wav_resample = librosa.resample(y=wav_resample.astype(np.float32), 
-                                           orig_sr=int(sr), 
-                                           target_sr=self.model.speaker_encoder_sample_rate)
-
-            spk_emb = self.model.extract_speaker_embedding(audio=wav_resample,
-                                                           sr=self.model.speaker_encoder_sample_rate)
-
-            items.append(
-                VoiceClonePromptItem(
-                    ref_code=None if xvec_only else code,
-                    ref_spk_embedding=spk_emb,
-                    x_vector_only_mode=bool(xvec_only),
-                    icl_mode=bool(not xvec_only),
-                    ref_text=rtext,
-                )
-            )
-        return items
-
-    def _prompt_items_to_voice_clone_prompt(self, items: List[VoiceClonePromptItem]) -> Dict[str, Any]:
-        return dict(
-            ref_code=[it.ref_code for it in items],
-            ref_spk_embedding=[it.ref_spk_embedding for it in items],
-            x_vector_only_mode=[it.x_vector_only_mode for it in items],
-            icl_mode=[it.icl_mode for it in items],
-        )
-
-    # voice clone model
-    @torch.no_grad()
-    def generate_voice_clone(
-        self,
-        text: Union[str, List[str]],
-        language: Union[str, List[str]] = None,
-        ref_audio: Optional[Union[AudioLike, List[AudioLike]]] = None,
-        ref_text: Optional[Union[str, List[Optional[str]]]] = None,
-        x_vector_only_mode: Union[bool, List[bool]] = False,
-        voice_clone_prompt: Optional[Union[Dict[str, Any], List[VoiceClonePromptItem]]] = None,
-        **kwargs,
-    ) -> Tuple[List[np.ndarray], int]:
-        """
-        Voice clone speech using the Base model.
-
-        You can provide either:
-          - (ref_audio, ref_text, x_vector_only_mode) and let this method build the prompt, OR
-          - `VoiceClonePromptItem` returned by `create_voice_clone_prompt`, OR
-          - a list of `VoiceClonePromptItem` returned by `create_voice_clone_prompt`.
-        
-        `ref_audio` Supported forms:
-        - str: wav path / URL / base64 audio string
-        - (np.ndarray, sr): waveform + sampling rate
-        - list of the above
-
-        Input flexibility:
-          - text/language can be scalar or list.
-          - prompt can be single or batch.
-          - If batch mode (len(text)>1), lengths must match.
-
-        Args:
-            text:
-                Text(s) to synthesize.
-            language:
-                Language(s) for each sample.
-            ref_audio:
-                Reference audio(s) for prompt building. Required if voice_clone_prompt is not provided.
-            ref_text:
-                Reference text(s) used for ICL mode (required when x_vector_only_mode=False).
-            x_vector_only_mode:
-                If True, only speaker embedding is used (ignores ref_text/ref_code).
-                If False, ICL mode is used automatically.
-            voice_clone_prompt:
-                list[VoiceClonePromptItem] from `create_voice_clone_prompt`.
-            non_streaming_mode:
-                Using non-streaming text input, this option currently only simulates streaming text input when set to `false`, 
-                rather than enabling true streaming input or streaming generation.
-            do_sample:
-                Whether to use sampling, recommended to be set to `true` for most use cases.
-            top_k:
-                Top-k sampling parameter.
-            top_p:
-                Top-p sampling parameter.
-            temperature:
-                Sampling temperature; higher => more random.
-            repetition_penalty:
-                Penalty to reduce repeated tokens/codes.
-            subtalker_dosample:
-                Sampling switch for the sub-talker (only valid for qwen3-tts-tokenizer-v2) if applicable.
-            subtalker_top_k:
-                Top-k for sub-talker sampling (only valid for qwen3-tts-tokenizer-v2).
-            subtalker_top_p:
-                Top-p for sub-talker sampling (only valid for qwen3-tts-tokenizer-v2).
-            subtalker_temperature:
-                Temperature for sub-talker sampling (only valid for qwen3-tts-tokenizer-v2).
-            max_new_tokens:
-                Maximum number of new codec tokens to generate.
-            **kwargs:
-                Any other keyword arguments supported by HuggingFace Transformers `generate()` can be passed.
-                They will be forwarded to the underlying `Qwen3TTSForConditionalGeneration.generate(...)`.
-
-        Returns:
-            Tuple[List[np.ndarray], int]:
-                (wavs, sample_rate)
-
-        Raises:
-            ValueError:
-                If batch sizes mismatch or required prompt inputs are missing.
-        """
-        if self.model.tts_model_type != "base":
-            raise ValueError(
-                f"model with \ntokenizer_type: {self.model.tokenizer_type}\n"
-                f"tts_model_size: {self.model.tts_model_size}\n"
-                f"tts_model_type: {self.model.tts_model_type}\n"
-                "does not support generate_voice_clone, Please check Model Card or Readme for more details."
-            )
-        
-        texts = self._ensure_list(text)
-        languages = self._ensure_list(language) if isinstance(language, list) else ([language] * len(texts) if language is not None else ["Auto"] * len(texts))
-        if len(languages) == 1 and len(texts) > 1:
-            languages = languages * len(texts)
-        if len(texts) != len(languages):
-            raise ValueError(f"Batch size mismatch: text={len(texts)}, language={len(languages)}")
-
-        self._validate_languages(languages)
-
-        if voice_clone_prompt is None:
-            if ref_audio is None:
-                raise ValueError("Either `voice_clone_prompt` or `ref_audio` must be provided.")
-            prompt_items = self.create_voice_clone_prompt(ref_audio=ref_audio, ref_text=ref_text, x_vector_only_mode=x_vector_only_mode)
-            if len(prompt_items) == 1 and len(texts) > 1:
-                prompt_items = prompt_items * len(texts)
-            if len(prompt_items) != len(texts):
-                raise ValueError(f"Batch size mismatch: prompt={len(prompt_items)}, text={len(texts)}")
-            voice_clone_prompt_dict = self._prompt_items_to_voice_clone_prompt(prompt_items)
-            ref_texts_for_ids = [it.ref_text for it in prompt_items]
-        else:
-            if isinstance(voice_clone_prompt, list):
-                prompt_items = voice_clone_prompt
-                if len(prompt_items) == 1 and len(texts) > 1:
-                    prompt_items = prompt_items * len(texts)
-                if len(prompt_items) != len(texts):
-                    raise ValueError(f"Batch size mismatch: prompt={len(prompt_items)}, text={len(texts)}")
-                voice_clone_prompt_dict = self._prompt_items_to_voice_clone_prompt(prompt_items)
-                ref_texts_for_ids = [it.ref_text for it in prompt_items]
-            else:
-                voice_clone_prompt_dict = voice_clone_prompt
-                ref_texts_for_ids = None
-
-        input_texts = [self._build_assistant_text(t) for t in texts]
-        input_ids = self._tokenize_texts(input_texts)
-
-        ref_ids = None
-        if ref_texts_for_ids is not None:
-            ref_ids = []
-            for i, rt in enumerate(ref_texts_for_ids):
-                if rt is None or rt == "":
-                    ref_ids.append(None)
-                else:
-                    ref_tok = self._tokenize_texts([self._build_ref_text(rt)])[0]
-                    ref_ids.append(ref_tok)
-
-        gen_kwargs = self._merge_generate_kwargs(**kwargs)
-
-        talker_codes_list, _ = self.model.generate(
-            input_ids=input_ids,
-            ref_ids=ref_ids,
-            voice_clone_prompt=voice_clone_prompt_dict,
-            languages=languages,
-            **gen_kwargs,
-        )
-
-        codes_for_decode = []
-        for i, codes in enumerate(talker_codes_list):
-            ref_code_list = voice_clone_prompt_dict.get("ref_code", None)
-            if ref_code_list is not None and ref_code_list[i] is not None:
-                codes_for_decode.append(torch.cat([ref_code_list[i].to(codes.device), codes], dim=0))
-            else:
-                codes_for_decode.append(codes)
-
-        wavs_all, fs = self.model.speech_tokenizer.decode([{"audio_codes": c} for c in codes_for_decode])
-
-        wavs_out: List[np.ndarray] = []
-        for i, wav in enumerate(wavs_all):
-            ref_code_list = voice_clone_prompt_dict.get("ref_code", None)
-            if ref_code_list is not None and ref_code_list[i] is not None:
-                ref_len = int(ref_code_list[i].shape[0])
-                total_len = int(codes_for_decode[i].shape[0])
-                cut = int(ref_len / max(total_len, 1) * wav.shape[0])
-                wavs_out.append(wav[cut:])
-            else:
-                wavs_out.append(wav)
-
-        return wavs_out, fs
-
-    # voice design model
-    @torch.no_grad()
-    def generate_voice_design(
-        self,
-        text: Union[str, List[str]],
-        instruct: Union[str, List[str]],
-        language: Union[str, List[str]] = None,
-        **kwargs,
-    ) -> Tuple[List[np.ndarray], int]:
-        """
-        Generate speech with the VoiceDesign model using natural-language style instructions.
-
-        Args:
-            text:
-                Text(s) to synthesize.
-            language:
-                Language(s) for each sample.
-            instruct:
-                Instruction(s) describing desired voice/style. Empty string is allowed (treated as no instruction).
-            non_streaming_mode:
-                Using non-streaming text input, this option currently only simulates streaming text input when set to `false`, 
-                rather than enabling true streaming input or streaming generation.
-            do_sample:
-                Whether to use sampling, recommended to be set to `true` for most use cases.
-            top_k:
-                Top-k sampling parameter.
-            top_p:
-                Top-p sampling parameter.
-            temperature:
-                Sampling temperature; higher => more random.
-            repetition_penalty:
-                Penalty to reduce repeated tokens/codes.
-            subtalker_dosample:
-                Sampling switch for the sub-talker (only valid for qwen3-tts-tokenizer-v2) if applicable.
-            subtalker_top_k:
-                Top-k for sub-talker sampling (only valid for qwen3-tts-tokenizer-v2).
-            subtalker_top_p:
-                Top-p for sub-talker sampling (only valid for qwen3-tts-tokenizer-v2).
-            subtalker_temperature:
-                Temperature for sub-talker sampling (only valid for qwen3-tts-tokenizer-v2).
-            max_new_tokens:
-                Maximum number of new codec tokens to generate.
-            **kwargs:
-                Any other keyword arguments supported by HuggingFace Transformers `generate()` can be passed.
-                They will be forwarded to the underlying `Qwen3TTSForConditionalGeneration.generate(...)`.
-
-        Returns:
-            Tuple[List[np.ndarray], int]:
-                (wavs, sample_rate)
-        """
-        if self.model.tts_model_type != "voice_design":
-            raise ValueError(
-                f"model with \ntokenizer_type: {self.model.tokenizer_type}\n"
-                f"tts_model_size: {self.model.tts_model_size}\n"
-                f"tts_model_type: {self.model.tts_model_type}\n"
-                "does not support generate_voice_design, Please check Model Card or Readme for more details."
-            )
-        
-        texts = self._ensure_list(text)
-        languages = self._ensure_list(language) if isinstance(language, list) else ([language] * len(texts) if language is not None else ["Auto"] * len(texts))
-        instructs = self._ensure_list(instruct)
-
-        if len(languages) == 1 and len(texts) > 1:
-            languages = languages * len(texts)
-        if len(instructs) == 1 and len(texts) > 1:
-            instructs = instructs * len(texts)
-
-        if not (len(texts) == len(languages) == len(instructs)):
-            raise ValueError(f"Batch size mismatch: text={len(texts)}, language={len(languages)}, instruct={len(instructs)}")
-
-        self._validate_languages(languages)
-
-        input_ids = self._tokenize_texts([self._build_assistant_text(t) for t in texts])
-
-        instruct_ids: List[Optional[torch.Tensor]] = []
-        for ins in instructs:
-            if ins is None or ins == "":
-                instruct_ids.append(None)
-            else:
-                instruct_ids.append(self._tokenize_texts([self._build_instruct_text(ins)])[0])
-
-        gen_kwargs = self._merge_generate_kwargs(**kwargs)
-
-        talker_codes_list, _ = self.model.generate(
-            input_ids=input_ids,
-            instruct_ids=instruct_ids,
-            languages=languages,
-            **gen_kwargs,
-        )
-
-        wavs, fs = self.model.speech_tokenizer.decode([{"audio_codes": c} for c in talker_codes_list])
-        return wavs, fs
-
-    # custom voice model
-    @torch.no_grad()
-    def generate_custom_voice(
-        self,
-        text: Union[str, List[str]],
-        speaker: Union[str, List[str]],
-        language: Union[str, List[str]] = None,
-        instruct: Optional[Union[str, List[str]]] = None,
-        **kwargs,
-    ) -> Tuple[List[np.ndarray], int]:
-        """
-        Generate speech with the CustomVoice model using a predefined speaker id, optionally controlled by instruction text.
-
-        Args:
-            text:
-                Text(s) to synthesize.
-            language:
-                Language(s) for each sample.
-            speaker:
-                Speaker name(s). Will be validated against `model.get_supported_speakers()` (case-insensitive).
-            instruct:
-                Optional instruction(s). If None, treated as empty (no instruction).
-            non_streaming_mode:
-                Using non-streaming text input, this option currently only simulates streaming text input when set to `false`, 
-                rather than enabling true streaming input or streaming generation.
-            do_sample:
-                Whether to use sampling, recommended to be set to `true` for most use cases.
-            top_k:
-                Top-k sampling parameter.
-            top_p:
-                Top-p sampling parameter.
-            temperature:
-                Sampling temperature; higher => more random.
-            repetition_penalty:
-                Penalty to reduce repeated tokens/codes.
-            subtalker_dosample:
-                Sampling switch for the sub-talker (only valid for qwen3-tts-tokenizer-v2) if applicable.
-            subtalker_top_k:
-                Top-k for sub-talker sampling (only valid for qwen3-tts-tokenizer-v2).
-            subtalker_top_p:
-                Top-p for sub-talker sampling (only valid for qwen3-tts-tokenizer-v2).
-            subtalker_temperature:
-                Temperature for sub-talker sampling (only valid for qwen3-tts-tokenizer-v2).
-            max_new_tokens:
-                Maximum number of new codec tokens to generate.
-            **kwargs:
-                Any other keyword arguments supported by HuggingFace Transformers `generate()` can be passed.
-                They will be forwarded to the underlying `Qwen3TTSForConditionalGeneration.generate(...)`.
-
-        Returns:
-            Tuple[List[np.ndarray], int]:
-                (wavs, sample_rate)
-
-        Raises:
-            ValueError:
-                If any speaker/language is unsupported or batch sizes mismatch.
-        """
-        if self.model.tts_model_type != "custom_voice":
-            raise ValueError(
-                f"model with \ntokenizer_type: {self.model.tokenizer_type}\n"
-                f"tts_model_size: {self.model.tts_model_size}\n"
-                f"tts_model_type: {self.model.tts_model_type}\n"
-                "does not support generate_custom_voice, Please check Model Card or Readme for more details."
-            )
-
-        texts = self._ensure_list(text)
-        languages = self._ensure_list(language) if isinstance(language, list) else ([language] * len(texts) if language is not None else ["Auto"] * len(texts))
-        speakers = self._ensure_list(speaker)
-        if self.model.tts_model_size in "0b6": # for 0b6 model, instruct is not supported
-            instruct = None
-        instructs = self._ensure_list(instruct) if isinstance(instruct, list) else ([instruct] * len(texts) if instruct is not None else [""] * len(texts))
-
-        if len(languages) == 1 and len(texts) > 1:
-            languages = languages * len(texts)
-        if len(speakers) == 1 and len(texts) > 1:
-            speakers = speakers * len(texts)
-        if len(instructs) == 1 and len(texts) > 1:
-            instructs = instructs * len(texts)
-
-        if not (len(texts) == len(languages) == len(speakers) == len(instructs)):
-            raise ValueError(
-                f"Batch size mismatch: text={len(texts)}, language={len(languages)}, speaker={len(speakers)}, instruct={len(instructs)}"
-            )
-
-        self._validate_languages(languages)
-        self._validate_speakers(speakers)
-
-        input_ids = self._tokenize_texts([self._build_assistant_text(t) for t in texts])
-
-        instruct_ids: List[Optional[torch.Tensor]] = []
-        for ins in instructs:
-            if ins is None or ins == "":
-                instruct_ids.append(None)
-            else:
-                instruct_ids.append(self._tokenize_texts([self._build_instruct_text(ins)])[0])
-
-        gen_kwargs = self._merge_generate_kwargs(**kwargs)
-
-        talker_codes_list, _ = self.model.generate(
-            input_ids=input_ids,
-            instruct_ids=instruct_ids,
-            languages=languages,
-            speakers=speakers,
-            **gen_kwargs,
-        )
-
-        wavs, fs = self.model.speech_tokenizer.decode([{"audio_codes": c} for c in talker_codes_list])
-        return wavs, fs
-
-
-    def get_supported_speakers(self) -> Optional[List[str]]:
-        """
-        List supported speaker names for the current model.
-
-        This is a convenience wrapper around `model.get_supported_speakers()`.
-        If the underlying model does not expose speaker constraints (returns None),
-        this method also returns None.
-
-        Returns:
-            Optional[List[str]]:
-                - A sorted list of supported speaker names (lowercased), if available.
-                - None if the model does not provide supported speakers.
-        """
-        supported = self._supported_speakers_set()
-        if supported is None:
-            return None
-        return sorted(supported)
-
-
-    def get_supported_languages(self) -> Optional[List[str]]:
-        """
-        List supported language names for the current model.
-
-        This is a convenience wrapper around `model.get_supported_languages()`.
-        If the underlying model does not expose language constraints (returns None),
-        this method also returns None.
-
-        Returns:
-            Optional[List[str]]:
-                - A sorted list of supported language names (lowercased), if available.
-                - None if the model does not provide supported languages.
-        """
-        supported = self._supported_languages_set()
-        if supported is None:
-            return None
-        return sorted(supported)

qwen_tts/inference/qwen3_tts_tokenizer.py

@@ -1,411 +0,0 @@
-# coding=utf-8
-# 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 typing import List, Optional, Tuple, Union
-from urllib.parse import urlparse
-
-import librosa
-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 ..core import (
-    Qwen3TTSTokenizerV1Config,
-    Qwen3TTSTokenizerV1Model,
-    Qwen3TTSTokenizerV2Config,
-    Qwen3TTSTokenizerV2Model,
-)
-
-AudioInput = Union[
-    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()
-
-        AutoConfig.register("qwen3_tts_tokenizer_25hz", Qwen3TTSTokenizerV1Config)
-        AutoModel.register(Qwen3TTSTokenizerV1Config, Qwen3TTSTokenizerV1Model)
-
-        AutoConfig.register("qwen3_tts_tokenizer_12hz", Qwen3TTSTokenizerV2Config)
-        AutoModel.register(Qwen3TTSTokenizerV2Config, Qwen3TTSTokenizerV2Model)
-
-        inst.feature_extractor = AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path)
-        inst.model = AutoModel.from_pretrained(pretrained_model_name_or_path, **kwargs)
-        inst.config = inst.model.config
-
-        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 = librosa.load(x, sr=None, mono=True)
-
-        if audio.ndim > 1:
-            audio = np.mean(audio, axis=-1)
-
-        if sr != target_sr:
-            audio = librosa.resample(y=audio, orig_sr=sr, target_sr=target_sr)
-
-        return audio.astype(np.float32)
-
-    def _normalize_audio_inputs(
-        self,
-        audios: AudioInput,
-        sr: Optional[int],
-    ) -> 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:
-                a = librosa.resample(y=a.astype(np.float32), orig_sr=int(sr), target_sr=target_sr)
-            out.append(a.astype(np.float32))
-        return out
-
-    def encode(
-        self,
-        audios: AudioInput,
-        sr: Optional[int] = None,
-        return_dict: bool = True,
-    ):
-        """
-        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:
-            25Hz:
-                Qwen3TTSTokenizerV1EncoderOutput (if return_dict=True) with fields:
-                  - audio_codes: List[torch.LongTensor] each (codes_len,)
-                  - xvectors:   List[torch.FloatTensor] each (xvector_dim,)
-                  - ref_mels:   List[torch.FloatTensor] each (mel_len, mel_dim)
-            12Hz:
-                Qwen3TTSTokenizerV2EncoderOutput (if return_dict=True) with fields:
-                  - audio_codes: List[torch.LongTensor] each (codes_len, num_quantizers)
-
-            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=0).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())

requirements.txt

@@ -1,17 +0,0 @@
-daggr==0.5.2
-torch
-transformers==4.57.3
-accelerate==1.12.0
-einops
-gradio
-librosa
-torchaudio
-soundfile
-sox
-nagisa==0.2.11
-soynlp==0.0.493
-onnxruntime
-scipy
-torch
-numpy
-huggingface_hub