Showing preview only (375K chars total). Download the full file or copy to clipboard to get everything.
Repository: resemble-ai/chatterbox
Branch: master
Commit: eaf93540a1dd
Files: 64
Total size: 354.7 KB
Directory structure:
gitextract_lols3utf/
├── .github/
│ └── workflows/
│ └── install_check.yml
├── .gitignore
├── LICENSE
├── README.md
├── example_for_mac.py
├── example_tts.py
├── example_tts_turbo.py
├── example_vc.py
├── gradio_tts_app.py
├── gradio_tts_turbo_app.py
├── gradio_vc_app.py
├── multilingual_app.py
├── pyproject.toml
└── src/
└── chatterbox/
├── __init__.py
├── models/
│ ├── __init__.py
│ ├── s3gen/
│ │ ├── __init__.py
│ │ ├── configs.py
│ │ ├── const.py
│ │ ├── decoder.py
│ │ ├── f0_predictor.py
│ │ ├── flow.py
│ │ ├── flow_matching.py
│ │ ├── hifigan.py
│ │ ├── matcha/
│ │ │ ├── decoder.py
│ │ │ ├── flow_matching.py
│ │ │ ├── text_encoder.py
│ │ │ └── transformer.py
│ │ ├── s3gen.py
│ │ ├── transformer/
│ │ │ ├── __init__.py
│ │ │ ├── activation.py
│ │ │ ├── attention.py
│ │ │ ├── convolution.py
│ │ │ ├── embedding.py
│ │ │ ├── encoder_layer.py
│ │ │ ├── positionwise_feed_forward.py
│ │ │ ├── subsampling.py
│ │ │ └── upsample_encoder.py
│ │ ├── utils/
│ │ │ ├── class_utils.py
│ │ │ ├── intmeanflow.py
│ │ │ ├── mask.py
│ │ │ └── mel.py
│ │ └── xvector.py
│ ├── s3tokenizer/
│ │ ├── __init__.py
│ │ └── s3tokenizer.py
│ ├── t3/
│ │ ├── __init__.py
│ │ ├── inference/
│ │ │ ├── alignment_stream_analyzer.py
│ │ │ └── t3_hf_backend.py
│ │ ├── llama_configs.py
│ │ ├── modules/
│ │ │ ├── cond_enc.py
│ │ │ ├── learned_pos_emb.py
│ │ │ ├── perceiver.py
│ │ │ └── t3_config.py
│ │ └── t3.py
│ ├── tokenizers/
│ │ ├── __init__.py
│ │ └── tokenizer.py
│ ├── utils.py
│ └── voice_encoder/
│ ├── __init__.py
│ ├── config.py
│ ├── melspec.py
│ └── voice_encoder.py
├── mtl_tts.py
├── tts.py
├── tts_turbo.py
└── vc.py
================================================
FILE CONTENTS
================================================
================================================
FILE: .github/workflows/install_check.yml
================================================
name: Test Installation
on:
push:
branches: [ "master" ]
pull_request:
branches: [ "master" ]
jobs:
build:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v3
- name: Set up Python 3.10
uses: actions/setup-python@v4
with:
python-version: "3.10"
- name: Test Standard Install
run: |
pip install -e .
================================================
FILE: .gitignore
================================================
.vscode
# Pylance
pyrightconfig.json
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
syn_out/
checkpoints/
.gradio
# Ignore generated sample .wav files
**/*.wav
================================================
FILE: LICENSE
================================================
MIT License
Copyright (c) 2025 Resemble AI
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
================================================
FILE: README.md
================================================
# Chatterbox TTS
[](https://resemble-ai.github.io/chatterbox_turbo_demopage/)
[](https://huggingface.co/spaces/ResembleAI/chatterbox-turbo-demo)
[](https://podonos.com/resembleai/chatterbox)
[](https://discord.gg/rJq9cRJBJ6)
*Made with ♥️ by* <a href="https://resemble.ai" target="_blank"><img width="100" alt="resemble-logo-horizontal" src="https://github.com/user-attachments/assets/35cf756b-3506-4943-9c72-c05ddfa4e525" /></a>
**Chatterbox** is a family of three state-of-the-art, open-source text-to-speech models by Resemble AI.
We are excited to introduce **Chatterbox-Turbo**, our most efficient model yet. Built on a streamlined 350M parameter architecture, **Turbo** delivers high-quality speech with less compute and VRAM than our previous models. We have also distilled the speech-token-to-mel decoder, previously a bottleneck, reducing generation from 10 steps to just **one**, while retaining high-fidelity audio output.
**Paralinguistic tags** are now native to the Turbo model, allowing you to use `[cough]`, `[laugh]`, `[chuckle]`, and more to add distinct realism. While Turbo was built primarily for low-latency voice agents, it excels at narration and creative workflows.
If you like the model but need to scale or tune it for higher accuracy, check out our competitively priced TTS service (<a href="https://resemble.ai">link</a>). It delivers reliable performance with ultra-low latency of sub 200ms—ideal for production use in agents, applications, or interactive media.
<img width="1200" height="600" alt="Podonos Turbo Eval" src="https://storage.googleapis.com/chatterbox-demo-samples/turbo/podonos_turbo.png" />
### ⚡ Model Zoo
Choose the right model for your application.
| Model | Size | Languages | Key Features | Best For | 🤗 | Examples |
|:----------------------------------------------------------------------------------------------------------------| :--- | :--- |:--------------------------------------------------------|:---------------------------------------------|:--------------------------------------------------------------------------| :--- |
| **Chatterbox-Turbo** | **350M** | **English** | Paralinguistic Tags (`[laugh]`), Lower Compute and VRAM | Zero-shot voice agents, Production | [Demo](https://huggingface.co/spaces/ResembleAI/chatterbox-turbo-demo) | [Listen](https://resemble-ai.github.io/chatterbox_turbo_demopage/) |
| Chatterbox-Multilingual [(Language list)](#supported-languages) | 500M | 23+ | Zero-shot cloning, Multiple Languages | Global applications, Localization | [Demo](https://huggingface.co/spaces/ResembleAI/Chatterbox-Multilingual-TTS) | [Listen](https://resemble-ai.github.io/chatterbox_demopage/) |
| Chatterbox [(Tips and Tricks)](#original-chatterbox-tips) | 500M | English | CFG & Exaggeration tuning | General zero-shot TTS with creative controls | [Demo](https://huggingface.co/spaces/ResembleAI/Chatterbox) | [Listen](https://resemble-ai.github.io/chatterbox_demopage/) |
## Installation
```shell
pip install chatterbox-tts
```
Alternatively, you can install from source:
```shell
# conda create -yn chatterbox python=3.11
# conda activate chatterbox
git clone https://github.com/resemble-ai/chatterbox.git
cd chatterbox
pip install -e .
```
We developed and tested Chatterbox on Python 3.11 on Debian 11 OS; the versions of the dependencies are pinned in `pyproject.toml` to ensure consistency. You can modify the code or dependencies in this installation mode.
## Usage
##### Chatterbox-Turbo
```python
import torchaudio as ta
import torch
from chatterbox.tts_turbo import ChatterboxTurboTTS
# Load the Turbo model
model = ChatterboxTurboTTS.from_pretrained(device="cuda")
# Generate with Paralinguistic Tags
text = "Hi there, Sarah here from MochaFone calling you back [chuckle], have you got one minute to chat about the billing issue?"
# Generate audio (requires a reference clip for voice cloning)
wav = model.generate(text, audio_prompt_path="your_10s_ref_clip.wav")
ta.save("test-turbo.wav", wav, model.sr)
```
##### Chatterbox and Chatterbox-Multilingual
```python
import torchaudio as ta
from chatterbox.tts import ChatterboxTTS
from chatterbox.mtl_tts import ChatterboxMultilingualTTS
# English example
model = ChatterboxTTS.from_pretrained(device="cuda")
text = "Ezreal and Jinx teamed up with Ahri, Yasuo, and Teemo to take down the enemy's Nexus in an epic late-game pentakill."
wav = model.generate(text)
ta.save("test-english.wav", wav, model.sr)
# Multilingual examples
multilingual_model = ChatterboxMultilingualTTS.from_pretrained(device=device)
french_text = "Bonjour, comment ça va? Ceci est le modèle de synthèse vocale multilingue Chatterbox, il prend en charge 23 langues."
wav_french = multilingual_model.generate(french_text, language_id="fr")
ta.save("test-french.wav", wav_french, model.sr)
chinese_text = "你好，今天天气真不错，希望你有一个愉快的周末。"
wav_chinese = multilingual_model.generate(chinese_text, language_id="zh")
ta.save("test-chinese.wav", wav_chinese, model.sr)
# If you want to synthesize with a different voice, specify the audio prompt
AUDIO_PROMPT_PATH = "YOUR_FILE.wav"
wav = model.generate(text, audio_prompt_path=AUDIO_PROMPT_PATH)
ta.save("test-2.wav", wav, model.sr)
```
See `example_tts.py` and `example_vc.py` for more examples.
## Supported Languages
Arabic (ar) • Danish (da) • German (de) • Greek (el) • English (en) • Spanish (es) • Finnish (fi) • French (fr) • Hebrew (he) • Hindi (hi) • Italian (it) • Japanese (ja) • Korean (ko) • Malay (ms) • Dutch (nl) • Norwegian (no) • Polish (pl) • Portuguese (pt) • Russian (ru) • Swedish (sv) • Swahili (sw) • Turkish (tr) • Chinese (zh)
## Original Chatterbox Tips
- **General Use (TTS and Voice Agents):**
- Ensure that the reference clip matches the specified language tag. Otherwise, language transfer outputs may inherit the accent of the reference clip’s language. To mitigate this, set `cfg_weight` to `0`.
- The default settings (`exaggeration=0.5`, `cfg_weight=0.5`) work well for most prompts across all languages.
- If the reference speaker has a fast speaking style, lowering `cfg_weight` to around `0.3` can improve pacing.
- **Expressive or Dramatic Speech:**
- Try lower `cfg_weight` values (e.g. `~0.3`) and increase `exaggeration` to around `0.7` or higher.
- Higher `exaggeration` tends to speed up speech; reducing `cfg_weight` helps compensate with slower, more deliberate pacing.
## Built-in PerTh Watermarking for Responsible AI
Every audio file generated by Chatterbox includes [Resemble AI's Perth (Perceptual Threshold) Watermarker](https://github.com/resemble-ai/perth) - imperceptible neural watermarks that survive MP3 compression, audio editing, and common manipulations while maintaining nearly 100% detection accuracy.
## Watermark extraction
You can look for the watermark using the following script.
```python
import perth
import librosa
AUDIO_PATH = "YOUR_FILE.wav"
# Load the watermarked audio
watermarked_audio, sr = librosa.load(AUDIO_PATH, sr=None)
# Initialize watermarker (same as used for embedding)
watermarker = perth.PerthImplicitWatermarker()
# Extract watermark
watermark = watermarker.get_watermark(watermarked_audio, sample_rate=sr)
print(f"Extracted watermark: {watermark}")
# Output: 0.0 (no watermark) or 1.0 (watermarked)
```
## Official Discord
👋 Join us on [Discord](https://discord.gg/rJq9cRJBJ6) and let's build something awesome together!
## Evaluation
Chatterbox Turbo was evaluated using Podonos, a platform for reproducible subjective speech evaluation.
We compared Chatterbox Turbo to competitive TTS systems using Podonos' standardized evaluation suite, focusing on overall preference, naturalness, and expressiveness.
Evaluation reports:
- [Chatterbox Turbo vs ElevenLabs Turbo v2.5](https://podonos.com/resembleai/chatterbox-turbo-vs-elevenlabs-turbo)
- [Chatterbox Turbo vs Cartesia Sonic 3](https://podonos.com/resembleai/chatterbox-turbo-vs-cartesia-sonic3)
- [Chatterbox Turbo vs VibeVoice 7B](https://podonos.com/resembleai/chatterbox-turbo-vs-vibevoice7b)
These evaluations were conducted under identical conditions and are publicly accessible via Podonos.
## Acknowledgements
- [Podonos](https://podonos.com) — for supporting reproducible subjective speech evaluation
- [Cosyvoice](https://github.com/FunAudioLLM/CosyVoice)
- [Real-Time-Voice-Cloning](https://github.com/CorentinJ/Real-Time-Voice-Cloning)
- [HiFT-GAN](https://github.com/yl4579/HiFTNet)
- [Llama 3](https://github.com/meta-llama/llama3)
- [S3Tokenizer](https://github.com/xingchensong/S3Tokenizer)
## Citation
If you find this model useful, please consider citing.
```
@misc{chatterboxtts2025,
author = {{Resemble AI}},
title = {{Chatterbox-TTS}},
year = {2025},
howpublished = {\url{https://github.com/resemble-ai/chatterbox}},
note = {GitHub repository}
}
```
## Disclaimer
Don't use this model to do bad things. Prompts are sourced from freely available data on the internet.
================================================
FILE: example_for_mac.py
================================================
import torch
import torchaudio as ta
from chatterbox.tts import ChatterboxTTS
# Detect device (Mac with M1/M2/M3/M4)
device = "mps" if torch.backends.mps.is_available() else "cpu"
map_location = torch.device(device)
torch_load_original = torch.load
def patched_torch_load(*args, **kwargs):
if 'map_location' not in kwargs:
kwargs['map_location'] = map_location
return torch_load_original(*args, **kwargs)
torch.load = patched_torch_load
model = ChatterboxTTS.from_pretrained(device=device)
text = "Today is the day. I want to move like a titan at dawn, sweat like a god forging lightning. No more excuses. From now on, my mornings will be temples of discipline. I am going to work out like the gods… every damn day."
# If you want to synthesize with a different voice, specify the audio prompt
AUDIO_PROMPT_PATH = "YOUR_FILE.wav"
wav = model.generate(
text,
audio_prompt_path=AUDIO_PROMPT_PATH,
exaggeration=2.0,
cfg_weight=0.5
)
ta.save("test-2.wav", wav, model.sr)
================================================
FILE: example_tts.py
================================================
import torchaudio as ta
import torch
from pathlib import Path
from chatterbox.tts import ChatterboxTTS
from chatterbox.mtl_tts import ChatterboxMultilingualTTS
# Automatically detect the best available device
if torch.cuda.is_available():
device = "cuda"
elif torch.backends.mps.is_available():
device = "mps"
else:
device = "cpu"
print(f"Using device: {device}")
model = ChatterboxTTS.from_pretrained(device=device)
text = "Ezreal and Jinx teamed up with Ahri, Yasuo, and Teemo to take down the enemy's Nexus in an epic late-game pentakill."
wav = model.generate(text)
ta.save("test-1.wav", wav, model.sr)
multilingual_model = ChatterboxMultilingualTTS.from_pretrained(device=device)
text = "Bonjour, comment ça va? Ceci est le modèle de synthèse vocale multilingue Chatterbox, il prend en charge 23 langues."
wav = multilingual_model.generate(text, language_id="fr")
ta.save("test-2.wav", wav, multilingual_model.sr)
# If you want to synthesize with a different voice, specify the audio prompt
AUDIO_PROMPT_PATH = "YOUR_FILE.wav"
if Path(AUDIO_PROMPT_PATH).exists():
wav = model.generate(text, audio_prompt_path=AUDIO_PROMPT_PATH)
ta.save("test-3.wav", wav, model.sr)
else:
print(f"Warning: audio prompt file '{AUDIO_PROMPT_PATH}' not found, skipping voice cloning example.")
================================================
FILE: example_tts_turbo.py
================================================
import torchaudio as ta
import torch
from chatterbox.tts_turbo import ChatterboxTurboTTS
# Load the Turbo model
model = ChatterboxTurboTTS.from_pretrained(device="cuda")
# Generate with Paralinguistic Tags
text = "Oh, that's hilarious! [chuckle] Um anyway, we do have a new model in store. It's the SkyNet T-800 series and it's got basically everything. Including AI integration with ChatGPT and all that jazz. Would you like me to get some prices for you?"
# Generate audio (requires a reference clip for voice cloning)
# wav = model.generate(text, audio_prompt_path="your_10s_ref_clip.wav")
wav = model.generate(text)
ta.save("test-turbo.wav", wav, model.sr)
================================================
FILE: example_vc.py
================================================
import torch
import torchaudio as ta
from chatterbox.vc import ChatterboxVC
# Automatically detect the best available device
if torch.cuda.is_available():
device = "cuda"
elif torch.backends.mps.is_available():
device = "mps"
else:
device = "cpu"
print(f"Using device: {device}")
AUDIO_PATH = "YOUR_FILE.wav"
TARGET_VOICE_PATH = "YOUR_FILE.wav"
model = ChatterboxVC.from_pretrained(device)
wav = model.generate(
audio=AUDIO_PATH,
target_voice_path=TARGET_VOICE_PATH,
)
ta.save("testvc.wav", wav, model.sr)
================================================
FILE: gradio_tts_app.py
================================================
import random
import numpy as np
import torch
import gradio as gr
from chatterbox.tts import ChatterboxTTS
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
def set_seed(seed: int):
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
random.seed(seed)
np.random.seed(seed)
def load_model():
model = ChatterboxTTS.from_pretrained(DEVICE)
return model
def generate(model, text, audio_prompt_path, exaggeration, temperature, seed_num, cfgw, min_p, top_p, repetition_penalty):
if model is None:
model = ChatterboxTTS.from_pretrained(DEVICE)
if seed_num != 0:
set_seed(int(seed_num))
wav = model.generate(
text,
audio_prompt_path=audio_prompt_path,
exaggeration=exaggeration,
temperature=temperature,
cfg_weight=cfgw,
min_p=min_p,
top_p=top_p,
repetition_penalty=repetition_penalty,
)
return (model.sr, wav.squeeze(0).numpy())
with gr.Blocks() as demo:
model_state = gr.State(None) # Loaded once per session/user
with gr.Row():
with gr.Column():
text = gr.Textbox(
value="Now let's make my mum's favourite. So three mars bars into the pan. Then we add the tuna and just stir for a bit, just let the chocolate and fish infuse. A sprinkle of olive oil and some tomato ketchup. Now smell that. Oh boy this is going to be incredible.",
label="Text to synthesize (max chars 300)",
max_lines=5
)
ref_wav = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Reference Audio File", value=None)
exaggeration = gr.Slider(0.25, 2, step=.05, label="Exaggeration (Neutral = 0.5, extreme values can be unstable)", value=.5)
cfg_weight = gr.Slider(0.0, 1, step=.05, label="CFG/Pace", value=0.5)
with gr.Accordion("More options", open=False):
seed_num = gr.Number(value=0, label="Random seed (0 for random)")
temp = gr.Slider(0.05, 5, step=.05, label="temperature", value=.8)
min_p = gr.Slider(0.00, 1.00, step=0.01, label="min_p || Newer Sampler. Recommend 0.02 > 0.1. Handles Higher Temperatures better. 0.00 Disables", value=0.05)
top_p = gr.Slider(0.00, 1.00, step=0.01, label="top_p || Original Sampler. 1.0 Disables(recommended). Original 0.8", value=1.00)
repetition_penalty = gr.Slider(1.00, 2.00, step=0.1, label="repetition_penalty", value=1.2)
run_btn = gr.Button("Generate", variant="primary")
with gr.Column():
audio_output = gr.Audio(label="Output Audio")
demo.load(fn=load_model, inputs=[], outputs=model_state)
run_btn.click(
fn=generate,
inputs=[
model_state,
text,
ref_wav,
exaggeration,
temp,
seed_num,
cfg_weight,
min_p,
top_p,
repetition_penalty,
],
outputs=audio_output,
)
if __name__ == "__main__":
demo.queue(
max_size=50,
default_concurrency_limit=1,
).launch(share=True)
================================================
FILE: gradio_tts_turbo_app.py
================================================
import random
import numpy as np
import torch
import gradio as gr
from chatterbox.tts_turbo import ChatterboxTurboTTS
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
EVENT_TAGS = [
"[clear throat]", "[sigh]", "[shush]", "[cough]", "[groan]",
"[sniff]", "[gasp]", "[chuckle]", "[laugh]"
]
# --- REFINED CSS ---
# 1. tag-container: Forces the row to wrap items instead of scrolling. Removes borders/backgrounds.
# 2. tag-btn: Sets the specific look (indigo theme) and stops them from stretching.
CUSTOM_CSS = """
.tag-container {
display: flex !important;
flex-wrap: wrap !important; /* This fixes the one-per-line issue */
gap: 8px !important;
margin-top: 5px !important;
margin-bottom: 10px !important;
border: none !important;
background: transparent !important;
}
.tag-btn {
min-width: fit-content !important;
width: auto !important;
height: 32px !important;
font-size: 13px !important;
background: #eef2ff !important;
border: 1px solid #c7d2fe !important;
color: #3730a3 !important;
border-radius: 6px !important;
padding: 0 10px !important;
margin: 0 !important;
box-shadow: none !important;
}
.tag-btn:hover {
background: #c7d2fe !important;
transform: translateY(-1px);
}
"""
INSERT_TAG_JS = """
(tag_val, current_text) => {
const textarea = document.querySelector('#main_textbox textarea');
if (!textarea) return current_text + " " + tag_val;
const start = textarea.selectionStart;
const end = textarea.selectionEnd;
let prefix = " ";
let suffix = " ";
if (start === 0) prefix = "";
else if (current_text[start - 1] === ' ') prefix = "";
if (end < current_text.length && current_text[end] === ' ') suffix = "";
return current_text.slice(0, start) + prefix + tag_val + suffix + current_text.slice(end);
}
"""
def set_seed(seed: int):
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
random.seed(seed)
np.random.seed(seed)
def load_model():
print(f"Loading Chatterbox-Turbo on {DEVICE}...")
model = ChatterboxTurboTTS.from_pretrained(DEVICE)
return model
def generate(
model,
text,
audio_prompt_path,
temperature,
seed_num,
min_p,
top_p,
top_k,
repetition_penalty,
norm_loudness
):
if model is None:
model = ChatterboxTurboTTS.from_pretrained(DEVICE)
if seed_num != 0:
set_seed(int(seed_num))
wav = model.generate(
text,
audio_prompt_path=audio_prompt_path,
temperature=temperature,
min_p=min_p,
top_p=top_p,
top_k=int(top_k),
repetition_penalty=repetition_penalty,
norm_loudness=norm_loudness,
)
return (model.sr, wav.squeeze(0).numpy())
with gr.Blocks(title="Chatterbox Turbo", css=CUSTOM_CSS) as demo:
gr.Markdown("# ⚡ Chatterbox Turbo")
model_state = gr.State(None)
with gr.Row():
with gr.Column():
text = gr.Textbox(
value="Oh, that's hilarious! [chuckle] Um anyway, we do have a new model in store. It's the SkyNet T-800 series and it's got basically everything. Including AI integration with ChatGPT and um all that jazz. Would you like me to get some prices for you?",
label="Text to synthesize (max chars 300)",
max_lines=5,
elem_id="main_textbox"
)
# --- Event Tags ---
# Switched back to Row, but applied specific CSS to force wrapping
with gr.Row(elem_classes=["tag-container"]):
for tag in EVENT_TAGS:
# elem_classes targets the button specifically
btn = gr.Button(tag, elem_classes=["tag-btn"])
btn.click(
fn=None,
inputs=[btn, text],
outputs=text,
js=INSERT_TAG_JS
)
ref_wav = gr.Audio(
sources=["upload", "microphone"],
type="filepath",
label="Reference Audio File",
value="https://storage.googleapis.com/chatterbox-demo-samples/prompts/female_random_podcast.wav"
)
run_btn = gr.Button("Generate ⚡", variant="primary")
with gr.Column():
audio_output = gr.Audio(label="Output Audio")
with gr.Accordion("Advanced Options", open=False):
seed_num = gr.Number(value=0, label="Random seed (0 for random)")
temp = gr.Slider(0.05, 2.0, step=.05, label="Temperature", value=0.8)
top_p = gr.Slider(0.00, 1.00, step=0.01, label="Top P", value=0.95)
top_k = gr.Slider(0, 1000, step=10, label="Top K", value=1000)
repetition_penalty = gr.Slider(1.00, 2.00, step=0.05, label="Repetition Penalty", value=1.2)
min_p = gr.Slider(0.00, 1.00, step=0.01, label="Min P (Set to 0 to disable)", value=0.00)
norm_loudness = gr.Checkbox(value=True, label="Normalize Loudness (-27 LUFS)")
demo.load(fn=load_model, inputs=[], outputs=model_state)
run_btn.click(
fn=generate,
inputs=[
model_state,
text,
ref_wav,
temp,
seed_num,
min_p,
top_p,
top_k,
repetition_penalty,
norm_loudness,
],
outputs=audio_output,
)
if __name__ == "__main__":
demo.queue(
max_size=50,
default_concurrency_limit=1,
).launch(share=True)
================================================
FILE: gradio_vc_app.py
================================================
import torch
import gradio as gr
from chatterbox.vc import ChatterboxVC
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
model = ChatterboxVC.from_pretrained(DEVICE)
def generate(audio, target_voice_path):
wav = model.generate(
audio, target_voice_path=target_voice_path,
)
return model.sr, wav.squeeze(0).numpy()
demo = gr.Interface(
generate,
[
gr.Audio(sources=["upload", "microphone"], type="filepath", label="Input audio file"),
gr.Audio(sources=["upload", "microphone"], type="filepath", label="Target voice audio file (if none, the default voice is used)", value=None),
],
"audio",
)
if __name__ == "__main__":
demo.launch()
================================================
FILE: multilingual_app.py
================================================
import random
import numpy as np
import torch
from chatterbox.mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUAGES
import gradio as gr
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🚀 Running on device: {DEVICE}")
# --- Global Model Initialization ---
MODEL = None
LANGUAGE_CONFIG = {
"ar": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/ar_f/ar_prompts2.flac",
"text": "في الشهر الماضي، وصلنا إلى معلم جديد بمليارين من المشاهدات على قناتنا على يوتيوب."
},
"da": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/da_m1.flac",
"text": "Sidste måned nåede vi en ny milepæl med to milliarder visninger på vores YouTube-kanal."
},
"de": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/de_f1.flac",
"text": "Letzten Monat haben wir einen neuen Meilenstein erreicht: zwei Milliarden Aufrufe auf unserem YouTube-Kanal."
},
"el": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/el_m.flac",
"text": "Τον περασμένο μήνα, φτάσαμε σε ένα νέο ορόσημο με δύο δισεκατομμύρια προβολές στο κανάλι μας στο YouTube."
},
"en": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/en_f1.flac",
"text": "Last month, we reached a new milestone with two billion views on our YouTube channel."
},
"es": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/es_f1.flac",
"text": "El mes pasado alcanzamos un nuevo hito: dos mil millones de visualizaciones en nuestro canal de YouTube."
},
"fi": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/fi_m.flac",
"text": "Viime kuussa saavutimme uuden virstanpylvään kahden miljardin katselukerran kanssa YouTube-kanavallamme."
},
"fr": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/fr_f1.flac",
"text": "Le mois dernier, nous avons atteint un nouveau jalon avec deux milliards de vues sur notre chaîne YouTube."
},
"he": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/he_m1.flac",
"text": "בחודש שעבר הגענו לאבן דרך חדשה עם שני מיליארד צפיות בערוץ היוטיוב שלנו."
},
"hi": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/hi_f1.flac",
"text": "पिछले महीने हमने एक नया मील का पत्थर छुआ: हमारे YouTube चैनल पर दो अरब व्यूज़।"
},
"it": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/it_m1.flac",
"text": "Il mese scorso abbiamo raggiunto un nuovo traguardo: due miliardi di visualizzazioni sul nostro canale YouTube."
},
"ja": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/ja/ja_prompts1.flac",
"text": "先月、私たちのYouTubeチャンネルで二十億回の再生回数という新たなマイルストーンに到達しました。"
},
"ko": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/ko_f.flac",
"text": "지난달 우리는 유튜브 채널에서 이십억 조회수라는 새로운 이정표에 도달했습니다."
},
"ms": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/ms_f.flac",
"text": "Bulan lepas, kami mencapai pencapaian baru dengan dua bilion tontonan di saluran YouTube kami."
},
"nl": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/nl_m.flac",
"text": "Vorige maand bereikten we een nieuwe mijlpaal met twee miljard weergaven op ons YouTube-kanaal."
},
"no": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/no_f1.flac",
"text": "Forrige måned nådde vi en ny milepæl med to milliarder visninger på YouTube-kanalen vår."
},
"pl": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/pl_m.flac",
"text": "W zeszłym miesiącu osiągnęliśmy nowy kamień milowy z dwoma miliardami wyświetleń na naszym kanale YouTube."
},
"pt": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/pt_m1.flac",
"text": "No mês passado, alcançámos um novo marco: dois mil milhões de visualizações no nosso canal do YouTube."
},
"ru": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/ru_m.flac",
"text": "В прошлом месяце мы достигли нового рубежа: два миллиарда просмотров на нашем YouTube-канале."
},
"sv": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/sv_f.flac",
"text": "Förra månaden nådde vi en ny milstolpe med två miljarder visningar på vår YouTube-kanal."
},
"sw": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/sw_m.flac",
"text": "Mwezi uliopita, tulifika hatua mpya ya maoni ya bilioni mbili kweny kituo chetu cha YouTube."
},
"tr": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/tr_m.flac",
"text": "Geçen ay YouTube kanalımızda iki milyar görüntüleme ile yeni bir dönüm noktasına ulaştık."
},
"zh": {
"audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/zh_f2.flac",
"text": "上个月，我们达到了一个新的里程碑. 我们的YouTube频道观看次数达到了二十亿次，这绝对令人难以置信。"
},
}
# --- UI Helpers ---
def default_audio_for_ui(lang: str) -> str | None:
return LANGUAGE_CONFIG.get(lang, {}).get("audio")
def default_text_for_ui(lang: str) -> str:
return LANGUAGE_CONFIG.get(lang, {}).get("text", "")
def get_supported_languages_display() -> str:
"""Generate a formatted display of all supported languages."""
language_items = []
for code, name in sorted(SUPPORTED_LANGUAGES.items()):
language_items.append(f"**{name}** (`{code}`)")
# Split into 2 lines
mid = len(language_items) // 2
line1 = " • ".join(language_items[:mid])
line2 = " • ".join(language_items[mid:])
return f"""
### 🌍 Supported Languages ({len(SUPPORTED_LANGUAGES)} total)
{line1}
{line2}
"""
def get_or_load_model():
"""Loads the ChatterboxMultilingualTTS model if it hasn't been loaded already,
and ensures it's on the correct device."""
global MODEL
if MODEL is None:
print("Model not loaded, initializing...")
try:
MODEL = ChatterboxMultilingualTTS.from_pretrained(DEVICE)
if hasattr(MODEL, 'to') and str(MODEL.device) != DEVICE:
MODEL.to(DEVICE)
print(f"Model loaded successfully. Internal device: {getattr(MODEL, 'device', 'N/A')}")
except Exception as e:
print(f"Error loading model: {e}")
raise
return MODEL
# Attempt to load the model at startup.
try:
get_or_load_model()
except Exception as e:
print(f"CRITICAL: Failed to load model on startup. Application may not function. Error: {e}")
def set_seed(seed: int):
"""Sets the random seed for reproducibility across torch, numpy, and random."""
torch.manual_seed(seed)
if DEVICE == "cuda":
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
random.seed(seed)
np.random.seed(seed)
def resolve_audio_prompt(language_id: str, provided_path: str | None) -> str | None:
"""
Decide which audio prompt to use:
- If user provided a path (upload/mic/url), use it.
- Else, fall back to language-specific default (if any).
"""
if provided_path and str(provided_path).strip():
return provided_path
return LANGUAGE_CONFIG.get(language_id, {}).get("audio")
def generate_tts_audio(
text_input: str,
language_id: str,
audio_prompt_path_input: str = None,
exaggeration_input: float = 0.5,
temperature_input: float = 0.8,
seed_num_input: int = 0,
cfgw_input: float = 0.5
) -> tuple[int, np.ndarray]:
"""
Generate high-quality speech audio from text using Chatterbox Multilingual model with optional reference audio styling.
Supported languages: English, French, German, Spanish, Italian, Portuguese, and Hindi.
This tool synthesizes natural-sounding speech from input text. When a reference audio file
is provided, it captures the speaker's voice characteristics and speaking style. The generated audio
maintains the prosody, tone, and vocal qualities of the reference speaker, or uses default voice if no reference is provided.
Args:
text_input (str): The text to synthesize into speech (maximum 300 characters)
language_id (str): The language code for synthesis (eg. en, fr, de, es, it, pt, hi)
audio_prompt_path_input (str, optional): File path or URL to the reference audio file that defines the target voice style. Defaults to None.
exaggeration_input (float, optional): Controls speech expressiveness (0.25-2.0, neutral=0.5, extreme values may be unstable). Defaults to 0.5.
temperature_input (float, optional): Controls randomness in generation (0.05-5.0, higher=more varied). Defaults to 0.8.
seed_num_input (int, optional): Random seed for reproducible results (0 for random generation). Defaults to 0.
cfgw_input (float, optional): CFG/Pace weight controlling generation guidance (0.2-1.0). Defaults to 0.5, 0 for language transfer.
Returns:
tuple[int, np.ndarray]: A tuple containing the sample rate (int) and the generated audio waveform (numpy.ndarray)
"""
current_model = get_or_load_model()
if current_model is None:
raise RuntimeError("TTS model is not loaded.")
if seed_num_input != 0:
set_seed(int(seed_num_input))
print(f"Generating audio for text: '{text_input[:50]}...'")
# Handle optional audio prompt
chosen_prompt = audio_prompt_path_input or default_audio_for_ui(language_id)
generate_kwargs = {
"exaggeration": exaggeration_input,
"temperature": temperature_input,
"cfg_weight": cfgw_input,
}
if chosen_prompt:
generate_kwargs["audio_prompt_path"] = chosen_prompt
print(f"Using audio prompt: {chosen_prompt}")
else:
print("No audio prompt provided; using default voice.")
wav = current_model.generate(
text_input[:300], # Truncate text to max chars
language_id=language_id,
**generate_kwargs
)
print("Audio generation complete.")
return (current_model.sr, wav.squeeze(0).numpy())
with gr.Blocks() as demo:
gr.Markdown(
"""
# Chatterbox Multilingual Demo
Generate high-quality multilingual speech from text with reference audio styling, supporting 23 languages.
"""
)
# Display supported languages
gr.Markdown(get_supported_languages_display())
with gr.Row():
with gr.Column():
initial_lang = "fr"
text = gr.Textbox(
value=default_text_for_ui(initial_lang),
label="Text to synthesize (max chars 300)",
max_lines=5
)
language_id = gr.Dropdown(
choices=list(ChatterboxMultilingualTTS.get_supported_languages().keys()),
value=initial_lang,
label="Language",
info="Select the language for text-to-speech synthesis"
)
ref_wav = gr.Audio(
sources=["upload", "microphone"],
type="filepath",
label="Reference Audio File (Optional)",
value=default_audio_for_ui(initial_lang)
)
gr.Markdown(
"💡 **Note**: Ensure that the reference clip matches the specified language tag. Otherwise, language transfer outputs may inherit the accent of the reference clip's language. To mitigate this, set the CFG weight to 0.",
elem_classes=["audio-note"]
)
exaggeration = gr.Slider(
0.25, 2, step=.05, label="Exaggeration (Neutral = 0.5, extreme values can be unstable)", value=.5
)
cfg_weight = gr.Slider(
0.2, 1, step=.05, label="CFG/Pace", value=0.5
)
with gr.Accordion("More options", open=False):
seed_num = gr.Number(value=0, label="Random seed (0 for random)")
temp = gr.Slider(0.05, 5, step=.05, label="Temperature", value=.8)
run_btn = gr.Button("Generate", variant="primary")
with gr.Column():
audio_output = gr.Audio(label="Output Audio")
def on_language_change(lang, current_ref, current_text):
return default_audio_for_ui(lang), default_text_for_ui(lang)
language_id.change(
fn=on_language_change,
inputs=[language_id, ref_wav, text],
outputs=[ref_wav, text],
show_progress=False
)
run_btn.click(
fn=generate_tts_audio,
inputs=[
text,
language_id,
ref_wav,
exaggeration,
temp,
seed_num,
cfg_weight,
],
outputs=[audio_output],
)
demo.launch(mcp_server=True)
================================================
FILE: pyproject.toml
================================================
[project]
name = "chatterbox-tts"
version = "0.1.6"
description = "Chatterbox: Open Source TTS and Voice Conversion by Resemble AI"
readme = "README.md"
requires-python = ">=3.10"
license = {file = "LICENSE"}
authors = [
{name = "resemble-ai", email = "engineering@resemble.ai"}
]
dependencies = [
"numpy>=1.24.0,<1.26.0",
"librosa==0.11.0",
"s3tokenizer",
"torch==2.6.0",
"torchaudio==2.6.0",
"transformers==5.2.0",
"diffusers==0.29.0",
"resemble-perth @ git+https://github.com/resemble-ai/Perth.git@master",
"conformer==0.3.2",
"safetensors==0.5.3",
"spacy-pkuseg",
"pykakasi==2.3.0",
"gradio==6.8.0",
"pyloudnorm",
"omegaconf"
]
[project.urls]
Homepage = "https://github.com/resemble-ai/chatterbox"
Repository = "https://github.com/resemble-ai/chatterbox"
[build-system]
requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"
[tool.setuptools.packages.find]
where = ["src"]
================================================
FILE: src/chatterbox/__init__.py
================================================
try:
from importlib.metadata import version
except ImportError:
from importlib_metadata import version # For Python <3.8
__version__ = version("chatterbox-tts")
from .tts import ChatterboxTTS
from .vc import ChatterboxVC
from .mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUAGES
================================================
FILE: src/chatterbox/models/__init__.py
================================================
================================================
FILE: src/chatterbox/models/s3gen/__init__.py
================================================
from .s3gen import S3Token2Wav as S3Gen
from .const import S3GEN_SR
================================================
FILE: src/chatterbox/models/s3gen/configs.py
================================================
from ..utils import AttrDict
CFM_PARAMS = AttrDict({
"sigma_min": 1e-06,
"solver": "euler",
"t_scheduler": "cosine",
"training_cfg_rate": 0.2,
"inference_cfg_rate": 0.7,
"reg_loss_type": "l1"
})
================================================
FILE: src/chatterbox/models/s3gen/const.py
================================================
S3GEN_SR = 24000
S3GEN_SIL = 4299
================================================
FILE: src/chatterbox/models/s3gen/decoder.py
================================================
# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import pack, rearrange, repeat
from .utils.mask import add_optional_chunk_mask
from .matcha.decoder import SinusoidalPosEmb, Block1D, ResnetBlock1D, Downsample1D, \
TimestepEmbedding, Upsample1D
from .matcha.transformer import BasicTransformerBlock
from .utils.intmeanflow import get_intmeanflow_time_mixer
def mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
assert mask.dtype == torch.bool
assert dtype in [torch.float32, torch.bfloat16, torch.float16]
mask = mask.to(dtype)
# attention mask bias
# NOTE(Mddct): torch.finfo jit issues
# chunk_masks = (1.0 - chunk_masks) * torch.finfo(dtype).min
mask = (1.0 - mask) * -1.0e+10
return mask
class Transpose(torch.nn.Module):
def __init__(self, dim0: int, dim1: int):
super().__init__()
self.dim0 = dim0
self.dim1 = dim1
def forward(self, x: torch.Tensor):
x = torch.transpose(x, self.dim0, self.dim1)
return x
class CausalBlock1D(Block1D):
def __init__(self, dim: int, dim_out: int):
super(CausalBlock1D, self).__init__(dim, dim_out)
self.block = torch.nn.Sequential(
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):
output = self.block(x * mask)
return output * mask
class CausalResnetBlock1D(ResnetBlock1D):
def __init__(self, dim: int, dim_out: int, time_emb_dim: int, groups: int = 8):
super(CausalResnetBlock1D, self).__init__(dim, dim_out, time_emb_dim, groups)
self.block1 = CausalBlock1D(dim, dim_out)
self.block2 = CausalBlock1D(dim_out, dim_out)
class CausalConv1d(torch.nn.Conv1d):
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(CausalConv1d, self).__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(CausalConv1d, self).forward(x)
return x
class ConditionalDecoder(nn.Module):
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",
meanflow=False,
):
"""
This decoder requires an input with the same shape of the target. So, if your text content
is shorter or longer than the outputs, please re-sampling it before feeding to the decoder.
"""
super().__init__()
channels = tuple(channels)
self.meanflow = meanflow
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([])
# NOTE jrm: `static_chunk_size` is missing?
self.static_chunk_size = 0
output_channel = in_channels
for i in range(len(channels)): # pylint: disable=consider-using-enumerate
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) if self.causal else \
ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim)
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,
)
for _ in range(n_blocks)
]
)
downsample = (
Downsample1D(output_channel) if not is_last else
CausalConv1d(output_channel, output_channel, 3) if self.causal else nn.Conv1d(output_channel, output_channel, 3, padding=1)
)
self.down_blocks.append(nn.ModuleList([resnet, transformer_blocks, downsample]))
for _ in range(num_mid_blocks):
input_channel = channels[-1]
out_channels = channels[-1]
resnet = CausalResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) if self.causal else \
ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim)
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,
)
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,
) if self.causal else ResnetBlock1D(
dim=input_channel,
dim_out=output_channel,
time_emb_dim=time_embed_dim,
)
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,
)
for _ in range(n_blocks)
]
)
upsample = (
Upsample1D(output_channel, use_conv_transpose=True)
if not is_last
else CausalConv1d(output_channel, output_channel, 3) if self.causal else nn.Conv1d(output_channel, output_channel, 3, padding=1)
)
self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample]))
self.final_block = CausalBlock1D(channels[-1], channels[-1]) if self.causal else Block1D(channels[-1], channels[-1])
self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1)
self.initialize_weights()
self.time_embed_mixer = None
if self.meanflow:
self.time_embed_mixer = get_intmeanflow_time_mixer(time_embed_dim)
@property
def dtype(self):
return self.final_proj.weight.dtype
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.GroupNorm):
nn.init.constant_(m.weight, 1)
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, r=None):
"""Forward pass of the UNet1DConditional model.
Args:
x: (B, 80, T)
mask (_type_)
t (_type_): shape (batch_size)
spks (_type_, optional) Defaults to None.
cond (_type_, optional)
r: end time for meanflow mode (shape (1,) tensor)
Raises:
ValueError: _description_
ValueError: _description_
Returns:
_type_: _description_
"""
t = self.time_embeddings(t).to(t.dtype)
t = self.time_mlp(t)
if self.meanflow:
r = self.time_embeddings(r).to(t.dtype)
r = self.time_mlp(r)
concat_embed = torch.cat([t, r], dim=1)
t = self.time_embed_mixer(concat_embed)
x = pack([x, mu], "b * t")[0]
if spks is not None:
spks = repeat(spks, "b c -> b c t", t=x.shape[-1])
x = pack([x, spks], "b * t")[0]
if cond is not None:
x = pack([x, cond], "b * t")[0]
hiddens = []
masks = [mask]
for resnet, transformer_blocks, downsample in self.down_blocks:
mask_down = masks[-1]
x = resnet(x, mask_down, t)
x = rearrange(x, "b c t -> b t c").contiguous()
# attn_mask = torch.matmul(mask_down.transpose(1, 2).contiguous(), mask_down)
attn_mask = add_optional_chunk_mask(x, mask_down.bool(), False, False, 0, self.static_chunk_size, -1)
attn_mask = mask_to_bias(attn_mask == 1, x.dtype)
for transformer_block in transformer_blocks:
x = transformer_block(
hidden_states=x,
attention_mask=attn_mask,
timestep=t,
)
x = rearrange(x, "b t c -> b c t").contiguous()
hiddens.append(x) # Save hidden states for skip connections
x = downsample(x * mask_down)
masks.append(mask_down[:, :, ::2])
masks = masks[:-1]
mask_mid = masks[-1]
for resnet, transformer_blocks in self.mid_blocks:
x = resnet(x, mask_mid, t)
x = rearrange(x, "b c t -> b t c").contiguous()
# attn_mask = torch.matmul(mask_mid.transpose(1, 2).contiguous(), mask_mid)
attn_mask = add_optional_chunk_mask(x, mask_mid.bool(), False, False, 0, self.static_chunk_size, -1)
attn_mask = mask_to_bias(attn_mask == 1, x.dtype)
for transformer_block in transformer_blocks:
x = transformer_block(
hidden_states=x,
attention_mask=attn_mask,
timestep=t,
)
x = rearrange(x, "b t c -> b c t").contiguous()
for resnet, transformer_blocks, upsample in self.up_blocks:
mask_up = masks.pop()
skip = hiddens.pop()
x = pack([x[:, :, :skip.shape[-1]], skip], "b * t")[0]
x = resnet(x, mask_up, t)
x = rearrange(x, "b c t -> b t c").contiguous()
# attn_mask = torch.matmul(mask_up.transpose(1, 2).contiguous(), mask_up)
attn_mask = add_optional_chunk_mask(x, mask_up.bool(), False, False, 0, self.static_chunk_size, -1)
attn_mask = mask_to_bias(attn_mask == 1, x.dtype)
for transformer_block in transformer_blocks:
x = transformer_block(
hidden_states=x,
attention_mask=attn_mask,
timestep=t,
)
x = rearrange(x, "b t c -> b c t").contiguous()
x = upsample(x * mask_up)
x = self.final_block(x, mask_up)
output = self.final_proj(x * mask_up)
return output * mask
================================================
FILE: src/chatterbox/models/s3gen/f0_predictor.py
================================================
# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
import torch.nn as nn
from torch.nn.utils.parametrizations import weight_norm
class ConvRNNF0Predictor(nn.Module):
def __init__(self,
num_class: int = 1,
in_channels: int = 80,
cond_channels: int = 512
):
super().__init__()
self.num_class = num_class
self.condnet = nn.Sequential(
weight_norm(
nn.Conv1d(in_channels, cond_channels, kernel_size=3, padding=1)
),
nn.ELU(),
weight_norm(
nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)
),
nn.ELU(),
weight_norm(
nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)
),
nn.ELU(),
weight_norm(
nn.Conv1d(cond_channels, cond_channels, kernel_size=3, padding=1)
),
nn.ELU(),
weight_norm(
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:
x = self.condnet(x)
x = x.transpose(1, 2)
return torch.abs(self.classifier(x).squeeze(-1))
================================================
FILE: src/chatterbox/models/s3gen/flow.py
================================================
# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import random
from typing import Dict, Optional
logger = logging.getLogger(__name__)
import torch
import torch.nn as nn
from torch.nn import functional as F
from .utils.mask import make_pad_mask
from .configs import CFM_PARAMS
from omegaconf import DictConfig
logger = logging.getLogger(__name__)
def _repeat_batch_dim(tnsr, B, ndim):
"repeat batch dimension if it's equal to 1"
if tnsr is not None:
# add missing batch dim if needed
while tnsr.ndim < ndim:
tnsr = tnsr[None]
# repeat batch dim as needed
if B > 1 and tnsr.size(0) == 1:
tnsr = tnsr.repeat(B, *([1] * (ndim - 1)))
assert tnsr.ndim == ndim, f"Expected {ndim=}, got {tnsr.ndim=}"
return tnsr
class CausalMaskedDiffWithXvec(torch.nn.Module):
def __init__(self,
input_size: int = 512,
output_size: int = 80,
spk_embed_dim: int = 192,
output_type: str = "mel",
vocab_size: int = 6561,
input_frame_rate: int = 25,
only_mask_loss: bool = True,
token_mel_ratio: int = 2,
pre_lookahead_len: int = 3,
encoder: torch.nn.Module = None,
decoder: torch.nn.Module = None,
decoder_conf: Dict = {'in_channels': 240, 'out_channel': 80, 'spk_emb_dim': 80, 'n_spks': 1,
'cfm_params': DictConfig(
{'sigma_min': 1e-06, 'solver': 'euler', 't_scheduler': 'cosine',
'training_cfg_rate': 0.2, 'inference_cfg_rate': 0.7,
'reg_loss_type': 'l1'}),
'decoder_params': {'channels': [256, 256], 'dropout': 0.0,
'attention_head_dim': 64,
'n_blocks': 4, 'num_mid_blocks': 12, 'num_heads': 8,
'act_fn': 'gelu'}},
mel_feat_conf: Dict = {'n_fft': 1024, 'num_mels': 80, 'sampling_rate': 22050,
'hop_size': 256, 'win_size': 1024, 'fmin': 0, 'fmax': 8000}):
super().__init__()
self.input_size = input_size
self.output_size = output_size
self.decoder_conf = decoder_conf
self.mel_feat_conf = mel_feat_conf
self.vocab_size = vocab_size
self.output_type = output_type
self.input_frame_rate = input_frame_rate
logging.info(f"input frame rate={self.input_frame_rate}")
self.input_embedding = nn.Embedding(vocab_size, input_size)
self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size)
self.encoder = encoder
self.encoder_proj = torch.nn.Linear(self.encoder.output_size(), output_size)
self.decoder = decoder
self.only_mask_loss = only_mask_loss
self.token_mel_ratio = token_mel_ratio
self.pre_lookahead_len = pre_lookahead_len
# NOTE: copied in from cosyvoice repo
def compute_loss(
self,
batch: dict,
device: torch.device,
) -> Dict[str, Optional[torch.Tensor]]:
token = batch['speech_token'].to(device)
token_len = batch['speech_token_len'].to(device)
feat = batch['speech_feat'].to(device) # (B, 80, T)
feat_len = batch['speech_feat_len'].to(device)
embedding = batch['embedding'].to(device)
# NOTE unified training, static_chunk_size > 0 or = 0
# streaming = True if random.random() < 0.5 else False
# xvec projection
embedding = F.normalize(embedding, dim=1)
embedding = self.spk_embed_affine_layer(embedding)
# concat text and prompt_text
mask = (~make_pad_mask(token_len)).float().unsqueeze(-1).to(device) # (B, T, 1)
token = self.input_embedding(torch.clamp(token, min=0)) * mask # (B, T, emb)
# text encode
h, h_lengths = self.encoder(token, token_len) # (B, T, C) -> (B, 2T, C)
h = self.encoder_proj(h)
# get conditions
conds = torch.zeros(feat.shape, device=token.device)
for i, j in enumerate(feat_len):
if random.random() < 0.5:
continue
index = random.randint(0, int(0.3 * j))
conds[i, :, :index] = feat[i, :, :index]
mask = (~make_pad_mask(h_lengths.sum(dim=-1).squeeze(dim=1))).to(h)
loss, _ = self.decoder.compute_loss(
feat.contiguous(),
mask.unsqueeze(1),
h.transpose(1, 2).contiguous(),
embedding,
cond=conds,
# streaming=streaming,
)
return {'loss': loss}
@torch.inference_mode()
def inference(self,
token,
token_len,
prompt_token,
prompt_token_len,
prompt_feat,
prompt_feat_len,
embedding,
finalize,
n_timesteps=10,
noised_mels=None,
meanflow=False):
# token: (B, n_toks)
# token_len: (B,)
B = token.size(0)
# xvec projection
embedding = torch.atleast_2d(embedding)
embedding = F.normalize(embedding, dim=1)
embedding = self.spk_embed_affine_layer(embedding) # (1 or B, emb_dim)
# adjust shapes (batching logic)
prompt_token = _repeat_batch_dim(prompt_token, B, ndim=2) # (B, n_prompt)
prompt_token_len = _repeat_batch_dim(prompt_token_len, B, ndim=1) # (B,)
prompt_feat = _repeat_batch_dim(prompt_feat, B, ndim=3) # (B, n_feat, feat_dim=80)
prompt_feat_len = _repeat_batch_dim(prompt_feat_len, B, ndim=1) # (B,) or None
embedding = _repeat_batch_dim(embedding, B, ndim=2) # (B, emb_dim)
# concat text and prompt_text
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)
if (token >= self.vocab_size).any():
logger.error(f"{token.max()}>{self.vocab_size}\n out-of-range special tokens found in flow, fix inputs!")
token = self.input_embedding(token.long()) * mask
# text encode
h, h_masks = self.encoder(token, token_len)
if finalize is False:
h = h[:, :-self.pre_lookahead_len * self.token_mel_ratio]
h_lengths = h_masks.sum(dim=-1).squeeze(dim=-1)
mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1]
h = self.encoder_proj(h)
# # get conditions
conds = torch.zeros([B, 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(h_lengths)).unsqueeze(1).to(h)
if mask.shape[0] != B:
mask = mask.repeat(B, 1, 1)
feat, _ = self.decoder(
mu=h.transpose(1, 2).contiguous(),
mask=mask,
spks=embedding,
cond=conds,
n_timesteps=n_timesteps,
noised_mels=noised_mels,
meanflow=meanflow,
)
feat = feat[:, :, mel_len1:]
assert feat.shape[2] == mel_len2
return feat, None # NOTE jrm: why are they returning None here?
================================================
FILE: src/chatterbox/models/s3gen/flow_matching.py
================================================
# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import threading
import torch
import torch.nn.functional as F
from .matcha.flow_matching import BASECFM
from .configs import CFM_PARAMS
from tqdm import tqdm
def cast_all(*args, dtype):
return [a if (not a.dtype.is_floating_point) or a.dtype == dtype else a.to(dtype) for a in args]
class ConditionalCFM(BASECFM):
def __init__(self, in_channels, cfm_params, n_spks=1, spk_emb_dim=64, estimator: torch.nn.Module = None):
super().__init__(
n_feats=in_channels,
cfm_params=cfm_params,
n_spks=n_spks,
spk_emb_dim=spk_emb_dim,
)
self.t_scheduler = cfm_params.t_scheduler
self.training_cfg_rate = cfm_params.training_cfg_rate
self.inference_cfg_rate = cfm_params.inference_cfg_rate
in_channels = in_channels + (spk_emb_dim if n_spks > 0 else 0)
# Just change the architecture of the estimator here
self.estimator = estimator
@torch.inference_mode()
def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, prompt_len=0, flow_cache=torch.zeros(1, 80, 0, 2)):
"""Forward diffusion
Args:
mu (torch.Tensor): output of encoder
shape: (batch_size, n_feats, mel_timesteps)
mask (torch.Tensor): output_mask
shape: (batch_size, 1, mel_timesteps)
n_timesteps (int): number of diffusion steps
temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
spks (torch.Tensor, optional): speaker ids. Defaults to None.
shape: (batch_size, spk_emb_dim)
cond: Not used but kept for future purposes
Returns:
sample: generated mel-spectrogram
shape: (batch_size, n_feats, mel_timesteps)
"""
raise NotImplementedError("unused, needs updating for meanflow model")
z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature
cache_size = flow_cache.shape[2]
# fix prompt and overlap part mu and z
if cache_size != 0:
z[:, :, :cache_size] = flow_cache[:, :, :, 0]
mu[:, :, :cache_size] = flow_cache[:, :, :, 1]
z_cache = torch.concat([z[:, :, :prompt_len], z[:, :, -34:]], dim=2)
mu_cache = torch.concat([mu[:, :, :prompt_len], mu[:, :, -34:]], dim=2)
flow_cache = torch.stack([z_cache, mu_cache], dim=-1)
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), flow_cache
def solve_euler(self, x, t_span, mu, mask, spks, cond, meanflow=False):
"""
Fixed euler solver for ODEs.
Args:
x (torch.Tensor): random noise
t_span (torch.Tensor): n_timesteps interpolated
shape: (n_timesteps + 1,)
mu (torch.Tensor): output of encoder
shape: (batch_size, n_feats, mel_timesteps)
mask (torch.Tensor): output_mask
shape: (batch_size, 1, mel_timesteps)
spks (torch.Tensor, optional): speaker ids. Defaults to None.
shape: (batch_size, spk_emb_dim)
cond: Not used but kept for future purposes
meanflow: meanflow mode
"""
in_dtype = x.dtype
x, t_span, mu, mask, spks, cond = cast_all(x, t_span, mu, mask, spks, cond, dtype=self.estimator.dtype)
# Duplicated batch dims are for CFG
# Do not use concat, it may cause memory format changed and trt infer with wrong results!
B, T = mu.size(0), x.size(2)
x_in = torch.zeros([2 * B, 80, T], device=x.device, dtype=x.dtype)
mask_in = torch.zeros([2 * B, 1, T], device=x.device, dtype=x.dtype)
mu_in = torch.zeros([2 * B, 80, T], device=x.device, dtype=x.dtype)
t_in = torch.zeros([2 * B ], device=x.device, dtype=x.dtype)
spks_in = torch.zeros([2 * B, 80 ], device=x.device, dtype=x.dtype)
cond_in = torch.zeros([2 * B, 80, T], device=x.device, dtype=x.dtype)
r_in = torch.zeros([2 * B ], device=x.device, dtype=x.dtype) # (only used for meanflow)
for t, r in zip(t_span[:-1], t_span[1:]):
t = t.unsqueeze(dim=0)
r = r.unsqueeze(dim=0)
# Shapes:
# x_in ( 2B, 80, T )
# mask_in ( 2B, 1, T )
# mu_in ( 2B, 80, T )
# t_in ( 2B, )
# spks_in ( 2B, 80, )
# cond_in ( 2B, 80, T )
# r_in ( 2B, )
# x ( B, 80, T )
# mask ( B, 1, T )
# mu ( B, 80, T )
# t ( B, )
# spks ( B, 80, )
# cond ( B, 80, T )
# r ( B, )
x_in[:B] = x_in[B:] = x
mask_in[:B] = mask_in[B:] = mask
mu_in[:B] = mu
t_in[:B] = t_in[B:] = t
spks_in[:B] = spks
cond_in[:B] = cond
r_in[:B] = r_in[B:] = r # (only used for meanflow)
dxdt = self.estimator.forward(
x=x_in, mask=mask_in, mu=mu_in, t=t_in, spks=spks_in, cond=cond_in,
r=r_in if meanflow else None,
)
dxdt, cfg_dxdt = torch.split(dxdt, [B, B], dim=0)
dxdt = ((1.0 + self.inference_cfg_rate) * dxdt - self.inference_cfg_rate * cfg_dxdt)
dt = r - t
x = x + dt * dxdt
return x.to(in_dtype)
def compute_loss(self, x1, mask, mu, spks=None, cond=None):
"""Computes diffusion loss
Args:
x1 (torch.Tensor): Target
shape: (batch_size, n_feats, mel_timesteps)
mask (torch.Tensor): target mask
shape: (batch_size, 1, mel_timesteps)
mu (torch.Tensor): output of encoder
shape: (batch_size, n_feats, mel_timesteps)
spks (torch.Tensor, optional): speaker embedding. Defaults to None.
shape: (batch_size, spk_emb_dim)
Returns:
loss: conditional flow matching loss
y: conditional flow
shape: (batch_size, n_feats, mel_timesteps)
"""
b, _, t = mu.shape
# random timestep
t = torch.rand([b, 1, 1], device=mu.device, dtype=mu.dtype)
if self.t_scheduler == 'cosine':
t = 1 - torch.cos(t * 0.5 * torch.pi)
# sample noise p(x_0)
z = torch.randn_like(x1)
y = (1 - (1 - self.sigma_min) * t) * z + t * x1
u = x1 - (1 - self.sigma_min) * z
# during training, we randomly drop condition to trade off mode coverage and sample fidelity
if self.training_cfg_rate > 0:
cfg_mask = torch.rand(b, device=x1.device) > self.training_cfg_rate
mu = mu * cfg_mask.view(-1, 1, 1)
spks = spks * cfg_mask.view(-1, 1)
cond = cond * cfg_mask.view(-1, 1, 1)
pred = self.estimator(y, mask, mu, t.squeeze(), spks, cond)
loss = F.mse_loss(pred * mask, u * mask, reduction="sum") / (torch.sum(mask) * u.shape[1])
return loss, y
class CausalConditionalCFM(ConditionalCFM):
def __init__(self, in_channels=240, cfm_params=CFM_PARAMS, n_spks=1, spk_emb_dim=80, estimator=None):
super().__init__(in_channels, cfm_params, n_spks, spk_emb_dim, estimator)
# TODO: BAD BAD IDEA - IT'LL MESS UP DISTILLATION - SETTING TO NONE
self.rand_noise = None
@torch.inference_mode()
def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, noised_mels=None, meanflow=False):
"""Forward diffusion
Args:
mu (torch.Tensor): output of encoder
shape: (batch_size, n_feats, mel_timesteps)
mask (torch.Tensor): output_mask
shape: (batch_size, 1, mel_timesteps)
n_timesteps (int): number of diffusion steps
temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
spks (torch.Tensor, optional): speaker ids. Defaults to None.
shape: (batch_size, spk_emb_dim)
cond: Not used but kept for future purposes
noised_mels: gt mels noised a time t
Returns:
sample: generated mel-spectrogram
shape: (batch_size, n_feats, mel_timesteps)
"""
B = mu.size(0)
z = torch.randn_like(mu)
if noised_mels is not None:
prompt_len = mu.size(2) - noised_mels.size(2)
z[..., prompt_len:] = noised_mels
# time steps for reverse diffusion
t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
if (not meanflow) and (self.t_scheduler == 'cosine'):
t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
# NOTE: right now, the only meanflow models are also distilled models, which don't need CFG
# because they were distilled with CFG outputs. We would need to add another hparam and
# change the conditional logic here if we want to use CFG inference with a meanflow model.
if meanflow:
return self.basic_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None
return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond, meanflow=meanflow), None
def basic_euler(self, x, t_span, mu, mask, spks, cond):
in_dtype = x.dtype
x, t_span, mu, mask, spks, cond = cast_all(x, t_span, mu, mask, spks, cond, dtype=self.estimator.dtype)
print("S3 Token -> Mel Inference...")
for t, r in tqdm(zip(t_span[..., :-1], t_span[..., 1:]), total=t_span.shape[-1] - 1):
t, r = t[None], r[None]
dxdt = self.estimator.forward(x, mask=mask, mu=mu, t=t, spks=spks, cond=cond, r=r)
dt = r - t
x = x + dt * dxdt
return x.to(in_dtype)
================================================
FILE: src/chatterbox/models/s3gen/hifigan.py
================================================
# jrm: adapted from CosyVoice/cosyvoice/hifigan/generator.py
# most modules should be reusable, but I found their SineGen changed a git.
# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu)
#
# 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.
"""HIFI-GAN"""
from typing import Dict, Optional, List
import numpy as np
from scipy.signal import get_window
import torch
import torch.nn.functional as F
from torch.nn import Conv1d
from torch.nn import ConvTranspose1d
from torch.nn.utils import remove_weight_norm
from torch.nn.utils.parametrizations import weight_norm
from torch.distributions.uniform import Uniform
from torch import nn, sin, pow
from torch.nn import Parameter
class Snake(nn.Module):
'''
Implementation of a sine-based periodic activation function
Shape:
- Input: (B, C, T)
- Output: (B, C, T), same shape as the input
Parameters:
- alpha - trainable parameter
References:
- This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
https://arxiv.org/abs/2006.08195
Examples:
>>> a1 = snake(256)
>>> x = torch.randn(256)
>>> x = a1(x)
'''
def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
'''
Initialization.
INPUT:
- in_features: shape of the input
- alpha: trainable parameter
alpha is initialized to 1 by default, higher values = higher-frequency.
alpha will be trained along with the rest of your model.
'''
super(Snake, self).__init__()
self.in_features = in_features
# initialize alpha
self.alpha_logscale = alpha_logscale
if self.alpha_logscale: # log scale alphas initialized to zeros
self.alpha = Parameter(torch.zeros(in_features) * alpha)
else: # linear scale alphas initialized to ones
self.alpha = Parameter(torch.ones(in_features) * alpha)
self.alpha.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.
Snake ∶= x + 1/a * sin^2 (xa)
'''
alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
if self.alpha_logscale:
alpha = torch.exp(alpha)
x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
return x
def get_padding(kernel_size, dilation=1):
return int((kernel_size * dilation - dilation) / 2)
def init_weights(m, mean=0.0, std=0.01):
classname = m.__class__.__name__
if classname.find("Conv") != -1:
m.weight.data.normal_(mean, std)
"""hifigan based generator implementation.
This code is modified from https://github.com/jik876/hifi-gan
,https://github.com/kan-bayashi/ParallelWaveGAN and
https://github.com/NVIDIA/BigVGAN
"""
class ResBlock(torch.nn.Module):
"""Residual block module in HiFiGAN/BigVGAN."""
def __init__(
self,
channels: int = 512,
kernel_size: int = 3,
dilations: List[int] = [1, 3, 5],
):
super(ResBlock, self).__init__()
self.convs1 = nn.ModuleList()
self.convs2 = nn.ModuleList()
for dilation in dilations:
self.convs1.append(
weight_norm(
Conv1d(
channels,
channels,
kernel_size,
1,
dilation=dilation,
padding=get_padding(kernel_size, dilation)
)
)
)
self.convs2.append(
weight_norm(
Conv1d(
channels,
channels,
kernel_size,
1,
dilation=1,
padding=get_padding(kernel_size, 1)
)
)
)
self.convs1.apply(init_weights)
self.convs2.apply(init_weights)
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 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(torch.nn.Module):
""" Definition of sine generator
SineGen(samp_rate, harmonic_num = 0,
sine_amp = 0.1, noise_std = 0.003,
voiced_threshold = 0,
flag_for_pulse=False)
samp_rate: sampling rate in Hz
harmonic_num: number of harmonic overtones (default 0)
sine_amp: amplitude of sine-wavefrom (default 0.1)
noise_std: std of Gaussian noise (default 0.003)
voiced_thoreshold: F0 threshold for U/V classification (default 0)
flag_for_pulse: this SinGen is used inside PulseGen (default False)
Note: when flag_for_pulse is True, the first time step of a voiced
segment is always sin(np.pi) or cos(0)
"""
def __init__(self, samp_rate, harmonic_num=0,
sine_amp=0.1, noise_std=0.003,
voiced_threshold=0):
super(SineGen, self).__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):
# generate uv signal
uv = (f0 > self.voiced_threshold).type(torch.float32)
return uv
@torch.no_grad()
def forward(self, f0):
"""
:param f0: [B, 1, sample_len], Hz
:return: [B, 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(torch.nn.Module):
""" SourceModule for hn-nsf
SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
add_noise_std=0.003, voiced_threshod=0)
sampling_rate: sampling_rate in Hz
harmonic_num: number of harmonic 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)
note that amplitude of noise in unvoiced is decided
by sine_amp
voiced_threshold: threhold to set U/V given F0 (default: 0)
Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
F0_sampled (batchsize, length, 1)
Sine_source (batchsize, length, 1)
noise_source (batchsize, length 1)
uv (batchsize, length, 1)
"""
def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1,
add_noise_std=0.003, voiced_threshod=0):
super(SourceModuleHnNSF, self).__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_threshod)
# to merge source harmonics into a single excitation
self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
self.l_tanh = torch.nn.Tanh()
def forward(self, x):
"""
Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
F0_sampled (batchsize, length, 1)
Sine_source (batchsize, length, 1)
noise_source (batchsize, length 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 HiFTGenerator(nn.Module):
"""
HiFTNet Generator: Neural Source Filter + ISTFTNet
https://arxiv.org/abs/2309.09493
"""
def __init__(
self,
in_channels: int = 80,
base_channels: int = 512,
nb_harmonics: int = 8,
sampling_rate: int = 22050,
nsf_alpha: float = 0.1,
nsf_sigma: float = 0.003,
nsf_voiced_threshold: float = 10,
upsample_rates: List[int] = [8, 8],
upsample_kernel_sizes: List[int] = [16, 16],
istft_params: Dict[str, int] = {"n_fft": 16, "hop_len": 4},
resblock_kernel_sizes: List[int] = [3, 7, 11],
resblock_dilation_sizes: List[List[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
source_resblock_kernel_sizes: List[int] = [7, 11],
source_resblock_dilation_sizes: List[List[int]] = [[1, 3, 5], [1, 3, 5]],
lrelu_slope: float = 0.1,
audio_limit: float = 0.99,
f0_predictor: torch.nn.Module = None,
):
super(HiFTGenerator, self).__init__()
self.out_channels = 1
self.nb_harmonics = nb_harmonics
self.sampling_rate = sampling_rate
self.istft_params = istft_params
self.lrelu_slope = lrelu_slope
self.audio_limit = audio_limit
self.num_kernels = len(resblock_kernel_sizes)
self.num_upsamples = len(upsample_rates)
self.m_source = SourceModuleHnNSF(
sampling_rate=sampling_rate,
upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"],
harmonic_num=nb_harmonics,
sine_amp=nsf_alpha,
add_noise_std=nsf_sigma,
voiced_threshod=nsf_voiced_threshold)
self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * istft_params["hop_len"])
self.conv_pre = weight_norm(
Conv1d(in_channels, base_channels, 7, 1, padding=3)
)
# Up
self.ups = nn.ModuleList()
for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
self.ups.append(
weight_norm(
ConvTranspose1d(
base_channels // (2**i),
base_channels // (2**(i + 1)),
k,
u,
padding=(k - u) // 2,
)
)
)
# Down
self.source_downs = nn.ModuleList()
self.source_resblocks = nn.ModuleList()
downsample_rates = [1] + upsample_rates[::-1][:-1]
downsample_cum_rates = np.cumprod(downsample_rates)
for i, (u, k, d) in enumerate(zip(downsample_cum_rates[::-1], source_resblock_kernel_sizes, source_resblock_dilation_sizes)):
if u == 1:
self.source_downs.append(
Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), 1, 1)
)
else:
self.source_downs.append(
Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), u * 2, u, padding=(u // 2))
)
self.source_resblocks.append(
ResBlock(base_channels // (2 ** (i + 1)), k, d)
)
self.resblocks = nn.ModuleList()
for i in range(len(self.ups)):
ch = base_channels // (2**(i + 1))
for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
self.resblocks.append(ResBlock(ch, k, d))
self.conv_post = weight_norm(Conv1d(ch, istft_params["n_fft"] + 2, 7, 1, padding=3))
self.ups.apply(init_weights)
self.conv_post.apply(init_weights)
self.reflection_pad = nn.ReflectionPad1d((1, 0))
self.stft_window = torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32))
self.f0_predictor = f0_predictor
def remove_weight_norm(self):
print('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)
self.m_source.remove_weight_norm()
for l in self.source_downs:
remove_weight_norm(l)
for l in self.source_resblocks:
l.remove_weight_norm()
def _stft(self, x):
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, phase):
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 = torch.zeros(1, 1, 0)) -> torch.Tensor:
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
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:, :]) # actually, sin is redundancy
x = self._istft(magnitude, phase)
x = torch.clamp(x, -self.audio_limit, self.audio_limit)
return x
def forward(
self,
batch: dict,
device: torch.device,
) -> Dict[str, Optional[torch.Tensor]]:
speech_feat = batch['speech_feat'].transpose(1, 2).to(device)
# mel->f0
f0 = self.f0_predictor(speech_feat)
# f0->source
s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t
s, _, _ = self.m_source(s)
s = s.transpose(1, 2)
# mel+source->speech
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: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor:
# mel->f0
f0 = self.f0_predictor(speech_feat)
# f0->source
s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t
s, _, _ = self.m_source(s)
s = s.transpose(1, 2)
# use cache_source to avoid glitch
if cache_source.shape[2] != 0:
s[:, :, :cache_source.shape[2]] = cache_source
generated_speech = self.decode(x=speech_feat, s=s)
return generated_speech, s
================================================
FILE: src/chatterbox/models/s3gen/matcha/decoder.py
================================================
import math
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from conformer import ConformerBlock
from diffusers.models.activations import get_activation
from einops import pack, rearrange, repeat
from .transformer import BasicTransformerBlock
class SinusoidalPosEmb(torch.nn.Module):
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 Block1D(torch.nn.Module):
def __init__(self, dim, dim_out, groups=8):
super().__init__()
self.block = torch.nn.Sequential(
torch.nn.Conv1d(dim, dim_out, 3, padding=1),
torch.nn.GroupNorm(groups, dim_out),
nn.Mish(),
)
def forward(self, x, mask):
output = self.block(x * mask)
return output * mask
class ResnetBlock1D(torch.nn.Module):
def __init__(self, dim, dim_out, time_emb_dim, groups=8):
super().__init__()
self.mlp = torch.nn.Sequential(nn.Mish(), torch.nn.Linear(time_emb_dim, dim_out))
self.block1 = Block1D(dim, dim_out, groups=groups)
self.block2 = Block1D(dim_out, dim_out, groups=groups)
self.res_conv = torch.nn.Conv1d(dim, dim_out, 1)
def forward(self, x, mask, time_emb):
h = self.block1(x, mask)
h += self.mlp(time_emb).unsqueeze(-1)
h = self.block2(h, mask)
output = h + self.res_conv(x * mask)
return output
class Downsample1D(nn.Module):
def __init__(self, dim):
super().__init__()
self.conv = torch.nn.Conv1d(dim, dim, 3, 2, 1)
def forward(self, x):
return self.conv(x)
class TimestepEmbedding(nn.Module):
def __init__(
self,
in_channels: int,
time_embed_dim: int,
act_fn: str = "silu",
out_dim: int = None,
post_act_fn: Optional[str] = None,
cond_proj_dim=None,
):
super().__init__()
self.linear_1 = nn.Linear(in_channels, time_embed_dim)
if cond_proj_dim is not None:
self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)
else:
self.cond_proj = None
self.act = get_activation(act_fn)
if out_dim is not None:
time_embed_dim_out = out_dim
else:
time_embed_dim_out = time_embed_dim
self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out)
if post_act_fn is None:
self.post_act = None
else:
self.post_act = get_activation(post_act_fn)
def forward(self, sample, condition=None):
if condition is not None:
sample = sample + self.cond_proj(condition)
sample = self.linear_1(sample)
if self.act is not None:
sample = self.act(sample)
sample = self.linear_2(sample)
if self.post_act is not None:
sample = self.post_act(sample)
return sample
class Upsample1D(nn.Module):
"""A 1D upsampling layer with an optional convolution.
Parameters:
channels (`int`):
number of channels in the inputs and outputs.
use_conv (`bool`, default `False`):
option to use a convolution.
use_conv_transpose (`bool`, default `False`):
option to use a convolution transpose.
out_channels (`int`, optional):
number of output channels. Defaults to `channels`.
"""
def __init__(self, channels, use_conv=False, use_conv_transpose=True, out_channels=None, name="conv"):
super().__init__()
self.channels = channels
self.out_channels = out_channels or channels
self.use_conv = use_conv
self.use_conv_transpose = use_conv_transpose
self.name = name
self.conv = None
if use_conv_transpose:
self.conv = nn.ConvTranspose1d(channels, self.out_channels, 4, 2, 1)
elif use_conv:
self.conv = nn.Conv1d(self.channels, self.out_channels, 3, padding=1)
def forward(self, inputs):
assert inputs.shape[1] == self.channels
if self.use_conv_transpose:
return self.conv(inputs)
outputs = F.interpolate(inputs, scale_factor=2.0, mode="nearest")
if self.use_conv:
outputs = self.conv(outputs)
return outputs
class ConformerWrapper(ConformerBlock):
def __init__( # pylint: disable=useless-super-delegation
self,
*,
dim,
dim_head=64,
heads=8,
ff_mult=4,
conv_expansion_factor=2,
conv_kernel_size=31,
attn_dropout=0,
ff_dropout=0,
conv_dropout=0,
conv_causal=False,
):
super().__init__(
dim=dim,
dim_head=dim_head,
heads=heads,
ff_mult=ff_mult,
conv_expansion_factor=conv_expansion_factor,
conv_kernel_size=conv_kernel_size,
attn_dropout=attn_dropout,
ff_dropout=ff_dropout,
conv_dropout=conv_dropout,
conv_causal=conv_causal,
)
def forward(
self,
hidden_states,
attention_mask,
encoder_hidden_states=None,
encoder_attention_mask=None,
timestep=None,
):
return super().forward(x=hidden_states, mask=attention_mask.bool())
class Decoder(nn.Module):
def __init__(
self,
in_channels,
out_channels,
channels=(256, 256),
dropout=0.05,
attention_head_dim=64,
n_blocks=1,
num_mid_blocks=2,
num_heads=4,
act_fn="snake",
down_block_type="transformer",
mid_block_type="transformer",
up_block_type="transformer",
):
super().__init__()
channels = tuple(channels)
self.in_channels = in_channels
self.out_channels = out_channels
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([])
output_channel = in_channels
for i in range(len(channels)): # pylint: disable=consider-using-enumerate
input_channel = output_channel
output_channel = channels[i]
is_last = i == len(channels) - 1
resnet = ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim)
transformer_blocks = nn.ModuleList(
[
self.get_block(
down_block_type,
output_channel,
attention_head_dim,
num_heads,
dropout,
act_fn,
)
for _ in range(n_blocks)
]
)
downsample = (
Downsample1D(output_channel) if not is_last else nn.Conv1d(output_channel, output_channel, 3, padding=1)
)
self.down_blocks.append(nn.ModuleList([resnet, transformer_blocks, downsample]))
for i in range(num_mid_blocks):
input_channel = channels[-1]
out_channels = channels[-1]
resnet = ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim)
transformer_blocks = nn.ModuleList(
[
self.get_block(
mid_block_type,
output_channel,
attention_head_dim,
num_heads,
dropout,
act_fn,
)
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]
output_channel = channels[i + 1]
is_last = i == len(channels) - 2
resnet = ResnetBlock1D(
dim=2 * input_channel,
dim_out=output_channel,
time_emb_dim=time_embed_dim,
)
transformer_blocks = nn.ModuleList(
[
self.get_block(
up_block_type,
output_channel,
attention_head_dim,
num_heads,
dropout,
act_fn,
)
for _ in range(n_blocks)
]
)
upsample = (
Upsample1D(output_channel, use_conv_transpose=True)
if not is_last
else nn.Conv1d(output_channel, output_channel, 3, padding=1)
)
self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample]))
self.final_block = Block1D(channels[-1], channels[-1])
self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1)
self.initialize_weights()
# nn.init.normal_(self.final_proj.weight)
@staticmethod
def get_block(block_type, dim, attention_head_dim, num_heads, dropout, act_fn):
if block_type == "conformer":
block = ConformerWrapper(
dim=dim,
dim_head=attention_head_dim,
heads=num_heads,
ff_mult=1,
conv_expansion_factor=2,
ff_dropout=dropout,
attn_dropout=dropout,
conv_dropout=dropout,
conv_kernel_size=31,
)
elif block_type == "transformer":
block = BasicTransformerBlock(
dim=dim,
num_attention_heads=num_heads,
attention_head_dim=attention_head_dim,
dropout=dropout,
activation_fn=act_fn,
)
else:
raise ValueError(f"Unknown block type {block_type}")
return block
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.GroupNorm):
nn.init.constant_(m.weight, 1)
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):
"""Forward pass of the UNet1DConditional model.
Args:
x (torch.Tensor): shape (batch_size, in_channels, time)
mask (_type_): shape (batch_size, 1, time)
t (_type_): shape (batch_size)
spks (_type_, optional): shape: (batch_size, condition_channels). Defaults to None.
cond (_type_, optional): placeholder for future use. Defaults to None.
Raises:
ValueError: _description_
ValueError: _description_
Returns:
_type_: _description_
"""
t = self.time_embeddings(t)
t = self.time_mlp(t)
x = pack([x, mu], "b * t")[0]
if spks is not None:
spks = repeat(spks, "b c -> b c t", t=x.shape[-1])
x = pack([x, spks], "b * t")[0]
hiddens = []
masks = [mask]
for resnet, transformer_blocks, downsample in self.down_blocks:
mask_down = masks[-1]
x = resnet(x, mask_down, t)
x = rearrange(x, "b c t -> b t c")
mask_down = rearrange(mask_down, "b 1 t -> b t")
for transformer_block in transformer_blocks:
x = transformer_block(
hidden_states=x,
attention_mask=mask_down,
timestep=t,
)
x = rearrange(x, "b t c -> b c t")
mask_down = rearrange(mask_down, "b t -> b 1 t")
hiddens.append(x) # Save hidden states for skip connections
x = downsample(x * mask_down)
masks.append(mask_down[:, :, ::2])
masks = masks[:-1]
mask_mid = masks[-1]
for resnet, transformer_blocks in self.mid_blocks:
x = resnet(x, mask_mid, t)
x = rearrange(x, "b c t -> b t c")
mask_mid = rearrange(mask_mid, "b 1 t -> b t")
for transformer_block in transformer_blocks:
x = transformer_block(
hidden_states=x,
attention_mask=mask_mid,
timestep=t,
)
x = rearrange(x, "b t c -> b c t")
mask_mid = rearrange(mask_mid, "b t -> b 1 t")
for resnet, transformer_blocks, upsample in self.up_blocks:
mask_up = masks.pop()
x = resnet(pack([x, hiddens.pop()], "b * t")[0], mask_up, t)
x = rearrange(x, "b c t -> b t c")
mask_up = rearrange(mask_up, "b 1 t -> b t")
for transformer_block in transformer_blocks:
x = transformer_block(
hidden_states=x,
attention_mask=mask_up,
timestep=t,
)
x = rearrange(x, "b t c -> b c t")
mask_up = rearrange(mask_up, "b t -> b 1 t")
x = upsample(x * mask_up)
x = self.final_block(x, mask_up)
output = self.final_proj(x * mask_up)
return output * mask
================================================
FILE: src/chatterbox/models/s3gen/matcha/flow_matching.py
================================================
from abc import ABC
import torch
import torch.nn.functional as F
from .decoder import Decoder
class BASECFM(torch.nn.Module, ABC):
def __init__(
self,
n_feats,
cfm_params,
n_spks=1,
spk_emb_dim=128,
):
super().__init__()
self.n_feats = n_feats
self.n_spks = n_spks
self.spk_emb_dim = spk_emb_dim
self.solver = cfm_params.solver
if hasattr(cfm_params, "sigma_min"):
self.sigma_min = cfm_params.sigma_min
else:
self.sigma_min = 1e-4
self.estimator = None
@torch.inference_mode()
def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None):
"""Forward diffusion
Args:
mu (torch.Tensor): output of encoder
shape: (batch_size, n_feats, mel_timesteps)
mask (torch.Tensor): output_mask
shape: (batch_size, 1, mel_timesteps)
n_timesteps (int): number of diffusion steps
temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
spks (torch.Tensor, optional): speaker ids. Defaults to None.
shape: (batch_size, spk_emb_dim)
cond: Not used but kept for future purposes
Returns:
sample: generated mel-spectrogram
shape: (batch_size, n_feats, mel_timesteps)
"""
z = torch.randn_like(mu) * temperature
t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device)
return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond)
def solve_euler(self, x, t_span, mu, mask, spks, cond):
"""
Fixed euler solver for ODEs.
Args:
x (torch.Tensor): random noise
t_span (torch.Tensor): n_timesteps interpolated
shape: (n_timesteps + 1,)
mu (torch.Tensor): output of encoder
shape: (batch_size, n_feats, mel_timesteps)
mask (torch.Tensor): output_mask
shape: (batch_size, 1, mel_timesteps)
spks (torch.Tensor, optional): speaker ids. Defaults to None.
shape: (batch_size, spk_emb_dim)
cond: Not used but kept for future purposes
"""
t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
# I am storing this because I can later plot it by putting a debugger here and saving it to a file
# Or in future might add like a return_all_steps flag
sol = []
for step in range(1, len(t_span)):
dphi_dt = self.estimator(x, mask, mu, t, spks, cond)
x = x + dt * dphi_dt
t = t + dt
sol.append(x)
if step < len(t_span) - 1:
dt = t_span[step + 1] - t
return sol[-1]
def compute_loss(self, x1, mask, mu, spks=None, cond=None):
"""Computes diffusion loss
Args:
x1 (torch.Tensor): Target
shape: (batch_size, n_feats, mel_timesteps)
mask (torch.Tensor): target mask
shape: (batch_size, 1, mel_timesteps)
mu (torch.Tensor): output of encoder
shape: (batch_size, n_feats, mel_timesteps)
spks (torch.Tensor, optional): speaker embedding. Defaults to None.
shape: (batch_size, spk_emb_dim)
Returns:
loss: conditional flow matching loss
y: conditional flow
shape: (batch_size, n_feats, mel_timesteps)
"""
b, _, t = mu.shape
# random timestep
t = torch.rand([b, 1, 1], device=mu.device, dtype=mu.dtype)
# sample noise p(x_0)
z = torch.randn_like(x1)
y = (1 - (1 - self.sigma_min) * t) * z + t * x1
u = x1 - (1 - self.sigma_min) * z
loss = F.mse_loss(self.estimator(y, mask, mu, t.squeeze(), spks), u, reduction="sum") / (
torch.sum(mask) * u.shape[1]
)
return loss, y
class CFM(BASECFM):
def __init__(self, in_channels, out_channel, cfm_params, decoder_params, n_spks=1, spk_emb_dim=64):
super().__init__(
n_feats=in_channels,
cfm_params=cfm_params,
n_spks=n_spks,
spk_emb_dim=spk_emb_dim,
)
in_channels = in_channels + (spk_emb_dim if n_spks > 1 else 0)
# Just change the architecture of the estimator here
self.estimator = Decoder(in_channels=in_channels, out_channels=out_channel, **decoder_params)
================================================
FILE: src/chatterbox/models/s3gen/matcha/text_encoder.py
================================================
""" from https://github.com/jaywalnut310/glow-tts """
import math
import torch
import torch.nn as nn
from einops import rearrange
def sequence_mask(length, max_length=None):
if max_length is None:
max_length = length.max()
x = torch.arange(max_length, dtype=length.dtype, device=length.device)
return x.unsqueeze(0) < length.unsqueeze(1)
class LayerNorm(nn.Module):
def __init__(self, channels, eps=1e-4):
super().__init__()
self.channels = channels
self.eps = eps
self.gamma = torch.nn.Parameter(torch.ones(channels))
self.beta = torch.nn.Parameter(torch.zeros(channels))
def forward(self, x):
n_dims = len(x.shape)
mean = torch.mean(x, 1, keepdim=True)
variance = torch.mean((x - mean) ** 2, 1, keepdim=True)
x = (x - mean) * torch.rsqrt(variance + self.eps)
shape = [1, -1] + [1] * (n_dims - 2)
x = x * self.gamma.view(*shape) + self.beta.view(*shape)
return x
class ConvReluNorm(nn.Module):
def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
super().__init__()
self.in_channels = in_channels
self.hidden_channels = hidden_channels
self.out_channels = out_channels
self.kernel_size = kernel_size
self.n_layers = n_layers
self.p_dropout = p_dropout
self.conv_layers = torch.nn.ModuleList()
self.norm_layers = torch.nn.ModuleList()
self.conv_layers.append(torch.nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size // 2))
self.norm_layers.append(LayerNorm(hidden_channels))
self.relu_drop = torch.nn.Sequential(torch.nn.ReLU(), torch.nn.Dropout(p_dropout))
for _ in range(n_layers - 1):
self.conv_layers.append(
torch.nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size // 2)
)
self.norm_layers.append(LayerNorm(hidden_channels))
self.proj = torch.nn.Conv1d(hidden_channels, out_channels, 1)
self.proj.weight.data.zero_()
self.proj.bias.data.zero_()
def forward(self, x, x_mask):
x_org = x
for i in range(self.n_layers):
x = self.conv_layers[i](x * x_mask)
x = self.norm_layers[i](x)
x = self.relu_drop(x)
x = x_org + self.proj(x)
return x * x_mask
class DurationPredictor(nn.Module):
def __init__(self, in_channels, filter_channels, kernel_size, p_dropout):
super().__init__()
self.in_channels = in_channels
self.filter_channels = filter_channels
self.p_dropout = p_dropout
self.drop = torch.nn.Dropout(p_dropout)
self.conv_1 = torch.nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size // 2)
self.norm_1 = LayerNorm(filter_channels)
self.conv_2 = torch.nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size // 2)
self.norm_2 = LayerNorm(filter_channels)
self.proj = torch.nn.Conv1d(filter_channels, 1, 1)
def forward(self, x, x_mask):
x = self.conv_1(x * x_mask)
x = torch.relu(x)
x = self.norm_1(x)
x = self.drop(x)
x = self.conv_2(x * x_mask)
x = torch.relu(x)
x = self.norm_2(x)
x = self.drop(x)
x = self.proj(x * x_mask)
return x * x_mask
class RotaryPositionalEmbeddings(nn.Module):
"""
## RoPE module
Rotary encoding transforms pairs of features by rotating in the 2D plane.
That is, it organizes the $d$ features as $\frac{d}{2}$ pairs.
Each pair can be considered a coordinate in a 2D plane, and the encoding will rotate it
by an angle depending on the position of the token.
"""
def __init__(self, d: int, base: int = 10_000):
r"""
* `d` is the number of features $d$
* `base` is the constant used for calculating $\Theta$
"""
super().__init__()
self.base = base
self.d = int(d)
self.cos_cached = None
self.sin_cached = None
def _build_cache(self, x: torch.Tensor):
r"""
Cache $\cos$ and $\sin$ values
"""
# Return if cache is already built
if self.cos_cached is not None and x.shape[0] <= self.cos_cached.shape[0]:
return
# Get sequence length
seq_len = x.shape[0]
# $\Theta = {\theta_i = 10000^{-\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
theta = 1.0 / (self.base ** (torch.arange(0, self.d, 2).float() / self.d)).to(x.device)
# Create position indexes `[0, 1, ..., seq_len - 1]`
seq_idx = torch.arange(seq_len, device=x.device).float().to(x.device)
# Calculate the product of position index and $\theta_i$
idx_theta = torch.einsum("n,d->nd", seq_idx, theta)
# Concatenate so that for row $m$ we have
# $[m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}, m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}]$
idx_theta2 = torch.cat([idx_theta, idx_theta], dim=1)
# Cache them
self.cos_cached = idx_theta2.cos()[:, None, None, :]
self.sin_cached = idx_theta2.sin()[:, None, None, :]
def _neg_half(self, x: torch.Tensor):
# $\frac{d}{2}$
d_2 = self.d // 2
# Calculate $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., x^{(\frac{d}{2})}]$
return torch.cat([-x[:, :, :, d_2:], x[:, :, :, :d_2]], dim=-1)
def forward(self, x: torch.Tensor):
"""
* `x` is the Tensor at the head of a key or a query with shape `[seq_len, batch_size, n_heads, d]`
"""
# Cache $\cos$ and $\sin$ values
x = rearrange(x, "b h t d -> t b h d")
self._build_cache(x)
# Split the features, we can choose to apply rotary embeddings only to a partial set of features.
x_rope, x_pass = x[..., : self.d], x[..., self.d :]
# Calculate
# $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., x^{(\frac{d}{2})}]$
neg_half_x = self._neg_half(x_rope)
x_rope = (x_rope * self.cos_cached[: x.shape[0]]) + (neg_half_x * self.sin_cached[: x.shape[0]])
return rearrange(torch.cat((x_rope, x_pass), dim=-1), "t b h d -> b h t d")
class MultiHeadAttention(nn.Module):
def __init__(
self,
channels,
out_channels,
n_heads,
heads_share=True,
p_dropout=0.0,
proximal_bias=False,
proximal_init=False,
):
super().__init__()
assert channels % n_heads == 0
self.channels = channels
self.out_channels = out_channels
self.n_heads = n_heads
self.heads_share = heads_share
self.proximal_bias = proximal_bias
self.p_dropout = p_dropout
self.attn = None
self.k_channels = channels // n_heads
self.conv_q = torch.nn.Conv1d(channels, channels, 1)
self.conv_k = torch.nn.Conv1d(channels, channels, 1)
self.conv_v = torch.nn.Conv1d(channels, channels, 1)
# from https://nn.labml.ai/transformers/rope/index.html
self.query_rotary_pe = RotaryPositionalEmbeddings(self.k_channels * 0.5)
self.key_rotary_pe = RotaryPositionalEmbeddings(self.k_channels * 0.5)
self.conv_o = torch.nn.Conv1d(channels, out_channels, 1)
self.drop = torch.nn.Dropout(p_dropout)
torch.nn.init.xavier_uniform_(self.conv_q.weight)
torch.nn.init.xavier_uniform_(self.conv_k.weight)
if proximal_init:
self.conv_k.weight.data.copy_(self.conv_q.weight.data)
self.conv_k.bias.data.copy_(self.conv_q.bias.data)
torch.nn.init.xavier_uniform_(self.conv_v.weight)
def forward(self, x, c, attn_mask=None):
q = self.conv_q(x)
k = self.conv_k(c)
v = self.conv_v(c)
x, self.attn = self.attention(q, k, v, mask=attn_mask)
x = self.conv_o(x)
return x
def attention(self, query, key, value, mask=None):
b, d, t_s, t_t = (*key.size(), query.size(2))
query = rearrange(query, "b (h c) t-> b h t c", h=self.n_heads)
key = rearrange(key, "b (h c) t-> b h t c", h=self.n_heads)
value = rearrange(value, "b (h c) t-> b h t c", h=self.n_heads)
query = self.query_rotary_pe(query)
key = self.key_rotary_pe(key)
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.k_channels)
if self.proximal_bias:
assert t_s == t_t, "Proximal bias is only available for self-attention."
scores = scores + self._attention_bias_proximal(t_s).to(device=scores.device, dtype=scores.dtype)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e4)
p_attn = torch.nn.functional.softmax(scores, dim=-1)
p_attn = self.drop(p_attn)
output = torch.matmul(p_attn, value)
output = output.transpose(2, 3).contiguous().view(b, d, t_t)
return output, p_attn
@staticmethod
def _attention_bias_proximal(length):
r = torch.arange(length, dtype=torch.float32)
diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
class FFN(nn.Module):
def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0.0):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.filter_channels = filter_channels
self.kernel_size = kernel_size
self.p_dropout = p_dropout
self.conv_1 = torch.nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size // 2)
self.conv_2 = torch.nn.Conv1d(filter_channels, out_channels, kernel_size, padding=kernel_size // 2)
self.drop = torch.nn.Dropout(p_dropout)
def forward(self, x, x_mask):
x = self.conv_1(x * x_mask)
x = torch.relu(x)
x = self.drop(x)
x = self.conv_2(x * x_mask)
return x * x_mask
class Encoder(nn.Module):
def __init__(
self,
hidden_channels,
filter_channels,
n_heads,
n_layers,
kernel_size=1,
p_dropout=0.0,
**kwargs,
):
super().__init__()
self.hidden_channels = hidden_channels
self.filter_channels = filter_channels
self.n_heads = n_heads
self.n_layers = n_layers
self.kernel_size = kernel_size
self.p_dropout = p_dropout
self.drop = torch.nn.Dropout(p_dropout)
self.attn_layers = torch.nn.ModuleList()
self.norm_layers_1 = torch.nn.ModuleList()
self.ffn_layers = torch.nn.ModuleList()
self.norm_layers_2 = torch.nn.ModuleList()
for _ in range(self.n_layers):
self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout))
self.norm_layers_1.append(LayerNorm(hidden_channels))
self.ffn_layers.append(
FFN(
hidden_channels,
hidden_channels,
filter_channels,
kernel_size,
p_dropout=p_dropout,
)
)
self.norm_layers_2.append(LayerNorm(hidden_channels))
def forward(self, x, x_mask):
attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
for i in range(self.n_layers):
x = x * x_mask
y = self.attn_layers[i](x, x, attn_mask)
y = self.drop(y)
x = self.norm_layers_1[i](x + y)
y = self.ffn_layers[i](x, x_mask)
y = self.drop(y)
x = self.norm_layers_2[i](x + y)
x = x * x_mask
return x
class TextEncoder(nn.Module):
def __init__(
self,
encoder_type,
encoder_params,
duration_predictor_params,
n_vocab,
n_spks=1,
spk_emb_dim=128,
):
super().__init__()
self.encoder_type = encoder_type
self.n_vocab = n_vocab
self.n_feats = encoder_params.n_feats
self.n_channels = encoder_params.n_channels
self.spk_emb_dim = spk_emb_dim
self.n_spks = n_spks
self.emb = torch.nn.Embedding(n_vocab, self.n_channels)
torch.nn.init.normal_(self.emb.weight, 0.0, self.n_channels**-0.5)
if encoder_params.prenet:
self.prenet = ConvReluNorm(
self.n_channels,
self.n_channels,
self.n_channels,
kernel_size=5,
n_layers=3,
p_dropout=0.5,
)
else:
self.prenet = lambda x, x_mask: x
self.encoder = Encoder(
encoder_params.n_channels + (spk_emb_dim if n_spks > 1 else 0),
encoder_params.filter_channels,
encoder_params.n_heads,
encoder_params.n_layers,
encoder_params.kernel_size,
encoder_params.p_dropout,
)
self.proj_m = torch.nn.Conv1d(self.n_channels + (spk_emb_dim if n_spks > 1 else 0), self.n_feats, 1)
self.proj_w = DurationPredictor(
self.n_channels + (spk_emb_dim if n_spks > 1 else 0),
duration_predictor_params.filter_channels_dp,
duration_predictor_params.kernel_size,
duration_predictor_params.p_dropout,
)
def forward(self, x, x_lengths, spks=None):
"""Run forward pass to the transformer based encoder and duration predictor
Args:
x (torch.Tensor): text input
shape: (batch_size, max_text_length)
x_lengths (torch.Tensor): text input lengths
shape: (batch_size,)
spks (torch.Tensor, optional): speaker ids. Defaults to None.
shape: (batch_size,)
Returns:
mu (torch.Tensor): average output of the encoder
shape: (batch_size, n_feats, max_text_length)
logw (torch.Tensor): log duration predicted by the duration predictor
shape: (batch_size, 1, max_text_length)
x_mask (torch.Tensor): mask for the text input
shape: (batch_size, 1, max_text_length)
"""
x = self.emb(x) * math.sqrt(self.n_channels)
x = torch.transpose(x, 1, -1)
x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
x = self.prenet(x, x_mask)
if self.n_spks > 1:
x = torch.cat([x, spks.unsqueeze(-1).repeat(1, 1, x.shape[-1])], dim=1)
x = self.encoder(x, x_mask)
mu = self.proj_m(x) * x_mask
x_dp = torch.detach(x)
logw = self.proj_w(x_dp, x_mask)
return mu, logw, x_mask
================================================
FILE: src/chatterbox/models/s3gen/matcha/transformer.py
================================================
from typing import Any, Dict, Optional
import torch
import torch.nn as nn
from diffusers.models.attention import (
GEGLU,
GELU,
AdaLayerNorm,
AdaLayerNormZero,
ApproximateGELU,
)
from diffusers.models.attention_processor import Attention
from diffusers.models.lora import LoRACompatibleLinear
from diffusers.utils.torch_utils import maybe_allow_in_graph
class SnakeBeta(nn.Module):
"""
A modified Snake function which uses separate parameters for the magnitude of the periodic components
Shape:
- Input: (B, C, T)
- Output: (B, C, T), same shape as the input
Parameters:
- alpha - trainable parameter that controls frequency
- beta - trainable parameter that controls magnitude
References:
- This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
https://arxiv.org/abs/2006.08195
Examples:
>>> a1 = snakebeta(256)
>>> x = torch.randn(256)
>>> x = a1(x)
"""
def __init__(self, in_features, out_features, alpha=1.0, alpha_trainable=True, alpha_logscale=True):
"""
Initialization.
INPUT:
- in_features: shape of the input
- alpha - trainable parameter that controls frequency
- beta - trainable parameter that controls magnitude
alpha is initialized to 1 by default, higher values = higher-frequency.
beta is initialized to 1 by default, higher values = higher-magnitude.
alpha will be trained along with the rest of your model.
"""
super().__init__()
self.in_features = out_features if isinstance(out_features, list) else [out_features]
self.proj = LoRACompatibleLinear(in_features, out_features)
# initialize alpha
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
class FeedForward(nn.Module):
r"""
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)
if activation_fn == "gelu-approximate":
act_fn = GELU(dim, inner_dim, approximate="tanh")
elif activation_fn == "geglu":
act_fn = GEGLU(dim, inner_dim)
elif activation_fn == "geglu-approximate":
act_fn = ApproximateGELU(dim, inner_dim)
elif activation_fn == "snakebeta":
act_fn = SnakeBeta(dim, inner_dim)
self.net = nn.ModuleList([])
# project in
self.net.append(act_fn)
# project dropout
self.net.append(nn.Dropout(dropout))
# project out
self.net.append(LoRACompatibleLinear(inner_dim, dim_out))
# 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
@maybe_allow_in_graph
class BasicTransformerBlock(nn.Module):
r"""
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,
# scale_qk=False, # uncomment this to not to use flash attention
) # is self-attn if encoder_hidden_states is none
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: Dict[str, Any] = 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
================================================
FILE: src/chatterbox/models/s3gen/s3gen.py
================================================
# Modified from CosyVoice https://github.com/FunAudioLLM/CosyVoice
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import numpy as np
import torch
import torchaudio as ta
from functools import lru_cache
from typing import Optional
from ..s3tokenizer import S3_SR, SPEECH_VOCAB_SIZE, S3Tokenizer
from .const import S3GEN_SR
from .flow import CausalMaskedDiffWithXvec
from .xvector import CAMPPlus
from .utils.mel import mel_spectrogram
from .f0_predictor import ConvRNNF0Predictor
from .hifigan import HiFTGenerator
from .transformer.upsample_encoder import UpsampleConformerEncoder
from .flow_matching import CausalConditionalCFM
from .decoder import ConditionalDecoder
from .configs import CFM_PARAMS
def drop_invalid_tokens(x):
assert len(x.shape) <= 2 and x.shape[0] == 1, "only batch size of one allowed for now"
return x[x < SPEECH_VOCAB_SIZE]
# TODO: global resampler cache
@lru_cache(100)
def get_resampler(src_sr, dst_sr, device):
return ta.transforms.Resample(src_sr, dst_sr).to(device)
class S3Token2Mel(torch.nn.Module):
"""
S3Gen's CFM decoder maps S3 speech tokens to mel-spectrograms.
TODO: make these modules configurable?
"""
def __init__(self, meanflow=False):
super().__init__()
self.tokenizer = S3Tokenizer("speech_tokenizer_v2_25hz")
self.mel_extractor = mel_spectrogram # TODO: make it a torch module?
self.speaker_encoder = CAMPPlus(
# NOTE: This doesn't affect inference. It turns off activation checkpointing
# (a training optimization), which causes a crazy DDP error with accelerate
memory_efficient=False,
)
self.meanflow = meanflow
encoder = UpsampleConformerEncoder(
output_size=512,
attention_heads=8,
linear_units=2048,
num_blocks=6,
dropout_rate=0.1,
positional_dropout_rate=0.1,
attention_dropout_rate=0.1,
normalize_before=True,
input_layer='linear',
pos_enc_layer_type='rel_pos_espnet',
selfattention_layer_type='rel_selfattn',
input_size=512,
use_cnn_module=False,
macaron_style=False,
)
estimator = ConditionalDecoder(
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',
meanflow=self.meanflow,
)
cfm_params = CFM_PARAMS
decoder = CausalConditionalCFM(
spk_emb_dim=80,
cfm_params=cfm_params,
estimator=estimator,
)
self.flow = CausalMaskedDiffWithXvec(
encoder=encoder,
decoder=decoder
)
self.resamplers = {}
@property
def device(self):
params = self.tokenizer.parameters()
return next(params).device
@property
def dtype(self):
params = self.flow.parameters()
return next(params).dtype
def embed_ref(
self,
ref_wav: torch.Tensor,
ref_sr: int,
device="auto",
ref_fade_out=True,
):
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) # (B, L)
if ref_wav.size(1) > 10 * ref_sr:
print("WARNING: s3gen received ref longer than 10s")
ref_wav_24 = ref_wav
if ref_sr != S3GEN_SR:
ref_wav_24 = get_resampler(ref_sr, S3GEN_SR, device)(ref_wav)
ref_wav_24 = ref_wav_24.to(device=device, dtype=self.dtype)
ref_mels_24 = self.mel_extractor(ref_wav_24).transpose(1, 2).to(dtype=self.dtype)
ref_mels_24_len = None
# Resample to 16kHz
ref_wav_16 = ref_wav
if ref_sr != S3_SR:
ref_wav_16 = get_resampler(ref_sr, S3_SR, device)(ref_wav)
# Speaker embedding
ref_x_vector = self.speaker_encoder.inference(ref_wav_16.to(dtype=self.dtype))
# Tokenize 16khz reference
ref_speech_tokens, ref_speech_token_lens = self.tokenizer(ref_wav_16.float())
# Make sure mel_len = 2 * stoken_len (happens when the input is not padded to multiple of 40ms)
if ref_mels_24.shape[1] != 2 * ref_speech_tokens.shape[1]:
logging.warning(
"Reference mel length is not equal to 2 * reference token length.\n"
)
ref_speech_tokens = ref_speech_tokens[:, :ref_mels_24.shape[1] // 2]
ref_speech_token_lens[0] = ref_speech_tokens.shape[1]
return dict(
prompt_token=ref_speech_tokens.to(device),
prompt_token_len=ref_speech_token_lens,
prompt_feat=ref_mels_24,
prompt_feat_len=ref_mels_24_len,
embedding=ref_x_vector,
)
def forward(
self,
speech_tokens: torch.LongTensor,
# locally-computed ref embedding (mutex with ref_dict)
ref_wav: Optional[torch.Tensor],
ref_sr: Optional[int],
# pre-computed ref embedding (prod API)
ref_dict: Optional[dict] = None,
n_cfm_timesteps = None,
finalize: bool = False,
speech_token_lens=None,
noised_mels=None,
):
"""
Generate waveforms from S3 speech tokens and a reference waveform, which the speaker timbre is inferred from.
NOTE:
- The speaker encoder accepts 16 kHz waveform.
- S3TokenizerV2 accepts 16 kHz waveform.
- The mel-spectrogram for the reference assumes 24 kHz input signal.
- This function is designed for batch_size=1 only.
Args
----
- `speech_tokens`: S3 speech tokens [B=1, T]
- `ref_wav`: reference waveform (`torch.Tensor` with shape=[B=1, T])
- `ref_sr`: reference sample rate
- `finalize`: whether streaming is finished or not. Note that if False, the last 3 tokens will be ignored.
"""
assert (ref_wav is None) ^ (ref_dict is None), f"Must provide exactly one of ref_wav or ref_dict (got {ref_wav} and {ref_dict})"
if ref_dict is None:
ref_dict = self.embed_ref(ref_wav, ref_sr)
else:
# type/device casting (all values will be numpy if it's from a prod API call)
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(device=self.device, dtype=self.dtype)
speech_tokens = torch.atleast_2d(speech_tokens)
# backcompat
if speech_token_lens is None:
speech_token_lens = torch.LongTensor([st.size(-1) for st in speech_tokens]).to(self.device)
output_mels, _ = self.flow.inference(
token=speech_tokens,
token_len=speech_token_lens,
finalize=finalize,
noised_mels=noised_mels,
n_timesteps=n_cfm_timesteps,
meanflow=self.meanflow,
**ref_dict,
)
return output_mels
class S3Token2Wav(S3Token2Mel):
"""
The decoder of S3Gen is a concat of token-to-mel (CFM) and a mel-to-waveform (HiFiGAN) modules.
TODO: make these modules configurable?
"""
ignore_state_dict_missing = ("tokenizer._mel_filters", "tokenizer.window")
def __init__(self, meanflow=False):
super().__init__(meanflow)
f0_predictor = ConvRNNF0Predictor()
self.mel2wav = HiFTGenerator(
sampling_rate=S3GEN_SR,
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]],
f0_predictor=f0_predictor,
)
# silence out a few ms and fade audio in to reduce artifacts
n_trim = S3GEN_SR // 50 # 20ms = half of a frame
trim_fade = torch.zeros(2 * n_trim)
trim_fade[n_trim:] = (torch.cos(torch.linspace(torch.pi, 0, n_trim)) + 1) / 2
self.register_buffer("trim_fade", trim_fade, persistent=False) # (buffers get automatic device casting)
self.estimator_dtype = "fp32"
def forward(
self,
speech_tokens,
# locally-computed ref embedding (mutex with ref_dict)
ref_wav: Optional[torch.Tensor],
ref_sr: Optional[int],
# pre-computed ref embedding (prod API)
ref_dict: Optional[dict] = None,
finalize: bool = False,
speech_token_lens=None,
skip_vocoder=False,
n_cfm_timesteps=None,
noised_mels=None,
):
"""
Generate waveforms from S3 speech tokens and a reference waveform, which the speaker timbre is inferred from.
NOTE: used for sync synthesis only. Please use `S3GenStreamer` for streaming synthesis.
"""
output_mels = super().forward(
speech_tokens, speech_token_lens=speech_token_lens, ref_wav=ref_wav,
ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize,
n_cfm_timesteps=n_cfm_timesteps, noised_mels=noised_mels,
)
if skip_vocoder:
return output_mels
# TODO jrm: ignoring the speed control (mel interpolation) and the HiFTGAN caching mechanisms for now.
hift_cache_source = torch.zeros(1, 1, 0).to(self.device)
output_wavs, *_ = self.mel2wav.inference(speech_feat=output_mels, cache_source=hift_cache_source)
if not self.training:
# NOTE: ad-hoc method to reduce "spillover" from the reference clip.
output_wavs[:, :len(self.trim_fade)] *= self.trim_fade
return output_wavs
@torch.inference_mode()
def flow_inference(
self,
speech_tokens,
# locally-computed ref embedding (mutex with ref_dict)
ref_wav: Optional[torch.Tensor] = None,
ref_sr: Optional[int] = None,
# pre-computed ref embedding (prod API)
ref_dict: Optional[dict] = None,
n_cfm_timesteps = None,
finalize: bool = False,
speech_token_lens=None,
):
n_cfm_timesteps = n_cfm_timesteps or (2 if self.meanflow else 10)
noise = None
if self.meanflow:
noise = torch.randn(1, 80, speech_tokens.size(-1) * 2, dtype=self.dtype, device=self.device)
output_mels = super().forward(
speech_tokens, speech_token_lens=speech_token_lens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict,
n_cfm_timesteps=n_cfm_timesteps, finalize=finalize, noised_mels=noise,
)
return output_mels
@torch.inference_mode()
def hift_inference(self, speech_feat, cache_source: torch.Tensor = None):
if cache_source is None:
cache_source = torch.zeros(1, 1, 0).to(device=self.device, dtype=self.dtype)
return self.mel2wav.inference(speech_feat=speech_feat, cache_source=cache_source)
@torch.inference_mode()
def inference(
self,
speech_tokens,
# locally-computed ref embedding (mutex with ref_dict)
ref_wav: Optional[torch.Tensor] = None,
ref_sr: Optional[int] = None,
# pre-computed ref embedding (prod API)
ref_dict: Optional[dict] = None,
# left as a kwarg because this can change input/output size ratio
drop_invalid_tokens=True,
n_cfm_timesteps=None,
speech_token_lens=None,
):
# hallucination prevention, drop special tokens
# if drop_invalid_tokens:
# speech_tokens, speech_token_lens = drop_invalid(speech_tokens, pad=S3_QUIET_PAD)
output_mels = self.flow_inference(
speech_tokens,
speech_token_lens=speech_token_lens,
ref_wav=ref_wav,
ref_sr=ref_sr,
ref_dict=ref_dict,
n_cfm_timesteps=n_cfm_timesteps,
finalize=True,
)
output_mels = output_mels.to(dtype=self.dtype) # FIXME (fp16 mode) is this still needed?
output_wavs, output_sources = self.hift_inference(output_mels, None)
# NOTE: ad-hoc method to reduce "spillover" from the reference clip.
output_wavs[:, :len(self.trim_fade)] *= self.trim_fade
return output_wavs, output_sources
================================================
FILE: src/chatterbox/models/s3gen/transformer/__init__.py
================================================
================================================
FILE: src/chatterbox/models/s3gen/transformer/activation.py
================================================
# Copyright (c) 2020 Johns Hopkins University (Shinji Watanabe)
# 2020 Northwestern Polytechnical University (Pengcheng Guo)
# 2020 Mobvoi Inc (Binbin Zhang)
# 2024 Alibaba Inc (Xiang Lyu)
#
# 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.
"""Swish() activation function for Conformer."""
import torch
from torch import nn, sin, pow
from torch.nn import Parameter
class Swish(torch.nn.Module):
"""Construct an Swish object."""
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Return Swish activation function."""
return x * torch.sigmoid(x)
# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license.
# LICENSE is in incl_licenses directory.
class Snake(nn.Module):
'''
Implementation of a sine-based periodic activation function
Shape:
- Input: (B, C, T)
- Output: (B, C, T), same shape as the input
Parameters:
- alpha - trainable parameter
References:
- This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
https://arxiv.org/abs/2006.08195
Examples:
>>> a1 = snake(256)
>>> x = torch.randn(256)
>>> x = a1(x)
'''
def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
'''
Initialization.
INPUT:
- in_features: shape of the input
- alpha: trainable parameter
alpha is initialized to 1 by default, higher values = higher-frequency.
alpha will be trained along with the rest of your model.
'''
super(Snake, self).__init__()
self.in_features = in_features
# initialize alpha
self.alpha_logscale = alpha_logscale
if self.alpha_logscale: # log scale alphas initialized to zeros
self.alpha = Parameter(torch.zeros(in_features) * alpha)
else: # linear scale alphas initialized to ones
self.alpha = Parameter(torch.ones(in_features) * alpha)
self.alpha.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.
Snake ∶= x + 1/a * sin^2 (xa)
'''
alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
if self.alpha_logscale:
alpha = torch.exp(alpha)
x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
return x
================================================
FILE: src/chatterbox/models/s3gen/transformer/attention.py
================================================
# Copyright (c) 2019 Shigeki Karita
# 2020 Mobvoi Inc (Binbin Zhang)
# 2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn)
# 2024 Alibaba Inc (Xiang Lyu)
#
# 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.
"""Multi-Head Attention layer definition."""
import math
from typing import Tuple
import torch
from torch import nn
class MultiHeadedAttention(nn.Module):
"""Multi-Head Attention layer.
Args:
n_head (int): The number of heads.
n_feat (int): The number of features.
dropout_rate (float): Dropout rate.
"""
def __init__(self,
n_head: int,
n_feat: int,
dropout_rate: float,
key_bias: bool = True):
"""Construct an MultiHeadedAttention object."""
super().__init__()
assert n_feat % n_head == 0
# We assume d_v always equals d_k
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
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Transform query, key and value.
Args:
query (torch.Tensor): Query tensor (#batch, time1, size).
key (torch.Tensor): Key tensor (#batch, time2, size).
value (torch.Tensor): Value tensor (#batch, time2, size).
Returns:
torch.Tensor: Transformed query tensor, size
(#batch, n_head, time1, d_k).
torch.Tensor: Transformed key tensor, size
(#batch, n_head, time2, d_k).
torch.Tensor: Transformed value tensor, size
(#batch, n_head, time2, d_k).
"""
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) # (batch, head, time1, d_k)
k = k.transpose(1, 2) # (batch, head, time2, d_k)
v = v.transpose(1, 2) # (batch, head, time2, d_k)
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:
"""Compute attention context vector.
Args:
value (torch.Tensor): Transformed value, size
(#batch, n_head, time2, d_k).
scores (torch.Tensor): Attention score, size
(#batch, n_head, time1, time2).
mask (torch.Tensor): Mask, size (#batch, 1, time2) or
(#batch, time1, time2), (0, 0, 0) means fake mask.
Returns:
torch.Tensor: Transformed value (#batch, time1, d_model)
weighted by the attention score (#batch, time1, time2).
"""
n_batch = value.size(0)
# NOTE(xcsong): When will `if mask.size(2) > 0` be True?
# 1. onnx(16/4) [WHY? Because we feed real cache & real mask for the
# 1st chunk to ease the onnx export.]
# 2. pytorch training
if mask.size(2) > 0: # time2 > 0
mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2)
# For last chunk, time2 might be larger than scores.size(-1)
mask = mask[:, :, :, :scores.size(-1)] # (batch, 1, *, time2)
scores = scores.masked_fill(mask, -float('inf'))
attn = torch.softmax(scores, dim=-1).masked_fill(
mask, 0.0) # (batch, head, time1, time2)
# NOTE(xcsong): When will `if mask.size(2) > 0` be False?
# 1. onnx(16/-1, -1/-1, 16/0)
# 2. jit (16/-1, -1/-1, 16/0, 16/4)
else:
attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2)
p_attn = self.dropout(attn)
x = torch.matmul(p_attn, value) # (batch, head, time1, d_k)
x = (x.transpose(1, 2).contiguous().view(n_batch, -1,
self.h * self.d_k)
) # (batch, time1, d_model)
return self.linear_out(x) # (batch, time1, d_model)
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]:
"""Compute scaled dot product attention.
Args:
query (torch.Tensor): Query tensor (#batch, time1, size).
key (torch.Tensor): Key tensor (#batch, time2, size).
value (torch.Tensor): Value tensor (#batch, time2, size).
mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
(#batch, time1, time2).
1.When applying cross attention between decoder and encoder,
the batch padding mask for input is in (#batch, 1, T) shape.
2.When applying self attention of encoder,
the mask is in (#batch, T, T) shape.
3.When applying self attention of decoder,
the mask is in (#batch, L, L) shape.
4.If the different position in decoder see different block
of the encoder, such as Mocha, the passed in mask could be
in (#batch, L, T) shape. But there is no such case in current
CosyVoice.
cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2),
where `cache_t == chunk_size * num_decoding_left_chunks`
and `head * d_k == size`
Returns:
torch.Tensor: Output tensor (#batch, time1, d_model).
torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2)
where `cache_t == chunk_size * num_decoding_left_chunks`
and `head * d_k == size`
"""
q, k, v = self.forward_qkv(query, key, value)
# NOTE(xcsong):
# when export onnx model, for 1st chunk, we feed
# cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode)
# or cache(1, head, real_cache_t, d_k * 2) (16/4 mode).
# In all modes, `if cache.size(0) > 0` will alwayse be `True`
# and we will always do splitting and
# concatnation(this will simplify onnx export). Note that
# it's OK to concat & split zero-shaped tensors(see code below).
# when export jit model, for 1st chunk, we always feed
# cache(0, 0, 0, 0) since jit supports dynamic if-branch.
# >>> a = torch.ones((1, 2, 0, 4))
# >>> b = torch.ones((1, 2, 3, 4))
# >>> c = torch.cat((a, b), dim=2)
# >>> torch.equal(b, c) # True
# >>> d = torch.split(a, 2, dim=-1)
# >>> torch.equal(d[0], d[1]) # True
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)
# NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
# non-trivial to calculate `next_cache_start` here.
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.
Paper: https://arxiv.org/abs/1901.02860
Args:
n_head (int): The number of heads.
n_feat (int): The number of features.
dropout_rate (float): Dropout rate.
"""
def __init__(self,
n_head: int,
n_feat: int,
dropout_rate: float,
key_bias: bool = True):
"""Construct an RelPositionMultiHeadedAttention object."""
super().__init__(n_head, n_feat, dropout_rate, key_bias)
# linear transformation for positional encoding
self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
# these two learnable bias are used in matrix c and matrix d
# as described in https://arxiv.org/abs/1901.02860 Section 3.3
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))
torch.nn.init.xavier_uniform_(self.pos_bias_u)
torch.nn.init.xavier_uniform_(self.pos_bias_v)
def rel_shift(self, x: torch.Tensor) -> torch.Tensor:
"""Compute relative positional encoding.
Args:
x (torch.Tensor): Input tensor (batch, head, time1, 2*time1-1).
time1 means the length of query vector.
Returns:
torch.Tensor: Output 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
] # only keep the positions from 0 to time2
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]:
"""Compute 'Scaled Dot Product Attention' with rel. positional encoding.
Args:
query (torch.Tensor): Query tensor (#batch, time1, size).
key (torch.Tensor): Key tensor (#batch, time2, size).
value (torch.Tensor): Value tensor (#batch, time2, size).
mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
(#batch, time1, time2), (0, 0, 0) means fake mask.
pos_emb (torch.Tensor): Positional embedding tensor
(#batch, time2, size).
cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2),
where `cache_t == chunk_size * num_decoding_left_chunks`
and `head * d_k == size`
Returns:
torch.Tensor: Output tensor (#batch, time1, d_model).
torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2)
where `cache_t == chunk_size * num_decoding_left_chunks`
and `head * d_k == size`
"""
q, k, v = self.forward_qkv(query, key, value)
q = q.transpose(1, 2) # (batch, time1, head, d_k)
# NOTE(xcsong):
# when export onnx model, for 1st chunk, we feed
# cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode)
# or cache(1, head, real_cache_t, d_k * 2) (16/4 mode).
# In all modes, `if cache.size(0) > 0` will alwayse be `True`
# and we will always do splitting and
# concatnation(this will simplify onnx export). Note that
# it's OK to concat & split zero-shaped tensors(see code below).
# when export jit model, for 1st chunk, we always feed
# cache(0, 0, 0, 0) since jit supports dynamic if-branch.
# >>> a = torch.ones((1, 2, 0, 4))
# >>> b = torch.ones((1, 2, 3, 4))
# >>> c = torch.cat((a, b), dim=2)
# >>> torch.equal(b, c) # True
# >>> d = torch.split(a, 2, dim=-1)
# >>> torch.equal(d[0], d[1]) # True
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)
# NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
# non-trivial to calculate `next_cache_start` here.
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) # (batch, head, time1, d_k)
# (batch, head, time1, d_k)
q_with_bias_u = (q + self.pos_bias_u.to(q.device)).transpose(1, 2)
# (batch, head, time1, d_k)
q_with_bias_v = (q + self.pos_bias_v.to(q.device)).transpose(1, 2)
# compute attention score
# first compute matrix a and matrix c
# as described in https://arxiv.org/abs/1901.02860 Section 3.3
# (batch, head, time1, time2)
matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
# compute matrix b and matrix d
# (batch, head, time1, time2)
matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
# NOTE(Xiang Lyu): Keep rel_shift since espnet rel_pos_emb is used
if matrix_ac.shape != matrix_bd.shape:
matrix_bd = self.rel_shift(matrix_bd)
scores = (matrix_ac + matrix_bd) / math.sqrt(
self.d_k) # (batch, head, time1, time2)
return self.forward_attention(v, scores, mask), new_cache
================================================
FILE: src/chatterbox/models/s3gen/transformer/convolution.py
================================================
# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Di Wu)
# 2024 Alibaba Inc (Xiang Lyu)
#
# 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.
# Modified from ESPnet(https://github.com/espnet/espnet)
"""ConvolutionModule definition."""
from typing import Tuple
import torch
from torch import nn
class ConvolutionModule(nn.Module):
"""ConvolutionModule in Conformer model."""
def __init__(self,
channels: int,
kernel_size: int = 15,
activation: nn.Module = nn.ReLU(),
norm: str = "batch_norm",
causal: bool = False,
bias: bool = True):
"""Construct an ConvolutionModule object.
Args:
channels (int): The number of channels of conv layers.
kernel_size (int): Kernel size of conv layers.
causal (int): Whether use causal convolution or not
"""
super().__init__()
self.pointwise_conv1 = nn.Conv1d(
channels,
2 * channels,
kernel_size=1,
stride=1,
padding=0,
bias=bias,
)
# self.lorder is used to distinguish if it's a causal convolution,
# if self.lorder > 0: it's a causal convolution, the input will be
# padded with self.lorder frames on the left in forward.
# else: it's a symmetrical convolution
if causal:
padding = 0
self.lorder = kernel_size - 1
else:
# kernel_size should be an odd number for none causal convolution
assert (kernel_size - 1) % 2 == 0
padding = (kernel_size - 1) // 2
self.lorder = 0
self.depthwise_conv = nn.Conv1d(
channels,
channels,
kernel_size,
stride=1,
padding=padding,
groups=channels,
bias=bias,
)
assert norm in ['batch_norm', 'layer_norm']
if norm == "batch_norm":
self.use_layer_norm = False
self.norm = nn.BatchNorm1d(channels)
else:
self.use_layer_norm = True
self.norm = nn.LayerNorm(channels)
self.pointwise_conv2 = nn.Conv1d(
channels,
channels,
kernel_size=1,
stride=1,
padding=0,
bias=bias,
)
self.activation = activation
def forward(
self,
x: torch.Tensor,
mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
cache: torch.Tensor = torch.zeros((0, 0, 0)),
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Compute convolution module.
Args:
x (torch.Tensor): Input tensor (#batch, time, channels).
mask_pad (torch.Tensor): used for batch padding (#batch, 1, time),
(0, 0, 0) means fake mask.
cache (torch.Tensor): left context cache, it is only
used in causal convolution (#batch, channels, cache_t),
(0, 0, 0) meas fake cache.
Returns:
torch.Tensor: Output tensor (#batch, time, channels).
"""
# exchange the temporal dimension and the feature dimension
x = x.transpose(1, 2) # (#batch, channels, time)
# mask batch padding
if mask_pad.size(2) > 0: # time > 0
x.masked_fill_(~mask_pad, 0.0)
if self.lorder > 0:
if cache.size(2) == 0: # cache_t == 0
x = nn.functional.pad(x, (self.lorder, 0), 'constant', 0.0)
else:
assert cache.size(0) == x.size(0) # equal batch
assert cache.size(1) == x.size(1) # equal channel
x = torch.cat((cache, x), dim=2)
assert (x.size(2) > self.lorder)
new_cache = x[:, :, -self.lorder:]
else:
# It's better we just return None if no cache is required,
# However, for JIT export, here we just fake one tensor instead of
# None.
new_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device)
# GLU mechanism
x = self.pointwise_conv1(x) # (batch, 2*channel, dim)
x = nn.functional.glu(x, dim=1) # (batch, channel, dim)
# 1D Depthwise Conv
x = self.depthwise_conv(x)
if self.use_layer_norm:
x = x.transpose(1, 2)
x = self.activation(self.norm(x))
if self.use_layer_norm:
x = x.transpose(1, 2)
x = self.pointwise_conv2(x)
# mask batch padding
if mask_pad.size(2) > 0: # time > 0
x.masked_fill_(~mask_pad, 0.0)
return x.transpose(1, 2), new_cache
================================================
FILE: src/chatterbox/models/s3gen/transformer/embedding.py
================================================
# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Di Wu)
# 2024 Alibaba Inc (Xiang Lyu)
#
# 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.
# Modified from ESPnet(https://github.com/espnet/espnet)
"""Positonal Encoding Module."""
import math
from typing import Tuple, Union
import torch
import torch.nn.functional as F
import numpy as np
class PositionalEncoding(torch.nn.Module):
"""Positional encoding.
:param int d_model: embedding dim
:param float dropout_rate: dropout rate
:param int max_len: maximum input length
PE(pos, 2i) = sin(pos/(10000^(2i/dmodel)))
PE(pos, 2i+1) = cos(pos/(10000^(2i/dmodel)))
"""
def __init__(self,
d_model: int,
dropout_rate: float,
max_len: int = 5000,
reverse: bool = False):
"""Construct an PositionalEncoding object."""
super().__init__()
self.d_model = d_model
self.xscale = math.sqrt(self.d_model)
self.dropout = torch.nn.Dropout(p=dropout_rate)
self.max_len = max_len
self.pe = torch.zeros(self.max_len, self.d_model)
position = torch.arange(0, self.max_len,
dtype=torch.float32).unsqueeze(1)
div_term = torch.exp(
torch.arange(0, self.d_model, 2, dtype=torch.float32) *
-(math.log(10000.0) / self.d_model))
self.pe[:, 0::2] = torch.sin(position * div_term)
self.pe[:, 1::2] = torch.cos(position * div_term)
self.pe = self.pe.unsqueeze(0)
def forward(self,
x: torch.Tensor,
offset: Union[int, torch.Tensor] = 0) \
-> Tuple[torch.Tensor, torch.Tensor]:
"""Add positional encoding.
Args:
x (torch.Tensor): Input. Its shape is (batch, time, ...)
offset (int, torch.tensor): position offset
Returns:
torch.Tensor: Encoded tensor. Its shape is (batch, time, ...)
torch.Tensor: for compatibility to RelPositionalEncoding
"""
self.pe = self.pe.to(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: Union[int, torch.Tensor],
size: int,
apply_dropout: bool = True) -> torch.Tensor:
""" For getting encoding in a streaming fashion
Attention!!!!!
we apply dropout only once at the whole utterance level in a none
streaming way, but will call this function several times with
increasing input size in a streaming scenario, so the dropout will
be applied several times.
Args:
offset (int or torch.tensor): start offset
size (int): required size of position encoding
Returns:
torch.Tensor: Corresponding encoding
"""
# How to subscript a Union type:
# https://github.com/pytorch/pytorch/issues/69434
if isinstance(offset, int):
assert offset + size <= self.max_len
pos_emb = self.pe[:, offset:offset + size]
elif isinstance(offset, torch.Tensor) and offset.dim() == 0: # scalar
assert offset + size <= self.max_len
pos_emb = self.pe[:, offset:offset + size]
else: # for batched streaming decoding on GPU
assert torch.max(offset) + size <= self.max_len
index = offset.unsqueeze(1) + \
torch.arange(0, size).to(offset.device) # B X T
flag = index > 0
# remove negative offset
index = index * flag
pos_emb = F.embedding(index, self.pe[0]) # B X T X d_model
if apply_dropout:
pos_emb = self.dropout(pos_emb)
return pos_emb
class RelPositionalEncoding(PositionalEncoding):
"""Relative positional encoding module.
See : Appendix B in https://arxiv.org/abs/1901.02860
Args:
d_model (int): Embedding dimension.
dropout_rate (float): Dropout rate.
max_len (int): Maximum input length.
"""
def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000):
"""Initialize class."""
super().__init__(d_model, dropout_rate, max_len, reverse=True)
def forward(self,
x: torch.Tensor,
offset: Union[int, torch.Tensor] = 0) \
-> Tuple[torch.Tensor, torch.Tensor]:
"""Compute positional encoding.
Args:
x (torch.Tensor): Input tensor (batch, time, `*`).
Returns:
torch.Tensor: Encoded tensor (batch, time, `*`).
torch.Tensor: Positional embedding tensor (1, time, `*`).
"""
self.pe = self.pe.to(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 WhisperPositionalEncoding(PositionalEncoding):
""" Sinusoids position encoding used in openai-whisper.encoder
"""
def __init__(self, d_model: int, dropout_rate: float, max_len: int = 1500):
super().__init__(d_model, dropout_rate, max_len)
self.xscale = 1.0
log_timescale_increment = np.log(10000) / (d_model // 2 - 1)
inv_timescales = torch.exp(-log_timescale_increment *
torch.arange(d_model // 2))
scaled_time = torch.arange(max_len)[:, np.newaxis] * \
inv_timescales[np.newaxis, :]
pe = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
delattr(self, "pe")
self.register_buffer("pe", pe.unsqueeze(0))
class LearnablePositionalEncoding(PositionalEncoding):
""" Learnable position encoding used in openai-whisper.decoder
"""
def __init__(self, d_model: int, dropout_rate: float, max_len: int = 448):
super().__init__(d_model, dropout_rate, max_len)
# NOTE(xcsong): overwrite self.pe & self.xscale
self.pe = torch.nn.Parameter(torch.empty(1, max_len, d_model))
self.xscale = 1.0
class NoPositionalEncoding(torch.nn.Module):
""" No position encoding
"""
def __init__(self, d_model: int, dropout_rate: float):
super().__init__()
self.d_model = d_model
self.dropout = torch.nn.Dropout(p=dropout_rate)
def forward(self,
x: torch.Tensor,
offset: Union[int, torch.Tensor] = 0) \
-> Tuple[torch.Tensor, torch.Tensor]:
""" Just return zero vector for interface compatibility
"""
pos_emb = torch.zeros(1, x.size(1), self.d_model).to(x.device)
return self.dropout(x), pos_emb
def position_encoding(self, offset: Union[int, torch.Tensor],
size: int) -> torch.Tensor:
return torch.zeros(1, size, self.d_model)
class EspnetRelPositionalEncoding(torch.nn.Module):
"""Relative positional encoding module (new implementation).
Details can be found in https://github.com/espnet/espnet/pull/2816.
See : Appendix B in https://arxiv.org/abs/1901.02860
Args:
d_model (int): Embedding dimension.
dropout_rate (float): Dropout rate.
max_len (int): Maximum input length.
"""
def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000):
"""Construct an PositionalEncoding object."""
super(EspnetRelPositionalEncoding, self).__init__()
self.d_model = d_model
self.xscale = math.sqrt(self.d_model)
self.dropout = torch.nn.Dropout(p=dropout_rate)
self.pe = None
self.extend_pe(torch.tensor(0.0).expand(1, max_len))
def extend_pe(self, x: torch.Tensor):
"""Reset the positional encodings."""
if self.pe is not None:
# self.pe contains both positive and negative parts
# the length of self.pe is 2 * input_len - 1
if 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
# Suppose `i` means to the position of query vecotr and `j` means the
# position of key vector. We use position relative positions when keys
# are to the left (i>j) and negative relative positions otherwise (i<j).
pe_positive = torch.zeros(x.size(1), self.d_model)
pe_negative = torch.zeros(x.size(1), self.d_model)
position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
div_term = torch.exp(
torch.arange(0, self.d_model, 2, dtype=torch.float32)
* -(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)
# Reserve the order of positive indices and concat both positive and
# negative indices. This is used to support the shifting trick
# as in https://arxiv.org/abs/1901.02860
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: Union[int, torch.Tensor] = 0) \
-> Tuple[torch.Tensor, torch.Tensor]:
"""Add positional encoding.
Args:
x (torch.Tensor): Input tensor (batch, time, `*`).
Returns:
torch.Tensor: Encoded tensor (batch, time, `*`).
"""
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: Union[int, torch.Tensor],
size: int) -> torch.Tensor:
""" For getting encoding in a streaming fashion
Attention!!!!!
we apply dropout only once at the whole utterance level in a none
streaming way, but will call this function several times with
increasing input size in a streaming scenario, so the dropout will
be applied several times.
Args:
offset (int or torch.tensor): start offset
size (int): required size of position encoding
Returns:
torch.Tensor: Corresponding encoding
"""
pos_emb = self.pe[
:,
self.pe.size(1) // 2 - size + 1: self.pe.size(1) // 2 + size,
]
return pos_emb
================================================
FILE: src/chatterbox/models/s3gen/transformer/encoder_layer.py
================================================
# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu)
# 2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn)
#
# 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.
# Modified from ESPnet(https://github.com/espnet/espnet)
"""Encoder self-attention layer definition."""
from typing import Optional, Tuple
import torch
from torch import nn
class TransformerEncoderLayer(nn.Module):
"""Encoder layer module.
Args:
size (int): Input dimension.
self_attn (torch.nn.Module): Self-attention module instance.
`MultiHeadedAttention` or `RelPositionMultiHeadedAttention`
instance can be used as the argument.
feed_forward (torch.nn.Module): Feed-forward module instance.
`PositionwiseFeedForward`, instance can be used as the argument.
dropout_rate (float): Dropout rate.
normalize_before (bool):
True: use layer_norm before each sub-block.
False: to use layer_norm after each sub-block.
"""
def __init__(
self,
size: int,
self_attn: torch.nn.Module,
feed_forward: torch.nn.Module,
dropout_rate: float,
normalize_before: bool = True,
):
"""Construct an EncoderLayer object."""
super().__init__()
self.self_attn = self_attn
self.feed_forward = feed_forward
self.norm1 = nn.LayerNorm(size, eps=1e-12)
self.norm2 = 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)),
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Compute encoded features.
Args:
x (torch.Tensor): (#batch, time, size)
mask (torch.Tensor): Mask tensor for the input (#batch, time，time),
(0, 0, 0) means fake mask.
pos_emb (torch.Tensor): just for interface compatibility
to ConformerEncoderLayer
mask_pad (torch.Tensor): does not used in transformer layer,
just for unified api with conformer.
att_cache (torch.Tensor): Cache tensor of the KEY & VALUE
(#batch=1, head, cache_t1, d_k * 2), head * d_k == size.
cnn_cache (torch.Tensor): Convolution cache in conformer layer
(#batch=1, size, cache_t2), not used here, it's for interface
compatibility to ConformerEncoderLayer.
Returns:
torch.Tensor: Output tensor (#batch, time, size).
torch.Tensor: Mask tensor (#batch, time, time).
torch.Tensor: att_cache tensor,
(#batch=1, head, cache_t1 + time, d_k * 2).
torch.Tensor: cnn_cahce tensor (#batch=1, size, cache_t2).
"""
residual = x
if self.normalize_before:
x = self.norm1(x)
x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb=pos_emb, cache=att_cache)
x = residual + self.dropout(x_att)
if not self.normalize_before:
x = self.norm1(x)
residual = x
if self.normalize_before:
x = self.norm2(x)
x = residual + self.dropout(self.feed_forward(x))
if not self.normalize_before:
x = self.norm2(x)
fake_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device)
return x, mask, new_att_cache, fake_cnn_cache
class ConformerEncoderLayer(nn.Module):
"""Encoder layer module.
Args:
size (int): Input dimension.
self_attn (torch.nn.Module): Self-attention module instance.
`MultiHeadedAttention` or `RelPositionMultiHeadedAttention`
instance can be used as the argument.
feed_forward (torch.nn.Module): Feed-forward module instance.
`PositionwiseFeedForward` instance can be used as the argument.
feed_forward_macaron (torch.nn.Module): Additional feed-forward module
instance.
`PositionwiseFeedForward` instance can be used as the argument.
conv_module (torch.nn.Module): Convolution module instance.
`ConvlutionModule` instance can be used as the argument.
dropout_rate (float): Dropout rate.
normalize_before (bool):
True: use layer_norm before each sub-block.
False: use layer_norm after each sub-block.
"""
def __init__(
self,
size: int,
self_attn: torch.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,
):
"""Construct an EncoderLayer object."""
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) # for the FNN module
self.norm_mha = nn.LayerNorm(size, eps=1e-12) # for the MHA module
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) # for the CNN module
self.norm_final = nn.LayerNorm(
size, eps=1e-12) # for the final output of the block
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)),
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Compute encoded features.
Args:
x (torch.Tensor): (#batch, time, size)
mask (torch.Tensor): Mask tensor for the input (#batch, time，time),
(0, 0, 0) means fake mask.
pos_emb (torch.Tensor): positional encoding, must not be None
for ConformerEncoderLayer.
mask_pad (torch.Tensor): batch padding mask used for conv module.
(#batch, 1，time), (0, 0, 0) means fake mask.
att_cache (torch.Tensor): Cache tensor of the KEY & VALUE
(#batch=1, head, cache_t1, d_k * 2), head * d_k == size.
cnn_cache (torch.Tensor): Convolution cache in conformer layer
(#batch=1, size, cache_t2)
Returns:
torch.Tensor: Output tensor (#batch, time, size).
torch.Tensor: Mask tensor (#batch, time, time).
torch.Tensor: att_cache tensor,
(#batch=1, head, cache_t1 + time, d_k * 2).
torch.Tensor: cnn_cahce tensor (#batch, size, cache_t2).
"""
# whether to use macaron style
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)
# multi-headed self-attention module
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)
# convolution module
# Fake new cnn cache here, and then change it in conv_module
new_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device)
if
gitextract_lols3utf/
├── .github/
│ └── workflows/
│ └── install_check.yml
├── .gitignore
├── LICENSE
├── README.md
├── example_for_mac.py
├── example_tts.py
├── example_tts_turbo.py
├── example_vc.py
├── gradio_tts_app.py
├── gradio_tts_turbo_app.py
├── gradio_vc_app.py
├── multilingual_app.py
├── pyproject.toml
└── src/
└── chatterbox/
├── __init__.py
├── models/
│ ├── __init__.py
│ ├── s3gen/
│ │ ├── __init__.py
│ │ ├── configs.py
│ │ ├── const.py
│ │ ├── decoder.py
│ │ ├── f0_predictor.py
│ │ ├── flow.py
│ │ ├── flow_matching.py
│ │ ├── hifigan.py
│ │ ├── matcha/
│ │ │ ├── decoder.py
│ │ │ ├── flow_matching.py
│ │ │ ├── text_encoder.py
│ │ │ └── transformer.py
│ │ ├── s3gen.py
│ │ ├── transformer/
│ │ │ ├── __init__.py
│ │ │ ├── activation.py
│ │ │ ├── attention.py
│ │ │ ├── convolution.py
│ │ │ ├── embedding.py
│ │ │ ├── encoder_layer.py
│ │ │ ├── positionwise_feed_forward.py
│ │ │ ├── subsampling.py
│ │ │ └── upsample_encoder.py
│ │ ├── utils/
│ │ │ ├── class_utils.py
│ │ │ ├── intmeanflow.py
│ │ │ ├── mask.py
│ │ │ └── mel.py
│ │ └── xvector.py
│ ├── s3tokenizer/
│ │ ├── __init__.py
│ │ └── s3tokenizer.py
│ ├── t3/
│ │ ├── __init__.py
│ │ ├── inference/
│ │ │ ├── alignment_stream_analyzer.py
│ │ │ └── t3_hf_backend.py
│ │ ├── llama_configs.py
│ │ ├── modules/
│ │ │ ├── cond_enc.py
│ │ │ ├── learned_pos_emb.py
│ │ │ ├── perceiver.py
│ │ │ └── t3_config.py
│ │ └── t3.py
│ ├── tokenizers/
│ │ ├── __init__.py
│ │ └── tokenizer.py
│ ├── utils.py
│ └── voice_encoder/
│ ├── __init__.py
│ ├── config.py
│ ├── melspec.py
│ └── voice_encoder.py
├── mtl_tts.py
├── tts.py
├── tts_turbo.py
└── vc.py
SYMBOL INDEX (440 symbols across 45 files)
FILE: example_for_mac.py
function patched_torch_load (line 10) | def patched_torch_load(*args, **kwargs):
FILE: gradio_tts_app.py
function set_seed (line 11) | def set_seed(seed: int):
function load_model (line 19) | def load_model():
function generate (line 24) | def generate(model, text, audio_prompt_path, exaggeration, temperature, ...
FILE: gradio_tts_turbo_app.py
function set_seed (line 69) | def set_seed(seed: int):
function load_model (line 77) | def load_model():
function generate (line 83) | def generate(
FILE: gradio_vc_app.py
function generate (line 10) | def generate(audio, target_voice_path):
FILE: multilingual_app.py
function default_audio_for_ui (line 109) | def default_audio_for_ui(lang: str) -> str | None:
function default_text_for_ui (line 113) | def default_text_for_ui(lang: str) -> str:
function get_supported_languages_display (line 117) | def get_supported_languages_display() -> str:
function get_or_load_model (line 136) | def get_or_load_model():
function set_seed (line 158) | def set_seed(seed: int):
function resolve_audio_prompt (line 167) | def resolve_audio_prompt(language_id: str, provided_path: str | None) ->...
function generate_tts_audio (line 178) | def generate_tts_audio(
function on_language_change (line 293) | def on_language_change(lang, current_ref, current_text):
FILE: src/chatterbox/models/s3gen/decoder.py
function mask_to_bias (line 26) | def mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
class Transpose (line 38) | class Transpose(torch.nn.Module):
method __init__ (line 39) | def __init__(self, dim0: int, dim1: int):
method forward (line 44) | def forward(self, x: torch.Tensor):
class CausalBlock1D (line 49) | class CausalBlock1D(Block1D):
method __init__ (line 50) | def __init__(self, dim: int, dim_out: int):
method forward (line 60) | def forward(self, x: torch.Tensor, mask: torch.Tensor):
class CausalResnetBlock1D (line 65) | class CausalResnetBlock1D(ResnetBlock1D):
method __init__ (line 66) | def __init__(self, dim: int, dim_out: int, time_emb_dim: int, groups: ...
class CausalConv1d (line 72) | class CausalConv1d(torch.nn.Conv1d):
method __init__ (line 73) | def __init__(
method forward (line 95) | def forward(self, x: torch.Tensor):
class ConditionalDecoder (line 99) | class ConditionalDecoder(nn.Module):
method __init__ (line 100) | def __init__(
method dtype (line 226) | def dtype(self):
method initialize_weights (line 229) | def initialize_weights(self):
method forward (line 243) | def forward(self, x, mask, mu, t, spks=None, cond=None, r=None):
FILE: src/chatterbox/models/s3gen/f0_predictor.py
class ConvRNNF0Predictor (line 19) | class ConvRNNF0Predictor(nn.Module):
method __init__ (line 20) | def __init__(self,
method forward (line 52) | def forward(self, x: torch.Tensor) -> torch.Tensor:
FILE: src/chatterbox/models/s3gen/flow.py
function _repeat_batch_dim (line 30) | def _repeat_batch_dim(tnsr, B, ndim):
class CausalMaskedDiffWithXvec (line 43) | class CausalMaskedDiffWithXvec(torch.nn.Module):
method __init__ (line 44) | def __init__(self,
method compute_loss (line 86) | def compute_loss(
method inference (line 132) | def inference(self,
FILE: src/chatterbox/models/s3gen/flow_matching.py
function cast_all (line 22) | def cast_all(*args, dtype):
class ConditionalCFM (line 26) | class ConditionalCFM(BASECFM):
method __init__ (line 27) | def __init__(self, in_channels, cfm_params, n_spks=1, spk_emb_dim=64, ...
method forward (line 42) | def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, c...
method solve_euler (line 78) | def solve_euler(self, x, t_span, mu, mask, spks, cond, meanflow=False):
method compute_loss (line 147) | def compute_loss(self, x1, mask, mu, spks=None, cond=None):
class CausalConditionalCFM (line 189) | class CausalConditionalCFM(ConditionalCFM):
method __init__ (line 190) | def __init__(self, in_channels=240, cfm_params=CFM_PARAMS, n_spks=1, s...
method forward (line 196) | def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, c...
method basic_euler (line 235) | def basic_euler(self, x, t_span, mu, mask, spks, cond):
FILE: src/chatterbox/models/s3gen/hifigan.py
class Snake (line 34) | class Snake(nn.Module):
method __init__ (line 50) | def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha...
method forward (line 73) | def forward(self, x):
function get_padding (line 88) | def get_padding(kernel_size, dilation=1):
function init_weights (line 91) | def init_weights(m, mean=0.0, std=0.01):
class ResBlock (line 106) | class ResBlock(torch.nn.Module):
method __init__ (line 108) | def __init__(
method forward (line 154) | def forward(self, x: torch.Tensor) -> torch.Tensor:
method remove_weight_norm (line 163) | def remove_weight_norm(self):
class SineGen (line 169) | class SineGen(torch.nn.Module):
method __init__ (line 185) | def __init__(self, samp_rate, harmonic_num=0,
method _f02uv (line 195) | def _f02uv(self, f0):
method forward (line 201) | def forward(self, f0):
class SourceModuleHnNSF (line 234) | class SourceModuleHnNSF(torch.nn.Module):
method __init__ (line 252) | def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine...
method forward (line 267) | def forward(self, x):
class HiFTGenerator (line 286) | class HiFTGenerator(nn.Module):
method __init__ (line 291) | def __init__(
method remove_weight_norm (line 382) | def remove_weight_norm(self):
method _stft (line 396) | def _stft(self, x):
method _istft (line 404) | def _istft(self, magnitude, phase):
method decode (line 412) | def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, ...
method forward (line 446) | def forward(
method inference (line 463) | def inference(self, speech_feat: torch.Tensor, cache_source: torch.Ten...
FILE: src/chatterbox/models/s3gen/matcha/decoder.py
class SinusoidalPosEmb (line 14) | class SinusoidalPosEmb(torch.nn.Module):
method __init__ (line 15) | def __init__(self, dim):
method forward (line 20) | def forward(self, x, scale=1000):
class Block1D (line 32) | class Block1D(torch.nn.Module):
method __init__ (line 33) | def __init__(self, dim, dim_out, groups=8):
method forward (line 41) | def forward(self, x, mask):
class ResnetBlock1D (line 46) | class ResnetBlock1D(torch.nn.Module):
method __init__ (line 47) | def __init__(self, dim, dim_out, time_emb_dim, groups=8):
method forward (line 56) | def forward(self, x, mask, time_emb):
class Downsample1D (line 64) | class Downsample1D(nn.Module):
method __init__ (line 65) | def __init__(self, dim):
method forward (line 69) | def forward(self, x):
class TimestepEmbedding (line 73) | class TimestepEmbedding(nn.Module):
method __init__ (line 74) | def __init__(
method forward (line 105) | def forward(self, sample, condition=None):
class Upsample1D (line 120) | class Upsample1D(nn.Module):
method __init__ (line 134) | def __init__(self, channels, use_conv=False, use_conv_transpose=True, ...
method forward (line 148) | def forward(self, inputs):
class ConformerWrapper (line 161) | class ConformerWrapper(ConformerBlock):
method __init__ (line 162) | def __init__( # pylint: disable=useless-super-delegation
method forward (line 189) | def forward(
class Decoder (line 200) | class Decoder(nn.Module):
method __init__ (line 201) | def __init__(
method get_block (line 319) | def get_block(block_type, dim, attention_head_dim, num_heads, dropout,...
method initialize_weights (line 345) | def initialize_weights(self):
method forward (line 363) | def forward(self, x, mask, mu, t, spks=None, cond=None):
FILE: src/chatterbox/models/s3gen/matcha/flow_matching.py
class BASECFM (line 9) | class BASECFM(torch.nn.Module, ABC):
method __init__ (line 10) | def __init__(
method forward (line 30) | def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, c...
method solve_euler (line 52) | def solve_euler(self, x, t_span, mu, mask, spks, cond):
method compute_loss (line 84) | def compute_loss(self, x1, mask, mu, spks=None, cond=None):
class CFM (line 118) | class CFM(BASECFM):
method __init__ (line 119) | def __init__(self, in_channels, out_channel, cfm_params, decoder_param...
FILE: src/chatterbox/models/s3gen/matcha/text_encoder.py
function sequence_mask (line 10) | def sequence_mask(length, max_length=None):
class LayerNorm (line 18) | class LayerNorm(nn.Module):
method __init__ (line 19) | def __init__(self, channels, eps=1e-4):
method forward (line 27) | def forward(self, x):
class ConvReluNorm (line 39) | class ConvReluNorm(nn.Module):
method __init__ (line 40) | def __init__(self, in_channels, hidden_channels, out_channels, kernel_...
method forward (line 63) | def forward(self, x, x_mask):
class DurationPredictor (line 73) | class DurationPredictor(nn.Module):
method __init__ (line 74) | def __init__(self, in_channels, filter_channels, kernel_size, p_dropout):
method forward (line 87) | def forward(self, x, x_mask):
class RotaryPositionalEmbeddings (line 100) | class RotaryPositionalEmbeddings(nn.Module):
method __init__ (line 110) | def __init__(self, d: int, base: int = 10_000):
method _build_cache (line 122) | def _build_cache(self, x: torch.Tensor):
method _neg_half (line 150) | def _neg_half(self, x: torch.Tensor):
method forward (line 157) | def forward(self, x: torch.Tensor):
class MultiHeadAttention (line 178) | class MultiHeadAttention(nn.Module):
method __init__ (line 179) | def __init__(
method forward (line 219) | def forward(self, x, c, attn_mask=None):
method attention (line 229) | def attention(self, query, key, value, mask=None):
method _attention_bias_proximal (line 252) | def _attention_bias_proximal(length):
class FFN (line 258) | class FFN(nn.Module):
method __init__ (line 259) | def __init__(self, in_channels, out_channels, filter_channels, kernel_...
method forward (line 271) | def forward(self, x, x_mask):
class Encoder (line 279) | class Encoder(nn.Module):
method __init__ (line 280) | def __init__(
method forward (line 317) | def forward(self, x, x_mask):
class TextEncoder (line 331) | class TextEncoder(nn.Module):
method __init__ (line 332) | def __init__(
method forward (line 381) | def forward(self, x, x_lengths, spks=None):
FILE: src/chatterbox/models/s3gen/matcha/transformer.py
class SnakeBeta (line 17) | class SnakeBeta(nn.Module):
method __init__ (line 35) | def __init__(self, in_features, out_features, alpha=1.0, alpha_trainab...
method forward (line 64) | def forward(self, x):
class FeedForward (line 83) | class FeedForward(nn.Module):
method __init__ (line 96) | def __init__(
method forward (line 131) | def forward(self, hidden_states):
class BasicTransformerBlock (line 138) | class BasicTransformerBlock(nn.Module):
method __init__ (line 159) | def __init__(
method set_chunk_feed_forward (line 238) | def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int):
method forward (line 243) | def forward(
FILE: src/chatterbox/models/s3gen/s3gen.py
function drop_invalid_tokens (line 36) | def drop_invalid_tokens(x):
function get_resampler (line 43) | def get_resampler(src_sr, dst_sr, device):
class S3Token2Mel (line 47) | class S3Token2Mel(torch.nn.Module):
method __init__ (line 53) | def __init__(self, meanflow=False):
method device (line 109) | def device(self):
method dtype (line 114) | def dtype(self):
method embed_ref (line 118) | def embed_ref(
method forward (line 173) | def forward(
class S3Token2Wav (line 232) | class S3Token2Wav(S3Token2Mel):
method __init__ (line 241) | def __init__(self, meanflow=False):
method forward (line 261) | def forward(
method flow_inference (line 301) | def flow_inference(
method hift_inference (line 324) | def hift_inference(self, speech_feat, cache_source: torch.Tensor = None):
method inference (line 330) | def inference(
FILE: src/chatterbox/models/s3gen/transformer/activation.py
class Swish (line 24) | class Swish(torch.nn.Module):
method forward (line 27) | def forward(self, x: torch.Tensor) -> torch.Tensor:
class Snake (line 34) | class Snake(nn.Module):
method __init__ (line 50) | def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha...
method forward (line 73) | def forward(self, x):
FILE: src/chatterbox/models/s3gen/transformer/attention.py
class MultiHeadedAttention (line 26) | class MultiHeadedAttention(nn.Module):
method __init__ (line 36) | def __init__(self,
method forward_qkv (line 53) | def forward_qkv(
method forward_attention (line 82) | def forward_attention(
method forward (line 129) | def forward(
class RelPositionMultiHeadedAttention (line 200) | class RelPositionMultiHeadedAttention(MultiHeadedAttention):
method __init__ (line 209) | def __init__(self,
method rel_shift (line 225) | def rel_shift(self, x: torch.Tensor) -> torch.Tensor:
method forward (line 249) | def forward(
FILE: src/chatterbox/models/s3gen/transformer/convolution.py
class ConvolutionModule (line 24) | class ConvolutionModule(nn.Module):
method __init__ (line 27) | def __init__(self,
method forward (line 90) | def forward(
FILE: src/chatterbox/models/s3gen/transformer/embedding.py
class PositionalEncoding (line 26) | class PositionalEncoding(torch.nn.Module):
method __init__ (line 37) | def __init__(self,
method forward (line 59) | def forward(self,
method position_encoding (line 79) | def position_encoding(self,
class RelPositionalEncoding (line 120) | class RelPositionalEncoding(PositionalEncoding):
method __init__ (line 129) | def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5...
method forward (line 133) | def forward(self,
class WhisperPositionalEncoding (line 150) | class WhisperPositionalEncoding(PositionalEncoding):
method __init__ (line 154) | def __init__(self, d_model: int, dropout_rate: float, max_len: int = 1...
class LearnablePositionalEncoding (line 167) | class LearnablePositionalEncoding(PositionalEncoding):
method __init__ (line 171) | def __init__(self, d_model: int, dropout_rate: float, max_len: int = 4...
class NoPositionalEncoding (line 178) | class NoPositionalEncoding(torch.nn.Module):
method __init__ (line 182) | def __init__(self, d_model: int, dropout_rate: float):
method forward (line 187) | def forward(self,
method position_encoding (line 196) | def position_encoding(self, offset: Union[int, torch.Tensor],
class EspnetRelPositionalEncoding (line 201) | class EspnetRelPositionalEncoding(torch.nn.Module):
method __init__ (line 215) | def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5...
method extend_pe (line 224) | def extend_pe(self, x: torch.Tensor):
method forward (line 256) | def forward(self, x: torch.Tensor, offset: Union[int, torch.Tensor] = ...
method position_encoding (line 272) | def position_encoding(self,
FILE: src/chatterbox/models/s3gen/transformer/encoder_layer.py
class TransformerEncoderLayer (line 24) | class TransformerEncoderLayer(nn.Module):
method __init__ (line 40) | def __init__(
method forward (line 58) | def forward(
class ConformerEncoderLayer (line 109) | class ConformerEncoderLayer(nn.Module):
method __init__ (line 129) | def __init__(
method forward (line 160) | def forward(
FILE: src/chatterbox/models/s3gen/transformer/positionwise_feed_forward.py
class PositionwiseFeedForward (line 20) | class PositionwiseFeedForward(torch.nn.Module):
method __init__ (line 33) | def __init__(
method forward (line 47) | def forward(self, xs: torch.Tensor) -> torch.Tensor:
class MoEFFNLayer (line 58) | class MoEFFNLayer(torch.nn.Module):
method __init__ (line 75) | def __init__(
method forward (line 91) | def forward(self, xs: torch.Tensor) -> torch.Tensor:
FILE: src/chatterbox/models/s3gen/transformer/subsampling.py
class BaseSubsampling (line 23) | class BaseSubsampling(torch.nn.Module):
method __init__ (line 25) | def __init__(self):
method position_encoding (line 30) | def position_encoding(self, offset: Union[int, torch.Tensor],
class EmbedinigNoSubsampling (line 35) | class EmbedinigNoSubsampling(BaseSubsampling):
method __init__ (line 39) | def __init__(self, idim: int, odim: int, dropout_rate: float,
method forward (line 45) | def forward(
class LinearNoSubsampling (line 69) | class LinearNoSubsampling(BaseSubsampling):
method __init__ (line 79) | def __init__(self, idim: int, odim: int, dropout_rate: float,
method forward (line 92) | def forward(
class Conv1dSubsampling2 (line 116) | class Conv1dSubsampling2(BaseSubsampling):
method __init__ (line 128) | def __init__(self, idim: int, odim: int, dropout_rate: float,
method forward (line 145) | def forward(
class Conv2dSubsampling4 (line 173) | class Conv2dSubsampling4(BaseSubsampling):
method __init__ (line 183) | def __init__(self, idim: int, odim: int, dropout_rate: float,
method forward (line 202) | def forward(
class Conv2dSubsampling6 (line 230) | class Conv2dSubsampling6(BaseSubsampling):
method __init__ (line 239) | def __init__(self, idim: int, odim: int, dropout_rate: float,
method forward (line 256) | def forward(
class Conv2dSubsampling8 (line 282) | class Conv2dSubsampling8(BaseSubsampling):
method __init__ (line 292) | def __init__(self, idim: int, odim: int, dropout_rate: float,
method forward (line 311) | def forward(
class LegacyLinearNoSubsampling (line 338) | class LegacyLinearNoSubsampling(BaseSubsampling):
method __init__ (line 348) | def __init__(self, idim: int, odim: int, dropout_rate: float,
method forward (line 362) | def forward(
FILE: src/chatterbox/models/s3gen/transformer/upsample_encoder.py
class Upsample1D (line 37) | class Upsample1D(nn.Module):
method __init__ (line 51) | def __init__(self, channels: int, out_channels: int, stride: int = 2):
method forward (line 59) | def forward(self, inputs: torch.Tensor, input_lengths: torch.Tensor):
class PreLookaheadLayer (line 66) | class PreLookaheadLayer(nn.Module):
method __init__ (line 67) | def __init__(self, channels: int, pre_lookahead_len: int = 1):
method forward (line 81) | def forward(self, inputs: torch.Tensor) -> torch.Tensor:
class UpsampleConformerEncoder (line 99) | class UpsampleConformerEncoder(torch.nn.Module):
method __init__ (line 101) | def __init__(
method output_size (line 234) | def output_size(self) -> int:
method forward (line 237) | def forward(
method forward_layers (line 306) | def forward_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor,
method forward_up_layers (line 313) | def forward_up_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor,
FILE: src/chatterbox/models/s3gen/utils/intmeanflow.py
function get_intmeanflow_time_mixer (line 5) | def get_intmeanflow_time_mixer(dims):
FILE: src/chatterbox/models/s3gen/utils/mask.py
function subsequent_chunk_mask (line 54) | def subsequent_chunk_mask(
function add_optional_chunk_mask (line 89) | def add_optional_chunk_mask(xs: torch.Tensor,
function make_pad_mask (line 167) | def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
FILE: src/chatterbox/models/s3gen/utils/mel.py
function dynamic_range_compression_torch (line 15) | def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
function spectral_normalize_torch (line 19) | def spectral_normalize_torch(magnitudes):
function mel_spectrogram (line 36) | def mel_spectrogram(y, n_fft=1920, num_mels=80, sampling_rate=24000, hop...
FILE: src/chatterbox/models/s3gen/xvector.py
function pad_list (line 15) | def pad_list(xs, pad_value):
function extract_feature (line 45) | def extract_feature(audio):
class BasicResBlock (line 61) | class BasicResBlock(torch.nn.Module):
method __init__ (line 64) | def __init__(self, in_planes, planes, stride=1):
method forward (line 86) | def forward(self, x):
class FCM (line 94) | class FCM(torch.nn.Module):
method __init__ (line 95) | def __init__(self, block=BasicResBlock, num_blocks=[2, 2], m_channels=...
method _make_layer (line 110) | def _make_layer(self, block, planes, num_blocks, stride):
method forward (line 118) | def forward(self, x):
function get_nonlinear (line 130) | def get_nonlinear(config_str, channels):
function statistics_pooling (line 146) | def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, eps=1e-2):
class StatsPool (line 155) | class StatsPool(torch.nn.Module):
method forward (line 156) | def forward(self, x):
class TDNNLayer (line 160) | class TDNNLayer(torch.nn.Module):
method __init__ (line 161) | def __init__(
method forward (line 189) | def forward(self, x):
class CAMLayer (line 195) | class CAMLayer(torch.nn.Module):
method __init__ (line 196) | def __init__(
method forward (line 214) | def forward(self, x):
method seg_pooling (line 221) | def seg_pooling(self, x, seg_len=100, stype="avg"):
class CAMDenseTDNNLayer (line 234) | class CAMDenseTDNNLayer(torch.nn.Module):
method __init__ (line 235) | def __init__(
method bn_function (line 266) | def bn_function(self, x):
method forward (line 269) | def forward(self, x):
class CAMDenseTDNNBlock (line 278) | class CAMDenseTDNNBlock(torch.nn.ModuleList):
method __init__ (line 279) | def __init__(
method forward (line 307) | def forward(self, x):
class TransitLayer (line 313) | class TransitLayer(torch.nn.Module):
method __init__ (line 314) | def __init__(self, in_channels, out_channels, bias=True, config_str="b...
method forward (line 319) | def forward(self, x):
class DenseLayer (line 325) | class DenseLayer(torch.nn.Module):
method __init__ (line 326) | def __init__(self, in_channels, out_channels, bias=False, config_str="...
method forward (line 331) | def forward(self, x):
class CAMPPlus (line 340) | class CAMPPlus(torch.nn.Module):
method __init__ (line 341) | def __init__(
method forward (line 417) | def forward(self, x):
method inference (line 425) | def inference(self, audio_list):
FILE: src/chatterbox/models/s3tokenizer/__init__.py
function drop_invalid_tokens (line 16) | def drop_invalid_tokens(x):
FILE: src/chatterbox/models/s3tokenizer/s3tokenizer.py
class S3Tokenizer (line 22) | class S3Tokenizer(S3TokenizerV2):
method __init__ (line 31) | def __init__(
method pad (line 54) | def pad(self, wavs, sr) -> List[torch.Tensor]:
method _prepare_audio (line 78) | def _prepare_audio(self, wavs):
method forward (line 91) | def forward(
method log_mel_spectrogram (line 128) | def log_mel_spectrogram(
FILE: src/chatterbox/models/t3/inference/alignment_stream_analyzer.py
class AlignmentAnalysisResult (line 17) | class AlignmentAnalysisResult:
class AlignmentStreamAnalyzer (line 32) | class AlignmentStreamAnalyzer:
method __init__ (line 33) | def __init__(self, tfmr, queue, text_tokens_slice, alignment_layer_idx...
method _add_attention_spy (line 67) | def _add_attention_spy(self, tfmr, buffer_idx, layer_idx, head_idx):
method step (line 92) | def step(self, logits, next_token=None):
FILE: src/chatterbox/models/t3/inference/t3_hf_backend.py
class T3HuggingfaceBackend (line 9) | class T3HuggingfaceBackend(LlamaPreTrainedModel, GenerationMixin):
method __init__ (line 17) | def __init__(
method prepare_inputs_for_generation (line 36) | def prepare_inputs_for_generation(
method forward (line 73) | def forward(
FILE: src/chatterbox/models/t3/modules/cond_enc.py
class T3Cond (line 12) | class T3Cond:
method to (line 24) | def to(self, *, device=None, dtype=None):
method save (line 32) | def save(self, fpath):
method load (line 36) | def load(fpath, map_location="cpu"):
class T3CondEnc (line 41) | class T3CondEnc(nn.Module):
method __init__ (line 46) | def __init__(self, hp: T3Config):
method forward (line 64) | def forward(self, cond: T3Cond):
FILE: src/chatterbox/models/t3/modules/learned_pos_emb.py
class LearnedPositionEmbeddings (line 7) | class LearnedPositionEmbeddings(nn.Module):
method __init__ (line 8) | def __init__(self, seq_len, model_dim, init=.02):
method forward (line 14) | def forward(self, x):
method get_fixed_embedding (line 21) | def get_fixed_embedding(self, idx: 'Union[int, Tensor]'):
FILE: src/chatterbox/models/t3/modules/perceiver.py
class RelativePositionBias (line 12) | class RelativePositionBias(nn.Module):
method __init__ (line 13) | def __init__(self, scale, causal=False, num_buckets=32, max_distance=1...
method _relative_position_bucket (line 22) | def _relative_position_bucket(relative_position, causal=True, num_buck...
method forward (line 43) | def forward(self, qk_dots):
class AttentionQKV (line 55) | class AttentionQKV(nn.Module):
method __init__ (line 56) | def __init__(self, n_heads, head_dim, dropout_rate=0.1, scale=None, fl...
method setup_flash_config (line 66) | def setup_flash_config(self):
method forward (line 75) | def forward(self, q, k, v, mask=None):
method scaled_dot_product_attention (line 84) | def scaled_dot_product_attention(self, q, k, v, mask=None):
method flash_attention (line 92) | def flash_attention(self, q, k, v, mask=None):
method split_heads (line 102) | def split_heads(self, x):
method combine_heads (line 107) | def combine_heads(self, x):
class AttentionBlock2 (line 113) | class AttentionBlock2(nn.Module):
method __init__ (line 119) | def __init__(
method forward (line 156) | def forward(self, x1, x2, mask=None):
class Perceiver (line 173) | class Perceiver(nn.Module):
method __init__ (line 175) | def __init__(self, pre_attention_query_token=32, pre_attention_query_s...
method forward (line 200) | def forward(self, h):
FILE: src/chatterbox/models/t3/modules/t3_config.py
class T3Config (line 4) | class T3Config:
method __init__ (line 5) | def __init__(self, text_tokens_dict_size=704):
method n_channels (line 26) | def n_channels(self):
method is_multilingual (line 30) | def is_multilingual(self):
method english_only (line 34) | def english_only(cls):
method multilingual (line 39) | def multilingual(cls):
FILE: src/chatterbox/models/t3/t3.py
function _ensure_BOT_EOT (line 34) | def _ensure_BOT_EOT(text_tokens: Tensor, hp):
class T3 (line 40) | class T3(nn.Module):
method __init__ (line 50) | def __init__(self, hp=None):
method device (line 90) | def device(self):
method prepare_conditioning (line 93) | def prepare_conditioning(self, t3_cond: T3Cond):
method prepare_input_embeds (line 103) | def prepare_input_embeds(
method forward (line 133) | def forward(
method loss (line 190) | def loss(
method inference (line 227) | def inference(
method inference_turbo (line 415) | def inference_turbo(self, t3_cond, text_tokens, temperature=0.8, top_k...
FILE: src/chatterbox/models/tokenizers/tokenizer.py
class EnTokenizer (line 20) | class EnTokenizer:
method __init__ (line 21) | def __init__(self, vocab_file_path):
method check_vocabset_sot_eot (line 25) | def check_vocabset_sot_eot(self):
method text_to_tokens (line 30) | def text_to_tokens(self, text: str):
method encode (line 35) | def encode(self, txt: str):
method decode (line 44) | def decode(self, seq):
function is_kanji (line 65) | def is_kanji(c: str) -> bool:
function is_katakana (line 70) | def is_katakana(c: str) -> bool:
function hiragana_normalize (line 75) | def hiragana_normalize(text: str) -> str:
function add_hebrew_diacritics (line 117) | def add_hebrew_diacritics(text: str) -> str:
function korean_normalize (line 136) | def korean_normalize(text: str) -> str:
class ChineseCangjieConverter (line 157) | class ChineseCangjieConverter:
method __init__ (line 160) | def __init__(self, model_dir=None):
method _load_cangjie_mapping (line 167) | def _load_cangjie_mapping(self, model_dir=None):
method _init_segmenter (line 190) | def _init_segmenter(self):
method _cangjie_encode (line 199) | def _cangjie_encode(self, glyph: str):
method __call__ (line 211) | def __call__(self, text):
function add_russian_stress (line 237) | def add_russian_stress(text: str) -> str:
class MTLTokenizer (line 256) | class MTLTokenizer:
method __init__ (line 257) | def __init__(self, vocab_file_path):
method check_vocabset_sot_eot (line 263) | def check_vocabset_sot_eot(self):
method preprocess_text (line 268) | def preprocess_text(self, raw_text: str, language_id: str = None, lowe...
method text_to_tokens (line 280) | def text_to_tokens(self, text: str, language_id: str = None, lowercase...
method encode (line 285) | def encode(self, txt: str, language_id: str = None, lowercase: bool = ...
method decode (line 307) | def decode(self, seq):
FILE: src/chatterbox/models/utils.py
class AttrDict (line 1) | class AttrDict(dict):
method __init__ (line 2) | def __init__(self, *args, **kwargs):
FILE: src/chatterbox/models/voice_encoder/config.py
class VoiceEncConfig (line 1) | class VoiceEncConfig:
FILE: src/chatterbox/models/voice_encoder/melspec.py
function mel_basis (line 9) | def mel_basis(hp):
function preemphasis (line 19) | def preemphasis(wav, hp):
function melspectrogram (line 26) | def melspectrogram(wav, hp, pad=True):
function _stft (line 54) | def _stft(y, hp, pad=True):
function _amp_to_db (line 67) | def _amp_to_db(x, hp):
function _db_to_amp (line 71) | def _db_to_amp(x):
function _normalize (line 75) | def _normalize(s, hp, headroom_db=15):
FILE: src/chatterbox/models/voice_encoder/voice_encoder.py
function pack (line 16) | def pack(arrays, seq_len: int=None, pad_value=0):
function get_num_wins (line 54) | def get_num_wins(
function get_frame_step (line 69) | def get_frame_step(
function stride_as_partials (line 84) | def stride_as_partials(
class VoiceEncoder (line 119) | class VoiceEncoder(nn.Module):
method __init__ (line 120) | def __init__(self, hp=VoiceEncConfig()):
method device (line 136) | def device(self):
method forward (line 139) | def forward(self, mels: torch.FloatTensor):
method inference (line 162) | def inference(self, mels: torch.Tensor, mel_lens, overlap=0.5, rate: f...
method utt_to_spk_embed (line 202) | def utt_to_spk_embed(utt_embeds: np.ndarray):
method voice_similarity (line 212) | def voice_similarity(embeds_x: np.ndarray, embeds_y: np.ndarray):
method embeds_from_mels (line 220) | def embeds_from_mels(
method embeds_from_wavs (line 246) | def embeds_from_wavs(
FILE: src/chatterbox/mtl_tts.py
function punc_norm (line 51) | def punc_norm(text: str) -> str:
class Conditionals (line 94) | class Conditionals:
method to (line 113) | def to(self, device):
method save (line 120) | def save(self, fpath: Path):
method load (line 128) | def load(cls, fpath, map_location="cpu"):
class ChatterboxMultilingualTTS (line 135) | class ChatterboxMultilingualTTS:
method __init__ (line 139) | def __init__(
method get_supported_languages (line 158) | def get_supported_languages(cls):
method from_local (line 163) | def from_local(cls, ckpt_dir, device) -> 'ChatterboxMultilingualTTS':
method from_pretrained (line 202) | def from_pretrained(cls, device: torch.device) -> 'ChatterboxMultiling...
method prepare_conditionals (line 222) | def prepare_conditionals(self, wav_fpath, exaggeration=0.5):
method generate (line 249) | def generate(
FILE: src/chatterbox/tts.py
function punc_norm (line 22) | def punc_norm(text: str) -> str:
class Conditionals (line 65) | class Conditionals:
method to (line 84) | def to(self, device):
method save (line 91) | def save(self, fpath: Path):
method load (line 99) | def load(cls, fpath, map_location="cpu"):
class ChatterboxTTS (line 106) | class ChatterboxTTS:
method __init__ (line 110) | def __init__(
method from_local (line 129) | def from_local(cls, ckpt_dir, device) -> 'ChatterboxTTS':
method from_pretrained (line 168) | def from_pretrained(cls, device) -> 'ChatterboxTTS':
method prepare_conditionals (line 182) | def prepare_conditionals(self, wav_fpath, exaggeration=0.5):
method generate (line 208) | def generate(
FILE: src/chatterbox/tts_turbo.py
function punc_norm (line 29) | def punc_norm(text: str) -> str:
class Conditionals (line 69) | class Conditionals:
method to (line 88) | def to(self, device):
method save (line 95) | def save(self, fpath: Path):
method load (line 103) | def load(cls, fpath, map_location="cpu"):
class ChatterboxTurboTTS (line 110) | class ChatterboxTurboTTS:
method __init__ (line 114) | def __init__(
method from_local (line 133) | def from_local(cls, ckpt_dir, device) -> 'ChatterboxTurboTTS':
method from_pretrained (line 186) | def from_pretrained(cls, device) -> 'ChatterboxTurboTTS':
method norm_loudness (line 204) | def norm_loudness(self, wav, sr, target_lufs=-27):
method prepare_conditionals (line 217) | def prepare_conditionals(self, wav_fpath, exaggeration=0.5, norm_loudn...
method generate (line 248) | def generate(
FILE: src/chatterbox/vc.py
class ChatterboxVC (line 16) | class ChatterboxVC:
method __init__ (line 20) | def __init__(
method from_local (line 39) | def from_local(cls, ckpt_dir, device) -> 'ChatterboxVC':
method from_pretrained (line 62) | def from_pretrained(cls, device) -> 'ChatterboxVC':
method set_target_voice (line 76) | def set_target_voice(self, wav_fpath):
method generate (line 83) | def generate(
Condensed preview — 64 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (380K chars).
[
{
"path": ".github/workflows/install_check.yml",
"chars": 382,
"preview": "name: Test Installation\n\non:\n push:\n branches: [ \"master\" ]\n pull_request:\n branches: [ \"master\" ]\n\njobs:\n buil"
},
{
"path": ".gitignore",
"chars": 623,
"preview": ".vscode\n\n# Pylance\npyrightconfig.json\n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C ex"
},
{
"path": "LICENSE",
"chars": 1067,
"preview": "MIT License\n\nCopyright (c) 2025 Resemble AI\n\nPermission is hereby granted, free of charge, to any person obtaining a cop"
},
{
"path": "README.md",
"chars": 9918,
"preview": "\n\n\n# Chatterbox TTS\n\n[![Alt Text](https://img.shields.io/badge/listen-d"
},
{
"path": "example_for_mac.py",
"chars": 1009,
"preview": "import torch\nimport torchaudio as ta\nfrom chatterbox.tts import ChatterboxTTS\n\n# Detect device (Mac with M1/M2/M3/M4)\nde"
},
{
"path": "example_tts.py",
"chars": 1309,
"preview": "import torchaudio as ta\nimport torch\nfrom pathlib import Path\nfrom chatterbox.tts import ChatterboxTTS\nfrom chatterbox.m"
},
{
"path": "example_tts_turbo.py",
"chars": 664,
"preview": "import torchaudio as ta\nimport torch\nfrom chatterbox.tts_turbo import ChatterboxTurboTTS\n\n# Load the Turbo model\nmodel ="
},
{
"path": "example_vc.py",
"chars": 531,
"preview": "import torch\nimport torchaudio as ta\n\nfrom chatterbox.vc import ChatterboxVC\n\n# Automatically detect the best available "
},
{
"path": "gradio_tts_app.py",
"chars": 3216,
"preview": "import random\nimport numpy as np\nimport torch\nimport gradio as gr\nfrom chatterbox.tts import ChatterboxTTS\n\n\nDEVICE = \"c"
},
{
"path": "gradio_tts_turbo_app.py",
"chars": 5693,
"preview": "import random\nimport numpy as np\nimport torch\nimport gradio as gr\nfrom chatterbox.tts_turbo import ChatterboxTurboTTS\n\nD"
},
{
"path": "gradio_vc_app.py",
"chars": 702,
"preview": "import torch\nimport gradio as gr\nfrom chatterbox.vc import ChatterboxVC\n\n\nDEVICE = \"cuda\" if torch.cuda.is_available() e"
},
{
"path": "multilingual_app.py",
"chars": 13370,
"preview": "import random\nimport numpy as np\nimport torch\nfrom chatterbox.mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUA"
},
{
"path": "pyproject.toml",
"chars": 961,
"preview": "[project]\nname = \"chatterbox-tts\"\nversion = \"0.1.6\"\ndescription = \"Chatterbox: Open Source TTS and Voice Conversion by R"
},
{
"path": "src/chatterbox/__init__.py",
"chars": 300,
"preview": "try:\n from importlib.metadata import version\nexcept ImportError:\n from importlib_metadata import version # For Py"
},
{
"path": "src/chatterbox/models/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "src/chatterbox/models/s3gen/__init__.py",
"chars": 68,
"preview": "from .s3gen import S3Token2Wav as S3Gen\nfrom .const import S3GEN_SR\n"
},
{
"path": "src/chatterbox/models/s3gen/configs.py",
"chars": 220,
"preview": "from ..utils import AttrDict\n\nCFM_PARAMS = AttrDict({\n \"sigma_min\": 1e-06,\n \"solver\": \"euler\",\n \"t_scheduler\": "
},
{
"path": "src/chatterbox/models/s3gen/const.py",
"chars": 34,
"preview": "S3GEN_SR = 24000\nS3GEN_SIL = 4299\n"
},
{
"path": "src/chatterbox/models/s3gen/decoder.py",
"chars": 13171,
"preview": "# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du)\n#\n# Licensed under the Apache License, Version 2.0 (the"
},
{
"path": "src/chatterbox/models/s3gen/f0_predictor.py",
"chars": 1993,
"preview": "# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu)\n#\n# Licensed under the Apache License, Version 2.0 (the \"L"
},
{
"path": "src/chatterbox/models/s3gen/flow.py",
"chars": 8178,
"preview": "# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du)\n#\n# Licensed under the Apache License, Version 2.0 (the"
},
{
"path": "src/chatterbox/models/s3gen/flow_matching.py",
"chars": 10786,
"preview": "# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du)\n#\n# Licensed under the Apache License, Version 2.0 (the"
},
{
"path": "src/chatterbox/models/s3gen/hifigan.py",
"chars": 17650,
"preview": "# jrm: adapted from CosyVoice/cosyvoice/hifigan/generator.py\n# most modules should be reusable, but I found their S"
},
{
"path": "src/chatterbox/models/s3gen/matcha/decoder.py",
"chars": 14435,
"preview": "import math\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom conform"
},
{
"path": "src/chatterbox/models/s3gen/matcha/flow_matching.py",
"chars": 4556,
"preview": "from abc import ABC\n\nimport torch\nimport torch.nn.functional as F\n\nfrom .decoder import Decoder\n\n\nclass BASECFM(torch.nn"
},
{
"path": "src/chatterbox/models/s3gen/matcha/text_encoder.py",
"chars": 14965,
"preview": "\"\"\" from https://github.com/jaywalnut310/glow-tts \"\"\"\n\nimport math\n\nimport torch\nimport torch.nn as nn\nfrom einops impor"
},
{
"path": "src/chatterbox/models/s3gen/matcha/transformer.py",
"chars": 13235,
"preview": "from typing import Any, Dict, Optional\n\nimport torch\nimport torch.nn as nn\nfrom diffusers.models.attention import (\n "
},
{
"path": "src/chatterbox/models/s3gen/s3gen.py",
"chars": 13189,
"preview": "# Modified from CosyVoice https://github.com/FunAudioLLM/CosyVoice\n#\n# Licensed under the Apache License, Version 2.0 (t"
},
{
"path": "src/chatterbox/models/s3gen/transformer/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "src/chatterbox/models/s3gen/transformer/activation.py",
"chars": 3087,
"preview": "# Copyright (c) 2020 Johns Hopkins University (Shinji Watanabe)\n# 2020 Northwestern Polytechnical Universi"
},
{
"path": "src/chatterbox/models/s3gen/transformer/attention.py",
"chars": 14415,
"preview": "# Copyright (c) 2019 Shigeki Karita\n# 2020 Mobvoi Inc (Binbin Zhang)\n# 2022 Xingchen Song (s"
},
{
"path": "src/chatterbox/models/s3gen/transformer/convolution.py",
"chars": 5230,
"preview": "# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Di Wu)\n# 2024 Alibaba Inc (Xiang Lyu)\n#\n# License"
},
{
"path": "src/chatterbox/models/s3gen/transformer/embedding.py",
"chars": 11399,
"preview": "# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Di Wu)\n# 2024 Alibaba Inc (Xiang Lyu)\n#\n# License"
},
{
"path": "src/chatterbox/models/s3gen/transformer/encoder_layer.py",
"chars": 9596,
"preview": "# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu)\n# 2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn)\n#"
},
{
"path": "src/chatterbox/models/s3gen/transformer/positionwise_feed_forward.py",
"chars": 4219,
"preview": "# Copyright (c) 2019 Shigeki Karita\n# 2020 Mobvoi Inc (Binbin Zhang)\n#\n# Licensed under the Apache License"
},
{
"path": "src/chatterbox/models/s3gen/transformer/subsampling.py",
"chars": 12666,
"preview": "# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu)\n# 2024 Alibaba Inc (Xiang Lyu)\n#\n# Licensed under th"
},
{
"path": "src/chatterbox/models/s3gen/transformer/upsample_encoder.py",
"chars": 13694,
"preview": "# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu)\n# 2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn)\n#"
},
{
"path": "src/chatterbox/models/s3gen/utils/class_utils.py",
"chars": 2334,
"preview": "# Copyright [2023-11-28] <sxc19@mails.tsinghua.edu.cn, Xingchen Song>\n# 2024 Alibaba Inc (authors: Xiang Lyu)"
},
{
"path": "src/chatterbox/models/s3gen/utils/intmeanflow.py",
"chars": 951,
"preview": "import torch\nimport torch.nn as nn\n\n\ndef get_intmeanflow_time_mixer(dims):\n \"\"\"\"\n Diagonal init as described in 3."
},
{
"path": "src/chatterbox/models/s3gen/utils/mask.py",
"chars": 7666,
"preview": "# Copyright (c) 2019 Shigeki Karita\n# 2020 Mobvoi Inc (Binbin Zhang)\n# 2024 Alibaba Inc (aut"
},
{
"path": "src/chatterbox/models/s3gen/utils/mel.py",
"chars": 2556,
"preview": "\"\"\"mel-spectrogram extraction in Matcha-TTS\"\"\"\nimport logging\nfrom librosa.filters import mel as librosa_mel_fn\nimport t"
},
{
"path": "src/chatterbox/models/s3gen/xvector.py",
"chars": 14141,
"preview": "#!/usr/bin/env python3\n# -*- encoding: utf-8 -*-\n# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). Al"
},
{
"path": "src/chatterbox/models/s3tokenizer/__init__.py",
"chars": 595,
"preview": "from .s3tokenizer import (\n S3_SR,\n S3_HOP,\n S3_TOKEN_HOP,\n S3_TOKEN_RATE,\n SPEECH_VOCAB_SIZE,\n S3Toke"
},
{
"path": "src/chatterbox/models/s3tokenizer/s3tokenizer.py",
"chars": 5155,
"preview": "from typing import List, Tuple\n\nimport numpy as np\nimport librosa\nimport torch\nimport torch.nn.functional as F\nfrom s3to"
},
{
"path": "src/chatterbox/models/t3/__init__.py",
"chars": 19,
"preview": "from .t3 import T3\n"
},
{
"path": "src/chatterbox/models/t3/inference/alignment_stream_analyzer.py",
"chars": 8416,
"preview": "# Copyright (c) 2025 Resemble AI\n# Author: John Meade, Jeremy Hsu\n# MIT License\nimport logging\nimport torch\nfrom datacla"
},
{
"path": "src/chatterbox/models/t3/inference/t3_hf_backend.py",
"chars": 4271,
"preview": "from typing import Optional\n\nimport torch\nfrom torch import nn as nn\nfrom transformers import LlamaConfig, LlamaModel, L"
},
{
"path": "src/chatterbox/models/t3/llama_configs.py",
"chars": 1809,
"preview": "LLAMA_520M_CONFIG_DICT = dict(\n # Arbitrary small number that won't cause problems when loading.\n # These param ar"
},
{
"path": "src/chatterbox/models/t3/modules/cond_enc.py",
"chars": 3293,
"preview": "from dataclasses import dataclass\nfrom typing import Optional\n\nimport torch\nfrom torch import nn, Tensor\n\nfrom .perceive"
},
{
"path": "src/chatterbox/models/t3/modules/learned_pos_emb.py",
"chars": 1099,
"preview": "from typing import Union\n\nimport torch\nfrom torch import nn, Tensor\n\n\nclass LearnedPositionEmbeddings(nn.Module):\n de"
},
{
"path": "src/chatterbox/models/t3/modules/perceiver.py",
"chars": 7719,
"preview": "# Copyright (c) 2025 Resemble AI\n# Author: Manmay Nakhashi\n# MIT License\nimport math\n\nimport torch\nfrom torch import nn\n"
},
{
"path": "src/chatterbox/models/t3/modules/t3_config.py",
"chars": 1245,
"preview": "from ..llama_configs import LLAMA_CONFIGS\n\n\nclass T3Config:\n def __init__(self, text_tokens_dict_size=704):\n s"
},
{
"path": "src/chatterbox/models/t3/t3.py",
"chars": 19429,
"preview": "# Copyright (c) 2025 Resemble AI\n# MIT License\nimport logging\nfrom typing import Union, Optional, List\n\nlogger = logging"
},
{
"path": "src/chatterbox/models/tokenizers/__init__.py",
"chars": 48,
"preview": "from .tokenizer import EnTokenizer, MTLTokenizer"
},
{
"path": "src/chatterbox/models/tokenizers/tokenizer.py",
"chars": 10198,
"preview": "import logging\nimport json\n\nimport torch\nfrom pathlib import Path\nfrom unicodedata import category, normalize\nfrom token"
},
{
"path": "src/chatterbox/models/utils.py",
"chars": 148,
"preview": "class AttrDict(dict):\n def __init__(self, *args, **kwargs):\n super(AttrDict, self).__init__(*args, **kwargs)\n "
},
{
"path": "src/chatterbox/models/voice_encoder/__init__.py",
"chars": 56,
"preview": "from .voice_encoder import VoiceEncoder, VoiceEncConfig\n"
},
{
"path": "src/chatterbox/models/voice_encoder/config.py",
"chars": 406,
"preview": "class VoiceEncConfig:\n num_mels = 40\n sample_rate = 16000\n speaker_embed_size = 256\n ve_hidden_size = 256\n "
},
{
"path": "src/chatterbox/models/voice_encoder/melspec.py",
"chars": 1940,
"preview": "from functools import lru_cache\n\nfrom scipy import signal\nimport numpy as np\nimport librosa\n\n\n@lru_cache()\ndef mel_basis"
},
{
"path": "src/chatterbox/models/voice_encoder/voice_encoder.py",
"chars": 10555,
"preview": "# Adapted from https://github.com/CorentinJ/Real-Time-Voice-Cloning\n# MIT License\nfrom typing import List, Union, Option"
},
{
"path": "src/chatterbox/mtl_tts.py",
"chars": 10411,
"preview": "from dataclasses import dataclass\nfrom pathlib import Path\nimport os\n\nimport librosa\nimport torch\nimport perth\nimport to"
},
{
"path": "src/chatterbox/tts.py",
"chars": 8910,
"preview": "from dataclasses import dataclass\nfrom pathlib import Path\n\nimport librosa\nimport torch\nimport perth\nimport torch.nn.fun"
},
{
"path": "src/chatterbox/tts_turbo.py",
"chars": 9814,
"preview": "import os\nimport math\nfrom dataclasses import dataclass\nfrom pathlib import Path\n\nimport librosa\nimport torch\nimport per"
},
{
"path": "src/chatterbox/vc.py",
"chars": 3520,
"preview": "from pathlib import Path\n\nimport librosa\nimport torch\nimport perth\nfrom huggingface_hub import hf_hub_download\nfrom safe"
}
]
About this extraction
This page contains the full source code of the resemble-ai/chatterbox GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 64 files (354.7 KB), approximately 91.6k tokens, and a symbol index with 440 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.