voice_worker.py

"""Voice synthesis worker — runs Kokoro-82M in a SEPARATE process with CUDA
hidden, so the CPU TTS never initializes a CUDA context in the main app process.

On ZeroGPU that isolation is mandatory: any CUDA init in the main process
poisons the @spaces.GPU worker fork ("No CUDA GPUs are available") and every
model turn dies. Kokoro probes/initializes CUDA on import even with
device="cpu" (the `spaces` patch makes torch.cuda.is_available() True), so the
only reliable fix is to keep it in its own process where CUDA_VISIBLE_DEVICES="".

Protocol (binary, line-based over stdio): the parent writes one line =
base64(utf-8 cleaned text); the worker replies one line = base64(float32 PCM
bytes) or the literal b"ERR". A first b"READY" line signals Kokoro has loaded.
"""
import os

# hide every GPU from this process BEFORE torch is imported (Kokoro -> torch)
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ.setdefault("OMP_NUM_THREADS", "4")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "4")
os.environ.setdefault("MKL_NUM_THREADS", "4")

# Reserve the REAL stdout (fd 1) as the protocol pipe to the parent, then point
# fd 1 at stderr so Kokoro/torch/tqdm prints during model load + synth go to
# stderr and can't corrupt the line-based protocol the parent reads on stdout.
_PROTO_FD = os.dup(1)
os.dup2(2, 1)

import base64
import sys

import numpy as np

_VOICE = "af_nicole"
_SPEED = 0.94


def _out(line: bytes) -> None:
    os.write(_PROTO_FD, line + b"\n")


try:
    from kokoro import KPipeline
    _PIPELINE = KPipeline(lang_code="a", device="cpu")
except Exception as e:                       # no kokoro/model -> parent goes silent
    sys.stderr.write(f"[voice-worker] disabled ({e!r})\n")
    sys.exit(1)


_out(b"READY")

for raw in sys.stdin.buffer:                 # one request per line
    raw = raw.strip()
    if not raw:
        continue
    try:
        text = base64.b64decode(raw).decode("utf-8")
        audio = None
        for _, _, chunk in _PIPELINE(text, voice=_VOICE, speed=_SPEED):
            c = np.asarray(chunk, dtype=np.float32)
            audio = c if audio is None else np.concatenate([audio, c])
        if audio is None or audio.size == 0:
            _out(b"ERR")
        else:
            _out(base64.b64encode(np.asarray(audio, dtype=np.float32).tobytes()))
    except Exception as e:
        sys.stderr.write(f"[voice-worker] synth failed: {e!r}\n")
        _out(b"ERR")