mind/backends.py

"""AIRWAVES GPU-window plumbing: the VoxCPM2 hype-man bank in ONE @spaces.GPU call.

ZeroGPU pickles every ARGUMENT into the forked worker, so the voice actor (which
holds the loaded model) is resolved from module state the worker inherits via
fork — only picklable args (the lines + voice id) cross the boundary.
register_gpu_actors() runs before the first GPU call; `spaces` is imported (and
the decorator applied) at import time, which also keeps it ahead of any torch
import. Locally (no `spaces`) the bare function runs.
"""
from __future__ import annotations

import logging
import threading
from typing import Optional

from .voice import make_voice

logger = logging.getLogger("airwaves")

# Pre-fork actor: the loaded voice the GPU worker uses (resolved from module state).
_GPU_ACTORS: dict = {"voice": None}


def register_gpu_actors(voice) -> None:
    _GPU_ACTORS["voice"] = voice


def _bank_entry(lines: list, voice_id: str) -> list:
    """Synthesize every line; ONLY picklable args cross into the worker."""
    voice = _GPU_ACTORS["voice"] or _defaults()
    return [voice.speak(line, voice_id) for line in lines]


# Decorate at import (ZeroGPU discovers @spaces.GPU functions at startup) and
# keep `import spaces` ahead of torch. ~5 short lines fit well inside 120s.
try:  # pragma: no cover — `spaces` exists only on the HF runtime
    import spaces  # noqa: PLC0415 - MUST precede torch

    _bank_pipeline = spaces.GPU(duration=120)(_bank_entry)
except Exception:  # noqa: BLE001 - local dev / tests
    _bank_pipeline = _bank_entry


_LOCK = threading.Lock()
_VOICE = None


def _defaults():
    """Process-wide voice singleton resolved from AIRWAVES_BACKEND."""
    global _VOICE
    with _LOCK:
        if _VOICE is None:
            _VOICE = make_voice()
        return _VOICE


def warm_voice() -> None:
    """Eagerly build + preload the voice (call at app startup) so the forked GPU
    worker inherits a loaded model and the first /api/hype stays fast."""
    voice = _defaults()
    preload = getattr(voice, "preload", None)
    if callable(preload):
        preload()


def synth_bank(lines: dict, voice_id: str = "club_mc") -> dict:
    """Synthesize every hype line in ONE GPU window; returns {key: wav_bytes}."""
    voice = _defaults()
    register_gpu_actors(voice)  # pin the actor before the worker forks
    keys = list(lines.keys())
    texts = [lines[k] for k in keys]
    if getattr(voice, "name", "") == "vox":
        wavs = _bank_pipeline(texts, voice_id)   # @spaces.GPU window
    else:
        wavs = _bank_entry(texts, voice_id)      # mock / local
    return dict(zip(keys, wavs))