mind/voice.py

"""Voice synthesis for the AIRWAVES hype-man: MockVoice (stdlib beep) + VoxVoice.

Cloned from pareidolia's verified VoxCPM2 wrapper:
- Voice design needs NO reference audio — prefix the line with a parenthesized
  persona description: ``"(A booming club MC…) Drop it!"``.
- ``VoxCPM.from_pretrained("openbmb/VoxCPM2", load_denoiser=False,
  optimize=False)``; ``generate(text=…, cfg_value=2.0, inference_timesteps=10)``.
- ``TORCHDYNAMO_DISABLE=1`` must precede ANY torch import (torch.compile warmup
  breaks ZeroGPU) — set unconditionally at module top.

All ML imports live inside VoxVoice's lazy path; MockVoice touches nothing
heavier than the stdlib.
"""
from __future__ import annotations

import io
import math
import os
import struct
import threading
import wave
from typing import Optional

os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

# Club-MC personas — voice design from text (no reference audio). One designed
# voice keeps the hype-man consistent across the set.
VOICE_DESIGNS: dict[str, str] = {
    "club_mc": (
        "A booming, hyped club MC — breathless and electric, commanding the "
        "crowd like a ringmaster, with a little Spanglish swagger"
    ),
}
DEFAULT_VOICE_ID = "club_mc"


def _beep_wav(freq: float = 660.0, dur: float = 0.4, sr: int = 24_000) -> bytes:
    """A short decaying sine — the mock 'voice', so the frontend ducking/timing
    develops against real audio with zero ML and zero asset files."""
    buf = io.BytesIO()
    w = wave.open(buf, "wb")
    w.setnchannels(1); w.setsampwidth(2); w.setframerate(sr)
    n = int(sr * dur)
    frames = bytearray()
    for i in range(n):
        env = math.exp(-3.0 * i / n)
        frames += struct.pack("<h", int(32767 * 0.3 * env * math.sin(2 * math.pi * freq * i / sr)))
    w.writeframes(bytes(frames)); w.close()
    return buf.getvalue()


class MockVoice:
    """Canned beep voice — every line is the same short tone. Dependency-free;
    lets the whole app + the ducking path run on a laptop with AIRWAVES_BACKEND=mock."""

    name = "mock"
    model_id = "airwaves-mock-beep"
    sample_rate = 24_000
    _cached: Optional[bytes] = None

    def preload(self) -> None:  # nothing to load
        ...

    def speak(self, line: str, voice_id: str) -> bytes:
        if MockVoice._cached is None:
            MockVoice._cached = _beep_wav()
        return MockVoice._cached


_VOX_MODEL = None
_VOX_LOCK = threading.Lock()


def _ensure_vox():
    """Load VoxCPM2 once per process (module-level singleton). On ZeroGPU this
    MUST be reached via VoxVoice.preload() at startup so the forked GPU worker
    inherits a loaded model instead of re-paying the 2.3B load in-window."""
    global _VOX_MODEL
    with _VOX_LOCK:
        if _VOX_MODEL is None:
            from voxcpm import VoxCPM  # heavy import, guarded by design

            _VOX_MODEL = VoxCPM.from_pretrained(
                "openbmb/VoxCPM2",
                load_denoiser=False,
                optimize=False,  # torch.compile warmup breaks ZeroGPU
            )
        return _VOX_MODEL


class VoxVoice:
    """VoxCPM2 voice-design synthesis — the AIRWAVES hype-man's actual voice.

    ``generate_fn`` is a test seam: a ``(text) -> waveform`` callable that
    replaces the real model so the wav-encoding path is testable with no ML.
    """

    name = "vox"
    model_id = "openbmb/VoxCPM2"
    sample_rate = 48_000

    def __init__(self, generate_fn=None):
        self._generate_fn = generate_fn

    def preload(self) -> None:
        if self._generate_fn is None:
            _ensure_vox()

    def speak(self, line: str, voice_id: str) -> bytes:
        """Synthesize ``line`` in the designed voice; return 48 kHz WAV bytes.
        cfg_value=2.0 / inference_timesteps=10 are the verified speed settings."""
        design = VOICE_DESIGNS.get(voice_id) or VOICE_DESIGNS[DEFAULT_VOICE_ID]
        text = f"({design}) {line}"
        if self._generate_fn is not None:
            waveform = self._generate_fn(text)
        else:
            waveform = _ensure_vox().generate(text=text, cfg_value=2.0, inference_timesteps=10)
        return _waveform_to_wav_bytes(waveform, self.sample_rate)


def _waveform_to_wav_bytes(waveform, sample_rate: int) -> bytes:
    import numpy as np
    import soundfile as sf

    data = np.asarray(waveform)
    if data.ndim > 1:
        data = data.squeeze()
    buf = io.BytesIO()
    sf.write(buf, data, sample_rate, format="WAV", subtype="PCM_16")
    return buf.getvalue()


def make_voice(backend_name: Optional[str] = None):
    raw = (backend_name or os.environ.get("AIRWAVES_BACKEND") or "mock").strip().lower()
    if raw in ("mock", ""):
        return MockVoice()
    if raw in ("zerogpu", "zero-gpu", "vox", "voxcpm"):
        return VoxVoice()
    raise ValueError(f"unknown AIRWAVES_BACKEND {raw!r} (expected mock | zerogpu)")