Spaces:
Running on Zero
Running on Zero
File size: 8,949 Bytes
6d1acb0 255bed9 6d1acb0 1b2da16 6d1acb0 1b2da16 6d1acb0 255bed9 6d1acb0 255bed9 6d1acb0 255bed9 6d1acb0 1b2da16 6d1acb0 1b2da16 6d1acb0 255bed9 6d1acb0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | """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)
@contextmanager
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
|