mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-15 04:40:40 -07:00
James Pine
9beb9d7fec
chatterbox-tts 0.1.6 pins numpy<1.26 and torch==2.6 which are incompatible with Python 3.12+. Install with --no-deps and list its sub-dependencies explicitly in requirements.txt. Also removes HFProgressTracker from chatterbox backend to avoid 'generator didn't stop after throw()' errors from tqdm patching.
316 lines
10 KiB
Python
316 lines
10 KiB
Python
"""
|
|
Chatterbox TTS backend implementation.
|
|
|
|
Wraps ChatterboxMultilingualTTS from chatterbox-tts for zero-shot
|
|
voice cloning. Supports 23 languages including Hebrew. Forces CPU
|
|
on macOS due to known MPS tensor issues.
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
import platform
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import ClassVar, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
from . import TTSBackend
|
|
from ..utils.audio import normalize_audio, load_audio
|
|
from ..utils.progress import get_progress_manager
|
|
from ..utils.tasks import get_task_manager
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
CHATTERBOX_HF_REPO = "ResembleAI/chatterbox"
|
|
|
|
# Files that must be present for the multilingual model
|
|
_MTL_WEIGHT_FILES = [
|
|
"t3_mtl23ls_v2.safetensors",
|
|
"s3gen.pt",
|
|
"ve.pt",
|
|
]
|
|
|
|
|
|
class ChatterboxTTSBackend:
|
|
"""Chatterbox Multilingual TTS backend for voice cloning."""
|
|
|
|
# Class-level lock for torch.load monkey-patching
|
|
_load_lock: ClassVar[threading.Lock] = threading.Lock()
|
|
|
|
def __init__(self):
|
|
self.model = None
|
|
self.model_size = "default"
|
|
self._device = None
|
|
self._model_load_lock = asyncio.Lock()
|
|
|
|
def _get_device(self) -> str:
|
|
"""Get the best available device. Forces CPU on macOS (MPS issue)."""
|
|
if platform.system() == "Darwin":
|
|
return "cpu"
|
|
try:
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
except ImportError:
|
|
pass
|
|
return "cpu"
|
|
|
|
def is_loaded(self) -> bool:
|
|
return self.model is not None
|
|
|
|
def _get_model_path(self, model_size: str = "default") -> str:
|
|
return CHATTERBOX_HF_REPO
|
|
|
|
def _is_model_cached(self, model_size: str = "default") -> bool:
|
|
"""Check if the Chatterbox multilingual model is cached locally."""
|
|
try:
|
|
from huggingface_hub import constants as hf_constants
|
|
|
|
repo_cache = Path(hf_constants.HF_HUB_CACHE) / (
|
|
"models--" + CHATTERBOX_HF_REPO.replace("/", "--")
|
|
)
|
|
|
|
if not repo_cache.exists():
|
|
return False
|
|
|
|
blobs_dir = repo_cache / "blobs"
|
|
if blobs_dir.exists() and any(blobs_dir.glob("*.incomplete")):
|
|
return False
|
|
|
|
# Check for multilingual weight files
|
|
snapshots_dir = repo_cache / "snapshots"
|
|
if snapshots_dir.exists():
|
|
for fname in _MTL_WEIGHT_FILES:
|
|
if not any(snapshots_dir.rglob(fname)):
|
|
return False
|
|
return True
|
|
|
|
return False
|
|
except Exception as e:
|
|
logger.warning(f"Error checking Chatterbox cache: {e}")
|
|
return False
|
|
|
|
async def load_model(self, model_size: str = "default") -> None:
|
|
"""Load the Chatterbox multilingual model."""
|
|
if self.model is not None:
|
|
return
|
|
async with self._model_load_lock:
|
|
if self.model is not None:
|
|
return
|
|
await asyncio.to_thread(self._load_model_sync)
|
|
|
|
def _load_model_sync(self):
|
|
"""Synchronous model loading."""
|
|
progress_manager = get_progress_manager()
|
|
task_manager = get_task_manager()
|
|
model_name = "chatterbox-tts"
|
|
|
|
is_cached = self._is_model_cached()
|
|
|
|
if not is_cached:
|
|
task_manager.start_download(model_name)
|
|
progress_manager.update_progress(
|
|
model_name=model_name,
|
|
current=0,
|
|
total=0,
|
|
filename="Downloading Chatterbox model...",
|
|
status="downloading",
|
|
)
|
|
|
|
try:
|
|
device = self._get_device()
|
|
self._device = device
|
|
|
|
logger.info(f"Loading Chatterbox Multilingual TTS on {device}...")
|
|
|
|
import torch
|
|
from chatterbox.mtl_tts import ChatterboxMultilingualTTS
|
|
|
|
# Monkey-patch torch.load for CPU loading. The model's .pt files
|
|
# were saved on CUDA; from_pretrained() doesn't pass map_location
|
|
# so loading on CPU fails without this.
|
|
if device == "cpu":
|
|
_orig_torch_load = torch.load
|
|
|
|
def _patched_load(*args, **kwargs):
|
|
kwargs.setdefault("map_location", "cpu")
|
|
return _orig_torch_load(*args, **kwargs)
|
|
|
|
with ChatterboxTTSBackend._load_lock:
|
|
torch.load = _patched_load
|
|
try:
|
|
self.model = ChatterboxMultilingualTTS.from_pretrained(
|
|
device=device,
|
|
)
|
|
finally:
|
|
torch.load = _orig_torch_load
|
|
else:
|
|
self.model = ChatterboxMultilingualTTS.from_pretrained(
|
|
device=device,
|
|
)
|
|
|
|
# Fix: transformers >= 4.36 defaults LlamaModel to sdpa attention
|
|
# which doesn't support output_attentions=True (needed by
|
|
# Chatterbox's AlignmentStreamAnalyzer). Force eager attention.
|
|
t3_tfmr = self.model.t3.tfmr
|
|
if hasattr(t3_tfmr, "config") and hasattr(
|
|
t3_tfmr.config, "_attn_implementation"
|
|
):
|
|
t3_tfmr.config._attn_implementation = "eager"
|
|
for layer in getattr(t3_tfmr, "layers", []):
|
|
if hasattr(layer, "self_attn"):
|
|
layer.self_attn._attn_implementation = "eager"
|
|
|
|
if not is_cached:
|
|
progress_manager.mark_complete(model_name)
|
|
task_manager.complete_download(model_name)
|
|
|
|
logger.info("Chatterbox Multilingual TTS loaded successfully")
|
|
|
|
except ImportError as e:
|
|
logger.error(
|
|
"chatterbox-tts package not found. "
|
|
"Install with: pip install chatterbox-tts"
|
|
)
|
|
if not is_cached:
|
|
progress_manager.mark_error(model_name, str(e))
|
|
task_manager.error_download(model_name, str(e))
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Failed to load Chatterbox: {e}")
|
|
if not is_cached:
|
|
progress_manager.mark_error(model_name, str(e))
|
|
task_manager.error_download(model_name, str(e))
|
|
raise
|
|
|
|
def unload_model(self) -> None:
|
|
"""Unload model to free memory."""
|
|
if self.model is not None:
|
|
device = self._device
|
|
del self.model
|
|
self.model = None
|
|
self._device = None
|
|
if device == "cuda":
|
|
import torch
|
|
|
|
torch.cuda.empty_cache()
|
|
logger.info("Chatterbox unloaded")
|
|
|
|
async def create_voice_prompt(
|
|
self,
|
|
audio_path: str,
|
|
reference_text: str,
|
|
use_cache: bool = True,
|
|
) -> Tuple[dict, bool]:
|
|
"""
|
|
Create voice prompt from reference audio.
|
|
|
|
Chatterbox processes reference audio at generation time, so the
|
|
prompt just stores the file path. The actual audio is loaded by
|
|
model.generate() via audio_prompt_path.
|
|
"""
|
|
voice_prompt = {
|
|
"ref_audio": str(audio_path),
|
|
"ref_text": reference_text,
|
|
}
|
|
return voice_prompt, False
|
|
|
|
async def combine_voice_prompts(
|
|
self,
|
|
audio_paths: List[str],
|
|
reference_texts: List[str],
|
|
) -> Tuple[np.ndarray, str]:
|
|
"""Combine multiple reference samples."""
|
|
combined_audio = []
|
|
for path in audio_paths:
|
|
audio, _sr = load_audio(path)
|
|
audio = normalize_audio(audio)
|
|
combined_audio.append(audio)
|
|
|
|
mixed = np.concatenate(combined_audio)
|
|
mixed = normalize_audio(mixed)
|
|
combined_text = " ".join(reference_texts)
|
|
return mixed, combined_text
|
|
|
|
# Per-language generation defaults. Lower temp + higher cfg = clearer speech.
|
|
_LANG_DEFAULTS: ClassVar[dict] = {
|
|
"he": {
|
|
"exaggeration": 0.4,
|
|
"cfg_weight": 0.7,
|
|
"temperature": 0.65,
|
|
"repetition_penalty": 2.5,
|
|
},
|
|
}
|
|
_GLOBAL_DEFAULTS: ClassVar[dict] = {
|
|
"exaggeration": 0.5,
|
|
"cfg_weight": 0.5,
|
|
"temperature": 0.8,
|
|
"repetition_penalty": 2.0,
|
|
}
|
|
|
|
async def generate(
|
|
self,
|
|
text: str,
|
|
voice_prompt: dict,
|
|
language: str = "en",
|
|
seed: Optional[int] = None,
|
|
instruct: Optional[str] = None,
|
|
) -> Tuple[np.ndarray, int]:
|
|
"""
|
|
Generate audio using Chatterbox Multilingual TTS.
|
|
|
|
Args:
|
|
text: Text to synthesize
|
|
voice_prompt: Dict with ref_audio path
|
|
language: BCP-47 language code
|
|
seed: Random seed for reproducibility
|
|
instruct: Unused (protocol compatibility)
|
|
|
|
Returns:
|
|
Tuple of (audio_array, sample_rate)
|
|
"""
|
|
await self.load_model()
|
|
|
|
ref_audio = voice_prompt.get("ref_audio")
|
|
if ref_audio and not Path(ref_audio).exists():
|
|
logger.warning(f"Reference audio not found: {ref_audio}")
|
|
ref_audio = None
|
|
|
|
# Merge language-specific defaults with global defaults
|
|
lang_defaults = self._LANG_DEFAULTS.get(language, self._GLOBAL_DEFAULTS)
|
|
|
|
def _generate_sync():
|
|
import torch
|
|
|
|
if seed is not None:
|
|
torch.manual_seed(seed)
|
|
|
|
logger.info(f"[Chatterbox] Generating: lang={language}")
|
|
|
|
wav = self.model.generate(
|
|
text,
|
|
language_id=language,
|
|
audio_prompt_path=ref_audio,
|
|
exaggeration=lang_defaults["exaggeration"],
|
|
cfg_weight=lang_defaults["cfg_weight"],
|
|
temperature=lang_defaults["temperature"],
|
|
repetition_penalty=lang_defaults["repetition_penalty"],
|
|
)
|
|
|
|
# Convert tensor -> numpy
|
|
if isinstance(wav, torch.Tensor):
|
|
audio = wav.squeeze().cpu().numpy().astype(np.float32)
|
|
else:
|
|
audio = np.asarray(wav, dtype=np.float32)
|
|
|
|
sample_rate = (
|
|
getattr(self.model, "sr", None)
|
|
or getattr(self.model, "sample_rate", 24000)
|
|
)
|
|
|
|
return audio, sample_rate
|
|
|
|
return await asyncio.to_thread(_generate_sync)
|