voice.py

"""Hollow's voice: Kokoro-82M whispered through a locked DSP recipe (the
validated 'af_nicole, 28% fog' cast). CPU only — no GPU budget. Import is
guarded: where kokoro is unavailable (e.g. the Python 3.13 dev venv) the app
runs silent instead of crashing."""
import os
# cap CPU math threads BEFORE numpy/torch load — on many-core machines under
# memory pressure OpenBLAS spawns one buffer-hungry thread per core and the CPU
# allocator fails, so Kokoro won't load. 4 is plenty for our CPU-side TTS work.
os.environ.setdefault("OMP_NUM_THREADS", "4")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "4")
os.environ.setdefault("MKL_NUM_THREADS", "4")

import base64
import importlib.util
import re
import subprocess
import sys
import threading

import numpy as np

from sting import _wav

_SR = 24000
_VOICE = "af_nicole"
_SPEED = 0.94   # a touch quicker than 0.88 so the finale recital isn't dense; still a slow whisper
_FILTER_INTENSITY = 0.28   # the "verylight" setting the user picked

# Kokoro runs in a SEPARATE process (voice_worker.py) with CUDA hidden, so the
# CPU TTS never initializes a CUDA context in THIS process. On ZeroGPU that is
# mandatory: any CUDA init in the main app process poisons the @spaces.GPU worker
# fork ("No CUDA GPUs are available") and every model turn dies. The DSP below
# stays here (pure numpy); only the torch model lives in the subprocess.
# `_PIPELINE` is kept as a public availability flag (None = silent); speak() and
# the tests gate on it. HOLLOW_VOICE_OFF=1 forces silence.
_VOICE_OFF = os.environ.get("HOLLOW_VOICE_OFF") == "1"
_KOKORO_AVAILABLE = (not _VOICE_OFF) and importlib.util.find_spec("kokoro") is not None
_PIPELINE = True if _KOKORO_AVAILABLE else None
if _VOICE_OFF:
    print("[voice] disabled (HOLLOW_VOICE_OFF=1)")
elif not _KOKORO_AVAILABLE:
    print("[voice] disabled (kokoro not installed)")

_WORKER = None
_WORKER_LOCK = threading.Lock()
_WORKER_DEAD = False        # set if the worker can't start, so we don't retry forever


def _start_worker():
    """Lazily spawn the CUDA-free Kokoro worker and reuse it across calls.
    Returns the Popen handle, or None if voice is unavailable / it won't start."""
    global _WORKER, _WORKER_DEAD
    if not _KOKORO_AVAILABLE or _WORKER_DEAD:
        return None
    if _WORKER is not None and _WORKER.poll() is None:
        return _WORKER
    try:
        worker = os.path.join(os.path.dirname(os.path.abspath(__file__)), "voice_worker.py")
        env = dict(os.environ)
        env["CUDA_VISIBLE_DEVICES"] = ""        # the worker also sets it; belt-and-suspenders
        # stderr is inherited (worker noise -> the app log) so a Kokoro/torch
        # failure in the worker is visible; the protocol pipe (stdout) stays clean.
        p = subprocess.Popen([sys.executable, worker], stdin=subprocess.PIPE,
                             stdout=subprocess.PIPE, stderr=None,
                             bufsize=0, env=env)
        ready = b""
        for _ in range(200):                    # skip any stray noise, wait for READY
            ready = p.stdout.readline()         # blocks until Kokoro has loaded
            if not ready or b"READY" in ready:
                break
        if b"READY" not in ready:
            try:
                p.kill()
            except Exception:
                pass
            _WORKER_DEAD = True
            print("[voice] worker did not report READY; voice disabled")
            return None
        print("[voice] worker ready")
        _WORKER = p
        return _WORKER
    except Exception as e:
        print(f"[voice] worker failed to start: {e!r}")
        _WORKER_DEAD = True
        return None


def _synth(cleaned: str):
    """Send cleaned text to the worker; return raw float32 audio (or None)."""
    global _WORKER
    p = _start_worker()
    if p is None:
        return None
    try:
        with _WORKER_LOCK:
            p.stdin.write(base64.b64encode(cleaned.encode("utf-8")) + b"\n")
            p.stdin.flush()
            line = p.stdout.readline()
        if not line or line.strip() == b"ERR":
            return None
        return np.frombuffer(base64.b64decode(line.strip()), dtype=np.float32)
    except Exception as e:
        print(f"[voice] worker synth failed: {e!r}")
        _WORKER = None          # drop the (maybe broken) worker; next call respawns
        return None


def _pitch_up(x, factor):
    n = int(len(x) / factor)
    idx = np.linspace(0, len(x) - 1, n)
    return np.interp(idx, np.arange(len(x)), x).astype(np.float32)


def _breath(x, amount):
    rng = np.random.default_rng(3)
    env = np.abs(x)
    k = np.hanning(441); k /= k.sum()
    env = np.convolve(env, k, mode="same")
    noise = rng.standard_normal(len(x)).astype(np.float32) * env
    return x * (1 - amount) + noise * amount * 2.0


def _fog_reverb(x, taps):
    out = x.copy()
    for delay_s, gain in taps:
        d = int(delay_s * _SR)
        if d < len(x):
            echo = np.zeros_like(x)
            echo[d:] = x[:-d] * gain
            out = out + echo
    return out


def _fog(x, intensity=_FILTER_INTENSITY, pitch=1.16):
    y = _pitch_up(x, pitch)
    y = _breath(y, amount=0.38 * intensity)
    taps = [(0.06, 0.5 * intensity), (0.13, 0.32 * intensity),
            (0.23, 0.18 * intensity)]
    y = _fog_reverb(y, taps)
    return (y / (np.max(np.abs(y)) + 1e-9) * 0.85).astype(np.float32)


def _clean_for_tts(text: str) -> str:
    """Kokoro verbalizes dashes ('—', '-') as words. Replace em/en dashes with a
    comma pause and hyphens with a space, then collapse whitespace."""
    text = text.replace("—", ", ").replace("–", ", ").replace("-", " ")
    return re.sub(r"\s+", " ", text).strip()


def speak(text: str) -> str | None:
    """Synthesize `text` as Hollow's whisper; return base64 WAV, or None if
    voice is unavailable or anything fails (never raises)."""
    if _PIPELINE is None or not text or not text.strip():
        return None
    try:
        cleaned = _clean_for_tts(text)
        if not cleaned:
            return None
        audio = _synth(cleaned)
        if audio is None or audio.size == 0:
            return None
        audio = np.asarray(audio, dtype=np.float32)
        audio = audio / (np.max(np.abs(audio)) + 1e-9) * 0.9
        return base64.b64encode(_wav(_fog(audio), _SR)).decode()
    except Exception as e:
        print(f"[voice] speak failed: {e!r}")
        return None