- 
+ 
ComfyUI Chatterbox
diff --git a/__init__.py b/__init__.py
index 6181aae..709d2c2 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,7 +1,59 @@
import os
+import sys
+import types
+import logging
import folder_paths
+from importlib import metadata as importlib_metadata
+
+# Configure a logger for the entire custom node package
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+# Add a handler if none exist to avoid duplicate logs
+if not logger.hasHandlers():
+ # Use stdout for info, stderr for errors
+ handler = logging.StreamHandler(sys.stdout)
+ formatter = logging.Formatter(f"[%(name)s] %(message)s")
+ handler.setFormatter(formatter)
+ logger.addHandler(handler)
+
+
+# Monkey-Patch for 'chatterbox-tts' version
+# original_importlib_version_func = importlib_metadata.version
+# def patched_version_lookup(package_name):
+# if package_name == "chatterbox-tts":
+# return "local-vendored"
+# return original_importlib_version_func(package_name)
+# importlib_metadata.version = patched_version_lookup
+
+
+# easily check if the real 'perth' is being used.
+try:
+ import perth
+ # A simple check to ensure it's not our mock from a previous run
+ if not hasattr(perth, '_is_mock'):
+ logger.info("Found and using 'resemble-perth' library for watermarking.")
+except ImportError:
+ logger.warning("'resemble-perth' not found. Watermarking will be unavailable.")
+ class DummyPerthImplicitWatermarker:
+ def apply_watermark(self, wav, sample_rate):
+ logger.warning("Watermarking skipped: 'resemble-perth' is not installed.")
+ return wav
+ perth_mock = types.ModuleType('perth')
+ perth_mock.PerthImplicitWatermarker = DummyPerthImplicitWatermarker
+ # Flag to identify our mock module
+ perth_mock._is_mock = True
+ sys.modules['perth'] = perth_mock
+
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+src_dir = os.path.join(current_dir, "src")
+if src_dir not in sys.path:
+ sys.path.insert(0, src_dir)
+
+
from .nodes import ChatterboxTTSNode, ChatterboxVCNode
-from .modules.chatterbox_handler import CHATTERBOX_MODEL_SUBDIR, DEFAULT_MODEL_PACK_NAME
+from .modules.chatterbox_handler import CHATTERBOX_MODEL_SUBDIR
NODE_CLASS_MAPPINGS = {
"ChatterboxTTS": ChatterboxTTSNode,
@@ -13,25 +65,21 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"ChatterboxVC": "Chatterbox Voice Conversion 🗣️",
}
-# WEB_DIRECTORY = "./js"
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
+# Model path setup for ComfyUI
chatterbox_models_full_path = os.path.join(folder_paths.models_dir, CHATTERBOX_MODEL_SUBDIR)
if not os.path.exists(chatterbox_models_full_path):
try:
os.makedirs(chatterbox_models_full_path, exist_ok=True)
- #print(f"ChatterboxTTS/VC Init: Created models directory at {chatterbox_models_full_path}")
except OSError as e:
- print(f"ChatterboxTTS/VC Init: Error creating models directory {chatterbox_models_full_path}: {e}")
+ logger.error(f"Error creating models directory {chatterbox_models_full_path}: {e}")
-if CHATTERBOX_MODEL_SUBDIR not in folder_paths.folder_names_and_paths:
- folder_paths.folder_names_and_paths[CHATTERBOX_MODEL_SUBDIR] = (
- [chatterbox_models_full_path],
- folder_paths.supported_pt_extensions
- )
- #print(f"ChatterboxTTS/VC Init: Registered '{CHATTERBOX_MODEL_SUBDIR}' with ComfyUI model paths: {chatterbox_models_full_path}")
+# Register the tts/chatterbox path with ComfyUI
+tts_chatterbox_path = os.path.join(folder_paths.models_dir, "tts")
+if "tts" not in folder_paths.folder_names_and_paths:
+ supported_exts = folder_paths.supported_pt_extensions.union({".safetensors"})
+ folder_paths.folder_names_and_paths["tts"] = ([tts_chatterbox_path], supported_exts)
else:
- if chatterbox_models_full_path not in folder_paths.folder_names_and_paths[CHATTERBOX_MODEL_SUBDIR][0]:
- folder_paths.folder_names_and_paths[CHATTERBOX_MODEL_SUBDIR][0].append(chatterbox_models_full_path)
- #print(f"ChatterboxTTS/VC Init: Appended path for '{CHATTERBOX_MODEL_SUBDIR}': {chatterbox_models_full_path}")
-
+ if tts_chatterbox_path not in folder_paths.folder_names_and_paths["tts"][0]:
+ folder_paths.folder_names_and_paths["tts"][0].append(tts_chatterbox_path)
\ No newline at end of file
diff --git a/modules/chatterbox_handler.py b/modules/chatterbox_handler.py
index 26b6c13..092c569 100644
--- a/modules/chatterbox_handler.py
+++ b/modules/chatterbox_handler.py
@@ -1,222 +1,98 @@
import os
-import sys
+import logging
import torch
import random
import numpy as np
import folder_paths
-import hashlib
from huggingface_hub import hf_hub_download
-current_script_path = os.path.dirname(os.path.abspath(__file__))
-chatterbox_src_path = os.path.abspath(os.path.join(current_script_path, '..', 'src'))
-if chatterbox_src_path not in sys.path:
- sys.path.insert(0, chatterbox_src_path)
-
-
from chatterbox.tts import ChatterboxTTS
-from chatterbox.vc import ChatterboxVC # <--- Added VC import
+from chatterbox.vc import ChatterboxVC
-TTS_MODEL_CACHE = {}
-VC_MODEL_CACHE = {}
-CHATTERBOX_MODEL_SUBDIR = "chatterbox_tts"
+logger = logging.getLogger(__name__)
+
+CHATTERBOX_MODEL_SUBDIR = os.path.join("tts", "chatterbox")
CHATTERBOX_REPO_ID = "ResembleAI/chatterbox"
-
-CHATTERBOX_FILES_TO_DOWNLOAD = ["ve.pt", "t3_cfg.pt", "s3gen.pt", "tokenizer.json", "conds.pt"]
+CHATTERBOX_FILES_TO_DOWNLOAD = ["ve.safetensors", "t3_cfg.safetensors", "s3gen.safetensors", "tokenizer.json", "conds.pt"]
DEFAULT_MODEL_PACK_NAME = "resembleai_default_voice"
def get_chatterbox_model_pack_names():
- """Returns a list of available Chatterbox model pack names (subdirectories)."""
chatterbox_models_base_path = os.path.join(folder_paths.models_dir, CHATTERBOX_MODEL_SUBDIR)
if not os.path.isdir(chatterbox_models_base_path):
os.makedirs(chatterbox_models_base_path, exist_ok=True)
- #print(f"ChatterboxTTS/VC: Created models directory at {chatterbox_models_base_path}")
- return []
-
+ return [DEFAULT_MODEL_PACK_NAME]
packs = [d for d in os.listdir(chatterbox_models_base_path) if os.path.isdir(os.path.join(chatterbox_models_base_path, d))]
- if not packs:
- return []
- return packs
+ # Ensure default is first if it exists
+ if DEFAULT_MODEL_PACK_NAME in packs:
+ packs.insert(0, packs.pop(packs.index(DEFAULT_MODEL_PACK_NAME)))
+
+ # Return default even if folder doesn't exist yet, to prompt the user to download it.
+ return packs if packs else [DEFAULT_MODEL_PACK_NAME]
def get_model_pack_path(model_pack_name):
- """Gets the full path to a specific model pack."""
- if not model_pack_name:
+ # Added check for None or empty string to prevent errors
+ if not model_pack_name:
return None
return os.path.join(folder_paths.models_dir, CHATTERBOX_MODEL_SUBDIR, model_pack_name)
-def _download_file_from_hf(repo_id, filename, local_dir, local_filename=None):
- if local_filename is None:
- local_filename = filename
- destination = os.path.join(local_dir, local_filename)
-
+def _download_file_from_hf(repo_id, filename, local_dir):
+ destination = os.path.join(local_dir, filename)
if not os.path.exists(destination):
- print(f"ChatterboxTTS/VC: Downloading '{filename}' from '{repo_id}' to '{destination}'...")
+ logger.info(f"Downloading '{filename}' from '{repo_id}'...")
try:
- hf_hub_download(
- repo_id=repo_id,
- filename=filename,
- local_dir=local_dir,
- )
- if filename != local_filename and os.path.exists(os.path.join(local_dir, filename)) and not os.path.exists(destination):
- os.rename(os.path.join(local_dir, filename), destination)
- print(f"ChatterboxTTS/VC: Successfully downloaded '{local_filename}'.")
+ hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir, local_dir_use_symlinks=False, resume_download=True)
+ logger.info(f"Successfully downloaded '{filename}'.")
return True
except Exception as e:
- print(f"ChatterboxTTS/VC: Failed to download '{filename}' from '{repo_id}'. Error: {e}")
- if os.path.exists(destination + ".incomplete"):
- os.remove(destination + ".incomplete")
+ logger.error(f"Failed to download '{filename}': {e}")
+ if os.path.exists(destination + ".incomplete"): os.remove(destination + ".incomplete")
return False
- else:
- return True
+ return True
def download_chatterbox_model_pack_if_missing(model_pack_name):
- """Downloads all necessary model files for a given pack name if they don't exist."""
ckpt_dir = get_model_pack_path(model_pack_name)
if not ckpt_dir:
- print(f"ChatterboxTTS/VC: Invalid model pack name '{model_pack_name}', cannot download.")
+ logger.warning(f"Invalid model pack name '{model_pack_name}', cannot download.")
return False
-
os.makedirs(ckpt_dir, exist_ok=True)
-
- all_files_successfully_managed = True
+ all_files_ok = all(_download_file_from_hf(CHATTERBOX_REPO_ID, f, ckpt_dir) for f in CHATTERBOX_FILES_TO_DOWNLOAD)
+ if not all_files_ok:
+ logger.error(f"Some files failed to download for model pack '{model_pack_name}'. Check logs.")
+ return all_files_ok
- vc_required_files = ["s3gen.pt", "conds.pt"]
-
- files_to_check = CHATTERBOX_FILES_TO_DOWNLOAD
-
- for file_name in files_to_check:
- if not _download_file_from_hf(CHATTERBOX_REPO_ID, file_name, ckpt_dir):
- all_files_successfully_managed = False
- if file_name in vc_required_files:
- print(f"ChatterboxTTS/VC: Critical file '{file_name}' for VC failed to download for pack '{model_pack_name}'.")
- # return False
-
- if all_files_successfully_managed:
- pass #print(f"ChatterboxTTS/VC: Verified/Downloaded all files for model pack '{model_pack_name}' into '{ckpt_dir}'.")
- else:
- print(f"ChatterboxTTS/VC: Some files failed to download for model pack '{model_pack_name}'. Check logs.")
- return all_files_successfully_managed
-
-
-def load_chatterbox_tts_model(model_pack_name, device_str="cuda"):
+def load_chatterbox_models(model_pack_name, device):
+ """Loads both TTS and VC models for a given pack onto a specified device."""
ckpt_dir = get_model_pack_path(model_pack_name)
if not ckpt_dir:
- raise ValueError(f"ChatterboxTTS: Invalid model_pack_name: {model_pack_name}")
- #print(f"ChatterboxTTS: Attempting to load TTS model pack '{model_pack_name}' from '{ckpt_dir}'.")
+ raise ValueError(f"Invalid model_pack_name: {model_pack_name}")
+
if not download_chatterbox_model_pack_if_missing(model_pack_name):
- print(f"ChatterboxTTS: Warning - Not all TTS model files could be verified/downloaded for '{model_pack_name}'. Loading may fail.")
- if not os.path.isdir(ckpt_dir):
- raise FileNotFoundError(f"ChatterboxTTS: Model pack directory '{model_pack_name}' not found and could not be created at '{ckpt_dir}'.")
- #print(f"ChatterboxTTS: Loading ChatterboxTTS model from local directory: {ckpt_dir} onto device: {device_str}")
- try:
- model = ChatterboxTTS.from_local(ckpt_dir, device=device_str)
- except Exception as e:
- print(f"ChatterboxTTS: Error during ChatterboxTTS.from_local('{ckpt_dir}', device='{device_str}'): {e}")
- print(f"ChatterboxTTS: Please ensure all required model files ({', '.join(CHATTERBOX_FILES_TO_DOWNLOAD)}) are present in {ckpt_dir} or can be downloaded from {CHATTERBOX_REPO_ID}.")
- raise
- return model
-
-def get_cached_chatterbox_tts_model(model_pack_name, device_str="cuda"):
- """Loads and caches the ChatterboxTTS model."""
- if not model_pack_name:
- available_packs = get_chatterbox_model_pack_names()
- model_pack_name = available_packs[0] if available_packs else DEFAULT_MODEL_PACK_NAME
- print(f"ChatterboxTTS: No model pack specified for TTS, using '{model_pack_name}'.")
-
- cache_key = (model_pack_name, device_str, "tts")
- current_model = TTS_MODEL_CACHE.get(cache_key)
- model_device_correct = False
- if current_model is not None and hasattr(current_model, 'device'):
- try:
- if str(current_model.device) == device_str:
- model_device_correct = True
- except Exception as e:
- print(f"ChatterboxTTS: Error checking cached TTS model device: {e}. Will reload.")
-
- if current_model is not None and model_device_correct:
- return current_model
- else:
- if current_model is not None and not model_device_correct:
- print(f"ChatterboxTTS: Device mismatch for cached TTS model '{model_pack_name}'. Reloading.")
- else:
- print(f"ChatterboxTTS: TTS Model for '{model_pack_name}' on '{device_str}' not in cache. Loading...")
+ logger.warning(f"Not all model files could be verified for '{model_pack_name}'. Loading may fail.")
- TTS_MODEL_CACHE[cache_key] = load_chatterbox_tts_model(model_pack_name, device_str)
- return TTS_MODEL_CACHE[cache_key]
-
-
-def load_chatterbox_vc_model(model_pack_name, device_str="cuda"):
- """Loads the ChatterboxVC model from a specified pack."""
- ckpt_dir = get_model_pack_path(model_pack_name)
- if not ckpt_dir:
- raise ValueError(f"ChatterboxVC: Invalid model_pack_name: {model_pack_name}")
-
- #print(f"ChatterboxVC: Attempting to load VC model pack '{model_pack_name}' from '{ckpt_dir}'.")
-
- if not download_chatterbox_model_pack_if_missing(model_pack_name):
- vc_min_files = ["s3gen.pt", "conds.pt"]
- for f_name in vc_min_files:
- if not os.path.exists(os.path.join(ckpt_dir, f_name)):
- print(f"ChatterboxVC: Critical file '{f_name}' for VC is missing from pack '{model_pack_name}' after download attempt. Loading will likely fail.")
- break
- else:
- print(f"ChatterboxVC: Warning - Not all model files could be verified/downloaded for '{model_pack_name}', but attempting to load VC with available files.")
-
-
if not os.path.isdir(ckpt_dir):
- raise FileNotFoundError(f"ChatterboxVC: Model pack directory '{model_pack_name}' not found and could not be created at '{ckpt_dir}'.")
-
- # print(f"ChatterboxVC: Loading ChatterboxVC model from local directory: {ckpt_dir} onto device: {device_str}")
+ raise FileNotFoundError(f"Model pack directory '{model_pack_name}' not found at '{ckpt_dir}'.")
+
try:
- model = ChatterboxVC.from_local(ckpt_dir, device=device_str)
+ logger.info(f"Loading Chatterbox TTS model from {ckpt_dir} onto {device}")
+ tts_model = ChatterboxTTS.from_local(ckpt_dir, device=device)
except Exception as e:
- print(f"ChatterboxVC: Error during ChatterboxVC.from_local('{ckpt_dir}', device='{device_str}'): {e}")
- print(f"ChatterboxVC: Please ensure at least 's3gen.pt' and optionally 'conds.pt' (for default voice) are present in {ckpt_dir} or can be downloaded from {CHATTERBOX_REPO_ID}.")
+ logger.error(f"Error loading ChatterboxTTS from '{ckpt_dir}': {e}", exc_info=True)
+ raise
+
+ try:
+ logger.info(f"Loading Chatterbox VC model from {ckpt_dir} onto {device}")
+ vc_model = ChatterboxVC.from_local(ckpt_dir, device=device)
+ except Exception as e:
+ logger.error(f"Error loading ChatterboxVC from '{ckpt_dir}': {e}", exc_info=True)
raise
- return model
-
-def get_cached_chatterbox_vc_model(model_pack_name, device_str="cuda"):
- """Loads and caches the ChatterboxVC model."""
- if not model_pack_name:
- available_packs = get_chatterbox_model_pack_names()
- model_pack_name = available_packs[0] if available_packs else DEFAULT_MODEL_PACK_NAME
- print(f"ChatterboxVC: No model pack specified for VC, using '{model_pack_name}'.")
-
- cache_key = (model_pack_name, device_str, "vc")
- current_model = VC_MODEL_CACHE.get(cache_key)
- model_device_correct = False
- if current_model is not None and hasattr(current_model, 'device'):
- try:
- if str(current_model.device) == device_str:
- model_device_correct = True
- except Exception as e:
- print(f"ChatterboxVC: Error checking cached VC model device: {e}. Will reload.")
-
- if current_model is not None and model_device_correct:
- return current_model
- else:
- if current_model is not None and not model_device_correct:
- print(f"ChatterboxVC: Device mismatch for cached VC model '{model_pack_name}'. Reloading.")
- else:
- print(f"ChatterboxVC: VC Model for '{model_pack_name}' on '{device_str}' not in cache. Loading...")
- VC_MODEL_CACHE[cache_key] = load_chatterbox_vc_model(model_pack_name, device_str)
- return VC_MODEL_CACHE[cache_key]
+ return tts_model, vc_model
def set_chatterbox_seed(seed: int):
- """Sets the seed for Chatterbox TTS/VC."""
MAX_NUMPY_SEED = 2**32 - 1
- if seed == 0:
- actual_seed_for_torch_random = random.randint(1, 0xffffffffffffffff)
- actual_seed_for_numpy = random.randint(1, MAX_NUMPY_SEED)
- else:
- actual_seed_for_torch_random = seed
- # fast fix range for seed, sorry
- actual_seed_for_numpy = seed % MAX_NUMPY_SEED
-
+ actual_seed_for_torch_random = random.randint(1, 0xffffffffffffffff) if seed == 0 else seed
+ actual_seed_for_numpy = random.randint(1, MAX_NUMPY_SEED) if seed == 0 else (seed % MAX_NUMPY_SEED)
torch.manual_seed(actual_seed_for_torch_random)
- if torch.cuda.is_available():
- torch.cuda.manual_seed(actual_seed_for_torch_random)
- torch.cuda.manual_seed_all(actual_seed_for_torch_random)
+ if torch.cuda.is_available(): torch.cuda.manual_seed_all(actual_seed_for_torch_random)
random.seed(actual_seed_for_torch_random)
- np.random.seed(actual_seed_for_numpy)
-
\ No newline at end of file
+ np.random.seed(actual_seed_for_numpy)
\ No newline at end of file
diff --git a/nodes.py b/nodes.py
index 05833e8..99f3f52 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1,231 +1,325 @@
import os
import torch
-import folder_paths
import tempfile
import soundfile as sf
import numpy as np
+import logging
+import perth
+
+import comfy.model_management as mm
+import comfy.model_patcher
+from comfy.utils import ProgressBar
from .modules.chatterbox_handler import (
get_chatterbox_model_pack_names,
- get_cached_chatterbox_tts_model,
- get_cached_chatterbox_vc_model,
+ load_chatterbox_models,
set_chatterbox_seed,
- CHATTERBOX_MODEL_SUBDIR,
DEFAULT_MODEL_PACK_NAME
)
+logger = logging.getLogger(__name__)
+
+CHATTERBOX_PATCHER_CACHE = {}
+
+class ChatterboxModelWrapper(torch.nn.Module):
+ """
+ A simple torch.nn.Module wrapper for the Chatterbox models.
+ This allows ComfyUI's model management to treat our custom models like any other
+ torch module, enabling device placement (.to()) and other standard operations.
+ """
+ def __init__(self, model_pack_name):
+ super().__init__()
+ self.model_pack_name = model_pack_name
+ self.tts_model = None
+ self.vc_model = None
+
+ def load_model(self, device):
+ self.tts_model, self.vc_model = load_chatterbox_models(self.model_pack_name, device)
+
+class ChatterboxPatcher(comfy.model_patcher.ModelPatcher):
+ """
+ Custom ModelPatcher for Chatterbox. This class hooks into ComfyUI's
+ model management system (loading, offloading) to handle our non-standard models.
+ """
+ def __init__(self, model, *args, **kwargs):
+ super().__init__(model, *args, **kwargs)
+
+ def patch_model(self, device_to=None, *args, **kwargs):
+ """
+ This method is called by ComfyUI's model manager when it's time to load
+ the model onto the target device (usually the GPU). Our responsibility here
+ is to ensure the model weights are loaded from disk if they haven't been already.
+ """
+ target_device = self.load_device
+
+ # The core loading logic: If the model isn't in memory, load it from disk.
+ if self.model.tts_model is None:
+ logger.info(f"Loading Chatterbox models for '{self.model.model_pack_name}' to {target_device}...")
+ self.model.load_model(target_device)
+ self.model.model_loaded_weight_memory = self.size
+ else:
+ logger.info(f"Chatterbox models for '{self.model.model_pack_name}' already in memory.")
+
+ return super().patch_model(device_to=target_device, *args, **kwargs)
+
+ def unpatch_model(self, device_to=None, unpatch_weights=True, *args, **kwargs):
+ """
+ This method is called by ComfyUI's model manager to offload the model
+ (usually to the CPU) and free up VRAM.
+ """
+ if unpatch_weights:
+ logger.info(f"Offloading Chatterbox models for '{self.model.model_pack_name}' to {device_to}...")
+ self.model.tts_model = None
+ self.vc_model = None
+ # Reset memory footprint
+ self.model.model_loaded_weight_memory = 0
+ mm.soft_empty_cache()
+ return super().unpatch_model(device_to, unpatch_weights, *args, **kwargs)
+
class ChatterboxTTSNode:
@classmethod
def INPUT_TYPES(cls):
- available_model_packs = get_chatterbox_model_pack_names()
- displayed_packs = [DEFAULT_MODEL_PACK_NAME] + [p for p in available_model_packs if p != DEFAULT_MODEL_PACK_NAME]
- if not displayed_packs:
- displayed_packs = [DEFAULT_MODEL_PACK_NAME]
+ return {"required": {
+ "model_pack_name": (get_chatterbox_model_pack_names(), {
+ "default": DEFAULT_MODEL_PACK_NAME,
+ "tooltip": "Select the Chatterbox voice model pack to use for generation."
+ }),
+ "text": ("STRING", {
+ "multiline": True,
+ "default": "Hello, this is a test of Chatterbox TTS in ComfyUI.",
+ "tooltip": "Text to be synthesized into speech."
+ }),
+ "max_new_tokens": ("INT", {
+ "default": 1000, "min": 16, "max": 4000, "step": 8,
+ "tooltip": "Maximum number of audio tokens to generate. 25 tokens ≈ 1 second. The hard limit is 4096 tokens (≈ 163 seconds)."
+ }),
+ "flow_cfg_scale": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "CFG scale for the mel spectrogram decoder (flow matching). Higher values increase adherence to content and timbre but may reduce naturalness."}),
+ "exaggeration": ("FLOAT", {
+ "default": 0.5, "min": 0.25, "max": 2.0, "step": 0.05,
+ "tooltip": "Controls the expressiveness and emotional intensity. Higher values lead to more exaggerated prosody."
+ }),
+ "temperature": ("FLOAT", {
+ "default": 0.8, "min": 0.05, "max": 5.0, "step": 0.05,
+ "tooltip": "Controls the randomness of the sampling process. Higher values produce more diverse speech, while lower values are more deterministic."
+ }),
+ "cfg_weight": ("FLOAT", {
+ "default": 0.5, "min": 0.2, "max": 1.0, "step": 0.05,
+ "tooltip": "Classifier-Free Guidance (CFG) weight. Controls how strongly the model adheres to the text prompt. Higher values may reduce naturalness."
+ }),
+ "repetition_penalty": ("FLOAT", {
+ "default": 1.2, "min": 1.0, "max": 2.0, "step": 0.1,
+ "tooltip": "Penalizes repeated tokens to discourage monotonous or repetitive speech. A value of 1.0 means no penalty."
+ }),
+ "min_p": ("FLOAT", {
+ "default": 0.05, "min": 0.0, "max": 1.0, "step": 0.01,
+ "tooltip": "Sets a minimum probability threshold for nucleus sampling (Min-P). Filters out tokens with very low probability."
+ }),
+ "top_p": ("FLOAT", {
+ "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01,
+ "tooltip": "Nucleus sampling (Top-P) parameter. The model samples from the smallest set of tokens whose cumulative probability exceeds this value."
+ }),
+ "seed": ("INT", {
+ "default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True,
+ "tooltip": "Seed for random number generation. A value of 0 will use a random seed."
+ }),
+ "use_watermark": ("BOOLEAN", {
+ "default": False,
+ "tooltip": "Enable or disable the audio watermark. Requires 'resemble-perth' to be installed."
+ }),
+ }, "optional": {"audio_prompt": ("AUDIO",),}}
- return {
- "required": {
- "model_pack_name": (displayed_packs, {"default": DEFAULT_MODEL_PACK_NAME if DEFAULT_MODEL_PACK_NAME in displayed_packs else (displayed_packs[0] if displayed_packs else DEFAULT_MODEL_PACK_NAME)}),
- "text": ("STRING", {"multiline": True, "default": "Hello, this is a test of Chatterbox TTS in ComfyUI."}),
- "exaggeration": ("FLOAT", {"default": 0.5, "min": 0.25, "max": 2.0, "step": 0.05}),
- "temperature": ("FLOAT", {"default": 0.8, "min": 0.05, "max": 5.0, "step": 0.05}),
- "cfg_weight": ("FLOAT", {"default": 0.5, "min": 0.2, "max": 1.0, "step": 0.05}),
- "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}),
- "device": (["cuda", "cpu"], {"default": "cuda" if torch.cuda.is_available() else "cpu"}),
- },
- "optional": {
- "audio_prompt": ("AUDIO",),
- }
- }
+ RETURN_TYPES = ("AUDIO",); RETURN_NAMES = ("audio",); FUNCTION = "synthesize"; CATEGORY = "audio/generation"; OUTPUT_NODE = True
- RETURN_TYPES = ("AUDIO",)
- RETURN_NAMES = ("audio",)
- FUNCTION = "synthesize"
- CATEGORY = "audio/generation"
- OUTPUT_NODE = True
-
- def synthesize(self, model_pack_name, text, exaggeration, temperature, cfg_weight, seed, device, audio_prompt=None):
+ def synthesize(self, model_pack_name, text, max_new_tokens, flow_cfg_scale, exaggeration, temperature, cfg_weight, repetition_penalty, min_p, top_p, seed, use_watermark, audio_prompt=None):
if not text.strip():
- #print("Chatterbox TTS: Empty text provided, returning silent audio.")
- dummy_sr = 24000
- silent_waveform = torch.zeros((1, dummy_sr), dtype=torch.float32, device="cpu")
+ logger.info("Empty text provided, returning silent audio.")
+ dummy_sr = 24000; silent_waveform = torch.zeros((1, dummy_sr), dtype=torch.float32, device="cpu")
return ({"waveform": silent_waveform.unsqueeze(0), "sample_rate": dummy_sr},)
- try:
- chatterbox_model = get_cached_chatterbox_tts_model(model_pack_name, device_str=device)
- except Exception as e:
- print(f"ChatterboxTTS: Error loading/downloading TTS model pack '{model_pack_name}': {e}")
- dummy_sr = 24000
- silent_waveform = torch.zeros((1, dummy_sr), dtype=torch.float32, device="cpu")
+ cache_key = model_pack_name
+ if cache_key not in CHATTERBOX_PATCHER_CACHE:
+ load_device = mm.get_torch_device()
+ logger.info(f"Creating Chatterbox ModelPatcher for {model_pack_name} on device {load_device}")
+ model_wrapper = ChatterboxModelWrapper(model_pack_name)
+ patcher = ChatterboxPatcher(
+ model=model_wrapper,
+ load_device=load_device,
+ offload_device=mm.unet_offload_device(),
+ size=int(1.5 * 1024**3)
+ )
+ CHATTERBOX_PATCHER_CACHE[cache_key] = patcher
+
+ patcher = CHATTERBOX_PATCHER_CACHE[cache_key]
+
+ mm.load_model_gpu(patcher)
+ tts_model = patcher.model.tts_model
+
+ if tts_model is None:
+ logger.error("TTS model failed to load. Please check logs for download or loading errors.")
+ dummy_sr = 24000; silent_waveform = torch.zeros((1, dummy_sr), dtype=torch.float32, device="cpu")
return ({"waveform": silent_waveform.unsqueeze(0), "sample_rate": dummy_sr},)
set_chatterbox_seed(seed)
+
+ is_perth_installed = not getattr(perth, '_is_mock', False)
+ if use_watermark and not is_perth_installed:
+ logger.warning("Watermarking is enabled, but 'resemble-perth' is not installed. Output will not be watermarked.")
+
+ original_watermarker = tts_model.watermarker
+ if not use_watermark:
+ class TmpDummyWatermarker:
+ def apply_watermark(self, wav, sample_rate): return wav
+ tts_model.watermarker = TmpDummyWatermarker()
+ if is_perth_installed: logger.info("Watermarking disabled by user.")
+
+ wav_tensor_chatterbox = None; audio_prompt_path_temp = None
- audio_prompt_path_temp = None
- if audio_prompt is not None and \
- audio_prompt.get("waveform") is not None and \
- audio_prompt["waveform"].numel() > 0:
-
- waveform_in = audio_prompt["waveform"]
- sample_rate_in = audio_prompt["sample_rate"]
- waveform_cpu = waveform_in.cpu()
- if waveform_cpu.shape[0] > 1:
- print(f"ChatterboxTTS: Audio prompt has batch size {waveform_cpu.shape[0]}, using first item.")
- current_waveform = waveform_cpu[0]
- if current_waveform.shape[0] > 1:
- current_waveform = torch.mean(current_waveform, dim=0)
- else:
- current_waveform = current_waveform.squeeze(0)
-
- processed_audio_prompt = current_waveform.numpy().astype(np.float32)
- try:
- with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_wav:
- audio_prompt_path_temp = tmp_wav.name
- sf.write(audio_prompt_path_temp, processed_audio_prompt, sample_rate_in)
- #print(f"ChatterboxTTS: Using audio prompt from temp file: {audio_prompt_path_temp}")
- except Exception as e:
- print(f"ChatterboxTTS: Error writing temp audio prompt file: {e}")
- audio_prompt_path_temp = None
+ pbar = ProgressBar(max_new_tokens)
try:
- wav_tensor_chatterbox = chatterbox_model.generate(
+ if audio_prompt and audio_prompt.get("waveform") is not None and audio_prompt["waveform"].numel() > 0:
+ with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_wav:
+ audio_prompt_path_temp = tmp_wav.name
+ waveform_in = audio_prompt["waveform"]; sample_rate_in = audio_prompt["sample_rate"]
+ waveform_cpu = waveform_in.cpu()[0]
+ current_waveform = torch.mean(waveform_cpu, dim=0) if waveform_cpu.shape[0] > 1 else waveform_cpu.squeeze(0)
+ sf.write(audio_prompt_path_temp, current_waveform.numpy().astype(np.float32), sample_rate_in)
+
+ wav_tensor_chatterbox = tts_model.generate(
text,
audio_prompt_path=audio_prompt_path_temp,
exaggeration=exaggeration,
temperature=temperature,
- cfg_weight=cfg_weight
- )
+ cfg_weight=cfg_weight,
+ repetition_penalty=repetition_penalty,
+ min_p=min_p,
+ top_p=top_p,
+ pbar=pbar,
+ max_new_tokens=max_new_tokens,
+ flow_cfg_scale=flow_cfg_scale
+ )
except Exception as e:
- print(f"ChatterboxTTS: Error during TTS generation: {e}")
- dummy_sr = 24000
- silent_waveform = torch.zeros((1, dummy_sr), dtype=torch.float32, device="cpu")
+ logger.error(f"Error during TTS generation: {e}", exc_info=True)
+ dummy_sr = 24000; silent_waveform = torch.zeros((1, dummy_sr), dtype=torch.float32, device="cpu")
return ({"waveform": silent_waveform.unsqueeze(0), "sample_rate": dummy_sr},)
finally:
+ tts_model.watermarker = original_watermarker
if audio_prompt_path_temp and os.path.exists(audio_prompt_path_temp):
- try:
- os.remove(audio_prompt_path_temp)
- except Exception as e:
- print(f"ChatterboxTTS: Error removing temp audio prompt file {audio_prompt_path_temp}: {e}")
-
- wav_tensor_comfy = wav_tensor_chatterbox.cpu().unsqueeze(0)
- return ({"waveform": wav_tensor_comfy, "sample_rate": chatterbox_model.sr},)
+ try: os.remove(audio_prompt_path_temp)
+ except Exception as e: logger.error(f"Error removing temp audio prompt file: {e}")
+ wav_tensor_comfy = wav_tensor_chatterbox.cpu().unsqueeze(0)
+ return ({"waveform": wav_tensor_comfy, "sample_rate": tts_model.sr},)
class ChatterboxVCNode:
@classmethod
def INPUT_TYPES(cls):
- available_model_packs = get_chatterbox_model_pack_names()
- displayed_packs = [DEFAULT_MODEL_PACK_NAME] + [p for p in available_model_packs if p != DEFAULT_MODEL_PACK_NAME]
- if not displayed_packs:
- displayed_packs = [DEFAULT_MODEL_PACK_NAME]
+ return {"required": {
+ "model_pack_name": (get_chatterbox_model_pack_names(), {
+ "default": DEFAULT_MODEL_PACK_NAME,
+ "tooltip": "Select the Chatterbox voice model pack to use for conversion."
+ }),
+ "source_audio": ("AUDIO", {
+ "tooltip": "The audio containing the speech content to be converted."
+ }),
+ "n_timesteps": ("INT", {
+ "default": 10, "min": 2, "max": 50, "step": 1,
+ "tooltip": "Number of diffusion steps for the flow matching process. Higher values may improve quality at the cost of speed."
+ }),
+ "temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05,
+ "tooltip": "Controls the randomness of the initial noise. 1.0 is standard. Lower values are more deterministic."}),
+ "flow_cfg_scale": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.05,
+ "tooltip": "CFG scale for the mel spectrogram decoder. Higher values increase adherence to the target voice but may reduce naturalness."}),
+ "use_watermark": ("BOOLEAN", {
+ "default": False,
+ "tooltip": "Enable or disable the audio watermark. Requires 'resemble-perth' to be installed."
+ }),
+ },
+ "optional": {"target_voice_audio": ("AUDIO", {
+ "tooltip": "The audio file containing the target voice timbre. If not provided, the default voice from the model pack will be used."
+ }), }}
- return {
- "required": {
- "model_pack_name": (displayed_packs, {"default": DEFAULT_MODEL_PACK_NAME if DEFAULT_MODEL_PACK_NAME in displayed_packs else (displayed_packs[0] if displayed_packs else DEFAULT_MODEL_PACK_NAME)}),
- "source_audio": ("AUDIO",),
- "device": (["cuda", "cpu"], {"default": "cuda" if torch.cuda.is_available() else "cpu"}),
- },
- "optional": {
- "target_voice_audio": ("AUDIO",), # Optional: if not provided, uses default voice from conds.pt
- }
- }
-
- RETURN_TYPES = ("AUDIO",)
- RETURN_NAMES = ("converted_audio",)
- FUNCTION = "convert_voice"
- CATEGORY = "audio/generation"
- OUTPUT_NODE = True
+ RETURN_TYPES = ("AUDIO",); RETURN_NAMES = ("converted_audio",); FUNCTION = "convert_voice"; CATEGORY = "audio/generation"; OUTPUT_NODE = True
def _save_audio_to_temp_file(self, audio_data, prefix=""):
- """Helper to save ComfyUI AUDIO dict to a temporary WAV file."""
- if audio_data is None or \
- audio_data.get("waveform") is None or \
- audio_data["waveform"].numel() == 0:
- return None
+ if audio_data is None or audio_data.get("waveform") is None or audio_data["waveform"].numel() == 0: return None
+ with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_wav:
+ try:
+ waveform_in = audio_data["waveform"]; sample_rate_in = audio_data["sample_rate"]
+ waveform_cpu = waveform_in.cpu()[0]
+ current_waveform = torch.mean(waveform_cpu, dim=0) if waveform_cpu.shape[0] > 1 else waveform_cpu.squeeze(0)
+ sf.write(tmp_wav.name, current_waveform.numpy().astype(np.float32), sample_rate_in)
+ return tmp_wav.name
+ except Exception as e:
+ logger.error(f"Error writing temp {prefix}audio file: {e}", exc_info=True)
+ return None
- waveform_in = audio_data["waveform"]
- sample_rate_in = audio_data["sample_rate"]
- waveform_cpu = waveform_in.cpu()
-
- if waveform_cpu.shape[0] > 1:
- print(f"ChatterboxVC: {prefix}Audio has batch size {waveform_cpu.shape[0]}, using first item.")
- current_waveform = waveform_cpu[0]
-
- if current_waveform.shape[0] > 1: # If C > 1 (stereo or more)
- print(f"ChatterboxVC: {prefix}Audio has {current_waveform.shape[0]} channels, converting to mono by averaging.")
- current_waveform = torch.mean(current_waveform, dim=0)
- else: # If C == 1
- current_waveform = current_waveform.squeeze(0)
-
- processed_audio = current_waveform.numpy().astype(np.float32)
-
- temp_file_path = None
- try:
- with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_wav:
- temp_file_path = tmp_wav.name
- sf.write(temp_file_path, processed_audio, sample_rate_in)
- #print(f"ChatterboxVC: Saved {prefix}audio to temp file: {temp_file_path}")
- return temp_file_path
- except Exception as e:
- print(f"ChatterboxVC: Error writing temp {prefix}audio file: {e}")
- if temp_file_path and os.path.exists(temp_file_path):
- os.remove(temp_file_path)
- return None
-
- def convert_voice(self, model_pack_name, source_audio, device, target_voice_audio=None):
+ def convert_voice(self, model_pack_name, source_audio, n_timesteps, temperature, flow_cfg_scale, use_watermark, target_voice_audio=None):
if source_audio is None or source_audio.get("waveform") is None or source_audio["waveform"].numel() == 0:
- print("ChatterboxVC: No source audio provided, returning silent audio.")
- dummy_sr = 24000
- silent_waveform = torch.zeros((1, dummy_sr), dtype=torch.float32, device="cpu")
+ logger.warning("No source audio provided, returning silent audio.")
+ dummy_sr = 24000; silent_waveform = torch.zeros((1, dummy_sr), dtype=torch.float32, device="cpu")
return ({"waveform": silent_waveform.unsqueeze(0), "sample_rate": dummy_sr},)
- try:
- vc_model = get_cached_chatterbox_vc_model(model_pack_name, device_str=device)
- except Exception as e:
- print(f"ChatterboxVC: Error loading/downloading VC model pack '{model_pack_name}': {e}")
- dummy_sr = 24000
- silent_waveform = torch.zeros((1, dummy_sr), dtype=torch.float32, device="cpu")
+ cache_key = model_pack_name
+ if cache_key not in CHATTERBOX_PATCHER_CACHE:
+ load_device = mm.get_torch_device()
+ logger.info(f"Creating Chatterbox ModelPatcher for {model_pack_name} on device {load_device}")
+ model_wrapper = ChatterboxModelWrapper(model_pack_name)
+ patcher = ChatterboxPatcher(
+ model=model_wrapper,
+ load_device=load_device,
+ offload_device=mm.unet_offload_device(),
+ size=int(1.5 * 1024**3)
+ )
+ CHATTERBOX_PATCHER_CACHE[cache_key] = patcher
+
+ patcher = CHATTERBOX_PATCHER_CACHE[cache_key]
+
+ mm.load_model_gpu(patcher)
+ vc_model = patcher.model.vc_model
+
+ if vc_model is None:
+ logger.error("VC model failed to load. Please check logs for download or loading errors.")
+ dummy_sr = 24000; silent_waveform = torch.zeros((1, dummy_sr), dtype=torch.float32, device="cpu")
return ({"waveform": silent_waveform.unsqueeze(0), "sample_rate": dummy_sr},)
- source_audio_path_temp = None
- target_voice_path_temp = None
+ is_perth_installed = not getattr(perth, '_is_mock', False)
+ if use_watermark and not is_perth_installed:
+ logger.warning("Watermarking is enabled, but 'resemble-perth' is not installed. Output will not be watermarked.")
+
+ original_watermarker = vc_model.watermarker
+ if not use_watermark:
+ class TmpDummyWatermarker:
+ def apply_watermark(self, wav, sample_rate): return wav
+ vc_model.watermarker = TmpDummyWatermarker()
+ if is_perth_installed: logger.info("Watermarking disabled by user.")
+
+ source_audio_path_temp = None; target_voice_path_temp = None
+
+ pbar = ProgressBar(n_timesteps)
try:
source_audio_path_temp = self._save_audio_to_temp_file(source_audio, prefix="Source ")
- if not source_audio_path_temp:
- raise ValueError("Failed to process source audio.")
-
- if target_voice_audio is not None and \
- target_voice_audio.get("waveform") is not None and \
- target_voice_audio["waveform"].numel() > 0:
+ if not source_audio_path_temp: raise ValueError("Failed to process source audio.")
+ if target_voice_audio and target_voice_audio.get("waveform") is not None and target_voice_audio["waveform"].numel() > 0:
target_voice_path_temp = self._save_audio_to_temp_file(target_voice_audio, prefix="Target ")
- #print(f"ChatterboxVC: Using target voice from temp file: {target_voice_path_temp}")
- else:
- print("ChatterboxVC: No target voice audio provided or it's empty. Using default reference from model pack if available.")
- # ChatterboxVC.generate expects file paths
converted_wav_tensor = vc_model.generate(
audio=source_audio_path_temp,
- target_voice_path=target_voice_path_temp # This will be None if no target_voice_audio was provided or saving it failed
- ) # Expected output: (1, num_samples)
-
+ target_voice_path=target_voice_path_temp,
+ n_timesteps=n_timesteps,
+ pbar=pbar,
+ temperature=temperature,
+ flow_cfg_scale=flow_cfg_scale
+ )
except Exception as e:
- print(f"ChatterboxVC: Error during voice conversion: {e}")
- dummy_sr = 24000
- silent_waveform = torch.zeros((1, dummy_sr), dtype=torch.float32, device="cpu")
+ logger.error(f"Error during voice conversion: {e}", exc_info=True)
+ dummy_sr = 24000; silent_waveform = torch.zeros((1, dummy_sr), dtype=torch.float32, device="cpu")
return ({"waveform": silent_waveform.unsqueeze(0), "sample_rate": dummy_sr},)
finally:
- if source_audio_path_temp and os.path.exists(source_audio_path_temp):
- try:
- os.remove(source_audio_path_temp)
- except Exception as e:
- print(f"ChatterboxVC: Error removing temp source audio file {source_audio_path_temp}: {e}")
- if target_voice_path_temp and os.path.exists(target_voice_path_temp):
- try:
- os.remove(target_voice_path_temp)
- except Exception as e:
- print(f"ChatterboxVC: Error removing temp target audio file {target_voice_path_temp}: {e}")
+ vc_model.watermarker = original_watermarker
+ if source_audio_path_temp and os.path.exists(source_audio_path_temp): os.remove(source_audio_path_temp)
+ if target_voice_path_temp and os.path.exists(target_voice_path_temp): os.remove(target_voice_path_temp)
- # ComfyUI AUDIO format: {"waveform": tensor (B, C, T), "sample_rate": int}
- # ChatterboxVC output: tensor (1, T)
- vc_wav_tensor_comfy = converted_wav_tensor.cpu().unsqueeze(0) # (1, T) -> (1, 1, T)
- return ({"waveform": vc_wav_tensor_comfy, "sample_rate": vc_model.sr},)
\ No newline at end of file
+ vc_wav_tensor_comfy = converted_wav_tensor.cpu().unsqueeze(0)
+ return ({"waveform": vc_wav_tensor_comfy, "sample_rate": vc_model.sr},)
diff --git a/src/chatterbox/__init__.py b/src/chatterbox/__init__.py
index 20b4391..6ae2ab1 100644
--- a/src/chatterbox/__init__.py
+++ b/src/chatterbox/__init__.py
@@ -1,2 +1,9 @@
+try:
+ from importlib.metadata import version
+except ImportError:
+ from importlib_metadata import version # For Python <3.8
+
+# __version__ = version("chatterbox-tts")
+
from .tts import ChatterboxTTS
from .vc import ChatterboxVC
diff --git a/src/chatterbox/models/__init__.py b/src/chatterbox/models/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/chatterbox/models/s3gen/configs.py b/src/chatterbox/models/s3gen/configs.py
new file mode 100644
index 0000000..b09b2e5
--- /dev/null
+++ b/src/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"
+})
diff --git a/src/chatterbox/models/s3gen/flow.py b/src/chatterbox/models/s3gen/flow.py
index fee2ec9..41b90f5 100644
--- a/src/chatterbox/models/s3gen/flow.py
+++ b/src/chatterbox/models/s3gen/flow.py
@@ -14,32 +14,54 @@
import logging
import random
from typing import Dict, Optional
+
import torch
import torch.nn as nn
from torch.nn import functional as F
-from omegaconf import DictConfig
from .utils.mask import make_pad_mask
+from .configs import CFM_PARAMS
+logger = logging.getLogger(__name__)
class MaskedDiffWithXvec(torch.nn.Module):
- def __init__(self,
- input_size: int = 512,
- output_size: int = 80,
- spk_embed_dim: int = 192,
- output_type: str = "mel",
- vocab_size: int = 4096,
- input_frame_rate: int = 50,
- only_mask_loss: bool = True,
- encoder: torch.nn.Module = None,
- length_regulator: torch.nn.Module = None,
- decoder: torch.nn.Module = None,
- decoder_conf: Dict = {'in_channels': 240, 'out_channel': 80, 'spk_emb_dim': 80, 'n_spks': 1,
- 'cfm_params': DictConfig({'sigma_min': 1e-06, 'solver': 'euler', 't_scheduler': 'cosine',
- 'training_cfg_rate': 0.2, 'inference_cfg_rate': 0.7, 'reg_loss_type': 'l1'}),
- 'decoder_params': {'channels': [256, 256], 'dropout': 0.0, 'attention_head_dim': 64,
- 'n_blocks': 4, 'num_mid_blocks': 12, 'num_heads': 8, 'act_fn': 'gelu'}},
- mel_feat_conf: Dict = {'n_fft': 1024, 'num_mels': 80, 'sampling_rate': 22050,
- 'hop_size': 256, 'win_size': 1024, 'fmin': 0, 'fmax': 8000}):
+ def __init__(
+ self,
+ input_size: int = 512,
+ output_size: int = 80,
+ spk_embed_dim: int = 192,
+ output_type: str = "mel",
+ vocab_size: int = 4096,
+ input_frame_rate: int = 50,
+ only_mask_loss: bool = True,
+ encoder: torch.nn.Module = None,
+ length_regulator: torch.nn.Module = None,
+ decoder: torch.nn.Module = None,
+ decoder_conf: Dict = {
+ 'in_channels': 240,
+ 'out_channel': 80,
+ 'spk_emb_dim': 80,
+ 'n_spks': 1,
+ 'cfm_params': CFM_PARAMS,
+ 'decoder_params': {
+ 'channels': [256, 256],
+ 'dropout': 0.0,
+ 'attention_head_dim': 64,
+ 'n_blocks': 4,
+ 'num_mid_blocks': 12,
+ 'num_heads': 8,
+ 'act_fn': 'gelu',
+ }
+ },
+ mel_feat_conf: Dict = {
+ 'n_fft': 1024,
+ 'num_mels': 80,
+ 'sampling_rate': 22050,
+ 'hop_size': 256,
+ 'win_size': 1024,
+ 'fmin': 0,
+ 'fmax': 8000
+ }
+ ):
super().__init__()
self.input_size = input_size
self.output_size = output_size
@@ -48,7 +70,7 @@ class MaskedDiffWithXvec(torch.nn.Module):
self.vocab_size = vocab_size
self.output_type = output_type
self.input_frame_rate = input_frame_rate
- logging.info(f"input frame rate={self.input_frame_rate}")
+ # logger.info(f"input frame rate={self.input_frame_rate}")
self.input_embedding = nn.Embedding(vocab_size, input_size)
self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size)
self.encoder = encoder
@@ -110,8 +132,13 @@ class MaskedDiffWithXvec(torch.nn.Module):
prompt_feat,
prompt_feat_len,
embedding,
- flow_cache):
- if self.fp16 is True:
+ flow_cache,
+ n_timesteps=10,
+ pbar=None,
+ temperature=1.0,
+ flow_cfg_scale=0.7
+ ):
+ if hasattr(self, 'fp16') and self.fp16 is True:
prompt_feat = prompt_feat.half()
embedding = embedding.half()
@@ -143,9 +170,12 @@ class MaskedDiffWithXvec(torch.nn.Module):
mask=mask.unsqueeze(1),
spks=embedding,
cond=conds,
- n_timesteps=10,
+ n_timesteps=n_timesteps,
prompt_len=mel_len1,
- flow_cache=flow_cache
+ flow_cache=flow_cache,
+ pbar=pbar,
+ temperature=temperature,
+ flow_cfg_scale=flow_cfg_scale
)
feat = feat[:, :, mel_len1:]
assert feat.shape[2] == mel_len2
@@ -153,25 +183,45 @@ class MaskedDiffWithXvec(torch.nn.Module):
class CausalMaskedDiffWithXvec(torch.nn.Module):
- def __init__(self,
- input_size: int = 512,
- output_size: int = 80,
- spk_embed_dim: int = 192,
- output_type: str = "mel",
- vocab_size: int = 6561,
- input_frame_rate: int = 25,
- only_mask_loss: bool = True,
- token_mel_ratio: int = 2,
- pre_lookahead_len: int = 3,
- encoder: torch.nn.Module = None,
- decoder: torch.nn.Module = None,
- decoder_conf: Dict = {'in_channels': 240, 'out_channel': 80, 'spk_emb_dim': 80, 'n_spks': 1,
- 'cfm_params': DictConfig({'sigma_min': 1e-06, 'solver': 'euler', 't_scheduler': 'cosine',
- 'training_cfg_rate': 0.2, 'inference_cfg_rate': 0.7, 'reg_loss_type': 'l1'}),
- 'decoder_params': {'channels': [256, 256], 'dropout': 0.0, 'attention_head_dim': 64,
- 'n_blocks': 4, 'num_mid_blocks': 12, 'num_heads': 8, 'act_fn': 'gelu'}},
- mel_feat_conf: Dict = {'n_fft': 1024, 'num_mels': 80, 'sampling_rate': 22050,
- 'hop_size': 256, 'win_size': 1024, 'fmin': 0, 'fmax': 8000}):
+ def __init__(
+ self,
+ input_size: int = 512,
+ output_size: int = 80,
+ spk_embed_dim: int = 192,
+ output_type: str = "mel",
+ vocab_size: int = 6561,
+ input_frame_rate: int = 25,
+ only_mask_loss: bool = True,
+ token_mel_ratio: int = 2,
+ pre_lookahead_len: int = 3,
+ encoder: torch.nn.Module = None,
+ decoder: torch.nn.Module = None,
+ decoder_conf: Dict = {
+ 'in_channels': 240,
+ 'out_channel': 80,
+ 'spk_emb_dim': 80,
+ 'n_spks': 1,
+ 'cfm_params': CFM_PARAMS,
+ 'decoder_params': {
+ 'channels': [256, 256],
+ 'dropout': 0.0,
+ 'attention_head_dim': 64,
+ 'n_blocks': 4,
+ 'num_mid_blocks': 12,
+ 'num_heads': 8,
+ 'act_fn': 'gelu',
+ }
+ },
+ mel_feat_conf: Dict = {
+ 'n_fft': 1024,
+ 'num_mels': 80,
+ 'sampling_rate': 22050,
+ 'hop_size': 256,
+ 'win_size': 1024,
+ 'fmin': 0,
+ 'fmax': 8000
+ }
+ ):
super().__init__()
self.input_size = input_size
self.output_size = output_size
@@ -180,7 +230,7 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
self.vocab_size = vocab_size
self.output_type = output_type
self.input_frame_rate = input_frame_rate
- logging.info(f"input frame rate={self.input_frame_rate}")
+ # logger.info(f"input frame rate={self.input_frame_rate}")
self.input_embedding = nn.Embedding(vocab_size, input_size)
self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size)
self.encoder = encoder
@@ -202,8 +252,13 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
prompt_feat,
prompt_feat_len,
embedding,
- finalize):
- if self.fp16 is True:
+ finalize,
+ n_timesteps=10,
+ pbar=None,
+ temperature=1.0,
+ flow_cfg_scale=0.7
+ ):
+ if hasattr(self, 'fp16') and self.fp16 is True:
prompt_feat = prompt_feat.half()
embedding = embedding.half()
@@ -235,7 +290,10 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
mask=mask.unsqueeze(1),
spks=embedding,
cond=conds,
- n_timesteps=10
+ n_timesteps=n_timesteps,
+ pbar=pbar,
+ temperature=temperature,
+ flow_cfg_scale=flow_cfg_scale,
)
feat = feat[:, :, mel_len1:]
assert feat.shape[2] == mel_len2
diff --git a/src/chatterbox/models/s3gen/flow_matching.py b/src/chatterbox/models/s3gen/flow_matching.py
index 74fc66f..458c17b 100644
--- a/src/chatterbox/models/s3gen/flow_matching.py
+++ b/src/chatterbox/models/s3gen/flow_matching.py
@@ -14,18 +14,9 @@
import threading
import torch
import torch.nn.functional as F
+from tqdm import tqdm
from .matcha.flow_matching import BASECFM
-from omegaconf import OmegaConf
-
-
-CFM_PARAMS = OmegaConf.create({
- "sigma_min": 1e-06,
- "solver": "euler",
- "t_scheduler": "cosine",
- "training_cfg_rate": 0.2,
- "inference_cfg_rate": 0.7,
- "reg_loss_type": "l1"
-})
+from .configs import CFM_PARAMS
class ConditionalCFM(BASECFM):
@@ -45,7 +36,7 @@ class ConditionalCFM(BASECFM):
self.lock = threading.Lock()
@torch.inference_mode()
- def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, prompt_len=0, flow_cache=torch.zeros(1, 80, 0, 2)):
+ def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, prompt_len=0, flow_cache=torch.zeros(1, 80, 0, 2), pbar=None, flow_cfg_scale=None):
"""Forward diffusion
Args:
@@ -58,11 +49,15 @@ class ConditionalCFM(BASECFM):
spks (torch.Tensor, optional): speaker ids. Defaults to None.
shape: (batch_size, spk_emb_dim)
cond: Not used but kept for future purposes
+ pbar: ComfyUI ProgressBar instance
Returns:
sample: generated mel-spectrogram
shape: (batch_size, n_feats, mel_timesteps)
"""
+
+ if flow_cfg_scale is not None:
+ self.inference_cfg_rate = flow_cfg_scale
z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature
cache_size = flow_cache.shape[2]
@@ -77,9 +72,10 @@ class ConditionalCFM(BASECFM):
t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
if self.t_scheduler == 'cosine':
t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
- return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), flow_cache
+
+ return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond, pbar=pbar), flow_cache
- def solve_euler(self, x, t_span, mu, mask, spks, cond):
+ def solve_euler(self, x, t_span, mu, mask, spks, cond, pbar=None):
"""
Fixed euler solver for ODEs.
Args:
@@ -93,6 +89,7 @@ class ConditionalCFM(BASECFM):
spks (torch.Tensor, optional): speaker ids. Defaults to None.
shape: (batch_size, spk_emb_dim)
cond: Not used but kept for future purposes
+ pbar: ComfyUI ProgressBar instance
"""
t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
t = t.unsqueeze(dim=0)
@@ -100,7 +97,8 @@ class ConditionalCFM(BASECFM):
# I am storing this because I can later plot it by putting a debugger here and saving it to a file
# Or in future might add like a return_all_steps flag
sol = []
-
+ iterator = tqdm(range(1, len(t_span)), desc="Voice Converting", dynamic_ncols=True)
+
# Do not use concat, it may cause memory format changed and trt infer with wrong results!
x_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=x.dtype)
mask_in = torch.zeros([2, 1, x.size(2)], device=x.device, dtype=x.dtype)
@@ -108,7 +106,7 @@ class ConditionalCFM(BASECFM):
t_in = torch.zeros([2], device=x.device, dtype=x.dtype)
spks_in = torch.zeros([2, 80], device=x.device, dtype=x.dtype)
cond_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=x.dtype)
- for step in range(1, len(t_span)):
+ for step in iterator:
# Classifier-Free Guidance inference introduced in VoiceBox
x_in[:] = x
mask_in[:] = mask
@@ -129,6 +127,9 @@ class ConditionalCFM(BASECFM):
sol.append(x)
if step < len(t_span) - 1:
dt = t_span[step + 1] - t
+
+ if pbar:
+ pbar.update(1)
return sol[-1].float()
@@ -201,7 +202,7 @@ class CausalConditionalCFM(ConditionalCFM):
self.rand_noise = torch.randn([1, 80, 50 * 300])
@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, pbar=None, flow_cfg_scale=None, **kwargs):
"""Forward diffusion
Args:
@@ -220,9 +221,12 @@ class CausalConditionalCFM(ConditionalCFM):
shape: (batch_size, n_feats, mel_timesteps)
"""
+ if flow_cfg_scale is not None:
+ self.inference_cfg_rate = flow_cfg_scale
+
z = self.rand_noise[:, :, :mu.size(2)].to(mu.device).to(mu.dtype) * temperature
# fix prompt and overlap part mu and z
t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
if self.t_scheduler == 'cosine':
t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
- return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None
+ return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond, pbar=pbar), None
diff --git a/src/chatterbox/models/s3gen/matcha/transformer.py b/src/chatterbox/models/s3gen/matcha/transformer.py
index 4f6762b..5517e1f 100644
--- a/src/chatterbox/models/s3gen/matcha/transformer.py
+++ b/src/chatterbox/models/s3gen/matcha/transformer.py
@@ -2,15 +2,21 @@ from typing import Any, Dict, Optional
import torch
import torch.nn as nn
-from diffusers.models.attention import (
+from diffusers.models.activations import (
GEGLU,
GELU,
- AdaLayerNorm,
- AdaLayerNormZero,
ApproximateGELU,
)
+from diffusers.models.normalization import (
+ AdaLayerNorm,
+ AdaLayerNormZero,
+)
+
from diffusers.models.attention_processor import Attention
-from diffusers.models.lora import LoRACompatibleLinear
+
+# deprecated LoRACompatibleLinear
+# from diffusers.models.lora import LoRACompatibleLinear
+
from diffusers.utils.torch_utils import maybe_allow_in_graph
@@ -45,7 +51,9 @@ class SnakeBeta(nn.Module):
"""
super().__init__()
self.in_features = out_features if isinstance(out_features, list) else [out_features]
- self.proj = LoRACompatibleLinear(in_features, out_features)
+
+ # Switched to standard torch.nn.Linear
+ self.proj = nn.Linear(in_features, out_features)
# initialize alpha
self.alpha_logscale = alpha_logscale
@@ -123,7 +131,9 @@ class FeedForward(nn.Module):
# project dropout
self.net.append(nn.Dropout(dropout))
# project out
- self.net.append(LoRACompatibleLinear(inner_dim, dim_out))
+
+ # Switched to standard torch.nn.Linear
+ self.net.append(nn.Linear(inner_dim, dim_out))
# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
if final_dropout:
self.net.append(nn.Dropout(dropout))
diff --git a/src/chatterbox/models/s3gen/s3gen.py b/src/chatterbox/models/s3gen/s3gen.py
index c611344..fb352e9 100644
--- a/src/chatterbox/models/s3gen/s3gen.py
+++ b/src/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,7 +30,9 @@ from .hifigan import HiFTGenerator
from .transformer.upsample_encoder import UpsampleConformerEncoder
from .flow_matching import CausalConditionalCFM
from .decoder import ConditionalDecoder
+from .configs import CFM_PARAMS
+logger = logging.getLogger(__name__)
def drop_invalid_tokens(x):
assert len(x.shape) <= 2 and x.shape[0] == 1, "only batch size of one allowed for now"
@@ -85,14 +86,7 @@ class S3Token2Mel(torch.nn.Module):
num_heads=8,
act_fn='gelu',
)
- cfm_params = DictConfig({
- "sigma_min": 1e-06,
- "solver": 'euler',
- "t_scheduler": 'cosine',
- "training_cfg_rate": 0.2,
- "inference_cfg_rate": 0.7,
- "reg_loss_type": 'l1',
- })
+ cfm_params = CFM_PARAMS
decoder = CausalConditionalCFM(
spk_emb_dim=80,
cfm_params=cfm_params,
@@ -129,7 +123,7 @@ class S3Token2Mel(torch.nn.Module):
ref_wav = ref_wav.unsqueeze(0) # (B, L)
if ref_wav.size(1) > 10 * ref_sr:
- print("WARNING: cosydec received ref longer than 10s")
+ logger.warning("WARNING: cosydec received ref longer than 10s")
ref_wav_24 = ref_wav
if ref_sr != S3GEN_SR:
@@ -149,7 +143,7 @@ class S3Token2Mel(torch.nn.Module):
# Make sure mel_len = 2 * stoken_len (happens when the input is not padded to multiple of 40ms)
if ref_mels_24.shape[1] != 2 * ref_speech_tokens.shape[1]:
- logging.warning(
+ logger.warning(
"Reference mel length is not equal to 2 * reference token length.\n"
)
ref_speech_tokens = ref_speech_tokens[:, :ref_mels_24.shape[1] // 2]
@@ -172,6 +166,8 @@ class S3Token2Mel(torch.nn.Module):
# pre-computed ref embedding (prod API)
ref_dict: Optional[dict] = None,
finalize: bool = False,
+ n_timesteps: int = 10,
+ pbar=None
):
"""
Generate waveforms from S3 speech tokens and a reference waveform, which the speaker timbre is inferred from.
@@ -211,11 +207,15 @@ class S3Token2Mel(torch.nn.Module):
token=speech_tokens,
token_len=speech_token_lens,
finalize=finalize,
+ n_timesteps=n_timesteps,
+ pbar=pbar,
**ref_dict,
)
return output_mels
+
+
class S3Token2Wav(S3Token2Mel):
"""
The decoder of CosyVoice2 is a concat of token-to-mel (CFM) and a mel-to-waveform (HiFiGAN) modules.
@@ -250,9 +250,11 @@ class S3Token2Wav(S3Token2Mel):
ref_sr: Optional[int],
# pre-computed ref embedding (prod API)
ref_dict: Optional[dict] = None,
- finalize: bool = False
+ finalize: bool = False,
+ n_timesteps: int = 10,
+ pbar=None
):
- output_mels = super().forward(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize)
+ output_mels = super().forward(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize, n_timesteps=n_timesteps, pbar=pbar)
# 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)
@@ -275,8 +277,41 @@ class S3Token2Wav(S3Token2Mel):
# pre-computed ref embedding (prod API)
ref_dict: Optional[dict] = None,
finalize: bool = False,
+ n_timesteps: int = 10,
+ pbar=None,
+ temperature: float = 1.0,
+ flow_cfg_scale: float = 0.7
):
- return super().forward(speech_tokens, ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict, finalize=finalize)
+ # This method in the base class now needs to accept and pass the new params
+ # (The base class `S3Token2Mel` needs this change)
+
+ assert (ref_wav is None) ^ (ref_dict is None), f"Must provide exactly one of ref_wav or ref_dict (got {ref_wav} and {ref_dict})"
+
+ if ref_dict is None:
+ ref_dict = self.embed_ref(ref_wav, ref_sr)
+ else:
+ for rk in list(ref_dict):
+ if isinstance(ref_dict[rk], np.ndarray):
+ ref_dict[rk] = torch.from_numpy(ref_dict[rk])
+ if torch.is_tensor(ref_dict[rk]):
+ ref_dict[rk] = ref_dict[rk].to(self.device)
+
+ if len(speech_tokens.shape) == 1:
+ speech_tokens = speech_tokens.unsqueeze(0)
+
+ speech_token_lens = torch.LongTensor([speech_tokens.size(1)]).to(self.device)
+
+ output_mels, _ = self.flow.inference(
+ token=speech_tokens,
+ token_len=speech_token_lens,
+ finalize=finalize,
+ n_timesteps=n_timesteps,
+ pbar=pbar,
+ temperature=temperature,
+ flow_cfg_scale=flow_cfg_scale,
+ **ref_dict,
+ )
+ return output_mels
@torch.inference_mode()
def hift_inference(self, speech_feat, cache_source: torch.Tensor = None):
@@ -295,8 +330,17 @@ class S3Token2Wav(S3Token2Mel):
ref_dict: Optional[dict] = None,
cache_source: torch.Tensor = None, # NOTE: this arg is for streaming, it can probably be removed here
finalize: bool = True,
+ n_timesteps: int = 10,
+ pbar=None,
+ temperature: float = 1.0,
+ flow_cfg_scale: float = 0.7
):
- 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(
+ speech_tokens,
+ ref_wav=ref_wav, ref_sr=ref_sr, ref_dict=ref_dict,
+ finalize=finalize, n_timesteps=n_timesteps, pbar=pbar,
+ temperature=temperature, flow_cfg_scale=flow_cfg_scale
+ )
output_wavs, output_sources = self.hift_inference(output_mels, cache_source)
# NOTE: ad-hoc method to reduce "spillover" from the reference clip.
diff --git a/src/chatterbox/models/t3/llama_configs.py b/src/chatterbox/models/t3/llama_configs.py
index 6247ead..c8484cd 100644
--- a/src/chatterbox/models/t3/llama_configs.py
+++ b/src/chatterbox/models/t3/llama_configs.py
@@ -1,17 +1,19 @@
-# FILE: src/chatterbox/models/t3/llama_configs.py
LLAMA_520M_CONFIG_DICT = dict(
# Arbitrary small number that won't cause problems when loading.
# These param are unused due to custom input layers.
vocab_size=8,
- # This defines the maximum sequence length the RoPE mechanism is initially configured for
- # if rope_scaling is None. 8192 is a common value for Llama models.
- max_position_embeddings=8192,
+ # default params needed for loading most pretrained 1B weights
+ max_position_embeddings=131072,
hidden_size=1024,
intermediate_size=4096,
num_hidden_layers=30,
num_attention_heads=16,
- attn_implementation="sdpa", # Use "eager" if "sdpa" causes issues with older torch/transformers
- head_dim=64, # hidden_size // num_attention_heads
+ # This is required because AlignmentStreamAnalyzer needs to inspect attention weights,
+ # which is not supported by optimized backends like FlashAttention (SDPA).
+ # This informs transformers that the fallback to the slower path is intentional, silencing the warning.
+ attn_implementation="eager",
+ #attn_implementation="sdpa",
+ head_dim=64,
tie_word_embeddings=False,
hidden_act="silu",
attention_bias=False,
@@ -19,15 +21,21 @@ LLAMA_520M_CONFIG_DICT = dict(
initializer_range=0.02,
mlp_bias=False,
model_type="llama",
- num_key_value_heads=16, # For Llama GQA/MQA. For standard MHA, num_key_value_heads = num_attention_heads
+ num_key_value_heads=16,
pretraining_tp=1,
rms_norm_eps=1e-05,
- rope_scaling=None, # MODIFICATION: Explicitly set to None
- rope_theta=500000.0, # This is crucial for Llama 3 style RoPE
- torch_dtype="bfloat16", # Consider "float16" or "float32" if bf16 is not supported or causes issues
+ rope_scaling=dict(
+ factor=8.0,
+ high_freq_factor=4.0,
+ low_freq_factor=1.0,
+ original_max_position_embeddings=8192,
+ rope_type="llama3"
+ ),
+ rope_theta=500000.0,
+ torch_dtype="bfloat16",
use_cache=True,
)
LLAMA_CONFIGS = {
"Llama_520M": LLAMA_520M_CONFIG_DICT,
-}
\ No newline at end of file
+}
diff --git a/src/chatterbox/models/t3/modules/perceiver.py b/src/chatterbox/models/t3/modules/perceiver.py
index 8bed6d6..a47e813 100644
--- a/src/chatterbox/models/t3/modules/perceiver.py
+++ b/src/chatterbox/models/t3/modules/perceiver.py
@@ -7,7 +7,15 @@ import torch
from torch import nn
import torch.nn.functional as F
from einops import rearrange
-from torch.nn.attention import SDPBackend, sdpa_kernel
+
+try:
+ from torch.nn.attention import SDPBackend, sdpa_kernel
+ # Flag to indicate that the modern SDPA API is available.
+ SUPPORTS_SDPA_KERNEL = True
+except (ImportError, AttributeError):
+ # Fallback for older PyTorch versions.
+ SUPPORTS_SDPA_KERNEL = False
+
class RelativePositionBias(nn.Module):
def __init__(self, scale, causal=False, num_buckets=32, max_distance=128, heads=8):
@@ -66,6 +74,7 @@ class AttentionQKV(nn.Module):
def setup_flash_config(self):
# Setup flash attention configuration
flash_config = {
+ 'enable_flash': True,
'enable_math': True,
'enable_mem_efficient': True
}
@@ -89,14 +98,32 @@ class AttentionQKV(nn.Module):
return torch.einsum("bhts,bhls->bhlt", attn, v)
def flash_attention(self, q, k, v, mask=None):
- config = self.flash_config if self.flash_config else {}
- with sdpa_kernel(backends=[SDPBackend.EFFICIENT_ATTENTION]):
- out = F.scaled_dot_product_attention(
- q, k, v,
- attn_mask=mask,
- dropout_p=self.dropout_rate if self.training else 0.
- )
+ if SUPPORTS_SDPA_KERNEL:
+ # Modern PyTorch (>= 2.2) API for enabling specific backends.
+ # This replaces the deprecated `torch.backends.cuda.sdp_kernel`.
+ backends = []
+ if self.flash_config.get('enable_flash', True):
+ backends.append(SDPBackend.FLASH_ATTENTION)
+ if self.flash_config.get('enable_mem_efficient', True):
+ backends.append(SDPBackend.EFFICIENT_ATTENTION)
+ if self.flash_config.get('enable_math', True):
+ backends.append(SDPBackend.MATH)
+ with sdpa_kernel(backends=backends):
+ out = F.scaled_dot_product_attention(
+ q, k, v,
+ attn_mask=mask,
+ dropout_p=self.dropout_rate if self.training else 0.
+ )
+ else:
+ # Fallback for older PyTorch versions.
+ config = self.flash_config if self.flash_config else {}
+ with torch.backends.cuda.sdp_kernel(**config):
+ out = F.scaled_dot_product_attention(
+ q, k, v,
+ attn_mask=mask,
+ dropout_p=self.dropout_rate if self.training else 0.
+ )
return out
def split_heads(self, x):
diff --git a/src/chatterbox/models/t3/t3.py b/src/chatterbox/models/t3/t3.py
index 32fc6bc..70b21bd 100644
--- a/src/chatterbox/models/t3/t3.py
+++ b/src/chatterbox/models/t3/t3.py
@@ -8,7 +8,7 @@ import torch
import torch.nn.functional as F
from torch import nn, Tensor
from transformers import LlamaModel, LlamaConfig
-from transformers.generation.logits_process import TopPLogitsWarper, RepetitionPenaltyLogitsProcessor
+from transformers.generation.logits_process import MinPLogitsWarper, RepetitionPenaltyLogitsProcessor, TopPLogitsWarper
from .modules.learned_pos_emb import LearnedPositionEmbeddings
@@ -16,18 +16,12 @@ 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
+from ..utils import AttrDict
logger = logging.getLogger(__name__)
-class AttrDict(dict):
- def __init__(self, *args, **kwargs):
- super(AttrDict, self).__init__(*args, **kwargs)
- self.__dict__ = self
-
-
def _ensure_BOT_EOT(text_tokens: Tensor, hp):
B = text_tokens.size(0)
assert (text_tokens == hp.start_text_token).int().sum() >= B, "missing start_text_token"
@@ -89,11 +83,13 @@ class T3(nn.Module):
t3_cond: T3Cond,
text_tokens: torch.LongTensor,
speech_tokens: torch.LongTensor,
+ cfg_weight: float = 0.0,
):
# prepare input embeddings (skip backbone tranformer embeddings)
cond_emb = self.prepare_conditioning(t3_cond) # (B, len_cond, dim)
text_emb = self.text_emb(text_tokens) # (B, len_text, dim)
- text_emb[1].zero_() # CFG uncond
+ if cfg_weight > 0.0:
+ text_emb[1].zero_() # CFG uncond
speech_emb = self.speech_emb(speech_tokens) # (B, len_speech, dim)
if self.hp.input_pos_emb == "learned":
@@ -217,14 +213,16 @@ class T3(nn.Module):
# HF generate args
num_return_sequences=1,
- max_new_tokens=None,
+ max_new_tokens,
stop_on_eos=True,
do_sample=True,
temperature=0.8,
- top_p=0.8,
+ min_p=0.05,
+ top_p=1.0,
length_penalty=1.0,
- repetition_penalty=2.0,
+ repetition_penalty=1.2,
cfg_weight=0,
+ pbar=None
):
"""
Args:
@@ -235,6 +233,9 @@ class T3(nn.Module):
_ensure_BOT_EOT(text_tokens, self.hp)
text_tokens = torch.atleast_2d(text_tokens).to(dtype=torch.long, device=self.device)
+ if max_new_tokens is None:
+ max_new_tokens = self.hp.max_speech_tokens
+
# Default initial speech to a single start-of-speech token
if initial_speech_tokens is None:
initial_speech_tokens = self.hp.start_speech_token * torch.ones_like(text_tokens[:, :1])
@@ -244,6 +245,7 @@ class T3(nn.Module):
t3_cond=t3_cond,
text_tokens=text_tokens,
speech_tokens=initial_speech_tokens,
+ cfg_weight=cfg_weight,
)
# In order to use the standard HF generate method, we need to extend some methods to inject our custom logic
@@ -254,19 +256,12 @@ class T3(nn.Module):
# TODO? synchronize the expensive compile function
# with self.compile_lock:
if not self.compiled:
- alignment_stream_analyzer = AlignmentStreamAnalyzer(
- self.tfmr,
- None,
- text_tokens_slice=(len_cond, len_cond + text_tokens.size(-1)),
- alignment_layer_idx=9, # TODO: hparam or something?
- 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
@@ -281,7 +276,7 @@ class T3(nn.Module):
# max_new_tokens=max_new_tokens or self.hp.max_speech_tokens,
# num_return_sequences=num_return_sequences,
# temperature=temperature,
- # top_p=top_p,
+ # min_p=min_p,
# length_penalty=length_penalty,
# repetition_penalty=repetition_penalty,
# do_sample=do_sample,
@@ -297,18 +292,22 @@ class T3(nn.Module):
# batch_size=2 for CFG
bos_embed = torch.cat([bos_embed, bos_embed])
- # Combine condition and BOS token for the initial input
- inputs_embeds = torch.cat([embeds, bos_embed], dim=1)
+ # Combine condition and BOS token for the initial input if cfg_weight > 0
+ if cfg_weight > 0:
+ inputs_embeds = torch.cat([embeds, bos_embed], dim=1)
+ else:
+ inputs_embeds = embeds
# Track generated token ids; start with the BOS token.
generated_ids = bos_token.clone()
predicted = [] # To store the predicted tokens
# Instantiate the logits processors.
+ min_p_warper = MinPLogitsWarper(min_p=min_p)
top_p_warper = TopPLogitsWarper(top_p=top_p)
- repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(penalty=repetition_penalty)
+ repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(penalty=float(repetition_penalty))
- # ---- Initial Forward Pass (no kv_cache yet) ----
+ # Initial Forward Pass (no kv_cache yet)
output = self.patched_model(
inputs_embeds=inputs_embeds,
past_key_values=None,
@@ -320,14 +319,17 @@ class T3(nn.Module):
# Initialize kv_cache with the full context.
past = output.past_key_values
- # ---- Generation Loop using kv_cache ----
- for i in tqdm(range(max_new_tokens), desc="Sampling", dynamic_ncols=True):
+ iterator = tqdm(range(max_new_tokens), desc="Chatterbox Sampling", dynamic_ncols=True)
+
+ for i in iterator:
logits = output.logits[:, -1, :]
# CFG
- logits_cond = logits[0:1]
- logits_uncond = logits[1:2]
- logits = logits_cond + cfg_weight * (logits_cond - logits_uncond)
+ if cfg_weight > 0.0:
+ logits_cond = logits[0:1]
+ logits_uncond = logits[1:2]
+ logits = logits_cond + cfg_weight * (logits_cond - logits_uncond)
+
logits = logits.squeeze(1)
# Apply temperature scaling.
@@ -336,6 +338,7 @@ class T3(nn.Module):
# Apply repetition penalty and top‑p filtering.
logits = repetition_penalty_processor(generated_ids, logits)
+ logits = min_p_warper(None, logits)
logits = top_p_warper(None, logits)
# Convert logits to probabilities and sample the next token.
@@ -345,8 +348,7 @@ class T3(nn.Module):
predicted.append(next_token)
generated_ids = torch.cat([generated_ids, next_token], dim=1)
- # Check for EOS token.
- if next_token.view(-1) == self.hp.stop_speech_token:
+ if stop_on_eos and next_token.view(-1) == self.hp.stop_speech_token:
break
# Get embedding for the new token.
@@ -354,7 +356,8 @@ class T3(nn.Module):
next_token_embed = next_token_embed + self.speech_pos_emb.get_fixed_embedding(i + 1)
# For CFG
- next_token_embed = torch.cat([next_token_embed, next_token_embed])
+ if cfg_weight > 0.0:
+ next_token_embed = torch.cat([next_token_embed, next_token_embed])
# Forward pass with only the new token and the cached past.
output = self.patched_model(
@@ -367,6 +370,9 @@ class T3(nn.Module):
# Update the kv_cache.
past = output.past_key_values
+ if pbar:
+ pbar.update(1)
+
# Concatenate all predicted tokens along the sequence dimension.
predicted_tokens = torch.cat(predicted, dim=1) # shape: (B, num_tokens)
return predicted_tokens
diff --git a/src/chatterbox/models/utils.py b/src/chatterbox/models/utils.py
new file mode 100644
index 0000000..a4abce5
--- /dev/null
+++ b/src/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
diff --git a/src/chatterbox/models/voice_encoder/voice_encoder.py b/src/chatterbox/models/voice_encoder/voice_encoder.py
index b0ed2df..7fc00ca 100644
--- a/src/chatterbox/models/voice_encoder/voice_encoder.py
+++ b/src/chatterbox/models/voice_encoder/voice_encoder.py
@@ -259,7 +259,7 @@ class VoiceEncoder(nn.Module):
"""
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_best")
for wav in wavs
]
diff --git a/src/chatterbox/tts.py b/src/chatterbox/tts.py
index 3371dae..2adbaca 100644
--- a/src/chatterbox/tts.py
+++ b/src/chatterbox/tts.py
@@ -1,3 +1,4 @@
+import logging
from dataclasses import dataclass
from pathlib import Path
@@ -6,14 +7,16 @@ import torch
import perth
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
from .models.t3 import T3
-from .models.s3tokenizer import S3_SR, drop_invalid_tokens
+from .models.s3tokenizer import S3_SR, S3_TOKEN_RATE, drop_invalid_tokens
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
+logger = logging.getLogger(__name__)
REPO_ID = "ResembleAI/chatterbox"
@@ -96,9 +99,27 @@ class Conditionals:
@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'])
+# helper function for padding
+def _pad_wav_to_40ms_multiple(wav: torch.Tensor, sr: int) -> torch.Tensor:
+ """
+ Pads a waveform to be a multiple of 40ms to prevent rounding errors between
+ the mel spectrogram (20ms hop) and the speech tokenizer (40ms hop).
+ """
+ S3_TOKEN_DURATION_S = 1 / S3_TOKEN_RATE # 0.04 seconds
+ samples_per_token = int(sr * S3_TOKEN_DURATION_S)
+ current_samples = wav.shape[-1]
+ remainder = current_samples % samples_per_token
+ if remainder != 0:
+ padding_needed = samples_per_token - remainder
+ padded_wav = F.pad(wav, (0, padding_needed))
+ return padded_wav
+ return wav
+
class ChatterboxTTS:
ENC_COND_LEN = 6 * S3_SR
@@ -126,14 +147,20 @@ class ChatterboxTTS:
def from_local(cls, ckpt_dir, device) -> 'ChatterboxTTS':
ckpt_dir = Path(ckpt_dir)
+ # Always load to CPU first for non-CUDA devices to handle CUDA-saved models
+ 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")
+ load_file(ckpt_dir / "ve.safetensors")
)
ve.to(device).eval()
t3 = T3()
- t3_state = torch.load(ckpt_dir / "t3_cfg.pt")
+ t3_state = load_file(ckpt_dir / "t3_cfg.safetensors")
if "model" in t3_state.keys():
t3_state = t3_state["model"][0]
t3.load_state_dict(t3_state)
@@ -141,7 +168,7 @@ class ChatterboxTTS:
s3gen = S3Gen()
s3gen.load_state_dict(
- torch.load(ckpt_dir / "s3gen.pt")
+ load_file(ckpt_dir / "s3gen.safetensors"), strict=False
)
s3gen.to(device).eval()
@@ -151,13 +178,21 @@ class ChatterboxTTS:
conds = None
if (builtin_voice := ckpt_dir / "conds.pt").exists():
- conds = Conditionals.load(builtin_voice).to(device)
+ 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) -> 'ChatterboxTTS':
- for fpath in ["ve.pt", "t3_cfg.pt", "s3gen.pt", "tokenizer.json", "conds.pt"]:
+ # Check if MPS is available on macOS
+ if device == "mps" and not torch.backends.mps.is_available():
+ if not torch.backends.mps.is_built():
+ logger.warning("MPS not available because the current PyTorch install was not built with MPS enabled.")
+ else:
+ logger.warning("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"]:
local_path = hf_hub_download(repo_id=REPO_ID, filename=fpath)
return cls.from_local(Path(local_path).parent, device)
@@ -165,11 +200,18 @@ class ChatterboxTTS:
def prepare_conditionals(self, wav_fpath, exaggeration=0.5):
## Load reference wav
s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
+
+ # Convert to tensor and pad to a 40ms boundary
+ s3gen_ref_wav = torch.from_numpy(s3gen_ref_wav).float().unsqueeze(0)
+ s3gen_ref_wav = _pad_wav_to_40ms_multiple(s3gen_ref_wav, S3GEN_SR)
+
+ # Now convert back to numpy for librosa, or use torch audio for resampling
+ s3gen_ref_wav_np = s3gen_ref_wav.squeeze(0).numpy()
- ref_16k_wav = librosa.resample(s3gen_ref_wav, orig_sr=S3GEN_SR, target_sr=S3_SR)
+ ref_16k_wav = librosa.resample(s3gen_ref_wav_np, 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)
+ s3gen_ref_wav_np = s3gen_ref_wav_np[:self.DEC_COND_LEN]
+ s3gen_ref_dict = self.s3gen.embed_ref(s3gen_ref_wav_np, S3GEN_SR, device=self.device)
# Speech cond prompt tokens
if plen := self.t3.hp.speech_cond_prompt_len:
@@ -191,10 +233,16 @@ class ChatterboxTTS:
def generate(
self,
text,
+ repetition_penalty=1.2,
+ min_p=0.05,
+ top_p=1.0,
audio_prompt_path=None,
exaggeration=0.5,
cfg_weight=0.5,
temperature=0.8,
+ pbar=None,
+ max_new_tokens=1000,
+ flow_cfg_scale=0.7
):
if audio_prompt_path:
self.prepare_conditionals(audio_prompt_path, exaggeration=exaggeration)
@@ -213,7 +261,9 @@ class ChatterboxTTS:
# Norm and tokenize text
text = punc_norm(text)
text_tokens = self.tokenizer.text_to_tokens(text).to(self.device)
- text_tokens = torch.cat([text_tokens, text_tokens], dim=0) # Need two seqs for CFG
+
+ if cfg_weight > 0.0:
+ 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
@@ -224,20 +274,28 @@ class ChatterboxTTS:
speech_tokens = self.t3.inference(
t3_cond=self.conds.t3,
text_tokens=text_tokens,
- max_new_tokens=1000, # TODO: use the value in config
+ max_new_tokens=max_new_tokens,
temperature=temperature,
cfg_weight=cfg_weight,
+ repetition_penalty=repetition_penalty,
+ min_p=min_p,
+ top_p=top_p,
+ pbar=pbar
)
# Extract only the conditional batch.
speech_tokens = speech_tokens[0]
# TODO: output becomes 1D
speech_tokens = drop_invalid_tokens(speech_tokens)
+
+ speech_tokens = speech_tokens[speech_tokens < 6561]
+
speech_tokens = speech_tokens.to(self.device)
wav, _ = self.s3gen.inference(
speech_tokens=speech_tokens,
ref_dict=self.conds.gen,
+ flow_cfg_scale=flow_cfg_scale
)
wav = wav.squeeze(0).detach().cpu().numpy()
watermarked_wav = self.watermarker.apply_watermark(wav, sample_rate=self.sr)
diff --git a/src/chatterbox/vc.py b/src/chatterbox/vc.py
index 629b686..ef05b10 100644
--- a/src/chatterbox/vc.py
+++ b/src/chatterbox/vc.py
@@ -1,16 +1,34 @@
+import logging
from pathlib import Path
import librosa
import torch
import perth
+import torch.nn.functional as F
from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
-from .models.s3tokenizer import S3_SR
+from .models.s3tokenizer import S3_SR, S3_TOKEN_RATE
from .models.s3gen import S3GEN_SR, S3Gen
+logger = logging.getLogger(__name__)
REPO_ID = "ResembleAI/chatterbox"
+def _pad_wav_to_40ms_multiple(wav: torch.Tensor, sr: int) -> torch.Tensor:
+ """
+ Pads a waveform to be a multiple of 40ms to prevent rounding errors between
+ the mel spectrogram (20ms hop) and the speech tokenizer (40ms hop).
+ """
+ S3_TOKEN_DURATION_S = 1 / S3_TOKEN_RATE # 0.04 seconds
+ samples_per_token = int(sr * S3_TOKEN_DURATION_S)
+ current_samples = wav.shape[-1]
+ remainder = current_samples % samples_per_token
+ if remainder != 0:
+ padding_needed = samples_per_token - remainder
+ padded_wav = F.pad(wav, (0, padding_needed))
+ return padded_wav
+ return wav
class ChatterboxVC:
ENC_COND_LEN = 6 * S3_SR
@@ -37,14 +55,21 @@ class ChatterboxVC:
@classmethod
def from_local(cls, ckpt_dir, device) -> 'ChatterboxVC':
ckpt_dir = Path(ckpt_dir)
+
+ # Always load to CPU first for non-CUDA devices to handle CUDA-saved models
+ if device in ["cpu", "mps"]:
+ map_location = torch.device('cpu')
+ else:
+ map_location = None
+
ref_dict = None
if (builtin_voice := ckpt_dir / "conds.pt").exists():
- states = torch.load(builtin_voice)
+ states = torch.load(builtin_voice, map_location=map_location)
ref_dict = states['gen']
s3gen = S3Gen()
s3gen.load_state_dict(
- torch.load(ckpt_dir / "s3gen.pt")
+ load_file(ckpt_dir / "s3gen.safetensors"), strict=False
)
s3gen.to(device).eval()
@@ -52,7 +77,15 @@ class ChatterboxVC:
@classmethod
def from_pretrained(cls, device) -> 'ChatterboxVC':
- for fpath in ["s3gen.pt", "conds.pt"]:
+ # Check if MPS is available on macOS
+ if device == "mps" and not torch.backends.mps.is_available():
+ if not torch.backends.mps.is_built():
+ logger.warning("MPS not available because the current PyTorch install was not built with MPS enabled.")
+ else:
+ logger.warning("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.safetensors", "conds.pt"]:
local_path = hf_hub_download(repo_id=REPO_ID, filename=fpath)
return cls.from_local(Path(local_path).parent, device)
@@ -61,18 +94,27 @@ class ChatterboxVC:
## Load reference wav
s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
- 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)
+ # Convert to tensor and pad to a 40ms boundary
+ s3gen_ref_wav = torch.from_numpy(s3gen_ref_wav).float().unsqueeze(0)
+ s3gen_ref_wav = _pad_wav_to_40ms_multiple(s3gen_ref_wav, S3GEN_SR)
+ s3gen_ref_wav_np = s3gen_ref_wav.squeeze(0).numpy()
+
+ s3gen_ref_wav_np = s3gen_ref_wav_np[:self.DEC_COND_LEN]
+ self.ref_dict = self.s3gen.embed_ref(s3gen_ref_wav_np, S3GEN_SR, device=self.device)
def generate(
self,
audio,
target_voice_path=None,
+ n_timesteps=10,
+ pbar=None,
+ temperature=1.0,
+ flow_cfg_scale=0.7
):
if target_voice_path:
self.set_target_voice(target_voice_path)
else:
- assert self.ref_dict is not None, "Please `prepare_conditionals` first or specify `target_voice_path`"
+ assert self.ref_dict is not None, "Please `set_target_voice` first or specify `target_voice_path`"
with torch.inference_mode():
audio_16, _ = librosa.load(audio, sr=S3_SR)
@@ -82,6 +124,10 @@ class ChatterboxVC:
wav, _ = self.s3gen.inference(
speech_tokens=s3_tokens,
ref_dict=self.ref_dict,
+ n_timesteps=n_timesteps,
+ pbar=pbar,
+ temperature=temperature,
+ flow_cfg_scale=flow_cfg_scale
)
wav = wav.squeeze(0).detach().cpu().numpy()
watermarked_wav = self.watermarker.apply_watermark(wav, sample_rate=self.sr)