voicebox/backend/backends/chatterbox_backend.py

"""
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

            # 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

                    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)

            # 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 = 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)

            # 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
            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:
            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)