Spaces:
Running on Zero
Running on Zero
| """Orchestrate CPU preparation, the synchronous GPU call, and a single CPU export.""" | |
| from contextlib import contextmanager | |
| import html | |
| import json | |
| import logging | |
| from pathlib import Path | |
| import threading | |
| import time | |
| import gradio as gr | |
| from audio_files import export_take, prune_outputs, read_reference | |
| from lifecycle import AlreadyRunning, Cancelled, Jobs, check_cancelled | |
| from requests_schema import InputError, validate | |
| log = logging.getLogger("zonos.controller") | |
| def message(text, error=False, busy=False): | |
| return f'<p id="status-state" data-busy="{str(busy).lower()}" role="{"alert" if error else "status"}">{html.escape(text)}</p>' | |
| class Controller: | |
| def __init__(self, engine, runtime_dir, voice_names, aligner=None): | |
| self.engine = engine | |
| self.aligner = aligner | |
| self.root = Path(runtime_dir) | |
| self.outputs = self.root / "outputs" | |
| self.outputs.mkdir(parents=True, exist_ok=True) | |
| self.jobs = Jobs(self.root / "jobs") | |
| self.gate = threading.Semaphore(1) | |
| self.voice_names = set(voice_names) | |
| def inference_slot(self, job): | |
| acquired = False | |
| try: | |
| while not acquired: | |
| check_cancelled(job.marker) | |
| acquired = self.gate.acquire(timeout=.1) | |
| check_cancelled(job.marker) | |
| yield | |
| finally: | |
| if acquired: | |
| self.gate.release() | |
| def cancel(self, request: gr.Request): | |
| owner = request.session_hash if request is not None else None | |
| if not owner: | |
| return gr.skip(), gr.skip() | |
| if self.jobs.cancel(owner): | |
| return message("Cancellation requested. Waiting for the current operation to stop.", busy=True), gr.update(interactive=False) | |
| return gr.skip(), gr.skip() | |
| def on_disconnect(self, request: gr.Request): | |
| if request is not None and request.session_hash: | |
| self.jobs.cancel(request.session_hash) | |
| def generate(self, text, picker, source, clip, consent, normalize, seed, randomize, | |
| temperature, top_p, top_k, repetition_penalty, rate, request: gr.Request, | |
| progress=gr.Progress(track_tqdm=False)): | |
| # Outputs: waveform/player data, status, main button, cancel button, seed. | |
| job = None | |
| final = None | |
| pending_error = None | |
| exported = None | |
| keep_export = False | |
| try: | |
| owner = request.session_hash if request is not None else None | |
| if not owner: | |
| raise InputError("Missing session. Reload the page or use gradio_client.Client.") | |
| request = validate(text, picker, source, seed, randomize, normalize, temperature, | |
| top_p, top_k, repetition_penalty, rate, self.voice_names, | |
| bool(clip), consent) | |
| job = self.jobs.begin(owner) | |
| yield (gr.skip(), message("Request received.", busy=True), gr.update(interactive=False, value="Working…"), | |
| gr.update(visible=True, interactive=True), gr.skip()) | |
| progress(None, desc="Preparing text") | |
| reference = read_reference(clip) if request.source == "Voice clip" else None | |
| check_cancelled(job.marker) | |
| prepared = self.engine.prepare(request, reference) | |
| del reference | |
| progress(None, desc="Waiting for an inference slot") | |
| with self.inference_slot(job): | |
| progress(None, desc="Requesting GPU") | |
| outcome = self.engine.run(prepared, job.marker, progress) | |
| # The decorated call has returned; its SDK release path completed. | |
| log.info("gpu_call_returned request=%s", job.id) | |
| del prepared | |
| check_cancelled(job.marker) | |
| if outcome.cancelled: | |
| raise Cancelled() | |
| yield (gr.skip(), message("Finalizing audio…", busy=True), gr.update(interactive=False, value="Finalizing…"), | |
| gr.update(visible=True, interactive=False), gr.skip()) | |
| progress(None, desc="Finalizing audio") | |
| check_cancelled(job.marker) | |
| exported, result = export_take(outcome.pcm, self.outputs, seed=request.seed, | |
| request_seconds=time.monotonic()-job.started, | |
| inference_seconds=outcome.seconds) | |
| result.update(text=request.text, words=[], alignment="unavailable") | |
| if self.aligner is not None: | |
| yield (gr.skip(), message("Synchronizing words…", busy=True), gr.update(interactive=False, value="Synchronizing…"), | |
| gr.update(visible=True, interactive=True), gr.skip()) | |
| try: | |
| from alignment import attach_timings | |
| progress(None, desc="Preparing word synchronization") | |
| alignment_inputs = self.aligner.prepare(outcome.pcm, request.text) | |
| check_cancelled(job.marker) | |
| with self.inference_slot(job): | |
| progress(None, desc="Requesting GPU for word synchronization") | |
| words = self.aligner.run(alignment_inputs, job.marker, progress) | |
| log.info("alignment_gpu_returned request=%s", job.id) | |
| check_cancelled(job.marker) | |
| result["words"] = attach_timings(request.text, words, result["duration"]) | |
| result["alignment"] = "ready" | |
| except Cancelled: | |
| raise | |
| except Exception as exc: | |
| check_cancelled(job.marker) | |
| log.warning("alignment_unavailable request=%s error_type=%s", job.id, type(exc).__name__) | |
| finally: | |
| alignment_inputs = None | |
| del outcome | |
| check_cancelled(job.marker) | |
| prune_outputs(self.outputs) | |
| result["request_seconds"] = round(time.monotonic()-job.started, 2) | |
| final = (json.dumps(result), message("Ready to play." if result["alignment"] == "ready" or self.aligner is None else "Audio ready. Word synchronization unavailable for this take."), | |
| gr.update(interactive=True, value="Generate speech"), | |
| gr.update(visible=False, interactive=True), request.seed) | |
| keep_export = True | |
| except AlreadyRunning: | |
| # Do not re-enable controls belonging to the already-running request. | |
| yield gr.skip(), message("A request is already running in this session.", busy=True), gr.skip(), gr.skip(), gr.skip() | |
| raise gr.Error("A request is already running in this session.") | |
| except Cancelled: | |
| final = (gr.skip(), message("Cancelled. Your previous audio is still available."), | |
| gr.update(interactive=True, value="Generate speech"), | |
| gr.update(visible=False, interactive=True), gr.skip()) | |
| except InputError as exc: | |
| pending_error = gr.Error(str(exc)) | |
| final = (gr.skip(), message(str(exc), error=True), gr.update(interactive=True, value="Generate speech"), | |
| gr.update(visible=False, interactive=True), gr.skip()) | |
| except Exception as exc: | |
| # Do not echo exception bodies that may include user text, audio or paths. | |
| log.error("request_failed request=%s error_type=%s", job.id if job else "validation", type(exc).__name__) | |
| text = "Generation failed. Your previous audio is still available. Try again or shorten the text." | |
| if job and job.marker.exists(): | |
| text = "Cancelled. Your previous audio is still available." | |
| elif isinstance(exc, gr.Error): | |
| # ZeroGPU owns quota/authentication errors and its native modal. | |
| text = "The GPU request could not complete. Check the Hugging Face notice, then try again." | |
| if not (job and job.marker.exists()): | |
| pending_error = exc if isinstance(exc, gr.Error) else gr.Error(text) | |
| final = (gr.skip(), message(text, error=True), gr.update(interactive=True, value="Generate speech"), | |
| gr.update(visible=False, interactive=True), gr.skip()) | |
| finally: | |
| if exported is not None and not keep_export: | |
| exported.unlink(missing_ok=True) | |
| if job is not None: | |
| self.jobs.finish(job) | |
| # Finishing happens BEFORE the final yield, even if the consumer never | |
| # resumes or explicitly closes the generator after receiving the result. | |
| if final is not None: | |
| yield final | |
| if pending_error is not None: | |
| # Preserve the native error notice and an API failure, after controls | |
| # have been restored and the previous audio has been preserved. | |
| raise pending_error | |