Labyricorn/voicebox
1
0
Fork 0
mirror of https://github.com/jamiepine/voicebox.git synced 2026-09-20 15:20:39 -07:00
Code Issues Packages Projects Releases Wiki Activity

refactor: remove dead code, deduplicate backends

Browse Source 
Phase 1 - delete dead code:
- studio.py, migrate_add_instruct.py, utils/validation.py
- duplicate _profile_to_response in main.py, duplicate asyncio import
- pointless _get_profiles_dir/_get_generations_dir wrappers
- duplicate LANGUAGE_CODE_TO_NAME and WHISPER_HF_REPOS constants

Phase 2 - extract backends/base.py with shared utilities:
- is_model_cached() replaces 7 copy-pasted HF cache checks
- get_torch_device() replaces 5 device detection methods
- combine_voice_prompts() replaces 5 identical implementations
- model_load_progress() ctx manager replaces progress boilerplate in all backends
- patch_chatterbox_f32() replaces identical monkey-patches in both chatterbox backends

net -1078 lines across the backend
This commit is contained in:
Jamie Pine
2026-03-16 01:10:02 -07:00
parent 9514c6596c
commit 0813a3d9d6
16 changed files with 480 additions and 1281 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/backends/chatterbox_backend.py
+28 -159
View File
@@ -8,7 +8,6 @@ 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
@@ -16,9 +15,13 @@ 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
from .base import (
is_model_cached,
get_torch_device,
combine_voice_prompts as _combine_voice_prompts,
model_load_progress,
patch_chatterbox_f32,
)
logger = logging.getLogger(__name__)
@@ -45,17 +48,7 @@ class ChatterboxTTSBackend:
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"
return get_torch_device(force_cpu_on_mac=True)
def is_loaded(self) -> bool:
return self.model is not None
@@ -64,33 +57,7 @@ class ChatterboxTTSBackend:
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
return is_model_cached(CHATTERBOX_HF_REPO, required_files=_MTL_WEIGHT_FILES)
async def load_model(self, model_size: str = "default") -> None:
"""Load the Chatterbox multilingual model."""
@@ -103,133 +70,45 @@ class ChatterboxTTSBackend:
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:
with model_load_progress(model_name, is_cached):
device = self._get_device()
self._device = device
logger.info(f"Loading Chatterbox Multilingual TTS on {device}...")
import torch
from chatterbox.mtl_tts import ChatterboxMultilingualTTS
# Load into a local variable first, apply all patches, then
# assign to self.model. This avoids leaving a half-initialised
# model on self.model if any patch step raises an exception.
#
# 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
if device == "cpu":
_orig_torch_load = torch.load
def _patched_load(*args, **kwargs):
kwargs.setdefault("map_location", "cpu")
return _orig_torch_load(*args, **kwargs)
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:
model = ChatterboxMultilingualTTS.from_pretrained(
device=device,
)
finally:
torch.load = _orig_torch_load
else:
model = ChatterboxMultilingualTTS.from_pretrained(
device=device,
)
finally:
tracker_context.__exit__(None, None, None)
with ChatterboxTTSBackend._load_lock:
torch.load = _patched_load
try:
model = ChatterboxMultilingualTTS.from_pretrained(device=device)
finally:
torch.load = _orig_torch_load
else:
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.
# Fix sdpa attention for output_attentions support
t3_tfmr = model.t3.tfmr
if hasattr(t3_tfmr, "config") and hasattr(
t3_tfmr.config, "_attn_implementation"
):
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)
# Patch float64 → float32 dtype mismatches in upstream chatterbox.
# librosa.load returns float64 numpy; multiple upstream code paths
# convert it to a torch tensor via torch.from_numpy() without
# casting, then matmul it against float32 model weights.
import types
# Patch S3Tokenizer (used by s3gen.tokenizer)
_tokzr = model.s3gen.tokenizer
_orig_log_mel = _tokzr.log_mel_spectrogram.__func__
def _f32_log_mel(self_tokzr, audio, padding=0):
import torch as _torch
if _torch.is_tensor(audio):
audio = audio.float()
return _orig_log_mel(self_tokzr, audio, padding)
_tokzr.log_mel_spectrogram = types.MethodType(_f32_log_mel, _tokzr)
# Patch VoiceEncoder
_ve = model.ve
_orig_ve_forward = _ve.forward.__func__
def _f32_ve_forward(self_ve, mels):
return _orig_ve_forward(self_ve, mels.float())
_ve.forward = types.MethodType(_f32_ve_forward, _ve)
# All patches applied successfully — publish the model
patch_chatterbox_f32(model)
self.model = model
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:
import traceback
logger.error(f"Failed to load Chatterbox: {e}\n{traceback.format_exc()}")
if not is_cached:
progress_manager.mark_error(model_name, str(e))
task_manager.error_download(model_name, str(e))
raise
logger.info("Chatterbox Multilingual TTS loaded successfully")
def unload_model(self) -> None:
"""Unload model to free memory."""
@@ -268,17 +147,7 @@ class ChatterboxTTSBackend:
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
return await _combine_voice_prompts(audio_paths, reference_texts)
# Per-language generation defaults. Lower temp + higher cfg = clearer speech.
_LANG_DEFAULTS: ClassVar[dict] = {