Fix optional Perth imports

Scope the torch.load map_location wrapper instead of replacing it globally

README.md

@@ -1,43 +1,129 @@
-# ComfyUI_Fill-ChatterBox
+# FL ChatterBox
 
-A custom node extension for ComfyUI that adds text-to-speech (TTS) and voice conversion (VC) capabilities using the Chatterbox library.
+High-quality text-to-speech nodes for ComfyUI powered by ResembleAI's Chatterbox models. Features voice cloning, multilingual synthesis, paralinguistic expressions, and voice conversion.
 
-![ChatterBox Example](web/image.png)
+[![Chatterbox](https://img.shields.io/badge/Chatterbox-Original%20Repo-blue?style=for-the-badge&logo=github&logoColor=white)](https://github.com/resemble-ai/chatterbox)
+[![Patreon](https://img.shields.io/badge/Patreon-Support%20Me-F96854?style=for-the-badge&logo=patreon&logoColor=white)](https://www.patreon.com/Machinedelusions)
+
+![Workflow Preview](assets/workflow_preview.png)
+
+## Features
+
+- **Zero-Shot Voice Cloning** - Clone any voice from a few seconds of reference audio
+- **3 TTS Models** - Standard, Turbo (faster), and Multilingual variants
+- **23 Languages** - Arabic, Chinese, Danish, Dutch, English, Finnish, French, German, Greek, Hebrew, Hindi, Italian, Japanese, Korean, Malay, Norwegian, Polish, Portuguese, Russian, Spanish, Swahili, Swedish, Turkish
+- **Paralinguistic Tags** - Express emotions with tags like `[laugh]`, `[sigh]`, `[gasp]`, `[chuckle]` (Turbo model)
+- **Voice Conversion** - Transform one voice to sound like another
+- **Dialog Synthesis** - Multi-speaker conversations with up to 4 voices
+- **Model Caching** - Keep models loaded between runs for faster iteration
+
+## Nodes
+
+| Node | Description |
+|------|-------------|
+| **FL Chatterbox TTS** | Standard high-quality text-to-speech with voice cloning |
+| **FL Chatterbox Turbo TTS** | Faster GPT2-based TTS with paralinguistic tag support |
+| **FL Chatterbox Multilingual TTS** | 23-language TTS with voice cloning |
+| **FL Chatterbox VC** | Voice conversion - transform source audio to target voice |
+| **FL Chatterbox Dialog TTS** | Multi-speaker dialog synthesis with up to 4 voices |
 
 ## Installation
 
-1. Clone this repository into your ComfyUI custom_nodes directory:
+### ComfyUI Manager
+Search for "FL ChatterBox" and install.
+
+### Manual
 ```bash
-   cd /path/to/ComfyUI/custom_nodes
+cd ComfyUI/custom_nodes
 git clone https://github.com/filliptm/ComfyUI_Fill-ChatterBox.git
+cd ComfyUI_Fill-ChatterBox
+pip install -r requirements.txt
 ```
 
-2. Install the base dependencies:
+### Optional: Watermarking Support
 ```bash
-   pip install -r ComfyUI_Fill-ChatterBox/requirements.txt
+pip install resemble-perth
 ```
+**Note**: The `resemble-perth` package may have compatibility issues with Python 3.12+. Nodes will function without watermarking if import fails.
 
-3. Install chatterbox-tts WITHOUT its dependencies:
+## Quick Start
-   ```bash
+
-   pip install chatterbox-tts --no-deps
+1. Add **FL Chatterbox TTS** (or Turbo/Multilingual variant)
+2. Enter your text in the text field
+3. Optionally connect reference audio for voice cloning
+4. Set `keep_model_loaded = True` for faster subsequent runs
+5. Generate!
+
+### Turbo Model with Expressions
 ```
+Hello there! [laugh] Isn't this amazing? [sigh] I just love text to speech.
+```
+Supported tags: `[laugh]`, `[sigh]`, `[gasp]`, `[chuckle]`, `[cough]`, `[sniff]`, `[groan]`, `[shush]`, `[clear throat]`
 
-   ⚠️ The `--no-deps` flag is crucial to prevent conflicts with ComfyUI's PyTorch installation!
+## Models
 
-## Usage
+| Model | Speed | Languages | Notes |
+|-------|-------|-----------|-------|
+| Standard | Normal | English | Highest quality |
+| Turbo | Fast | English | Paralinguistic tags, GPT2-based |
+| Multilingual | Normal | 23 languages | Cross-lingual voice cloning |
 
-### Text-to-Speech Node (FL Chatterbox TTS)
+Models download automatically on first use to `ComfyUI/models/chatterbox/`.
-- Add the "FL Chatterbox TTS" node to your workflow
-- Configure text input and parameters (exaggeration, cfg_weight, temperature)
-- Optionally provide an audio prompt for voice cloning
 
-### Voice Conversion Node (FL Chatterbox VC)
+## Parameters
-- Add the "FL Chatterbox VC" node to your workflow
-- Connect input audio and target voice
-- Both nodes support CPU fallback if CUDA errors occur
 
-## Change Log
+### TTS Parameters
+| Parameter | Range | Description |
+|-----------|-------|-------------|
+| `exaggeration` | 0.25-2.0 | Emotion intensity |
+| `cfg_weight` | 0.2-1.0 | Pace/classifier-free guidance |
+| `temperature` | 0.05-5.0 | Randomness in generation |
+| `seed` | 0-4.29B | Reproducible generation |
+| `keep_model_loaded` | bool | Cache model between runs |
 
-### 5/31/2025
+### Turbo Parameters
-- Added Persistent model loading, and loading bar functionality
+| Parameter | Range | Description |
-- Added Mac support (needs to be tested so HMU)
+|-----------|-------|-------------|
+| `temperature` | 0.05-2.0 | Randomness in generation |
+| `top_k` | 1-5000 | Top-k sampling |
+| `top_p` | 0.1-1.0 | Nucleus sampling threshold |
+| `repetition_penalty` | 1.0-3.0 | Token repetition penalty |
+
+## Limitations
+
+- Maximum audio length: ~40 seconds per generation
+- Reference audio: Minimum 5-6 seconds recommended
+- Turbo paralinguistic tags: English only
+
+## Requirements
+
+- Python 3.10+
+- 8GB RAM minimum (16GB+ recommended)
+- NVIDIA GPU with 8GB+ VRAM recommended
+- CPU and Mac MPS supported
+
+## License
+
+MIT License - See [Chatterbox repo](https://github.com/resemble-ai/chatterbox) for model licenses.
+
+## Changelog
+
+### 2025-12-28
+- Added Turbo TTS node (faster, GPT2-based with paralinguistic tags)
+- Added Multilingual TTS node (23 languages)
+- Improved model caching using module-level globals
+- Centralized model downloads to `ComfyUI/models/chatterbox/`
+
+### 2025-07-24
+- Added Dialog TTS node for multi-speaker conversations (up to 4 speakers)
+- Extended all nodes with seed parameters for reproducible generation
+- Isolated audio track outputs per speaker
+
+### 2025-06-24
+- Added seed parameter for reproducible generation
+- Made Perth watermarking optional for Python 3.12+ compatibility
+
+### 2025-05-31
+- Added persistent model loading and loading bar
+- Added Mac MPS support
+- Native inference code (removed chatterbox-tts library dependency)

__init__.py

@@ -1,4 +1,21 @@
-from .chatterbox_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+from .chatterbox_node import NODE_CLASS_MAPPINGS as BASE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as BASE_DISPLAY_NAME_MAPPINGS
+from .chatterbox_dialog_node import NODE_CLASS_MAPPINGS as DIALOG_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as DIALOG_DISPLAY_NAME_MAPPINGS
+
+NODE_CLASS_MAPPINGS = {}
+NODE_CLASS_MAPPINGS.update(BASE_CLASS_MAPPINGS)
+NODE_CLASS_MAPPINGS.update(DIALOG_CLASS_MAPPINGS)
+
+NODE_DISPLAY_NAME_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS.update(BASE_DISPLAY_NAME_MAPPINGS)
+NODE_DISPLAY_NAME_MAPPINGS.update(DIALOG_DISPLAY_NAME_MAPPINGS)
 
 WEB_DIRECTORY = "./web"
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
+
+# ASCII art banner
+ascii_art = """
+⣏⡉ ⡇    ⡎⠑ ⣇⡀ ⢀⣀ ⣰⡀ ⣰⡀ ⢀⡀ ⡀⣀ ⣏⡱ ⢀⡀ ⡀⢀
+⠇  ⠧⠤   ⠣⠔ ⠇⠸ ⠣⠼ ⠘⠤ ⠘⠤ ⠣⠭ ⠏  ⠧⠜ ⠣⠜ ⠜⠣
+"""
+print(f"\033[35m{ascii_art}\033[0m")
+print("FL ChatterBox Custom Nodes Loaded")

chatterbox_dialog_node.py

@@ -0,0 +1,197 @@
+import os
+import torch
+import torchaudio
+import tempfile
+
+from .local_chatterbox.chatterbox.tts import ChatterboxTTS
+from comfy.utils import ProgressBar
+
+class FL_ChatterboxDialogTTSNode:
+    """
+    TTS Node that accepts dialog with speaker labels and generates audio using separate voice prompts.
+    """
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "dialog_text": ("STRING", {"multiline": True, "default": "SPEAKER A: Test test\nSPEAKER B: 1 2 3"}),
+                "speaker_A_Audio": ("AUDIO",),
+                "speaker_B_Audio": ("AUDIO",),
+                "exaggeration": ("FLOAT", {"default": 0.5, "min": 0.25, "max": 2.0, "step": 0.05}),
+                "cfg_weight": ("FLOAT", {"default": 0.5, "min": 0.2, "max": 1.0, "step": 0.05}),
+                "temperature": ("FLOAT", {"default": 0.8, "min": 0.05, "max": 5.0, "step": 0.05}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 4294967295}),
+            },
+            "optional": {
+                "speaker_C_Audio": ("AUDIO",),
+                "speaker_D_Audio": ("AUDIO",),
+                "use_cpu": ("BOOLEAN", {"default": False}),
+                "keep_model_loaded": ("BOOLEAN", {"default": False}),
+            }
+        }
+
+    RETURN_TYPES = ("AUDIO", "AUDIO", "AUDIO", "AUDIO", "AUDIO", "STRING")
+    RETURN_NAMES = ("dialog_audio", "speaker_a_audio", "speaker_b_audio", "speaker_c_audio", "speaker_d_audio", "message")
+    FUNCTION = "generate_dialog"
+    CATEGORY = "ChatterBox"
+
+    _model = None
+    _device = None
+
+    def generate_dialog(self, dialog_text, speaker_A_Audio, speaker_B_Audio,
+                        exaggeration, cfg_weight, temperature, seed,
+                        speaker_C_Audio=None, speaker_D_Audio=None,
+                        use_cpu=False, keep_model_loaded=False):
+        # Set random seeds for reproducibility
+        torch.manual_seed(seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+        if torch.backends.mps.is_available():
+            torch.mps.manual_seed(seed)
+        import numpy as np
+        import random
+        np.random.seed(seed)
+        random.seed(seed)
+        
+        device = "cpu" if use_cpu else ("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
+        pbar = ProgressBar(100)
+        message = f"Running on {device}"
+
+        def save_temp_audio(audio_data):
+            path = tempfile.NamedTemporaryFile(suffix='.wav', delete=False).name
+            torchaudio.save(path, audio_data['waveform'].squeeze(0), audio_data['sample_rate'])
+            return path
+
+        prompt_a_path = save_temp_audio(speaker_A_Audio)
+        prompt_b_path = save_temp_audio(speaker_B_Audio)
+        temp_files = [prompt_a_path, prompt_b_path]
+        
+        # Handle optional speakers C and D
+        prompt_c_path = None
+        prompt_d_path = None
+        if speaker_C_Audio is not None:
+            prompt_c_path = save_temp_audio(speaker_C_Audio)
+            temp_files.append(prompt_c_path)
+        if speaker_D_Audio is not None:
+            prompt_d_path = save_temp_audio(speaker_D_Audio)
+            temp_files.append(prompt_d_path)
+
+        if self._model is None or self._device != device:
+            self._model = ChatterboxTTS.from_pretrained(device=device)
+            self._device = device
+        tts = self._model
+
+        lines = dialog_text.strip().splitlines()
+        speaker_a_waveforms = []
+        speaker_b_waveforms = []
+        speaker_c_waveforms = []
+        speaker_d_waveforms = []
+        combined_dialog_waveforms = []
+
+        for i, line in enumerate(lines):
+            wav = None
+            if line.startswith("SPEAKER A:"):
+                content = line[len("SPEAKER A:"):].strip()
+                prompt_path = prompt_a_path
+                pbar.update_absolute(int((i / len(lines)) * 80))
+                current_speaker_wav = tts.generate(
+                    text=content,
+                    audio_prompt_path=prompt_path,
+                    exaggeration=exaggeration,
+                    cfg_weight=cfg_weight,
+                    temperature=temperature
+                )
+                speaker_a_waveforms.append(current_speaker_wav)
+                combined_dialog_waveforms.append(current_speaker_wav)
+                # Add silence to other speakers' tracks
+                silence = torch.zeros_like(current_speaker_wav)
+                speaker_b_waveforms.append(silence)
+                speaker_c_waveforms.append(silence)
+                speaker_d_waveforms.append(silence)
+            elif line.startswith("SPEAKER B:"):
+                content = line[len("SPEAKER B:"):].strip()
+                prompt_path = prompt_b_path
+                pbar.update_absolute(int((i / len(lines)) * 80))
+                current_speaker_wav = tts.generate(
+                    text=content,
+                    audio_prompt_path=prompt_path,
+                    exaggeration=exaggeration,
+                    cfg_weight=cfg_weight,
+                    temperature=temperature
+                )
+                speaker_b_waveforms.append(current_speaker_wav)
+                combined_dialog_waveforms.append(current_speaker_wav)
+                # Add silence to other speakers' tracks
+                silence = torch.zeros_like(current_speaker_wav)
+                speaker_a_waveforms.append(silence)
+                speaker_c_waveforms.append(silence)
+                speaker_d_waveforms.append(silence)
+            elif line.startswith("SPEAKER C:") and prompt_c_path is not None:
+                content = line[len("SPEAKER C:"):].strip()
+                prompt_path = prompt_c_path
+                pbar.update_absolute(int((i / len(lines)) * 80))
+                current_speaker_wav = tts.generate(
+                    text=content,
+                    audio_prompt_path=prompt_path,
+                    exaggeration=exaggeration,
+                    cfg_weight=cfg_weight,
+                    temperature=temperature
+                )
+                speaker_c_waveforms.append(current_speaker_wav)
+                combined_dialog_waveforms.append(current_speaker_wav)
+                # Add silence to other speakers' tracks
+                silence = torch.zeros_like(current_speaker_wav)
+                speaker_a_waveforms.append(silence)
+                speaker_b_waveforms.append(silence)
+                speaker_d_waveforms.append(silence)
+            elif line.startswith("SPEAKER D:") and prompt_d_path is not None:
+                content = line[len("SPEAKER D:"):].strip()
+                prompt_path = prompt_d_path
+                pbar.update_absolute(int((i / len(lines)) * 80))
+                current_speaker_wav = tts.generate(
+                    text=content,
+                    audio_prompt_path=prompt_path,
+                    exaggeration=exaggeration,
+                    cfg_weight=cfg_weight,
+                    temperature=temperature
+                )
+                speaker_d_waveforms.append(current_speaker_wav)
+                combined_dialog_waveforms.append(current_speaker_wav)
+                # Add silence to other speakers' tracks
+                silence = torch.zeros_like(current_speaker_wav)
+                speaker_a_waveforms.append(silence)
+                speaker_b_waveforms.append(silence)
+                speaker_c_waveforms.append(silence)
+            else:
+                continue  # skip malformed line or missing prompt
+
+        if not combined_dialog_waveforms:
+            empty_audio = {"waveform": torch.zeros((1, 1, 1)), "sample_rate": tts.sr if tts else 16000}
+            return (empty_audio, empty_audio, empty_audio, empty_audio, empty_audio, "No valid dialog lines found.")
+
+        combined_waveform = torch.cat(combined_dialog_waveforms, dim=-1)
+        speaker_a_track = torch.cat(speaker_a_waveforms, dim=-1)
+        speaker_b_track = torch.cat(speaker_b_waveforms, dim=-1)
+        speaker_c_track = torch.cat(speaker_c_waveforms, dim=-1)
+        speaker_d_track = torch.cat(speaker_d_waveforms, dim=-1)
+
+        dialog_audio = {"waveform": combined_waveform.unsqueeze(0), "sample_rate": tts.sr}
+        speaker_a_audio = {"waveform": speaker_a_track.unsqueeze(0), "sample_rate": tts.sr}
+        speaker_b_audio = {"waveform": speaker_b_track.unsqueeze(0), "sample_rate": tts.sr}
+        speaker_c_audio = {"waveform": speaker_c_track.unsqueeze(0), "sample_rate": tts.sr}
+        speaker_d_audio = {"waveform": speaker_d_track.unsqueeze(0), "sample_rate": tts.sr}
+
+        for f in temp_files:
+            os.unlink(f)
+
+        pbar.update_absolute(100)
+        return (dialog_audio, speaker_a_audio, speaker_b_audio, speaker_c_audio, speaker_d_audio, "Dialog synthesized successfully.")
+
+NODE_CLASS_MAPPINGS = {
+    "FL_ChatterboxDialogTTS": FL_ChatterboxDialogTTSNode,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "FL_ChatterboxDialogTTS": "FL Chatterbox Dialog TTS",
+}

local_chatterbox/chatterbox/__init__.py

@@ -1,2 +1,4 @@
 from .tts import ChatterboxTTS
+from .tts_turbo import ChatterboxTurboTTS
+from .mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUAGES
 from .vc import ChatterboxVC

local_chatterbox/chatterbox/models/s3gen/configs.py

@@ -0,0 +1,10 @@
+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"
+})

local_chatterbox/chatterbox/models/s3gen/const.py

@@ -1 +1,2 @@
 S3GEN_SR = 24000
+S3GEN_SIL = 4299  # Silence token for Turbo model

local_chatterbox/chatterbox/models/s3gen/decoder.py

@@ -20,6 +20,7 @@ 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:
@@ -110,6 +111,7 @@ class ConditionalDecoder(nn.Module):
         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
@@ -117,6 +119,7 @@ class ConditionalDecoder(nn.Module):
         """
         super().__init__()
         channels = tuple(channels)
+        self.meanflow = meanflow
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.causal = causal
@@ -215,6 +218,14 @@ class ConditionalDecoder(nn.Module):
         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():
@@ -230,7 +241,7 @@ class ConditionalDecoder(nn.Module):
                 if m.bias is not None:
                     nn.init.constant_(m.bias, 0)
 
-    def forward(self, x, mask, mu, t, spks=None, cond=None):
+    def forward(self, x, mask, mu, t, spks=None, cond=None, r=None):
         """Forward pass of the UNet1DConditional model.
 
         Args:
@@ -239,6 +250,7 @@ class ConditionalDecoder(nn.Module):
             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.
+            r: end time for meanflow mode (shape (1,) tensor)
 
         Raises:
             ValueError: _description_
@@ -247,10 +259,15 @@ class ConditionalDecoder(nn.Module):
         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:

local_chatterbox/chatterbox/models/s3gen/flow.py

@@ -20,6 +20,24 @@ from torch.nn import functional as F
 from omegaconf import DictConfig
 from .utils.mask import make_pad_mask
 
+logger = logging.getLogger(__name__)
+
+
+def _repeat_batch_dim(x, B, ndim):
+    """Repeat tensor along batch dimension if needed."""
+    if x is None:
+        return x
+    if x.size(0) == B:
+        return x
+    if x.size(0) == 1:
+        if ndim == 1:
+            return x.repeat(B)
+        elif ndim == 2:
+            return x.repeat(B, 1)
+        elif ndim == 3:
+            return x.repeat(B, 1, 1)
+    raise ValueError(f"Cannot repeat tensor of size {x.size(0)} to batch size {B}")
+
 
 class MaskedDiffWithXvec(torch.nn.Module):
     def __init__(self,
@@ -202,41 +220,62 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
                   prompt_feat,
                   prompt_feat_len,
                   embedding,
-                  finalize):
+                  finalize,
-        if self.fp16 is True:
+                  n_timesteps=10,
-            prompt_feat = prompt_feat.half()
+                  noised_mels=None,
-            embedding = embedding.half()
+                  meanflow=False):
+        # token: (B, n_toks)
+        # token_len: (B,)
+        B = token.size(0)
 
-        assert token.shape[0] == 1
         # xvec projection
+        embedding = torch.atleast_2d(embedding)
         embedding = F.normalize(embedding, dim=1)
-        embedding = self.spk_embed_affine_layer(embedding)
+        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
+        token, token_len = torch.concat([prompt_token, token], dim=1), (prompt_token_len + token_len).long()
         mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding)
-        token = self.input_embedding(torch.clamp(token, min=0)) * mask
+
+        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_lengths = self.encoder(token, token_len)
+        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).long()
         mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1]
         h = self.encoder_proj(h)
 
-        # get conditions
+        # # get conditions
-        conds = torch.zeros([1, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype)
+        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(torch.tensor([mel_len1 + mel_len2]))).to(h)
+        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.unsqueeze(1),
+            mask=mask,
             spks=embedding,
             cond=conds,
-            n_timesteps=10
+            n_timesteps=n_timesteps,
+            noised_mels=noised_mels,
+            meanflow=meanflow,
         )
         feat = feat[:, :, mel_len1:]
         assert feat.shape[2] == mel_len2
-        return feat.float(), None  # NOTE jrm: why are they returning None here?
+        return feat, None  # NOTE jrm: why are they returning None here?

local_chatterbox/chatterbox/models/s3gen/flow_matching.py

@@ -14,10 +14,16 @@
 import threading
 import torch
 import torch.nn.functional as F
+from tqdm import tqdm
 from .matcha.flow_matching import BASECFM
 from omegaconf import OmegaConf
 
 
+def cast_all(*tensors, dtype):
+    """Cast all tensors to specified dtype."""
+    return [t.to(dtype) if t is not None else None for t in tensors]
+
+
 CFM_PARAMS = OmegaConf.create({
     "sigma_min": 1e-06,
     "solver": "euler",
@@ -198,10 +204,11 @@ class ConditionalCFM(BASECFM):
 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)
-        self.rand_noise = torch.randn([1, 80, 50 * 300])
+        # 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):
+    def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, noised_mels=None, meanflow=False):
         """Forward diffusion
 
         Args:
@@ -214,15 +221,43 @@ class CausalConditionalCFM(ConditionalCFM):
             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 at time t
+            meanflow: whether to use meanflow mode
 
         Returns:
             sample: generated mel-spectrogram
                 shape: (batch_size, n_feats, mel_timesteps)
         """
 
-        z = self.rand_noise[:, :, :mu.size(2)].to(mu.device).to(mu.dtype) * temperature
+        B = mu.size(0)
-        # fix prompt and overlap part mu and z
+        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 self.t_scheduler == 'cosine':
+        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), 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)

local_chatterbox/chatterbox/models/s3gen/s3gen.py

@@ -19,7 +19,6 @@ import torch
 import torchaudio as ta
 from functools import lru_cache
 from typing import Optional
-from omegaconf import DictConfig
 
 from ..s3tokenizer import S3_SR, SPEECH_VOCAB_SIZE, S3Tokenizer
 from .const import S3GEN_SR
@@ -31,6 +30,7 @@ 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):
@@ -46,15 +46,20 @@ def get_resampler(src_sr, dst_sr, device):
 
 class S3Token2Mel(torch.nn.Module):
     """
-    CosyVoice2's CFM decoder maps S3 speech tokens to mel-spectrograms.
+    S3Gen's CFM decoder maps S3 speech tokens to mel-spectrograms.
 
     TODO: make these modules configurable?
     """
-    def __init__(self):
+    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()  # use default args
+        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,
@@ -84,15 +89,9 @@ class S3Token2Mel(torch.nn.Module):
             num_mid_blocks=12,
             num_heads=8,
             act_fn='gelu',
+            meanflow=self.meanflow,
         )
-        cfm_params = DictConfig({
+        cfm_params = CFM_PARAMS
-            "sigma_min": 1e-06,
-            "solver": 'euler',
-            "t_scheduler": 'cosine',
-            "training_cfg_rate": 0.2,
-            "inference_cfg_rate": 0.7,
-            "reg_loss_type": 'l1',
-        })
         decoder = CausalConditionalCFM(
             spk_emb_dim=80,
             cfm_params=cfm_params,
@@ -111,6 +110,11 @@ class S3Token2Mel(torch.nn.Module):
         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,
@@ -171,7 +175,10 @@ class S3Token2Mel(torch.nn.Module):
         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.
@@ -199,18 +206,21 @@ class S3Token2Mel(torch.nn.Module):
                 if isinstance(ref_dict[rk], np.ndarray):
                     ref_dict[rk] = torch.from_numpy(ref_dict[rk])
                 if torch.is_tensor(ref_dict[rk]):
-                    ref_dict[rk] = ref_dict[rk].to(self.device)
+                    ref_dict[rk] = ref_dict[rk].to(device=self.device, dtype=self.dtype)
 
-        if len(speech_tokens.shape) == 1:
+        speech_tokens = torch.atleast_2d(speech_tokens)
-            speech_tokens = speech_tokens.unsqueeze(0)
 
-        # assert speech_tokens.shape[0] == 1, "only batch size of one allowed for now"
+        # backcompat
-        speech_token_lens = torch.LongTensor([speech_tokens.size(1)]).to(self.device)
+        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
@@ -218,13 +228,15 @@ class S3Token2Mel(torch.nn.Module):
 
 class S3Token2Wav(S3Token2Mel):
     """
-    The decoder of CosyVoice2 is a concat of token-to-mel (CFM) and a mel-to-waveform (HiFiGAN) modules.
+    The decoder of S3Gen is a concat of token-to-mel (CFM) and a mel-to-waveform (HiFiGAN) modules.
 
     TODO: make these modules configurable?
     """
 
-    def __init__(self):
+    ignore_state_dict_missing = ("tokenizer._mel_filters", "tokenizer.window")
-        super().__init__()
+
+    def __init__(self, meanflow=False):
+        super().__init__(meanflow)
 
         f0_predictor = ConvRNNF0Predictor()
         self.mel2wav = HiFTGenerator(
@@ -241,6 +253,7 @@ class S3Token2Wav(S3Token2Mel):
         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,
@@ -250,9 +263,24 @@ class S3Token2Wav(S3Token2Mel):
         ref_sr: Optional[int],
         # pre-computed ref embedding (prod API)
         ref_dict: Optional[dict] = None,
-        finalize: bool = False
+        finalize: bool = False,
+        speech_token_lens=None,
+        skip_vocoder=False,
+        n_cfm_timesteps=None,
+        noised_mels=None,
     ):
-        output_mels = super().forward(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize)
+        """
+        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)
@@ -274,14 +302,24 @@ class S3Token2Wav(S3Token2Mel):
         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,
     ):
-        return super().forward(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize)
+        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(self.device)
+            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()
@@ -293,11 +331,22 @@ class S3Token2Wav(S3Token2Mel):
         ref_sr: Optional[int] = None,
         # pre-computed ref embedding (prod API)
         ref_dict: Optional[dict] = None,
-        cache_source: torch.Tensor = None, # NOTE: this arg is for streaming, it can probably be removed here
+        # left as a kwarg because this can change input/output size ratio
-        finalize: bool = True,
+        drop_invalid_tokens=True,
+        n_cfm_timesteps=None,
+        speech_token_lens=None,
     ):
-        output_mels = self.flow_inference(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize)
+        output_mels = self.flow_inference(
-        output_wavs, output_sources = self.hift_inference(output_mels, cache_source)
+            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

local_chatterbox/chatterbox/models/s3gen/utils/intmeanflow.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+
+
+def get_intmeanflow_time_mixer(dims):
+    """"
+    Diagonal init as described in 3.3 https://arxiv.org/pdf/2510.07979
+    """
+    layer = nn.Linear(dims * 2, dims, bias=False)
+
+    with torch.no_grad():
+        target_weight = torch.zeros(dims, 2 * dims)
+        target_weight[:, 0:dims] = torch.eye(dims)
+        layer.weight.data = target_weight
+
+    return layer

local_chatterbox/chatterbox/models/s3gen/utils/mask.py

@@ -182,7 +182,7 @@ def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
                  [0, 0, 1, 1, 1]]
     """
     batch_size = lengths.size(0)
-    max_len = max_len if max_len > 0 else lengths.max().item()
+    max_len = max_len if max_len > 0 else int(lengths.max().item())
     seq_range = torch.arange(0,
                              max_len,
                              dtype=torch.int64,

local_chatterbox/chatterbox/models/t3/llama_configs.py

@@ -32,6 +32,40 @@ LLAMA_520M_CONFIG_DICT = dict(
     use_cache=True,
 )
 
+GPT2_MEDIUM_CONFIG = {
+    "activation_function": "gelu_new",
+    "architectures": ["GPT2LMHeadModel"],
+    "attn_pdrop": 0.1,
+    "bos_token_id": 50256,
+    "embd_pdrop": 0.1,
+    "eos_token_id": 50256,
+    "initializer_range": 0.02,
+    "layer_norm_epsilon": 1e-05,
+    "model_type": "gpt2",
+    "n_ctx": 8196,
+    "n_embd": 1024,
+    "hidden_size": 1024,
+    "n_head": 16,
+    "n_layer": 24,
+    "n_positions": 8196,
+    "n_special": 0,
+    "predict_special_tokens": True,
+    "resid_pdrop": 0.1,
+    "summary_activation": None,
+    "summary_first_dropout": 0.1,
+    "summary_proj_to_labels": True,
+    "summary_type": "cls_index",
+    "summary_use_proj": True,
+    "task_specific_params": {
+        "text-generation": {
+            "do_sample": True,
+            "max_length": 50
+        }
+    },
+    "vocab_size": 50276,
+}
+
 LLAMA_CONFIGS = {
     "Llama_520M": LLAMA_520M_CONFIG_DICT,
+    "GPT2_medium": GPT2_MEDIUM_CONFIG,
 }

local_chatterbox/chatterbox/models/t3/modules/t3_config.py

@@ -22,6 +22,36 @@ class T3Config:
     use_perceiver_resampler = True
     emotion_adv = True
 
+    def __init__(self, text_tokens_dict_size=None):
+        """Initialize T3Config with optional overrides."""
+        if text_tokens_dict_size is not None:
+            self.text_tokens_dict_size = text_tokens_dict_size
+
+    @classmethod
+    def multilingual(cls) -> 'T3Config':
+        """Return config for multilingual model (23 languages)."""
+        config = cls()
+        config.text_tokens_dict_size = 2454
+        config.llama_config_name = "Llama_520M"
+        config.input_pos_emb = "learned"
+        config.speech_cond_prompt_len = 150
+        config.use_perceiver_resampler = True
+        config.emotion_adv = True
+        return config
+
+    @classmethod
+    def turbo(cls) -> 'T3Config':
+        """Return config for turbo model (faster, GPT2-based)."""
+        config = cls()
+        config.text_tokens_dict_size = 50276
+        config.llama_config_name = "GPT2_medium"
+        config.speech_tokens_dict_size = 6563
+        config.input_pos_emb = None
+        config.speech_cond_prompt_len = 375
+        config.use_perceiver_resampler = False
+        config.emotion_adv = False
+        return config
+
     @property
     def n_channels(self):
         return LLAMA_CONFIGS[self.llama_config_name]["hidden_size"]

local_chatterbox/chatterbox/models/t3/t3.py

@@ -7,8 +7,15 @@ from tqdm import tqdm
 import torch
 import torch.nn.functional as F
 from torch import nn, Tensor
-from transformers import LlamaModel, LlamaConfig
+from transformers import LlamaModel, LlamaConfig, GPT2Config, GPT2Model
-from transformers.generation.logits_process import TopPLogitsWarper, RepetitionPenaltyLogitsProcessor
+from transformers.generation.logits_process import (
+    LogitsProcessorList,
+    TopPLogitsWarper,
+    RepetitionPenaltyLogitsProcessor,
+    TemperatureLogitsWarper,
+    TopKLogitsWarper,
+    MinPLogitsWarper,
+)
 
 from .modules.learned_pos_emb import LearnedPositionEmbeddings
 
@@ -16,7 +23,7 @@ from .modules.cond_enc import T3CondEnc, T3Cond
 from .modules.t3_config import T3Config
 from .llama_configs import LLAMA_CONFIGS
 from .inference.t3_hf_backend import T3HuggingfaceBackend
-from .inference.alignment_stream_analyzer import AlignmentStreamAnalyzer
+# AlignmentStreamAnalyzer import removed - disabled due to SDPA incompatibility
 
 
 logger = logging.getLogger(__name__)
@@ -44,11 +51,23 @@ class T3(nn.Module):
             different PE embedding space for speech.
     """
 
-    def __init__(self, hp=T3Config()):
+    def __init__(self, hp=None):
+        if hp is None:
+            hp = T3Config()
         super().__init__()
         self.hp = hp
-        self.cfg = LlamaConfig(**LLAMA_CONFIGS[hp.llama_config_name])
+
+        # Determine if we're using GPT2 or LLaMA backend
+        config_dict = LLAMA_CONFIGS[hp.llama_config_name]
+        self.is_gpt = config_dict.get("model_type") == "gpt2"
+
+        if self.is_gpt:
+            self.cfg = GPT2Config(**config_dict)
+            self.tfmr = GPT2Model(self.cfg)
+        else:
+            self.cfg = LlamaConfig(**config_dict)
             self.tfmr = LlamaModel(self.cfg)
+
         self.dim = self.cfg.hidden_size
         self.deepspeed_patch_applied = False
 
@@ -57,7 +76,9 @@ class T3(nn.Module):
         self.text_emb = nn.Embedding(hp.text_tokens_dict_size, self.dim)
         self.speech_emb = nn.Embedding(hp.speech_tokens_dict_size, self.dim)
 
-        # custom position embedding
+        # custom position embedding (only for LLaMA with learned PE)
+        self.text_pos_emb = None
+        self.speech_pos_emb = None
         if hp.input_pos_emb == "learned":
             max_text_seq_len = hp.max_text_tokens + 2
             self.text_pos_emb = LearnedPositionEmbeddings(max_text_seq_len, self.dim)
@@ -65,9 +86,9 @@ class T3(nn.Module):
             max_mel_seq_len = hp.max_speech_tokens + 2 + 2
             self.speech_pos_emb = LearnedPositionEmbeddings(max_mel_seq_len, self.dim)
 
-        # logit projection
+        # logit projection (GPT2 speech_head has bias, LLaMA doesn't)
         self.text_head = nn.Linear(self.cfg.hidden_size, hp.text_tokens_dict_size, bias=False)
-        self.speech_head = nn.Linear(self.cfg.hidden_size, hp.speech_tokens_dict_size, bias=False)
+        self.speech_head = nn.Linear(self.cfg.hidden_size, hp.speech_tokens_dict_size, bias=self.is_gpt)
         self.compiled = False
 
     @property
@@ -79,7 +100,10 @@ class T3(nn.Module):
         Token cond data needs to be embedded, so that needs to be here instead of in `T3CondEnc`.
         """
         if t3_cond.cond_prompt_speech_tokens is not None and t3_cond.cond_prompt_speech_emb is None:
-            t3_cond.cond_prompt_speech_emb = self.speech_emb(t3_cond.cond_prompt_speech_tokens) + \
+            t3_cond.cond_prompt_speech_emb = self.speech_emb(t3_cond.cond_prompt_speech_tokens)
+            # Only add position embeddings for LLaMA (not GPT2)
+            if not self.is_gpt and self.speech_pos_emb is not None:
+                t3_cond.cond_prompt_speech_emb = t3_cond.cond_prompt_speech_emb + \
                     self.speech_pos_emb(t3_cond.cond_prompt_speech_tokens)
         return self.cond_enc(t3_cond)  # (B, len_cond, dim)
 
@@ -98,7 +122,8 @@ class T3(nn.Module):
             text_emb[1].zero_()  # CFG uncond
 
         speech_emb = self.speech_emb(speech_tokens)  # (B, len_speech, dim)
-        if self.hp.input_pos_emb == "learned":
+        # Only add position embeddings for LLaMA with learned PE (not GPT2)
+        if self.hp.input_pos_emb == "learned" and not self.is_gpt:
             text_emb = text_emb + self.text_pos_emb(text_tokens)
             speech_emb = speech_emb + self.speech_pos_emb(speech_tokens)
         len_cond = cond_emb.size(1)
@@ -223,10 +248,11 @@ class T3(nn.Module):
         stop_on_eos=True,
         do_sample=True,
         temperature=0.8,
-        top_p=0.8,
+        top_p=0.95,
+        min_p=0.05,
         length_penalty=1.0,
-        repetition_penalty=2.0,
+        repetition_penalty=1.2,
-        cfg_weight=0,
+        cfg_weight=0.5,
     ):
         """
         Args:
@@ -257,19 +283,17 @@ class T3(nn.Module):
         # TODO? synchronize the expensive compile function
         # with self.compile_lock:
         if not self.compiled:
-            alignment_stream_analyzer = AlignmentStreamAnalyzer(
+            # NOTE: AlignmentStreamAnalyzer disabled - it requires attention weights from
-                self.tfmr,
+            # output_attentions=True, which is incompatible with SDPA (the default in
-                None,
+            # transformers >=4.36). Since the analyzer's step() method was already commented
-                text_tokens_slice=(len_cond, len_cond + text_tokens.size(-1)),
+            # out in t3_hf_backend.py, we skip creating it entirely to avoid the crash.
-                alignment_layer_idx=9, # TODO: hparam or something?
+            # See: https://github.com/resemble-ai/chatterbox/issues/106
-                eos_idx=self.hp.stop_speech_token,
-            )
             patched_model = T3HuggingfaceBackend(
                 config=self.cfg,
                 llama=self.tfmr,
                 speech_enc=self.speech_emb,
                 speech_head=self.speech_head,
-                alignment_stream_analyzer=alignment_stream_analyzer,
+                alignment_stream_analyzer=None,
             )
             self.patched_model = patched_model
             self.compiled = True
@@ -312,7 +336,8 @@ class T3(nn.Module):
 
         # Instantiate the logits processors.
         top_p_warper = TopPLogitsWarper(top_p=top_p)
-        repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(penalty=repetition_penalty)
+        min_p_warper = MinPLogitsWarper(min_p=min_p)
+        repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(penalty=float(repetition_penalty))
 
         # ---- Initial Forward Pass (no kv_cache yet) ----
         output = self.patched_model(
@@ -342,9 +367,11 @@ class T3(nn.Module):
             if temperature != 1.0:
                 logits = logits / temperature
 
-            # Apply repetition penalty and top‑p filtering.
+            # Apply repetition penalty, min_p, and top_p filtering.
-            logits = repetition_penalty_processor(generated_ids, logits)
+            ids_for_proc = generated_ids[:1, ...]   # batch = 1
-            logits = top_p_warper(None, logits)
+            logits = repetition_penalty_processor(ids_for_proc, logits)
+            logits = min_p_warper(ids_for_proc, logits)
+            logits = top_p_warper(ids_for_proc, logits)
 
             # Convert logits to probabilities and sample the next token.
             probs = torch.softmax(logits, dim=-1)
@@ -379,3 +406,83 @@ class T3(nn.Module):
         # Concatenate all predicted tokens along the sequence dimension.
         predicted_tokens = torch.cat(predicted, dim=1)  # shape: (B, num_tokens)
         return predicted_tokens
+
+    @torch.inference_mode()
+    def inference_turbo(self, t3_cond, text_tokens, temperature=0.8, top_k=1000, top_p=0.95, repetition_penalty=1.2,
+                        max_gen_len=1000):
+        """
+        Turbo inference method using GPT2-style generation.
+        Used by ChatterboxTurboTTS for faster inference.
+        """
+        logits_processors = LogitsProcessorList()
+        if temperature > 0 and temperature != 1.0:
+            logits_processors.append(TemperatureLogitsWarper(temperature))
+        if top_k > 0:
+            logits_processors.append(TopKLogitsWarper(top_k))
+        if top_p < 1.0:
+            logits_processors.append(TopPLogitsWarper(top_p))
+        if repetition_penalty != 1.0:
+            logits_processors.append(RepetitionPenaltyLogitsProcessor(repetition_penalty))
+
+        speech_start_token = self.hp.start_speech_token * torch.ones_like(text_tokens[:, :1])
+        embeds, _ = self.prepare_input_embeds(
+            t3_cond=t3_cond,
+            text_tokens=text_tokens,
+            speech_tokens=speech_start_token,
+            cfg_weight=0.0,
+        )
+
+        generated_speech_tokens = []
+
+        llm_outputs = self.tfmr(
+            inputs_embeds=embeds,
+            use_cache=True
+        )
+
+        hidden_states = llm_outputs[0]
+        past_key_values = llm_outputs.past_key_values
+
+        speech_hidden = hidden_states[:, -1:]
+        speech_logits = self.speech_head(speech_hidden)
+
+        processed_logits = logits_processors(speech_start_token, speech_logits[:, -1, :])
+        probs = F.softmax(processed_logits, dim=-1)
+        next_speech_token = torch.multinomial(probs, num_samples=1)
+
+        generated_speech_tokens.append(next_speech_token)
+        current_speech_token = next_speech_token
+
+        for _ in tqdm(range(max_gen_len), desc="Turbo Sampling", dynamic_ncols=True):
+            current_speech_embed = self.speech_emb(current_speech_token)
+
+            llm_outputs = self.tfmr(
+                inputs_embeds=current_speech_embed,
+                past_key_values=past_key_values,
+                use_cache=True
+            )
+
+            hidden_states = llm_outputs[0]
+            past_key_values = llm_outputs.past_key_values
+            speech_logits = self.speech_head(hidden_states)
+
+            input_ids = torch.cat(generated_speech_tokens, dim=1)
+            processed_logits = logits_processors(input_ids, speech_logits[:, -1, :])
+            if torch.all(processed_logits == -float("inf")):
+                print("Warning: All logits are -inf")
+                break
+
+            probs = F.softmax(processed_logits, dim=-1)
+            next_speech_token = torch.multinomial(probs, num_samples=1)
+
+            generated_speech_tokens.append(next_speech_token)
+            current_speech_token = next_speech_token
+            if torch.all(next_speech_token == self.hp.stop_speech_token):
+                break
+
+        all_tokens = torch.cat(generated_speech_tokens, dim=1)
+
+        # Remove EOS token if present
+        if all_tokens.size(1) > 0 and all_tokens[0, -1] == self.hp.stop_speech_token:
+            all_tokens = all_tokens[:, :-1]
+
+        return all_tokens

local_chatterbox/chatterbox/models/tokenizers/__init__.py

@@ -1 +1 @@
-from .tokenizer import EnTokenizer
+from .tokenizer import EnTokenizer, MTLTokenizer

local_chatterbox/chatterbox/models/tokenizers/tokenizer.py

@@ -1,6 +1,9 @@
 import logging
+import json
 
 import torch
+from pathlib import Path
+from unicodedata import category, normalize
 from tokenizers import Tokenizer
 
 
@@ -13,6 +16,7 @@ SPECIAL_TOKENS = [SOT, EOT, UNK, SPACE, "[PAD]", "[SEP]", "[CLS]", "[MASK]"]
 
 logger = logging.getLogger(__name__)
 
+
 class EnTokenizer:
     def __init__(self, vocab_file_path):
         self.tokenizer: Tokenizer = Tokenizer.from_file(vocab_file_path)
@@ -41,10 +45,242 @@ class EnTokenizer:
         if isinstance(seq, torch.Tensor):
             seq = seq.cpu().numpy()
 
-        txt: str = self.tokenizer.decode(seq,
+        txt: str = self.tokenizer.decode(seq, skip_special_tokens=False)
-        skip_special_tokens=False)
         txt = txt.replace(' ', '')
         txt = txt.replace(SPACE, ' ')
         txt = txt.replace(EOT, '')
         txt = txt.replace(UNK, '')
         return txt
+
+
+# Global instances for optional language-specific dependencies
+_kakasi = None
+_dicta = None
+_russian_stresser = None
+
+
+def is_kanji(c: str) -> bool:
+    """Check if character is kanji."""
+    return 19968 <= ord(c) <= 40959
+
+
+def is_katakana(c: str) -> bool:
+    """Check if character is katakana."""
+    return 12449 <= ord(c) <= 12538
+
+
+def hiragana_normalize(text: str) -> str:
+    """Japanese text normalization: converts kanji to hiragana."""
+    global _kakasi
+
+    try:
+        if _kakasi is None:
+            import pykakasi
+            _kakasi = pykakasi.kakasi()
+
+        result = _kakasi.convert(text)
+        out = []
+
+        for r in result:
+            inp = r['orig']
+            hira = r["hira"]
+
+            if any([is_kanji(c) for c in inp]):
+                if hira and hira[0] in ["は", "へ"]:
+                    hira = " " + hira
+                out.append(hira)
+            elif all([is_katakana(c) for c in inp]) if inp else False:
+                out.append(r['orig'])
+            else:
+                out.append(inp)
+
+        normalized_text = "".join(out)
+        import unicodedata
+        normalized_text = unicodedata.normalize('NFKD', normalized_text)
+        return normalized_text
+
+    except ImportError:
+        logger.warning("pykakasi not available - Japanese text processing skipped")
+        return text
+
+
+def add_hebrew_diacritics(text: str) -> str:
+    """Hebrew text normalization: adds diacritics."""
+    global _dicta
+
+    try:
+        if _dicta is None:
+            from dicta_onnx import Dicta
+            _dicta = Dicta()
+        return _dicta.add_diacritics(text)
+    except ImportError:
+        logger.warning("dicta_onnx not available - Hebrew text processing skipped")
+        return text
+    except Exception as e:
+        logger.warning(f"Hebrew diacritization failed: {e}")
+        return text
+
+
+def korean_normalize(text: str) -> str:
+    """Korean text normalization: decompose syllables into Jamo."""
+
+    def decompose_hangul(char):
+        if not ('\uac00' <= char <= '\ud7af'):
+            return char
+        base = ord(char) - 0xAC00
+        initial = chr(0x1100 + base // (21 * 28))
+        medial = chr(0x1161 + (base % (21 * 28)) // 28)
+        final = chr(0x11A7 + base % 28) if base % 28 > 0 else ''
+        return initial + medial + final
+
+    result = ''.join(decompose_hangul(char) for char in text)
+    return result.strip()
+
+
+def add_russian_stress(text: str) -> str:
+    """Russian text normalization: adds stress marks."""
+    global _russian_stresser
+
+    try:
+        if _russian_stresser is None:
+            from russian_text_stresser.text_stresser import RussianTextStresser
+            _russian_stresser = RussianTextStresser()
+        return _russian_stresser.stress_text(text)
+    except ImportError:
+        logger.warning("russian_text_stresser not available - Russian stress labeling skipped")
+        return text
+    except Exception as e:
+        logger.warning(f"Russian stress labeling failed: {e}")
+        return text
+
+
+class ChineseCangjieConverter:
+    """Converts Chinese characters to Cangjie codes for tokenization."""
+
+    def __init__(self, model_dir=None):
+        self.word2cj = {}
+        self.cj2word = {}
+        self.segmenter = None
+        self._load_cangjie_mapping(model_dir)
+        self._init_segmenter()
+
+    def _load_cangjie_mapping(self, model_dir=None):
+        """Load Cangjie mapping from model directory."""
+        try:
+            if model_dir:
+                cangjie_file = Path(model_dir) / "Cangjie5_TC.json"
+                if cangjie_file.exists():
+                    with open(cangjie_file, "r", encoding="utf-8") as fp:
+                        data = json.load(fp)
+
+                    for entry in data:
+                        word, code = entry.split("\t")[:2]
+                        self.word2cj[word] = code
+                        if code not in self.cj2word:
+                            self.cj2word[code] = [word]
+                        else:
+                            self.cj2word[code].append(word)
+        except Exception as e:
+            logger.warning(f"Could not load Cangjie mapping: {e}")
+
+    def _init_segmenter(self):
+        """Initialize pkuseg segmenter."""
+        try:
+            from spacy_pkuseg import pkuseg
+            self.segmenter = pkuseg()
+        except ImportError:
+            logger.warning("pkuseg not available - Chinese segmentation will be skipped")
+            self.segmenter = None
+
+    def _cangjie_encode(self, glyph: str):
+        """Encode a single Chinese glyph to Cangjie code."""
+        code = self.word2cj.get(glyph, None)
+        if code is None:
+            return None
+        index = self.cj2word[code].index(glyph)
+        index = str(index) if index > 0 else ""
+        return code + str(index)
+
+    def __call__(self, text):
+        """Convert Chinese characters in text to Cangjie tokens."""
+        output = []
+        if self.segmenter is not None:
+            segmented_words = self.segmenter.cut(text)
+            full_text = " ".join(segmented_words)
+        else:
+            full_text = text
+
+        for t in full_text:
+            if category(t) == "Lo":
+                cangjie = self._cangjie_encode(t)
+                if cangjie is None:
+                    output.append(t)
+                    continue
+                code = []
+                for c in cangjie:
+                    code.append(f"[cj_{c}]")
+                code.append("[cj_.]")
+                code = "".join(code)
+                output.append(code)
+            else:
+                output.append(t)
+        return "".join(output)
+
+
+class MTLTokenizer:
+    """Multilingual tokenizer supporting 23 languages."""
+
+    def __init__(self, vocab_file_path):
+        self.tokenizer: Tokenizer = Tokenizer.from_file(vocab_file_path)
+        model_dir = Path(vocab_file_path).parent
+        self.cangjie_converter = ChineseCangjieConverter(model_dir)
+        self.check_vocabset_sot_eot()
+
+    def check_vocabset_sot_eot(self):
+        voc = self.tokenizer.get_vocab()
+        assert SOT in voc
+        assert EOT in voc
+
+    def preprocess_text(self, raw_text: str, language_id: str = None, lowercase: bool = True, nfkd_normalize: bool = True):
+        """Text preprocessor for lowercase and NFKD normalization."""
+        preprocessed_text = raw_text
+        if lowercase:
+            preprocessed_text = preprocessed_text.lower()
+        if nfkd_normalize:
+            preprocessed_text = normalize("NFKD", preprocessed_text)
+        return preprocessed_text
+
+    def text_to_tokens(self, text: str, language_id: str = None, lowercase: bool = True, nfkd_normalize: bool = True):
+        text_tokens = self.encode(text, language_id=language_id, lowercase=lowercase, nfkd_normalize=nfkd_normalize)
+        text_tokens = torch.IntTensor(text_tokens).unsqueeze(0)
+        return text_tokens
+
+    def encode(self, txt: str, language_id: str = None, lowercase: bool = True, nfkd_normalize: bool = True):
+        txt = self.preprocess_text(txt, language_id=language_id, lowercase=lowercase, nfkd_normalize=nfkd_normalize)
+
+        # Language-specific text processing
+        if language_id == 'zh':
+            txt = self.cangjie_converter(txt)
+        elif language_id == 'ja':
+            txt = hiragana_normalize(txt)
+        elif language_id == 'he':
+            txt = add_hebrew_diacritics(txt)
+        elif language_id == 'ko':
+            txt = korean_normalize(txt)
+        elif language_id == 'ru':
+            txt = add_russian_stress(txt)
+
+        # Prepend language token
+        if language_id:
+            txt = f"[{language_id.lower()}]{txt}"
+
+        txt = txt.replace(' ', SPACE)
+        return self.tokenizer.encode(txt).ids
+
+    def decode(self, seq):
+        if isinstance(seq, torch.Tensor):
+            seq = seq.cpu().numpy()
+
+        txt = self.tokenizer.decode(seq, skip_special_tokens=False)
+        txt = txt.replace(' ', '').replace(SPACE, ' ').replace(EOT, '').replace(UNK, '')
+        return txt

local_chatterbox/chatterbox/models/utils.py

@@ -0,0 +1,4 @@
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self

local_chatterbox/chatterbox/models/voice_encoder/voice_encoder.py

@@ -257,14 +257,17 @@ class VoiceEncoder(nn.Module):
 
         :param trim_top_db: this argument was only added for the sake of compatibility with metavoice's implementation
         """
+        # Ensure all wavs are float32 (librosa operations can return float64)
+        wavs = [wav.astype(np.float32) if wav.dtype != np.float32 else wav for wav in wavs]
+
         if sample_rate != self.hp.sample_rate:
             wavs = [
-                librosa.resample(wav, orig_sr=sample_rate, target_sr=self.hp.sample_rate, res_type="kaiser_fast")
+                librosa.resample(wav, orig_sr=sample_rate, target_sr=self.hp.sample_rate, res_type="kaiser_fast").astype(np.float32)
                 for wav in wavs
             ]
 
         if trim_top_db:
-            wavs = [librosa.effects.trim(wav, top_db=trim_top_db)[0] for wav in wavs]
+            wavs = [librosa.effects.trim(wav, top_db=trim_top_db)[0].astype(np.float32) for wav in wavs]
 
         if "rate" not in kwargs:
             kwargs["rate"] = 1.3  # Resemble's default value.

local_chatterbox/chatterbox/mtl_tts.py

@@ -0,0 +1,376 @@
+"""
+Chatterbox Multilingual TTS - Text-to-speech supporting 23 languages.
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import librosa
+import numpy as np
+import torch
+import torch.nn.functional as F
+from safetensors.torch import load_file as load_safetensors
+
+# Optional Perth watermarking - gracefully handle import failure
+try:
+    import perth
+    PERTH_AVAILABLE = True
+except (ImportError, AttributeError):
+    PERTH_AVAILABLE = False
+    print("Warning: Perth watermarking not available. Audio will be generated without watermarking.")
+
+from .models.t3 import T3
+from .models.t3.modules.t3_config import T3Config
+from .models.s3tokenizer import S3_SR, drop_invalid_tokens
+from .models.s3gen import S3GEN_SR, S3Gen
+from .models.tokenizers import MTLTokenizer
+from .models.voice_encoder import VoiceEncoder
+from .models.t3.modules.cond_enc import T3Cond
+from .paths import get_chatterbox_multilingual_dir, download_to_local
+
+
+REPO_ID = "ResembleAI/chatterbox"
+MTL_MODEL_FILES = [
+    "ve.pt",
+    "t3_mtl23ls_v2.safetensors",
+    "s3gen.pt",
+    "grapheme_mtl_merged_expanded_v1.json",
+    "conds.pt",
+    "Cangjie5_TC.json",
+]
+
+# Supported languages for the multilingual model
+SUPPORTED_LANGUAGES = {
+    "ar": "Arabic",
+    "da": "Danish",
+    "de": "German",
+    "el": "Greek",
+    "en": "English",
+    "es": "Spanish",
+    "fi": "Finnish",
+    "fr": "French",
+    "he": "Hebrew",
+    "hi": "Hindi",
+    "it": "Italian",
+    "ja": "Japanese",
+    "ko": "Korean",
+    "ms": "Malay",
+    "nl": "Dutch",
+    "no": "Norwegian",
+    "pl": "Polish",
+    "pt": "Portuguese",
+    "ru": "Russian",
+    "sv": "Swedish",
+    "sw": "Swahili",
+    "tr": "Turkish",
+    "zh": "Chinese",
+}
+
+
+def punc_norm(text: str) -> str:
+    """Punctuation normalization for multilingual model."""
+    if len(text) == 0:
+        return "You need to add some text for me to talk."
+
+    if text[0].islower():
+        text = text[0].upper() + text[1:]
+
+    text = " ".join(text.split())
+
+    punc_to_replace = [
+        ("...", ", "),
+        ("…", ", "),
+        (":", ","),
+        (" - ", ", "),
+        (";", ", "),
+        ("—", "-"),
+        ("–", "-"),
+        (" ,", ","),
+        (""", "\""),
+        (""", "\""),
+        ("'", "'"),
+        ("'", "'"),
+    ]
+    for old_char_sequence, new_char in punc_to_replace:
+        text = text.replace(old_char_sequence, new_char)
+
+    text = text.rstrip(" ")
+    sentence_enders = {".", "!", "?", "-", ",", "、", "，", "。", "？", "！"}
+    if not any(text.endswith(p) for p in sentence_enders):
+        text += "."
+
+    return text
+
+
+@dataclass
+class Conditionals:
+    """Conditionals for T3 and S3Gen."""
+    t3: T3Cond
+    gen: dict
+
+    def to(self, device):
+        self.t3 = self.t3.to(device=device)
+        for k, v in self.gen.items():
+            if torch.is_tensor(v):
+                self.gen[k] = v.to(device=device)
+        return self
+
+    def save(self, fpath: Path):
+        arg_dict = dict(
+            t3=self.t3.__dict__,
+            gen=self.gen
+        )
+        torch.save(arg_dict, fpath)
+
+    @classmethod
+    def load(cls, fpath, map_location="cpu"):
+        kwargs = torch.load(fpath, map_location=map_location, weights_only=True)
+        return cls(T3Cond(**kwargs['t3']), kwargs['gen'])
+
+
+class ChatterboxMultilingualTTS:
+    """
+    Chatterbox Multilingual TTS - Text-to-speech supporting 23 languages.
+
+    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)
+    """
+    ENC_COND_LEN = 6 * S3_SR
+    DEC_COND_LEN = 10 * S3GEN_SR
+
+    def __init__(
+        self,
+        t3: T3,
+        s3gen: S3Gen,
+        ve: VoiceEncoder,
+        tokenizer: MTLTokenizer,
+        device: str,
+        conds: Conditionals = None,
+    ):
+        self.sr = S3GEN_SR
+        self.t3 = t3
+        self.s3gen = s3gen
+        self.ve = ve
+        self.tokenizer = tokenizer
+        self.device = device
+        self.conds = conds
+        self.watermarker = perth.PerthImplicitWatermarker() if PERTH_AVAILABLE else None
+
+    @classmethod
+    def get_supported_languages(cls):
+        """Return dictionary of supported language codes and names."""
+        return SUPPORTED_LANGUAGES.copy()
+
+    @classmethod
+    def from_local(cls, ckpt_dir, device) -> 'ChatterboxMultilingualTTS':
+        ckpt_dir = Path(ckpt_dir)
+
+        if device in ["cpu", "mps"]:
+            map_location = torch.device('cpu')
+        else:
+            map_location = None
+
+        ve = VoiceEncoder()
+        ve.load_state_dict(torch.load(ckpt_dir / "ve.pt", weights_only=True, map_location=map_location))
+        ve.to(device).eval()
+
+        t3 = T3(T3Config.multilingual())
+        t3_state = load_safetensors(ckpt_dir / "t3_mtl23ls_v2.safetensors")
+        if "model" in t3_state.keys():
+            t3_state = t3_state["model"][0]
+        t3.load_state_dict(t3_state)
+        t3.to(device).eval()
+
+        s3gen = S3Gen()
+        s3gen.load_state_dict(torch.load(ckpt_dir / "s3gen.pt", weights_only=True, map_location=map_location))
+        s3gen.to(device).eval()
+
+        tokenizer = MTLTokenizer(str(ckpt_dir / "grapheme_mtl_merged_expanded_v1.json"))
+
+        conds = None
+        if (builtin_voice := ckpt_dir / "conds.pt").exists():
+            conds = Conditionals.load(builtin_voice, map_location=map_location).to(device)
+
+        return cls(t3, s3gen, ve, tokenizer, device, conds=conds)
+
+    @classmethod
+    def from_pretrained(cls, device) -> 'ChatterboxMultilingualTTS':
+        if device == "mps" and not torch.backends.mps.is_available():
+            if not torch.backends.mps.is_built():
+                print("MPS not available because the current PyTorch install was not built with MPS enabled.")
+            else:
+                print("MPS not available because the current MacOS version is not 12.3+.")
+            device = "cpu"
+
+        # Download models to centralized location
+        local_dir = get_chatterbox_multilingual_dir()
+        print(f"[FL Chatterbox Multilingual] Model download path: {local_dir}")
+        download_to_local(REPO_ID, MTL_MODEL_FILES, local_dir)
+
+        return cls.from_local(local_dir, device)
+
+    def _trim_trailing_silence(self, wav, threshold_db=-40, min_silence_duration=0.5):
+        """
+        Trim trailing silence/noise from audio.
+
+        Args:
+            wav: Audio waveform as numpy array
+            threshold_db: Silence threshold in dB (default -40dB)
+            min_silence_duration: Minimum silence duration to trigger trimming (seconds)
+
+        Returns:
+            Trimmed audio waveform
+        """
+        # Convert threshold from dB to linear amplitude
+        threshold = 10 ** (threshold_db / 20)
+
+        # Calculate RMS energy in windows
+        window_size = int(0.02 * self.sr)  # 20ms windows
+        hop_size = window_size // 2
+
+        # Compute RMS for each window
+        num_windows = (len(wav) - window_size) // hop_size + 1
+        if num_windows <= 0:
+            return wav
+
+        rms = np.zeros(num_windows)
+        for i in range(num_windows):
+            start = i * hop_size
+            end = start + window_size
+            rms[i] = np.sqrt(np.mean(wav[start:end] ** 2))
+
+        # Find the last window above threshold
+        above_threshold = rms > threshold
+        if not np.any(above_threshold):
+            # All silence, return minimal audio
+            return wav[:int(0.1 * self.sr)]
+
+        last_voice_idx = np.where(above_threshold)[0][-1]
+
+        # Convert back to sample index and add a small buffer
+        trim_sample = (last_voice_idx + 1) * hop_size + window_size
+        trim_sample = min(trim_sample + int(0.1 * self.sr), len(wav))  # Add 100ms buffer
+
+        return wav[:trim_sample]
+
+    def prepare_conditionals(self, wav_fpath, exaggeration=0.5):
+        """Prepare voice conditionals from reference audio."""
+        s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
+
+        ref_16k_wav = librosa.resample(s3gen_ref_wav, orig_sr=S3GEN_SR, target_sr=S3_SR)
+
+        s3gen_ref_wav = s3gen_ref_wav[:self.DEC_COND_LEN]
+        s3gen_ref_dict = self.s3gen.embed_ref(s3gen_ref_wav, S3GEN_SR, device=self.device)
+
+        t3_cond_prompt_tokens = None
+        if plen := self.t3.hp.speech_cond_prompt_len:
+            s3_tokzr = self.s3gen.tokenizer
+            t3_cond_prompt_tokens, _ = s3_tokzr.forward([ref_16k_wav[:self.ENC_COND_LEN]], max_len=plen)
+            t3_cond_prompt_tokens = torch.atleast_2d(t3_cond_prompt_tokens).to(self.device)
+
+        ve_embed = torch.from_numpy(self.ve.embeds_from_wavs([ref_16k_wav], sample_rate=S3_SR))
+        ve_embed = ve_embed.mean(axis=0, keepdim=True).to(self.device)
+
+        t3_cond = T3Cond(
+            speaker_emb=ve_embed,
+            cond_prompt_speech_tokens=t3_cond_prompt_tokens,
+            emotion_adv=exaggeration * torch.ones(1, 1, 1),
+        ).to(device=self.device)
+        self.conds = Conditionals(t3_cond, s3gen_ref_dict)
+
+    def generate(
+        self,
+        text,
+        language_id,
+        audio_prompt_path=None,
+        exaggeration=0.5,
+        cfg_weight=0.5,
+        temperature=0.8,
+        repetition_penalty=2.0,
+        min_p=0.05,
+        top_p=1.0,
+    ):
+        """
+        Generate speech from text in specified language.
+
+        Args:
+            text: Text to speak
+            language_id: Two-letter language code (e.g., 'en', 'fr', 'ja', 'zh')
+            audio_prompt_path: Path to reference voice audio (min 6 seconds)
+            exaggeration: Emotion intensity (0.0-1.0+, default 0.5)
+            cfg_weight: Classifier-free guidance weight (default 0.5)
+            temperature: Sampling temperature (default 0.8)
+            repetition_penalty: Penalty for token repetition (default 2.0)
+            min_p: Minimum probability threshold (default 0.05)
+            top_p: Nucleus sampling threshold (default 1.0)
+
+        Returns:
+            Tensor of audio waveform
+        """
+        # Validate language_id
+        if language_id and language_id.lower() not in SUPPORTED_LANGUAGES:
+            supported_langs = ", ".join(SUPPORTED_LANGUAGES.keys())
+            raise ValueError(
+                f"Unsupported language_id '{language_id}'. "
+                f"Supported languages: {supported_langs}"
+            )
+
+        if audio_prompt_path:
+            self.prepare_conditionals(audio_prompt_path, exaggeration=exaggeration)
+        else:
+            assert self.conds is not None, "Please `prepare_conditionals` first or specify `audio_prompt_path`"
+
+        # Update exaggeration if needed
+        if float(exaggeration) != float(self.conds.t3.emotion_adv[0, 0, 0].item()):
+            _cond: T3Cond = self.conds.t3
+            self.conds.t3 = T3Cond(
+                speaker_emb=_cond.speaker_emb,
+                cond_prompt_speech_tokens=_cond.cond_prompt_speech_tokens,
+                emotion_adv=exaggeration * torch.ones(1, 1, 1),
+            ).to(device=self.device)
+
+        text = punc_norm(text)
+        text_tokens = self.tokenizer.text_to_tokens(
+            text, language_id=language_id.lower() if language_id else None
+        ).to(self.device)
+        text_tokens = torch.cat([text_tokens, text_tokens], dim=0)  # Need two seqs for CFG
+
+        sot = self.t3.hp.start_text_token
+        eot = self.t3.hp.stop_text_token
+        text_tokens = F.pad(text_tokens, (1, 0), value=sot)
+        text_tokens = F.pad(text_tokens, (0, 1), value=eot)
+
+        with torch.inference_mode():
+            speech_tokens = self.t3.inference(
+                t3_cond=self.conds.t3,
+                text_tokens=text_tokens,
+                max_new_tokens=1000,
+                temperature=temperature,
+                cfg_weight=cfg_weight,
+                repetition_penalty=repetition_penalty,
+                min_p=min_p,
+                top_p=top_p,
+            )
+            speech_tokens = speech_tokens[0]
+            speech_tokens = drop_invalid_tokens(speech_tokens)
+            speech_tokens = speech_tokens.to(self.device)
+
+            wav, _ = self.s3gen.inference(
+                speech_tokens=speech_tokens,
+                ref_dict=self.conds.gen,
+            )
+            wav = wav.squeeze(0).detach().cpu().numpy()
+
+            # Trim trailing silence/noise - workaround for known multilingual model issue
+            # See: https://github.com/resemble-ai/chatterbox/issues/287
+            wav = self._trim_trailing_silence(wav)
+
+            if self.watermarker is not None:
+                watermarked_wav = self.watermarker.apply_watermark(wav, sample_rate=self.sr)
+                return torch.from_numpy(watermarked_wav).unsqueeze(0)
+            else:
+                return torch.from_numpy(wav).unsqueeze(0)

local_chatterbox/chatterbox/paths.py

@@ -0,0 +1,84 @@
+"""
+Centralized path management for Chatterbox models.
+All models download to: ComfyUI/models/chatterbox/
+"""
+
+from pathlib import Path
+
+
+def get_chatterbox_models_dir() -> Path:
+    """
+    Get the centralized models directory for all Chatterbox models.
+
+    Returns:
+        Path to ComfyUI/models/chatterbox/
+    """
+    # Navigate: paths.py -> chatterbox -> local_chatterbox -> ComfyUI_Fill-ChatterBox -> custom_nodes -> ComfyUI
+    current_dir = Path(__file__).parent
+    comfyui_root = current_dir.parent.parent.parent.parent
+
+    models_dir = comfyui_root / "models" / "chatterbox"
+
+    # Verify we're in a valid ComfyUI structure
+    if not (comfyui_root / "custom_nodes").exists():
+        # Fallback: use a local directory if not in ComfyUI structure
+        models_dir = current_dir.parent.parent / "models"
+
+    models_dir.mkdir(parents=True, exist_ok=True)
+    return models_dir
+
+
+def get_chatterbox_tts_dir() -> Path:
+    """Get the directory for standard Chatterbox TTS models (English, 500M)."""
+    path = get_chatterbox_models_dir() / "chatterbox"
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def get_chatterbox_turbo_dir() -> Path:
+    """Get the directory for Chatterbox Turbo TTS models (350M, faster)."""
+    path = get_chatterbox_models_dir() / "chatterbox_turbo"
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def get_chatterbox_multilingual_dir() -> Path:
+    """Get the directory for Chatterbox Multilingual TTS models (23 languages)."""
+    path = get_chatterbox_models_dir() / "chatterbox_multilingual"
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def download_to_local(repo_id: str, filenames: list, local_dir: Path) -> Path:
+    """
+    Download model files from HuggingFace to a local directory.
+
+    Args:
+        repo_id: HuggingFace repository ID (e.g., "ResembleAI/chatterbox")
+        filenames: List of filenames to download
+        local_dir: Local directory to download to
+
+    Returns:
+        Path to the local directory containing the downloaded files
+    """
+    from huggingface_hub import hf_hub_download
+    import shutil
+
+    local_dir.mkdir(parents=True, exist_ok=True)
+
+    for filename in filenames:
+        local_path = local_dir / filename
+        if not local_path.exists():
+            print(f"[FL Chatterbox] Downloading {filename} to {local_dir}...")
+            try:
+                # Download to HF cache first, then copy to our location
+                cached_path = hf_hub_download(repo_id=repo_id, filename=filename)
+                # Copy from cache to our directory
+                shutil.copy2(cached_path, local_path)
+            except Exception as e:
+                print(f"[FL Chatterbox] Error downloading {filename}: {e}")
+                raise
+        else:
+            print(f"[FL Chatterbox] Using cached {filename}")
+
+    return local_dir

local_chatterbox/chatterbox/tts.py

@@ -1,11 +1,18 @@
 from dataclasses import dataclass
 from pathlib import Path
 
+import numpy as np
 import librosa
 import torch
-import perth
 import torch.nn.functional as F
-from huggingface_hub import hf_hub_download
+
+# Optional Perth watermarking - gracefully handle import failure
+try:
+    import perth
+    PERTH_AVAILABLE = True
+except (ImportError, AttributeError):
+    PERTH_AVAILABLE = False
+    print("Warning: Perth watermarking not available. Audio will be generated without watermarking.")
 from safetensors.torch import load_file
 
 from .models.t3 import T3
@@ -14,9 +21,11 @@ from .models.s3gen import S3GEN_SR, S3Gen
 from .models.tokenizers import EnTokenizer
 from .models.voice_encoder import VoiceEncoder
 from .models.t3.modules.cond_enc import T3Cond
+from .paths import get_chatterbox_tts_dir, download_to_local
 
 
 REPO_ID = "ResembleAI/chatterbox"
+TTS_MODEL_FILES = ["ve.safetensors", "t3_cfg.safetensors", "s3gen.safetensors", "tokenizer.json", "conds.pt"]
 
 
 def punc_norm(text: str) -> str:
@@ -123,7 +132,7 @@ class ChatterboxTTS:
         self.tokenizer = tokenizer
         self.device = device
         self.conds = conds
-        self.watermarker = perth.PerthImplicitWatermarker()
+        self.watermarker = perth.PerthImplicitWatermarker() if PERTH_AVAILABLE else None
 
     @classmethod
     def from_local(cls, ckpt_dir, device) -> 'ChatterboxTTS':
@@ -174,14 +183,18 @@ class ChatterboxTTS:
                 print("MPS not available because the current MacOS version is not 12.3+ and/or you do not have an MPS-enabled device on this machine.")
             device = "cpu"
 
-        for fpath in ["ve.safetensors", "t3_cfg.safetensors", "s3gen.safetensors", "tokenizer.json", "conds.pt"]:
+        # Download models to centralized location: ComfyUI/models/chatterbox/chatterbox/
-            local_path = hf_hub_download(repo_id=REPO_ID, filename=fpath)
+        local_dir = get_chatterbox_tts_dir()
+        print(f"[FL Chatterbox] Model download path: {local_dir}")
+        download_to_local(REPO_ID, TTS_MODEL_FILES, local_dir)
 
-        return cls.from_local(Path(local_path).parent, device)
+        return cls.from_local(local_dir, device)
 
     def prepare_conditionals(self, wav_fpath, exaggeration=0.5):
         ## Load reference wav
         s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
+        # Ensure float32 for PyTorch compatibility
+        s3gen_ref_wav = s3gen_ref_wav.astype(np.float32)
 
         ref_16k_wav = librosa.resample(s3gen_ref_wav, orig_sr=S3GEN_SR, target_sr=S3_SR)
 
@@ -259,5 +272,8 @@ class ChatterboxTTS:
                 ref_dict=self.conds.gen,
             )
             wav = wav.squeeze(0).detach().cpu().numpy()
+            if self.watermarker is not None:
                 watermarked_wav = self.watermarker.apply_watermark(wav, sample_rate=self.sr)
                 return torch.from_numpy(watermarked_wav).unsqueeze(0)
+            else:
+                return torch.from_numpy(wav).unsqueeze(0)

local_chatterbox/chatterbox/tts_turbo.py

@@ -0,0 +1,326 @@
+"""
+Chatterbox Turbo TTS - Faster GPT2-based text-to-speech.
+
+Supports paralinguistic tags: [laugh], [sigh], [gasp], [chuckle], etc.
+"""
+
+import math
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import librosa
+import torch
+from safetensors.torch import load_file
+from transformers import AutoTokenizer
+
+# Optional Perth watermarking - gracefully handle import failure
+try:
+    import perth
+    PERTH_AVAILABLE = True
+except (ImportError, AttributeError):
+    PERTH_AVAILABLE = False
+    print("Warning: Perth watermarking not available. Audio will be generated without watermarking.")
+
+# Optional loudness normalization
+try:
+    import pyloudnorm as ln
+    PYLOUDNORM_AVAILABLE = True
+except ImportError:
+    PYLOUDNORM_AVAILABLE = False
+    print("Warning: pyloudnorm not available. Loudness normalization disabled.")
+
+from .models.t3 import T3
+from .models.s3tokenizer import S3_SR
+from .models.s3gen import S3GEN_SR, S3Gen
+from .models.voice_encoder import VoiceEncoder
+from .models.t3.modules.cond_enc import T3Cond
+from .models.t3.modules.t3_config import T3Config
+from .models.s3gen.const import S3GEN_SIL
+from .paths import get_chatterbox_turbo_dir, download_to_local
+
+logger = logging.getLogger(__name__)
+
+REPO_ID = "ResembleAI/chatterbox-turbo"
+TURBO_MODEL_FILES = [
+    "ve.safetensors",
+    "t3_turbo_v1.safetensors",
+    "s3gen_meanflow.safetensors",
+    # GPT2-style tokenizer files (no tokenizer.json - uses vocab.json + merges.txt)
+    "tokenizer_config.json",
+    "special_tokens_map.json",
+    "vocab.json",
+    "merges.txt",
+    "added_tokens.json",
+    "conds.pt",
+]
+
+
+def punc_norm(text: str) -> str:
+    """Punctuation normalization for Turbo model."""
+    if len(text) == 0:
+        return "You need to add some text for me to talk."
+
+    if text[0].islower():
+        text = text[0].upper() + text[1:]
+
+    text = " ".join(text.split())
+
+    punc_to_replace = [
+        ("…", ", "),
+        (":", ","),
+        ("—", "-"),
+        ("–", "-"),
+        (" ,", ","),
+        (""", "\""),
+        (""", "\""),
+        ("'", "'"),
+        ("'", "'"),
+    ]
+    for old_char_sequence, new_char in punc_to_replace:
+        text = text.replace(old_char_sequence, new_char)
+
+    text = text.rstrip(" ")
+    sentence_enders = {".", "!", "?", "-", ","}
+    if not any(text.endswith(p) for p in sentence_enders):
+        text += "."
+
+    return text
+
+
+@dataclass
+class Conditionals:
+    """Conditionals for T3 and S3Gen."""
+    t3: T3Cond
+    gen: dict
+
+    def to(self, device):
+        self.t3 = self.t3.to(device=device)
+        for k, v in self.gen.items():
+            if torch.is_tensor(v):
+                self.gen[k] = v.to(device=device)
+        return self
+
+    def save(self, fpath: Path):
+        arg_dict = dict(
+            t3=self.t3.__dict__,
+            gen=self.gen
+        )
+        torch.save(arg_dict, fpath)
+
+    @classmethod
+    def load(cls, fpath, map_location="cpu"):
+        if isinstance(map_location, str):
+            map_location = torch.device(map_location)
+        kwargs = torch.load(fpath, map_location=map_location, weights_only=True)
+        return cls(T3Cond(**kwargs['t3']), kwargs['gen'])
+
+
+class ChatterboxTurboTTS:
+    """
+    Chatterbox Turbo TTS - Fast GPT2-based text-to-speech with paralinguistic tags.
+
+    Supports tags like: [laugh], [sigh], [gasp], [chuckle], [cough], [sniff], [groan], [shush], [clear throat]
+    """
+    ENC_COND_LEN = 15 * S3_SR
+    DEC_COND_LEN = 10 * S3GEN_SR
+
+    def __init__(
+        self,
+        t3: T3,
+        s3gen: S3Gen,
+        ve: VoiceEncoder,
+        tokenizer,
+        device: str,
+        conds: Conditionals = None,
+    ):
+        self.sr = S3GEN_SR
+        self.t3 = t3
+        self.s3gen = s3gen
+        self.ve = ve
+        self.tokenizer = tokenizer
+        self.device = device
+        self.conds = conds
+        self.watermarker = perth.PerthImplicitWatermarker() if PERTH_AVAILABLE else None
+
+    @classmethod
+    def from_local(cls, ckpt_dir, device) -> 'ChatterboxTurboTTS':
+        ckpt_dir = Path(ckpt_dir)
+
+        if device in ["cpu", "mps"]:
+            map_location = torch.device('cpu')
+        else:
+            map_location = None
+
+        ve = VoiceEncoder()
+        ve.load_state_dict(load_file(ckpt_dir / "ve.safetensors"))
+        ve.to(device).eval()
+
+        # Turbo-specific T3 config
+        hp = T3Config.turbo()
+
+        t3 = T3(hp)
+        t3_state = load_file(ckpt_dir / "t3_turbo_v1.safetensors")
+        if "model" in t3_state.keys():
+            t3_state = t3_state["model"][0]
+        t3.load_state_dict(t3_state)
+        # Remove unused embedding layer
+        if hasattr(t3.tfmr, 'wte'):
+            del t3.tfmr.wte
+        t3.to(device).eval()
+
+        s3gen = S3Gen(meanflow=True)
+        weights = load_file(ckpt_dir / "s3gen_meanflow.safetensors")
+        s3gen.load_state_dict(weights, strict=True)
+        s3gen.to(device).eval()
+
+        tokenizer = AutoTokenizer.from_pretrained(ckpt_dir)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        if len(tokenizer) != 50276:
+            print(f"WARNING: Tokenizer len {len(tokenizer)} != 50276")
+
+        conds = None
+        builtin_voice = ckpt_dir / "conds.pt"
+        if builtin_voice.exists():
+            conds = Conditionals.load(builtin_voice, map_location=map_location).to(device)
+
+        return cls(t3, s3gen, ve, tokenizer, device, conds=conds)
+
+    @classmethod
+    def from_pretrained(cls, device) -> 'ChatterboxTurboTTS':
+        if device == "mps" and not torch.backends.mps.is_available():
+            if not torch.backends.mps.is_built():
+                print("MPS not available because the current PyTorch install was not built with MPS enabled.")
+            else:
+                print("MPS not available because the current MacOS version is not 12.3+.")
+            device = "cpu"
+
+        # Download models to centralized location
+        local_dir = get_chatterbox_turbo_dir()
+        print(f"[FL Chatterbox Turbo] Model download path: {local_dir}")
+        download_to_local(REPO_ID, TURBO_MODEL_FILES, local_dir)
+
+        return cls.from_local(local_dir, device)
+
+    def norm_loudness(self, wav, sr, target_lufs=-27):
+        """Normalize audio loudness to target LUFS."""
+        if not PYLOUDNORM_AVAILABLE:
+            return wav
+
+        try:
+            meter = ln.Meter(sr)
+            loudness = meter.integrated_loudness(wav)
+            gain_db = target_lufs - loudness
+            gain_linear = 10.0 ** (gain_db / 20.0)
+            if math.isfinite(gain_linear) and gain_linear > 0.0:
+                wav = wav * gain_linear
+        except Exception as e:
+            print(f"Warning: Error in norm_loudness, skipping: {e}")
+
+        return wav
+
+    def prepare_conditionals(self, wav_fpath, exaggeration=0.5, norm_loudness=True):
+        """Prepare voice conditionals from reference audio."""
+        s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
+        # Ensure float32 for PyTorch compatibility
+        s3gen_ref_wav = s3gen_ref_wav.astype(np.float32)
+
+        assert len(s3gen_ref_wav) / _sr > 5.0, "Audio prompt must be longer than 5 seconds!"
+
+        if norm_loudness:
+            s3gen_ref_wav = self.norm_loudness(s3gen_ref_wav, _sr)
+            # Ensure float32 after loudness normalization (multiplication can promote to float64)
+            s3gen_ref_wav = s3gen_ref_wav.astype(np.float32)
+
+        ref_16k_wav = librosa.resample(s3gen_ref_wav, orig_sr=S3GEN_SR, target_sr=S3_SR).astype(np.float32)
+
+        s3gen_ref_wav = s3gen_ref_wav[:self.DEC_COND_LEN]
+        s3gen_ref_dict = self.s3gen.embed_ref(s3gen_ref_wav, S3GEN_SR, device=self.device)
+
+        if plen := self.t3.hp.speech_cond_prompt_len:
+            s3_tokzr = self.s3gen.tokenizer
+            t3_cond_prompt_tokens, _ = s3_tokzr.forward([ref_16k_wav[:self.ENC_COND_LEN]], max_len=plen)
+            t3_cond_prompt_tokens = torch.atleast_2d(t3_cond_prompt_tokens).to(self.device)
+
+        ve_embed = torch.from_numpy(self.ve.embeds_from_wavs([ref_16k_wav], sample_rate=S3_SR))
+        ve_embed = ve_embed.mean(axis=0, keepdim=True).to(self.device)
+
+        t3_cond = T3Cond(
+            speaker_emb=ve_embed,
+            cond_prompt_speech_tokens=t3_cond_prompt_tokens,
+            emotion_adv=exaggeration * torch.ones(1, 1, 1),
+        ).to(device=self.device)
+        self.conds = Conditionals(t3_cond, s3gen_ref_dict)
+
+    def generate(
+        self,
+        text,
+        repetition_penalty=1.2,
+        min_p=0.00,
+        top_p=0.95,
+        audio_prompt_path=None,
+        exaggeration=0.0,
+        cfg_weight=0.0,
+        temperature=0.8,
+        top_k=1000,
+        norm_loudness=True,
+    ):
+        """
+        Generate speech from text using Turbo model.
+
+        Args:
+            text: Text to speak. Supports paralinguistic tags like [laugh], [sigh], etc.
+            repetition_penalty: Penalty for token repetition (default 1.2)
+            min_p: Minimum probability threshold (ignored in Turbo)
+            top_p: Nucleus sampling threshold (default 0.95)
+            audio_prompt_path: Path to reference voice audio (min 5 seconds)
+            exaggeration: Ignored in Turbo model
+            cfg_weight: Ignored in Turbo model
+            temperature: Sampling temperature (default 0.8)
+            top_k: Top-k sampling (default 1000)
+            norm_loudness: Normalize output loudness to -27 LUFS
+
+        Returns:
+            Tensor of audio waveform
+        """
+        if audio_prompt_path:
+            self.prepare_conditionals(audio_prompt_path, exaggeration=exaggeration, norm_loudness=norm_loudness)
+        else:
+            assert self.conds is not None, "Please `prepare_conditionals` first or specify `audio_prompt_path`"
+
+        if cfg_weight > 0.0 or exaggeration > 0.0 or min_p > 0.0:
+            logger.warning("CFG, min_p and exaggeration are not supported by Turbo version and will be ignored.")
+
+        text = punc_norm(text)
+        text_tokens = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+        text_tokens = text_tokens.input_ids.to(self.device)
+
+        speech_tokens = self.t3.inference_turbo(
+            t3_cond=self.conds.t3,
+            text_tokens=text_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+        )
+
+        # Remove OOV tokens and add silence to end
+        speech_tokens = speech_tokens[speech_tokens < 6561]
+        speech_tokens = speech_tokens.to(self.device)
+        silence = torch.tensor([S3GEN_SIL, S3GEN_SIL, S3GEN_SIL]).long().to(self.device)
+        speech_tokens = torch.cat([speech_tokens, silence])
+
+        wav, _ = self.s3gen.inference(
+            speech_tokens=speech_tokens,
+            ref_dict=self.conds.gen,
+            n_cfm_timesteps=2,
+        )
+        wav = wav.squeeze(0).detach().cpu().numpy()
+
+        if self.watermarker is not None:
+            watermarked_wav = self.watermarker.apply_watermark(wav, sample_rate=self.sr)
+            return torch.from_numpy(watermarked_wav).unsqueeze(0)
+        else:
+            return torch.from_numpy(wav).unsqueeze(0)

local_chatterbox/chatterbox/vc.py

@@ -1,15 +1,24 @@
 from pathlib import Path
 
+import numpy as np
 import librosa
 import torch
+
+# Optional Perth watermarking - gracefully handle import failure
+try:
     import perth
-from huggingface_hub import hf_hub_download
+    PERTH_AVAILABLE = True
+except (ImportError, AttributeError):
+    PERTH_AVAILABLE = False
+    print("Warning: Perth watermarking not available. Audio will be generated without watermarking.")
 
 from .models.s3tokenizer import S3_SR
 from .models.s3gen import S3GEN_SR, S3Gen
+from .paths import get_chatterbox_tts_dir, download_to_local
 
 
 REPO_ID = "ResembleAI/chatterbox"
+VC_MODEL_FILES = ["s3gen.pt", "conds.pt"]
 
 
 class ChatterboxVC:
@@ -25,7 +34,7 @@ class ChatterboxVC:
         self.sr = S3GEN_SR
         self.s3gen = s3gen
         self.device = device
-        self.watermarker = perth.PerthImplicitWatermarker()
+        self.watermarker = perth.PerthImplicitWatermarker() if PERTH_AVAILABLE else None
         if ref_dict is None:
             self.ref_dict = None
         else:
@@ -67,14 +76,18 @@ class ChatterboxVC:
                 print("MPS not available because the current MacOS version is not 12.3+ and/or you do not have an MPS-enabled device on this machine.")
             device = "cpu"
 
-        for fpath in ["s3gen.pt", "conds.pt"]:
+        # Download models to centralized location: ComfyUI/models/chatterbox/chatterbox/
-            local_path = hf_hub_download(repo_id=REPO_ID, filename=fpath)
+        local_dir = get_chatterbox_tts_dir()
+        print(f"[FL Chatterbox VC] Model download path: {local_dir}")
+        download_to_local(REPO_ID, VC_MODEL_FILES, local_dir)
 
-        return cls.from_local(Path(local_path).parent, device)
+        return cls.from_local(local_dir, device)
 
     def set_target_voice(self, wav_fpath):
         ## Load reference wav
         s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
+        # Ensure float32 for PyTorch compatibility
+        s3gen_ref_wav = s3gen_ref_wav.astype(np.float32)
 
         s3gen_ref_wav = s3gen_ref_wav[:self.DEC_COND_LEN]
         self.ref_dict = self.s3gen.embed_ref(s3gen_ref_wav, S3GEN_SR, device=self.device)
@@ -99,5 +112,9 @@ class ChatterboxVC:
                 ref_dict=self.ref_dict,
             )
             wav = wav.squeeze(0).detach().cpu().numpy()
+
+            if self.watermarker is not None:
                 watermarked_wav = self.watermarker.apply_watermark(wav, sample_rate=self.sr)
                 return torch.from_numpy(watermarked_wav).unsqueeze(0)
+            else:
+                return torch.from_numpy(wav).unsqueeze(0)

pyproject.toml

@@ -1,9 +1,9 @@
 [project]
 name = "comfyui_fill-chatterbox"
 description = "Voice Clone and TTS model."
-version = "1.0.0"
+version = "1.0.5"
 license = "LICENSE"
-dependencies = ["numpy", "resampy", "librosa", "s3tokenizer", "transformers", "diffusers", "resemble-perth", "omegaconf", "conformer"]
+dependencies = ["numpy", "resampy", "librosa", "s3tokenizer", "transformers", "diffusers", "omegaconf", "conformer", "safetensors", "soundfile"]
 
 [project.urls]
 Repository = "https://github.com/filliptm/ComfyUI_Fill-ChatterBox"
@@ -12,4 +12,4 @@ Repository = "https://github.com/filliptm/ComfyUI_Fill-ChatterBox"
 [tool.comfy]
 PublisherId = "machinedelusions"
 DisplayName = "ComfyUI_Fill-ChatterBox"
-Icon = ""
+Icon = "https://avatars.githubusercontent.com/u/55672949"

requirements.txt

@@ -4,7 +4,9 @@ librosa
 s3tokenizer
 transformers
 diffusers
-resemble-perth
 omegaconf
 conformer
 safetensors
+soundfile
+# Optional watermarking (may have Python 3.12+ compatibility issues)
+# resemble-perth

tests/test_optional_perth.py

@@ -0,0 +1,43 @@
+import builtins
+import importlib
+import sys
+import tomllib
+import unittest
+from pathlib import Path
+from unittest import mock
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+
+class OptionalPerthTests(unittest.TestCase):
+    def test_import_continues_when_perth_raises_python_312_attribute_error(self):
+        original_import = builtins.__import__
+
+        def import_with_broken_perth(name, *args, **kwargs):
+            if name == "perth":
+                raise AttributeError("module 'pkgutil' has no attribute 'ImpImporter'")
+            return original_import(name, *args, **kwargs)
+
+        with mock.patch("builtins.__import__", side_effect=import_with_broken_perth):
+            package = importlib.import_module("local_chatterbox.chatterbox")
+
+        for module_name in ("tts", "tts_turbo", "mtl_tts", "vc"):
+            module = getattr(package, module_name)
+            self.assertFalse(module.PERTH_AVAILABLE)
+
+    def test_dependency_manifests_match_without_optional_perth(self):
+        requirements = {
+            line.strip()
+            for line in (ROOT / "requirements.txt").read_text(encoding="utf-8").splitlines()
+            if line.strip() and not line.lstrip().startswith("#")
+        }
+        project = tomllib.loads((ROOT / "pyproject.toml").read_text(encoding="utf-8"))["project"]
+
+        self.assertEqual(requirements, set(project["dependencies"]))
+        self.assertNotIn("resemble-perth", requirements)
+
+
+if __name__ == "__main__":
+    unittest.main()

workflows/Chatterbox.json

@@ -0,0 +1,665 @@
+{
+  "id": "d6814114-45ec-4432-8aa8-15a1d08dd55c",
+  "revision": 0,
+  "last_node_id": 26,
+  "last_link_id": 33,
+  "nodes": [
+    {
+      "id": 16,
+      "type": "FL_ChatterboxTTS",
+      "pos": [
+        604.4141939259806,
+        -616.6611297524448
+      ],
+      "size": [
+        400,
+        276
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "audio_prompt",
+          "shape": 7,
+          "type": "AUDIO",
+          "link": 18
+        }
+      ],
+      "outputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": [
+            19
+          ]
+        },
+        {
+          "name": "message",
+          "type": "STRING",
+          "links": null
+        }
+      ],
+      "properties": {
+        "aux_id": "filliptm/ComfyUI_Fill-ChatterBox",
+        "ver": "e0f3a489f9f72d3347f61b5fe8da281430eb20df",
+        "Node name for S&R": "FL_ChatterboxTTS"
+      },
+      "widgets_values": [
+        "Hello, this is a test.",
+        0.5,
+        0.5,
+        0.8,
+        2223229701,
+        "randomize",
+        false,
+        true
+      ],
+      "color": "#16727c",
+      "bgcolor": "#4F0074"
+    },
+    {
+      "id": 13,
+      "type": "PreviewAudio",
+      "pos": [
+        1048.4384618228423,
+        -616.0986702686058
+      ],
+      "size": [
+        348.560906069926,
+        88
+      ],
+      "flags": {},
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 19
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.5.1",
+        "Node name for S&R": "PreviewAudio"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 18,
+      "type": "PreviewAudio",
+      "pos": [
+        1067.0988910964618,
+        -221.83018695692084
+      ],
+      "size": [
+        348.560906069926,
+        88
+      ],
+      "flags": {},
+      "order": 8,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 21
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.5.1",
+        "Node name for S&R": "PreviewAudio"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 15,
+      "type": "FL_ChatterboxTurboTTS",
+      "pos": [
+        603.7173176061074,
+        -219.77300439909237
+      ],
+      "size": [
+        400,
+        300
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "audio_prompt",
+          "shape": 7,
+          "type": "AUDIO",
+          "link": 20
+        }
+      ],
+      "outputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": [
+            21
+          ]
+        },
+        {
+          "name": "message",
+          "type": "STRING",
+          "links": null
+        }
+      ],
+      "properties": {
+        "aux_id": "filliptm/ComfyUI_Fill-ChatterBox",
+        "ver": "e0f3a489f9f72d3347f61b5fe8da281430eb20df",
+        "Node name for S&R": "FL_ChatterboxTurboTTS"
+      },
+      "widgets_values": [
+        "Hello, this is a test. [laugh] Isn't that funny?",
+        0.8,
+        1000,
+        0.95,
+        1.2,
+        42,
+        "fixed",
+        false,
+        true
+      ],
+      "color": "#16727c",
+      "bgcolor": "#4F0074"
+    },
+    {
+      "id": 19,
+      "type": "FL_ChatterboxVC",
+      "pos": [
+        612.2979451213336,
+        145.98831998769114
+      ],
+      "size": [
+        277.015625,
+        150
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "input_audio",
+          "type": "AUDIO",
+          "link": 23
+        },
+        {
+          "name": "target_voice",
+          "type": "AUDIO",
+          "link": 28
+        }
+      ],
+      "outputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": [
+            25
+          ]
+        },
+        {
+          "name": "message",
+          "type": "STRING",
+          "links": null
+        }
+      ],
+      "properties": {
+        "aux_id": "filliptm/ComfyUI_Fill-ChatterBox",
+        "ver": "e0f3a489f9f72d3347f61b5fe8da281430eb20df",
+        "Node name for S&R": "FL_ChatterboxVC"
+      },
+      "widgets_values": [
+        644843572,
+        "randomize",
+        false,
+        false
+      ],
+      "color": "#16727c",
+      "bgcolor": "#4F0074"
+    },
+    {
+      "id": 20,
+      "type": "LoadAudio",
+      "pos": [
+        -35.50478505338169,
+        147.91684742338992
+      ],
+      "size": [
+        282.798828125,
+        136
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "AUDIO",
+          "type": "AUDIO",
+          "links": [
+            28,
+            30
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.5.1",
+        "Node name for S&R": "LoadAudio"
+      },
+      "widgets_values": [
+        "SCD_VO_BarbKing_01.wav",
+        null,
+        null
+      ]
+    },
+    {
+      "id": 12,
+      "type": "LoadAudio",
+      "pos": [
+        -34.20996863529223,
+        -207.3536826807498
+      ],
+      "size": [
+        282.798828125,
+        136
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "AUDIO",
+          "type": "AUDIO",
+          "links": [
+            18,
+            20,
+            23,
+            29,
+            32
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.5.1",
+        "Node name for S&R": "LoadAudio"
+      },
+      "widgets_values": [
+        "1225(1).mp4",
+        null,
+        null
+      ]
+    },
+    {
+      "id": 25,
+      "type": "FL_ChatterboxMultilingualTTS",
+      "pos": [
+        608.6063072713163,
+        -1064.8343442640428
+      ],
+      "size": [
+        400,
+        372
+      ],
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "audio_prompt",
+          "shape": 7,
+          "type": "AUDIO",
+          "link": 32
+        }
+      ],
+      "outputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": [
+            33
+          ]
+        },
+        {
+          "name": "message",
+          "type": "STRING",
+          "links": null
+        }
+      ],
+      "properties": {
+        "aux_id": "filliptm/ComfyUI_Fill-ChatterBox",
+        "ver": "e0f3a489f9f72d3347f61b5fe8da281430eb20df",
+        "Node name for S&R": "FL_ChatterboxMultilingualTTS"
+      },
+      "widgets_values": [
+        "Hello, this is a multilingual test.",
+        "English (en)",
+        0.5,
+        0.5,
+        0.8,
+        2,
+        0.05,
+        1,
+        0,
+        "randomize",
+        false,
+        false
+      ],
+      "color": "#16727c",
+      "bgcolor": "#4F0074"
+    },
+    {
+      "id": 26,
+      "type": "PreviewAudio",
+      "pos": [
+        1064.9230540382507,
+        -1064.4969036586663
+      ],
+      "size": [
+        348.560906069926,
+        88
+      ],
+      "flags": {},
+      "order": 11,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 33
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.5.1",
+        "Node name for S&R": "PreviewAudio"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 23,
+      "type": "FL_ChatterboxDialogTTS",
+      "pos": [
+        615.3417155903248,
+        379.15894178321156
+      ],
+      "size": [
+        400,
+        356
+      ],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "speaker_A_Audio",
+          "type": "AUDIO",
+          "link": 29
+        },
+        {
+          "name": "speaker_B_Audio",
+          "type": "AUDIO",
+          "link": 30
+        },
+        {
+          "name": "speaker_C_Audio",
+          "shape": 7,
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "speaker_D_Audio",
+          "shape": 7,
+          "type": "AUDIO",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "dialog_audio",
+          "type": "AUDIO",
+          "links": [
+            31
+          ]
+        },
+        {
+          "name": "speaker_a_audio",
+          "type": "AUDIO",
+          "links": null
+        },
+        {
+          "name": "speaker_b_audio",
+          "type": "AUDIO",
+          "links": null
+        },
+        {
+          "name": "speaker_c_audio",
+          "type": "AUDIO",
+          "links": null
+        },
+        {
+          "name": "speaker_d_audio",
+          "type": "AUDIO",
+          "links": null
+        },
+        {
+          "name": "message",
+          "type": "STRING",
+          "links": null
+        }
+      ],
+      "properties": {
+        "aux_id": "filliptm/ComfyUI_Fill-ChatterBox",
+        "ver": "e0f3a489f9f72d3347f61b5fe8da281430eb20df",
+        "Node name for S&R": "FL_ChatterboxDialogTTS"
+      },
+      "widgets_values": [
+        "SPEAKER A: Test test\nSPEAKER B: 1 2 3",
+        0.5,
+        0.5,
+        0.8,
+        0,
+        "randomize",
+        false,
+        false
+      ],
+      "color": "#16727c",
+      "bgcolor": "#4F0074"
+    },
+    {
+      "id": 24,
+      "type": "PreviewAudio",
+      "pos": [
+        1098.8488879717245,
+        378.6384238512216
+      ],
+      "size": [
+        348.560906069926,
+        88
+      ],
+      "flags": {},
+      "order": 10,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 31
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.5.1",
+        "Node name for S&R": "PreviewAudio"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 21,
+      "type": "PreviewAudio",
+      "pos": [
+        1023.4075617775408,
+        145.45042516865564
+      ],
+      "size": [
+        348.560906069926,
+        88
+      ],
+      "flags": {},
+      "order": 9,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 25
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "cnr_id": "comfy-core",
+        "ver": "0.5.1",
+        "Node name for S&R": "PreviewAudio"
+      },
+      "widgets_values": []
+    }
+  ],
+  "links": [
+    [
+      18,
+      12,
+      0,
+      16,
+      0,
+      "AUDIO"
+    ],
+    [
+      19,
+      16,
+      0,
+      13,
+      0,
+      "AUDIO"
+    ],
+    [
+      20,
+      12,
+      0,
+      15,
+      0,
+      "AUDIO"
+    ],
+    [
+      21,
+      15,
+      0,
+      18,
+      0,
+      "AUDIO"
+    ],
+    [
+      23,
+      12,
+      0,
+      19,
+      0,
+      "AUDIO"
+    ],
+    [
+      25,
+      19,
+      0,
+      21,
+      0,
+      "AUDIO"
+    ],
+    [
+      28,
+      20,
+      0,
+      19,
+      1,
+      "AUDIO"
+    ],
+    [
+      29,
+      12,
+      0,
+      23,
+      0,
+      "AUDIO"
+    ],
+    [
+      30,
+      20,
+      0,
+      23,
+      1,
+      "AUDIO"
+    ],
+    [
+      31,
+      23,
+      0,
+      24,
+      0,
+      "AUDIO"
+    ],
+    [
+      32,
+      12,
+      0,
+      25,
+      0,
+      "AUDIO"
+    ],
+    [
+      33,
+      25,
+      0,
+      26,
+      0,
+      "AUDIO"
+    ]
+  ],
+  "groups": [
+    {
+      "id": 1,
+      "title": "Group",
+      "bounding": [
+        -45.50478505338169,
+        -1138.4343442640427,
+        1502.9145790950322,
+        1883.5932860472542
+      ],
+      "color": "#3f789e",
+      "font_size": 24,
+      "flags": {}
+    }
+  ],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.4815769309148436,
+      "offset": [
+        732.4394040290586,
+        1134.7214387004035
+      ]
+    },
+    "workflowRendererVersion": "LG",
+    "frontendVersion": "1.34.9"
+  },
+  "version": 0.4
+}