Spaces:
Running on Zero
Running on Zero
| """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") | |
| 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 | |
| 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 | |