diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 231eb15262ca..9f3a9ee6d184 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -931,6 +931,10 @@ title: PE Audio - local: model_doc/pop2piano title: Pop2Piano + - local: model_doc/s3gen + title: S3Gen + - local: model_doc/s3tokenizer + title: S3Tokenizer - local: model_doc/seamless_m4t title: Seamless-M4T - local: model_doc/seamless_m4t_v2 diff --git a/docs/source/en/model_doc/chatterbox.md b/docs/source/en/model_doc/chatterbox.md new file mode 100644 index 000000000000..7bcc317017ba --- /dev/null +++ b/docs/source/en/model_doc/chatterbox.md @@ -0,0 +1,285 @@ +# Chatterbox + +## Overview + +Chatterbox is a complete text-to-speech (TTS) pipeline that combines three specialized models to convert text directly to natural-sounding speech. It was introduced in the [chatterbox repository](https://github.com/resemble-ai/chatterbox) and provides a unified interface for high-quality voice cloning and speech synthesis. + +The pipeline consists of three main components: + +1. **T3 Model**: Converts text tokens to speech tokens using a language model approach +2. **S3Gen Model**: Generates mel spectrograms from speech tokens using speaker-conditioned Conditional Flow Matching (CFM) +3. **HiFTNet Vocoder**: Converts mel spectrograms to high-fidelity waveforms + +Chatterbox enables zero-shot voice cloning by conditioning the generation on a reference audio sample, allowing you to synthesize speech in any voice from just a few seconds of audio. + +## Model Architecture + +The Chatterbox model follows this end-to-end pipeline: + +``` +Text → Text Tokenizer → Text Tokens +Text Tokens + Style → T3 → Speech Tokens +Speech Tokens + Reference Audio → S3Gen → Mel Spectrograms +Mel Spectrograms → HiFTNet → Waveforms +``` + +### Key Features + +- **End-to-end TTS**: Complete pipeline from text to waveform in a single model +- **Zero-shot voice cloning**: Clone any voice from a short reference audio sample +- **Multilingual support**: Supports both English-only and multilingual configurations +- **High-quality synthesis**: Uses state-of-the-art conditional flow matching and neural vocoding +- **Flexible configuration**: Each component can be configured independently +- **Style control**: Optional style conditioning for expressive speech synthesis + +## Usage + +### Basic Text-to-Speech + +```python +from transformers import ChatterboxModel +import torch +import torchaudio + +# Load model +model = ChatterboxModel.from_pretrained("ResembleAI/chatterbox-hf") +model = model.to("cuda") # or "cpu" +model.eval() + +# Load text tokenizer +model.load_text_tokenizer("path/to/tokenizer.json") + +# Load reference audio for voice cloning +ref_wav, ref_sr = torchaudio.load("reference.wav") + +# Convert to mono if needed +if ref_wav.shape[0] > 1: + ref_wav = ref_wav.mean(dim=0, keepdim=True) + +# Convert to numpy array +ref_audio = ref_wav.squeeze().numpy() + +# Generate speech from text +text = "Hello, this is a text-to-speech demo using Chatterbox." +waveform = model.generate( + text=text, + reference_wav=ref_audio, + reference_sr=ref_sr, + exaggeration=0.5, + temperature=0.8, + top_p=0.95, + min_p=0.05, + repetition_penalty=1.2, + cfg_weight=0.5, + max_new_tokens=1000, +) + +# Save output +torchaudio.save("output.wav", waveform.cpu().unsqueeze(0), 24000) +``` + +### Advanced: Two-Stage Generation + +For more control, you can prepare the conditionals and run stages separately: + +```python +import numpy as np + +# Prepare conditionals (speaker embeddings and prompts) +conds = model.prepare_conditionals( + reference_wav=ref_audio, # numpy array + reference_sr=ref_sr, + exaggeration=0.5 +) + +# Stage 1: Prepare text tokens +text_tokens = model.prepare_text_tokens(text) + +# Stage 2: Generate speech tokens using T3 +with torch.no_grad(): + speech_tokens = model.t3.inference( + t3_cond=conds.t3, + text_tokens=text_tokens[0], + max_new_tokens=1000, + temperature=0.8, + top_p=0.95, + ) + +# Stage 3: Generate waveform using S3Gen +with torch.no_grad(): + waveform, _ = model.s3gen.inference( + speech_tokens=speech_tokens[0], + ref_dict=conds.gen, + finalize=True + ) +``` + +### Pre-computed Conditionals for Batch Generation + +For production use where you're generating multiple utterances with the same voice: + +```python +import numpy as np + +# Prepare conditionals once +conds = model.prepare_conditionals( + reference_wav=ref_audio, + reference_sr=ref_sr, + exaggeration=0.5 +) + +# Generate multiple outputs efficiently +texts = ["First sentence.", "Second sentence.", "Third sentence."] +waveforms = [] + +for text in texts: + # Prepare text tokens + text_tokens = model.prepare_text_tokens(text) + + # Generate speech tokens + with torch.no_grad(): + speech_tokens = model.t3.inference( + t3_cond=conds.t3, + text_tokens=text_tokens[0], + max_new_tokens=1000, + temperature=0.8, + ) + + # Generate waveform with cached embeddings + waveform, _ = model.s3gen.inference( + speech_tokens=speech_tokens[0], + ref_dict=conds.gen, + finalize=True + ) + waveforms.append(waveform.squeeze(0)) +``` + +### Generation with Return Intermediates + +You can retrieve intermediate outputs (text tokens, speech tokens) for debugging: + +```python +waveform, intermediates = model.generate( + text=text, + reference_wav=ref_audio, + reference_sr=ref_sr, + exaggeration=0.5, + temperature=0.8, + return_intermediates=True, +) + +print(f"Text tokens: {intermediates['text_tokens'].shape}") +print(f"Speech tokens: {intermediates['speech_tokens'].shape}") +print(f"Waveform: {waveform.shape}") +``` + +## Model Details + +### Configuration + +The model can be configured via [`ChatterboxConfig`]: + +```python +from transformers import ChatterboxConfig + +# English-only configuration +config = ChatterboxConfig.english_only() + +# Multilingual configuration +config = ChatterboxConfig.multilingual() + +# Custom configuration +config = ChatterboxConfig( + t3_config={"num_layers": 12, "num_heads": 16}, + s3gen_config={"encoder_num_blocks": 6}, + hiftnet_config={"upsample_rates": [5, 5, 4, 2]}, + is_multilingual=False, +) +``` + +### Input Requirements + +For the `generate()` method: +- **Text**: String input (automatically normalized for punctuation via `punc_norm()`) +- **Reference Audio**: NumPy array of shape `(audio_length,)` - mono audio +- **Reference Sample Rate**: Integer (will be resampled internally to 16kHz for speaker encoder and 24kHz for mel extraction) +- **Tokenizer**: Optional - if not provided via parameter, must be loaded via `load_text_tokenizer()` + +Generation parameters: +- **exaggeration** (float, default 0.5): Emotion/expressiveness level (0.0 to 1.0) +- **temperature** (float, default 0.8): Sampling temperature for token generation +- **top_p** (float, default 0.95): Top-p (nucleus) sampling threshold +- **min_p** (float, default 0.05): Minimum probability threshold +- **repetition_penalty** (float, default 1.2): Penalty for repeating tokens +- **cfg_weight** (float, default 0.5): Classifier-free guidance weight +- **max_new_tokens** (int, default 1000): Maximum speech tokens to generate + +### Output + +- **Waveforms**: Float tensor of shape `(batch_size, audio_samples)` at 24kHz sample rate + +### Text Normalization + +The model automatically applies text normalization via the `punc_norm()` function, which: +- Capitalizes the first letter +- Normalizes punctuation (replaces uncommon characters like "…" with ", ") +- Converts colons and semicolons to commas +- Replaces em-dashes and en-dashes with hyphens +- Ensures proper sentence ending (adds period if missing) +- Removes multiple spaces + +Example: +```python +from transformers.models.chatterbox.modeling_chatterbox import punc_norm + +text = "hello world... this is a test: with semicolons; and dashes—like this" +normalized = punc_norm(text) +# Output: "Hello world, this is a test, with semicolons, and dashes-like this." +``` + +## Limitations + +- Requires GPU for real-time performance +- Reference audio quality directly affects output quality +- Reference audio should be 6-10 seconds for best results +- English-only model works best for English text; use multilingual model for other languages +- Text tokenizer must be loaded separately via `load_text_tokenizer()` if not passed as parameter +- Generated speech quality depends on T3 sampling parameters (temperature, top_p, etc.) + +## Citation + +If you use Chatterbox in your research, please cite: + +```bibtex +@misc{chatterbox2025, + title={Chatterbox: High-Quality Text-to-Speech Synthesis}, + author={Resemble AI}, + year={2025}, + publisher={GitHub}, + url={https://github.com/resemble-ai/chatterbox} +} +``` + +## ChatterboxConfig + +[[autodoc]] ChatterboxConfig + - english_only + - multilingual + +## T3Config + +[[autodoc]] T3Config + +## ChatterboxFeatureExtractor + +[[autodoc]] ChatterboxFeatureExtractor + +## ChatterboxModel + +[[autodoc]] ChatterboxModel + - forward + - generate + - prepare_text_tokens + - prepare_conditionals + - load_text_tokenizer + diff --git a/docs/source/en/model_doc/s3gen.md b/docs/source/en/model_doc/s3gen.md new file mode 100644 index 000000000000..dba958d52509 --- /dev/null +++ b/docs/source/en/model_doc/s3gen.md @@ -0,0 +1,166 @@ +# S3Gen + +## Overview + +S3Gen is a complete text-to-speech model that converts speech tokens to waveforms using speaker-conditioned Conditional Flow Matching (CFM) and HiFTNet vocoder. It was introduced in the [CosyVoice2 paper](https://arxiv.org/abs/2409.15939) and is part of the [chatterbox](https://github.com/resemble-ai/chatterbox) TTS. + +The model consists of four main components: + +1. **S3 Tokenizer**: Tokenizes reference audio to extract speech tokens +2. **CAMPPlus Speaker Encoder**: Extracts speaker embeddings from reference audio +3. **CFM Flow Decoder**: Generates mel spectrograms from speech tokens using conditional flow matching +4. **HiFTNet Vocoder**: Converts mel spectrograms to waveforms + +S3Gen enables zero-shot voice cloning by conditioning the generation on a reference audio sample. The model uses a causal architecture suitable for streaming applications. + +## Model Architecture + +The S3Gen model follows this pipeline: + +``` +Reference Audio → S3 Tokenizer + CAMPPlus → Speaker Embeddings +Speech Tokens + Speaker Embeddings → CFM Decoder → Mel Spectrograms +Mel Spectrograms → HiFTNet → Waveforms +``` + +### Key Features + +- **Zero-shot voice cloning**: Clone any voice from a short reference audio sample +- **High-quality synthesis**: Uses conditional flow matching for natural mel spectrogram generation +- **Neural source-filter vocoder**: HiFTNet provides high-fidelity waveform synthesis +- **Causal architecture**: Supports streaming inference +- **Speaker conditioning**: Robust speaker embedding extraction via CAMPPlus encoder + +## Usage + +### Basic Usage + +```python +from transformers import S3GenModel +import torch +import torchaudio + +# Load model +model = S3GenModel.from_pretrained("ResembleAI/s3gen") +model.eval() + +# Load reference audio +ref_wav, ref_sr = torchaudio.load("reference.wav") + +# Create speech tokens (from your TTS frontend or S3 tokenizer) +speech_tokens = torch.randint(0, 6561, (1, 100)) # Example tokens + +# Generate waveform +with torch.no_grad(): + waveform, _ = model.inference( + speech_tokens=speech_tokens, + ref_wav=ref_wav, + ref_sr=ref_sr, + finalize=True + ) + +# Save output +torchaudio.save("output.wav", waveform.cpu(), 24000) +``` + +### Two-Stage Generation + +You can also run the model in two stages for more control: + +```python +# Stage 1: Generate mel spectrograms +with torch.no_grad(): + mel_spectrograms = model( + speech_tokens=speech_tokens, + ref_wav=ref_wav, + ref_sr=ref_sr, + finalize=True + ) + +# Stage 2: Generate waveforms from mels +from transformers import HiFTNetModel + +hiftnet = model.mel2wav # or load separately +cache_source = torch.zeros(1, 1, 0) +with torch.no_grad(): + waveform, _ = hiftnet.inference( + speech_feat=mel_spectrograms, + cache_source=cache_source + ) +``` + +### Pre-computed Reference Embeddings + +For production use, you can pre-compute reference embeddings: + +```python +# Extract reference embeddings once +ref_dict = model.embed_ref(ref_wav, ref_sr) + +# Reuse for multiple generations +with torch.no_grad(): + mel1 = model(tokens1, ref_dict=ref_dict, finalize=True) + mel2 = model(tokens2, ref_dict=ref_dict, finalize=True) +``` + +## Model Details + +### Input Requirements + +- **Speech Tokens**: Integer tensor of shape `(batch_size, sequence_length)` with values in range `[0, 6560]` +- **Reference Audio**: Float tensor of shape `(batch_size, audio_length)` or `(audio_length,)` +- **Reference Sample Rate**: Integer (will be resampled internally to 16kHz for speaker encoder and 24kHz for mel extraction) + +### Output + +- **Mel Spectrograms**: Float tensor of shape `(batch_size, mel_bins, time_steps)` with `mel_bins=80` +- **Waveforms**: Float tensor of shape `(batch_size, audio_samples)` at 24kHz sample rate + +### Configuration + +The model can be configured via [`S3GenConfig`]: + +```python +from transformers import S3GenConfig + +config = S3GenConfig( + vocab_size=6561, + encoder_num_blocks=6, + decoder_num_mid_blocks=12, + sampling_rate=24000, + mel_bins=80, +) +``` + +## Limitations + +- Requires GPU for real-time performance +- Reference audio quality affects output quality +- Token sequence length affects generation time + +## Citation + +```bibtex +@article{cosyvoice2, + title={CosyVoice 2: Scalable Streaming Speech Synthesis with Large Language Models}, + author={Du, Zhihao and others}, + journal={arXiv preprint arXiv:2409.15939}, + year={2024} +} +``` + +## S3GenConfig + +[[autodoc]] S3GenConfig + +## HiFTNetConfig + +[[autodoc]] HiFTNetConfig + +## S3GenModel + +[[autodoc]] S3GenModel + - forward + - inference + - embed_ref + diff --git a/docs/source/en/model_doc/s3tokenizer.md b/docs/source/en/model_doc/s3tokenizer.md new file mode 100644 index 000000000000..b8fbe4c3294c --- /dev/null +++ b/docs/source/en/model_doc/s3tokenizer.md @@ -0,0 +1,72 @@ + + +# S3Tokenizer + +
+PyTorch +
+ +## Overview + +The S3Tokenizer model is a speech tokenizer that converts raw audio into discrete tokens at 25 tokens/second. +It uses a mel-spectrogram encoder with Finite Scalar Quantization (FSQ) to produce high-quality speech representations suitable for speech language models. + +This model was contributed by [xingchensong](https://github.com/xingchensong). +The original code can be found [here](https://github.com/xingchensong/S3Tokenizer). + +## Usage example + +Here is a quick example of how to tokenize audio using this model: + +```python +>>> from transformers import S3TokenizerModel +>>> import torch + +>>> # Load the model +>>> model = S3TokenizerModel.from_pretrained("path/to/model") + +>>> # Prepare audio (16kHz sample rate expected) +>>> audio = torch.randn(1, 16000) # 1 second of audio at 16kHz + +>>> # Tokenize the audio +>>> outputs = model(audio) +>>> speech_tokens = outputs.speech_tokens # Discrete tokens +>>> speech_token_lens = outputs.speech_token_lens # Length of token sequence +``` + +## S3TokenizerConfig + +[[autodoc]] S3TokenizerConfig + +## S3TokenizerFeatureExtractor + +[[autodoc]] S3TokenizerFeatureExtractor + - __call__ + +## S3TokenizerModel + +[[autodoc]] S3TokenizerModel + - forward + +## S3TokenizerOutput + +[[autodoc]] transformers.models.s3tokenizer.modeling_s3tokenizer.S3TokenizerOutput + +## Utilities + +[[autodoc]] transformers.models.s3tokenizer.modeling_s3tokenizer.drop_invalid_tokens + diff --git a/src/transformers/models/__init__.py b/src/transformers/models/__init__.py index 9e2261d2bc8f..1b586a32f418 100644 --- a/src/transformers/models/__init__.py +++ b/src/transformers/models/__init__.py @@ -58,6 +58,7 @@ from .camembert import * from .canine import * from .chameleon import * + from .chatterbox import * from .chinese_clip import * from .clap import * from .clip import * @@ -330,6 +331,8 @@ from .rt_detr import * from .rt_detr_v2 import * from .rwkv import * + from .s3gen import * + from .s3tokenizer import * from .sam import * from .sam2 import * from .sam2_video import * @@ -364,6 +367,7 @@ from .swin2sr import * from .swinv2 import * from .switch_transformers import * + from .t3 import * from .t5 import * from .t5gemma import * from .t5gemma2 import * diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py index 55dd1b820073..f66d30fb307d 100644 --- a/src/transformers/models/auto/configuration_auto.py +++ b/src/transformers/models/auto/configuration_auto.py @@ -73,6 +73,7 @@ ("camembert", "CamembertConfig"), ("canine", "CanineConfig"), ("chameleon", "ChameleonConfig"), + ("chatterbox", "ChatterboxConfig"), ("chinese_clip", "ChineseCLIPConfig"), ("chinese_clip_vision_model", "ChineseCLIPVisionConfig"), ("clap", "ClapConfig"), @@ -366,6 +367,8 @@ ("rt_detr_resnet", "RTDetrResNetConfig"), ("rt_detr_v2", "RTDetrV2Config"), ("rwkv", "RwkvConfig"), + ("s3gen", "S3GenConfig"), + ("s3tokenizer", "S3TokenizerConfig"), ("sam", "SamConfig"), ("sam2", "Sam2Config"), ("sam2_hiera_det_model", "Sam2HieraDetConfig"), @@ -519,6 +522,7 @@ ("camembert", "CamemBERT"), ("canine", "CANINE"), ("chameleon", "Chameleon"), + ("chatterbox", "Chatterbox"), ("chinese_clip", "Chinese-CLIP"), ("chinese_clip_vision_model", "ChineseCLIPVisionModel"), ("clap", "CLAP"), @@ -834,6 +838,8 @@ ("rt_detr_resnet", "RT-DETR-ResNet"), ("rt_detr_v2", "RT-DETRv2"), ("rwkv", "RWKV"), + ("s3gen", "S3Gen"), + ("s3tokenizer", "S3Tokenizer"), ("sam", "SAM"), ("sam2", "SAM2"), ("sam2_hiera_det_model", "Sam2HieraDetModel"), diff --git a/src/transformers/models/auto/feature_extraction_auto.py b/src/transformers/models/auto/feature_extraction_auto.py index b36fe91f720a..aa34f53544bd 100644 --- a/src/transformers/models/auto/feature_extraction_auto.py +++ b/src/transformers/models/auto/feature_extraction_auto.py @@ -39,6 +39,7 @@ [ ("audio-spectrogram-transformer", "ASTFeatureExtractor"), ("audioflamingo3", "WhisperFeatureExtractor"), + ("chatterbox", "ChatterboxFeatureExtractor"), ("clap", "ClapFeatureExtractor"), ("clvp", "ClvpFeatureExtractor"), ("csm", "EncodecFeatureExtractor"), diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index 8b151b68e1df..13ac0a0f7099 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -81,6 +81,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("camembert", "CamembertModel"), ("canine", "CanineModel"), ("chameleon", "ChameleonModel"), + ("chatterbox", "ChatterboxModel"), ("chinese_clip", "ChineseCLIPModel"), ("chinese_clip_vision_model", "ChineseCLIPVisionModel"), ("clap", "ClapModel"), @@ -1665,6 +1666,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): [ # Model for Text-To-Waveform mapping ("bark", "BarkModel"), + ("chatterbox", "ChatterboxModel"), ("csm", "CsmForConditionalGeneration"), ("fastspeech2_conformer", "FastSpeech2ConformerWithHifiGan"), ("fastspeech2_conformer_with_hifigan", "FastSpeech2ConformerWithHifiGan"), diff --git a/src/transformers/models/chatterbox/__init__.py b/src/transformers/models/chatterbox/__init__.py new file mode 100644 index 000000000000..c58652d5db2e --- /dev/null +++ b/src/transformers/models/chatterbox/__init__.py @@ -0,0 +1,29 @@ +# coding=utf-8 +# Copyright 2025 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 typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_chatterbox import * + from .feature_extraction_chatterbox import * + from .modeling_chatterbox import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/src/transformers/models/chatterbox/configuration_chatterbox.py b/src/transformers/models/chatterbox/configuration_chatterbox.py new file mode 100644 index 000000000000..6a4ed995c154 --- /dev/null +++ b/src/transformers/models/chatterbox/configuration_chatterbox.py @@ -0,0 +1,301 @@ +# coding=utf-8 +# Copyright 2025 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. +"""Chatterbox model configuration""" + +from ...configuration_utils import PretrainedConfig + + +# ============================================================================ +# T3 Configuration +# ============================================================================ + +# Llama 520M configuration for T3 backbone +LLAMA_520M_CONFIG_DICT = { + # Arbitrary small number that won't cause problems when loading. + # These params are unused due to custom input layers. + "vocab_size": 8, + # Default params needed for loading most pretrained 1B weights + "max_position_embeddings": 131072, + "hidden_size": 1024, + "intermediate_size": 4096, + "num_hidden_layers": 30, + "num_attention_heads": 16, + "attn_implementation": "sdpa", + "head_dim": 64, + "tie_word_embeddings": False, + "hidden_act": "silu", + "attention_bias": False, + "attention_dropout": 0.0, + "initializer_range": 0.02, + "mlp_bias": False, + "model_type": "llama", + "num_key_value_heads": 16, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": { + "factor": 8.0, + "high_freq_factor": 4.0, + "low_freq_factor": 1.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3", + }, + "rope_theta": 500000.0, + "torch_dtype": "bfloat16", + "use_cache": True, +} + + +class T3Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a T3 component. It is used internally by + ChatterboxModel for the T3 (Token-To-Token) TTS component. + + T3 (Token-To-Token) is a TTS model that uses a LLaMA transformer backbone to generate speech tokens from text tokens. + The speech tokens can then be decoded by S3Gen to produce mel spectrograms and finally waveforms. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + text_tokens_dict_size (`int`, *optional*, defaults to 704): + Size of the text token vocabulary. Use 704 for English-only, 2454 for multilingual. + speech_tokens_dict_size (`int`, *optional*, defaults to 8194): + Size of the speech token vocabulary. Includes special tokens (start: 6561, stop: 6562). + start_text_token (`int`, *optional*, defaults to 255): + Token ID for start-of-text marker. + stop_text_token (`int`, *optional*, defaults to 0): + Token ID for end-of-text marker. + start_speech_token (`int`, *optional*, defaults to 6561): + Token ID for start-of-speech marker. + stop_speech_token (`int`, *optional*, defaults to 6562): + Token ID for end-of-speech marker. + max_text_tokens (`int`, *optional*, defaults to 2048): + Maximum number of text tokens in a sequence. + max_speech_tokens (`int`, *optional*, defaults to 4096): + Maximum number of speech tokens in a sequence. + llama_config_name (`str`, *optional*, defaults to `"Llama_520M"`): + Name of the LLaMA configuration to use as backbone. + hidden_size (`int`, *optional*, defaults to 1024): + Hidden size of the transformer backbone (from LLaMA config). + input_pos_emb (`str`, *optional*, defaults to `"learned"`): + Type of positional embeddings. Currently only "learned" is supported. + speech_cond_prompt_len (`int`, *optional*, defaults to 150): + Length of speech conditioning prompt tokens. + encoder_type (`str`, *optional*, defaults to `"voice_encoder"`): + Type of speaker encoder to use. + speaker_embed_size (`int`, *optional*, defaults to 256): + Dimension of speaker embeddings from the voice encoder. + use_perceiver_resampler (`bool`, *optional*, defaults to `True`): + Whether to use perceiver resampler for conditioning prompts. + perceiver_num_latents (`int`, *optional*, defaults to 32): + Number of latent query tokens in the perceiver resampler. + perceiver_latent_dim (`int`, *optional*, defaults to 1024): + Dimension of latent tokens in the perceiver resampler. + perceiver_num_heads (`int`, *optional*, defaults to 4): + Number of attention heads in the perceiver resampler. + emotion_adv (`bool`, *optional*, defaults to `True`): + Whether to use emotion/exaggeration conditioning. + use_alignment_analyzer (`bool`, *optional*): + Whether to use alignment stream analyzer for multilingual models. If None, automatically enabled for multilingual. + alignment_layer_idx (`int`, *optional*, defaults to 9): + Layer index to use for attention-based alignment analysis in multilingual models. + """ + + model_type = "t3" + + def __init__( + self, + # Token vocabulary sizes + text_tokens_dict_size=704, + speech_tokens_dict_size=8194, + # Special tokens + start_text_token=255, + stop_text_token=0, + start_speech_token=6561, + stop_speech_token=6562, + # Sequence lengths + max_text_tokens=2048, + max_speech_tokens=4096, + # LLaMA backbone config + llama_config_name="Llama_520M", + hidden_size=1024, + # Positional embeddings + input_pos_emb="learned", + # Conditioning + speech_cond_prompt_len=150, + encoder_type="voice_encoder", + speaker_embed_size=256, + use_perceiver_resampler=True, + perceiver_num_latents=32, + perceiver_latent_dim=1024, + perceiver_num_heads=4, + emotion_adv=True, + # Multilingual support + use_alignment_analyzer=None, + alignment_layer_idx=9, + **kwargs, + ): + self.text_tokens_dict_size = text_tokens_dict_size + self.speech_tokens_dict_size = speech_tokens_dict_size + self.start_text_token = start_text_token + self.stop_text_token = stop_text_token + self.start_speech_token = start_speech_token + self.stop_speech_token = stop_speech_token + self.max_text_tokens = max_text_tokens + self.max_speech_tokens = max_speech_tokens + self.llama_config_name = llama_config_name + self.hidden_size = hidden_size + self.input_pos_emb = input_pos_emb + self.speech_cond_prompt_len = speech_cond_prompt_len + self.encoder_type = encoder_type + self.speaker_embed_size = speaker_embed_size + self.use_perceiver_resampler = use_perceiver_resampler + self.perceiver_num_latents = perceiver_num_latents + self.perceiver_latent_dim = perceiver_latent_dim + self.perceiver_num_heads = perceiver_num_heads + self.emotion_adv = emotion_adv + self.alignment_layer_idx = alignment_layer_idx + + # Auto-detect multilingual based on vocab size if not explicitly set + if use_alignment_analyzer is None: + self.use_alignment_analyzer = self.is_multilingual + else: + self.use_alignment_analyzer = use_alignment_analyzer + + # Store LLaMA config dict + self.llama_config_dict = LLAMA_520M_CONFIG_DICT.copy() + self.llama_config_dict["hidden_size"] = hidden_size + + super().__init__(**kwargs) + + @property + def is_multilingual(self): + """Check if this is a multilingual configuration based on vocab size.""" + return self.text_tokens_dict_size == 2454 + + @classmethod + def english_only(cls): + """Create configuration for English-only TTS model.""" + return cls(text_tokens_dict_size=704) + + @classmethod + def multilingual(cls): + """Create configuration for multilingual TTS model.""" + return cls(text_tokens_dict_size=2454) + + +# ============================================================================ +# Chatterbox Configuration +# ============================================================================ + + +class ChatterboxConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`ChatterboxModel`]. It is used to instantiate a + Chatterbox model according to the specified arguments, defining the model architecture. Instantiating a + configuration with the defaults will yield a similar configuration to that of the + [ResembleAI/chatterbox-hf](https://huggingface.co/ResembleAI/chatterbox-hf). + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Chatterbox is a complete TTS pipeline that combines T3, S3Gen, and HiFTNet models. + + Args: + t3_config (`dict` or `T3Config`, *optional*): + Dictionary or config object for T3 model. If not provided, uses English-only defaults. + s3gen_config (`dict` or `S3GenConfig`, *optional*): + Dictionary or config object for S3Gen model. If not provided, uses defaults. + hiftnet_config (`dict` or `HiFTNetConfig`, *optional*): + Dictionary or config object for HiFTNet model. If not provided, uses defaults. + is_multilingual (`bool`, *optional*, defaults to `False`): + Whether to use multilingual configuration. + + ```python + >>> from transformers import ChatterboxConfig, ChatterboxModel + + >>> # Initializing a Chatterbox configuration + >>> configuration = ChatterboxConfig() + + >>> # Initializing a model from the configuration + >>> model = ChatterboxModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "chatterbox" + is_composition = True + + def __init__( + self, + t3_config=None, + s3gen_config=None, + hiftnet_config=None, + is_multilingual=False, + **kwargs, + ): + super().__init__(**kwargs) + + from ...models.s3gen.configuration_s3gen import HiFTNetConfig, S3GenConfig + + # Initialize sub-model configs + # Handle both dict and Config object inputs + if t3_config is None: + if is_multilingual: + self.t3_config = T3Config.multilingual() + else: + self.t3_config = T3Config.english_only() + elif isinstance(t3_config, dict): + self.t3_config = T3Config(**t3_config) + else: + self.t3_config = t3_config + + if s3gen_config is None: + self.s3gen_config = S3GenConfig() + elif isinstance(s3gen_config, dict): + self.s3gen_config = S3GenConfig(**s3gen_config) + else: + self.s3gen_config = s3gen_config + + if hiftnet_config is None: + self.hiftnet_config = HiFTNetConfig() + elif isinstance(hiftnet_config, dict): + self.hiftnet_config = HiFTNetConfig(**hiftnet_config) + else: + self.hiftnet_config = hiftnet_config + + self.is_multilingual = is_multilingual + + @classmethod + def english_only(cls): + """Create English-only configuration.""" + return cls(is_multilingual=False) + + @classmethod + def multilingual(cls): + """Create multilingual configuration.""" + return cls(is_multilingual=True) + + def to_dict(self): + """Serialize to dict.""" + output = super().to_dict() + output["t3_config"] = self.t3_config.to_dict() + output["s3gen_config"] = self.s3gen_config.to_dict() + output["hiftnet_config"] = self.hiftnet_config.to_dict() + return output + + +__all__ = ["ChatterboxConfig"] diff --git a/src/transformers/models/chatterbox/convert_chatterbox_to_hf.py b/src/transformers/models/chatterbox/convert_chatterbox_to_hf.py new file mode 100644 index 000000000000..a233d8f72352 --- /dev/null +++ b/src/transformers/models/chatterbox/convert_chatterbox_to_hf.py @@ -0,0 +1,200 @@ +# coding=utf-8 +# Copyright 2025 The Resemble AI and 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. +"""Converts a Chatterbox model in Resemble AI format to Hugging Face format.""" + +import argparse +import os +from pathlib import Path + +import torch +from huggingface_hub import snapshot_download + +from transformers import ChatterboxConfig, ChatterboxModel + + +def convert_chatterbox_model_to_hf(checkpoint_path, pytorch_dump_folder_path, verbose=False): + """ + Converts a Chatterbox model in Resemble AI format to Hugging Face format. + """ + # Download from HF Hub + checkpoint_dir = snapshot_download(repo_id=checkpoint_path) + print(f"Downloaded checkpoint from Hugging Face Hub: {checkpoint_dir}") + + # Load original checkpoints + print("Loading original checkpoints...") + s3gen_path = os.path.join(checkpoint_dir, "s3gen.pt") + t3_path = os.path.join(checkpoint_dir, "t3.pt") + ve_path = os.path.join(checkpoint_dir, "ve.pt") + + # Check if files exist + if not os.path.exists(t3_path): + # Fallback to other possible names + t3_path = os.path.join(checkpoint_dir, "t3_cfg.pt") + if not os.path.exists(t3_path): + # Maybe it is in a different snapshot or just t3_23lang.safetensors + t3_path = next(Path(checkpoint_dir).glob("t3*.safetensors"), None) or next( + Path(checkpoint_dir).glob("t3*.pt"), None + ) + if t3_path: + t3_path = str(t3_path) + + print(f"Using T3 path: {t3_path}") + print(f"Using S3Gen path: {s3gen_path}") + print(f"Using VE path: {ve_path}") + + s3gen_sd = torch.load(s3gen_path, map_location="cpu") + t3_sd = ( + torch.load(t3_path, map_location="cpu") if t3_path.endswith(".pt") else None + ) # Will handle safetensors later if needed + if t3_path.endswith(".safetensors"): + from safetensors.torch import load_file + + t3_sd = load_file(t3_path) + + ve_sd = torch.load(ve_path, map_location="cpu") if ve_path.endswith(".pt") else None + if ve_path.endswith(".safetensors"): + from safetensors.torch import load_file + + ve_sd = load_file(ve_path) + + # Handle 'model' wrapper in T3 if present + if "model" in t3_sd and len(t3_sd) == 1: + print("Unwrapping T3 state dict from 'model' key...") + t3_sd = t3_sd["model"] + if isinstance(t3_sd, list): + t3_sd = t3_sd[0] + + # Initialize HF model + # Detect if multilingual from t3_path name or vocab size if possible + is_multilingual = "23lang" in t3_path or "multilingual" in t3_path + + config = ChatterboxConfig(is_multilingual=is_multilingual) + model = ChatterboxModel(config) + + hf_sd = model.state_dict() + new_sd = {} + + if verbose: + print(f"HF model has {len(hf_sd)} keys") + # print("First 20 HF keys:", sorted(list(hf_sd.keys()))[:20]) + + print("Mapping T3 weights...") + # T3 Mapping + for key, value in t3_sd.items(): + # original T3: tfmr.layers.0... -> HF Chatterbox: t3.layers.0... + new_key = key + if new_key.startswith("tfmr."): + new_key = new_key.replace("tfmr.", "t3.", 1) + elif not new_key.startswith("t3."): + new_key = "t3." + new_key + + if new_key in hf_sd: + new_sd[new_key] = value + elif new_key.replace("t3.model.", "t3.") in hf_sd: + new_sd[new_key.replace("t3.model.", "t3.")] = value + else: + if verbose: + print(f"Skipping T3 key: {key} (mapped to {new_key})") + + print("Mapping S3Gen weights...") + # S3Gen Mapping + for key, value in s3gen_sd.items(): + # original S3Gen: flow..., mel2wav..., tokenizer... -> HF Chatterbox: s3gen.flow..., s3gen.mel2wav..., s3gen.tokenizer... + new_key = "s3gen." + key + + # Handle S3Tokenizer mapping difference + if new_key.startswith("s3gen.tokenizer."): + new_key = new_key.replace("s3gen.tokenizer.", "s3gen.tokenizer.s3_model.", 1) + + if new_key in hf_sd: + new_sd[new_key] = value + else: + if verbose: + print(f"Skipping S3Gen key: {key} (mapped to {new_key})") + + print("Mapping Voice Encoder weights...") + # VE Mapping + for key, value in ve_sd.items(): + # original VE: lstm..., proj... -> HF Chatterbox: t3.voice_encoder.lstm..., t3.voice_encoder.proj... + new_key = "t3.voice_encoder." + key + if new_key in hf_sd: + new_sd[new_key] = value + else: + if verbose: + print(f"Skipping VE key: {key} (mapped to {new_key})") + + # Check for missing keys + missing_keys = set(hf_sd.keys()) - set(new_sd.keys()) + # Filter out known computed/buffer keys that might not be in checkpoints + missing_keys = { + k + for k in missing_keys + if not any(x in k for x in ["inv_freq", "stft_window", "trim_fade", "window", "_mel_filters", "embed_tokens"]) + } + + if missing_keys: + print(f"Warning: Missing keys in new state dict: {len(missing_keys)}") + if verbose: + for k in sorted(missing_keys)[:20]: + print(f" Missing: {k}") + + # Load state dict + model.load_state_dict(new_sd, strict=False) + + # Remove weight norm for saving + print("Removing weight norm...") + try: + if hasattr(model.s3gen, "mel2wav") and hasattr(model.s3gen.mel2wav, "remove_weight_norm"): + model.s3gen.mel2wav.remove_weight_norm() + except Exception as e: + print(f"Warning: Could not remove weight norm from s3gen.mel2wav: {e}") + + # Save model + print(f"Saving model to {pytorch_dump_folder_path}...") + model.save_pretrained(pytorch_dump_folder_path) + + # Copy tokenizer.json if it exists + tokenizer_path = os.path.join(checkpoint_dir, "tokenizer.json") + if os.path.exists(tokenizer_path): + import shutil + + shutil.copy(tokenizer_path, os.path.join(pytorch_dump_folder_path, "tokenizer.json")) + print(f"Copied tokenizer.json to {pytorch_dump_folder_path}") + + print("Conversion completed successfully!") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--checkpoint_path", + type=str, + default="ResembleAI/chatterbox", + help="Path to the downloaded checkpoints", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default="chatterbox-hf", + type=str, + help="Path to the output PyTorch model.", + ) + parser.add_argument( + "--verbose", + action="store_true", + help="Whether or not to log information during conversion.", + ) + args = parser.parse_args() + + convert_chatterbox_model_to_hf(args.checkpoint_path, args.pytorch_dump_folder_path, args.verbose) diff --git a/src/transformers/models/chatterbox/feature_extraction_chatterbox.py b/src/transformers/models/chatterbox/feature_extraction_chatterbox.py new file mode 100644 index 000000000000..0d3b43277f31 --- /dev/null +++ b/src/transformers/models/chatterbox/feature_extraction_chatterbox.py @@ -0,0 +1,341 @@ +# coding=utf-8 +# Copyright 2025 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. +"""Feature extractor class for Chatterbox.""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Any, Union + +import numpy as np +from numpy.lib.stride_tricks import as_strided + +from ...feature_extraction_sequence_utils import SequenceFeatureExtractor +from ...models.s3tokenizer.feature_extraction_s3tokenizer import S3TokenizerFeatureExtractor +from ...utils import is_librosa_available, is_torch_available +from ...utils.import_utils import requires + + +if is_torch_available(): + import torch +else: + torch = None + +if is_librosa_available(): + import librosa +else: + librosa = None + + +class VoiceEncConfig: + """ + Configuration for Voice Encoder. + + NOTE: This is intentionally aligned with `chatterbox.models.voice_encoder.config.VoiceEncConfig`. + """ + + num_mels = 40 + sample_rate = 16000 + speaker_embed_size = 256 + ve_hidden_size = 256 + flatten_lstm_params = False + n_fft = 400 + hop_size = 160 + win_size = 400 + fmax = 8000 + fmin = 0 + preemphasis = 0.0 + mel_power = 2.0 + mel_type = "amp" + normalized_mels = False + ve_partial_frames = 160 + ve_final_relu = True + stft_magnitude_min = 1e-4 + + +@lru_cache +def _mel_basis_voice_encoder(hp: VoiceEncConfig): + if librosa is None: + raise ImportError( + "librosa is required to compute mel filters for Chatterbox voice encoder preprocessing. " + "Please install it with `pip install librosa`." + ) + assert hp.fmax <= hp.sample_rate // 2 + return librosa.filters.mel( + sr=hp.sample_rate, + n_fft=hp.n_fft, + n_mels=hp.num_mels, + fmin=hp.fmin, + fmax=hp.fmax, + ) # (n_mels, 1 + n_fft//2) + + +def _preemphasis_voice_encoder(wav: np.ndarray, hp: VoiceEncConfig) -> np.ndarray: + # Matches a simple lfilter([1, -preemph], [1], wav) without depending on scipy. + assert hp.preemphasis != 0 + out = np.empty_like(wav) + out[0] = wav[0] + out[1:] = wav[1:] - hp.preemphasis * wav[:-1] + return np.clip(out, -1, 1) + + +def _amp_to_db_voice_encoder(x: np.ndarray, hp: VoiceEncConfig) -> np.ndarray: + return 20 * np.log10(np.maximum(hp.stft_magnitude_min, x)) + + +def _normalize_voice_encoder(s: np.ndarray, hp: VoiceEncConfig, headroom_db: float = 15) -> np.ndarray: + min_level_db = 20 * np.log10(hp.stft_magnitude_min) + return (s - min_level_db) / (-min_level_db + headroom_db) + + +def _stft_voice_encoder(y: np.ndarray, hp: VoiceEncConfig, pad: bool = True): + if librosa is None: + raise ImportError( + "librosa is required to compute STFT for Chatterbox voice encoder preprocessing. " + "Please install it with `pip install librosa`." + ) + # Match chatterbox: pad_mode="reflect" for historical/streaming consistency. + return librosa.stft( + y, + n_fft=hp.n_fft, + hop_length=hp.hop_size, + win_length=hp.win_size, + center=pad, + pad_mode="reflect", + ) + + +def melspectrogram_voice_encoder(wav: np.ndarray, hp: VoiceEncConfig, pad: bool = True) -> np.ndarray: + """ + Voice encoder mel extraction aligned with `chatterbox.models.voice_encoder.melspec.melspectrogram`. + + Returns: + np.ndarray of shape (num_mels, T) + """ + if hp.preemphasis > 0: + wav = _preemphasis_voice_encoder(wav, hp) + assert np.abs(wav).max() - 1 < 1e-07 + + spec_complex = _stft_voice_encoder(wav, hp, pad=pad) + spec_magnitudes = np.abs(spec_complex) + + if hp.mel_power != 1.0: + spec_magnitudes **= hp.mel_power + + mel = np.dot(_mel_basis_voice_encoder(hp), spec_magnitudes) + if hp.mel_type == "db": + mel = _amp_to_db_voice_encoder(mel, hp) + + if hp.normalized_mels: + mel = _normalize_voice_encoder(mel, hp).astype(np.float32) + + assert not pad or mel.shape[1] == 1 + len(wav) // hp.hop_size + return mel + + +def stride_as_partials(mel: np.ndarray, hp: VoiceEncConfig, overlap=0.5, rate: float | None = None, min_coverage=0.8): + """Stride mel spectrogram into overlapping partials.""" + + def get_frame_step(overlap, rate, hp): + assert 0 <= overlap < 1 + if rate is None: + frame_step = int(np.round(hp.ve_partial_frames * (1 - overlap))) + else: + frame_step = int(np.round((hp.sample_rate / rate) / hp.ve_partial_frames)) + assert 0 < frame_step <= hp.ve_partial_frames + return frame_step + + def get_num_wins(n_frames, step, min_coverage, hp): + assert n_frames > 0 + win_size = hp.ve_partial_frames + n_wins, remainder = divmod(max(n_frames - win_size + step, 0), step) + if n_wins == 0 or (remainder + (win_size - step)) / win_size >= min_coverage: + n_wins += 1 + target_n = win_size + step * (n_wins - 1) + return n_wins, target_n + + assert 0 < min_coverage <= 1 + frame_step = get_frame_step(overlap, rate, hp) + n_partials, target_len = get_num_wins(len(mel), frame_step, min_coverage, hp) + + # Trim or pad + if target_len > len(mel): + mel = np.concatenate((mel, np.full((target_len - len(mel), hp.num_mels), 0))) + elif target_len < len(mel): + mel = mel[:target_len] + + mel = mel.astype(np.float32, order="C") + shape = (n_partials, hp.ve_partial_frames, hp.num_mels) + strides = (mel.strides[0] * frame_step, mel.strides[0], mel.strides[1]) + partials = as_strided(mel, shape, strides) + return partials + + +@requires(backends=("torch",)) +class ChatterboxFeatureExtractor(SequenceFeatureExtractor): + """ + Constructs a Chatterbox feature extractor. + + This feature extractor is responsible for preparing *conditioning* inputs for `ChatterboxModel`, including: + - resampling and truncation of reference audio, + - `S3GenModel.embed_ref(...)` inputs, + - voice-encoder speaker embedding inputs and extraction, + - optional speech prompt tokenization via S3Tokenizer. + + Notes: + This feature extractor is intentionally "model-assisted": it requires the instantiated sub-modules + (`s3gen`, `voice_encoder`) to compute conditioning tensors aligned with the model weights. + """ + + # Not used for padding in the usual sense, but keep a minimal, consistent base configuration. + model_input_names = [ + "speaker_emb", + "cond_prompt_speech_tokens", + "emotion_adv", + "s3gen_ref_dict", + ] + + def __init__( + self, + feature_size: int = 1, + sampling_rate: int = 16000, + padding_value: float = 0.0, + s3gen_sampling_rate: int = 24000, + s3gen_ref_seconds: int = 10, + t3_prompt_seconds: int = 6, + **kwargs, + ): + super().__init__( + feature_size=feature_size, + sampling_rate=sampling_rate, + padding_value=padding_value, + **kwargs, + ) + self.s3gen_sampling_rate = int(s3gen_sampling_rate) + self.s3gen_ref_seconds = int(s3gen_ref_seconds) + self.t3_prompt_seconds = int(t3_prompt_seconds) + + # Internal helper for prompt tokenization (16kHz mel features). + self._s3_feature_extractor = S3TokenizerFeatureExtractor(sampling_rate=self.sampling_rate) + + @staticmethod + def _to_1d_float32_np(reference_wav: Union[np.ndarray, list[float]]) -> np.ndarray: + ref_np = np.asarray(reference_wav, dtype=np.float32) + if ref_np.ndim == 0: + raise ValueError("`reference_wav` must be a 1D waveform array, got a scalar.") + if ref_np.ndim > 1: + ref_np = ref_np.squeeze() + if ref_np.ndim != 1: + raise ValueError(f"`reference_wav` must be 1D after squeeze, got shape {ref_np.shape}.") + return ref_np + + def extract_conditioning( + self, + reference_wav: Union[np.ndarray, list[float]], + reference_sr: int, + *, + s3gen: Any, + voice_encoder: Any, + device: Union[str, torch.device], + exaggeration: float = 0.5, + speech_cond_prompt_len: int = 150, + ) -> dict[str, Any]: + """ + Extract conditioning for Chatterbox from a reference waveform. + + Args: + reference_wav: Reference audio waveform (mono) as 1D numpy array or list of floats. + reference_sr: Sampling rate (Hz) of `reference_wav`. + s3gen: Instantiated S3Gen model. Must expose `embed_ref(...)` and `.tokenizer(...)`. + voice_encoder: Instantiated voice encoder module. Must expose `embeds_from_wavs(...)`. + device: Target device for returned torch tensors. + exaggeration: Emotion/expressiveness level. + speech_cond_prompt_len: Maximum length of speech conditioning prompt tokens. If <= 0, prompt is disabled. + + Returns: + Dict with keys: + - `speaker_emb`: torch.FloatTensor (1, speaker_embed_size) + - `cond_prompt_speech_tokens`: Optional[torch.LongTensor] (1, prompt_len) + - `emotion_adv`: torch.FloatTensor (1, 1, 1) + - `s3gen_ref_dict`: dict produced by `s3gen.embed_ref(...)` + """ + if reference_sr is None: + raise ValueError("`reference_sr` must be provided for Chatterbox conditioning.") + reference_sr = int(reference_sr) + if reference_sr <= 0: + raise ValueError(f"`reference_sr` must be > 0, got {reference_sr}.") + + ref_np = self._to_1d_float32_np(reference_wav) + + # Prepare audio for S3Gen (24kHz) and T3 components (16kHz). + if librosa is None: + raise ImportError( + "librosa is required to resample reference audio for Chatterbox conditioning. " + "Please install it with `pip install librosa`." + ) + if reference_sr != self.s3gen_sampling_rate: + ref_24k = librosa.resample(ref_np, orig_sr=reference_sr, target_sr=self.s3gen_sampling_rate) + else: + ref_24k = ref_np + + if reference_sr != self.sampling_rate: + ref_16k = librosa.resample(ref_np, orig_sr=reference_sr, target_sr=self.sampling_rate) + else: + ref_16k = ref_np + + # Truncate for conditioning lengths. + dec_len = self.s3gen_ref_seconds * self.s3gen_sampling_rate + enc_len = self.t3_prompt_seconds * self.sampling_rate + ref_24k = ref_24k[:dec_len] + ref_16k_prompt = ref_16k[:enc_len] + + # Compute S3Gen conditioning dict. + ref_tensor_24k = torch.from_numpy(ref_24k).unsqueeze(0).to(device) + with torch.no_grad(): + s3gen_ref_dict = s3gen.embed_ref(ref_tensor_24k, self.s3gen_sampling_rate, device=device) + + # Voice encoder speaker embedding. + ve_embed = voice_encoder.embeds_from_wavs([ref_16k], sample_rate=self.sampling_rate) + speaker_emb = torch.from_numpy(ve_embed).to(device) + + # Speech prompt tokens for T3. + cond_prompt_speech_tokens = None + if int(speech_cond_prompt_len) > 0: + features = self._s3_feature_extractor( + ref_16k_prompt, sampling_rate=self.sampling_rate, return_tensors="pt" + ) + features = features.to(device) + with torch.no_grad(): + prompt_tokens, _ = s3gen.tokenizer( + input_features=features.input_features, + attention_mask=features.attention_mask, + return_dict=False, + max_len=int(speech_cond_prompt_len), + ) + cond_prompt_speech_tokens = prompt_tokens.to(device) + + emotion_adv = float(exaggeration) * torch.ones(1, 1, 1, device=device) + + return { + "speaker_emb": speaker_emb, + "cond_prompt_speech_tokens": cond_prompt_speech_tokens, + "emotion_adv": emotion_adv, + "s3gen_ref_dict": s3gen_ref_dict, + } + + +__all__ = [ + "ChatterboxFeatureExtractor", +] diff --git a/src/transformers/models/chatterbox/modeling_chatterbox.py b/src/transformers/models/chatterbox/modeling_chatterbox.py new file mode 100644 index 000000000000..9e024c98ae83 --- /dev/null +++ b/src/transformers/models/chatterbox/modeling_chatterbox.py @@ -0,0 +1,1215 @@ +# coding=utf-8 +# Copyright 2025 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 Chatterbox model - Complete TTS Pipeline.""" + +import logging +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Optional, Union + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from tokenizers import Tokenizer +from torch import Tensor + +from ...generation.utils import GenerationMixin +from ...modeling_outputs import CausalLMOutputWithCrossAttentions +from ...modeling_utils import PreTrainedModel +from ...models.s3gen.modeling_s3gen import S3GenModel +from ...models.s3tokenizer.modeling_s3tokenizer import drop_invalid_tokens +from ...utils import auto_docstring, is_librosa_available +from ..llama.modeling_llama import LlamaConfig, LlamaModel, LlamaPreTrainedModel +from .configuration_chatterbox import ChatterboxConfig +from .feature_extraction_chatterbox import ( + ChatterboxFeatureExtractor, + VoiceEncConfig, + melspectrogram_voice_encoder, + stride_as_partials, +) + + +if is_librosa_available(): + import librosa +else: + librosa = None + + +logger = logging.getLogger(__name__) + + +def punc_norm(text: str) -> str: + """ + Quick cleanup func for punctuation from LLMs or containing chars not seen often in the dataset. + """ + if len(text) == 0: + return "You need to add some text for me to talk." + + # Capitalize first letter + if text[0].islower(): + text = text[0].upper() + text[1:] + + # Remove multiple space chars + text = " ".join(text.split()) + + # Replace uncommon/llm punc + punc_to_replace = [ + ("...", ", "), + ("…", ", "), + (":", ","), + (" - ", ", "), + (";", ", "), + ("—", "-"), + ("–", "-"), + (" ,", ","), + ( + """, '"'), + (""", + '"', + ), + ("'", "'"), + ("'", "'"), + ] + for old_char_sequence, new_char in punc_to_replace: + text = text.replace(old_char_sequence, new_char) + + # Add full stop if no ending punc + text = text.rstrip(" ") + sentence_enders = {".", "!", "?", "-", ","} + if not any(text.endswith(p) for p in sentence_enders): + text += "." + + return text + + +class VoiceEncoder(nn.Module): + """Voice encoder for speaker embedding extraction.""" + + def __init__(self, config: VoiceEncConfig = None): + super().__init__() + self.config = config if config is not None else VoiceEncConfig() + + self.lstm = nn.LSTM(self.config.num_mels, self.config.ve_hidden_size, num_layers=3, batch_first=True) + self.proj = nn.Linear(self.config.ve_hidden_size, self.config.speaker_embed_size) + + # Cosine similarity scaling + self.similarity_weight = nn.Parameter(torch.tensor([10.0]), requires_grad=True) + self.similarity_bias = nn.Parameter(torch.tensor([-5.0]), requires_grad=True) + + @property + def device(self): + return next(self.parameters()).device + + def forward(self, mels: torch.FloatTensor): + """Compute embeddings from mel spectrograms.""" + _, (hidden, _) = self.lstm(mels) + raw_embeds = self.proj(hidden[-1]) + if self.config.ve_final_relu: + raw_embeds = F.relu(raw_embeds) + return raw_embeds / torch.linalg.norm(raw_embeds, dim=1, keepdim=True) + + def embeds_from_wavs( + self, wavs: list[np.ndarray], sample_rate: int, overlap=0.5, rate: float = 1.3, batch_size=32 + ): + """Extract embeddings from waveforms.""" + if librosa is None: + raise ImportError( + "librosa is required for Chatterbox voice encoder preprocessing (resampling + trimming). " + "Please install it with `pip install librosa`." + ) + if sample_rate != self.config.sample_rate: + wavs = [ + librosa.resample(wav, orig_sr=sample_rate, target_sr=self.config.sample_rate, res_type="kaiser_fast") + for wav in wavs + ] + + wavs = [librosa.effects.trim(wav, top_db=20)[0] for wav in wavs] + + # Extract mel spectrograms + mels = [melspectrogram_voice_encoder(w, self.config).T for w in wavs] + + # Stride into partials + all_partials = [] + n_partials_per_wav = [] + for mel in mels: + partials = stride_as_partials(mel, self.config, overlap=overlap, rate=rate) + all_partials.append(torch.from_numpy(partials)) + n_partials_per_wav.append(len(partials)) + + # Stack and process + all_partials = torch.cat(all_partials, dim=0).to(self.device) + + # Forward in batches + n_chunks = int(np.ceil(len(all_partials) / batch_size)) + partial_embeds = [] + for chunk in all_partials.chunk(n_chunks): + with torch.inference_mode(): + partial_embeds.append(self(chunk)) + partial_embeds = torch.cat(partial_embeds, dim=0).cpu() + + # Aggregate partials per wav + slices = np.concatenate(([0], np.cumsum(n_partials_per_wav))) + embeds = [] + for start, end in zip(slices[:-1], slices[1:]): + raw_embed = torch.mean(partial_embeds[start:end], dim=0) + embeds.append(raw_embed / torch.linalg.norm(raw_embed)) + + return torch.stack(embeds).numpy() + + +# ============================================================================ +# Learned Position Embeddings +# ============================================================================ + + +class LearnedPositionEmbeddings(nn.Module): + """Learned positional embeddings.""" + + def __init__(self, seq_len, model_dim, init=0.02): + super().__init__() + self.emb = nn.Embedding(seq_len, model_dim) + self.emb.weight.data.normal_(mean=0.0, std=init) + + def forward(self, x): + """Returns positional embeddings for index 0 up to the length of x.""" + sl = x.shape[1] + return self.emb(torch.arange(0, sl, device=x.device)) + + def get_fixed_embedding(self, idx: Union[int, Tensor]): + """Get embeddings for specific indices.""" + device = self.emb.weight.device + idx = idx.to(device) if torch.is_tensor(idx) else torch.tensor(idx, device=device) + idx = torch.atleast_2d(idx) + assert idx.ndim == 2 + return self.emb(idx) + + +# ============================================================================ +# Perceiver Resampler +# ============================================================================ + + +class AttentionQKV(nn.Module): + """Attention module with separate Q, K, V projections.""" + + def __init__(self, n_heads, head_dim, dropout_rate=0.1, scale=None, flash=False): + super().__init__() + self.n_heads = n_heads + self.head_dim = head_dim + self.scale = scale if scale is not None else head_dim**-0.5 + self.flash = flash + self.dropout_rate = dropout_rate + self.dropout = nn.Dropout(dropout_rate) + + def forward(self, q, k, v, mask=None): + q, k, v = [self.split_heads(tensor) for tensor in [q, k, v]] + if self.flash and hasattr(F, "scaled_dot_product_attention"): + out = F.scaled_dot_product_attention( + q, k, v, attn_mask=mask, dropout_p=self.dropout_rate if self.training else 0.0 + ) + else: + out = self.scaled_dot_product_attention(q, k, v, mask=mask) + return self.combine_heads(out) + + def scaled_dot_product_attention(self, q, k, v, mask=None): + sim = torch.einsum("bhlt,bhls->bhts", q, k) * self.scale + if mask is not None: + sim = sim.masked_fill(mask == 0, float("-inf")) + attn = torch.softmax(sim, dim=-1) + attn = self.dropout(attn) + return torch.einsum("bhts,bhls->bhlt", attn, v) + + def split_heads(self, x): + bs, length, _ = x.shape + x = x.view(bs, length, self.n_heads, self.head_dim) + return x.permute(0, 2, 1, 3) + + def combine_heads(self, x): + bs, _, length, _ = x.shape + x = x.permute(0, 2, 1, 3).contiguous() + return x.view(bs, length, -1) + + +class AttentionBlock(nn.Module): + """Cross-attention block for perceiver.""" + + def __init__( + self, channels, num_heads=1, num_head_channels=-1, flash_attention=True, dropout_rate=0.2, scale=None + ): + super().__init__() + self.channels = channels + + if num_head_channels == -1: + self.num_heads = num_heads + else: + assert channels % num_head_channels == 0 + self.num_heads = channels // num_head_channels + + self.norm = nn.LayerNorm(channels) + self.to_q = nn.Linear(channels, channels) + self.to_k = nn.Linear(channels, channels) + self.to_v = nn.Linear(channels, channels) + self.attention = AttentionQKV( + self.num_heads, channels // self.num_heads, dropout_rate=dropout_rate, flash=flash_attention, scale=scale + ) + self.proj_out = nn.Linear(channels, channels) + + def forward(self, x1, x2, mask=None): + b1, c1, *spatial1 = x1.shape + x1_norm = self.norm(x1) + x2_norm = self.norm(x2) + + q = self.to_q(x1_norm) + k = self.to_k(x2_norm) + v = self.to_v(x2_norm) + + h = self.attention(q, k, v, mask=mask) + h = self.proj_out(h) + return (x1 + h).reshape(b1, c1, *spatial1) + + +class Perceiver(nn.Module): + """Perceiver resampler for conditioning.""" + + def __init__( + self, pre_attention_query_token=32, pre_attention_query_size=1024, embedding_dim=1024, num_attn_heads=4 + ): + super().__init__() + self.pre_attention_query = nn.Parameter(torch.empty(1, pre_attention_query_token, pre_attention_query_size)) + query_variance = math.sqrt(3.0) * math.sqrt(2.0 / (pre_attention_query_token + pre_attention_query_token)) + self.pre_attention_query.data.uniform_(-query_variance, query_variance) + self.attn = AttentionBlock(embedding_dim, num_attn_heads) + + def forward(self, h): + query_ = self.pre_attention_query.expand(h.shape[0], -1, -1) + pre_att = self.attn(query_, h) + attn = self.attn(pre_att, pre_att) + return attn + + +# ============================================================================ +# T3 Conditioning +# ============================================================================ + + +@dataclass +class T3Cond: + """Dataclass container for T3 conditioning information.""" + + speaker_emb: Tensor + clap_emb: Optional[Tensor] = None + cond_prompt_speech_tokens: Optional[Tensor] = None + cond_prompt_speech_emb: Optional[Tensor] = None + emotion_adv: Optional[Tensor] = None + + def to(self, *, device=None, dtype=None): + """Cast to a device and dtype.""" + for k, v in self.__dict__.items(): + if torch.is_tensor(v): + is_fp = v.dtype not in [torch.long, torch.int, torch.int32, torch.int64] + setattr(self, k, v.to(device=device, dtype=dtype if is_fp else None)) + return self + + +class T3CondEnc(nn.Module): + """Encoder for T3 conditioning (speaker, emotion, prompts).""" + + def __init__(self, config): + super().__init__() + self.config = config + + if config.encoder_type == "voice_encoder": + self.spkr_enc = nn.Linear(config.speaker_embed_size, config.hidden_size) + else: + raise NotImplementedError(str(config.encoder_type)) + + self.emotion_adv_fc = None + if config.emotion_adv: + self.emotion_adv_fc = nn.Linear(1, config.hidden_size, bias=False) + + self.perceiver = None + if config.use_perceiver_resampler: + self.perceiver = Perceiver( + pre_attention_query_token=config.perceiver_num_latents, + pre_attention_query_size=config.perceiver_latent_dim, + embedding_dim=config.hidden_size, + num_attn_heads=config.perceiver_num_heads, + ) + + def forward(self, cond: T3Cond): + assert (cond.cond_prompt_speech_tokens is None) == (cond.cond_prompt_speech_emb is None), ( + "no embeddings for cond_prompt_speech_tokens" + ) + + # Speaker embedding projection + cond_spkr = self.spkr_enc(cond.speaker_emb.view(-1, self.config.speaker_embed_size))[:, None] + empty = torch.zeros_like(cond_spkr[:, :0]) + + # CLAP (not implemented) + assert cond.clap_emb is None, "clap_embed not implemented" + cond_clap = empty + + # Conditioning prompt + cond_prompt_speech_emb = cond.cond_prompt_speech_emb + if cond_prompt_speech_emb is None: + cond_prompt_speech_emb = empty + elif self.config.use_perceiver_resampler: + cond_prompt_speech_emb = self.perceiver(cond_prompt_speech_emb) + + # Emotion + cond_emotion_adv = empty + if self.config.emotion_adv: + assert cond.emotion_adv is not None + cond_emotion_adv = self.emotion_adv_fc(cond.emotion_adv.view(-1, 1, 1)) + + # Concatenate + cond_embeds = torch.cat((cond_spkr, cond_clap, cond_prompt_speech_emb, cond_emotion_adv), dim=1) + return cond_embeds + + +# ============================================================================ +# Alignment Stream Analyzer (for multilingual) +# ============================================================================ + +LLAMA_ALIGNED_HEADS = [(12, 15), (13, 11), (9, 2)] + + +class AlignmentStreamAnalyzer: + """Alignment analyzer for detecting hallucinations in multilingual models.""" + + def __init__(self, tfmr, text_tokens_slice, alignment_layer_idx=9, eos_idx=0): + self.text_tokens_slice = (i, j) = text_tokens_slice + self.eos_idx = eos_idx + self.alignment = torch.zeros(0, j - i) + self.curr_frame_pos = 0 + self.text_position = 0 + self.started = False + self.started_at = None + self.complete = False + self.completed_at = None + self.generated_tokens = [] + self.last_aligned_attns = [] + + for i, (layer_idx, head_idx) in enumerate(LLAMA_ALIGNED_HEADS): + self.last_aligned_attns += [None] + self._add_attention_spy(tfmr, i, layer_idx, head_idx) + + def _add_attention_spy(self, tfmr, buffer_idx, layer_idx, head_idx): + """Add forward hook to collect attention weights.""" + + def attention_forward_hook(module, input, output): + if isinstance(output, tuple) and len(output) > 1 and output[1] is not None: + step_attention = output[1].cpu() + self.last_aligned_attns[buffer_idx] = step_attention[0, head_idx] + + target_layer = tfmr.layers[layer_idx].self_attn + target_layer.register_forward_hook(attention_forward_hook) + if hasattr(tfmr, "config") and hasattr(tfmr.config, "output_attentions"): + tfmr.config.output_attentions = True + + def step(self, logits, next_token=None): + """Analyze alignment and potentially modify logits.""" + aligned_attn = torch.stack(self.last_aligned_attns).mean(dim=0) + i, j = self.text_tokens_slice + if self.curr_frame_pos == 0: + A_chunk = aligned_attn[j:, i:j].clone().cpu() + else: + A_chunk = aligned_attn[:, i:j].clone().cpu() + + A_chunk[:, self.curr_frame_pos + 1 :] = 0 + self.alignment = torch.cat((self.alignment, A_chunk), dim=0) + + A = self.alignment + T, S = A.shape + + cur_text_posn = A_chunk[-1].argmax() + discontinuity = not (-4 < cur_text_posn - self.text_position < 7) + if not discontinuity: + self.text_position = cur_text_posn + + false_start = (not self.started) and (A[-2:, -2:].max() > 0.1 or A[:, :4].max() < 0.5) + self.started = not false_start + if self.started and self.started_at is None: + self.started_at = T + + self.complete = self.complete or self.text_position >= S - 3 + if self.complete and self.completed_at is None: + self.completed_at = T + + long_tail = self.complete and (A[self.completed_at :, -3:].sum(dim=0).max() >= 5) + alignment_repetition = self.complete and (A[self.completed_at :, :-5].max(dim=1).values.sum() > 5) + + # Track tokens + if next_token is not None: + if isinstance(next_token, torch.Tensor): + token_id = next_token.item() if next_token.numel() == 1 else next_token.view(-1)[0].item() + else: + token_id = next_token + self.generated_tokens.append(token_id) + if len(self.generated_tokens) > 8: + self.generated_tokens = self.generated_tokens[-8:] + + token_repetition = len(self.generated_tokens) >= 3 and len(set(self.generated_tokens[-2:])) == 1 + + # Suppress EOS early + if cur_text_posn < S - 3 and S > 5: + logits[..., self.eos_idx] = -(2**15) + + # Force EOS on bad endings + if long_tail or alignment_repetition or token_repetition: + logger.warning(f"Forcing EOS: {long_tail=}, {alignment_repetition=}, {token_repetition=}") + logits = -(2**15) * torch.ones_like(logits) + logits[..., self.eos_idx] = 2**15 + + self.curr_frame_pos += 1 + return logits + + +# ============================================================================ +# T3 Model +# ============================================================================ + + +class T3PreTrainedModel(LlamaPreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = None # Will be set by ChatterboxConfig + base_model_prefix = "t3" + supports_gradient_checkpointing = True + _no_split_modules = ["LlamaDecoderLayer"] + + +@auto_docstring +class T3Model(T3PreTrainedModel, LlamaModel, GenerationMixin): + """ + T3 (Token-To-Token) TTS model using LLaMA as backbone. + + This model generates speech tokens from text tokens, which can then be decoded by S3Gen. + """ + + def __init__(self, config): + # Create LLaMA backbone config and initialize parent LlamaModel + llama_config = LlamaConfig(**config.llama_config_dict) + super().__init__(llama_config) + + # Store the full T3 config for T3-specific settings + self.t3_config = config + self.dim = llama_config.hidden_size + + # llama_config_name is stored in config for reference + _ = config.llama_config_name + + # Conditioning encoder + self.cond_enc = T3CondEnc(self.t3_config) + + # Text and speech embeddings + self.text_emb = nn.Embedding(self.t3_config.text_tokens_dict_size, self.dim) + self.speech_emb = nn.Embedding(self.t3_config.speech_tokens_dict_size, self.dim) + + # Positional embeddings + if self.t3_config.input_pos_emb == "learned": + max_text_seq_len = self.t3_config.max_text_tokens + 2 + self.text_pos_emb = LearnedPositionEmbeddings(max_text_seq_len, self.dim) + + max_mel_seq_len = self.t3_config.max_speech_tokens + 2 + 2 + self.speech_pos_emb = LearnedPositionEmbeddings(max_mel_seq_len, self.dim) + + # Output heads + self.text_head = nn.Linear(self.dim, self.t3_config.text_tokens_dict_size, bias=False) + self.speech_head = nn.Linear(self.dim, self.t3_config.speech_tokens_dict_size, bias=False) + + # Voice encoder for speaker conditioning + self.voice_encoder = VoiceEncoder() + + # Set main input name for generation + self.main_input_name = "inputs_embeds" + + # Initialize weights + self.post_init() + + # Generation state + self._decoder_cond = None + self._added_cond = False + self._current_position = 0 + self.alignment_stream_analyzer = None + + def prepare_conditioning(self, t3_cond: T3Cond): + """Prepare conditioning embeddings.""" + if t3_cond.cond_prompt_speech_tokens is not None and t3_cond.cond_prompt_speech_emb is None: + t3_cond.cond_prompt_speech_emb = self.speech_emb(t3_cond.cond_prompt_speech_tokens) + self.speech_pos_emb( + t3_cond.cond_prompt_speech_tokens + ) + return self.cond_enc(t3_cond) + + def prepare_input_embeds( + self, + *, + t3_cond: T3Cond, + text_tokens: torch.LongTensor, + speech_tokens: torch.LongTensor, + cfg_weight: float = 0.0, + ): + """Prepare input embeddings for the transformer.""" + cond_emb = self.prepare_conditioning(t3_cond) + text_emb = self.text_emb(text_tokens) + if cfg_weight > 0.0 and text_emb.size(0) > 1: + text_emb[1].zero_() # CFG uncond + + speech_emb = self.speech_emb(speech_tokens) + if self.t3_config.input_pos_emb == "learned": + text_emb = text_emb + self.text_pos_emb(text_tokens) + speech_emb = speech_emb + self.speech_pos_emb(speech_tokens) + + len_cond = cond_emb.size(1) + + if cond_emb.size(0) != text_emb.size(0): + cond_emb = cond_emb.expand(text_emb.size(0), -1, -1) + + embeds = torch.stack([torch.cat((ce, te, se)) for ce, te, se in zip(cond_emb, text_emb, speech_emb)]) + return embeds, len_cond + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + # Training-specific arguments + t3_cond: Optional[T3Cond] = None, + text_tokens: Optional[torch.LongTensor] = None, + text_token_lens: Optional[torch.LongTensor] = None, + speech_tokens: Optional[torch.LongTensor] = None, + speech_token_lens: Optional[torch.LongTensor] = None, + ): + """ + Forward pass of T3 model. + + Supports both training mode (with t3_cond, text_tokens, speech_tokens) and + generation mode (with inputs_embeds from prepare_inputs_for_generation). + """ + return_dict = return_dict if return_dict is not None else self.t3_config.use_return_dict + + # Training/evaluation mode + if t3_cond is not None and text_tokens is not None and speech_tokens is not None: + embeds, len_cond = self.prepare_input_embeds( + t3_cond=t3_cond, text_tokens=text_tokens, speech_tokens=speech_tokens + ) + + tfmr_out = super().forward( + inputs_embeds=embeds, + output_hidden_states=True, + return_dict=True, + use_cache=False, + ) + hidden_states = tfmr_out.hidden_states[-1] + + # Split hidden states + len_text = text_tokens.size(1) + len_speech = speech_tokens.size(1) + B, _, dim = hidden_states.shape + device, dtype = hidden_states.device, hidden_states.dtype + + text_latents = torch.zeros(B, len_text, dim, dtype=dtype, device=device) + speech_latents = torch.zeros(B, len_speech, dim, dtype=dtype, device=device) + + ttl, stl = text_token_lens, speech_token_lens + for i in range(B): + text_end = len_cond + ttl[i].item() + speech_start = len_cond + text_tokens.size(1) + speech_end = speech_start + stl[i].item() + text_latents[i, : ttl[i]] = hidden_states[i, len_cond:text_end] + speech_latents[i, : stl[i]] = hidden_states[i, speech_start:speech_end] + + text_logits = self.text_head(text_latents) + speech_logits = self.speech_head(speech_latents) + + return { + "text_logits": text_logits, + "text_latents": text_latents, + "speech_logits": speech_logits, + "speech_latents": speech_latents, + "hidden_states": hidden_states, + } + + # Generation mode + else: + output_attentions = output_attentions if output_attentions is not None else False + output_hidden_states = output_hidden_states if output_hidden_states is not None else True + use_cache = use_cache if use_cache is not None else True + + tfmr_out = super().forward( + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + ) + + hidden_states = tfmr_out.hidden_states[-1] + logits = self.speech_head(hidden_states) + + if not return_dict: + return (logits, tfmr_out.past_key_values, hidden_states, tfmr_out.attentions) + + return CausalLMOutputWithCrossAttentions( + logits=logits, + past_key_values=tfmr_out.past_key_values, + hidden_states=tfmr_out.hidden_states, + attentions=tfmr_out.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids: torch.LongTensor, + past_key_values: Optional[tuple[tuple[torch.Tensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + **kwargs, + ): + """ + Prepare inputs for generation step. + + This method is called by HuggingFace's generate() at each step. + """ + # First call: add conditioning embeddings + if past_key_values is None: + # Initial call - input_ids is the BOS token(s) + inputs_embeds = self.speech_emb(input_ids) + inputs_embeds = inputs_embeds + self.speech_pos_emb.get_fixed_embedding(0) + + # Prepend decoder conditioning if available + if self._decoder_cond is not None: + inputs_embeds = torch.cat([self._decoder_cond, inputs_embeds], dim=1) + self._current_position = 0 + else: + # Subsequent calls - only process the new token + input_ids = input_ids[:, -1:] + inputs_embeds = self.speech_emb(input_ids) + self._current_position += 1 + inputs_embeds = inputs_embeds + self.speech_pos_emb.get_fixed_embedding(self._current_position) + + return { + "inputs_embeds": inputs_embeds, + "past_key_values": past_key_values, + "use_cache": use_cache if use_cache is not None else True, + "output_attentions": self.t3_config.use_alignment_analyzer + if hasattr(self.t3_config, "use_alignment_analyzer") + else False, + "output_hidden_states": True, + } + + def loss( + self, + *, + t3_cond: T3Cond, + text_tokens: torch.LongTensor, + text_token_lens: torch.LongTensor, + speech_tokens: torch.LongTensor, + speech_token_lens: torch.LongTensor, + ): + """Compute training loss.""" + len_text = text_tokens.size(1) + len_speech = speech_tokens.size(1) + assert len_text == text_token_lens.max() + assert len_speech == speech_token_lens.max() + + out = self.forward( + t3_cond=t3_cond, + text_tokens=text_tokens, + text_token_lens=text_token_lens, + speech_tokens=speech_tokens, + speech_token_lens=speech_token_lens, + ) + + IGNORE_ID = -100 + device = out["text_logits"].device + mask_text = torch.arange(len_text, device=device)[None] >= text_token_lens[:, None] + mask_speech = torch.arange(len_speech, device=device)[None] >= speech_token_lens[:, None] + masked_text = text_tokens.masked_fill(mask_text, IGNORE_ID) + masked_speech = speech_tokens.masked_fill(mask_speech, IGNORE_ID) + + loss_text = F.cross_entropy(out["text_logits"].transpose(1, 2), masked_text, ignore_index=IGNORE_ID) + loss_speech = F.cross_entropy(out["speech_logits"].transpose(1, 2), masked_speech, ignore_index=IGNORE_ID) + + return loss_text, loss_speech + + @torch.inference_mode() + def inference( + self, + *, + t3_cond: T3Cond, + text_tokens: Tensor, + initial_speech_tokens: Optional[Tensor] = None, + num_return_sequences=1, + max_new_tokens=None, + stop_on_eos=True, + do_sample=True, + temperature=0.8, + top_p=0.95, + min_p=0.05, + repetition_penalty=1.2, + cfg_weight=0.5, + ): + """Generate speech tokens from text tokens.""" + from transformers.generation.logits_process import ( + MinPLogitsWarper, + RepetitionPenaltyLogitsProcessor, + TemperatureLogitsWarper, + TopPLogitsWarper, + ) + + text_tokens = torch.atleast_2d(text_tokens).to(dtype=torch.long, device=self.device) + + if initial_speech_tokens is None: + initial_speech_tokens = self.t3_config.start_speech_token * torch.ones_like(text_tokens[:, :1]) + + # Prepare conditioning embeddings (text + initial speech) + embeds, len_cond = self.prepare_input_embeds( + t3_cond=t3_cond, text_tokens=text_tokens, speech_tokens=initial_speech_tokens, cfg_weight=cfg_weight + ) + + # Setup alignment analyzer if needed + if self.t3_config.use_alignment_analyzer: + alignment_analyzer = AlignmentStreamAnalyzer( + self, + text_tokens_slice=(len_cond, len_cond + text_tokens.size(-1)), + alignment_layer_idx=self.t3_config.alignment_layer_idx, + eos_idx=self.t3_config.stop_speech_token, + ) + else: + alignment_analyzer = None + + max_steps = max_new_tokens or self.t3_config.max_speech_tokens + device = embeds.device + use_cfg = cfg_weight > 0.0 and embeds.size(0) > 1 + + # If using CFG, we need manual generation loop (HF generate doesn't support CFG batching) + if use_cfg: + # Manual generation loop for CFG + bos_token = torch.tensor([[self.t3_config.start_speech_token]], dtype=torch.long, device=device) + bos_embed = self.speech_emb(bos_token) + bos_embed = bos_embed + self.speech_pos_emb.get_fixed_embedding(0) + bos_embed = torch.cat([bos_embed, bos_embed], dim=0) # Duplicate for CFG + + inputs_embeds = torch.cat([embeds, bos_embed], dim=1) + generated_ids = bos_token.clone() + + top_p_warper = TopPLogitsWarper(top_p=top_p) + min_p_warper = MinPLogitsWarper(min_p=min_p) + repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(penalty=float(repetition_penalty)) + + output = self( + inputs_embeds=inputs_embeds, + past_key_values=None, + use_cache=True, + output_attentions=alignment_analyzer is not None, + output_hidden_states=True, + return_dict=True, + ) + past = output.past_key_values + + predicted = [] + + for i in range(max_steps): + logits_step = output.logits[:, -1, :] + + # Apply CFG + cond = logits_step[0:1, :] + uncond = logits_step[1:2, :] + cfg = torch.as_tensor(cfg_weight, device=cond.device, dtype=cond.dtype) + logits = cond + cfg * (cond - uncond) + + # Apply alignment analyzer + if alignment_analyzer is not None: + if logits.dim() == 1: + logits = logits.unsqueeze(0) + last_token = generated_ids[0, -1].item() if generated_ids.size(1) > 0 else None + logits = alignment_analyzer.step(logits, next_token=last_token) + + ids_for_proc = generated_ids[:1, ...] + logits = repetition_penalty_processor(ids_for_proc, logits) + + if temperature != 1.0: + logits = logits / temperature + + logits = min_p_warper(ids_for_proc, logits) + logits = top_p_warper(ids_for_proc, logits) + + probs = torch.softmax(logits, dim=-1) + next_token = torch.multinomial(probs, num_samples=1) + predicted.append(next_token) + generated_ids = torch.cat([generated_ids, next_token], dim=1) + + if stop_on_eos and next_token.view(-1) == self.t3_config.stop_speech_token: + logger.info(f"EOS token detected at step {i + 1}") + break + + next_token_embed = self.speech_emb(next_token) + next_token_embed = next_token_embed + self.speech_pos_emb.get_fixed_embedding(i + 1) + next_token_embed = torch.cat([next_token_embed, next_token_embed], dim=0) + + output = self( + inputs_embeds=next_token_embed, + past_key_values=past, + use_cache=True, + output_attentions=alignment_analyzer is not None, + output_hidden_states=True, + return_dict=True, + ) + past = output.past_key_values + + if predicted: + predicted_tokens = torch.cat(predicted, dim=1) + else: + predicted_tokens = torch.empty((1, 0), dtype=torch.long, device=device) + + return predicted_tokens[0] + + else: + # Use HuggingFace's generate() for non-CFG case + # Store decoder conditioning for prepare_inputs_for_generation + self._decoder_cond = embeds + self._current_position = 0 + self.alignment_stream_analyzer = alignment_analyzer + + # Create custom logits processor for alignment analyzer + class CustomLogitsProcessor: + def __init__(self, alignment_analyzer_inst): + self.alignment_analyzer = alignment_analyzer_inst + self.generated_tokens = [] + + def __call__(self, input_ids, scores): + # Apply alignment analyzer + if self.alignment_analyzer is not None: + last_token = input_ids[0, -1].item() if input_ids.size(1) > 0 else None + scores = self.alignment_analyzer.step(scores, next_token=last_token) + + return scores + + # Build logits processors + from transformers.generation.logits_process import LogitsProcessorList + + logits_processors = LogitsProcessorList() + + logits_processors.append(CustomLogitsProcessor(alignment_analyzer)) + logits_processors.append(RepetitionPenaltyLogitsProcessor(penalty=float(repetition_penalty))) + + if temperature != 1.0: + logits_processors.append(TemperatureLogitsWarper(temperature)) + + logits_processors.append(MinPLogitsWarper(min_p=min_p)) + logits_processors.append(TopPLogitsWarper(top_p=top_p)) + + # Generate using HuggingFace's generate (batch size 1) + bos_token = torch.tensor([[self.t3_config.start_speech_token]], dtype=torch.long, device=device) + + generated_ids = self.generate( + input_ids=bos_token, + max_new_tokens=max_steps, + do_sample=do_sample, + logits_processor=logits_processors, + bos_token_id=self.t3_config.start_speech_token, + eos_token_id=self.t3_config.stop_speech_token if stop_on_eos else None, + pad_token_id=self.t3_config.stop_speech_token, + num_return_sequences=num_return_sequences, + output_attentions=self.t3_config.use_alignment_analyzer, + output_hidden_states=True, + return_dict_in_generate=False, + use_cache=True, + ) + + # Extract generated tokens (remove BOS) + predicted_tokens = generated_ids[:, 1:] + + # Clean up generation state + self._decoder_cond = None + self._current_position = 0 + + return predicted_tokens[0] + + @property + def device(self): + """Get device of the model.""" + return next(self.parameters()).device + + +# ============================================================================ +# Chatterbox Model +# ============================================================================ + + +class ChatterboxPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = ChatterboxConfig + base_model_prefix = "chatterbox" + supports_gradient_checkpointing = False + + +@dataclass +class Conditionals: + t3: T3Cond + gen: dict + + +@auto_docstring +class ChatterboxModel(ChatterboxPreTrainedModel): + """ + Complete Chatterbox TTS Pipeline Model. + + This model combines T3, S3Gen, and HiFTNet to provide a complete text-to-speech pipeline: + 1. Text tokens → T3 → Speech tokens + 2. Speech tokens → S3Gen → Mel spectrogram + 3. Mel spectrogram → HiFTNet → Waveform + """ + + def __init__(self, config: ChatterboxConfig): + super().__init__(config) + self.config = config + + # Store configuration for multilingual and hiftnet settings + # Note: hiftnet_config is embedded within s3gen_config, is_multilingual affects t3_config + self.is_multilingual = config.is_multilingual + _ = config.hiftnet_config # Stored in config for serialization + + # Initialize sub-models + logger.info("Initializing T3 model...") + self.t3 = T3Model(config.t3_config) + + logger.info("Initializing S3Gen model...") + self.s3gen = S3GenModel(config.s3gen_config) + + # Sampling rates + self.s3_sr = 16000 # S3 tokenizer sampling rate + self.s3gen_sr = 24000 # S3Gen output sampling rate + + # Text tokenizer + self.text_tokenizer = None + self.feature_extractor = ChatterboxFeatureExtractor( + sampling_rate=self.s3_sr, + s3gen_sampling_rate=self.s3gen_sr, + ) + + # Post init + self.post_init() + + def load_text_tokenizer(self, tokenizer_path): + """Load text tokenizer from tokenizer.json file.""" + tokenizer_path = Path(tokenizer_path) + if not tokenizer_path.exists(): + logger.warning(f"Tokenizer file not found: {tokenizer_path}") + return False + + try: + self.text_tokenizer = Tokenizer.from_file(str(tokenizer_path)) + logger.info(f"✓ Loaded text tokenizer from: {tokenizer_path}") + return True + except Exception as e: + logger.error(f"Failed to load tokenizer: {e}") + return False + + @property + def device(self): + """Get device of the model.""" + return next(self.parameters()).device + + def prepare_text_tokens(self, text: str, tokenizer=None) -> torch.Tensor: + """ + Prepare text tokens from raw text. + + Args: + text: Input text string + tokenizer: Text tokenizer. If None, uses `self.text_tokenizer`. + + Returns: + Text tokens with start/stop markers + """ + text = punc_norm(text) + # Match chatterbox `EnTokenizer`: replace spaces with a dedicated token before encoding. + text = text.replace(" ", "[SPACE]") + + # Use provided tokenizer, or self.text_tokenizer + if tokenizer is not None: + if hasattr(tokenizer, "encode"): + # Tokenizers-style: may return an Encoding with `.ids` or directly a list of ids. + encoding = tokenizer.encode(text) + ids = encoding.ids if hasattr(encoding, "ids") else encoding + text_tokens = torch.tensor([ids], dtype=torch.long) + else: + text_tokens = tokenizer.text_to_tokens(text) + elif self.text_tokenizer is not None: + # Use loaded tokenizer + encoding = self.text_tokenizer.encode(text) + text_tokens = torch.tensor([encoding.ids], dtype=torch.long) + else: + raise ValueError( + "No text tokenizer provided and `self.text_tokenizer` is not loaded. " + "Please pass a tokenizer or call `load_text_tokenizer()` first." + ) + + # Add start/stop tokens if not already present + sot = self.config.t3_config.start_text_token + eot = self.config.t3_config.stop_text_token + + # Check if start/stop tokens are already in the sequence + has_start = (text_tokens[0, 0] == sot).item() if text_tokens.numel() > 0 else False + has_stop = (text_tokens[0, -1] == eot).item() if text_tokens.numel() > 0 else False + + if not has_start: + text_tokens = F.pad(text_tokens, (1, 0), value=sot) + if not has_stop: + text_tokens = F.pad(text_tokens, (0, 1), value=eot) + + return text_tokens.to(self.device) + + def prepare_conditionals( + self, reference_wav: np.ndarray, reference_sr: int, exaggeration: float = 0.5 + ) -> Conditionals: + """ + Mirror the original Chatterbox prepare_conditionals method for parity. + """ + extracted = self.feature_extractor.extract_conditioning( + reference_wav, + reference_sr, + s3gen=self.s3gen, + voice_encoder=self.t3.voice_encoder, + device=self.device, + exaggeration=exaggeration, + speech_cond_prompt_len=self.config.t3_config.speech_cond_prompt_len, + ) + t3_cond = T3Cond( + speaker_emb=extracted["speaker_emb"], + cond_prompt_speech_tokens=extracted["cond_prompt_speech_tokens"], + emotion_adv=extracted["emotion_adv"], + ) + + return Conditionals(t3=t3_cond, gen=extracted["s3gen_ref_dict"]) + + @torch.inference_mode() + def generate( + self, + text: str, + reference_wav: np.ndarray, + reference_sr: int, + tokenizer=None, + exaggeration: float = 0.5, + temperature: float = 0.8, + top_p: float = 0.95, + min_p: float = 0.05, + repetition_penalty: float = 1.2, + cfg_weight: float = 0.5, + max_new_tokens: int = 1000, + return_intermediates: bool = False, + ): + """ + Generate speech from text using the complete pipeline. + + Args: + text: Input text to synthesize + reference_wav: Reference audio for voice cloning (numpy array) + reference_sr: Sampling rate of reference audio + tokenizer: Optional text tokenizer + exaggeration: Emotion/expressiveness level (0.0 to 1.0) + temperature: Sampling temperature for T3 + top_p: Top-p sampling for T3 + min_p: Min-p sampling for T3 + repetition_penalty: Repetition penalty for T3 + cfg_weight: Classifier-free guidance weight for T3 + max_new_tokens: Maximum speech tokens to generate + return_intermediates: Whether to return intermediate outputs (tokens, mel) + + Returns: + Waveform tensor, or tuple of (waveform, intermediates) if return_intermediates=True + """ + logger.info(f"Generating speech for text: '{text}'") + + # Step 1: Prepare text tokens + text_tokens = self.prepare_text_tokens(text, tokenizer) + + # Step 2: Prepare conditionals (T3 + S3Gen) as in original pipeline + conds = self.prepare_conditionals(reference_wav, reference_sr, exaggeration) + t3_cond = conds.t3 + + # For CFG: duplicate text tokens before passing to T3 + t3_text_tokens = text_tokens[0] # Remove batch dimension + if cfg_weight > 0.0: + t3_text_tokens = torch.cat([t3_text_tokens.unsqueeze(0), t3_text_tokens.unsqueeze(0)], dim=0) + + speech_tokens = self.t3.inference( + t3_cond=t3_cond, + text_tokens=t3_text_tokens, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + min_p=min_p, + repetition_penalty=repetition_penalty, + cfg_weight=cfg_weight, + ) + + # Extract conditional batch (first sequence) + if speech_tokens.dim() > 1: + speech_tokens = speech_tokens[0] + + # Clean up speech tokens + speech_tokens = drop_invalid_tokens(speech_tokens) + + # Additional safety check - ensure all tokens are valid + if speech_tokens.max() >= 6561: + speech_tokens = speech_tokens[speech_tokens < 6561] + + # Step 3: Generate waveform with S3Gen using prepared reference dict + wav, _ = self.s3gen.inference( + speech_tokens=speech_tokens, + ref_dict={k: (v.to(self.device) if torch.is_tensor(v) else v) for k, v in conds.gen.items()}, + finalize=True, + ) + + # Remove batch dimension + wav = wav.squeeze(0) + + if return_intermediates: + intermediates = { + "text_tokens": text_tokens, + "speech_tokens": speech_tokens, + } + return wav, intermediates + + return wav + + def forward( + self, + text: str, + reference_wav: np.ndarray, + reference_sr: int, + tokenizer=None, + **kwargs, + ): + """Forward pass - calls generate.""" + return self.generate(text, reference_wav, reference_sr, tokenizer, **kwargs) + + +__all__ = ["ChatterboxPreTrainedModel", "ChatterboxModel"] diff --git a/src/transformers/models/s3gen/__init__.py b/src/transformers/models/s3gen/__init__.py new file mode 100644 index 000000000000..9d4a0327013c --- /dev/null +++ b/src/transformers/models/s3gen/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2025 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 typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_s3gen import * + from .modeling_s3gen import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/src/transformers/models/s3gen/configuration_s3gen.py b/src/transformers/models/s3gen/configuration_s3gen.py new file mode 100644 index 000000000000..cf804e6966eb --- /dev/null +++ b/src/transformers/models/s3gen/configuration_s3gen.py @@ -0,0 +1,318 @@ +# coding=utf-8 +# Copyright 2025 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. +"""S3Gen model configuration""" + +from ...configuration_utils import PretrainedConfig +from ...utils import logging + + +logger = logging.get_logger(__name__) + + +# ============================================================================ +# HiFTNet Configuration +# ============================================================================ + + +class HiFTNetConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a HiFTNet vocoder component. It is used internally by + S3GenModel for the HiFTNet vocoder. + + HiFTNet is a neural vocoder that combines Neural Source Filter with ISTFTNet for high-quality speech synthesis. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + in_channels (`int`, *optional*, defaults to 80): + Number of input channels (mel spectrogram bins). + base_channels (`int`, *optional*, defaults to 512): + Base number of channels for the generator network. + nb_harmonics (`int`, *optional*, defaults to 8): + Number of harmonic overtones for the neural source filter. + sampling_rate (`int`, *optional*, defaults to 22050): + The sampling rate at which the output audio will be generated, expressed in hertz (Hz). + nsf_alpha (`float`, *optional*, defaults to 0.1): + Amplitude of sine-waveform in the neural source filter. + nsf_sigma (`float`, *optional*, defaults to 0.003): + Standard deviation of Gaussian noise in the neural source filter. + nsf_voiced_threshold (`float`, *optional*, defaults to 10.0): + F0 threshold for voiced/unvoiced classification. + upsample_rates (`list[int]`, *optional*, defaults to `[8, 8]`): + A list of integers defining the stride of each 1D convolutional layer in the upsampling network. + upsample_kernel_sizes (`list[int]`, *optional*, defaults to `[16, 16]`): + A list of integers defining the kernel size of each 1D convolutional layer in the upsampling network. + istft_n_fft (`int`, *optional*, defaults to 16): + FFT size for inverse STFT. + istft_hop_len (`int`, *optional*, defaults to 4): + Hop length for inverse STFT. + resblock_kernel_sizes (`list[int]`, *optional*, defaults to `[3, 7, 11]`): + A list of integers defining the kernel sizes of the 1D convolutional layers in the multi-receptive field + fusion (MRF) module. + resblock_dilation_sizes (`list[list[int]]`, *optional*, defaults to `[[1, 3, 5], [1, 3, 5], [1, 3, 5]]`): + A nested list of integers defining the dilation rates of the dilated 1D convolutional layers in the + multi-receptive field fusion (MRF) module. + source_resblock_kernel_sizes (`list[int]`, *optional*, defaults to `[7, 11]`): + A list of integers defining the kernel sizes for source residual blocks. + source_resblock_dilation_sizes (`list[list[int]]`, *optional*, defaults to `[[1, 3, 5], [1, 3, 5]]`): + A nested list of integers defining the dilation rates for source residual blocks. + lrelu_slope (`float`, *optional*, defaults to 0.1): + The slope of the leaky ReLU activation. + audio_limit (`float`, *optional*, defaults to 0.99): + Maximum absolute value for output audio clipping. + f0_predictor_in_channels (`int`, *optional*, defaults to 80): + Input channels for the F0 predictor (should match in_channels). + f0_predictor_cond_channels (`int`, *optional*, defaults to 512): + Conditional channels for the F0 predictor. + """ + + model_type = "hiftnet" + + def __init__( + self, + in_channels=80, + base_channels=512, + nb_harmonics=8, + sampling_rate=22050, + nsf_alpha=0.1, + nsf_sigma=0.003, + nsf_voiced_threshold=10.0, + upsample_rates=None, + upsample_kernel_sizes=None, + istft_n_fft=16, + istft_hop_len=4, + resblock_kernel_sizes=None, + resblock_dilation_sizes=None, + source_resblock_kernel_sizes=None, + source_resblock_dilation_sizes=None, + lrelu_slope=0.1, + audio_limit=0.99, + f0_predictor_in_channels=80, + f0_predictor_cond_channels=512, + **kwargs, + ): + # Set default values for list parameters + if upsample_rates is None: + upsample_rates = [8, 8] + if upsample_kernel_sizes is None: + upsample_kernel_sizes = [16, 16] + if resblock_kernel_sizes is None: + resblock_kernel_sizes = [3, 7, 11] + if resblock_dilation_sizes is None: + resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]] + if source_resblock_kernel_sizes is None: + source_resblock_kernel_sizes = [7, 11] + if source_resblock_dilation_sizes is None: + source_resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5]] + + self.in_channels = in_channels + self.base_channels = base_channels + self.nb_harmonics = nb_harmonics + self.sampling_rate = sampling_rate + self.nsf_alpha = nsf_alpha + self.nsf_sigma = nsf_sigma + self.nsf_voiced_threshold = nsf_voiced_threshold + self.upsample_rates = upsample_rates + self.upsample_kernel_sizes = upsample_kernel_sizes + self.istft_n_fft = istft_n_fft + self.istft_hop_len = istft_hop_len + self.resblock_kernel_sizes = resblock_kernel_sizes + self.resblock_dilation_sizes = resblock_dilation_sizes + self.source_resblock_kernel_sizes = source_resblock_kernel_sizes + self.source_resblock_dilation_sizes = source_resblock_dilation_sizes + self.lrelu_slope = lrelu_slope + self.audio_limit = audio_limit + self.f0_predictor_in_channels = f0_predictor_in_channels + self.f0_predictor_cond_channels = f0_predictor_cond_channels + # Add hidden_size for compatibility with common tests + self.hidden_size = base_channels + + super().__init__(**kwargs) + + +# ============================================================================ +# S3Gen Configuration +# ============================================================================ + + +class S3GenConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`S3GenModel`]. It is used to instantiate a S3Gen + model according to the specified arguments, defining the model architecture. Instantiating a configuration with the + defaults will yield a similar configuration to that of the S3Gen + [ResembleAI/chatterbox-hf](https://huggingface.co/ResembleAI/chatterbox-hf) 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 6561): + Vocabulary size of the S3 speech tokenizer. + token_embed_dim (`int`, *optional*, defaults to 512): + Dimension of the token embeddings. + speaker_feat_dim (`int`, *optional*, defaults to 80): + Number of mel bins for speaker encoder input. + speaker_embed_dim (`int`, *optional*, defaults to 192): + Dimension of the speaker embeddings. + encoder_output_size (`int`, *optional*, defaults to 512): + Output size of the conformer encoder. + encoder_attention_heads (`int`, *optional*, defaults to 8): + Number of attention heads in the conformer encoder. + encoder_linear_units (`int`, *optional*, defaults to 2048): + Dimension of the feedforward network in the conformer encoder. + encoder_num_blocks (`int`, *optional*, defaults to 6): + Number of conformer encoder blocks. + encoder_dropout_rate (`float`, *optional*, defaults to 0.1): + Dropout probability for the encoder. + decoder_in_channels (`int`, *optional*, defaults to 320): + Number of input channels for the conditional decoder (encoder output + speaker embedding). + decoder_out_channels (`int`, *optional*, defaults to 80): + Number of output channels for the decoder (mel bins). + decoder_channels (`List[int]`, *optional*, defaults to `[256]`): + List of channel dimensions for the decoder U-Net. + decoder_n_blocks (`int`, *optional*, defaults to 4): + Number of transformer blocks in each decoder stage. + decoder_num_mid_blocks (`int`, *optional*, defaults to 12): + Number of middle blocks in the decoder U-Net. + decoder_num_heads (`int`, *optional*, defaults to 8): + Number of attention heads in the decoder. + decoder_attention_head_dim (`int`, *optional*, defaults to 64): + Dimension of each attention head in the decoder. + decoder_act_fn (`str`, *optional*, defaults to `"gelu"`): + Activation function for the decoder. + cfm_sigma_min (`float`, *optional*, defaults to 1e-06): + Minimum sigma for conditional flow matching. + cfm_solver (`str`, *optional*, defaults to `"euler"`): + ODE solver for conditional flow matching. + cfm_t_scheduler (`str`, *optional*, defaults to `"cosine"`): + Time scheduler for conditional flow matching. + cfm_inference_cfg_rate (`float`, *optional*, defaults to 0.7): + Classifier-free guidance rate for inference. + sampling_rate (`int`, *optional*, defaults to 24000): + Audio sampling rate in Hz. + mel_bins (`int`, *optional*, defaults to 80): + Number of mel frequency bins. + n_fft (`int`, *optional*, defaults to 1920): + FFT size for mel spectrogram extraction (24kHz sampling rate). + hop_length (`int`, *optional*, defaults to 480): + Hop length for mel spectrogram extraction (24kHz sampling rate). + win_size (`int`, *optional*, defaults to 1920): + Window size for mel spectrogram extraction. + fmin (`int`, *optional*, defaults to 0): + Minimum frequency for mel filter bank. + fmax (`int`, *optional*, defaults to 8000): + Maximum frequency for mel filter bank. + input_frame_rate (`int`, *optional*, defaults to 25): + Frame rate of the S3 tokenizer (25 fps). + token_mel_ratio (`int`, *optional*, defaults to 2): + Ratio between mel frames and tokens (2 mel frames per token). + pre_lookahead_len (`int`, *optional*, defaults to 3): + Pre-lookahead length for causal streaming. + + Example: + + ```python + >>> from transformers import S3GenConfig, S3GenModel + + >>> # Initializing a S3Gen configuration + >>> configuration = S3GenConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = S3GenModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "s3gen" + + def __init__( + self, + # Token embedding + vocab_size=6561, + token_embed_dim=512, + # Speaker encoder (CAMPPlus) + speaker_feat_dim=80, + speaker_embed_dim=192, + # UpsampleConformerEncoder + encoder_output_size=512, + encoder_attention_heads=8, + encoder_linear_units=2048, + encoder_num_blocks=6, + encoder_dropout_rate=0.1, + # ConditionalDecoder (U-Net) + decoder_in_channels=320, + decoder_out_channels=80, + decoder_channels=[256], + decoder_n_blocks=4, + decoder_num_mid_blocks=12, + decoder_num_heads=8, + decoder_attention_head_dim=64, + decoder_act_fn="gelu", + # CFM params + cfm_sigma_min=1e-6, + cfm_solver="euler", + cfm_t_scheduler="cosine", + cfm_inference_cfg_rate=0.7, + # Audio params (mel extraction for reference audio at 24kHz) + sampling_rate=24000, + mel_bins=80, + n_fft=1920, + hop_length=480, + win_size=1920, + fmin=0, + fmax=8000, + # Flow params + input_frame_rate=25, + token_mel_ratio=2, + pre_lookahead_len=3, + **kwargs, + ): + self.vocab_size = vocab_size + self.token_embed_dim = token_embed_dim + self.speaker_feat_dim = speaker_feat_dim + self.speaker_embed_dim = speaker_embed_dim + self.encoder_output_size = encoder_output_size + self.encoder_attention_heads = encoder_attention_heads + self.encoder_linear_units = encoder_linear_units + self.encoder_num_blocks = encoder_num_blocks + self.encoder_dropout_rate = encoder_dropout_rate + self.decoder_in_channels = decoder_in_channels + self.decoder_out_channels = decoder_out_channels + self.decoder_channels = decoder_channels + self.decoder_n_blocks = decoder_n_blocks + self.decoder_num_mid_blocks = decoder_num_mid_blocks + self.decoder_num_heads = decoder_num_heads + self.decoder_attention_head_dim = decoder_attention_head_dim + self.decoder_act_fn = decoder_act_fn + self.cfm_sigma_min = cfm_sigma_min + self.cfm_solver = cfm_solver + self.cfm_t_scheduler = cfm_t_scheduler + self.cfm_inference_cfg_rate = cfm_inference_cfg_rate + self.sampling_rate = sampling_rate + self.mel_bins = mel_bins + self.n_fft = n_fft + self.hop_length = hop_length + self.win_size = win_size + self.fmin = fmin + self.fmax = fmax + self.input_frame_rate = input_frame_rate + self.token_mel_ratio = token_mel_ratio + self.pre_lookahead_len = pre_lookahead_len + super().__init__(**kwargs) + + +__all__ = ["S3GenConfig"] diff --git a/src/transformers/models/s3gen/modeling_s3gen.py b/src/transformers/models/s3gen/modeling_s3gen.py new file mode 100644 index 000000000000..697633a2b02d --- /dev/null +++ b/src/transformers/models/s3gen/modeling_s3gen.py @@ -0,0 +1,2747 @@ +# coding=utf-8 +# Copyright 2025 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 S3Gen model.""" + +import logging +import math +from typing import Optional + +import numpy as np +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 scipy.signal import get_window +from torch.distributions.uniform import Uniform +from torch.nn import Conv1d, ConvTranspose1d +from torch.nn.utils import remove_weight_norm +from torch.nn.utils.parametrizations import weight_norm + +from ...modeling_utils import PreTrainedModel +from ...utils import add_start_docstrings_to_model_forward, auto_docstring +from ..dac.modeling_dac import Snake1d +from ..s3tokenizer.configuration_s3tokenizer import S3TokenizerConfig +from ..s3tokenizer.feature_extraction_s3tokenizer import S3TokenizerFeatureExtractor +from ..s3tokenizer.modeling_s3tokenizer import S3TokenizerModel +from .configuration_s3gen import HiFTNetConfig, S3GenConfig + + +logger = logging.getLogger(__name__) + +# Global state for mel spectrogram computation +mel_basis = {} +hann_window = {} + + +# Utility functions +def pad_list(xs, pad_value): + """Perform padding for the list of tensors.""" + n_batch = len(xs) + max_len = max(x.size(0) for x in xs) + pad = xs[0].new(n_batch, max_len, *xs[0].size()[1:]).fill_(pad_value) + for i in range(n_batch): + pad[i, : xs[i].size(0)] = xs[i] + return pad + + +def extract_feature(audio): + """Extract fbank features for speaker encoder.""" + features = [] + feature_lengths = [] + for au in audio: + feature = Kaldi.fbank(au.unsqueeze(0), num_mel_bins=80) + feature = feature - feature.mean(dim=0, keepdim=True) + features.append(feature) + feature_lengths.append(feature.shape[0]) + features_padded = pad_list(features, pad_value=0) + return features_padded, feature_lengths + + +def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: + """Make mask tensor containing indices of padded part.""" + batch_size = lengths.size(0) + max_len = max_len if max_len > 0 else lengths.max().item() + seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device) + seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len) + seq_length_expand = lengths.unsqueeze(-1) + mask = seq_range_expand >= seq_length_expand + return mask + + +def mel_spectrogram( + y, n_fft=1024, num_mels=80, sampling_rate=24000, hop_size=256, win_size=1024, fmin=0, fmax=8000, center=False +): + """Extract mel spectrogram from audio.""" + if isinstance(y, np.ndarray): + y = torch.tensor(y).float() + if len(y.shape) == 1: + y = y[None,] + + global mel_basis, hann_window + if f"{str(fmax)}_{str(y.device)}" not in mel_basis: + mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax) + mel_basis[str(fmax) + "_" + str(y.device)] = torch.from_numpy(mel).float().to(y.device) + hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device) + + y = F.pad(y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode="reflect") + y = y.squeeze(1) + + spec = torch.view_as_real( + torch.stft( + y, + n_fft, + hop_length=hop_size, + win_length=win_size, + window=hann_window[str(y.device)], + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + ) + + spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-9) + spec = torch.matmul(mel_basis[str(fmax) + "_" + str(y.device)], spec) + spec = torch.log(torch.clamp(spec, min=1e-5)) + + return spec + + +# CAMPPlus Speaker Encoder Components +class BasicResBlock(nn.Module): + expansion = 1 + + def __init__(self, in_planes, planes, stride=1): + super().__init__() + self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + + self.shortcut = nn.ModuleList() + if stride != 1 or in_planes != self.expansion * planes: + self.shortcut.append( + nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=(stride, 1), bias=False) + ) + self.shortcut.append(nn.BatchNorm2d(self.expansion * planes)) + + def forward(self, x): + out = F.relu(self.bn1(self.conv1(x))) + out = self.bn2(self.conv2(out)) + + shortcut_out = x + for layer in self.shortcut: + shortcut_out = layer(shortcut_out) + out += shortcut_out + + out = F.relu(out) + return out + + +class FCM(nn.Module): + """Frequency Channel Masking module.""" + + def __init__(self, block=BasicResBlock, num_blocks=[2, 2], m_channels=32, feat_dim=80): + super().__init__() + self.in_planes = m_channels + self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(m_channels) + + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2) + self.layer2 = self._make_layer(block, m_channels, num_blocks[0], stride=2) + + self.conv2 = nn.Conv2d(m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(m_channels) + self.out_channels = m_channels * (feat_dim // 8) + + def _make_layer(self, block, planes, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_planes, planes, stride)) + self.in_planes = planes * block.expansion + return nn.ModuleList(layers) + + def forward(self, x): + x = x.unsqueeze(1) + out = F.relu(self.bn1(self.conv1(x))) + for layer in self.layer1: + out = layer(out) + for layer in self.layer2: + out = layer(out) + out = F.relu(self.bn2(self.conv2(out))) + shape = out.shape + out = out.reshape(shape[0], shape[1] * shape[2], shape[3]) + return out + + +def get_nonlinear(config_str, channels): + """Create non-linear activation module.""" + nonlinear = nn.ModuleDict() + if "batchnorm" in config_str: + affine = "batchnorm_" not in config_str + nonlinear["batchnorm"] = nn.BatchNorm1d(channels, affine=affine) + if "relu" in config_str: + nonlinear["relu"] = nn.ReLU(inplace=True) + return nonlinear + + +def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, eps=1e-2): + """Compute mean and standard deviation statistics.""" + mean = x.mean(dim=dim) + std = x.std(dim=dim, unbiased=unbiased) + stats = torch.cat([mean, std], dim=-1) + if keepdim: + stats = stats.unsqueeze(dim=dim) + return stats + + +class StatsPool(nn.Module): + def forward(self, x): + return statistics_pooling(x) + + +class TDNNLayer(nn.Module): + """Time Delay Neural Network layer.""" + + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + bias=False, + config_str="batchnorm-relu", + ): + super().__init__() + if padding < 0: + assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})" + padding = (kernel_size - 1) // 2 * dilation + self.linear = nn.Conv1d( + in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias + ) + self.nonlinear = get_nonlinear(config_str, out_channels) + + def forward(self, x): + x = self.linear(x) + for layer in self.nonlinear.values(): + x = layer(x) + return x + + +class CAMLayer(nn.Module): + """Context-Aware Masking layer.""" + + def __init__(self, bn_channels, out_channels, kernel_size, stride, padding, dilation, bias, reduction=2): + super().__init__() + self.linear_local = nn.Conv1d( + bn_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias + ) + self.linear1 = nn.Conv1d(bn_channels, bn_channels // reduction, 1) + self.relu = nn.ReLU(inplace=True) + self.linear2 = nn.Conv1d(bn_channels // reduction, out_channels, 1) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + y = self.linear_local(x) + context = x.mean(-1, keepdim=True) + self.seg_pooling(x) + context = self.relu(self.linear1(context)) + m = self.sigmoid(self.linear2(context)) + return y * m + + def seg_pooling(self, x, seg_len=100, stype="avg"): + if stype == "avg": + seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True) + elif stype == "max": + seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True) + else: + raise ValueError("Wrong segment pooling type.") + shape = seg.shape + seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1) + seg = seg[..., : x.shape[-1]] + return seg + + +class CAMDenseTDNNLayer(nn.Module): + """Dense TDNN layer with CAM.""" + + def __init__( + self, + in_channels, + out_channels, + bn_channels, + kernel_size, + stride=1, + dilation=1, + bias=False, + config_str="batchnorm-relu", + memory_efficient=False, + ): + super().__init__() + assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})" + padding = (kernel_size - 1) // 2 * dilation + self.memory_efficient = memory_efficient + self.nonlinear1 = get_nonlinear(config_str, in_channels) + self.linear1 = nn.Conv1d(in_channels, bn_channels, 1, bias=False) + self.nonlinear2 = get_nonlinear(config_str, bn_channels) + self.cam_layer = CAMLayer( + bn_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias + ) + + def bn_function(self, x): + for layer in self.nonlinear1.values(): + x = layer(x) + return self.linear1(x) + + def forward(self, x): + x = self.bn_function(x) + for layer in self.nonlinear2.values(): + x = layer(x) + x = self.cam_layer(x) + return x + + +class CAMDenseTDNNBlock(nn.ModuleList): + """Dense TDNN block with CAM.""" + + def __init__( + self, + num_layers, + in_channels, + out_channels, + bn_channels, + kernel_size, + stride=1, + dilation=1, + bias=False, + config_str="batchnorm-relu", + memory_efficient=False, + ): + super().__init__() + for i in range(num_layers): + layer = CAMDenseTDNNLayer( + in_channels=in_channels + i * out_channels, + out_channels=out_channels, + bn_channels=bn_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + bias=bias, + config_str=config_str, + memory_efficient=memory_efficient, + ) + self.add_module(f"tdnnd{i + 1}", layer) + + def forward(self, x): + for layer in self: + x = torch.cat([x, layer(x)], dim=1) + return x + + +class TransitLayer(nn.Module): + """Transition layer between blocks.""" + + def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"): + super().__init__() + self.nonlinear = get_nonlinear(config_str, in_channels) + self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias) + + def forward(self, x): + for layer in self.nonlinear.values(): + x = layer(x) + x = self.linear(x) + return x + + +class DenseLayer(nn.Module): + """Dense layer.""" + + def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"): + super().__init__() + self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias) + self.nonlinear = get_nonlinear(config_str, out_channels) + + def forward(self, x): + if len(x.shape) == 2: + x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1) + else: + x = self.linear(x) + for layer in self.nonlinear.values(): + x = layer(x) + return x + + +# ============================================================================ +# HiFTNet Vocoder Components +# ============================================================================ + + +def get_padding(kernel_size: int, dilation: int = 1) -> int: + """Calculate padding to maintain sequence length.""" + return int((kernel_size * dilation - dilation) / 2) + + +class Snake(Snake1d): + """ + Implementation of a sine-based periodic activation function. + Inherits from Snake1d for modularity. + """ + + def __init__( + self, in_features: int, alpha: float = 1.0, alpha_trainable: bool = True, alpha_logscale: bool = False + ): + super().__init__(in_features) + # Re-initialize to match Chatterbox original shapes (C,) instead of (1, C, 1) + # to ensure checkpoint compatibility. + self.alpha = nn.Parameter(torch.ones(in_features) * alpha) + self.alpha.requires_grad = alpha_trainable + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # We override forward to handle the (C,) -> (1, C, 1) unsqueeze + # while keeping the core logic identical to Snake1d. + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) + return x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) + + +class ResBlock(nn.Module): + """Residual block module in HiFiGAN/BigVGAN.""" + + def __init__( + self, + channels: int = 512, + kernel_size: int = 3, + dilations: tuple[int] = (1, 3, 5), + ): + super().__init__() + self.convs1 = nn.ModuleList() + self.convs2 = nn.ModuleList() + + for dilation in dilations: + self.convs1.append( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation, + padding=get_padding(kernel_size, dilation), + ) + ) + self.convs2.append( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ) + # Note: Weights will be initialized by _init_weights in post_init() + self.activations1 = nn.ModuleList([Snake(channels, alpha_logscale=False) for _ in range(len(self.convs1))]) + self.activations2 = nn.ModuleList([Snake(channels, alpha_logscale=False) for _ in range(len(self.convs2))]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for idx in range(len(self.convs1)): + xt = self.activations1[idx](x) + xt = self.convs1[idx](xt) + xt = self.activations2[idx](xt) + xt = self.convs2[idx](xt) + x = xt + x + return x + + def apply_weight_norm(self): + for idx in range(len(self.convs1)): + weight_norm(self.convs1[idx]) + weight_norm(self.convs2[idx]) + + def remove_weight_norm(self): + for idx in range(len(self.convs1)): + remove_weight_norm(self.convs1[idx]) + remove_weight_norm(self.convs2[idx]) + + +class SineGen(nn.Module): + """ + Definition of sine generator for neural source filter. + + SineGen(samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0) + + Args: + samp_rate: sampling rate in Hz + harmonic_num: number of harmonic overtones (default 0) + sine_amp: amplitude of sine-waveform (default 0.1) + noise_std: std of Gaussian noise (default 0.003) + voiced_threshold: F0 threshold for U/V classification (default 0) + """ + + def __init__( + self, + samp_rate: int, + harmonic_num: int = 0, + sine_amp: float = 0.1, + noise_std: float = 0.003, + voiced_threshold: float = 0, + ): + super().__init__() + self.sine_amp = sine_amp + self.noise_std = noise_std + self.harmonic_num = harmonic_num + self.sampling_rate = samp_rate + self.voiced_threshold = voiced_threshold + + def _f02uv(self, f0: torch.Tensor) -> torch.Tensor: + """Generate unvoiced/voiced signal.""" + uv = (f0 > self.voiced_threshold).type(torch.float32) + return uv + + @torch.no_grad() + def forward(self, f0: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Args: + f0: [B, 1, sample_len], Hz + + Returns: + sine_waves: [B, harmonic_num+1, sample_len] + uv: [B, 1, sample_len] + noise: [B, harmonic_num+1, sample_len] + """ + F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device) + for i in range(self.harmonic_num + 1): + F_mat[:, i : i + 1, :] = f0 * (i + 1) / self.sampling_rate + + theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1) + u_dist = Uniform(low=-np.pi, high=np.pi) + phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device) + phase_vec[:, 0, :] = 0 + + # generate sine waveforms + sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec) + + # generate uv signal + uv = self._f02uv(f0) + + # noise: for unvoiced should be similar to sine_amp + # std = self.sine_amp/3 -> max value ~ self.sine_amp + # for voiced regions is self.noise_std + noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 + noise = noise_amp * torch.randn_like(sine_waves) + + # first: set the unvoiced part to 0 by uv + # then: additive noise + sine_waves = sine_waves * uv + noise + return sine_waves, uv, noise + + +class SourceModuleHnNSF(nn.Module): + """ + Source Module for harmonic-plus-noise neural source filter. + + Args: + sampling_rate: sampling rate in Hz + upsample_scale: total upsampling scale factor + harmonic_num: number of harmonics above F0 (default: 0) + sine_amp: amplitude of sine source signal (default: 0.1) + add_noise_std: std of additive Gaussian noise (default: 0.003) + voiced_threshold: threshold to set U/V given F0 (default: 0) + """ + + def __init__( + self, + sampling_rate: int, + upsample_scale: int, + harmonic_num: int = 0, + sine_amp: float = 0.1, + add_noise_std: float = 0.003, + voiced_threshold: float = 0, + ): + super().__init__() + + self.sine_amp = sine_amp + self.noise_std = add_noise_std + + # to produce sine waveforms + self.l_sin_gen = SineGen(sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold) + + # to merge source harmonics into a single excitation + self.l_linear = nn.Linear(harmonic_num + 1, 1) + self.l_tanh = nn.Tanh() + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Args: + x: F0 sampled [B, T, 1] + + Returns: + sine_merge: [B, T, 1] + noise: [B, T, 1] + uv: [B, T, 1] + """ + # source for harmonic branch + with torch.no_grad(): + sine_wavs, uv, _ = self.l_sin_gen(x.transpose(1, 2)) + sine_wavs = sine_wavs.transpose(1, 2) + uv = uv.transpose(1, 2) + sine_merge = self.l_tanh(self.l_linear(sine_wavs)) + + # source for noise branch, in the same shape as uv + noise = torch.randn_like(uv) * self.sine_amp / 3 + return sine_merge, noise, uv + + +class ConvRNNF0Predictor(nn.Module): + """ + Convolutional RNN-based F0 predictor. + + Args: + num_class: number of output classes (default 1 for F0 regression) + in_channels: input feature dimension + cond_channels: conditional feature dimension + """ + + def __init__( + self, + num_class: int = 1, + in_channels: int = 80, + cond_channels: int = 512, + ): + super().__init__() + + self.num_class = num_class + # Using numeric keys to match state_dict + self.condnet = nn.ModuleList( + [ + nn.Conv1d(in_channels, cond_channels, kernel_size=3, padding=1), + nn.ELU(), + nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1), + nn.ELU(), + nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1), + nn.ELU(), + nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1), + nn.ELU(), + nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1), + nn.ELU(), + ] + ) + self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x: [B, C, T] mel spectrogram + + Returns: + f0: [B, T] predicted F0 + """ + for layer in self.condnet: + x = layer(x) + x = x.transpose(1, 2) + return torch.abs(self.classifier(x).squeeze(-1)) + + def apply_weight_norm(self): + for module in self.condnet: + if isinstance(module, nn.Conv1d): + weight_norm(module) + + def remove_weight_norm(self): + for module in self.condnet: + if isinstance(module, nn.Conv1d): + try: + remove_weight_norm(module) + except ValueError: + pass + + +class HiFTGenerator(nn.Module): + """ + HiFTNet Generator: Neural Source Filter + ISTFTNet. + + This is the core HiFTNet vocoder model that combines a neural source filter + with an inverse STFT network for high-quality speech synthesis. + + Reference: https://arxiv.org/abs/2309.09493 + """ + + def __init__(self, config: HiFTNetConfig): + super().__init__() + + self.config = config + self.out_channels = 1 + self.nb_harmonics = config.nb_harmonics + self.sampling_rate = config.sampling_rate + self.lrelu_slope = config.lrelu_slope + self.audio_limit = config.audio_limit + + # ISTFT parameters + self.istft_params = {"n_fft": config.istft_n_fft, "hop_len": config.istft_hop_len} + + self.num_kernels = len(config.resblock_kernel_sizes) + self.num_upsamples = len(config.upsample_rates) + + # Neural source filter + upsample_scale = int(np.prod(config.upsample_rates) * config.istft_hop_len) + self.m_source = SourceModuleHnNSF( + sampling_rate=config.sampling_rate, + upsample_scale=upsample_scale, + harmonic_num=config.nb_harmonics, + sine_amp=config.nsf_alpha, + add_noise_std=config.nsf_sigma, + voiced_threshold=config.nsf_voiced_threshold, + ) + self.f0_upsamp = nn.Upsample(scale_factor=upsample_scale) + + # F0 predictor + self.f0_predictor = ConvRNNF0Predictor( + num_class=1, + in_channels=config.f0_predictor_in_channels, + cond_channels=config.f0_predictor_cond_channels, + ) + + # Pre-convolution + self.conv_pre = Conv1d(config.in_channels, config.base_channels, 7, 1, padding=3) + + # Upsampling layers + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes)): + self.ups.append( + ConvTranspose1d( + config.base_channels // (2**i), + config.base_channels // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + + # Source downsampling and residual blocks + self.source_downs = nn.ModuleList() + self.source_resblocks = nn.ModuleList() + downsample_rates = [1] + config.upsample_rates[::-1][:-1] + downsample_cum_rates = np.cumprod(downsample_rates) + + for i, (u, k, d) in enumerate( + zip( + downsample_cum_rates[::-1], + config.source_resblock_kernel_sizes, + config.source_resblock_dilation_sizes, + ) + ): + if u == 1: + self.source_downs.append( + Conv1d(self.istft_params["n_fft"] + 2, config.base_channels // (2 ** (i + 1)), 1, 1) + ) + else: + self.source_downs.append( + Conv1d( + self.istft_params["n_fft"] + 2, + config.base_channels // (2 ** (i + 1)), + u * 2, + u, + padding=(u // 2), + ) + ) + + self.source_resblocks.append(ResBlock(config.base_channels // (2 ** (i + 1)), k, d)) + + # Main residual blocks + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = config.base_channels // (2 ** (i + 1)) + for k, d in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes): + self.resblocks.append(ResBlock(ch, k, d)) + + # Post-convolution + self.conv_post = Conv1d(ch, self.istft_params["n_fft"] + 2, 7, 1, padding=3) + # Note: Weights will be initialized by _init_weights in post_init() + + # Reflection padding and STFT window + self.reflection_pad = nn.ReflectionPad1d((1, 0)) + stft_window = torch.from_numpy(get_window("hann", self.istft_params["n_fft"], fftbins=True).astype(np.float32)) + self.register_buffer("stft_window", stft_window) + + # Apply weight normalization to match checkpoint + self.apply_weight_norm() + + def apply_weight_norm(self): + """Apply weight normalization to all relevant layers.""" + logger.info("Applying weight norm...") + for l in self.ups: + weight_norm(l) + for l in self.resblocks: + l.apply_weight_norm() + weight_norm(self.conv_pre) + weight_norm(self.conv_post) + for l in self.source_resblocks: + l.apply_weight_norm() + # Apply weight norm to F0 predictor + self.f0_predictor.apply_weight_norm() + + def remove_weight_norm(self): + """Remove weight normalization from all layers.""" + logger.info("Removing weight norm...") + for l in self.ups: + remove_weight_norm(l) + for l in self.resblocks: + l.remove_weight_norm() + remove_weight_norm(self.conv_pre) + remove_weight_norm(self.conv_post) + for l in self.source_resblocks: + l.remove_weight_norm() + # Remove weight norm from F0 predictor + self.f0_predictor.remove_weight_norm() + + def _stft(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Compute STFT.""" + spec = torch.stft( + x, + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=self.stft_window.to(x.device), + return_complex=True, + ) + spec = torch.view_as_real(spec) # [B, F, TT, 2] + return spec[..., 0], spec[..., 1] + + def _istft(self, magnitude: torch.Tensor, phase: torch.Tensor) -> torch.Tensor: + """Compute inverse STFT.""" + magnitude = torch.clip(magnitude, max=1e2) + real = magnitude * torch.cos(phase) + img = magnitude * torch.sin(phase) + inverse_transform = torch.istft( + torch.complex(real, img), + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=self.stft_window.to(magnitude.device), + ) + return inverse_transform + + def decode(self, x: torch.Tensor, s: torch.Tensor = None) -> torch.Tensor: + """ + Decode mel spectrogram to waveform. + + Args: + x: [B, C, T] mel spectrogram + s: [B, 1, T_audio] source signal (optional, defaults to zeros) + + Returns: + waveform: [B, T_audio] + """ + if s is None: + s = torch.zeros(x.size(0), 1, 0).to(x.device) + + s_stft_real, s_stft_imag = self._stft(s.squeeze(1)) + s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1) + + x = self.conv_pre(x) + for i in range(self.num_upsamples): + x = F.leaky_relu(x, self.lrelu_slope) + x = self.ups[i](x) + + if i == self.num_upsamples - 1: + x = self.reflection_pad(x) + + # Fusion with source + si = self.source_downs[i](s_stft) + si = self.source_resblocks[i](si) + x = x + si + + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + x = F.leaky_relu(x) + x = self.conv_post(x) + magnitude = torch.exp(x[:, : self.istft_params["n_fft"] // 2 + 1, :]) + phase = torch.sin(x[:, self.istft_params["n_fft"] // 2 + 1 :, :]) + + x = self._istft(magnitude, phase) + x = torch.clamp(x, -self.audio_limit, self.audio_limit) + return x + + def forward(self, speech_feat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Forward pass. + + Args: + speech_feat: [B, T, C] mel spectrogram (will be transposed internally) + + Returns: + generated_speech: [B, T_audio] + f0: [B, T] + """ + speech_feat = speech_feat.transpose(1, 2) # [B, C, T] + + # Predict F0 + f0 = self.f0_predictor(speech_feat) + + # Generate source signal + s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # [B, T_upsampled, 1] + s, _, _ = self.m_source(s) + s = s.transpose(1, 2) # [B, 1, T_upsampled] + + # Generate waveform + generated_speech = self.decode(x=speech_feat, s=s) + return generated_speech, f0 + + @torch.inference_mode() + def inference( + self, speech_feat: torch.Tensor, cache_source: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Inference method with source caching support. + + Args: + speech_feat: [B, C, T] mel spectrogram + cache_source: [B, 1, T_cache] cached source signal (optional) + + Returns: + generated_speech: [B, T_audio] + s: [B, 1, T_audio] source signal for caching + """ + if cache_source is None: + cache_source = torch.zeros(1, 1, 0).to(speech_feat.device) + + # Predict F0 + f0 = self.f0_predictor(speech_feat) + + # Generate source signal + s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # [B, T_upsampled, 1] + s, _, _ = self.m_source(s) + s = s.transpose(1, 2) # [B, 1, T_upsampled] + + # Use cache_source to avoid glitch + if cache_source.shape[2] != 0: + s[:, :, : cache_source.shape[2]] = cache_source + else: + # Smoothly fade-in the source when starting from scratch to reduce onset transients. + n_fade = min(int(self.sampling_rate // 40), s.size(2)) # ~25ms at 24kHz + if n_fade > 1: + fade = (torch.cos(torch.linspace(torch.pi, 0, n_fade, device=s.device, dtype=s.dtype)) + 1) / 2 + s[:, :, :n_fade] *= fade + + generated_speech = self.decode(x=speech_feat, s=s) + return generated_speech, s + + +# ============================================================================ +# Transformer Components (replacing diffusers dependency) +# ============================================================================ + + +# Decorator replacement +def maybe_allow_in_graph(cls): + """Decorator to allow class in graph (simplified version).""" + return cls + + +# Linear layer replacement +class LoRACompatibleLinear(nn.Linear): + """ + A Linear layer that can be used as a drop-in replacement for `torch.nn.Linear`. + Simplified version without LoRA support. + """ + + def __init__(self, in_features, out_features, bias=True): + super().__init__(in_features, out_features, bias=bias) + + +# Activation functions +class GELU(nn.Module): + """GELU activation function.""" + + def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out, bias=bias) + self.approximate = approximate + + def forward(self, hidden_states): + hidden_states = self.proj(hidden_states) + if self.approximate == "tanh": + return F.gelu(hidden_states, approximate="tanh") + return F.gelu(hidden_states) + + +class GEGLU(nn.Module): + """GEGLU activation function.""" + + def __init__(self, dim_in: int, dim_out: int, bias: bool = True): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias) + + def forward(self, hidden_states): + hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1) + return hidden_states * F.gelu(gate) + + +class ApproximateGELU(nn.Module): + """Approximate GEGLU activation function.""" + + def __init__(self, dim_in: int, dim_out: int, bias: bool = True): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias) + + def forward(self, hidden_states): + hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1) + return hidden_states * F.gelu(gate, approximate="tanh") + + +class SnakeBeta(nn.Module): + """ + A modified Snake function which uses separate parameters for the magnitude of the periodic components. + + Shape: + - Input: (B, C, T) or (B, T, C) + - Output: same shape as input + + Parameters: + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + """ + + def __init__(self, in_features, out_features, alpha=1.0, alpha_trainable=True, alpha_logscale=True): + super().__init__() + self.in_features = out_features if isinstance(out_features, list) else [out_features] + self.proj = LoRACompatibleLinear(in_features, out_features) + + # initialize alpha and beta + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = nn.Parameter(torch.zeros(self.in_features) * alpha) + self.beta = nn.Parameter(torch.zeros(self.in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = nn.Parameter(torch.ones(self.in_features) * alpha) + self.beta = nn.Parameter(torch.ones(self.in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + """ + Forward pass of the function. + Applies the function to the input elementwise. + SnakeBeta ∶= x + 1/b * sin^2 (xa) + """ + x = self.proj(x) + if self.alpha_logscale: + alpha = torch.exp(self.alpha) + beta = torch.exp(self.beta) + else: + alpha = self.alpha + beta = self.beta + + x = x + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(torch.sin(x * alpha), 2) + return x + + +# Normalization layers +class AdaLayerNorm(nn.Module): + """Adaptive Layer Normalization.""" + + def __init__(self, embedding_dim: int, num_embeddings: int): + super().__init__() + self.emb = nn.Embedding(num_embeddings, embedding_dim * 2) + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, embedding_dim * 2) + self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False) + + def forward(self, x, timestep): + emb = self.linear(self.silu(self.emb(timestep))) + scale, shift = torch.chunk(emb, 2, dim=-1) + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +class AdaLayerNormZero(nn.Module): + """Adaptive Layer Normalization Zero.""" + + def __init__(self, embedding_dim: int, num_embeddings: int): + super().__init__() + self.emb = nn.Embedding(num_embeddings, embedding_dim * 6) + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, embedding_dim * 6) + self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x, timestep, class_labels=None, hidden_dtype=None): + emb = self.linear(self.silu(self.emb(timestep))) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1) + x = self.norm(x) * (1 + scale_msa)[:, None, :] + shift_msa[:, None, :] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +# Attention processor +class Attention(nn.Module): + """Multi-head attention module.""" + + def __init__( + self, + query_dim: int, + cross_attention_dim: Optional[int] = None, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = False, + upcast_attention: bool = False, + out_bias: bool = True, + ): + super().__init__() + inner_dim = dim_head * heads + cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim + self.scale = dim_head**-0.5 + self.heads = heads + self.upcast_attention = upcast_attention + + self.to_q = LoRACompatibleLinear(query_dim, inner_dim, bias=bias) + self.to_k = LoRACompatibleLinear(cross_attention_dim, inner_dim, bias=bias) + self.to_v = LoRACompatibleLinear(cross_attention_dim, inner_dim, bias=bias) + + self.to_out = nn.ModuleList([]) + self.to_out.append(LoRACompatibleLinear(inner_dim, query_dim, bias=out_bias)) + self.to_out.append(nn.Dropout(dropout)) + + def forward( + self, + hidden_states, + encoder_hidden_states=None, + attention_mask=None, + **cross_attention_kwargs, + ): + batch_size, sequence_length, _ = hidden_states.shape + + encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states + + query = self.to_q(hidden_states) + key = self.to_k(encoder_hidden_states) + value = self.to_v(encoder_hidden_states) + + 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) + + # Compute attention + if self.upcast_attention: + query = query.float() + key = key.float() + + attention_scores = torch.matmul(query, key.transpose(-1, -2)) * self.scale + + if attention_mask is not None: + attention_scores = attention_scores + attention_mask + + attention_probs = F.softmax(attention_scores, dim=-1) + + if self.upcast_attention: + attention_probs = attention_probs.to(value.dtype) + + hidden_states = torch.matmul(attention_probs, value) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) + + # Linear projection + for layer in self.to_out: + hidden_states = layer(hidden_states) + + return hidden_states + + +# Feed-forward layer +class FeedForward(nn.Module): + """ + A feed-forward layer. + + Parameters: + dim (`int`): The number of channels in the input. + dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. + mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + final_dropout (`bool` *optional*, defaults to False): Apply a final dropout. + """ + + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + mult: int = 4, + dropout: float = 0.0, + activation_fn: str = "geglu", + final_dropout: bool = False, + ): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + if activation_fn == "gelu": + act_fn = GELU(dim, inner_dim, bias=True) + elif activation_fn == "gelu-approximate": + act_fn = GELU(dim, inner_dim, approximate="tanh", bias=True) + elif activation_fn == "geglu": + act_fn = GEGLU(dim, inner_dim, bias=True) + elif activation_fn == "geglu-approximate": + act_fn = ApproximateGELU(dim, inner_dim, bias=True) + elif activation_fn == "snakebeta": + act_fn = SnakeBeta(dim, inner_dim) + else: + act_fn = GEGLU(dim, inner_dim, bias=True) + + self.net = nn.ModuleList([]) + # project in + self.net.append(act_fn) + # project dropout + self.net.append(nn.Dropout(dropout)) + # project out + self.net.append(nn.Linear(inner_dim, dim_out, bias=True)) + # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout + if final_dropout: + self.net.append(nn.Dropout(dropout)) + + def forward(self, hidden_states): + for module in self.net: + hidden_states = module(hidden_states) + return hidden_states + + +# BasicTransformerBlock +@maybe_allow_in_graph +class BasicTransformerBlock(nn.Module): + """ + A basic Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + only_cross_attention (`bool`, *optional*): + Whether to use only cross-attention layers. In this case two cross attention layers are used. + double_self_attention (`bool`, *optional*): + Whether to use two self-attention layers. In this case no cross attention layers are used. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + attention_bias: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_elementwise_affine: bool = True, + norm_type: str = "layer_norm", + final_dropout: bool = False, + ): + super().__init__() + self.only_cross_attention = only_cross_attention + + self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" + self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" + + if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: + raise ValueError( + f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" + f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." + ) + + # Define 3 blocks. Each block has its own normalization layer. + # 1. Self-Attn + if self.use_ada_layer_norm: + self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) + elif self.use_ada_layer_norm_zero: + self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) + else: + self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) + + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + ) + + # 2. Cross-Attn + if cross_attention_dim is not None or double_self_attention: + # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. + # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during + # the second cross attention block. + self.norm2 = ( + AdaLayerNorm(dim, num_embeds_ada_norm) + if self.use_ada_layer_norm + else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) + ) + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim if not double_self_attention else None, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) + else: + self.norm2 = None + self.attn2 = None + + # 3. Feed-forward + self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine) + self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout) + + # let chunk size default to None + self._chunk_size = None + self._chunk_dim = 0 + + def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int): + # Sets chunk feed-forward + self._chunk_size = chunk_size + self._chunk_dim = dim + + def forward( + self, + hidden_states: torch.FloatTensor, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + timestep: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Optional[dict] = None, + class_labels: Optional[torch.LongTensor] = None, + ): + # Notice that normalization is always applied before the real computation in the following blocks. + # 1. Self-Attention + if self.use_ada_layer_norm: + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.use_ada_layer_norm_zero: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + else: + norm_hidden_states = self.norm1(hidden_states) + + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, + attention_mask=encoder_attention_mask if self.only_cross_attention else attention_mask, + **cross_attention_kwargs, + ) + if self.use_ada_layer_norm_zero: + attn_output = gate_msa.unsqueeze(1) * attn_output + hidden_states = attn_output + hidden_states + + # 2. Cross-Attention + if self.attn2 is not None: + norm_hidden_states = ( + self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states) + ) + + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 3. Feed-forward + norm_hidden_states = self.norm3(hidden_states) + + if self.use_ada_layer_norm_zero: + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + + if self._chunk_size is not None: + # "feed_forward_chunk_size" can be used to save memory + if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0: + raise ValueError( + f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." + ) + + num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size + ff_output = torch.cat( + [self.ff(hid_slice) for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim)], + dim=self._chunk_dim, + ) + else: + ff_output = self.ff(norm_hidden_states) + + if self.use_ada_layer_norm_zero: + ff_output = gate_mlp.unsqueeze(1) * ff_output + + hidden_states = ff_output + hidden_states + + return hidden_states + + +# ============================================================================ +# S3Gen Components +# ============================================================================ + + +class CAMPPlus(nn.Module): + """CAMPPlus speaker encoder.""" + + def __init__( + self, + feat_dim=80, + embedding_size=192, + growth_rate=32, + bn_size=4, + init_channels=128, + config_str="batchnorm-relu", + memory_efficient=True, + output_level="segment", + **kwargs, + ): + super().__init__() + + self.head = FCM(feat_dim=feat_dim) + channels = self.head.out_channels + self.output_level = output_level + + self.xvector = nn.ModuleDict() + self.xvector["tdnn"] = TDNNLayer( + channels, init_channels, 5, stride=2, dilation=1, padding=-1, config_str=config_str + ) + + channels = init_channels + for i, (num_layers, kernel_size, dilation) in enumerate(zip((12, 24, 16), (3, 3, 3), (1, 2, 2))): + block = CAMDenseTDNNBlock( + num_layers=num_layers, + in_channels=channels, + out_channels=growth_rate, + bn_channels=bn_size * growth_rate, + kernel_size=kernel_size, + dilation=dilation, + config_str=config_str, + memory_efficient=memory_efficient, + ) + self.xvector[f"block{i + 1}"] = block + channels = channels + num_layers * growth_rate + self.xvector[f"transit{i + 1}"] = TransitLayer(channels, channels // 2, bias=False, config_str=config_str) + channels //= 2 + + self.xvector["out_nonlinear"] = get_nonlinear(config_str, channels) + + if self.output_level == "segment": + self.xvector["stats"] = StatsPool() + self.xvector["dense"] = DenseLayer(channels * 2, embedding_size, config_str="batchnorm_") + else: + assert self.output_level == "frame", "`output_level` should be set to 'segment' or 'frame'." + + for m in self.modules(): + if isinstance(m, (nn.Conv1d, nn.Linear)): + nn.init.kaiming_normal_(m.weight.data) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x): + x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T) + x = self.head(x) + + # Manual forward through ModuleDict in correct order + for key in [ + "tdnn", + "block1", + "transit1", + "block2", + "transit2", + "block3", + "transit3", + "out_nonlinear", + "stats", + "dense", + ]: + if key in self.xvector: + module = self.xvector[key] + if isinstance(module, nn.ModuleDict): + for sublayer in module.values(): + x = sublayer(x) + else: + x = module(x) + + if self.output_level == "frame": + x = x.transpose(1, 2) + return x + + def inference(self, audio_list): + """Run inference on audio.""" + speech, speech_lengths = extract_feature(audio_list) + results = self.forward(speech.to(torch.float32)) + return results + + +# Conformer Encoder Components +class PositionalEncoding(nn.Module): + """Positional encoding.""" + + def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000, reverse: bool = False): + super().__init__() + self.d_model = d_model + self.xscale = math.sqrt(self.d_model) + self.dropout = nn.Dropout(p=dropout_rate) + self.max_len = max_len + # Lazily created on first forward to support meta-device initialization. + self.pe = None + + def _build_pe(self, device: torch.device): + pe = torch.zeros(self.max_len, self.d_model, device=device, dtype=torch.float32) + position = torch.arange(0, self.max_len, dtype=torch.float32, device=device).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, self.d_model, 2, dtype=torch.float32, device=device) * -(math.log(10000.0) / self.d_model) + ) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + # Keep in float32 (as originally) for numerical stability; do not follow `x.dtype` (often fp16/bf16). + self.pe = pe.unsqueeze(0) + + def forward(self, x: torch.Tensor, offset: int = 0) -> tuple[torch.Tensor, torch.Tensor]: + if self.pe is None or self.pe.is_meta or self.pe.device != x.device: + self._build_pe(device=x.device) + pos_emb = self.position_encoding(offset, x.size(1), False) + x = x * self.xscale + pos_emb + return self.dropout(x), self.dropout(pos_emb) + + def position_encoding(self, offset: int, size: int, apply_dropout: bool = True) -> torch.Tensor: + if isinstance(offset, int): + assert offset + size <= self.max_len + pos_emb = self.pe[:, offset : offset + size] + else: + pos_emb = self.pe[:, :size] + if apply_dropout: + pos_emb = self.dropout(pos_emb) + return pos_emb + + +class RelPositionalEncoding(PositionalEncoding): + """Relative positional encoding module.""" + + def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000): + super().__init__(d_model, dropout_rate, max_len, reverse=True) + + def forward(self, x: torch.Tensor, offset: int = 0) -> tuple[torch.Tensor, torch.Tensor]: + if self.pe is None or self.pe.is_meta or self.pe.device != x.device: + self._build_pe(device=x.device) + x = x * self.xscale + pos_emb = self.position_encoding(offset, x.size(1), False) + return self.dropout(x), self.dropout(pos_emb) + + +class EspnetRelPositionalEncoding(nn.Module): + """ESPnet-style relative positional encoding.""" + + def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000): + super().__init__() + self.d_model = d_model + self.xscale = math.sqrt(self.d_model) + self.dropout = nn.Dropout(p=dropout_rate) + self.pe = None + + def extend_pe(self, x: torch.Tensor): + if self.pe is not None: + if self.pe.is_meta: + self.pe = None + elif self.pe.size(1) >= x.size(1) * 2 - 1: + if self.pe.dtype != x.dtype or self.pe.device != x.device: + self.pe = self.pe.to(dtype=x.dtype, device=x.device) + return + pe_positive = torch.zeros(x.size(1), self.d_model, device=x.device, dtype=torch.float32) + pe_negative = torch.zeros(x.size(1), self.d_model, device=x.device, dtype=torch.float32) + position = torch.arange(0, x.size(1), dtype=torch.float32, device=x.device).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, self.d_model, 2, dtype=torch.float32, device=x.device) + * -(math.log(10000.0) / self.d_model) + ) + pe_positive[:, 0::2] = torch.sin(position * div_term) + pe_positive[:, 1::2] = torch.cos(position * div_term) + pe_negative[:, 0::2] = torch.sin(-1 * position * div_term) + pe_negative[:, 1::2] = torch.cos(-1 * position * div_term) + + pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0) + pe_negative = pe_negative[1:].unsqueeze(0) + pe = torch.cat([pe_positive, pe_negative], dim=1) + self.pe = pe.to(device=x.device, dtype=x.dtype) + + def forward(self, x: torch.Tensor, offset: int = 0) -> tuple[torch.Tensor, torch.Tensor]: + self.extend_pe(x) + x = x * self.xscale + pos_emb = self.position_encoding(size=x.size(1), offset=offset) + return self.dropout(x), self.dropout(pos_emb) + + def position_encoding(self, offset: int, size: int) -> torch.Tensor: + pos_emb = self.pe[:, self.pe.size(1) // 2 - size + 1 : self.pe.size(1) // 2 + size] + return pos_emb + + +class MultiHeadedAttention(nn.Module): + """Multi-Head Attention layer.""" + + def __init__(self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True): + super().__init__() + assert n_feat % n_head == 0 + self.d_k = n_feat // n_head + self.h = n_head + self.linear_q = nn.Linear(n_feat, n_feat) + self.linear_k = nn.Linear(n_feat, n_feat, bias=key_bias) + self.linear_v = nn.Linear(n_feat, n_feat) + self.linear_out = nn.Linear(n_feat, n_feat) + self.dropout = nn.Dropout(p=dropout_rate) + + def forward_qkv(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor): + n_batch = query.size(0) + q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) + k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) + v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + return q, k, v + + def forward_attention( + self, value: torch.Tensor, scores: torch.Tensor, mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool) + ) -> torch.Tensor: + n_batch = value.size(0) + if mask.size(2) > 0: + mask = mask.unsqueeze(1).eq(0) + mask = mask[:, :, :, : scores.size(-1)] + scores = scores.masked_fill(mask, -float("inf")) + attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0) + else: + attn = torch.softmax(scores, dim=-1) + + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) + x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) + return self.linear_out(x) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + pos_emb: torch.Tensor = torch.empty(0), + cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + ) -> tuple[torch.Tensor, torch.Tensor]: + q, k, v = self.forward_qkv(query, key, value) + if cache.size(0) > 0: + key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1) + k = torch.cat([key_cache, k], dim=2) + v = torch.cat([value_cache, v], dim=2) + new_cache = torch.cat((k, v), dim=-1) + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) + return self.forward_attention(v, scores, mask), new_cache + + +class RelPositionMultiHeadedAttention(MultiHeadedAttention): + """Multi-Head Attention layer with relative position encoding.""" + + def __init__(self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True): + super().__init__(n_head, n_feat, dropout_rate, key_bias) + self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) + self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k)) + self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k)) + nn.init.xavier_uniform_(self.pos_bias_u) + nn.init.xavier_uniform_(self.pos_bias_v) + + def rel_shift(self, x: torch.Tensor) -> torch.Tensor: + zero_pad = torch.zeros((x.size()[0], x.size()[1], x.size()[2], 1), device=x.device, dtype=x.dtype) + x_padded = torch.cat([zero_pad, x], dim=-1) + x_padded = x_padded.view(x.size()[0], x.size()[1], x.size(3) + 1, x.size(2)) + x = x_padded[:, :, 1:].view_as(x)[:, :, :, : x.size(-1) // 2 + 1] + return x + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + pos_emb: torch.Tensor = torch.empty(0), + cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + ) -> tuple[torch.Tensor, torch.Tensor]: + q, k, v = self.forward_qkv(query, key, value) + q = q.transpose(1, 2) + + if cache.size(0) > 0: + key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1) + k = torch.cat([key_cache, k], dim=2) + v = torch.cat([value_cache, v], dim=2) + new_cache = torch.cat((k, v), dim=-1) + + n_batch_pos = pos_emb.size(0) + p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) + p = p.transpose(1, 2) + + q_with_bias_u = (q + self.pos_bias_u.to(q.device)).transpose(1, 2) + q_with_bias_v = (q + self.pos_bias_v.to(q.device)).transpose(1, 2) + + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + if matrix_ac.shape != matrix_bd.shape: + matrix_bd = self.rel_shift(matrix_bd) + + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) + return self.forward_attention(v, scores, mask), new_cache + + +class PositionwiseFeedForward(nn.Module): + """Positionwise feed forward layer.""" + + def __init__(self, idim: int, hidden_units: int, dropout_rate: float, activation: nn.Module = nn.ReLU()): + super().__init__() + self.w_1 = nn.Linear(idim, hidden_units) + self.activation = activation + self.dropout = nn.Dropout(dropout_rate) + self.w_2 = nn.Linear(hidden_units, idim) + + def forward(self, xs: torch.Tensor) -> torch.Tensor: + return self.w_2(self.dropout(self.activation(self.w_1(xs)))) + + +class LinearNoSubsampling(nn.Module): + """Linear transform without subsampling.""" + + def __init__(self, idim: int, odim: int, dropout_rate: float, pos_enc_class: nn.Module): + super().__init__() + self.out = nn.ModuleList( + [ + nn.Linear(idim, odim), + nn.LayerNorm(odim, eps=1e-5), + nn.Dropout(dropout_rate), + ] + ) + self.pos_enc = pos_enc_class + self.right_context = 0 + self.subsampling_rate = 1 + + def forward(self, x: torch.Tensor, x_mask: torch.Tensor, offset: int = 0): + for layer in self.out: + x = layer(x) + x, pos_emb = self.pos_enc(x, offset) + return x, pos_emb, x_mask + + +class ConformerEncoderLayer(nn.Module): + """Encoder layer module.""" + + def __init__( + self, + size: int, + self_attn: nn.Module, + feed_forward: Optional[nn.Module] = None, + feed_forward_macaron: Optional[nn.Module] = None, + conv_module: Optional[nn.Module] = None, + dropout_rate: float = 0.1, + normalize_before: bool = True, + ): + super().__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.feed_forward_macaron = feed_forward_macaron + self.conv_module = conv_module + self.norm_ff = nn.LayerNorm(size, eps=1e-12) + self.norm_mha = nn.LayerNorm(size, eps=1e-12) + if feed_forward_macaron is not None: + self.norm_ff_macaron = nn.LayerNorm(size, eps=1e-12) + self.ff_scale = 0.5 + else: + self.ff_scale = 1.0 + if self.conv_module is not None: + self.norm_conv = nn.LayerNorm(size, eps=1e-12) + self.norm_final = nn.LayerNorm(size, eps=1e-12) + self.dropout = nn.Dropout(dropout_rate) + self.size = size + self.normalize_before = normalize_before + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor, + pos_emb: torch.Tensor, + mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + att_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + ): + if self.feed_forward_macaron is not None: + residual = x + if self.normalize_before: + x = self.norm_ff_macaron(x) + x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x)) + if not self.normalize_before: + x = self.norm_ff_macaron(x) + + residual = x + if self.normalize_before: + x = self.norm_mha(x) + x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb, att_cache) + x = residual + self.dropout(x_att) + if not self.normalize_before: + x = self.norm_mha(x) + + new_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device) + if self.conv_module is not None: + residual = x + if self.normalize_before: + x = self.norm_conv(x) + x, new_cnn_cache = self.conv_module(x, mask_pad, cnn_cache) + x = residual + self.dropout(x) + if not self.normalize_before: + x = self.norm_conv(x) + + residual = x + if self.normalize_before: + x = self.norm_ff(x) + x = residual + self.ff_scale * self.dropout(self.feed_forward(x)) + if not self.normalize_before: + x = self.norm_ff(x) + + if self.conv_module is not None: + x = self.norm_final(x) + + return x, mask, new_att_cache, new_cnn_cache + + +class Upsample1D(nn.Module): + """1D upsampling layer.""" + + def __init__(self, channels: int, out_channels: int, stride: int = 2): + super().__init__() + self.channels = channels + self.out_channels = out_channels + self.stride = stride + self.conv = nn.Conv1d(self.channels, self.out_channels, stride * 2 + 1, stride=1, padding=0) + + def forward(self, inputs: torch.Tensor, input_lengths: torch.Tensor): + outputs = F.interpolate(inputs, scale_factor=float(self.stride), mode="nearest") + outputs = F.pad(outputs, (self.stride * 2, 0), value=0.0) + outputs = self.conv(outputs) + return outputs, input_lengths * self.stride + + +class PreLookaheadLayer(nn.Module): + """Pre-lookahead layer for streaming.""" + + def __init__(self, channels: int, pre_lookahead_len: int = 1): + super().__init__() + self.channels = channels + self.pre_lookahead_len = pre_lookahead_len + self.conv1 = nn.Conv1d(channels, channels, kernel_size=pre_lookahead_len + 1, stride=1, padding=0) + self.conv2 = nn.Conv1d(channels, channels, kernel_size=3, stride=1, padding=0) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + outputs = inputs.transpose(1, 2).contiguous() + outputs = F.pad(outputs, (0, self.pre_lookahead_len), mode="constant", value=0.0) + outputs = F.leaky_relu(self.conv1(outputs)) + outputs = F.pad(outputs, (2, 0), mode="constant", value=0.0) + outputs = self.conv2(outputs) + outputs = outputs.transpose(1, 2).contiguous() + outputs = outputs + inputs + return outputs + + +class UpsampleConformerEncoder(nn.Module): + """Upsample Conformer Encoder.""" + + def __init__( + self, + input_size: int = 512, + output_size: int = 512, + attention_heads: int = 8, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.1, + positional_dropout_rate: float = 0.1, + attention_dropout_rate: float = 0.1, + input_layer: str = "linear", + pos_enc_layer_type: str = "rel_pos_espnet", + normalize_before: bool = True, + static_chunk_size: int = 0, + use_dynamic_chunk: bool = False, + global_cmvn: nn.Module = None, + use_dynamic_left_chunk: bool = False, + positionwise_conv_kernel_size: int = 1, + macaron_style: bool = False, + selfattention_layer_type: str = "rel_selfattn", + activation_type: str = "swish", + use_cnn_module: bool = False, + cnn_module_kernel: int = 15, + causal: bool = False, + cnn_module_norm: str = "batch_norm", + key_bias: bool = True, + gradient_checkpointing: bool = False, + ): + super().__init__() + self._output_size = output_size + self.global_cmvn = global_cmvn + + if pos_enc_layer_type == "rel_pos_espnet": + pos_enc_class = EspnetRelPositionalEncoding(output_size, positional_dropout_rate) + else: + pos_enc_class = RelPositionalEncoding(output_size, positional_dropout_rate) + + self.embed = LinearNoSubsampling(input_size, output_size, dropout_rate, pos_enc_class) + self.normalize_before = normalize_before + self.after_norm = nn.LayerNorm(output_size, eps=1e-5) + self.static_chunk_size = static_chunk_size + self.use_dynamic_chunk = use_dynamic_chunk + self.use_dynamic_left_chunk = use_dynamic_left_chunk + self.gradient_checkpointing = gradient_checkpointing + + activation = nn.SiLU() if activation_type == "swish" else nn.ReLU() + + encoder_selfattn_layer_args = (attention_heads, output_size, attention_dropout_rate, key_bias) + positionwise_layer_args = (output_size, linear_units, dropout_rate, activation) + + self.pre_lookahead_layer = PreLookaheadLayer(channels=512, pre_lookahead_len=3) + self.encoders = nn.ModuleList( + [ + ConformerEncoderLayer( + output_size, + RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args) + if selfattention_layer_type == "rel_selfattn" + else MultiHeadedAttention(*encoder_selfattn_layer_args), + PositionwiseFeedForward(*positionwise_layer_args), + PositionwiseFeedForward(*positionwise_layer_args) if macaron_style else None, + None, # No CNN module + dropout_rate, + normalize_before, + ) + for _ in range(num_blocks) + ] + ) + self.up_layer = Upsample1D(channels=512, out_channels=512, stride=2) + + if pos_enc_layer_type == "rel_pos_espnet": + up_pos_enc_class = EspnetRelPositionalEncoding(output_size, positional_dropout_rate) + else: + up_pos_enc_class = RelPositionalEncoding(output_size, positional_dropout_rate) + + self.up_embed = LinearNoSubsampling(input_size, output_size, dropout_rate, up_pos_enc_class) + self.up_encoders = nn.ModuleList( + [ + ConformerEncoderLayer( + output_size, + RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args) + if selfattention_layer_type == "rel_selfattn" + else MultiHeadedAttention(*encoder_selfattn_layer_args), + PositionwiseFeedForward(*positionwise_layer_args), + PositionwiseFeedForward(*positionwise_layer_args) if macaron_style else None, + None, + dropout_rate, + normalize_before, + ) + for _ in range(4) + ] + ) + + def output_size(self) -> int: + return self._output_size + + def forward( + self, xs: torch.Tensor, xs_lens: torch.Tensor, decoding_chunk_size: int = 0, num_decoding_left_chunks: int = -1 + ) -> tuple[torch.Tensor, torch.Tensor]: + T = xs.size(1) + masks = ~make_pad_mask(xs_lens, T).unsqueeze(1) + if self.global_cmvn is not None: + xs = self.global_cmvn(xs) + xs, pos_emb, masks = self.embed(xs, masks) + mask_pad = masks + chunk_masks = masks + + xs = self.pre_lookahead_layer(xs) + xs = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad) + + xs = xs.transpose(1, 2).contiguous() + xs, xs_lens = self.up_layer(xs, xs_lens) + xs = xs.transpose(1, 2).contiguous() + T = xs.size(1) + masks = ~make_pad_mask(xs_lens, T).unsqueeze(1) + xs, pos_emb, masks = self.up_embed(xs, masks) + mask_pad = masks + chunk_masks = masks + xs = self.forward_up_layers(xs, chunk_masks, pos_emb, mask_pad) + + if self.normalize_before: + xs = self.after_norm(xs) + return xs, masks + + def forward_layers( + self, xs: torch.Tensor, chunk_masks: torch.Tensor, pos_emb: torch.Tensor, mask_pad: torch.Tensor + ) -> torch.Tensor: + for layer in self.encoders: + xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad) + return xs + + def forward_up_layers( + self, xs: torch.Tensor, chunk_masks: torch.Tensor, pos_emb: torch.Tensor, mask_pad: torch.Tensor + ) -> torch.Tensor: + for layer in self.up_encoders: + xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad) + return xs + + +# CFM Decoder Components +class SinusoidalPosEmb(nn.Module): + """Sinusoidal positional embedding.""" + + def __init__(self, dim): + super().__init__() + self.dim = dim + assert self.dim % 2 == 0, "SinusoidalPosEmb requires dim to be even" + + def forward(self, x, scale=1000): + if x.ndim < 1: + x = x.unsqueeze(0) + device = x.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 * x.unsqueeze(1) * emb.unsqueeze(0) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +class TimestepEmbedding(nn.Module): + """Timestep embedding layer.""" + + def __init__(self, in_channels: int, time_embed_dim: int, act_fn: str = "silu"): + super().__init__() + self.linear_1 = nn.Linear(in_channels, time_embed_dim) + self.act = nn.SiLU() if act_fn == "silu" else nn.ReLU() + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim) + + def forward(self, sample): + sample = self.linear_1(sample) + if self.act is not None: + sample = self.act(sample) + sample = self.linear_2(sample) + return sample + + +class CausalConv1d(nn.Conv1d): + """Causal 1D convolution.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + device=None, + dtype=None, + ) -> None: + super().__init__( + in_channels, + out_channels, + kernel_size, + stride, + padding=0, + dilation=dilation, + groups=groups, + bias=bias, + padding_mode=padding_mode, + device=device, + dtype=dtype, + ) + assert stride == 1 + self.causal_padding = (kernel_size - 1, 0) + + def forward(self, x: torch.Tensor): + x = F.pad(x, self.causal_padding) + x = super().forward(x) + return x + + +class Transpose(nn.Module): + """Transpose module.""" + + def __init__(self, dim0: int, dim1: int): + super().__init__() + self.dim0 = dim0 + self.dim1 = dim1 + + def forward(self, x: torch.Tensor): + return torch.transpose(x, self.dim0, self.dim1) + + +class CausalBlock1D(nn.Module): + """Causal 1D block.""" + + def __init__(self, dim: int, dim_out: int): + super().__init__() + self.block = nn.ModuleList( + [ + CausalConv1d(dim, dim_out, 3), + Transpose(1, 2), + nn.LayerNorm(dim_out), + Transpose(1, 2), + nn.Mish(), + ] + ) + + def forward(self, x: torch.Tensor, mask: torch.Tensor): + x = x * mask + for layer in self.block: + x = layer(x) + return x * mask + + +class CausalResnetBlock1D(nn.Module): + """Causal ResNet block.""" + + def __init__(self, dim: int, dim_out: int, time_emb_dim: int, groups: int = 8): + super().__init__() + self.mlp = nn.ModuleList([nn.Mish(), nn.Linear(time_emb_dim, dim_out)]) + self.block1 = CausalBlock1D(dim, dim_out) + self.block2 = CausalBlock1D(dim_out, dim_out) + self.res_conv = CausalConv1d(dim, dim_out, 1) + + def forward(self, x, mask, time_emb): + h = self.block1(x, mask) + + mlp_out = time_emb + for layer in self.mlp: + mlp_out = layer(mlp_out) + h += mlp_out.unsqueeze(-1) + + h = self.block2(h, mask) + output = h + self.res_conv(x * mask) + return output + + +# ConditionalDecoder implementation for S3Gen + + +class ConditionalDecoder(nn.Module): + """Conditional decoder for CFM. Simplified U-Net architecture.""" + + def __init__( + self, + in_channels=320, + out_channels=80, + causal=True, + channels=[256], + dropout=0.0, + attention_head_dim=64, + n_blocks=4, + num_mid_blocks=12, + num_heads=8, + act_fn="gelu", + ): + super().__init__() + channels = tuple(channels) + self.in_channels = in_channels + self.out_channels = out_channels + self.causal = causal + self.time_embeddings = SinusoidalPosEmb(in_channels) + time_embed_dim = channels[0] * 4 + self.time_mlp = TimestepEmbedding(in_channels=in_channels, time_embed_dim=time_embed_dim, act_fn="silu") + self.down_blocks = nn.ModuleList([]) + self.mid_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + self.static_chunk_size = 0 + + output_channel = in_channels + for i in range(len(channels)): + input_channel = output_channel + output_channel = channels[i] + is_last = i == len(channels) - 1 + resnet = CausalResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) + + # Create transformer blocks + transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=output_channel, + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + attention_bias=False, + only_cross_attention=False, + upcast_attention=False, + ) + for _ in range(n_blocks) + ] + ) + + downsample = ( + CausalConv1d(output_channel, output_channel, 3) + if is_last + else nn.Conv1d(output_channel, output_channel // 2, 3, stride=2, padding=1) + ) + self.down_blocks.append(nn.ModuleList([resnet, transformer_blocks, downsample])) + + for _ in range(num_mid_blocks): + input_channel = channels[-1] + resnet = CausalResnetBlock1D(dim=input_channel, dim_out=channels[-1], time_emb_dim=time_embed_dim) + + # Create transformer blocks + transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=channels[-1], + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + attention_bias=False, + only_cross_attention=False, + upcast_attention=False, + ) + for _ in range(n_blocks) + ] + ) + + self.mid_blocks.append(nn.ModuleList([resnet, transformer_blocks])) + + channels = channels[::-1] + (channels[0],) + for i in range(len(channels) - 1): + input_channel = channels[i] * 2 + output_channel = channels[i + 1] + is_last = i == len(channels) - 2 + resnet = CausalResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) + + # Create transformer blocks + transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=output_channel, + num_attention_heads=num_heads, + attention_head_dim=attention_head_dim, + dropout=dropout, + activation_fn=act_fn, + attention_bias=False, + only_cross_attention=False, + upcast_attention=False, + ) + for _ in range(n_blocks) + ] + ) + + upsample = ( + CausalConv1d(output_channel, output_channel, 3) + if is_last + else nn.ConvTranspose1d(output_channel, output_channel, 4, 2, 1) + ) + self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample])) + + self.final_block = CausalBlock1D(channels[-1], channels[-1]) + self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1) + self.initialize_weights() + + def initialize_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv1d): + nn.init.kaiming_normal_(m.weight, nonlinearity="relu") + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Linear): + nn.init.kaiming_normal_(m.weight, nonlinearity="relu") + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x, mask, mu, t, spks=None, cond=None): + t = self.time_embeddings(t).to(t.dtype) + t = self.time_mlp(t) + + # Concatenate inputs + x = torch.cat([x, mu], dim=1) + if spks is not None: + spks = spks.unsqueeze(-1).expand(-1, -1, x.shape[-1]) + x = torch.cat([x, spks], dim=1) + if cond is not None: + x = torch.cat([x, cond], dim=1) + + hiddens = [] + masks = [mask] + for resnet, transformer_blocks, downsample in self.down_blocks: + mask_down = masks[-1] + x = resnet(x, mask_down, t) + # Transpose for transformer blocks: (B, C, T) -> (B, T, C) + for transformer_block in transformer_blocks: + x = transformer_block(x.transpose(1, 2)).transpose(1, 2) + hiddens.append(x) + x = downsample(x * mask_down) + masks.append(mask_down[:, :, ::2] if x.shape[-1] < mask_down.shape[-1] else mask_down) + + masks = masks[:-1] + mask_mid = masks[-1] + + for resnet, transformer_blocks in self.mid_blocks: + x = resnet(x, mask_mid, t) + # Transpose for transformer blocks: (B, C, T) -> (B, T, C) + for transformer_block in transformer_blocks: + x = transformer_block(x.transpose(1, 2)).transpose(1, 2) + + for resnet, transformer_blocks, upsample in self.up_blocks: + mask_up = masks.pop() + skip = hiddens.pop() + x = torch.cat([x[:, :, : skip.shape[-1]], skip], dim=1) + x = resnet(x, mask_up, t) + # Transpose for transformer blocks: (B, C, T) -> (B, T, C) + for transformer_block in transformer_blocks: + x = transformer_block(x.transpose(1, 2)).transpose(1, 2) + x = upsample(x * mask_up) + + x = self.final_block(x, mask_up) + output = self.final_proj(x * mask_up) + return output * mask + + +class CausalConditionalCFM(nn.Module): + """Causal Conditional Flow Matching.""" + + def __init__(self, in_channels=240, spk_emb_dim=80, estimator=None): + super().__init__() + self.n_feats = in_channels + self.spk_emb_dim = spk_emb_dim + self.solver = "euler" + self.sigma_min = 1e-6 + self.t_scheduler = "cosine" + self.training_cfg_rate = 0.2 + self.inference_cfg_rate = 0.7 + self.estimator = estimator + # Lazily materialized on first forward to support meta-device initialization. + self.rand_noise = None + + @torch.inference_mode() + def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None): + needed_len = mu.size(2) + if ( + self.rand_noise is None + or self.rand_noise.is_meta + or self.rand_noise.device != mu.device + or self.rand_noise.dtype != mu.dtype + or self.rand_noise.size(2) < needed_len + ): + # Keep a small cache so repeated calls don't reallocate for slightly different lengths. + cache_len = max(needed_len, 50 * 300) + self.rand_noise = torch.randn(1, 80, cache_len, device=mu.device, dtype=mu.dtype) + z = self.rand_noise[:, :, :needed_len] * temperature + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None + + def solve_euler(self, x, t_span, mu, mask, spks, cond): + t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0] + x_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=x.dtype) + mask_in = torch.zeros([2, 1, x.size(2)], device=x.device, dtype=x.dtype) + mu_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=x.dtype) + t_in = torch.zeros([2], device=x.device, dtype=x.dtype) + spks_in = torch.zeros([2, 80], device=x.device, dtype=x.dtype) + cond_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=x.dtype) + + for step in range(1, len(t_span)): + x_in[:] = x + mask_in[:] = mask + mu_in[0] = mu + t_in[:] = t.unsqueeze(0) + spks_in[0] = spks + cond_in[0] = cond + dphi_dt = self.estimator(x_in, mask_in, mu_in, t_in, spks_in, cond_in) + dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [x.size(0), x.size(0)], dim=0) + dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt + x = x + dt * dphi_dt + t = t + dt + if step < len(t_span) - 1: + dt = t_span[step + 1] - t + return x.float() + + +class CausalMaskedDiffWithXvec(nn.Module): + """Causal masked diffusion with speaker embedding.""" + + def __init__( + self, + input_size: int = 512, + output_size: int = 80, + spk_embed_dim: int = 192, + vocab_size: int = 6561, + input_frame_rate: int = 25, + token_mel_ratio: int = 2, + pre_lookahead_len: int = 3, + encoder: nn.Module = None, + decoder: nn.Module = None, + ): + super().__init__() + self.input_size = input_size + self.output_size = output_size + self.vocab_size = vocab_size + self.input_frame_rate = input_frame_rate + self.input_embedding = nn.Embedding(vocab_size, input_size) + self.spk_embed_affine_layer = nn.Linear(spk_embed_dim, output_size) + self.encoder = encoder + self.encoder_proj = nn.Linear(self.encoder.output_size(), output_size) + self.decoder = decoder + self.token_mel_ratio = token_mel_ratio + self.pre_lookahead_len = pre_lookahead_len + self.fp16 = False + + @torch.inference_mode() + def inference( + self, token, token_len, prompt_token, prompt_token_len, prompt_feat, prompt_feat_len, embedding, finalize + ): + if self.fp16: + prompt_feat = prompt_feat.half() + embedding = embedding.half() + + assert token.shape[0] == 1 + embedding = F.normalize(embedding, dim=1) + embedding = self.spk_embed_affine_layer(embedding) + + token, token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + token_len + mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding) + token = self.input_embedding(torch.clamp(token, min=0, max=self.input_embedding.num_embeddings - 1)) * mask + + h, h_lengths = self.encoder(token, token_len) + if not finalize: + h = h[:, : -self.pre_lookahead_len * self.token_mel_ratio] + mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1] + h = self.encoder_proj(h) + + conds = torch.zeros([1, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype) + conds[:, :mel_len1] = prompt_feat + conds = conds.transpose(1, 2) + + mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h) + feat, _ = self.decoder( + mu=h.transpose(1, 2).contiguous(), mask=mask.unsqueeze(1), spks=embedding, cond=conds, n_timesteps=10 + ) + feat = feat[:, :, mel_len1:] + assert feat.shape[2] == mel_len2 + return feat.float(), None + + +# Main Model +class S3GenPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = S3GenConfig + base_model_prefix = "s3gen" + main_input_name = "speech_tokens" + supports_gradient_checkpointing = False + + +S3GEN_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`S3GenConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +S3GEN_INPUTS_DOCSTRING = r""" + Args: + speech_tokens (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of speech tokens from S3 tokenizer. + ref_wav (`torch.FloatTensor` of shape `(batch_size, audio_length)`, *optional*): + Reference audio waveform for speaker conditioning. + ref_sr (`int`, *optional*): + Sample rate of the reference audio. + ref_dict (`dict`, *optional*): + Pre-computed reference embeddings dict (alternative to ref_wav). + finalize (`bool`, *optional*, defaults to `True`): + Whether this is the final chunk (for streaming). +""" + + +@auto_docstring +class S3GenModel(S3GenPreTrainedModel): + """ + The S3Gen Model for converting speech tokens to mel spectrograms and waveforms. + """ + + def __init__(self, config: S3GenConfig): + super().__init__(config) + self.config = config + + # S3 Tokenizer for reference audio (initialized locally, weights loaded from checkpoint) + tokenizer_config = S3TokenizerConfig() + self.tokenizer = S3TokenizerModel(tokenizer_config, name="speech_tokenizer_v2_25hz") + self.tokenizer_feature_extractor = S3TokenizerFeatureExtractor() + + # Speaker encoder + self.speaker_encoder = CAMPPlus( + feat_dim=config.speaker_feat_dim, + embedding_size=config.speaker_embed_dim, + ) + + # Conformer encoder + encoder = UpsampleConformerEncoder( + output_size=config.encoder_output_size, + attention_heads=config.encoder_attention_heads, + linear_units=config.encoder_linear_units, + num_blocks=config.encoder_num_blocks, + dropout_rate=config.encoder_dropout_rate, + positional_dropout_rate=config.encoder_dropout_rate, + attention_dropout_rate=config.encoder_dropout_rate, + normalize_before=True, + input_layer="linear", + pos_enc_layer_type="rel_pos_espnet", + selfattention_layer_type="rel_selfattn", + input_size=config.token_embed_dim, + use_cnn_module=False, + macaron_style=False, + ) + + # CFM decoder + estimator = ConditionalDecoder( + in_channels=config.decoder_in_channels, + out_channels=config.decoder_out_channels, + causal=True, + channels=config.decoder_channels, + dropout=0.0, + attention_head_dim=config.decoder_attention_head_dim, + n_blocks=config.decoder_n_blocks, + num_mid_blocks=config.decoder_num_mid_blocks, + num_heads=config.decoder_num_heads, + act_fn=config.decoder_act_fn, + ) + decoder = CausalConditionalCFM( + in_channels=config.decoder_in_channels, + spk_emb_dim=config.decoder_out_channels, + estimator=estimator, + ) + + self.flow = CausalMaskedDiffWithXvec( + encoder=encoder, + decoder=decoder, + input_size=config.token_embed_dim, + output_size=config.mel_bins, + spk_embed_dim=config.speaker_embed_dim, + vocab_size=config.vocab_size, + input_frame_rate=config.input_frame_rate, + token_mel_ratio=config.token_mel_ratio, + pre_lookahead_len=config.pre_lookahead_len, + ) + + # CFM parameters stored in config for future use + _ = (config.cfm_sigma_min, config.cfm_solver, config.cfm_t_scheduler, config.cfm_inference_cfg_rate) + + # HiFTNet vocoder + hiftnet_config = HiFTNetConfig( + sampling_rate=config.sampling_rate, + upsample_rates=[8, 5, 3], + upsample_kernel_sizes=[16, 11, 7], + source_resblock_kernel_sizes=[7, 7, 11], + source_resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]], + ) + self.mel2wav = HiFTGenerator(hiftnet_config) + + # Trim/fade buffer for reducing startup artifacts (glitches/clicks) from the vocoder. + # Use a short fade-in (no trimming) to avoid reintroducing discontinuities. + n_trim = 0 # ~10ms at 24kHz + # Smooth fade-in from 0 -> 1. + trim_fade = (torch.cos(torch.linspace(torch.pi, 0, n_trim)) + 1) / 2 + self.register_buffer("trim_fade", trim_fade, persistent=False) + + self.post_init() + + @property + def device(self): + return next(self.parameters()).device + + def embed_ref(self, ref_wav: torch.Tensor, ref_sr: int, device="auto"): + """Extract reference embeddings from audio.""" + device = self.device if device == "auto" else device + if isinstance(ref_wav, np.ndarray): + ref_wav = torch.from_numpy(ref_wav).float() + + if ref_wav.device != device: + ref_wav = ref_wav.to(device) + + if len(ref_wav.shape) == 1: + ref_wav = ref_wav.unsqueeze(0) + + # Resample to 24kHz for mel extraction + ref_wav_24 = ref_wav + if ref_sr != self.config.sampling_rate: + import torchaudio + + resampler = torchaudio.transforms.Resample(ref_sr, self.config.sampling_rate).to(device) + ref_wav_24 = resampler(ref_wav) + + ref_mels_24 = ( + mel_spectrogram( + ref_wav_24, + n_fft=self.config.n_fft, + num_mels=self.config.mel_bins, + sampling_rate=self.config.sampling_rate, + hop_size=self.config.hop_length, + win_size=self.config.win_size, + fmin=self.config.fmin, + fmax=self.config.fmax, + ) + .transpose(1, 2) + .to(device) + ) + + # Resample to 16kHz for speaker encoder + tokenizer + import torchaudio + + resampler_16 = torchaudio.transforms.Resample(ref_sr, 16000).to(device) + ref_wav_16 = resampler_16(ref_wav).to(device) + + # Speaker embedding + ref_x_vector = self.speaker_encoder.inference(ref_wav_16) + + # Tokenize reference (use feature extractor first) + features = self.tokenizer_feature_extractor( + ref_wav_16.cpu().numpy(), sampling_rate=16000, return_tensors="pt" + ).to(device) + ref_speech_tokens, ref_speech_token_lens = self.tokenizer( + input_features=features.input_features, attention_mask=features.attention_mask, return_dict=False + ) + + # Ensure mel_len = 2 * token_len + if ref_mels_24.shape[1] != 2 * ref_speech_tokens.shape[1]: + ref_speech_tokens = ref_speech_tokens[:, : ref_mels_24.shape[1] // 2] + ref_speech_token_lens[0] = ref_speech_tokens.shape[1] + + return { + "prompt_token": ref_speech_tokens.to(device), + "prompt_token_len": ref_speech_token_lens, + "prompt_feat": ref_mels_24, + "prompt_feat_len": None, + "embedding": ref_x_vector, + } + + @add_start_docstrings_to_model_forward(S3GEN_INPUTS_DOCSTRING) + def forward(self, speech_tokens, ref_wav=None, ref_sr=None, ref_dict=None, finalize=False, **kwargs): + """Generate mel spectrograms from tokens.""" + assert (ref_wav is None) ^ (ref_dict is None), "Must provide exactly one of ref_wav or ref_dict" + + if ref_dict is None: + ref_dict = self.embed_ref(ref_wav, ref_sr) + else: + for rk in list(ref_dict): + if isinstance(ref_dict[rk], np.ndarray): + ref_dict[rk] = torch.from_numpy(ref_dict[rk]) + if torch.is_tensor(ref_dict[rk]): + ref_dict[rk] = ref_dict[rk].to(self.device) + + if len(speech_tokens.shape) == 1: + speech_tokens = speech_tokens.unsqueeze(0) + + speech_token_lens = torch.LongTensor([speech_tokens.size(1)]).to(self.device) + + output_mels, _ = self.flow.inference( + token=speech_tokens, + token_len=speech_token_lens, + finalize=finalize, + **ref_dict, + ) + return output_mels + + @torch.inference_mode() + def inference(self, speech_tokens, ref_wav=None, ref_sr=None, ref_dict=None, cache_source=None, finalize=True): + """ + End-to-end inference: tokens → waveform. + + Args: + speech_tokens: Speech token sequence + ref_wav: Reference audio waveform (mutex with ref_dict) + ref_sr: Reference audio sample rate (required with ref_wav) + ref_dict: Pre-computed reference embeddings (mutex with ref_wav) + cache_source: Cached source for streaming + finalize: Whether to finalize generation + """ + output_mels = self.forward(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize) + + if cache_source is None: + cache_source = torch.zeros(1, 1, 0).to(self.device) + + output_wavs, output_sources = self.mel2wav.inference(speech_feat=output_mels, cache_source=cache_source) + + # Reduce spillover artifacts at the start (non-inplace to avoid InferenceMode error) + trim_fade = self.trim_fade.to(output_wavs.device) + output_wavs = output_wavs.clone() # Clone to allow inplace operation + n_fade = len(trim_fade) + if output_wavs.size(1) > n_fade: + output_wavs[:, :n_fade] *= trim_fade + + return output_wavs, output_sources + + @torch.inference_mode() + def generate(self, speech_tokens, ref_wav, ref_sr, cache_source=None, finalize=True): + """ + Generate audio from speech tokens. + + This is an alias for the inference method, provided for consistency with + HuggingFace generation API conventions. + + Args: + speech_tokens (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Speech tokens to convert to audio. + ref_wav (`torch.FloatTensor` of shape `(batch_size, audio_length)` or `(audio_length,)`): + Reference audio for speaker embedding extraction. + ref_sr (`int`): + Sample rate of the reference audio. + cache_source (`torch.FloatTensor`, *optional*): + Cached source for streaming generation. Defaults to None. + finalize (`bool`, *optional*, defaults to `True`): + Whether to finalize the generation (used for streaming). + + Returns: + `torch.FloatTensor`: Generated waveform of shape `(batch_size, audio_length)`. + """ + output_wavs, _ = self.inference( + speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, cache_source=cache_source, finalize=finalize + ) + return output_wavs + + +__all__ = ["S3GenPreTrainedModel", "S3GenModel"] diff --git a/src/transformers/models/s3tokenizer/__init__.py b/src/transformers/models/s3tokenizer/__init__.py new file mode 100644 index 000000000000..fe6a6b3bd267 --- /dev/null +++ b/src/transformers/models/s3tokenizer/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2024 The HuggingFace 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 typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_s3tokenizer import * + from .feature_extraction_s3tokenizer import * + from .modeling_s3tokenizer import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/src/transformers/models/s3tokenizer/configuration_s3tokenizer.py b/src/transformers/models/s3tokenizer/configuration_s3tokenizer.py new file mode 100644 index 000000000000..0b33e586c65c --- /dev/null +++ b/src/transformers/models/s3tokenizer/configuration_s3tokenizer.py @@ -0,0 +1,118 @@ +# coding=utf-8 +# Copyright 2024 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. + +"""S3Tokenizer model configuration""" + +from ...configuration_utils import PreTrainedConfig +from ...utils import logging + + +logger = logging.get_logger(__name__) + + +class S3TokenizerConfig(PreTrainedConfig): + r""" + This is the configuration class to store the configuration of a [`S3TokenizerModel`]. It is used to instantiate a + S3Tokenizer model according to the specified arguments, defining the model architecture. Instantiating a configuration + with the defaults will yield a similar configuration to that of the S3Tokenizer + [ResembleAI/s3tokenizer-v2](https://huggingface.co/ResembleAI/s3tokenizer-v2) architecture. + + Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PreTrainedConfig`] for more information. + + Args: + sampling_rate (`int`, *optional*, defaults to 16000): + The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz). + n_mels (`int`, *optional*, defaults to 128): + Number of mel-frequency bins for the mel-spectrogram. + n_fft (`int`, *optional*, defaults to 400): + Size of the FFT window for computing the mel-spectrogram. + vocab_size (`int`, *optional*, defaults to 6561): + Vocabulary size of the S3 tokenizer (3^8 for FSQ quantization). + n_audio_state (`int`, *optional*, defaults to 1280): + Hidden state dimension of the audio encoder. + n_audio_head (`int`, *optional*, defaults to 20): + Number of attention heads in the audio encoder. + n_audio_layer (`int`, *optional*, defaults to 6): + Number of transformer layers in the audio encoder. + use_sdpa (`bool`, *optional*, defaults to `False`): + Whether to use Scaled Dot Product Attention (SDPA) for faster inference. + num_attention_heads (``, *optional*): + num_key_value_heads (``, *optional*): + attention_bias (``, *optional*, defaults to `False`): + attention_dropout (``, *optional*, defaults to 0.0): + max_position_embeddings (``, *optional*, defaults to 2048): + rope_theta (``, *optional*, defaults to 10000.0): + rope_scaling (``, *optional*): + + Example: + + ```python + >>> from transformers import S3TokenizerModel, S3TokenizerConfig + + >>> # Initializing a S3Tokenizer configuration + >>> configuration = S3TokenizerConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = S3TokenizerModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "s3tokenizer" + + def __init__( + self, + sampling_rate=16000, + n_mels=128, + n_fft=400, + vocab_size=6561, + n_audio_state=1280, + n_audio_head=20, + n_audio_layer=6, + use_sdpa=False, + num_attention_heads=None, + num_key_value_heads=None, + attention_bias=False, + attention_dropout=0.0, + max_position_embeddings=2048, + rope_theta=10000.0, + rope_scaling=None, + **kwargs, + ): + self.sampling_rate = sampling_rate + self.n_mels = n_mels + self.n_fft = n_fft + self.vocab_size = vocab_size + self.n_audio_state = n_audio_state + self.n_audio_head = n_audio_head + self.n_audio_layer = n_audio_layer + self.use_sdpa = use_sdpa + # Add hidden_size as an alias for n_audio_state for compatibility with common tests + self.hidden_size = n_audio_state + self.num_attention_heads = num_attention_heads if num_attention_heads is not None else n_audio_head + self.num_key_value_heads = num_key_value_heads if num_key_value_heads is not None else self.num_attention_heads + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.max_position_embeddings = max_position_embeddings + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + + super().__init__(**kwargs) + self.convert_rope_params_to_dict() + + +__all__ = ["S3TokenizerConfig"] diff --git a/src/transformers/models/s3tokenizer/feature_extraction_s3tokenizer.py b/src/transformers/models/s3tokenizer/feature_extraction_s3tokenizer.py new file mode 100644 index 000000000000..9e313e96ae80 --- /dev/null +++ b/src/transformers/models/s3tokenizer/feature_extraction_s3tokenizer.py @@ -0,0 +1,201 @@ +# coding=utf-8 +# Copyright 2024 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. +"""Feature extractor class for S3Tokenizer.""" + +from typing import Optional, Union + +import numpy as np +import torch + +from ...feature_extraction_sequence_utils import SequenceFeatureExtractor +from ...feature_extraction_utils import BatchFeature +from ...utils import PaddingStrategy, TensorType, is_librosa_available, logging +from ...utils.import_utils import requires + + +logger = logging.get_logger(__name__) + + +if is_librosa_available(): + import librosa + + +@requires(backends=("torch",)) +class S3TokenizerFeatureExtractor(SequenceFeatureExtractor): + r""" + Constructs a S3Tokenizer feature extractor. + + This feature extractor inherits from [`~feature_extraction_sequence_utils.SequenceFeatureExtractor`] which contains + most of the main methods. Users should refer to this superclass for more information regarding those methods. + + Args: + feature_size (`int`, *optional*, defaults to 1): + The feature dimension of the extracted features. Use 1 for mono audio. + sampling_rate (`int`, *optional*, defaults to 16000): + The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz). + padding_value (`float`, *optional*, defaults to 0.0): + The value that is used to fill the padding values. + n_mels (`int`, *optional*, defaults to 128): + Number of mel-frequency bins for the mel-spectrogram. + n_fft (`int`, *optional*, defaults to 400): + Size of the FFT window for computing the mel-spectrogram. + hop_length (`int`, *optional*, defaults to 160): + Number of audio samples between adjacent STFT columns (10ms at 16kHz). + """ + + model_input_names = ["input_features", "attention_mask"] + + def __init__( + self, + feature_size: int = 1, + sampling_rate: int = 16000, + padding_value: float = 0.0, + n_mels: int = 128, + n_fft: int = 400, + hop_length: int = 160, + **kwargs, + ): + super().__init__( + feature_size=feature_size, + sampling_rate=sampling_rate, + padding_value=padding_value, + **kwargs, + ) + self.n_mels = n_mels + self.n_fft = n_fft + self.hop_length = hop_length + self._mel_filters = None + self._mel_filters_torch = None + self._window_torch = None + + def _get_mel_filters(self): + if self._mel_filters is None: + if not is_librosa_available(): + raise ImportError( + "librosa is required to compute mel filters in S3TokenizerFeatureExtractor. " + "Please install it with `pip install librosa`." + ) + self._mel_filters = librosa.filters.mel(sr=self.sampling_rate, n_fft=self.n_fft, n_mels=self.n_mels) + return self._mel_filters + + def _get_mel_filters_torch(self) -> torch.Tensor: + """ + Cached torch.Tensor version of mel filters for torch STFT path. + """ + if self._mel_filters_torch is None: + self._mel_filters_torch = torch.tensor(self._get_mel_filters(), dtype=torch.float32) + return self._mel_filters_torch + + def _get_window_torch(self) -> torch.Tensor: + """ + Cached torch.Tensor Hann window. Matches the model's `torch.hann_window` usage. + """ + if self._window_torch is None: + self._window_torch = torch.hann_window(self.n_fft, periodic=True, dtype=torch.float32) + return self._window_torch + + def _extract_mel_features(self, audio: np.ndarray) -> np.ndarray: + """ + Compute the log-Mel spectrogram of audio using torch STFT. + + This intentionally mirrors the model-side preprocessing (`S3Tokenizer.log_mel_spectrogram`): + - `torch.stft(..., center=True)` with default `pad_mode="reflect"` + - Hann window + - power spectrogram + mel projection + - log10 clamp + dynamic range compression + scaling + + Returns: + np.ndarray of shape (time, n_mels) for Transformers padding convention. + """ + audio_t = torch.as_tensor(audio, dtype=torch.float32) + window = self._get_window_torch() + stft = torch.stft( + audio_t, + self.n_fft, + self.hop_length, + window=window, + center=True, + return_complex=True, + ) + magnitudes = stft[..., :-1].abs().pow(2) + + mel_filters = self._get_mel_filters_torch() + mel_spec = mel_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 + + # Transpose to [time, n_mels] for padding convention, and return numpy. + return log_spec.transpose(0, 1).cpu().numpy() + + def __call__( + self, + raw_audio: Union[np.ndarray, list[float], list[np.ndarray], list[list[float]]], + padding: Union[bool, str, PaddingStrategy] = False, + truncation: bool = False, + max_length: Optional[int] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + sampling_rate: Optional[int] = None, + **kwargs, + ) -> BatchFeature: + """ + Main method to featurize and prepare for the model one or several sequence(s). + """ + if sampling_rate is not None: + if sampling_rate != self.sampling_rate: + raise ValueError( + f"The model corresponding to this feature extractor: {self} was trained using a sampling rate of" + f" {self.sampling_rate}. Please make sure that the provided `raw_audio` input was sampled with" + f" {self.sampling_rate} and not {sampling_rate}." + ) + else: + logger.warning( + "It is strongly recommended to pass the `sampling_rate` argument to this function. " + "Failing to do so can result in silent errors that might be hard to debug." + ) + + is_batched = bool( + isinstance(raw_audio, (list, tuple)) and (isinstance(raw_audio[0], (np.ndarray, tuple, list))) + ) + + if is_batched: + raw_audio = [np.asarray(audio, dtype=np.float32).squeeze() for audio in raw_audio] + else: + raw_audio = [np.asarray(raw_audio, dtype=np.float32).squeeze()] + + # Ensure all are 1D + raw_audio = [audio if audio.ndim == 1 else audio.flatten() for audio in raw_audio] + + # Extract features + input_features = [self._extract_mel_features(audio) for audio in raw_audio] + + # convert into correct format for padding + encoded_inputs = BatchFeature({"input_features": input_features}) + + padded_inputs = self.pad( + encoded_inputs, + padding=padding, + max_length=max_length, + truncation=truncation, + pad_to_multiple_of=None, + return_attention_mask=True, + return_tensors=return_tensors, + ) + + return padded_inputs + + +__all__ = ["S3TokenizerFeatureExtractor"] diff --git a/src/transformers/models/s3tokenizer/modeling_s3tokenizer.py b/src/transformers/models/s3tokenizer/modeling_s3tokenizer.py new file mode 100644 index 000000000000..6f9830ccf546 --- /dev/null +++ b/src/transformers/models/s3tokenizer/modeling_s3tokenizer.py @@ -0,0 +1,655 @@ +# coding=utf-8 +# Copyright 2024 Resemble AI, xingchensong and The HuggingFace Inc. team. All rights reserved. +# +# This code is adapted from: +# - Chatterbox S3Tokenizer implementation +# - xingchensong/S3Tokenizer repository: https://github.com/xingchensong/S3Tokenizer +# - Original Whisper model: https://github.com/openai/whisper +# +# 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 S3Tokenizer model - self-contained implementation.""" + +from dataclasses import dataclass +from typing import Optional, Union + +import torch +from torch.nn.utils.rnn import pad_sequence + +from ...modeling_utils import PreTrainedModel +from ...utils import ModelOutput, auto_docstring, logging +from ..llama.modeling_llama import LlamaAttention +from .configuration_s3tokenizer import S3TokenizerConfig + + +logger = logging.get_logger(__name__) + + +# Sampling rate and frame configuration for S3TokenizerV2 +S3_SR = 16_000 +S3_HOP = 160 # 100 frames/sec +S3_TOKEN_HOP = 640 # 25 tokens/sec +S3_TOKEN_RATE = 25 +SPEECH_VOCAB_SIZE = 6561 + +# Special tokens +SOS = SPEECH_VOCAB_SIZE +EOS = SPEECH_VOCAB_SIZE + 1 + + +def make_non_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: + """Make mask tensor containing indices of non-padded part.""" + batch_size = lengths.size(0) + max_len = max_len if max_len > 0 else lengths.max().item() + seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device) + seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len) + seq_length_expand = lengths.unsqueeze(-1) + mask = seq_range_expand >= seq_length_expand + return ~mask + + +def mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """Convert bool-tensor to float-tensor for flash attention.""" + assert mask.dtype == torch.bool + assert dtype in [torch.float32, torch.bfloat16, torch.float16] + mask = mask.to(dtype) + mask = (1.0 - mask) * -1.0e10 + return mask + + +def padding(data: list[torch.Tensor]): + """Padding the data into batch data""" + sample = data + assert isinstance(sample, list) + feats_lengths = torch.tensor([s.size(1) for s in sample], dtype=torch.int32) + feats = [s.t() for s in sample] + padded_feats = pad_sequence(feats, batch_first=True, padding_value=0) + + return padded_feats.transpose(1, 2), feats_lengths + + +def merge_tokenized_segments(tokenized_segments, overlap, token_rate): + """Merges tokenized outputs by keeping the middle and dropping half of the overlapped tokens.""" + merged_tokens = [] + overlap_tokens = (overlap // 2) * token_rate + + for i, tokens in enumerate(tokenized_segments): + l = 0 if i == 0 else overlap_tokens + r = -overlap_tokens if i != len(tokenized_segments) - 1 else len(tokens) + merged_tokens.extend(tokens[l:r]) + + return merged_tokens + + +def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, scaling=None): + """Precompute frequencies for rotary embeddings.""" + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + t = torch.arange(end, device=freqs.device) + if scaling is not None: + t = t * scaling + freqs = torch.outer(t, freqs).float() + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + return torch.cat((freqs_cis, freqs_cis), dim=-1) + + +def apply_rotary_emb( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply rotary embeddings to query and key tensors.""" + real = torch.view_as_real(freqs_cis) + cos, sin = real[:, :, 0], real[:, :, 1] + cos = cos.unsqueeze(0).unsqueeze(2) + sin = sin.unsqueeze(0).unsqueeze(2) + + D = xq.shape[-1] + half_l, half_r = xq[:, :, :, : D // 2], xq[:, :, :, D // 2 :] + xq_r = torch.cat((-half_r, half_l), dim=-1) + + D = xk.shape[-1] + half_l, half_r = xk[:, :, :, : D // 2], xk[:, :, :, D // 2 :] + xk_r = torch.cat((-half_r, half_l), dim=-1) + + return xq * cos + xq_r * sin, xk * cos + xk_r * sin + + +class FSMNMultiHeadAttention(LlamaAttention): + """FSMN (Feed-forward Sequential Memory Network) Multi-Head Attention.""" + + def __init__(self, config: S3TokenizerConfig, layer_idx: Optional[int] = None, kernel_size: int = 31): + super().__init__(config, layer_idx) + self.is_causal = False + self.n_head = config.num_attention_heads + self.attention_bias = config.attention_bias + self.attention_dropout = config.attention_dropout + self.max_position_embeddings = config.max_position_embeddings + self.rope_scaling = config.rope_scaling + self.n_audio_head = config.n_audio_head + self.n_audio_state = config.n_audio_state + self.num_key_value_heads = config.num_key_value_heads + + self.fsmn_block = torch.nn.Conv1d( + config.hidden_size, + config.hidden_size, + kernel_size, + stride=1, + padding=0, + groups=config.hidden_size, + bias=False, + ) + self.left_padding = (kernel_size - 1) // 2 + self.right_padding = kernel_size - 1 - self.left_padding + self.pad_fn = torch.nn.ConstantPad1d((self.left_padding, self.right_padding), 0.0) + + # Re-initialize to match Chatterbox biases and use standard nn.Linear + self.query = torch.nn.Linear(config.hidden_size, config.hidden_size, bias=True) + self.key = torch.nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.value = torch.nn.Linear(config.hidden_size, config.hidden_size, bias=True) + self.out = torch.nn.Linear(config.hidden_size, config.hidden_size, bias=True) + + # Remove the old names created by LlamaAttention from _modules + del self._modules["q_proj"] + del self._modules["k_proj"] + del self._modules["v_proj"] + del self._modules["o_proj"] + + def forward_fsmn(self, inputs: torch.Tensor, mask: Optional[torch.Tensor] = None): + b, t, _ = inputs.size() + if mask is not None and mask.size(2) > 0: + inputs = inputs * mask + x = inputs.transpose(1, 2) + x = self.pad_fn(x) + x = self.fsmn_block(x) + x = x.transpose(1, 2) + x += inputs + if mask is not None: + x = x * mask + return x + + def forward( + self, + x: torch.Tensor, + mask: Optional[torch.Tensor] = None, + mask_pad: Optional[torch.Tensor] = None, + freqs_cis: Optional[torch.Tensor] = None, + **kwargs, + ): + q = self.query(x) + k = self.key(x) + v = self.value(x) + + # Calculate fsm_memory BEFORE permuting v + fsm_memory = self.forward_fsmn(v, mask_pad) + + # Exact baseline logic from here + _, _, D = q.shape + scale = (D // self.n_head) ** -0.25 + q = q.view(*q.shape[:2], self.n_head, -1) + k = k.view(*k.shape[:2], self.n_head, -1) + v = v.view(*v.shape[:2], self.n_head, -1) + + if freqs_cis is not None: + q, k = apply_rotary_emb(q, k, freqs_cis=freqs_cis) + + q = q.permute(0, 2, 1, 3) * scale + v = v.permute(0, 2, 1, 3) + + if not getattr(self.config, "use_sdpa", False): + k = k.permute(0, 2, 3, 1) * scale + qk = q @ k + if mask is not None: + qk = qk + mask + qk = qk.float() + w = torch.nn.functional.softmax(qk, dim=-1).to(q.dtype) + wv = (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2) + return self.out(wv) + fsm_memory, qk.detach() + else: + k = k.permute(0, 2, 1, 3) * scale + output = torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=mask, + dropout_p=0.0, + scale=1.0, + ) + output = output.transpose(1, 2).contiguous().view(q.size(0), -1, D) + return self.out(output) + fsm_memory, None + + +class FSQCodebook(torch.nn.Module): + """Finite Scalar Quantization codebook.""" + + def __init__(self, dim: int, level: int = 3): + super().__init__() + self.project_down = torch.nn.Linear(dim, 8) + self.level = level + self.embed = None + + @torch.inference_mode() + def preprocess(self, x: torch.Tensor) -> torch.Tensor: + x = x.reshape(-1, x.shape[-1]) + return x + + @torch.inference_mode() + def encode(self, x: torch.Tensor) -> torch.Tensor: + x_shape = x.shape + x = self.preprocess(x) + h = self.project_down(x).float() + h = h.tanh() + h = h * 0.9990000128746033 + h = h.round() + 1 + powers = torch.pow(self.level, torch.arange(2**self.level, device=x.device, dtype=h.dtype)) + mu = torch.sum(h * powers.unsqueeze(0), dim=-1) + ind = mu.reshape(x_shape[0], x_shape[1]).int() + return ind + + @torch.inference_mode() + def decode(self, embed_ind: torch.Tensor) -> torch.Tensor: + raise NotImplementedError("There is no official up project component provided") + + +class FSQVectorQuantization(torch.nn.Module): + """FSQ Vector quantization implementation (inference-only).""" + + def __init__(self, dim: int, codebook_size: int): + super().__init__() + assert 3**8 == codebook_size + self._codebook = FSQCodebook(dim=dim, level=3) + self.codebook_size = codebook_size + + @property + def codebook(self): + return self._codebook.embed + + @torch.inference_mode() + def encode(self, x: torch.Tensor) -> torch.Tensor: + return self._codebook.encode(x) + + @torch.inference_mode() + def decode(self, embed_ind: torch.Tensor) -> torch.Tensor: + quantize = self._codebook.decode(embed_ind) + quantize = quantize.transpose(1, 2) + return quantize + + +class ResidualAttentionBlock(torch.nn.Module): + """Residual attention block with FSMN.""" + + def __init__(self, config: S3TokenizerConfig, layer_idx: Optional[int] = None, kernel_size: int = 31): + super().__init__() + self.attn = FSMNMultiHeadAttention(config, layer_idx, kernel_size) + self.attn_ln = torch.nn.LayerNorm(config.hidden_size, eps=1e-6) + n_mlp = config.hidden_size * 4 + + # Using numeric keys to match state_dict + self.mlp = torch.nn.ModuleList( + [torch.nn.Linear(config.hidden_size, n_mlp), torch.nn.GELU(), torch.nn.Linear(n_mlp, config.hidden_size)] + ) + self.mlp_ln = torch.nn.LayerNorm(config.hidden_size) + + def forward( + self, + x: torch.Tensor, + mask: Optional[torch.Tensor] = None, + mask_pad: Optional[torch.Tensor] = None, + freqs_cis: Optional[torch.Tensor] = None, + ): + # LN Stability: cast to float32 for computation + attn_input = self.attn_ln(x.float()).to(x.dtype) + attn_out, _ = self.attn( + attn_input, + mask=mask, + mask_pad=mask_pad, + freqs_cis=freqs_cis, + ) + x = x + attn_out + + # LN Stability: cast to float32 for computation + mlp_input = self.mlp_ln(x.float()).to(x.dtype) + mlp_out = mlp_input + for layer in self.mlp: + mlp_out = layer(mlp_out) + + x = x + mlp_out + return x + + +class AudioEncoderV2(torch.nn.Module): + """Audio encoder for S3TokenizerV2.""" + + def __init__(self, config: S3TokenizerConfig): + super().__init__() + self.stride = 2 # Hardcoded for V2 + self.conv1 = torch.nn.Conv1d(config.n_mels, config.hidden_size, kernel_size=3, stride=self.stride, padding=1) + self.conv2 = torch.nn.Conv1d(config.hidden_size, config.hidden_size, kernel_size=3, stride=2, padding=1) + + self.register_buffer("freqs_cis", precompute_freqs_cis(64, 1024 * 2), persistent=False) + + self.blocks = torch.nn.ModuleList( + [ResidualAttentionBlock(config, layer_idx=i) for i in range(config.n_audio_layer)] + ) + + def forward(self, x: torch.Tensor, x_len: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + mask = make_non_pad_mask(x_len).unsqueeze(1) + x = torch.nn.functional.gelu(self.conv1(x * mask)) + x_len = (x_len + 2 - 1 * (3 - 1) - 1) // self.stride + 1 + mask = make_non_pad_mask(x_len).unsqueeze(1) + x = torch.nn.functional.gelu(self.conv2(x * mask)) + x_len = (x_len + 2 - 1 * (3 - 1) - 1) // 2 + 1 + mask = make_non_pad_mask(x_len).unsqueeze(1) + x = x.permute(0, 2, 1) + + mask_pad = mask.transpose(1, 2) + mask_bias = mask_to_bias(mask, x.dtype) + freqs_cis = self.freqs_cis.to(x.device) + + for block in self.blocks: + x = block(x, mask=mask_bias.unsqueeze(1), mask_pad=mask_pad, freqs_cis=freqs_cis[: x.size(1)]) + + return x, x_len + + +class S3TokenizerV2Core(torch.nn.Module): + """Core S3 tokenizer v2 implementation.""" + + def __init__(self, config: S3TokenizerConfig): + super().__init__() + self.encoder = AudioEncoderV2(config) + self.quantizer = FSQVectorQuantization(config.hidden_size, config.vocab_size) + + def forward(self, mel: torch.Tensor, mel_len: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return self.quantize(mel, mel_len) + + @torch.inference_mode() + def quantize(self, mel: torch.Tensor, mel_len: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize mel spectrogram to tokens, with automatic long audio handling.""" + max_frames = 3000 + long_audio_mask = mel_len > max_frames + + if long_audio_mask.any(): + return self._quantize_mixed_batch(mel, mel_len, long_audio_mask, max_frames) + else: + hidden, code_len = self.encoder(mel, mel_len) + code = self.quantizer.encode(hidden) + return code, code_len + + @torch.inference_mode() + def _quantize_mixed_batch( + self, + mel: torch.Tensor, + mel_len: torch.Tensor, + long_audio_mask: torch.Tensor, + max_frames: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Handle mixed batch with both short and long audio using unified batch processing.""" + batch_size = mel.size(0) + sample_rate, hop_length = 16000, 160 + window_size, overlap = 30, 4 + frames_per_window = window_size * sample_rate // hop_length + frames_per_overlap = overlap * sample_rate // hop_length + frames_per_stride = frames_per_window - frames_per_overlap + + all_segments, all_segments_len, segment_info = [], [], [] + + for batch_idx in range(batch_size): + audio_mel = mel[batch_idx] + audio_mel_len = mel_len[batch_idx] + is_long_audio = long_audio_mask[batch_idx].item() + + if not is_long_audio: + segment = audio_mel[:, :audio_mel_len] + seg_len = audio_mel_len.item() + if seg_len < frames_per_window: + segment = torch.nn.functional.pad(segment, (0, frames_per_window - seg_len)) + all_segments.append(segment) + all_segments_len.append(torch.tensor(seg_len, device=mel.device)) + segment_info.append( + { + "batch_idx": batch_idx, + "is_long_audio": False, + "segment_idx": 0, + "total_segments": 1, + } + ) + else: + start, segment_idx = 0, 0 + while start < audio_mel_len: + end = min(start + frames_per_window, audio_mel_len) + segment = audio_mel[:, start:end] + seg_len = segment.size(1) + if seg_len < frames_per_window: + segment = torch.nn.functional.pad(segment, (0, frames_per_window - seg_len)) + all_segments.append(segment) + all_segments_len.append(torch.tensor(seg_len, device=mel.device)) + segment_info.append( + { + "batch_idx": batch_idx, + "is_long_audio": True, + "segment_idx": segment_idx, + "total_segments": None, + } + ) + segment_idx += 1 + start += frames_per_stride + + for info in segment_info: + if info["batch_idx"] == batch_idx and info["is_long_audio"]: + info["total_segments"] = segment_idx + + if not all_segments: + return torch.zeros(batch_size, 0, dtype=torch.long, device=mel.device), torch.zeros( + batch_size, dtype=torch.long, device=mel.device + ) + + unified_batch_mel = torch.stack(all_segments) + unified_batch_lens = torch.stack(all_segments_len) + hidden, code_len = self.encoder(unified_batch_mel, unified_batch_lens) + codes = self.quantizer.encode(hidden) + + results = {} + for seg_idx, info in enumerate(segment_info): + batch_idx = info["batch_idx"] + segment_code = codes[seg_idx, : code_len[seg_idx].item()].cpu().numpy().tolist() + if not info["is_long_audio"]: + code_tensor = torch.tensor(segment_code, dtype=torch.long, device=mel.device) + results[batch_idx] = (code_tensor, len(segment_code)) + else: + if batch_idx not in results: + results[batch_idx] = [] + results[batch_idx].append(segment_code) + + for batch_idx in range(batch_size): + if long_audio_mask[batch_idx].item(): + audio_codes = results[batch_idx] + merged_codes = merge_tokenized_segments(audio_codes, overlap=overlap, token_rate=25) + merged_codes_tensor = torch.tensor(merged_codes, dtype=torch.long, device=mel.device) + results[batch_idx] = (merged_codes_tensor, len(merged_codes)) + + max_code_len = max(code_info[1] for code_info in results.values()) + output_codes = torch.zeros(batch_size, max_code_len, dtype=torch.long, device=mel.device) + output_codes_len = torch.zeros(batch_size, dtype=torch.long, device=mel.device) + + for batch_idx, (code_tensor, code_len) in results.items(): + output_codes[batch_idx, :code_len] = code_tensor + output_codes_len[batch_idx] = code_len + + return output_codes, output_codes_len + + @property + def device(self): + return next(self.parameters()).device + + +@dataclass +@auto_docstring +class S3TokenizerOutput(ModelOutput): + r""" + speech_tokens (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Discrete speech tokens computed using `model.quantize`. + speech_token_lens (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Length of each speech token sequence. + """ + + speech_tokens: Optional[torch.LongTensor] = None + speech_token_lens: Optional[torch.LongTensor] = None + + +class S3TokenizerPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = S3TokenizerConfig + base_model_prefix = "s3tokenizer" + main_input_name = "input_features" + + def _init_weights(self, module): + """Initialize weights. + + S3Tokenizer models are expected to be loaded from pretrained checkpoints, but we still need to correctly + (re)initialize buffers when the model is created on meta device then materialized with `to_empty()`. + """ + # These buffers are registered in `__init__` to match checkpoint keys. When a model is initialized on meta + # device, then materialized via `to_empty()`, buffers may contain uninitialized values and need to be restored + # deterministically here. + if isinstance(module, S3TokenizerModel): + # During `from_pretrained`, core loading will set `_is_hf_initialized=True` on loaded tensors. + # Do not overwrite buffers that were loaded from the checkpoint. + if not getattr(module.window, "_is_hf_initialized", False): + module.window = torch.zeros( + module.config.n_fft, device=module.window.device, dtype=module.window.dtype + ) + if not getattr(module._mel_filters, "_is_hf_initialized", False): + module._mel_filters = torch.zeros( + module.config.n_mels, + module.config.n_fft // 2 + 1, + device=module._mel_filters.device, + dtype=module._mel_filters.dtype, + ) + elif isinstance(module, AudioEncoderV2): + if not getattr(module.freqs_cis, "_is_hf_initialized", False): + module.freqs_cis = precompute_freqs_cis(64, 1024 * 2).to(device=module.freqs_cis.device) + + +class S3TokenizerModel(S3TokenizerPreTrainedModel): + """ + S3Tokenizer model for speech tokenization. + + This model integrates the S3Tokenizer implementation from xingchensong/S3Tokenizer + repository into HuggingFace Transformers. + + Args: + config (`S3TokenizerConfig`): + name (`str`, *optional*, defaults to `"speech_tokenizer_v2_25hz"`): + """ + + all_tied_weights_keys = {} + + def __init__(self, config: S3TokenizerConfig, name: str = "speech_tokenizer_v2_25hz"): + super().__init__(config) + self.config = config + + self.s3_model = S3TokenizerV2Core(config) + + # Register buffers for STFT to match checkpoint keys + self.register_buffer("window", torch.zeros(config.n_fft)) + self.register_buffer("_mel_filters", torch.zeros(config.n_mels, config.n_fft // 2 + 1)) + + # Initialize weights and apply final processing + self.post_init() + + def forward( + self, + input_features: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + max_len: Optional[int] = None, + return_dict: Optional[bool] = None, + **kwargs, + ) -> Union[tuple, S3TokenizerOutput]: + """ + Args: + input_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, n_mels)`): + Float values of log-mel spectrogram features. + attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing operations on padding token indices. + max_len (`int`, *optional*): + Maximum length to truncate the output sequence to (25 token/sec). + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + + Returns: + `S3TokenizerOutput` or `tuple`: Speech tokens and their lengths. + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if max_len is not None: + input_features = input_features[..., : max_len * 4] + + # Get mel lengths from attention_mask or input_features shape + if attention_mask is not None: + mel_lens = attention_mask.sum(dim=-1).int() + else: + mel_lens = torch.full( + (input_features.size(0),), input_features.size(1), device=input_features.device, dtype=torch.int32 + ) + + # Transpose from [batch, time, n_mels] to [batch, n_mels, time] for conv layers + input_features = input_features.transpose(1, 2) + + speech_tokens, speech_token_lens = self.s3_model.quantize(input_features, mel_lens) + speech_tokens = speech_tokens.long().detach() + speech_token_lens = speech_token_lens.long().detach() + + if not return_dict: + return (speech_tokens, speech_token_lens) + + return S3TokenizerOutput( + speech_tokens=speech_tokens, + speech_token_lens=speech_token_lens, + ) + + def get_input_embeddings(self): + """S3Tokenizer does not use input embeddings in the traditional sense.""" + return None + + @property + def device(self): + return next(self.parameters()).device + + +def drop_invalid_tokens(x: torch.Tensor) -> torch.Tensor: + """Drop SoS and EoS tokens from speech token sequence.""" + assert len(x.shape) == 1 or (len(x.shape) == 2 and x.shape[0] == 1), "only batch size of one allowed for now" + + if SOS in x: + s = (x == SOS).nonzero(as_tuple=True)[0].squeeze(0) + 1 + else: + s = 0 + + if EOS in x: + e = (x == EOS).nonzero(as_tuple=True)[0].squeeze(0) + else: + e = None + + x = x[s:e] + return x + + +__all__ = [ + "S3TokenizerModel", + "S3TokenizerPreTrainedModel", + "S3TokenizerOutput", + "drop_invalid_tokens", +] diff --git a/src/transformers/utils/auto_docstring.py b/src/transformers/utils/auto_docstring.py index 2f435189c219..f2175e028b76 100644 --- a/src/transformers/utils/auto_docstring.py +++ b/src/transformers/utils/auto_docstring.py @@ -68,6 +68,7 @@ "donut": "DonutSwinConfig", "esmfold": "EsmConfig", "parakeet": "ParakeetCTCConfig", + "s3tokenizer": "S3TokenizerConfig", "lasr": "LasrCTCConfig", } diff --git a/tests/models/chatterbox/__init__.py b/tests/models/chatterbox/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/models/chatterbox/test_modeling_chatterbox.py b/tests/models/chatterbox/test_modeling_chatterbox.py new file mode 100644 index 000000000000..077086271cfb --- /dev/null +++ b/tests/models/chatterbox/test_modeling_chatterbox.py @@ -0,0 +1,104 @@ +# coding=utf-8 +# Copyright 2025 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. +"""Testing suite for the PyTorch Chatterbox model.""" + +import tempfile +import unittest + +from transformers.models.auto.feature_extraction_auto import FEATURE_EXTRACTOR_MAPPING_NAMES +from transformers.models.chatterbox.configuration_chatterbox import ChatterboxConfig +from transformers.models.chatterbox.feature_extraction_chatterbox import ChatterboxFeatureExtractor +from transformers.models.chatterbox.modeling_chatterbox import ChatterboxModel +from transformers.testing_utils import require_torch + + +@require_torch +class ChatterboxModelTest(unittest.TestCase): + def setUp(self): + """Set up test configuration.""" + self.config = ChatterboxConfig() + # Use smaller model for faster tests + self.config.t3_config.llama_config_dict["num_hidden_layers"] = 2 + self.config.t3_config.llama_config_dict["num_attention_heads"] = 4 + self.config.t3_config.hidden_size = 256 + self.config.s3gen_config.encoder_num_blocks = 2 + self.config.s3gen_config.decoder_n_blocks = 2 + + def test_model_initialization(self): + """Test that the model can be initialized.""" + model = ChatterboxModel(self.config) + self.assertIsInstance(model, ChatterboxModel) + + # Check that sub-modules exist + self.assertIsNotNone(model.t3) + self.assertIsNotNone(model.s3gen) + self.assertIsInstance(model.feature_extractor, ChatterboxFeatureExtractor) + + def test_config_attributes(self): + """Test that config attributes are properly set.""" + model = ChatterboxModel(self.config) + + # Check that sub-configs exist + self.assertIsNotNone(model.config.t3_config) + self.assertIsNotNone(model.config.s3gen_config) + + def test_save_and_load(self): + """Test saving and loading the model.""" + model = ChatterboxModel(self.config) + + with tempfile.TemporaryDirectory() as tmpdirname: + # Save model + model.save_pretrained(tmpdirname) + + # Check that files were created + import os + + self.assertTrue(os.path.exists(os.path.join(tmpdirname, "config.json"))) + self.assertTrue( + os.path.exists(os.path.join(tmpdirname, "model.safetensors")) + or os.path.exists(os.path.join(tmpdirname, "pytorch_model.bin")) + ) + + # Load model + loaded_model = ChatterboxModel.from_pretrained(tmpdirname) + self.assertIsInstance(loaded_model, ChatterboxModel) + + def test_auto_feature_extractor_mapping(self): + self.assertIn("chatterbox", FEATURE_EXTRACTOR_MAPPING_NAMES) + self.assertEqual(FEATURE_EXTRACTOR_MAPPING_NAMES["chatterbox"], "ChatterboxFeatureExtractor") + + @unittest.skip("Requires CUDA and full tokenizer setup") + def test_generate_basic(self): + """Test basic generation.""" + # Skipped: This test requires CUDA device and full tokenizer setup + pass + + def test_config_serialization(self): + """Test config serialization and deserialization.""" + config_dict = self.config.to_dict() + + # Check that nested configs are serialized + self.assertIn("t3_config", config_dict) + self.assertIn("s3gen_config", config_dict) + self.assertIn("hiftnet_config", config_dict) + + # Test reconstruction - just check it doesn't crash + # Note: exact equality may not hold due to nested config defaults + new_config = ChatterboxConfig.from_dict(config_dict) + self.assertIsInstance(new_config, ChatterboxConfig) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/models/s3gen/__init__.py b/tests/models/s3gen/__init__.py new file mode 100644 index 000000000000..ae8ccabe497b --- /dev/null +++ b/tests/models/s3gen/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2025 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. + diff --git a/tests/models/s3gen/test_modeling_s3gen.py b/tests/models/s3gen/test_modeling_s3gen.py new file mode 100644 index 000000000000..4827d5f884ee --- /dev/null +++ b/tests/models/s3gen/test_modeling_s3gen.py @@ -0,0 +1,210 @@ +# coding=utf-8 +# Copyright 2025 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. +"""Testing suite for the PyTorch S3Gen model.""" + +import unittest + +import torch + +from transformers.models.s3gen.configuration_s3gen import S3GenConfig +from transformers.models.s3gen.modeling_s3gen import S3GenModel +from transformers.testing_utils import require_torch, torch_device + + +@require_torch +class S3GenModelTest(unittest.TestCase): + def setUp(self): + self.config = S3GenConfig( + vocab_size=6561, + token_embed_dim=512, + speaker_embed_dim=192, + encoder_output_size=512, + encoder_attention_heads=8, + encoder_linear_units=2048, + encoder_num_blocks=6, + decoder_in_channels=320, + decoder_out_channels=80, + decoder_channels=[256], + decoder_n_blocks=4, + decoder_num_mid_blocks=12, + sampling_rate=24000, + mel_bins=80, + ) + + def test_model_initialization(self): + """Test that the model can be initialized.""" + model = S3GenModel(self.config) + self.assertIsInstance(model, S3GenModel) + + # Check that sub-modules exist + self.assertIsNotNone(model.tokenizer) + self.assertIsNotNone(model.speaker_encoder) + self.assertIsNotNone(model.flow) + self.assertIsNotNone(model.mel2wav) + + def test_forward_pass_with_ref_wav(self): + """Test forward pass with reference audio.""" + model = S3GenModel(self.config) + model.eval() + model.to(torch_device) + + batch_size = 1 + seq_len = 50 + audio_len = 8000 + + # Create dummy inputs + speech_tokens = torch.randint(0, self.config.vocab_size, (batch_size, seq_len), device=torch_device) + ref_wav = torch.randn(batch_size, audio_len, device=torch_device) + ref_sr = 16000 + + # Run forward pass + with torch.no_grad(): + output = model(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, finalize=True) + + # Check output shape + self.assertEqual(len(output.shape), 3) + self.assertEqual(output.shape[0], batch_size) + self.assertEqual(output.shape[1], self.config.mel_bins) + self.assertGreater(output.shape[2], 0) # Time dimension + + def test_forward_pass_with_ref_dict(self): + """Test forward pass with pre-computed reference embeddings.""" + model = S3GenModel(self.config) + model.eval() + model.to(torch_device) + + batch_size = 1 + seq_len = 50 + audio_len = 8000 + + # Create dummy inputs + speech_tokens = torch.randint(0, self.config.vocab_size, (batch_size, seq_len), device=torch_device) + ref_wav = torch.randn(batch_size, audio_len, device=torch_device) + ref_sr = 16000 + + # Extract reference embeddings + ref_dict = model.embed_ref(ref_wav, ref_sr) + + # Run forward pass with ref_dict + with torch.no_grad(): + output = model(speech_tokens, ref_dict=ref_dict, finalize=True) + + # Check output shape + self.assertEqual(len(output.shape), 3) + self.assertEqual(output.shape[0], batch_size) + + def test_embed_ref(self): + """Test reference embedding extraction.""" + model = S3GenModel(self.config) + model.eval() + model.to(torch_device) + + batch_size = 1 + audio_len = 8000 + + ref_wav = torch.randn(batch_size, audio_len, device=torch_device) + ref_sr = 16000 + + ref_dict = model.embed_ref(ref_wav, ref_sr) + + # Check that all required keys are present + self.assertIn("prompt_token", ref_dict) + self.assertIn("prompt_token_len", ref_dict) + self.assertIn("prompt_feat", ref_dict) + self.assertIn("embedding", ref_dict) + + # Check shapes + self.assertEqual(len(ref_dict["embedding"].shape), 2) + self.assertEqual(ref_dict["embedding"].shape[0], batch_size) + self.assertEqual(ref_dict["embedding"].shape[1], self.config.speaker_embed_dim) + + def test_end_to_end_inference(self): + """Test end-to-end inference (tokens → waveform).""" + model = S3GenModel(self.config) + model.eval() + model.to(torch_device) + + batch_size = 1 + seq_len = 50 + audio_len = 8000 + + speech_tokens = torch.randint(0, self.config.vocab_size, (batch_size, seq_len), device=torch_device) + ref_wav = torch.randn(batch_size, audio_len, device=torch_device) + ref_sr = 16000 + + # Run end-to-end inference + with torch.no_grad(): + wavs, sources = model.inference(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, finalize=True) + + # Check output shapes + self.assertEqual(len(wavs.shape), 2) + self.assertEqual(wavs.shape[0], batch_size) + self.assertGreater(wavs.shape[1], 0) # Audio samples + + self.assertEqual(len(sources.shape), 3) + self.assertEqual(sources.shape[0], batch_size) + + def test_save_and_load(self): + """Test saving and loading the model.""" + import os + import tempfile + + model = S3GenModel(self.config) + + with tempfile.TemporaryDirectory() as tmpdirname: + # Save model + model.save_pretrained(tmpdirname) + + # Check that files were created + self.assertTrue(os.path.exists(os.path.join(tmpdirname, "config.json"))) + self.assertTrue( + os.path.exists(os.path.join(tmpdirname, "model.safetensors")) + or os.path.exists(os.path.join(tmpdirname, "pytorch_model.bin")) + ) + + # Load model + loaded_model = S3GenModel.from_pretrained(tmpdirname) + self.assertIsInstance(loaded_model, S3GenModel) + + def test_different_token_lengths(self): + """Test with different token sequence lengths.""" + model = S3GenModel(self.config) + model.eval() + model.to(torch_device) + + ref_wav = torch.randn(1, 8000, device=torch_device) + ref_sr = 16000 + + for seq_len in [10, 50, 100]: + speech_tokens = torch.randint(0, self.config.vocab_size, (1, seq_len), device=torch_device) + + with torch.no_grad(): + output = model(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, finalize=True) + + self.assertEqual(output.shape[0], 1) + self.assertEqual(output.shape[1], self.config.mel_bins) + + def test_config_attributes(self): + """Test that config attributes are properly set.""" + model = S3GenModel(self.config) + + self.assertEqual(model.config.vocab_size, 6561) + self.assertEqual(model.config.speaker_embed_dim, 192) + self.assertEqual(model.config.mel_bins, 80) + self.assertEqual(model.config.sampling_rate, 24000) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/models/s3tokenizer/__init__.py b/tests/models/s3tokenizer/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/models/s3tokenizer/test_modeling_s3tokenizer.py b/tests/models/s3tokenizer/test_modeling_s3tokenizer.py new file mode 100644 index 000000000000..798bcf64a488 --- /dev/null +++ b/tests/models/s3tokenizer/test_modeling_s3tokenizer.py @@ -0,0 +1,212 @@ +# Copyright 2024 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. +"""Testing suite for the PyTorch S3Tokenizer model.""" + +import unittest + +from transformers import S3TokenizerConfig +from transformers.testing_utils import is_torch_available, require_torch, slow, torch_device + +from ...test_configuration_common import ConfigTester +from ...test_modeling_common import ModelTesterMixin, floats_tensor + + +if is_torch_available(): + from transformers import S3TokenizerModel + + +@require_torch +class S3TokenizerModelTester: + def __init__( + self, + parent, + batch_size=2, + seq_length=400, + is_training=False, + use_labels=False, + ): + self.parent = parent + self.batch_size = batch_size + self.seq_length = seq_length + self.is_training = is_training + self.use_labels = use_labels + + def prepare_config_and_inputs(self): + config = self.get_config() + input_features = floats_tensor([self.batch_size, self.seq_length, config.n_mels], scale=1.0) + inputs_dict = {"input_features": input_features} + return config, inputs_dict + + def prepare_config_and_inputs_for_common(self): + config, inputs_dict = self.prepare_config_and_inputs() + return config, inputs_dict + + def get_config(self): + return S3TokenizerConfig( + n_mels=80, + n_audio_state=512, + n_audio_head=8, + n_audio_layer=6, + vocab_size=6561, + n_fft=400, + hop_length=160, + sampling_rate=16000, + use_sdpa=False, + ) + + def create_and_check_model(self, config, input_features): + model = S3TokenizerModel(config=config) + model.to(torch_device) + model.eval() + result = model(input_features) + self.parent.assertIsNotNone(result.speech_tokens) + self.parent.assertIsNotNone(result.speech_token_lens) + + +@require_torch +class S3TokenizerModelTest(ModelTesterMixin, unittest.TestCase): + all_model_classes = (S3TokenizerModel,) if is_torch_available() else () + is_encoder_decoder = False + test_pruning = False + test_headmasking = False + test_resize_embeddings = False + test_torchscript = False + test_missing_keys = False + test_model_parallel = False + test_head_masking = False + + def setUp(self): + self.model_tester = S3TokenizerModelTester(self) + self.config_tester = ConfigTester( + self, config_class=S3TokenizerConfig, has_text_modality=False, common_properties=["hidden_size"] + ) + + def test_config(self): + self.config_tester.run_common_tests() + + def test_model(self): + config, inputs_dict = self.model_tester.prepare_config_and_inputs() + self.model_tester.create_and_check_model(config, inputs_dict["input_features"]) + + @unittest.skip(reason="S3Tokenizer does not output hidden states in the traditional sense") + def test_hidden_states_output(self): + pass + + @unittest.skip(reason="S3Tokenizer does not have attention weights in the traditional sense") + def test_attention_outputs(self): + pass + + @unittest.skip(reason="S3Tokenizer does not support input embeddings") + def test_inputs_embeds(self): + pass + + @unittest.skip(reason="S3Tokenizer does not support training mode") + def test_training(self): + pass + + @unittest.skip(reason="S3Tokenizer does not support training mode") + def test_training_gradient_checkpointing(self): + pass + + @unittest.skip(reason="S3Tokenizer does not support retain_grad") + def test_retain_grad_hidden_states_attentions(self): + pass + + @unittest.skip(reason="S3Tokenizer does not have typical model forward signature") + def test_forward_signature(self): + pass + + @unittest.skip(reason="S3Tokenizer model does not support typical model features") + def test_model_common_attributes(self): + pass + + @unittest.skip(reason="S3Tokenizer does not have input/output embeddings") + def test_model_get_set_embeddings(self): + pass + + @unittest.skip(reason="S3Tokenizer does not use feed forward chunking") + def test_feed_forward_chunking(self): + pass + + @unittest.skip(reason="S3Tokenizer model is too large for common tests") + def test_model_is_small(self): + pass + + @unittest.skip(reason="S3Tokenizer does not support output_hidden_states") + def test_model_outputs_equivalence(self): + pass + + @unittest.skip(reason="S3Tokenizer does not support init weights") + def test_can_init_all_missing_weights(self): + pass + + @unittest.skip(reason="S3Tokenizer does not support safetensors") + def test_can_use_safetensors(self): + return super().test_can_use_safetensors() + + @unittest.skip(reason="S3Tokenizer does not support tied weights") + def test_load_save_without_tied_weights(self): + return super().test_load_save_without_tied_weights() + + @unittest.skip(reason="S3Tokenizer does not support init weights") + def test_save_load(self): + return super().test_save_load() + + def test_window_buffer_loading(self): + """Test that the window buffer can be loaded from checkpoint if it exists.""" + import tempfile + + import torch + + config = self.model_tester.get_config() + model1 = S3TokenizerModel(config=config) + + # Modify the window to a custom value + custom_window = torch.ones_like(model1.window) * 0.5 + model1.window = custom_window + + # Save the model with the custom window + with tempfile.TemporaryDirectory() as tmp_dir: + model1.save_pretrained(tmp_dir) + + # Load the model and verify the window was loaded + model2 = S3TokenizerModel.from_pretrained(tmp_dir) + + # Verify the custom window was loaded from checkpoint + self.assertTrue(torch.allclose(model2.window, custom_window)) + + def test_window_buffer_missing_from_checkpoint(self): + """Test that the default window is used when not present in checkpoint.""" + import torch + + config = self.model_tester.get_config() + model = S3TokenizerModel(config=config) + + # Create a state dict without window + state_dict = {} + for key, value in model.state_dict().items(): + if key != "window": + state_dict[key] = value + + # Load the state dict (window should remain as default) + default_window = model.window.clone() + model.load_state_dict(state_dict, strict=False) + + # Verify the default window is still used + self.assertTrue(torch.allclose(model.window, default_window)) + + @slow + @require_torch + def test_model_from_pretrained(self): + pass diff --git a/utils/check_docstrings.py b/utils/check_docstrings.py index 2d64f8d27e4c..dbfb7dd83cce 100644 --- a/utils/check_docstrings.py +++ b/utils/check_docstrings.py @@ -436,6 +436,7 @@ class DecoratedItem: "SpeechT5Model", "SplinterConfig", "SplinterTokenizerFast", + "S3TokenizerModel", "SqueezeBertTokenizerFast", "SummarizationPipeline", "Swin2SRImageProcessor", @@ -518,6 +519,9 @@ class DecoratedItem: "Llama4TextConfig", "BltConfig", "BltPatcherConfig", + "T3Config", + "S3GenConfig", + "S3TokenizerConfig", } # In addition to the objects above, we also ignore objects with certain prefixes. If you add an item to the list # below, make sure to add a comment explaining why. diff --git a/utils/check_repo.py b/utils/check_repo.py index f36cda07dc51..0a3d5d010d37 100644 --- a/utils/check_repo.py +++ b/utils/check_repo.py @@ -106,6 +106,8 @@ "BltLocalDecoder", # Building part of bigger (tested) model. Tested implicitly through BLTForCausalLM. "BltGlobalTransformer", # Building part of bigger (tested) model. Tested implicitly through BLTForCausalLM. "Ovis2VisionModel", + "T3Model", # Building part of bigger (tested) model. + "T3PreTrainedModel", # Building part of bigger (tested) model. "PeAudioPreTrainedModel", "PeAudioVideoPreTrainedModel", "PeVideoPreTrainedModel", @@ -225,6 +227,9 @@ "models/vision_text_dual_encoder/test_modeling_vision_text_dual_encoder.py", "models/decision_transformer/test_modeling_decision_transformer.py", "models/bark/test_modeling_bark.py", + "models/s3gen/test_modeling_s3gen.py", + "models/t3/test_modeling_t3.py", + "models/chatterbox/test_modeling_chatterbox.py", "models/shieldgemma2/test_modeling_shieldgemma2.py", "models/llama4/test_modeling_llama4.py", "models/sam2_video/test_modeling_sam2_video.py", @@ -419,6 +424,8 @@ "Qwen3OmniMoeTalkerModel", # Building part of a bigger model "Qwen3OmniMoeThinkerForConditionalGeneration", # Building part of a bigger model "Qwen3OmniMoeThinkerTextModel", # Building part of a bigger model + "S3TokenizerModel", # Building part of a bigger model + "S3GenModel", # Building part of a bigger model "Ernie4_5_VL_MoeTextModel", # Building part of a bigger model "PeAudioFrameLevelModel", ] @@ -1025,6 +1032,7 @@ def find_all_documented_objects() -> list[str]: "VitPoseBackbone", # Internal module "VitPoseBackboneConfig", # Internal module "get_values", # Internal object + "T3Cond", # Internal conditioning class for T3 model ] # This list should be empty. Objects in it should get their own doc page.