mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-15 04:40:40 -07:00
327 lines
11 KiB
Python
327 lines
11 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."""
|
|
from ..utils.hf_progress import HFProgressTracker, create_hf_progress_callback
|
|
|
|
progress_manager = get_progress_manager()
|
|
task_manager = get_task_manager()
|
|
model_name = "chatterbox-tts"
|
|
|
|
is_cached = self._is_model_cached()
|
|
|
|
# Set up HF progress tracking (intercepts tqdm for file-level progress)
|
|
progress_callback = create_hf_progress_callback(model_name, progress_manager)
|
|
tracker = HFProgressTracker(progress_callback, filter_non_downloads=is_cached)
|
|
tracker_context = tracker.patch_download()
|
|
tracker_context.__enter__()
|
|
|
|
if not is_cached:
|
|
task_manager.start_download(model_name)
|
|
progress_manager.update_progress(
|
|
model_name=model_name,
|
|
current=0,
|
|
total=0,
|
|
filename="Connecting to HuggingFace...",
|
|
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.
|
|
try:
|
|
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,
|
|
)
|
|
finally:
|
|
tracker_context.__exit__(None, None, None)
|
|
|
|
# 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)
|