tts.py

"""tts.py — TTS integration for Marquee.

Supports two models selectable at runtime:
  chatterbox  (default) — 0.5B, MIT, fast CPU inference, emotion exaggeration
  orpheus              — 3B,   Apache 2.0, more expressive but slower

Collision fix:
  After all lines are generated, if line[i].t + line[i].duration > line[i+1].t
  we trim line[i]'s text to the number of words that fit in the available slot,
  then re-generate that line's audio. This guarantees no two voices ever overlap.

WAV encoding uses stdlib `wave` + numpy — no torchaudio / torchcodec needed.
"""

import io
import logging
import wave

import numpy as np
import spaces

log = logging.getLogger(__name__)

# ── Vibe → Chatterbox params ──────────────────────────────────────────────────
VIBE_PARAMS_CHATTERBOX: dict[str, dict] = {
    "football":   {"exaggeration": 0.85, "cfg_weight": 0.45},
    "diva":       {"exaggeration": 0.90, "cfg_weight": 0.40},
    "wildlife":   {"exaggeration": 0.20, "cfg_weight": 0.65},
    "boxing":     {"exaggeration": 0.80, "cfg_weight": 0.45},
    "masterchef": {"exaggeration": 0.70, "cfg_weight": 0.50},
}
_CB_DEFAULT = {"exaggeration": 0.55, "cfg_weight": 0.50}

# ── Vibe → Orpheus emotion tags ───────────────────────────────────────────────
# Orpheus uses inline tags: <laugh>, <chuckle>, <sigh>, <cough>, <gasp>, <groan>
VIBE_ORPHEUS_PREFIX: dict[str, str] = {
    "football":   "",           # already hyped by pacing
    "diva":       "",
    "wildlife":   "",
    "boxing":     "",
    "masterchef": "",
}
ORPHEUS_VOICE = "tara"   # options: tara, dan, emma, josh

# ── Model singletons ──────────────────────────────────────────────────────────
_cb_model      = None
_orpheus_model = None


def _get_chatterbox():
    global _cb_model
    if _cb_model is None:
        from chatterbox.tts import ChatterboxTTS
        log.info("[tts] Loading Chatterbox on CPU…")
        _cb_model = ChatterboxTTS.from_pretrained(device="cpu")
        log.info("[tts] Chatterbox ready.")
    return _cb_model


def _get_orpheus():
    global _orpheus_model
    if _orpheus_model is None:
        try:
            from orpheus_tts import OrpheusModel
            log.info("[tts] Loading Orpheus 3B (first call — slow)…")
            _orpheus_model = OrpheusModel(
                model_name="canopylabs/orpheus-tts-0.1-finetune-prod",
            )
            log.info("[tts] Orpheus ready.")
        except ImportError:
            raise RuntimeError(
                "orpheus-speech is not installed. "
                "Add 'orpheus-speech' to requirements.txt and redeploy.")
    return _orpheus_model


# ── WAV helpers ───────────────────────────────────────────────────────────────

def _tensor_to_wav_bytes(wav_tensor, sample_rate: int) -> bytes:
    pcm     = wav_tensor.squeeze().detach().cpu().numpy()
    pcm     = np.clip(pcm, -1.0, 1.0)
    pcm_i16 = (pcm * 32767).astype(np.int16)
    buf     = io.BytesIO()
    with wave.open(buf, "wb") as wf:
        wf.setnchannels(1); wf.setsampwidth(2); wf.setframerate(sample_rate)
        wf.writeframes(pcm_i16.tobytes())
    buf.seek(0)
    return buf.read()


def _pcm_iter_to_wav_bytes(pcm_chunks, sample_rate: int = 24000) -> bytes:
    """For Orpheus which yields int16 numpy chunks."""
    all_pcm = np.concatenate([c.astype(np.int16) for c in pcm_chunks])
    buf     = io.BytesIO()
    with wave.open(buf, "wb") as wf:
        wf.setnchannels(1); wf.setsampwidth(2); wf.setframerate(sample_rate)
        wf.writeframes(all_pcm.tobytes())
    buf.seek(0)
    return buf.read()


def _wav_duration(wav_bytes: bytes) -> float:
    with wave.open(io.BytesIO(wav_bytes)) as wf:
        return wf.getnframes() / wf.getframerate()


def _estimate_duration_sec(text: str, words_per_sec: float = 2.8) -> float:
    """Rough duration estimate from word count — used before audio is generated."""
    return len(text.split()) / words_per_sec


# ── Single-line generation ────────────────────────────────────────────────────

