chatterbex

#!/usr/bin/env python3
"""
Chatterbex Bridge - Python bridge for Elixir Chatterbex library.

This script communicates with the Elixir GenServer via stdin/stdout,
loading and running Chatterbox TTS models.
"""

import sys
import json
import base64
import io
import os

# Force eager attention to avoid SDPA compatibility issues with output_attentions
os.environ["ATTN_IMPLEMENTATION"] = "eager"

# Attempt to import torch and torchaudio early to catch missing deps
try:
    import torch
    import torchaudio as ta
except ImportError as e:
    print(json.dumps({"status": "error", "error": f"Missing dependency: {e}"}))
    sys.exit(1)

# Monkey patch torch.load to handle CUDA->CPU/MPS device mapping
# This is needed because Chatterbox models are saved with CUDA references
_original_torch_load = torch.load


def _patched_torch_load(f, map_location=None, **kwargs):
    """Patched torch.load that maps CUDA tensors to CPU for compatibility."""
    if map_location is None:
        map_location = "cpu"
    return _original_torch_load(f, map_location=map_location, **kwargs)


torch.load = _patched_torch_load


def _detect_device(requested_device: str) -> str:
    """
    Detect and validate the compute device.

    Returns the best available device, with fallback handling for MPS.
    """
    if requested_device == "cuda":
        if torch.cuda.is_available():
            return "cuda"
        return "cpu"
    elif requested_device == "mps":
        if torch.backends.mps.is_available():
            return "mps"
        return "cpu"
    else:
        return "cpu"


