velixio
1f03f5632c
Rendering 1126 previews ran one clip at a time, and only the first of a clip's five stages is on the GPU: render, then watermark embed, MP3 encode, decode, and detection. The card idled through four CPU stages per clip. The embed and detection are neural forward passes that ran inline on the event loop, so they held it for the whole clip -- concurrency would have queued behind a busy loop and bought nothing. They now go through asyncio.to_thread, which is what makes threads the right tool here: torch releases the GIL inside those passes, so there is no second model copy and no IPC for the tensors. Clips then build --jobs at a time (4 by default, 1 restores the old serial behaviour) under a semaphore, because every clip in flight holds decoded audio. A lost watermark still stops the entire run rather than only its own clip. --resume now also adopts MP3s already on disk. The manifest is written once, at the end, so a run interrupted at clip 900 left 900 correct files that --resume could not see and re-rendered every one of them. Everything an entry needs -- sha256, byte length, duration, featured flag -- is recoverable from the file and the catalog, so recover it. Also add a watermark preflight. Every clip was already verified individually, but only after the first full render, and the message blamed the bitrate when the cause can be unrelated to audio: on a host without python3-dev, AudioSeal's forward pass dies inside Inductor, embed_watermark catches it, and the clip is returned unmarked. Two seconds up front, with the actual cause named. It also warms the lazy generator/detector globals single-threaded, before --jobs fans out.
411 lines
17 KiB
Python
411 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""Render the voice gallery: every archetype preview, once, as a signed bundle.
|
|
|
|
The app renders archetype previews on the GPU on first request, which means a
|
|
fresh install hears nothing until the 2.4 GB TTS checkpoint has downloaded, and
|
|
every user pays a cold model load per voice. This script does that work once, on
|
|
a machine that already has the weights, and produces the artifacts
|
|
``backend/services/gallery.py`` downloads:
|
|
|
|
<out>/manifest.json schema + per-key filename/sha256/bytes/duration
|
|
<out>/previews/<key>.mp3 64 kbps mono, one per distinct preview key
|
|
<out>/featured.tar.gz the 51 featured previews, as one request
|
|
|
|
**Publishing is a manual owner step.** This script only builds the directory and
|
|
prints the two commands that follow: signing ``manifest.json`` with the existing
|
|
Tauri release key (the client verifies against the pubkey already baked into the
|
|
binary — no second trust root) and uploading the result. Nothing here talks to
|
|
GitHub.
|
|
|
|
MP3 at 64 kbps mono is a bytes decision, not a quality one: 1126 previews of a
|
|
sample script are ~110 MB as WAV and ~9 MB as MP3, and the featured tarball a
|
|
first run pulls is ~450 kB. The app already plays user-dropped MP3s.
|
|
|
|
Provenance: every clip is passed through ``mark_synthetic(force=True)`` at the
|
|
tensor stage before encoding — ``force`` because publication must not depend on
|
|
the *publisher's* watermark preference — and detection is then re-run **on the
|
|
decoded MP3**, so a bitrate that destroys the watermark fails the build instead
|
|
of shipping unmarked audio. That check is why the encode settings live here and
|
|
not in a shell one-liner.
|
|
|
|
Usage (from the repo root, with the model cached):
|
|
|
|
python3 scripts/render_gallery.py --out dist/gallery
|
|
python3 scripts/render_gallery.py --out dist/gallery --featured-only
|
|
|
|
Only the render step is on the GPU; marking, encoding, decoding and detection
|
|
are all CPU, so clips are built ``--jobs`` at a time (4 by default) and the card
|
|
does not sit idle through four stages per clip. ``--resume`` picks up whatever
|
|
is already in the output directory, including MP3s from a run that was
|
|
interrupted before it could write a manifest.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import hashlib
|
|
import gzip
|
|
import json
|
|
import sys
|
|
import tarfile
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parent.parent
|
|
BACKEND_DIR = REPO_ROOT / "backend"
|
|
if str(BACKEND_DIR) not in sys.path:
|
|
sys.path.insert(0, str(BACKEND_DIR))
|
|
|
|
#: Must match ``services.gallery.SCHEMA_VERSION`` — the client refuses anything
|
|
#: else rather than guessing at an unknown layout.
|
|
SCHEMA_VERSION = 1
|
|
|
|
MP3_BITRATE = "64k"
|
|
MP3_CHANNELS = 1
|
|
|
|
|
|
def _sha256(data: bytes) -> str:
|
|
return hashlib.sha256(data).hexdigest()
|
|
|
|
|
|
async def _encode_mp3(wav_path: Path, mp3_path: Path) -> None:
|
|
from services.ffmpeg_utils import find_ffmpeg, run_ffmpeg
|
|
|
|
cmd = [
|
|
find_ffmpeg(), "-y", "-hide_banner", "-loglevel", "error",
|
|
"-i", str(wav_path),
|
|
"-ac", str(MP3_CHANNELS), "-b:a", MP3_BITRATE,
|
|
# No metadata: the previews are identical for every user, and an
|
|
# encoder/date tag would make otherwise-identical bytes differ per run.
|
|
"-map_metadata", "-1",
|
|
str(mp3_path),
|
|
]
|
|
rc, _, err = await run_ffmpeg(cmd, timeout=120.0)
|
|
if rc != 0 or not mp3_path.is_file():
|
|
raise RuntimeError(f"mp3 encode failed: {err.decode('utf-8', 'replace')[:400]}")
|
|
|
|
|
|
async def _decode_wav(mp3_path: Path, wav_path: Path) -> None:
|
|
"""Decode back through ffmpeg so detection sees exactly the published bytes."""
|
|
from services.ffmpeg_utils import find_ffmpeg, run_ffmpeg
|
|
|
|
cmd = [
|
|
find_ffmpeg(), "-y", "-hide_banner", "-loglevel", "error",
|
|
"-i", str(mp3_path), "-c:a", "pcm_s16le", str(wav_path),
|
|
]
|
|
rc, _, err = await run_ffmpeg(cmd, timeout=120.0)
|
|
if rc != 0 or not wav_path.is_file():
|
|
raise RuntimeError(f"mp3 decode failed: {err.decode('utf-8', 'replace')[:400]}")
|
|
|
|
|
|
def _load(path: Path):
|
|
import torchaudio
|
|
|
|
wav, sr = torchaudio.load(str(path))
|
|
return wav, int(sr)
|
|
|
|
|
|
async def _build_one(archetype: dict, key: str, work: Path, out_previews: Path) -> dict:
|
|
"""Render, mark, encode, verify — returns the manifest entry for *key*."""
|
|
from api.routers.archetypes import _render_archetype_wav
|
|
from services.watermark import detect_watermark, mark_synthetic
|
|
|
|
raw_wav = work / f"{key}.raw.wav"
|
|
marked_wav = work / f"{key}.marked.wav"
|
|
mp3_path = out_previews / f"{key}.mp3"
|
|
|
|
await _render_archetype_wav(archetype, raw_wav)
|
|
|
|
# Everything below that is CPU-bound goes through asyncio.to_thread. The
|
|
# AudioSeal embed and detection are each a real neural forward pass, and run
|
|
# inline they hold the event loop for the whole clip — so --jobs above 1
|
|
# would queue work behind a busy loop and buy nothing. torch releases the
|
|
# GIL inside those passes, which is what makes threads (rather than
|
|
# processes) the right tool: no second model copy, no IPC for the tensors.
|
|
wav, sr = await asyncio.to_thread(_load, raw_wav)
|
|
|
|
# force=True: the published clip carries the mark regardless of whether the
|
|
# machine doing the publishing has invisible watermarking switched on. Same
|
|
# contract as persona_bundle's preview embed.
|
|
marked = await asyncio.to_thread(
|
|
mark_synthetic, wav, sr, force=True, context="gallery.publish"
|
|
)
|
|
from api.routers.generation import _safe_torchaudio_save
|
|
|
|
await asyncio.to_thread(_safe_torchaudio_save, str(marked_wav), marked, sr)
|
|
await _encode_mp3(marked_wav, mp3_path)
|
|
|
|
check_wav = work / f"{key}.check.wav"
|
|
await _decode_wav(mp3_path, check_wav)
|
|
decoded, decoded_sr = await asyncio.to_thread(_load, check_wav)
|
|
verdict = await asyncio.to_thread(detect_watermark, decoded, decoded_sr)
|
|
if not verdict.get("is_watermarked"):
|
|
mp3_path.unlink(missing_ok=True)
|
|
raise AssertionError(
|
|
f"watermark did not survive {MP3_BITRATE} mono encoding for {key} "
|
|
f"(confidence {verdict.get('confidence')}, {verdict.get('error', '')}) — "
|
|
"raise the bitrate or fix the embed before publishing"
|
|
)
|
|
|
|
data = mp3_path.read_bytes()
|
|
duration = round(decoded.shape[-1] / max(decoded_sr, 1), 3)
|
|
for scratch in (raw_wav, marked_wav, check_wav):
|
|
scratch.unlink(missing_ok=True)
|
|
return {
|
|
"filename": mp3_path.name,
|
|
"sha256": _sha256(data),
|
|
"bytes": len(data),
|
|
"duration": duration,
|
|
"featured": bool(archetype.get("is_featured")),
|
|
}
|
|
|
|
|
|
async def _preflight_watermark() -> None:
|
|
"""Prove the watermark works before rendering a thousand clips.
|
|
|
|
Every clip is verified individually, so a broken embed was always caught —
|
|
but only after the first full render, and the failure named the bitrate
|
|
("raise the bitrate or fix the embed") when the real cause can be nothing to
|
|
do with audio at all. On a machine missing ``python3-dev``, AudioSeal's
|
|
forward pass dies inside Inductor (``Python.h: No such file``),
|
|
``embed_watermark`` catches it, and the clip is returned *unmarked*. Five
|
|
seconds here beats discovering that at clip 1 of 1126.
|
|
|
|
Doubles as a single-threaded warm-up: the generator and detector are lazy
|
|
module globals, so touching them once before --jobs fans out avoids several
|
|
threads racing to load the same model.
|
|
"""
|
|
import torch
|
|
from services.watermark import detect_watermark, mark_synthetic
|
|
|
|
sample_rate = 24000
|
|
tone = torch.sin(
|
|
2 * 3.14159 * 220 * torch.arange(sample_rate * 2) / sample_rate
|
|
).unsqueeze(0) * 0.3
|
|
marked = await asyncio.to_thread(
|
|
mark_synthetic, tone, sample_rate, force=True, context="gallery.preflight"
|
|
)
|
|
verdict = await asyncio.to_thread(detect_watermark, marked, sample_rate)
|
|
if not verdict.get("is_watermarked"):
|
|
raise SystemExit(
|
|
"watermark preflight failed: mark_synthetic returned audio the "
|
|
f"detector does not recognise (confidence {verdict.get('confidence')}). "
|
|
"Publishing would ship unmarked audio, so this build stops here.\n"
|
|
"Most common cause: torch.compile/Inductor cannot build its helper "
|
|
"(missing Python headers — install python3-dev), which makes the "
|
|
"embed raise and silently pass the audio through unchanged. "
|
|
"TORCHDYNAMO_DISABLE=1 is the quick workaround."
|
|
)
|
|
|
|
|
|
def _resume_from_disk(out_previews: Path, by_key: dict) -> dict:
|
|
"""Rebuild manifest entries for previews already rendered.
|
|
|
|
The manifest is written once, at the end, so a run interrupted at clip 900
|
|
leaves 900 perfectly good MP3s that ``--resume`` cannot see — it keys off
|
|
the manifest, so it would render every one of them again. Everything an
|
|
entry needs is recoverable from the file itself, so recover it.
|
|
"""
|
|
recovered: dict = {}
|
|
for mp3_path in sorted(out_previews.glob("*.mp3")):
|
|
key = mp3_path.stem
|
|
archetype = by_key.get(key)
|
|
if archetype is None:
|
|
continue # a key from some older catalog — leave it out of the index
|
|
data = mp3_path.read_bytes()
|
|
if not data:
|
|
mp3_path.unlink(missing_ok=True)
|
|
continue
|
|
recovered[key] = {
|
|
"filename": mp3_path.name,
|
|
"sha256": _sha256(data),
|
|
"bytes": len(data),
|
|
"duration": _probe_duration(mp3_path),
|
|
"featured": bool(archetype.get("is_featured")),
|
|
}
|
|
return recovered
|
|
|
|
|
|
def _probe_duration(path: Path) -> float:
|
|
"""Duration in seconds, straight from the encoded file."""
|
|
import subprocess
|
|
|
|
from services.ffmpeg_utils import find_ffmpeg
|
|
|
|
ffprobe = str(Path(find_ffmpeg()).with_name("ffprobe"))
|
|
try:
|
|
out = subprocess.run(
|
|
[ffprobe, "-v", "error", "-show_entries", "format=duration",
|
|
"-of", "default=nw=1:nk=1", str(path)],
|
|
capture_output=True, text=True, timeout=30,
|
|
).stdout.strip()
|
|
return round(float(out), 3)
|
|
except (OSError, ValueError, subprocess.SubprocessError):
|
|
return 0.0
|
|
|
|
|
|
def _write_featured_tarball(out: Path, previews: dict) -> dict:
|
|
"""Bundle the featured previews so a first run costs one request, not 51."""
|
|
featured = sorted(k for k, e in previews.items() if e["featured"])
|
|
archive = out / "featured.tar.gz"
|
|
# mtime/uid/gid pinned so re-running with unchanged audio produces the same
|
|
# archive bytes — the client diffs on sha256, and a timestamp would make
|
|
# every rebuild look like a change worth re-downloading.
|
|
with archive.open("wb") as raw, gzip.GzipFile(
|
|
filename="", mode="wb", fileobj=raw, mtime=0
|
|
) as compressed, tarfile.open(fileobj=compressed, mode="w") as tar:
|
|
for key in featured:
|
|
path = out / "previews" / f"{key}.mp3"
|
|
info = tar.gettarinfo(str(path), arcname=f"previews/{key}.mp3")
|
|
info.mtime = 0
|
|
info.uid = info.gid = 0
|
|
info.uname = info.gname = ""
|
|
with path.open("rb") as fh:
|
|
tar.addfile(info, fh)
|
|
data = archive.read_bytes()
|
|
return {
|
|
"filename": archive.name,
|
|
"sha256": _sha256(data),
|
|
"bytes": len(data),
|
|
"count": len(featured),
|
|
}
|
|
|
|
|
|
async def _main(args: argparse.Namespace) -> int:
|
|
from core import archetypes as catalog
|
|
from core.version import APP_VERSION
|
|
from api.routers.archetypes import _preview_key
|
|
from services.tts_backend import active_backend_id
|
|
|
|
out = Path(args.out).resolve()
|
|
out_previews = out / "previews"
|
|
out_previews.mkdir(parents=True, exist_ok=True)
|
|
|
|
items = catalog.list_archetypes()
|
|
if args.featured_only:
|
|
items = [a for a in items if a["is_featured"]]
|
|
# One clip per distinct key: archetypes that resolve to the same
|
|
# (instruct, language) share a preview, exactly as the app's cache does.
|
|
by_key: dict[str, dict] = {}
|
|
for a in items:
|
|
by_key.setdefault(_preview_key(a), a)
|
|
keys = sorted(by_key)
|
|
if args.limit:
|
|
keys = keys[: args.limit]
|
|
|
|
previews: dict[str, dict] = {}
|
|
manifest_path = out / "manifest.json"
|
|
if args.resume:
|
|
if manifest_path.is_file():
|
|
previous = json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
previews = {
|
|
k: e for k, e in (previous.get("previews") or {}).items()
|
|
if (out_previews / f"{k}.mp3").is_file()
|
|
}
|
|
# Also adopt clips on disk the manifest never got to describe — an
|
|
# interrupted run has no manifest at all, and re-rendering audio that
|
|
# is already correct is the most expensive way to do nothing.
|
|
for key, entry in _resume_from_disk(out_previews, by_key).items():
|
|
previews.setdefault(key, entry)
|
|
if previews:
|
|
print(f"resuming: {len(previews)} preview(s) already rendered", flush=True)
|
|
|
|
await _preflight_watermark()
|
|
|
|
pending = [k for k in keys if k not in previews]
|
|
failures: list[str] = []
|
|
with tempfile.TemporaryDirectory(prefix="gallery-render-") as tmp:
|
|
work = Path(tmp)
|
|
# Bounded fan-out. Each clip is render → embed → encode → decode →
|
|
# detect, and only the first of those is on the GPU: with one clip in
|
|
# flight the card idles through four CPU stages. The cap keeps that
|
|
# overlap from turning into unbounded memory (every concurrent clip
|
|
# holds decoded audio) and matches how the app itself bounds GPU work.
|
|
limit = asyncio.Semaphore(max(1, args.jobs))
|
|
completed = 0
|
|
state = asyncio.Lock()
|
|
|
|
async def build(key: str) -> None:
|
|
nonlocal completed
|
|
archetype = by_key[key]
|
|
async with limit:
|
|
entry = await _build_one(archetype, key, work, out_previews)
|
|
async with state:
|
|
previews[key] = entry
|
|
completed += 1
|
|
print(f"[{completed}/{len(pending)}] {key} {archetype['name']}", flush=True)
|
|
|
|
tasks = [asyncio.create_task(build(key), name=key) for key in pending]
|
|
try:
|
|
for task in asyncio.as_completed(tasks):
|
|
try:
|
|
await task
|
|
except AssertionError:
|
|
# A lost watermark is a build failure, not a bad voice —
|
|
# stop the whole run rather than let the remaining jobs
|
|
# keep writing clips nobody has verified.
|
|
for other in tasks:
|
|
other.cancel()
|
|
raise
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as exc:
|
|
failures.append(f"{type(exc).__name__}: {exc}")
|
|
print(f" FAILED: {exc}", file=sys.stderr, flush=True)
|
|
finally:
|
|
for task in tasks:
|
|
task.cancel()
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
if not previews:
|
|
print("nothing rendered", file=sys.stderr)
|
|
return 1
|
|
|
|
manifest = {
|
|
"schema": SCHEMA_VERSION,
|
|
"generated_at": int(time.time()),
|
|
"engine": active_backend_id(),
|
|
"engine_version": APP_VERSION,
|
|
"format": {"codec": "mp3", "bitrate": MP3_BITRATE, "channels": MP3_CHANNELS},
|
|
"featured": _write_featured_tarball(out, previews),
|
|
"previews": dict(sorted(previews.items())),
|
|
}
|
|
manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n",
|
|
encoding="utf-8")
|
|
|
|
print(f"\n{len(previews)} previews → {out}")
|
|
if failures:
|
|
print(f"{len(failures)} archetype(s) did not render:", file=sys.stderr)
|
|
for line in failures:
|
|
print(f" {line}", file=sys.stderr)
|
|
print(
|
|
"\nNext (manual, owner):\n"
|
|
f" minisign -Sm {manifest_path} -s <tauri-release.key> "
|
|
f"-x {manifest_path}.minisig\n"
|
|
" gh release create gallery-v1 --repo debpalash/omnivoice-gallery "
|
|
f"{manifest_path} {manifest_path}.minisig {out / 'featured.tar.gz'} "
|
|
f"{out_previews}/*.mp3"
|
|
)
|
|
return 1 if failures else 0
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
|
|
parser.add_argument("--out", default="dist/gallery", help="output directory")
|
|
parser.add_argument("--featured-only", action="store_true",
|
|
help="render only the 51 featured archetypes")
|
|
parser.add_argument("--limit", type=int, default=0, help="stop after N keys")
|
|
parser.add_argument(
|
|
"--jobs", type=int, default=4,
|
|
help="clips built concurrently (default 4); 1 restores serial rendering",
|
|
)
|
|
parser.add_argument("--resume", action="store_true",
|
|
help="keep previews already in <out> (manifest entries "
|
|
"and any MP3s an interrupted run left behind)")
|
|
return asyncio.run(_main(parser.parse_args()))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|