Blame for engine.py

"""Pinned, plain-PyTorch inference. No UI code and no independent GPU thread."""
import base64
from dataclasses import dataclass
import gc
import json
import logging
import os
from pathlib import Path
import sys
import time

# The SDK must be imported before torch. These imports run at Space startup.
import spaces

from config import DECODER, MAX_FRAMES, MODEL_ID, MODEL_REVISION, ROOT, SAMPLE_RATE
from lifecycle import Cancelled, check_cancelled

os.environ["ZONOS3_RELEASE_REPO"] = MODEL_ID
os.environ["ZONOS3_RELEASE_REVISION"] = MODEL_REVISION
sys.path.insert(0, str(ROOT / "vendor"))

import numpy as np
import torch
from huggingface_hub import snapshot_download
from zonos3.reference.lm import Zonos3LM
from zonos3.reference.synth import ReferenceAssets, synthesize
from zonos3.tokenizer.textnorm import normalize_request_text
from zonos3.utils.release import RELEASE_LM

log = logging.getLogger("zonos.engine")
DEVICE = torch.device("cuda:0")


@dataclass
class Outcome:
    pcm: np.ndarray | None
    seconds: float
    cancelled: bool = False


def gpu_budget(request, marker, progress):
    # Reservation budget, never an ETA. spaces 0.51.3 multiplies large
    # Blackwell requests by 1.5; the former ceiling rejected long public takes
    # before execution. Reserve at most 60 seconds (90 after that SDK factor).
    return min(60, max(30, 10 + int(len(request["text"]) * .08)))


class NormalizationPrivacyFilter(logging.Filter):
    def filter(self, record):
        # Retain the upstream warning, but not its quoted user text (pinned contract).
        if record.msg == "TTS text normalization failed for lang=%s on %r; kept as typed":
            record.args = (record.args[0], "[input omitted]")
        return True


class Engine:
    def __init__(self):
        log.info("startup_model_loading revision=%s", MODEL_REVISION)
        self.lm = Zonos3LM.load(RELEASE_LM, device="cpu", dtype=torch.float32).to(DEVICE).eval()
        self.assets = ReferenceAssets.load()
        self.assets.decoder(self.lm.config.audio_codec, DEVICE, decoder=DECODER)
        if self.lm.config.speaker_embedding_codec != "ownmect_s2" or self.lm.config.speaker_embedding_dim != 192:
            raise RuntimeError("Unexpected speaker encoder in the pinned checkpoint")
        # Embedding mode, as in the official Lite app: no dual-codec teacher download.
        self.assets.speaker_encoder("ownmect_s2", 192, DEVICE)
        logging.getLogger("zonos3.tokenizer.textnorm").addFilter(NormalizationPrivacyFilter())
        normalizer = self.assets.normalizer()  # warmup=False: no background normalization thread.
        if normalizer is None:
            raise RuntimeError("This application requires the normalizer to be available")
        for lang in ("en", "fr"):
            normalizer.get(lang)
        catalog = json.loads((ROOT / "voices.json").read_text())["voices"]
        cache = Path(snapshot_download(MODEL_ID, revision=MODEL_REVISION,
                                      allow_patterns=[v["path"] for v in catalog]))
        self.voices = {}
        for voice in catalog:
            content = (cache / voice["path"]).read_bytes()
            import io
            vector = np.load(io.BytesIO(content), allow_pickle=False)
            if vector.shape != (192,) or not np.isfinite(vector).all():
                raise RuntimeError("Invalid preset embedding in pinned release")
            self.voices[voice["name"]] = base64.b64encode(content).decode("ascii")
        # CPU weight buffers used during startup can be reclaimed before the worker forks.
        gc.collect()
        log.info("startup_complete presets=%d decoder=%s conditioning=embedding", len(self.voices), DECODER)

    def prepare(self, request, reference_bytes=None):
        # Text normalization and validation happen outside the allocated GPU function.
        text = normalize_request_text(request.text, request.language,
                                      self.assets.normalizer() if request.normalize else None)
        result = {
            "text": text, "language": request.language, "seed": request.seed,
            "text_normalization": False, "speaker_mode": "embedding", "decoder": DECODER,
            "stream": False, "max_tokens": MAX_FRAMES,
            "temperature": request.temperature, "top_p": request.top_p,
            "topk": request.top_k, "repetition_penalty": request.repetition_penalty,
        }
        if request.source == "Voice clip":
            result["speaker_audio_base64"] = base64.b64encode(reference_bytes).decode("ascii")
        else:
            result["speaker_embedding_base64"] = self.voices[request.voice]
        if request.rate is not None:
            result["condition_values"] = {"speaking_rate": request.rate}
        return result

    def run(self, request, marker, progress):
        return _gpu_run(request, str(marker), progress)


ENGINE = None


@spaces.GPU(duration=gpu_budget)
def _gpu_run(request, marker, progress):
    started = time.perf_counter()
    identifier = Path(marker).stem
    log.info("gpu_function_enter request=%s", identifier)
    last_report = 0.0
    try:
        check_cancelled(marker)
        progress(None, desc="Preparing voice")

        def on_plan(plan):
            check_cancelled(marker)
            progress(None, desc="Synthesizing speech")
            log.info("inference_started request=%s", identifier)

        def on_audio(pcm):
            nonlocal last_report
            check_cancelled(marker)
            now = time.perf_counter()
            if now - last_report >= 1:
                progress(None, desc="Synthesizing speech")
                last_report = now

        # Callbacks are supplied to the actual upstream engine. No patched loops,
        # thread, wall-clock progress percentage, alignment estimate or torch.compile.
        result = synthesize(ENGINE.lm, request, assets=ENGINE.assets, device=DEVICE,
                            on_plan=on_plan, on_audio=on_audio)
        check_cancelled(marker)
        if result.eos_frame is None:
            raise ValueError("Audio reached the take limit. Please shorten the text.")
        pcm = np.asarray(result.pcm, dtype=np.float32).copy()
        del result
        return Outcome(pcm, time.perf_counter() - started)
    except Cancelled:
        return Outcome(None, time.perf_counter() - started, cancelled=True)
    finally:
        # Complete submitted kernels before returning control to spaces.GPU.
        # Returning from this regular function lets the pinned SDK release the lease.
        # A cached worker is not an active lease. empty_cache is not used as a substitute.
        torch.cuda.synchronize(DEVICE)
        log.info("gpu_function_exit request=%s", identifier)


def load_engine():
    global ENGINE
    if ENGINE is None:
        ENGINE = Engine()
    return ENGINE