class ChatterboxBridge:
    """Bridge between Elixir and Chatterbox TTS models."""

    def __init__(self):
        self.model = None
        self.model_type = None
        self.device = None

    def init_model(self, model_type: str, device: str = "cuda") -> dict:
        """Initialize the specified Chatterbox model."""
        try:
            # Validate and detect actual device
            actual_device = _detect_device(device)
            self.device = actual_device
            self.model_type = model_type

            # For MPS, load to CPU first then move components
            # This avoids "Placeholder storage has not been allocated on MPS device" errors
            load_device = "cpu" if actual_device == "mps" else actual_device

            if model_type == "turbo":
                from chatterbox.tts_turbo import ChatterboxTurboTTS
                self.model = ChatterboxTurboTTS.from_pretrained(device=load_device)

            elif model_type == "english":
                from chatterbox.tts import ChatterboxTTS
                self.model = ChatterboxTTS.from_pretrained(device=load_device)

            elif model_type == "multilingual":
                from chatterbox.mtl_tts import ChatterboxMultilingualTTS
                self.model = ChatterboxMultilingualTTS.from_pretrained(device=load_device)
                # Fix SDPA attention compatibility issue with output_attentions
                self._fix_attention_implementation()

            else:
                return {"status": "error", "error": f"Unknown model type: {model_type}"}

            # Move model components to MPS if requested
            if actual_device == "mps" and load_device == "cpu":
                self._move_to_mps()

            return {"status": "ok", "device": actual_device}

        except Exception as e:
            return {"status": "error", "error": str(e)}

    def _move_to_mps(self) -> None:
        """
        Move model components to MPS device.

        Chatterbox models have multiple subcomponents that need to be moved individually.
        If MPS fails for any component, fall back to CPU for stability.
        """
        try:
            # Common components across model types
            for component_name in ["t3", "s3gen", "ve", "model"]:
                if hasattr(self.model, component_name):
                    component = getattr(self.model, component_name)
                    if component is not None and hasattr(component, "to"):
                        setattr(self.model, component_name, component.to("mps"))

            # Update model's device reference if it has one
            if hasattr(self.model, "device"):
                self.model.device = "mps"

        except Exception:
            # MPS failed, fall back to CPU
            self.device = "cpu"
            if hasattr(self.model, "device"):
                self.model.device = "cpu"

    def _fix_attention_implementation(self) -> None:
        """
        Fix SDPA attention compatibility issues.

        Some transformer components use SDPA which doesn't support output_attentions.
        This patches them to use eager attention instead.
        """
        try:
            # Check common component names that might have transformer layers
            for component_name in ["t3", "model", "encoder", "decoder"]:
                if hasattr(self.model, component_name):
                    component = getattr(self.model, component_name)
                    if component is not None:
                        # Recursively find and fix attention implementation
                        self._set_eager_attention(component)
        except Exception:
            pass  # Best effort - if it fails, generation will show the original error

    def _set_eager_attention(self, module) -> None:
        """Recursively set attention implementation to eager on transformer modules."""
        if hasattr(module, "config") and hasattr(module.config, "_attn_implementation"):
            module.config._attn_implementation = "eager"
        if hasattr(module, "_attn_implementation"):
            module._attn_implementation = "eager"
        # Recurse into submodules
        if hasattr(module, "modules"):
            for submodule in module.modules():
                if submodule is not module:
                    if hasattr(submodule, "config") and hasattr(submodule.config, "_attn_implementation"):
                        submodule.config._attn_implementation = "eager"
                    if hasattr(submodule, "_attn_implementation"):
                        submodule._attn_implementation = "eager"

    def generate(self, text: str, **kwargs) -> dict:
        """Generate speech from text."""
        if self.model is None:
            return {"status": "error", "error": "Model not initialized"}

        try:
            # Build generation arguments
            gen_kwargs = {}

            if "audio_prompt" in kwargs and kwargs["audio_prompt"]:
                gen_kwargs["audio_prompt_path"] = kwargs["audio_prompt"]

            if self.model_type == "multilingual" and "language" in kwargs:
                gen_kwargs["language_id"] = kwargs["language"]

            if self.model_type == "english":
                if "exaggeration" in kwargs and kwargs["exaggeration"] is not None:
                    gen_kwargs["exaggeration"] = kwargs["exaggeration"]
                if "cfg_weight" in kwargs and kwargs["cfg_weight"] is not None:
                    gen_kwargs["cfg_weight"] = kwargs["cfg_weight"]

            # Generate audio
            wav = self.model.generate(text, **gen_kwargs)

            # Convert to WAV bytes
            audio_bytes = self._wav_to_bytes(wav, self.model.sr)

            # Encode as base64 for JSON transport
            audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")

            return {"status": "ok", "audio": audio_base64}

        except Exception as e:
            return {"status": "error", "error": str(e)}

    def _wav_to_bytes(self, wav: torch.Tensor, sample_rate: int) -> bytes:
        """Convert a PyTorch tensor to WAV bytes."""
        buffer = io.BytesIO()

        # Ensure correct shape (channels, samples)
        if wav.dim() == 1:
            wav = wav.unsqueeze(0)

        ta.save(buffer, wav.cpu(), sample_rate, format="wav")
        buffer.seek(0)
        return buffer.read()


def main():
    """Main loop - read JSON commands from stdin, write responses to stdout."""
    bridge = ChatterboxBridge()

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue

        try:
            request = json.loads(line)
        except json.JSONDecodeError as e:
            response = {"status": "error", "error": f"Invalid JSON: {e}"}
            print(json.dumps(response), flush=True)
            continue

        request_type = request.get("type")

        if request_type == "init":
            model_type = request.get("model", "turbo")
            device = request.get("device", "cuda")
            response = bridge.init_model(model_type, device)

        elif request_type == "generate":
            text = request.get("text", "")
            kwargs = {
                k: v for k, v in request.items()
                if k not in ("type", "text")
            }
            response = bridge.generate(text, **kwargs)

        elif request_type == "ping":
            response = {"status": "ok", "message": "pong"}

        else:
            response = {"status": "error", "error": f"Unknown request type: {request_type}"}

        print(json.dumps(response), flush=True)


if __name__ == "__main__":
    main()