def _generate_chatterbox(text: str, vibe: str) -> bytes:
    import time
    model  = _get_chatterbox()
    params = VIBE_PARAMS_CHATTERBOX.get(vibe, _CB_DEFAULT)
    t0     = time.time()
    wav    = model.generate(text=text,
                            exaggeration=params["exaggeration"],
                            cfg_weight=params["cfg_weight"])
    wav_bytes = _tensor_to_wav_bytes(wav, model.sr)
    log.info(f"[tts/cb] '{text[:40]}' → {time.time()-t0:.1f}s")
    return wav_bytes


@spaces.GPU(duration=60)
def _generate_orpheus(text: str, vibe: str) -> bytes:
    model  = _get_orpheus()
    chunks = list(model.generate_speech(prompt=text, voice=ORPHEUS_VOICE))
    return _pcm_iter_to_wav_bytes(chunks)


def generate_line(text: str, vibe: str, voice_model: str = "chatterbox") -> bytes:
    if voice_model == "orpheus":
        return _generate_orpheus(text, vibe)
    return _generate_chatterbox(text, vibe)


# ── Collision fix ─────────────────────────────────────────────────────────────

def _trim_text_to_duration(text: str, max_sec: float,
                           words_per_sec: float = 2.8) -> str:
    """Trim text to fit within max_sec at expected speaking rate."""
    max_words = max(2, int(max_sec * words_per_sec))
    words     = text.split()
    if len(words) <= max_words:
        return text
    trimmed = " ".join(words[:max_words])
    # prefer a natural break
    for punct in (".", "!", "?", "—", ","):
        idx = trimmed.rfind(punct)
        if idx > len(trimmed) // 2:
            return trimmed[:idx + 1]
    return trimmed + "…"


def _fix_collisions(lines: list[dict], vibe: str,
                    voice_model: str, gap_sec: float = 0.15) -> list[dict]:
    """Re-generate audio for any line that overlaps the next one's start time.

    gap_sec: minimum silence required between end of one line and start of next.
    """
    import base64

    for i in range(len(lines) - 1):
        cur  = lines[i]
        nxt  = lines[i + 1]
        dur  = cur.get("duration", 0)
        if dur == 0:
            continue
        available = nxt["t"] - cur["t"] - gap_sec
        if dur <= available:
            continue

        # Need to trim and re-generate
        log.info(f"[tts] Collision at t={cur['t']}: dur={dur:.2f}s "
                 f"available={available:.2f}s — trimming")
        new_text = _trim_text_to_duration(cur["text"], available)
        if new_text == cur["text"] and available < 0.5:
            # slot too short — skip this line's audio entirely
            cur.pop("audio_b64", None)
            cur["duration"] = 0
            continue
        try:
            wav_bytes = generate_line(new_text, vibe, voice_model)
            new_dur   = _wav_duration(wav_bytes)
            cur["text"]      = new_text
            cur["audio_b64"] = base64.b64encode(wav_bytes).decode()
            cur["duration"]  = round(new_dur, 3)
        except Exception as e:
            log.warning(f"[tts] Re-gen failed: {e}")

    return lines


# ── Public API ────────────────────────────────────────────────────────────────

def generate_script_audio(script: list[dict], vibe: str,
                          voice_model: str = "chatterbox") -> list[dict]:
    import time
    log.info(f"[tts] generate_script_audio: {len(script)} lines, "
             f"vibe={vibe}, model={voice_model}")
    _t0 = time.time()
    """Generate TTS for every line, then fix any audio collisions.

    Args:
        script:       list of {"t": float, "text": str}
        vibe:         football / diva / wildlife / boxing / masterchef
        voice_model:  "chatterbox" (default) or "orpheus"

    Returns:
        list of {"t", "text", "audio_b64"?, "duration"?}
    """
    import base64

    result = []
    for line in script:
        entry = {"t": line["t"], "text": line["text"]}
        try:
            wav_bytes        = generate_line(line["text"], vibe, voice_model)
            entry["audio_b64"] = base64.b64encode(wav_bytes).decode()
            entry["duration"]  = round(_wav_duration(wav_bytes), 3)
        except Exception as e:
            log.warning(f"[tts] Skipping line ('{line['text'][:30]}…'): {e}")
        result.append(entry)

    # Post-process: fix any overlapping lines
    result = _fix_collisions(result, vibe, voice_model)
    audio_ok = sum(1 for r in result if r.get("audio_b64"))
    log.info(f"[tts] All done in {time.time()-_t0:.1f}s — "
             f"{audio_ok}/{len(result)} lines have audio")
    return result