From a54c323952b2f36d854e9dd04d27394494984ee8 Mon Sep 17 00:00:00 2001 From: Gao Han Date: Thu, 22 Jan 2026 16:57:37 +0800 Subject: [PATCH] [Model] Support Qwen3-TTS model series (#895) Signed-off-by: Gaohan123 --- docs/.nav.yml | 1 + docs/api/README.md | 49 +- docs/mkdocs/hooks/generate_api_readme.py | 8 +- docs/models/supported_models.md | 3 + .../examples/offline_inference/qwen3_tts.md | 94 + .../examples/online_serving/text_to_image.md | 1 + .../offline_inference/qwen3_tts/README.md | 84 + .../offline_inference/qwen3_tts/end2end.py | 381 +++ .../test_omni_input_preprocessor.py | 59 + .../entrypoints/test_omni_new_request_data.py | 49 + .../test_gpu_generation_model_runner.py | 75 + .../core/sched/omni_generation_scheduler.py | 30 + vllm_omni/engine/arg_utils.py | 20 + vllm_omni/engine/input_processor.py | 8 +- vllm_omni/inputs/preprocess.py | 40 + .../models/qwen3_tts/__init__.py | 0 .../qwen3_tts/configuration_qwen3_tts.py | 516 ++++ .../models/qwen3_tts/modeling_qwen3_tts.py | 2296 +++++++++++++++++ .../models/qwen3_tts/processing_qwen3_tts.py | 102 + .../models/qwen3_tts/qwen3_tts.py | 1080 ++++++++ .../models/qwen3_tts/qwen3_tts_tokenizer.py | 417 +++ .../qwen3_tts/tokenizer_12hz/__init__.py | 1 + .../configuration_qwen3_tts_tokenizer_v2.py | 171 ++ .../modeling_qwen3_tts_tokenizer_v2.py | 1005 ++++++++ .../qwen3_tts/tokenizer_25hz/__init__.py | 1 + .../configuration_qwen3_tts_tokenizer_v1.py | 332 +++ .../modeling_qwen3_tts_tokenizer_v1.py | 1525 +++++++++++ .../qwen3_tts/tokenizer_25hz/vq/__init__.py | 1 + .../tokenizer_25hz/vq/assets/mel_filters.npz | Bin 0 -> 4271 bytes .../qwen3_tts/tokenizer_25hz/vq/core_vq.py | 515 ++++ .../qwen3_tts/tokenizer_25hz/vq/speech_vq.py | 403 +++ .../tokenizer_25hz/vq/whisper_encoder.py | 385 +++ vllm_omni/model_executor/models/registry.py | 5 + .../stage_configs/qwen3_tts.yaml | 22 + .../worker/gpu_generation_model_runner.py | 5 +- 35 files changed, 9642 insertions(+), 42 deletions(-) create mode 100644 docs/user_guide/examples/offline_inference/qwen3_tts.md create mode 100644 examples/offline_inference/qwen3_tts/README.md create mode 100644 examples/offline_inference/qwen3_tts/end2end.py create mode 100644 tests/entrypoints/test_omni_input_preprocessor.py create mode 100644 tests/entrypoints/test_omni_new_request_data.py create mode 100644 tests/worker/test_gpu_generation_model_runner.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/__init__.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/processing_qwen3_tts.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_tokenizer.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/__init__.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/configuration_qwen3_tts_tokenizer_v2.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/__init__.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/configuration_qwen3_tts_tokenizer_v1.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/__init__.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/assets/mel_filters.npz create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/whisper_encoder.py create mode 100644 vllm_omni/model_executor/stage_configs/qwen3_tts.yaml diff --git a/docs/.nav.yml b/docs/.nav.yml index be9306371..18e9c907e 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -14,6 +14,7 @@ nav: - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md + - Qwen3-TTS Offline Inference: user_guide/examples/offline_inference/qwen3_tts.md - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md - Online Serving: diff --git a/docs/api/README.md b/docs/api/README.md index a1f070111..258c3e7b5 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -36,7 +36,6 @@ Input data structures for multi-modal inputs. Engine classes for offline and online inference. -- [vllm_omni.diffusion.diffusion_engine.BackgroundResources][] - [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][] - [vllm_omni.engine.AdditionalInformationEntry][] - [vllm_omni.engine.AdditionalInformationPayload][] @@ -57,38 +56,11 @@ Core scheduling and caching components. - [vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler][] - [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][] - [vllm_omni.core.sched.output.OmniNewRequestData][] - -## Model Executor - -Model execution components. - -- [vllm_omni.model_executor.custom_process_mixin.CustomProcessMixin][] -- [vllm_omni.model_executor.models.output_templates.OmniOutput][] -- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.Qwen2_5OmniForConditionalGeneration][] -- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_talker.Qwen2_5OmniTalkerForConditionalGeneration][] -- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniConditionalGenerationMixin][] -- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniThinkerForConditionalGeneration][] -- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavBigVGANModel][] -- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavDiTModel][] -- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavForConditionalGenerationVLLM][] -- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavModel][] -- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2EmbeddingModel][] -- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2ForCausalLM][] -- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2Model][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_moe.Qwen3MoeForCausalLM][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni.Qwen3OmniMoeForConditionalGeneration][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav.Qwen3OmniMoeCode2Wav][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniCodePredictorBaseModel][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniMoeTalkerCodePredictor][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeModel][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerForConditionalGeneration][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerSharedExpertWrapper][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMForCausalLM][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMModel][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeConditionalGenerationMixin][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerForConditionalGeneration][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerMultiModalProcessor][] -- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerProcessingInfo][] +- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][] +- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][] +- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.EuclideanCodebook][] +- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.VectorQuantization][] +- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.preprocess][] ## Configuration @@ -98,6 +70,17 @@ Configuration classes. - [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][] - [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][] - [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][] +- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][] +- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSSpeakerEncoderConfig][] +- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerCodePredictorConfig][] +- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerConfig][] +- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2.Qwen3TTSTokenizerV2Config][] +- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2.Qwen3TTSTokenizerV2DecoderConfig][] +- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1Config][] +- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderBigVGANConfig][] +- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][] +- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][] +- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][] ## Workers diff --git a/docs/mkdocs/hooks/generate_api_readme.py b/docs/mkdocs/hooks/generate_api_readme.py index 57ddcadbf..9c344d3f1 100644 --- a/docs/mkdocs/hooks/generate_api_readme.py +++ b/docs/mkdocs/hooks/generate_api_readme.py @@ -34,10 +34,10 @@ CATEGORIES = { "name": "Core", "description": "Core scheduling and caching components.", }, - "model_executor": { - "name": "Model Executor", - "description": "Model execution components.", - }, + # "model_executor": { + # "name": "Model Executor", + # "description": "Model execution components.", + # }, "config": { "name": "Configuration", "description": "Configuration classes.", diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index cd70019de..e713d7934 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -34,6 +34,9 @@ th { |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | |`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | +|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | +|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | +|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ## List of Supported Models for NPU diff --git a/docs/user_guide/examples/offline_inference/qwen3_tts.md b/docs/user_guide/examples/offline_inference/qwen3_tts.md new file mode 100644 index 000000000..81bcb1c11 --- /dev/null +++ b/docs/user_guide/examples/offline_inference/qwen3_tts.md @@ -0,0 +1,94 @@ +# Qwen3-TTS Offline Inference + +Source . + + +This directory contains an offline demo for running Qwen3 TTS models with vLLM Omni. It builds task-specific inputs and generates WAV files locally. + +## Model Overview + +Qwen3 TTS provides multiple task variants for speech generation: + +- **CustomVoice**: Generate speech with a known speaker identity (speaker ID) and optional instruction. +- **VoiceDesign**: Generate speech from text plus a descriptive instruction that designs a new voice. +- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection. + +## Setup +Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. + +## Quick Start + +Run a single sample for a task: + +``` +python end2end.py --query-type CustomVoice +``` + +Generated audio files are saved to `output_audio/` by default. + +## Task Usage + +### CustomVoice + +Single sample: + +``` +python end2end.py --query-type CustomVoice +``` + +Batch sample (multiple prompts in one run): + +``` +python end2end.py --query-type CustomVoice --use-batch-sample +``` + +### VoiceDesign + +Single sample: + +``` +python end2end.py --query-type VoiceDesign +``` + +Batch sample: + +``` +python end2end.py --query-type VoiceDesign --use-batch-sample +``` + +### Base (Voice Clone) + +Single sample: + +``` +python end2end.py --query-type Base +``` + +Batch sample: + +``` +python end2end.py --query-type Base --use-batch-sample +``` + +Mode selection for Base: + +- `--mode-tag icl` (default): standard mode +- `--mode-tag xvec_only`: enable `x_vector_only_mode` in the request + +Examples: + +``` +python end2end.py --query-type Base --mode-tag icl +``` + +## Notes + +- The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs. +- Use `--output-dir` (preferred) or `--output-wav` to change the output folder. + +## Example materials + +??? abstract "end2end.py" + ``````py + --8<-- "examples/offline_inference/qwen3_tts/end2end.py" + `````` diff --git a/docs/user_guide/examples/online_serving/text_to_image.md b/docs/user_guide/examples/online_serving/text_to_image.md index 877f20a4c..2de2b4fcb 100644 --- a/docs/user_guide/examples/online_serving/text_to_image.md +++ b/docs/user_guide/examples/online_serving/text_to_image.md @@ -119,6 +119,7 @@ Use `extra_body` to pass generation parameters: | `seed` | int | None | Random seed (reproducible) | | `negative_prompt` | str | None | Negative prompt | | `num_outputs_per_prompt` | int | 1 | Number of images to generate | +| `--cfg-parallel-size`. | int | 1 | Number of GPUs for CFG parallelism | ## Response Format diff --git a/examples/offline_inference/qwen3_tts/README.md b/examples/offline_inference/qwen3_tts/README.md new file mode 100644 index 000000000..266227c4b --- /dev/null +++ b/examples/offline_inference/qwen3_tts/README.md @@ -0,0 +1,84 @@ +# Qwen3-TTS Offline Inference + +This directory contains an offline demo for running Qwen3 TTS models with vLLM Omni. It builds task-specific inputs and generates WAV files locally. + +## Model Overview + +Qwen3 TTS provides multiple task variants for speech generation: + +- **CustomVoice**: Generate speech with a known speaker identity (speaker ID) and optional instruction. +- **VoiceDesign**: Generate speech from text plus a descriptive instruction that designs a new voice. +- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection. + +## Setup +Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. + +## Quick Start + +Run a single sample for a task: + +``` +python end2end.py --query-type CustomVoice +``` + +Generated audio files are saved to `output_audio/` by default. + +## Task Usage + +### CustomVoice + +Single sample: + +``` +python end2end.py --query-type CustomVoice +``` + +Batch sample (multiple prompts in one run): + +``` +python end2end.py --query-type CustomVoice --use-batch-sample +``` + +### VoiceDesign + +Single sample: + +``` +python end2end.py --query-type VoiceDesign +``` + +Batch sample: + +``` +python end2end.py --query-type VoiceDesign --use-batch-sample +``` + +### Base (Voice Clone) + +Single sample: + +``` +python end2end.py --query-type Base +``` + +Batch sample: + +``` +python end2end.py --query-type Base --use-batch-sample +``` + +Mode selection for Base: + +- `--mode-tag icl` (default): standard mode +- `--mode-tag xvec_only`: enable `x_vector_only_mode` in the request + +Examples: + +``` +python end2end.py --query-type Base --mode-tag icl +``` + +## Notes + +- The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs. +- Use `--output-dir` (preferred) or `--output-wav` to change the output folder. diff --git a/examples/offline_inference/qwen3_tts/end2end.py b/examples/offline_inference/qwen3_tts/end2end.py new file mode 100644 index 000000000..7bdd0ae7d --- /dev/null +++ b/examples/offline_inference/qwen3_tts/end2end.py @@ -0,0 +1,381 @@ +"""Offline inference demo for Qwen3 TTS via vLLM Omni. + +Provides single and batch sample inputs for CustomVoice, VoiceDesign, and Base +tasks, then runs Omni generation and saves output wav files. +""" + +import os +from typing import NamedTuple + +import soundfile as sf + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +from vllm import SamplingParams +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni import Omni + + +class QueryResult(NamedTuple): + """Container for a prepared Omni request.""" + + inputs: dict + model_name: str + + +def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult: + """Build CustomVoice sample inputs. + + Args: + use_batch_sample: When True, return a batch of prompts; otherwise a single prompt. + + Returns: + QueryResult with Omni inputs and the CustomVoice model path. + """ + task_type = "CustomVoice" + if use_batch_sample: + texts = ["其实我真的有发现，我是一个特别善于观察别人情绪的人。", "She said she would be here by noon."] + instructs = ["", "Very happy."] + languages = ["Chinese", "English"] + speakers = ["Vivian", "Ryan"] + inputs = [] + for text, instruct, language, speaker in zip(texts, instructs, languages, speakers): + prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" + inputs.append( + { + "prompt": prompt, + "additional_information": { + "task_type": [task_type], + "text": [text], + "instruct": [instruct], + "language": [language], + "speaker": [speaker], + "max_new_tokens": [2048], + }, + } + ) + else: + text = "其实我真的有发现，我是一个特别善于观察别人情绪的人。" + language = "Chinese" + speaker = "Vivian" + instruct = "用特别愤怒的语气说" + prompts = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" + inputs = { + "prompt": prompts, + "additional_information": { + "task_type": [task_type], + "text": [text], + "language": [language], + "speaker": [speaker], + "instruct": [instruct], + "max_new_tokens": [2048], + }, + } + return QueryResult( + inputs=inputs, + model_name="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", + ) + + +def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult: + """Build VoiceDesign sample inputs. + + Args: + use_batch_sample: When True, return a batch of prompts; otherwise a single prompt. + + Returns: + QueryResult with Omni inputs and the VoiceDesign model path. + """ + task_type = "VoiceDesign" + if use_batch_sample: + texts = [ + "哥哥，你回来啦，人家等了你好久好久了，要抱抱！", + "It's in the top drawer... wait, it's empty? No way, that's impossible! I'm sure I put it there!", + ] + instructs = [ + "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。", + "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice.", + ] + languages = ["Chinese", "English"] + inputs = [] + for text, instruct, language in zip(texts, instructs, languages): + prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" + inputs.append( + { + "prompt": prompt, + "additional_information": { + "task_type": [task_type], + "text": [text], + "language": [language], + "instruct": [instruct], + "max_new_tokens": [2048], + "non_streaming_mode": [True], + }, + } + ) + else: + text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！" + instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。" + language = "Chinese" + prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" + inputs = { + "prompt": prompt, + "additional_information": { + "task_type": [task_type], + "text": [text], + "language": [language], + "instruct": [instruct], + "max_new_tokens": [2048], + "non_streaming_mode": [True], + }, + } + return QueryResult( + inputs=inputs, + model_name="Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign", + ) + + +def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> QueryResult: + """Build Base (voice clone) sample inputs. + + Args: + use_batch_sample: When True, return a batch of prompts (Case 2). + mode_tag: "icl" or "xvec_only" to control x_vector_only_mode behavior. + + Returns: + QueryResult with Omni inputs and the Base model path. + """ + task_type = "Base" + ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav" + ref_audio_single = ref_audio_path_1 + ref_text_single = ( + "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." + ) + syn_text_single = "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye." + syn_lang_single = "Auto" + x_vector_only_mode = mode_tag == "xvec_only" + if use_batch_sample: + syn_text_batch = [ + "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye.", + "其实我真的有发现，我是一个特别善于观察别人情绪的人。", + ] + syn_lang_batch = ["Chinese", "English"] + inputs = [] + for text, language in zip(syn_text_batch, syn_lang_batch): + prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" + inputs.append( + { + "prompt": prompt, + "additional_information": { + "task_type": [task_type], + "ref_audio": [ref_audio_single], + "ref_text": [ref_text_single], + "text": [text], + "language": [language], + "x_vector_only_mode": [x_vector_only_mode], + "max_new_tokens": [2048], + }, + } + ) + else: + prompt = f"<|im_start|>assistant\n{syn_text_single}<|im_end|>\n<|im_start|>assistant\n" + inputs = { + "prompt": prompt, + "additional_information": { + "task_type": [task_type], + "ref_audio": [ref_audio_single], + "ref_text": [ref_text_single], + "text": [syn_text_single], + "language": [syn_lang_single], + "x_vector_only_mode": [x_vector_only_mode], + "max_new_tokens": [2048], + }, + } + return QueryResult( + inputs=inputs, + model_name="Qwen/Qwen3-TTS-12Hz-1.7B-Base", + ) + + +def main(args): + """Run offline inference with Omni using prepared sample inputs. + + Args: + args: Parsed CLI args from parse_args(). + """ + query_func = query_map[args.query_type] + if args.query_type in {"CustomVoice", "VoiceDesign"}: + query_result = query_func(use_batch_sample=args.use_batch_sample) + elif args.query_type == "Base": + query_result = query_func( + use_batch_sample=args.use_batch_sample, + mode_tag=args.mode_tag, + ) + else: + query_result = query_func() + + model_name = query_result.model_name + omni = Omni( + model=model_name, + stage_configs_path=args.stage_configs_path, + log_stats=args.enable_stats, + stage_init_timeout=args.stage_init_timeout, + ) + + sampling_params = SamplingParams( + temperature=0.9, + top_p=1.0, + top_k=50, + max_tokens=2048, + seed=42, + detokenize=False, + repetition_penalty=1.05, + ) + + sampling_params_list = [ + sampling_params, + ] + + output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav + os.makedirs(output_dir, exist_ok=True) + + omni_generator = omni.generate(query_result.inputs, sampling_params_list) + for stage_outputs in omni_generator: + for output in stage_outputs.request_output: + request_id = output.request_id + audio_tensor = output.multimodal_output["audio"] + output_wav = os.path.join(output_dir, f"output_{request_id}.wav") + audio_samplerate = output.multimodal_output["sr"].item() + # Convert to numpy array and ensure correct format + audio_numpy = audio_tensor.float().detach().cpu().numpy() + + # Ensure audio is 1D (flatten if needed) + if audio_numpy.ndim > 1: + audio_numpy = audio_numpy.flatten() + + # Save audio file with explicit WAV format + sf.write(output_wav, audio_numpy, samplerate=audio_samplerate, format="WAV") + print(f"Request ID: {request_id}, Saved audio to {output_wav}") + + +def parse_args(): + """Parse CLI arguments for offline TTS inference. + + Returns: + argparse.Namespace with CLI options. + """ + parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") + parser.add_argument( + "--query-type", + "-q", + type=str, + default="CustomVoice", + choices=query_map.keys(), + help="Query type.", + ) + parser.add_argument( + "--enable-stats", + action="store_true", + default=False, + help="Enable writing detailed statistics (default: disabled)", + ) + parser.add_argument( + "--stage-init-timeout", + type=int, + default=300, + help="Timeout for initializing a single stage in seconds (default: 300)", + ) + parser.add_argument( + "--batch-timeout", + type=int, + default=5, + help="Timeout for batching in seconds (default: 5)", + ) + parser.add_argument( + "--init-timeout", + type=int, + default=300, + help="Timeout for initializing stages in seconds (default: 300)", + ) + parser.add_argument( + "--shm-threshold-bytes", + type=int, + default=65536, + help="Threshold for using shared memory in bytes (default: 65536)", + ) + parser.add_argument( + "--output-wav", + default="output_audio", + help="[Deprecated] Output wav directory (use --output-dir).", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=1, + help="Number of prompts to generate.", + ) + parser.add_argument( + "--txt-prompts", + type=str, + default=None, + help="Path to a .txt file with one prompt per line (preferred).", + ) + parser.add_argument( + "--stage-configs-path", + type=str, + default=None, + help="Path to a stage configs file.", + ) + parser.add_argument( + "--audio-path", + "-a", + type=str, + default=None, + help="Path to local audio file. If not provided, uses default audio asset.", + ) + parser.add_argument( + "--sampling-rate", + type=int, + default=16000, + help="Sampling rate for audio loading (default: 16000).", + ) + parser.add_argument( + "--log-dir", + type=str, + default="logs", + help="Log directory (default: logs).", + ) + parser.add_argument( + "--py-generator", + action="store_true", + default=False, + help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.", + ) + parser.add_argument( + "--use-batch-sample", + action="store_true", + default=False, + help="Use batch input sample for CustomVoice/VoiceDesign/Base query.", + ) + parser.add_argument( + "--mode-tag", + type=str, + default="icl", + choices=["icl", "xvec_only"], + help="Mode tag for Base query x_vector_only_mode (default: icl).", + ) + + return parser.parse_args() + + +query_map = { + "CustomVoice": get_custom_voice_query, + "VoiceDesign": get_voice_design_query, + "Base": get_base_query, +} + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/tests/entrypoints/test_omni_input_preprocessor.py b/tests/entrypoints/test_omni_input_preprocessor.py new file mode 100644 index 000000000..77c84f06b --- /dev/null +++ b/tests/entrypoints/test_omni_input_preprocessor.py @@ -0,0 +1,59 @@ +from vllm_omni.inputs.preprocess import OmniInputPreprocessor + + +def _make_preprocessor(monkeypatch): + preprocessor = object.__new__(OmniInputPreprocessor) + monkeypatch.setattr(preprocessor, "_truncate_inputs", lambda tokens, tokenization_kwargs=None: tokens) + monkeypatch.setattr( + preprocessor, + "_process_multimodal", + lambda *args, **kwargs: {"prompt_token_ids": [1, 2, 3]}, + ) + monkeypatch.setattr(preprocessor, "_tokenize_prompt", lambda prompt_text, tokenization_kwargs=None: [9, 8, 7]) + return preprocessor + + +def test_process_tokens_keeps_additional_information(monkeypatch): + preprocessor = _make_preprocessor(monkeypatch) + parsed = { + "prompt_token_ids": [1, 2, 3], + "prompt_embeds": "embeds", + "additional_information": {"task": ["tts"], "lang": ["auto"]}, + } + + inputs = OmniInputPreprocessor._process_tokens(preprocessor, parsed) + + assert inputs["prompt_token_ids"] == [1, 2, 3] + assert inputs["prompt_embeds"] == "embeds" + assert inputs["additional_information"] == {"task": ["tts"], "lang": ["auto"]} + + +def test_process_text_keeps_additional_information(monkeypatch): + preprocessor = _make_preprocessor(monkeypatch) + parsed = { + "prompt": "hello", + "prompt_embeds": "embeds", + "additional_information": {"speaker": ["alice"]}, + } + + inputs = OmniInputPreprocessor._process_text(preprocessor, parsed) + + assert inputs["prompt_token_ids"] == [9, 8, 7] + assert inputs["prompt_embeds"] == "embeds" + assert inputs["additional_information"] == {"speaker": ["alice"]} + + +def test_process_text_multimodal_skips_empty_payloads(monkeypatch): + preprocessor = _make_preprocessor(monkeypatch) + parsed = { + "prompt": "hello", + "multi_modal_data": {"image": "fake"}, + "prompt_embeds": None, + "additional_information": None, + } + + inputs = OmniInputPreprocessor._process_text(preprocessor, parsed) + + assert inputs["prompt_token_ids"] == [1, 2, 3] + assert "prompt_embeds" not in inputs + assert "additional_information" not in inputs diff --git a/tests/entrypoints/test_omni_new_request_data.py b/tests/entrypoints/test_omni_new_request_data.py new file mode 100644 index 000000000..81992fb39 --- /dev/null +++ b/tests/entrypoints/test_omni_new_request_data.py @@ -0,0 +1,49 @@ +from types import SimpleNamespace + +import torch + +from vllm_omni.core.sched.output import OmniNewRequestData + + +def test_omni_new_request_data_copies_payloads(): + prompt_embeds = torch.randn(2, 3) + additional_information = { + "speaker": ["test"], + "codes": torch.tensor([1, 2], dtype=torch.int64), + } + request = SimpleNamespace( + request_id="req-1", + prompt_token_ids=[101, 102], + mm_features=None, + sampling_params=None, + pooling_params=None, + num_computed_tokens=0, + lora_request=None, + prompt_embeds=prompt_embeds, + additional_information=additional_information, + ) + + data = OmniNewRequestData.from_request(request, ([0, 1],), prefill_token_ids=[101, 102]) + + assert data.prompt_embeds is prompt_embeds + assert data.additional_information is additional_information + assert data.prefill_token_ids == [101, 102] + + +def test_omni_new_request_data_allows_missing_payloads(): + request = SimpleNamespace( + request_id="req-2", + prompt_token_ids=[201, 202], + mm_features=None, + sampling_params=None, + pooling_params=None, + num_computed_tokens=0, + lora_request=None, + prompt_embeds=None, + additional_information=None, + ) + + data = OmniNewRequestData.from_request(request, ([0],), prefill_token_ids=None) + + assert data.prompt_embeds is None + assert data.additional_information is None diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py new file mode 100644 index 000000000..25ed1ae86 --- /dev/null +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -0,0 +1,75 @@ +import torch + +from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner + + +class _DummyInputBatch: + def __init__(self): + self.req_ids = ["req-1"] + self.req_id_to_index = {"req-1": 0} + self.num_reqs = 1 + self.vocab_size = 10 + + +def _make_runner(multimodal_outputs): + runner = object.__new__(GPUGenerationModelRunner) + runner.execute_model_state = ( + None, + None, + None, + None, + None, + None, + None, + None, + None, + multimodal_outputs, + ) + runner.kv_connector_output = None + runner.input_batch = _DummyInputBatch() + runner.use_async_scheduling = False + runner.device = torch.device("cpu") + runner.supports_mm_inputs = False + return runner + + +def test_sample_tokens_tensor_output(): + multimodal_outputs = torch.randn(1, 2, 3) + runner = _make_runner(multimodal_outputs) + + output = GPUGenerationModelRunner.sample_tokens(runner) + + assert len(output.pooler_output) == 1 + assert output.pooler_output[0]["model_outputs"].shape == (2, 3) + + +def test_sample_tokens_list_output(): + multimodal_outputs = [torch.randn(2, 1)] + runner = _make_runner(multimodal_outputs) + + output = GPUGenerationModelRunner.sample_tokens(runner) + + assert len(output.pooler_output) == 1 + assert output.pooler_output[0]["model_outputs"].shape == (2, 1) + + +def test_sample_tokens_list_allows_none_output(): + multimodal_outputs = [None] + runner = _make_runner(multimodal_outputs) + + output = GPUGenerationModelRunner.sample_tokens(runner) + + assert len(output.pooler_output) == 1 + assert output.pooler_output[0]["model_outputs"] is None + + +def test_sample_tokens_dict_output(): + multimodal_outputs = {"audio": torch.randn(1, 4), "unused": None} + runner = _make_runner(multimodal_outputs) + + output = GPUGenerationModelRunner.sample_tokens(runner) + + assert len(output.pooler_output) == 1 + assert "audio" in output.pooler_output[0] + assert "unused" not in output.pooler_output[0] + assert output.pooler_output[0]["audio"].shape == (1, 4) diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 2d04faeee..d9a4b12d3 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -2,6 +2,7 @@ import time from collections import defaultdict from vllm.distributed.kv_events import KVEventBatch +from vllm.logger import init_logger from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.request_queue import create_request_queue @@ -145,6 +146,35 @@ class OmniGenerationScheduler(VLLMScheduler): # Update internal state (advance num_computed_tokens, free encoder inputs, # etc.) self._update_after_schedule(scheduler_output) + + try: + # Rewrap base NewRequestData entries with OmniNewRequestData, + # enriching with request-level payloads + new_list = [] + for nr in scheduler_output.scheduled_new_reqs: + req_id = getattr(nr, "req_id", None) + request = self.requests.get(req_id) if req_id else None + # Build omni entry preserving all base fields + omni_nr = OmniNewRequestData( + req_id=nr.req_id, + prompt_token_ids=nr.prompt_token_ids, + mm_features=nr.mm_features, + sampling_params=nr.sampling_params, + pooling_params=nr.pooling_params, + block_ids=nr.block_ids, + num_computed_tokens=nr.num_computed_tokens, + lora_request=nr.lora_request, + # Enrich with omni payloads from the live request object + prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None), + additional_information=(getattr(request, "additional_information", None) if request else None), + ) + new_list.append(omni_nr) + + scheduler_output.scheduled_new_reqs = new_list # type: ignore[assignment] + except Exception: + # If anything goes wrong, leave the original output unchanged + init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData") + return scheduler_output """ diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 6b6a9c727..f58ab749f 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -11,11 +11,31 @@ from vllm_omni.config import OmniModelConfig logger = init_logger(__name__) +def _register_omni_hf_configs() -> None: + try: + from transformers import AutoConfig + + from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( + Qwen3TTSConfig, + ) + except Exception as exc: # pragma: no cover - best-effort optional registration + logger.warning("Skipping omni HF config registration due to import error: %s", exc) + return + + try: + AutoConfig.register("qwen3_tts", Qwen3TTSConfig) + except ValueError: + # Already registered elsewhere; ignore. + return + + def register_omni_models_to_vllm(): from vllm.model_executor.models import ModelRegistry from vllm_omni.model_executor.models.registry import _OMNI_MODELS + _register_omni_hf_configs() + supported_archs = ModelRegistry.get_supported_archs() for arch, (mod_folder, mod_relname, cls_name) in _OMNI_MODELS.items(): if arch not in supported_archs: diff --git a/vllm_omni/engine/input_processor.py b/vllm_omni/engine/input_processor.py index 688387452..eb81f38dc 100644 --- a/vllm_omni/engine/input_processor.py +++ b/vllm_omni/engine/input_processor.py @@ -232,8 +232,8 @@ class OmniInputProcessor(InputProcessor): # (direct-transfer path) prompt_embeds_payload: PromptEmbedsPayload | None = None additional_information_payload: AdditionalInformationPayload | None = None - if "prompt_embeds" in decoder_inputs: # type: ignore[operator] - pe: torch.Tensor = decoder_inputs["prompt_embeds"] # type: ignore[index] + pe: torch.Tensor | None = decoder_inputs.get("prompt_embeds") # type: ignore[operator] + if pe is not None: if pe.ndim != 2: raise ValueError("prompt_embeds must be of shape (seq_len, hidden_size)") # Move to CPU and ensure contiguous memory for stable serialization @@ -246,8 +246,8 @@ class OmniInputProcessor(InputProcessor): shape=[int(seq_len), int(hidden_size)], dtype=dtype_str, ) - if "additional_information" in decoder_inputs: # type: ignore[operator] - raw_info: dict[str, Any] = decoder_inputs["additional_information"] # type: ignore[index] # noqa: E501 + raw_info: dict[str, Any] | None = decoder_inputs.get("additional_information") # type: ignore[operator] + if raw_info is not None: entries: dict[str, AdditionalInformationEntry] = {} for key, value in raw_info.items(): if isinstance(value, torch.Tensor): diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 4c1ee5388..4eeadbff4 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -20,6 +20,46 @@ class OmniInputPreprocessor(InputPreprocessor): Supports processing tokens, embeddings, text, and multimodal inputs. """ + def _process_text( + self, + parsed_content: TextPrompt, + tokenization_kwargs: dict[str, Any] | None = None, + *, + mm_uuids: MultiModalUUIDDict | None = None, + ) -> OmniTokenInputs | MultiModalInputs: + prompt_text = parsed_content["prompt"] + + inputs: OmniTokenInputs | MultiModalInputs + if multi_modal_data := parsed_content.get("multi_modal_data"): + inputs = self._process_multimodal( + prompt_text, + multi_modal_data, + parsed_content.get("mm_processor_kwargs") or {}, + tokenization_kwargs=tokenization_kwargs, + mm_uuids=mm_uuids, + ) + prompt_embeds = parsed_content.get("prompt_embeds") + if prompt_embeds is not None: + inputs["prompt_embeds"] = prompt_embeds + additional_information = parsed_content.get("additional_information") + if additional_information is not None: + inputs["additional_information"] = additional_information + else: + prompt_token_ids = self._tokenize_prompt( + prompt_text, + tokenization_kwargs=tokenization_kwargs, + ) + inputs = token_inputs_omni( + prompt_token_ids, + prompt_embeds=parsed_content.get("prompt_embeds"), + additional_information=parsed_content.get("additional_information"), + ) + + if cache_salt := parsed_content.get("cache_salt"): + inputs["cache_salt"] = cache_salt + + return inputs + def _process_tokens( self, parsed_content: TokensPrompt, diff --git a/vllm_omni/model_executor/models/qwen3_tts/__init__.py b/vllm_omni/model_executor/models/qwen3_tts/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py new file mode 100644 index 000000000..dcf9d54ac --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py @@ -0,0 +1,516 @@ +# 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 + + def get_text_config(self, **kwargs): + # vLLM expects text config to expose hidden_size/num_attention_heads. + # For Qwen3 TTS, the talker config is the text model config. + config = self.talker_config + if hasattr(config, "rope_parameters"): + delattr(config, "rope_parameters") + return config + + +__all__ = ["Qwen3TTSConfig", "Qwen3TTSTalkerConfig", "Qwen3TTSSpeakerEncoderConfig"] diff --git a/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py new file mode 100644 index 000000000..ac9844f89 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/modeling_qwen3_tts.py @@ -0,0 +1,2296 @@ +# 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 collections.abc import Callable +from dataclasses import dataclass + +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 .configuration_qwen3_tts import ( + Qwen3TTSConfig, + Qwen3TTSSpeakerEncoderConfig, + Qwen3TTSTalkerCodePredictorConfig, + Qwen3TTSTalkerConfig, +) +from .qwen3_tts_tokenizer import Qwen3TTSTokenizer + +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: torch.Tensor | None, + 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. + 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: torch.Tensor | None, + past_key_values: Cache | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + 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: torch.FloatTensor | None = None + logits: torch.FloatTensor = None + past_key_values: list[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + generation_steps: int | None = 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: torch.Tensor | None, + past_key_values: Cache | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + 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: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + cache_position: torch.LongTensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, # necessary, but kept here for BC + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + 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: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + past_key_values: list[torch.FloatTensor] | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + past_hidden: torch.FloatTensor | None = None + generation_step: int | None = None + trailing_text_hidden: torch.FloatTensor | None = None + tts_pad_embed: torch.FloatTensor | None = 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: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + cache_position: torch.LongTensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, # necessary, but kept here for BC + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + """ + 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_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. + 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: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + cache_position: torch.LongTensor | None = 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: torch.Tensor | None = 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: + 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, 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: list[torch.Tensor] | None = None, + instruct_ids: list[torch.Tensor] | None = None, + ref_ids: list[torch.Tensor] | None = 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: int | None = 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 is 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()] is not 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", +] diff --git a/vllm_omni/model_executor/models/qwen3_tts/processing_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/processing_qwen3_tts.py new file mode 100644 index 000000000..5643a857c --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/processing_qwen3_tts.py @@ -0,0 +1,102 @@ +# 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"] diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py new file mode 100644 index 000000000..927ec0a90 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py @@ -0,0 +1,1080 @@ +# 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 collections.abc import Iterable +from dataclasses import dataclass +from typing import Any +from urllib.parse import urlparse + +import librosa +import numpy as np +import soundfile as sf +import torch +import torch.nn as nn +from transformers import AutoConfig, AutoModel, AutoProcessor +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .configuration_qwen3_tts import Qwen3TTSConfig +from .modeling_qwen3_tts import Qwen3TTSForConditionalGeneration +from .processing_qwen3_tts import Qwen3TTSProcessor + +logger = init_logger(__name__) + +AudioLike = ( + str # wav path, URL, base64 + | np.ndarray # waveform (requires sr) + | tuple[np.ndarray, int] # (waveform, sr) +) + +MaybeList = 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: torch.Tensor | None # (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: str | None = None + + +class Qwen3TTSModelForGeneration(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + model_path = vllm_config.model_config.model + + # Check if flash-attn is installed + try: + import flash_attn # noqa: F401 + + attn_kwargs = {"attn_implementation": "flash_attention_2"} + except ImportError: + logger.warning("Flash-Attn is not installed. Using default PyTorch attention implementation.") + attn_kwargs = {} + + self.model = Qwen3TTSModel.from_pretrained( + model_path, + torch_dtype=torch.bfloat16, + **attn_kwargs, + ) + self.task_type = model_path.split("-")[-1].strip("/") + # Mark that this model produces multimodal outputs + self.have_multimodal_outputs = True + + # Store vllm_config for potential future use + self.vllm_config = vllm_config + + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> OmniOutput: + """ + Forward pass for TTS generation model. + + Args: + input_ids: Input token IDs (required for TTS generation) + positions: Position IDs (not used for TTS, but required by runner) + intermediate_tensors: Intermediate tensors for pipeline parallelism (not used) + inputs_embeds: Input embeddings (not used for TTS, but required by runner) + **kwargs: Additional arguments including task_type, sampling_metadata, etc. + + Returns: + OmniOutput: Contains multimodal outputs with audio tensors + """ + + # Extract additional parameters from kwargs that the generation methods expect + + runtime_additional_information = kwargs.get("runtime_additional_information", [{}]) + if isinstance(runtime_additional_information, list) and len(runtime_additional_information) > 0: + runtime_additional_information = runtime_additional_information[0] + text = runtime_additional_information.pop("text", [""])[0] + # Extract task_type from kwargs, default to "instruct" + task_type = runtime_additional_information.pop("task_type", [self.task_type])[0] + speaker = runtime_additional_information.pop("speaker", ["uncle_fu"])[0] + language = runtime_additional_information.pop("language", ["Auto"])[0] + instruct = runtime_additional_information.pop("instruct", [""])[0] + for key, value in runtime_additional_information.items(): + if isinstance(value, list) and len(value) > 0: + runtime_additional_information[key] = value[0] + + # Call the appropriate generation method based on task_type + if task_type == "CustomVoice": + result = self.model.generate_custom_voice( + text, speaker=speaker, language=language, instruct=instruct, **runtime_additional_information + ) + elif task_type == "VoiceDesign": + result = self.model.generate_voice_design( + text, instruct=instruct, language=language, **runtime_additional_information + ) + elif task_type == "Base": + result = self.model.generate_voice_clone(text, language=language, **runtime_additional_information) + else: + raise ValueError(f"Invalid task type: {task_type}") + + # Convert result to OmniOutput format + return self.make_omni_output(result, **kwargs) + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput: + """ + Make an OmniOutput object from model outputs. + Args: + model_outputs: Model outputs (either OmniOutput, tuple of (audio_tensors, sr), or tensor) + """ + if isinstance(model_outputs, OmniOutput): + return model_outputs + + # Handle tuple format: (audio_tensors, sample_rate) + if isinstance(model_outputs, tuple) and len(model_outputs) == 2: + audio_tensors, sr = model_outputs + # audio_tensors is a list of numpy arrays, convert first one to tensor if needed + if isinstance(audio_tensors, list) and len(audio_tensors) > 0: + # Convert numpy array to tensor if needed + audio_tensor = audio_tensors[0] + if isinstance(audio_tensor, np.ndarray): + audio_tensor = torch.from_numpy(audio_tensor).float() + elif not isinstance(audio_tensor, torch.Tensor): + audio_tensor = torch.tensor(audio_tensor, dtype=torch.float32) + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audio_tensor, "sr": torch.tensor(sr, dtype=torch.int)}, + ) + + # If it's already a tensor, wrap it + if isinstance(model_outputs, torch.Tensor): + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": model_outputs}, + ) + + raise ValueError(f"Unsupported model_outputs type: {type(model_outputs)}") + + def make_empty_intermediate_tensors( + self, batch_size: int, dtype: torch.dtype, device: torch.device + ) -> IntermediateTensors: + """ + Create empty intermediate tensors for pipeline parallelism. + + For TTS generation models, pipeline parallelism is typically not used, + so this returns an empty dict. However, this method is required by the + runner infrastructure. + + Args: + batch_size: Batch size for the intermediate tensors + dtype: Data type for the tensors + device: Device for the tensors + + Returns: + IntermediateTensors: Empty dict (no PP support for TTS models) + """ + # TTS generation models typically don't use pipeline parallelism + # Return empty dict to satisfy the interface + return IntermediateTensors({}) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: Any = None, + is_multimodal: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + """ + Embed input token IDs into embeddings. + + This method is called by the runner when inputs_embeds are needed. + For TTS models, we typically work with input_ids directly, but this + method provides a fallback for cases where embeddings are required. + + Args: + input_ids: Input token IDs + multimodal_embeddings: Optional multimodal embeddings (not used for TTS) + is_multimodal: Optional mask indicating multimodal tokens (not used for TTS) + **kwargs: Additional arguments + + Returns: + torch.Tensor: Embedded representations of input_ids + """ + # For TTS models, we don't have a separate embedding layer exposed, + # so we return a dummy tensor. In practice, TTS models work with + # input_ids directly in the forward pass. + # This is a minimal implementation to bypass the function call. + return torch.zeros( + (input_ids.shape[0], input_ids.shape[1], 1024), # Dummy hidden size + dtype=torch.bfloat16, + device=input_ids.device, + ) + + def embed_multimodal(self, **kwargs: Any) -> Any: + """ + Embed multimodal inputs (e.g., images, audio). + + For TTS models, this is typically not used as they work with text input_ids. + This method provides a stub to satisfy the interface. + + Args: + **kwargs: Multimodal input arguments + + Returns: + None or empty list: TTS models don't use multimodal embeddings + """ + # TTS models work with text input_ids, not multimodal embeddings + # Return None to indicate no multimodal embeddings + return None + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> set[str]: + """Load weights into the wrapped HF model.""" + # params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + loaded_params.add(name) + + return loaded_params + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: Any = None, + ) -> torch.Tensor | None: + """Non-autoregressive TTS models do not compute token logits.""" + return 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: dict[str, Any] | None = 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: Any, + ) -> "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) -> set | None: + 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) -> set | None: + 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[str | None]) -> 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: 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: bool | None = None, + do_sample: bool | None = None, + top_k: int | None = None, + top_p: float | None = None, + temperature: float | None = None, + repetition_penalty: float | None = None, + subtalker_dosample: bool | None = None, + subtalker_top_k: int | None = None, + subtalker_top_p: float | None = None, + subtalker_temperature: float | None = None, + max_new_tokens: int | None = None, + **kwargs: Any, + ) -> 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: AudioLike | list[AudioLike], + ref_text: str | list[str | None] | None = None, + x_vector_only_mode: 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)}, " + f"ref_text={len(ref_text_list)}, " + f"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 == "": + rtext = "For profile run" + logger.warning( + f"ref_text is required when x_vector_only_mode=False (ICL mode). " + f"Bad index={i}. Please check if it is profile run or " + f"you missed to provide ref_text." + ) + # 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: str | list[str], + language: str | list[str] = None, + ref_audio: AudioLike | list[AudioLike] | None = None, + ref_text: str | list[str | None] | None = None, + x_vector_only_mode: bool | list[bool] = False, + voice_clone_prompt: dict[str, Any] | list[VoiceClonePromptItem] | None = None, + **kwargs: Any, + ) -> 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`. + **kwargs: + Additional generation options. Common keys include `non_streaming_mode`, `do_sample`, `top_k`, `top_p`, + `temperature`, `repetition_penalty`, `subtalker_dosample`, `subtalker_top_k`, `subtalker_top_p`, + `subtalker_temperature`, and `max_new_tokens`. Any other keyword arguments supported by HuggingFace + Transformers `generate()` can also be passed and will be forwarded to + `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: + # For profile run + sample_rate = int(self.model.speaker_encoder_sample_rate) + # Use a 1-second silent clip to satisfy padding requirements. + ref_audio = (np.zeros(sample_rate, dtype=np.float32), sample_rate) + logger.warning( + "ref_audio is not provided. Using a 1-second silent clip " + "to satisfy padding requirements. Please check if it is " + "profile run or you missed to provide ref_audio." + ) + 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: str | list[str], + instruct: str | list[str], + language: str | list[str] = None, + **kwargs: Any, + ) -> 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). + **kwargs: + Additional generation options. Common keys include `non_streaming_mode`, `do_sample`, `top_k`, `top_p`, + `temperature`, `repetition_penalty`, `subtalker_dosample`, `subtalker_top_k`, `subtalker_top_p`, + `subtalker_temperature`, and `max_new_tokens`. Any other keyword arguments supported by HuggingFace + Transformers `generate()` can also be passed and will be forwarded to + `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[torch.Tensor | None] = [] + 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: str | list[str], + speaker: str | list[str], + language: str | list[str] = None, + instruct: str | list[str] | None = None, + **kwargs: Any, + ) -> 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). + **kwargs: + Additional generation options. Common keys include `non_streaming_mode`, `do_sample`, `top_k`, `top_p`, + `temperature`, `repetition_penalty`, `subtalker_dosample`, `subtalker_top_k`, `subtalker_top_p`, + `subtalker_temperature`, and `max_new_tokens`. Any other keyword arguments supported by HuggingFace + Transformers `generate()` can also be passed and will be forwarded to + `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)}, " + f"language={len(languages)}, speaker={len(speakers)}, " + f"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[torch.Tensor | None] = [] + 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) -> list[str] | None: + """ + 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) -> list[str] | None: + """ + 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) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_tokenizer.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_tokenizer.py new file mode 100644 index 000000000..e6e502119 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_tokenizer.py @@ -0,0 +1,417 @@ +# Copyright 2026 The Alibaba Qwen team. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import base64 +import io +import urllib.request +from urllib.parse import urlparse + +import 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 .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config +from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2EncoderOutput, + Qwen3TTSTokenizerV2Model, +) +from .tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1 import Qwen3TTSTokenizerV1Config +from .tokenizer_25hz.modeling_qwen3_tts_tokenizer_v1 import ( + Qwen3TTSTokenizerV1EncoderOutput, + Qwen3TTSTokenizerV1Model, +) + +AudioInput = ( + str # wav path, or base64 string + | np.ndarray # 1-D float array + | list[str] + | list[np.ndarray] +) + + +class Qwen3TTSTokenizer: + """ + A wrapper for Qwen3 TTS Tokenizer 25Hz/12Hz with HuggingFace-style loading. + + - from_pretrained(): loads speech tokenizer model via AutoModel and feature_extractor via AutoFeatureExtractor. + - encode(): supports wav path(s), base64 audio string(s), numpy array(s). + - decode(): accepts either the raw model encode output, or a minimal dict/list-of-dicts. + + Notes: + - For numpy array input, you must pass `sr` so the audio can be resampled to model sample rate. + - Returned audio is float32 numpy arrays and the output sample rate. + """ + + def __init__(self): + self.model = None + self.feature_extractor = None + self.config = None + self.device = None + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs) -> "Qwen3TTSTokenizer": + """ + Initialize tokenizer with HuggingFace `from_pretrained` style. + + Args: + pretrained_model_name_or_path (str): + HuggingFace repo id or local directory. + **kwargs (Any): + Forwarded to `AutoModel.from_pretrained(...)` directly. + Typical examples: device_map="cuda:0", dtype=torch.bfloat16, attn_implementation="eager". + + Returns: + Qwen3TTSTokenizer: + Initialized instance with `model`, `feature_extractor`, `config`. + """ + inst = cls() + + 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: int | None, + ) -> list[np.ndarray]: + """ + Normalize all supported input types into a list of 1-D numpy float32 waveforms + at `self.feature_extractor.sampling_rate`. + + Args: + audios (AudioInput): + - str: wav path OR base64 audio string + - np.ndarray: raw waveform (sr must be provided) + - list[str] / list[np.ndarray] + sr (Optional[int]): + Sampling rate for raw numpy input. Required if input is np.ndarray or list[np.ndarray]. + + Returns: + List[np.ndarray]: + List of float32 waveforms resampled to model input SR. + """ + target_sr = int(self.feature_extractor.sampling_rate) + + if isinstance(audios, (str, np.ndarray)): + audios = [audios] + + if len(audios) == 0: + return [] + + if isinstance(audios[0], str): + # wav path list or base64 list + return [self.load_audio(x, target_sr=target_sr) for x in audios] # type: ignore[arg-type] + + # numpy list + if sr is None: + raise ValueError("For numpy waveform input, you must provide `sr` (original sampling rate).") + + out: list[np.ndarray] = [] + for a in audios: # type: ignore[assignment] + if not isinstance(a, np.ndarray): + raise TypeError("Mixed input types are not supported. Use all paths/base64 or all numpy arrays.") + if a.ndim > 1: + a = np.mean(a, axis=-1) + if int(sr) != target_sr: + 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: int | None = None, + return_dict: bool = True, + ) -> ( + Qwen3TTSTokenizerV1EncoderOutput + | Qwen3TTSTokenizerV2EncoderOutput + | tuple[list[torch.Tensor], list[torch.Tensor] | None, list[torch.Tensor] | None] + | tuple[list[torch.Tensor]] + ): + """ + Batch-encode audio into discrete codes (and optional conditioning, depending on 25Hz/12Hz). + + Args: + audios (AudioInput): + Supported forms: + - np.ndarray: waveform (requires sr) + - list[np.ndarray]: waveforms (requires sr) + - str: wav path OR base64 audio string + - list[str]: wav paths and/or base64 strings + sr (Optional[int], default=None): + Original sampling rate for numpy waveform input. + return_dict (bool, default=True): + Forwarded to model.encode(...). If True, returns ModelOutput. + + Returns: + Qwen3TTSTokenizerV1EncoderOutput | Qwen3TTSTokenizerV2EncoderOutput | tuple: + Encoder output or tuple returned by model.encode. If return_dict=True, + returns a model-specific encoder output. For 25Hz models, this includes + audio_codes/xvectors/ref_mels; for 12Hz models, this includes audio_codes. + If return_dict=False, returns the raw tuple from model.encode. + """ + wavs = self._normalize_audio_inputs(audios, sr=sr) + + inputs = self.feature_extractor( + raw_audio=wavs, + sampling_rate=int(self.feature_extractor.sampling_rate), + return_tensors="pt", + ) + inputs = inputs.to(self.device).to(self.model.dtype) + + with torch.inference_mode(): + # model.encode expects (B, T) and (B, T) + enc = self.model.encode( + inputs["input_values"].squeeze(1), + inputs["padding_mask"].squeeze(1), + return_dict=return_dict, + ) + return enc + + def decode( + self, + encoded, + ) -> tuple[list[np.ndarray], int]: + """ + Decode back to waveform. + + Usage: + 1) Pass the raw output of `encode(...)` directly (recommended). + - 25Hz: expects fields audio_codes, xvectors, ref_mels + - 12Hz: expects field audio_codes + 2) Pass a dict or list[dict] (minimal form) for custom pipelines: + - 25Hz dict keys: {"audio_codes", "xvectors", "ref_mels"} + - 12Hz dict keys: {"audio_codes"} + Values can be torch tensors or numpy arrays. + + Args: + encoded (Any): + - ModelOutput returned by `encode()`, OR + - dict, OR + - list[dict] + + Returns: + Tuple[List[np.ndarray], int]: + - wavs: list of 1-D float32 numpy arrays + - sample_rate: int, model output sampling rate + """ + model_type = self.model.get_model_type() + + def _to_tensor(x, dtype=None): + if isinstance(x, torch.Tensor): + return x + x = np.asarray(x) + t = torch.from_numpy(x) + if dtype is not None: + t = t.to(dtype) + return t + + # Normalize `encoded` into the same shapes as the official demo uses. + if hasattr(encoded, "audio_codes"): + # ModelOutput from encode() + audio_codes_list = encoded.audio_codes + xvectors_list = getattr(encoded, "xvectors", None) + ref_mels_list = getattr(encoded, "ref_mels", None) + elif isinstance(encoded, dict): + audio_codes_list = encoded["audio_codes"] + xvectors_list = encoded.get("xvectors", None) + ref_mels_list = encoded.get("ref_mels", None) + elif isinstance(encoded, list): + # list of dicts + audio_codes_list = [e["audio_codes"] for e in encoded] + xvectors_list = [e["xvectors"] for e in encoded] if ("xvectors" in encoded[0]) else None + ref_mels_list = [e["ref_mels"] for e in encoded] if ("ref_mels" in encoded[0]) else None + else: + raise TypeError("`encoded` must be an encode output, a dict, or a list of dicts.") + + # Ensure list form for per-sample tensors + if isinstance(audio_codes_list, torch.Tensor): + # Could be a single sample tensor or an already padded batch tensor. + t = audio_codes_list + if t.dim() == 1: + # 25Hz single sample: (C,) -> (1, C) + t = t.unsqueeze(0) + elif t.dim() == 2: + # 12Hz single sample: (C, Q) -> (1, C, Q) + t = t.unsqueeze(0) + audio_codes_padded = t.to(self.device) + else: + # List[Tensor/np] + audio_codes_list = [_to_tensor(c, dtype=torch.long) for c in audio_codes_list] + audio_codes_padded = pad_sequence(audio_codes_list, batch_first=True, padding_value=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()) diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/__init__.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/__init__.py new file mode 100644 index 000000000..db6952f1d --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/__init__.py @@ -0,0 +1 @@ +# Qwen3 TTS 12Hz tokenizer package. diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/configuration_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/configuration_qwen3_tts_tokenizer_v2.py new file mode 100644 index 000000000..6d3a6499e --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/configuration_qwen3_tts_tokenizer_v2.py @@ -0,0 +1,171 @@ +# 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 import MimiConfig +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +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"] diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py new file mode 100644 index 000000000..fb09f9b88 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py @@ -0,0 +1,1005 @@ +# 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 collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +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]`): + Discrete 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: torch.Tensor | None, + 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: torch.Tensor | None, + past_key_values: Cache | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + 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: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + cache_position: torch.LongTensor | None = None, + **kwargs, + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + """ + 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. + 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: int | None = 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: int | None = None, + output_dimension: int | None = 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_q_semantic (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: Any, + ): + super().__init__() + assert n_q > n_q_semantic, ( + f"Number of quantizers {n_q} must be larger 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: torch.Tensor | None = None, + return_dict: bool | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None] | 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: bool | None = None, + ) -> 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*): + Discrete 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[:length] for a, length in zip(audio_values, audio_lengths)] + + if not return_dict: + return (audio_values,) + + return Qwen3TTSTokenizerV2DecoderOutput(audio_values) + + +__all__ = ["Qwen3TTSTokenizerV2Model", "Qwen3TTSTokenizerV2PreTrainedModel"] diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/__init__.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/__init__.py new file mode 100644 index 000000000..5f49344f0 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/__init__.py @@ -0,0 +1 @@ +# Qwen3 TTS 25Hz tokenizer package. diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/configuration_qwen3_tts_tokenizer_v1.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/configuration_qwen3_tts_tokenizer_v1.py new file mode 100644 index 000000000..74272f936 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/configuration_qwen3_tts_tokenizer_v1.py @@ -0,0 +1,332 @@ +# 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", +] diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py new file mode 100644 index 000000000..bceafc98e --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/modeling_qwen3_tts_tokenizer_v1.py @@ -0,0 +1,1525 @@ +# 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 + +import numpy as np +import torch +from torch import nn +from torch.nn import Parameter +from torch.nn import functional as F +from torch.nn.utils.rnn import pad_sequence +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 .configuration_qwen3_tts_tokenizer_v1 import ( + Qwen3TTSTokenizerV1Config, + Qwen3TTSTokenizerV1DecoderBigVGANConfig, + Qwen3TTSTokenizerV1DecoderConfig, + Qwen3TTSTokenizerV1DecoderDiTConfig, + Qwen3TTSTokenizerV1EncoderConfig, +) +from .vq.speech_vq import WhisperEncoderVQ, XVectorExtractor +from .vq.whisper_encoder import get_mel_audio, get_T_after_cnn + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring +class Qwen3TTSTokenizerV1EncoderOutput(ModelOutput): + r""" + audio_codes (`List[torch.LongTensor]`): + Discrete 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: bool | None = False, + code_embed_uncond: bool | None = None, + apply_cfg: bool | None = 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) -> torch.Tensor: + """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, speeches): + 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 speeches + ] + 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, speeches): + mels = self.speech2mel(speeches) + 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: torch.Tensor | None = None, + return_dict: bool | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None] | 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[:length] for c, length 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: bool | None = None, + ) -> 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*): + Discrete 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[:length] for a, length in zip(audio_values, audio_lengths)] + + if not return_dict: + return (audio_values,) + + return Qwen3TTSTokenizerV1DecoderOutput(audio_values) + + +__all__ = ["Qwen3TTSTokenizerV1Model", "Qwen3TTSTokenizerV1PreTrainedModel"] diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/__init__.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/__init__.py new file mode 100644 index 000000000..d5f4e7196 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/__init__.py @@ -0,0 +1 @@ +# Qwen3 TTS 25Hz vector-quantization package. diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/assets/mel_filters.npz b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/assets/mel_filters.npz new file mode 100644 index 0000000000000000000000000000000000000000..28ea26909dbdfd608aef67afc4d74d7961ae4bb6 GIT binary patch literal 4271 zcmZ`-cQjmYw;lx1g6JcN7QKe3LG%_Oh!VX=^k~teM-XGQ(Mu4$_Y%?jkm$lFBkB+( z3yfKIgF zxGiAhze`A@t->QRNVV!%P+W=o}VHkB) z%g>qyRHfN1IQ4-=`Y@0T9qE#o+;4E3VQ!epW1Xt=ZG`I3U|62t?<>5h*W|9VvJc`KZ+)ghnA**Z~ET21Tjf_f8oe`vy zZQNtlOx?dDhS71hnOus5cqj)hfyF@H&4y?@9z{I#&cf>A+s2~~(I>TQF}SaR3_tqa z(7&ZdN^vR*t<~?{9DEoI>0PL@Sl?wa?Z{rGX`*eEx9Nh=z*J3HZL1*Py4z$TD#+;m zSSW(kcOTe(4hqgib_W6&xx+j~-u(p)Nn6?>a%wHk=h7Ay$%lcGoo;gAY zmVV7|!Nb;w(PlH@c24{ple2Y3<*9J@jE=sfLzwu_BiAFPE$0Axp`^Nq!H}eG0?r-X zFj@Pwp^al*p>K{@_Cz`q#(N0Y=OpZy^ z{P$KjLJuk_Y%I)$mh`b{uOW5C5Xcmxk!gt_Zg zw>}6fkD4zRK9!#ems~H%U$>V;_wK38Zf-baU$S!#i;7!HWsi}GuC>%@?lMdgkUGC& zh9gC?O-5BlS2#}?7x0?eP#bOL(cqE{M%LJD$CZnplD)CgQR#KCttD=dZK+Ck5R52; z*%5hZ+SXU7)8k%Y^_1U>yI*By(INn&+ir-_4$#dUwTlMNyR@iGQIaZ+eiYqucu)CB z#i{Ru1w+aU#}DHSyzjG_9c?ToB_YjU#f;N=qel98WBIjIc1!#ePwRR+(go&-by#}@ z+M+klVke5b@lWfZ+O&|c??YvRe)&W)qAgtc>t-IZtbRTG#X}49_Q$>P%-)=0W_QY-x%DPep2Vm9#ci zyQcCc4p2&dLtV1@rPe!%>Y^#9W8#ZH&}^@wJKT7N;R9A7cEq&;Y2CYvd@R+Mn&b5O zVyfS^*H#kD74=J5uhD)o`TXoX>>Si$!cT?TXRxj2pB)w_ljjhTby&Je;X|BESZZT= zC%G5!-$BJf&a~U78d_3zBjrvrkJ0CCl@Rfcf7I(`VTNPnI^B#B$zOfPW zG&mEd?R0+W<`l08O1dkcWKS8wB!Z*Cs%I1nMs-EeB-uu5?t@PuD3|z>je8DKi#X(B z{Z=Rz{4X%?-UnxnHQtkELIZ&=J;fK_t}yu8|IxG0(85e&K>H3!!~zlhyJrgti~o1i zzBS*jTgdG~Exp#B-T)6A+PB ztD-e`j^@XAx}|L&JSEFkRvS_%3b%m86z02#Hfn{Y+qIqQ_muywgt?roUA7oiS1xBD zFxmDMsj_cbBcn*^rn^KIMP{AlHM`NiVm*D&`z~7FH#hf<$L3HmJ+=NdiY5>W?nKD? z8Ox6{9dKyI1o8a-j9BtV-|=lm`<`v>tR^Cln&x1dMYzu{@wq5KW!#K14_QMnpH5K%Pavag+g6(i8i-#Eq zguc}rH3?BxH4SOqZW#7m*aT(U9-n#_Xn^Q19(}eH!xG`nI!GYziVQNcA0)`FDHD%~ zz2$HnxW4BQ{#*@u`dssbAa`|fESn$8i8FdxGZh48_Uf~_Q@tv?4in)6fwSed)k&ITqu|){^(WL~J z?Lb|0ro06J^>f>^2}^e-+$u5bU4IZNfO?75v8lstS15%XYw2ac^pkU34{QhDR(umt zPu~`w2?FP|nn3!RWZ3{?=77@teulahD9*S*k5KmY3*adlM)%{SR~bkZYlx1q@fkE= zI$7+kiw5!ha=dYlO>Z5KgxnZEJsaBm%v#nkX0MN-h%n&KA?N}xU3K3o-3Jpk?ANq2n9&Lh%K_CTvfiN ze>6w~NSSl8$#NEZ^t7h9YOxI=zcAG|a+m6AWei`3Jw7K;b;T${pJa^4RwRt%F>?>M zBmoQqm1`<_W7i!5P~THp-II)Ka^u;=z;}d{;SVj{G_4`9^HaEb!=@Pa;Dw)CH^DjsGxFqmb%o$Bkop$KnH8 zDYN)Bh)5=5!-*|f0Gh4)oZG=TEBr()g^DCtSQhmT3!ZN`Qd-E%@1cE}hm8&Vq5B+C zVF2_O)9IiZ(v(xzTwJIg5|}KVuE(;}|7dVIrT`$d=q_OG|3PY}x*URYkMXXJ6PT1$IFkNyvY_(9UglDi6TaeikPS(!Bnij z;Szn+)I_oxnRz7(WTYTp+IHSWQ?Xd~tQn(Q1r)kThM?NM< z?d6LaBG!H}R$zRy!Ij(}1?xe^+o+!;tqWJ3NgjHl1XNxzusxQ0I#6qzM(_00UPMw* zF*GWW_q&fqAN=uimSKgBu_@jD%MX3hpNY|*4r=e=k1lw2r**IyD(hcq?A+HtUgUy4Dqh5D7|G9q{)TsUj{g~c!xy>9wk^(LiXA4VKGz_zMvJMX#AgsR z34T3hhJ)#&sUaQ1+0PML(?YA~{5?=(MT}X^Vib%};uoI{qGW@wgJ&_M+8S8clsNz2 zPQkxMi`#3+Khwtl>>K>wxc{71{&!qGu&Zzz_wU(7TLTyG){PAu?!cXs?Dp-y0Ekcn AQvd(} literal 0 HcmV?d00001 diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py new file mode 100644 index 000000000..9c103a851 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py @@ -0,0 +1,515 @@ +# 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 math import ceil + +import numpy as np +import torch +import torch.nn.functional as F +from einops import rearrange, repeat +from torch import nn + + +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: int | None = 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.0, + ): + 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: list | None = 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: int | None = 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.0, device=device, dtype=torch.long) + null_loss = torch.full((1,), 0.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: int | None = 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: list | None = 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: int | None = 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: int | None = 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) diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py new file mode 100644 index 000000000..805feb81f --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py @@ -0,0 +1,403 @@ +# 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 copy +import operator +from itertools import accumulate + +import onnxruntime +import sox +import torch +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 torch import Tensor + +from .core_vq import DistributedGroupResidualVectorQuantization +from .whisper_encoder import Conv1d, ConvTranspose1d, WhisperEncoder + + +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'. + """ + + 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_pitches=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) + + 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 diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/whisper_encoder.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/whisper_encoder.py new file mode 100644 index 000000000..c1c3c987a --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/whisper_encoder.py @@ -0,0 +1,385 @@ +# 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 math +import operator +import os +from functools import cache +from itertools import accumulate + +import numpy as np +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +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 + + +@cache +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: str | np.ndarray | torch.Tensor, + n_mels: int = 80, + padding: int = 0, + device: str | torch.device | None = 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: Tensor | None) -> 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: Tensor | None) -> 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) + + 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) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 56bceae41..2d9095925 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -48,6 +48,11 @@ _OMNI_MODELS = { "qwen3_omni_code2wav", "Qwen3OmniMoeCode2Wav", ), + "Qwen3TTSForConditionalGeneration": ( + "qwen3_tts", + "qwen3_tts", + "Qwen3TTSModelForGeneration", + ), } _VLLM_OMNI_MODELS = { diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml new file mode 100644 index 000000000..40799fe7f --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml @@ -0,0 +1,22 @@ +stage_args: + - stage_id: 0 + stage_type: llm # Use llm stage type to launch OmniLLM + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: qwen3_tts + model_arch: Qwen3TTSForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: audio # Final output: audio waveform + gpu_memory_utilization: 0.1 + distributed_executor_backend: "mp" + max_num_batched_tokens: 1000000 + + final_output: true + final_output_type: audio diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 011c74d2e..7d5802483 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -308,8 +308,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None} ) elif isinstance(multimodal_outputs, dict): + mm_payload = {} for key, out in multimodal_outputs.items(): - pooler_output.append({key: out.detach().to("cpu").contiguous() if out is not None else None}) + if out is not None and isinstance(out, torch.Tensor): + mm_payload[key] = out.detach().to("cpu").contiguous() + pooler_output.append(mm_payload) else: raise RuntimeError("Unsupported diffusion output type") output = OmniModelRunnerOutput(