初始提交：3 款纯音频游戏 + 共享引擎 + 完整文档与资产管线
- 共享引擎（app/src/main）：GestureEngine 6 手势识别、AudioEngine 三层音频 （ExoPlayer BGM 交叉淡入 / 对白 ducking / SoundPool 低延迟）、Game/GameEngine 状态机接口、MainActivity 全屏封面 + TalkBack - 《声命线》storyline：末世互动广播剧，3 幕 8 抉择点 4 结局 - 《声东击西》reflex：听声辨位反应，12 波战斗 + BGM 三档 BPM 变速 - 《暗室》escape：3x3 网格密室逃脱，钥匙→铁柜→保险丝→电箱→大门谜题链 - 文档：setup（从零搭建）/ architecture（架构）/ engine-api（引擎 API）/ asset-pipeline（资产流水线）/ games/*（三款游戏设计文档） - 资产管线：Qwen3-TTS 对白 + 程序合成音效/BGM + Mureka BGM 的 gen 脚本与台账， 全部 ogg 资产入库可复现 - env.d 环境脚本：android/tts 入库；sfx/mureka 含密钥不入库，仅提交 .example.sh 模板
This commit is contained in:
SoundGame Dev
@@ -0,0 +1,724 @@
|
||||
"""SA3 TRT — gradio web UI (debug build).
|
||||
|
||||
Adds (vs the MVP):
|
||||
- Model picker: sm-music / sm-sfx / medium (hot-swap with cached SA3Inference
|
||||
instances; first switch ~5-7s, subsequent instant)
|
||||
- Engine-variant picker per model: auto-detects all `dit_*.trt` files in the
|
||||
model dir (canonical fixed, archived buggy, fp32, etc.) so you can A/B
|
||||
different precisions / quantizations without restarting
|
||||
- Spectrogram display: Underfit-style 3-band tinted stereo mel spectrogram
|
||||
rendered inline alongside the audio
|
||||
|
||||
Launch:
|
||||
./sa3-gradio # share=True by default, sm-music + same-s
|
||||
./sa3-gradio --dit medium
|
||||
./sa3-gradio --no-share # local-only
|
||||
|
||||
The previously-required `--dit ...` is now just the *initial* model — the
|
||||
runtime dropdown lets you switch between variants/models without restart.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import argparse
|
||||
import base64
|
||||
import math
|
||||
import sys
|
||||
import html as html_lib
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
SCRIPTS_DIR = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(SCRIPTS_DIR))
|
||||
|
||||
from sa3_trt import SA3Inference, SAMPLE_RATE, SAMPLES_PER_LATENT # noqa: E402
|
||||
import sa3_trt_core as canon # noqa: E402
|
||||
from spec import render_spectrogram_png # noqa: E402
|
||||
|
||||
|
||||
OUTPUT_DIR = SCRIPTS_DIR.parent / "output" / "gradio"
|
||||
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
KEEP_RECENT_N = 20
|
||||
|
||||
from gradio_ui import ( # backend-agnostic UI layer, forked from optimized/mlx
|
||||
# The other _JS_* constants (promote / pause-others / scroll / playhead) are wired
|
||||
# inside render_player and render_history themselves, so they stay private to that
|
||||
# module; only the two this file attaches to its own events are imported.
|
||||
_JS_FIX_SLIDERS, _save_wav, render_history, render_player, render_queue_status, verbose_basename,
|
||||
)
|
||||
|
||||
MODELS_ROOT = SCRIPTS_DIR.parent / "models" / canon.ARCH
|
||||
|
||||
|
||||
# Default decoder for each DiT bundle.
|
||||
DEFAULT_DECODERS = {"sm-music": "same-s", "sm-sfx": "same-s", "medium": "same-l"}
|
||||
|
||||
# DiT-name → on-disk subdir mapping (matches sa3_trt_core's DIT_CHOICES)
|
||||
DIT_SUBDIRS = {"sm-music": "sa3-sm-music", "sm-sfx": "sa3-sm-sfx", "medium": "sa3-m"}
|
||||
|
||||
|
||||
# Magic marker (matched against str(variant_path)) telling generate() to
|
||||
# dispatch to the PyTorch-eager backend instead of TRT. Only offered for
|
||||
# the medium+SAME-L+fp32 combo per the user's scope.
|
||||
PT_EAGER_VARIANT = "<PT_EAGER_FP32_GT>"
|
||||
|
||||
|
||||
def discover_variants(dit_name: str) -> list[tuple[str, Path]]:
|
||||
"""Return [(label, path)] of available DiT engine files for this model.
|
||||
|
||||
Scans models/<arch>/<dit_subdir>/dit_*.trt. The canonical engine
|
||||
(dit_fp16.trt) is always first if present; other variants follow
|
||||
alphabetically. For the medium DiT also appends a pseudo-variant
|
||||
"pytorch fp32 (GT)" that dispatches to the PyTorch-eager backend.
|
||||
"""
|
||||
subdir = DIT_SUBDIRS.get(dit_name)
|
||||
if subdir is None:
|
||||
return []
|
||||
d = MODELS_ROOT / subdir
|
||||
if not d.exists():
|
||||
return []
|
||||
files = [f for f in sorted(d.glob("dit_*.trt")) if f.name not in RETIRED_DITS]
|
||||
canonical_name = "dit_fp16.trt"
|
||||
canonical = [f for f in files if f.name == canonical_name]
|
||||
others = [f for f in files if f.name != canonical_name]
|
||||
out = []
|
||||
for f in canonical + others:
|
||||
label = f.name[len("dit_"):-len(".trt")] if f.name.startswith("dit_") and f.name.endswith(".trt") else f.name
|
||||
if f.name == canonical_name:
|
||||
label = "fp16 (canonical)"
|
||||
elif "buggy" in f.name:
|
||||
label = label + " ← old, broken"
|
||||
out.append((label, f))
|
||||
# Pseudo-variant: PT FP32 GT (currently medium only).
|
||||
if dit_name == "medium":
|
||||
out.append(("pytorch fp32 GT (slow, vanilla eager)", Path(PT_EAGER_VARIANT)))
|
||||
return out
|
||||
|
||||
|
||||
# Pre-chunkable engines. Removed from the model repo's download path and hidden from the picker.
|
||||
RETIRED_DECODERS = {"dec_dynamic_bf16.trt", "dec_dynamic_triton_swa.trt", "dec_dynamic_fp32.trt",
|
||||
"dec_fp8.trt", "dec_fp8_fast.trt", "dec_dynamic_fp16mixed.trt"}
|
||||
# Experimental DiT engines that are not part of the shipped set. "fp16mixed" is the retired name
|
||||
# for plain fp16 -- every tier is mixed-precision, so the qualifier only ever meant "old".
|
||||
RETIRED_DITS = {"dit_fp16mixed.trt", "dit_fp16_bands4096_32768.trt", "dit_fp16_maxL32768.trt"}
|
||||
|
||||
|
||||
def _precision_of(filename: str) -> str:
|
||||
"""Pull the precision token out of an engine filename: dec_fp16_chunkable_limiter -> fp16."""
|
||||
for tok in ("fp32", "fp16", "bf16", "fp8"):
|
||||
if f"_{tok}" in filename:
|
||||
return tok
|
||||
return ""
|
||||
|
||||
|
||||
def discover_decoder_variants(decoder_name: str) -> list[tuple[str, Path]]:
|
||||
"""Return [(label, path)] of decoder quantization tiers for this decoder —
|
||||
canonical first. Uses the known tier set (canonical / fp8; see quantize/README.md); engines
|
||||
auto-download from HF on selection if missing. Any extra local dec_*.trt not in the known set
|
||||
are appended, so a self-built variant shows up — except the RETIRED pre-chunkable engines,
|
||||
which are filtered out. A leftover copy of one of those in models/ would otherwise appear in
|
||||
the picker and silently serve pre-limiter audio with a 5.7-8.1 GB scratch reservation.
|
||||
"""
|
||||
out, known = [], set(RETIRED_DECODERS)
|
||||
for tier, fname in canon.DECODER_TIER_FILENAME.get(decoder_name, {}).items():
|
||||
# Name the tier by the precision it actually is -- SAME-L canonical is fp16, SAME-S is
|
||||
# bf16 -- rather than rendering "canonical (canonical)".
|
||||
prec = _precision_of(fname)
|
||||
label = f"{prec} (canonical)" if tier == "canonical" else (prec or tier)
|
||||
out.append((label, canon.ARCH_DIR / decoder_name / fname))
|
||||
known.add(fname)
|
||||
d = canon.ARCH_DIR / decoder_name
|
||||
if d.exists():
|
||||
for f in sorted(d.glob("dec_*.trt")):
|
||||
if f.name in known:
|
||||
continue
|
||||
stem = f.name
|
||||
label = (stem[len("dec_dynamic_"):-len(".trt")] if stem.startswith("dec_dynamic_")
|
||||
else stem[len("dec_"):-len(".trt")])
|
||||
out.append((label, f))
|
||||
return out
|
||||
|
||||
|
||||
def _prune_old_outputs():
|
||||
wavs = sorted(OUTPUT_DIR.glob("*.wav"), key=lambda p: p.stat().st_mtime, reverse=True)
|
||||
for old in wavs[KEEP_RECENT_N:]:
|
||||
try:
|
||||
old.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
# ── SA3Inference cache (model+variant → instance) ──────────────────────────
|
||||
# Each cached entry holds 5-15 GB of VRAM (engines + persistent graph buffers).
|
||||
# H100 80GB fits ~5 sm-music or ~3 medium combos before OOM. We LRU-evict to
|
||||
# stay under that limit. Switching back to an evicted variant re-loads (~5-7s).
|
||||
_inference_cache: dict[tuple[str, str, str, str, bool], SA3Inference] = {}
|
||||
_inference_lru: list[tuple[str, str, str, str]] = [] # MRU at end
|
||||
_INFERENCE_CACHE_MAX = 2
|
||||
|
||||
|
||||
def _mb(b):
|
||||
return f"{b / 2**20:,.0f} MB" if b else "—"
|
||||
|
||||
|
||||
def render_debug(inf, t_fused, t_stage):
|
||||
"""Per-stage ms + VRAM, and what the fused mega-graph costs by comparison.
|
||||
|
||||
⚠ The stage split is measured on the UNFUSED path, because a captured CUDA graph
|
||||
cannot be instrumented from inside: torch rejects Event.record() during capture, and
|
||||
even cudaEventRecordWithFlags(..., cudaEventRecordExternal) is accepted only to fail
|
||||
at capture_end(). Adding timers would mean splitting the mega-graph into per-stage
|
||||
graphs, which would change the very thing being measured. So debug mode times the
|
||||
same engines run stage-by-stage and reports the fused total beside it — the gap
|
||||
between the two IS the fusion win.
|
||||
|
||||
VRAM is split the same way it is actually paid: engine weights and scratch are
|
||||
resident for the whole process, while `transient` is what a stage allocates on top.
|
||||
"""
|
||||
fp = inf.engine_footprint()
|
||||
sv = t_stage.get("stage_vram", {})
|
||||
rows = [("encoder", t_stage.get("encode_ms", 0.0), sv.get("encode"), fp.get("enc")),
|
||||
("T5Gemma", t_stage.get("t5_ms", 0.0), sv.get("t5"), fp.get("t5")),
|
||||
(f"DiT x{t_stage.get('steps', 0)}",
|
||||
t_stage.get("sampling_ms", 0.0), sv.get("dit"), fp.get("dit")),
|
||||
("decoder", t_stage.get("decode_ms", 0.0), sv.get("decode"), fp.get("dec"))]
|
||||
th = ("padding:3px 10px 3px 0; text-align:right; font-weight:600; "
|
||||
"border-bottom:1px solid rgba(127,127,127,0.3)")
|
||||
td = "padding:2px 10px 2px 0; text-align:right; font-variant-numeric:tabular-nums"
|
||||
out = ['<div style="font-size:0.82em; margin-top:6px">',
|
||||
'<table style="border-collapse:collapse">',
|
||||
f'<tr><th style="{th}; text-align:left">stage</th><th style="{th}">ms</th>'
|
||||
f'<th style="{th}">transient</th><th style="{th}">weights</th>'
|
||||
f'<th style="{th}">scratch</th></tr>']
|
||||
for name, ms, tv, f in rows:
|
||||
if f is None and not ms:
|
||||
continue # encoder absent unless a2a/inpaint ran
|
||||
sc = _mb(f["scratch"]) + (" (shared)" if f and f["shared"] else "") if f else "—"
|
||||
out.append(f'<tr><td style="{td}; text-align:left">{html_lib.escape(name)}</td>'
|
||||
f'<td style="{td}">{ms:,.1f}</td><td style="{td}">{_mb(tv)}</td>'
|
||||
f'<td style="{td}">{_mb(f["weights"]) if f else "—"}</td>'
|
||||
f'<td style="{td}">{sc}</td></tr>')
|
||||
# Whatever the four stages do not account for: schedule build, noise, the int16 cast
|
||||
# and the GPU->CPU copy. Shown rather than hidden, so the column sums to the total.
|
||||
unfused = t_stage.get("inference_ms", 0.0)
|
||||
other = unfused - sum(r[1] for r in rows)
|
||||
if other > 0.5:
|
||||
out.append(f'<tr><td style="{td}; text-align:left; opacity:0.7">other '
|
||||
f'<span style="opacity:0.7">(schedule, noise, int16 cast, GPU→CPU)</span>'
|
||||
f'</td><td style="{td}; opacity:0.7">{other:,.1f}</td>'
|
||||
f'<td style="{td}">—</td><td style="{td}">—</td>'
|
||||
f'<td style="{td}">—</td></tr>')
|
||||
out.append("</table>")
|
||||
shared = fp.get("_shared_scratch", 0)
|
||||
note = []
|
||||
if t_fused.get("eager"):
|
||||
note.append("this render used the eager path (CFG / audio-to-audio / inpaint), "
|
||||
"so the breakdown IS the render — no second pass was run")
|
||||
else:
|
||||
fused = t_fused.get("inference_ms", 0.0)
|
||||
gain = f"{unfused / fused:.2f}x faster" if fused > 0 else ""
|
||||
note.append(f"unfused {unfused:,.1f} ms · "
|
||||
f"<b>fused mega-graph {fused:,.1f} ms</b> ({gain})")
|
||||
note.append(f"shared scratch {_mb(shared)} — one buffer for every engine above, so the "
|
||||
f"scratch column does not sum")
|
||||
note.append("a captured graph allocates nothing at replay, so the mega-graph's peak "
|
||||
"VRAM equals its resident footprint")
|
||||
out += [f'<div style="color:#888; margin-top:4px">{n}</div>' for n in note]
|
||||
out.append("</div>")
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def _evict_inference(key):
|
||||
"""Free the engines + caches held by a cached SA3Inference before dropping it.
|
||||
|
||||
The teardown ORDER is load-bearing and lives in SA3Inference.close(): captured CUDA
|
||||
graphs and the eager contexts must go before the runners whose contexts and shared
|
||||
scratch they point into. This used to free the runners first and clear the graphs
|
||||
after, which poisoned the CUDA context — the next render in a *different* instance
|
||||
died with `an illegal memory access was encountered`, so switching models past the
|
||||
cache limit bricked the server until restart.
|
||||
"""
|
||||
inf = _inference_cache.pop(key, None)
|
||||
if inf is None:
|
||||
return
|
||||
try:
|
||||
inf.close()
|
||||
except Exception as e:
|
||||
print(f" (eviction cleanup raised {type(e).__name__}: {e})")
|
||||
del inf
|
||||
import gc
|
||||
gc.collect()
|
||||
|
||||
|
||||
# Set once from --chunking in main(). A module-level default keeps it out of every call site
|
||||
# and out of the UI, which is right: it is a deployment choice (memory vs speed), not a per-render
|
||||
# one -- switching it means selecting a different optimization profile, so the engines reload.
|
||||
_CHUNKING = True
|
||||
|
||||
|
||||
def get_inference(dit: str, decoder: str, dit_variant_path: str,
|
||||
dec_variant_path: str,
|
||||
default_T_lat: int, default_steps: int,
|
||||
default_seconds: float, quiet: bool,
|
||||
chunking: bool | None = None) -> SA3Inference:
|
||||
"""Return a warm SA3Inference, building (and caching) if not yet loaded.
|
||||
|
||||
The cache key is (dit, decoder, dit_variant_path, dec_variant_path) so
|
||||
swapping either the DiT *or* the decoder variant triggers a fresh load.
|
||||
LRU eviction caps the cache at _INFERENCE_CACHE_MAX entries.
|
||||
"""
|
||||
chunking = _CHUNKING if chunking is None else chunking
|
||||
# chunking is part of the key: it selects the profile, so the two modes are different
|
||||
# contexts with different scratch and cannot share a cached SA3Inference.
|
||||
key = (dit, decoder, dit_variant_path, dec_variant_path, chunking)
|
||||
if key in _inference_cache:
|
||||
if key in _inference_lru:
|
||||
_inference_lru.remove(key)
|
||||
_inference_lru.append(key)
|
||||
return _inference_cache[key]
|
||||
|
||||
# LRU evict BEFORE loading new — frees VRAM so the new load doesn't OOM.
|
||||
while len(_inference_cache) >= _INFERENCE_CACHE_MAX:
|
||||
oldest = _inference_lru.pop(0)
|
||||
print(f"\n ← LRU-evicting SA3Inference{oldest[:2]} "
|
||||
f"dit={Path(oldest[2]).name} dec={Path(oldest[3]).name}")
|
||||
_evict_inference(oldest)
|
||||
|
||||
# Override the canonical engine path lookups before constructing.
|
||||
canon.DIT_CHOICES[dit]["engine"] = Path(dit_variant_path)
|
||||
canon.DECODER_PATHS[decoder] = Path(dec_variant_path)
|
||||
dec_tier = canon.decoder_tier_from_filename(decoder, Path(dec_variant_path).name)
|
||||
print(f"\n → loading SA3Inference({dit!r}, {decoder!r}, "
|
||||
f"dit={Path(dit_variant_path).name}, "
|
||||
f"dec={Path(dec_variant_path).name}, tier={dec_tier})")
|
||||
inf = SA3Inference(dit, decoder,
|
||||
dec_precision=dec_tier,
|
||||
chunking=chunking,
|
||||
# pass the picked files explicitly -- mutating canon.DIT_CHOICES /
|
||||
# DECODER_PATHS above does not reach the resolver, so the dropdowns
|
||||
# were silently ignored for the DiT.
|
||||
dit_engine=(None if str(dit_variant_path).startswith("<")
|
||||
else dit_variant_path),
|
||||
dec_engine=dec_variant_path,
|
||||
default_T_lat=default_T_lat,
|
||||
default_steps=default_steps,
|
||||
default_seconds=default_seconds,
|
||||
quiet=quiet)
|
||||
_inference_cache[key] = inf
|
||||
_inference_lru.append(key)
|
||||
return inf
|
||||
|
||||
|
||||
# ── Gradio UI ──────────────────────────────────────────────────────────────
|
||||
_css = ("#sa3-promote{display:none !important}"
|
||||
"#sa3-out{gap:4px !important}"
|
||||
"#sa3-out .html-container{padding:0 !important; margin:0 !important}"
|
||||
# the queued chip is clickable once a take is ready — say so on hover/focus
|
||||
".sa3-queued{transition:background 0.12s, color 0.12s}"
|
||||
".sa3-queued:hover,.sa3-queued:focus{background:rgba(127,127,127,0.22) !important;"
|
||||
"color:#bbb !important; outline:none}")
|
||||
|
||||
|
||||
def build_ui(initial_dit: str, initial_decoder: str, *,
|
||||
share: bool, quiet: bool,
|
||||
default_T_lat: int, default_steps: int, default_seconds: float):
|
||||
"""Same UI as the MLX gradio -- history, a pre-rendered next take, radio mode, the rich
|
||||
player -- over the TensorRT backend, plus the controls only TRT has (engine variant per
|
||||
model, decoder quantisation tier, chunked vs single-shot).
|
||||
|
||||
No LoRA panel: the TRT engines carry no LoRA branches, so there is nothing to hot-swap
|
||||
without first building refit-capable or branch-capable engines.
|
||||
"""
|
||||
import gradio as gr
|
||||
|
||||
initial_variants = discover_variants(initial_dit)
|
||||
if not initial_variants:
|
||||
raise RuntimeError(f"no DiT engines found for {initial_dit} under {MODELS_ROOT}")
|
||||
initial_variant_path = str(initial_variants[0][1])
|
||||
initial_dec_variants = discover_decoder_variants(initial_decoder)
|
||||
if not initial_dec_variants:
|
||||
raise RuntimeError(f"no decoder engines found for {initial_decoder} under {MODELS_ROOT}")
|
||||
initial_dec_variant_path = str(initial_dec_variants[0][1])
|
||||
get_inference(initial_dit, initial_decoder, initial_variant_path,
|
||||
initial_dec_variant_path, default_T_lat, default_steps,
|
||||
default_seconds, quiet)
|
||||
|
||||
DIT_OPTIONS = list(DEFAULT_DECODERS.keys())
|
||||
# The DiT profile runs L=1..4096 and every shipped chunkable autoencoder reports a
|
||||
# floor of 1 latent, so 1 s is the honest slider minimum. It used to be 3 s, derived
|
||||
# from the stale DECODER_MIN_L=32 that SA3Inference._dec_min_L() was returning because
|
||||
# of a TypeError it swallowed; the guard reads the real per-engine floor now, so an
|
||||
# older engine with a genuine 32 floor still gets rejected with a clear message.
|
||||
MIN_SECONDS = 1
|
||||
MAX_SECONDS = SA3Inference.DIT_MAX_L * SAMPLES_PER_LATENT / SAMPLE_RATE // 1 # 380 s
|
||||
|
||||
def on_dit_change(dit_name):
|
||||
variants = discover_variants(dit_name)
|
||||
suggested = DEFAULT_DECODERS.get(dit_name, "same-s")
|
||||
dec_variants = discover_decoder_variants(suggested)
|
||||
var_u = (gr.update(choices=[(l, str(p)) for l, p in variants], value=str(variants[0][1]))
|
||||
if variants else gr.update(choices=[], value=None))
|
||||
dec_u = (gr.update(choices=[(l, str(p)) for l, p in dec_variants],
|
||||
value=str(dec_variants[0][1]))
|
||||
if dec_variants else gr.update(choices=[], value=None))
|
||||
return var_u, gr.update(value=suggested), dec_u
|
||||
|
||||
def on_decoder_change(decoder_name):
|
||||
dv = discover_decoder_variants(decoder_name)
|
||||
if not dv:
|
||||
return gr.update(choices=[], value=None)
|
||||
return gr.update(choices=[(l, str(p)) for l, p in dv], value=str(dv[0][1]))
|
||||
|
||||
def on_seconds_change(sec):
|
||||
"""Keep the inpaint range inside the clip length."""
|
||||
return gr.update(maximum=float(sec)), gr.update(maximum=float(sec))
|
||||
|
||||
# ── one generation, packaged as a history entry ──────────────────────────────────────────
|
||||
def _entry(dit_name, decoder_name, variant_path, dec_variant_path, prompt, negative_prompt,
|
||||
seconds, steps, seed_text, cfg, apg, sigma_max, chunking, fmt,
|
||||
a2a_path, a2a_sigma, inpaint_path, inp_start, inp_end, debug=False):
|
||||
# An empty prompt is a legitimate request, not an error: it generates
|
||||
# unconditionally. The encoder takes the empty string happily, and both
|
||||
# renderers already caption it "(no prompt)".
|
||||
prompt = (prompt or "").strip()
|
||||
try:
|
||||
seed = int(seed_text.strip()) if seed_text and seed_text.strip() else None
|
||||
except ValueError:
|
||||
return None, "seed must be an integer"
|
||||
if str(variant_path) == PT_EAGER_VARIANT:
|
||||
return None, "PT FP32 GT is not wired through this backend yet"
|
||||
try:
|
||||
inf = get_inference(dit_name, decoder_name, str(variant_path), str(dec_variant_path),
|
||||
default_T_lat, default_steps, default_seconds, quiet,
|
||||
chunking=bool(chunking))
|
||||
except Exception as e:
|
||||
return None, f"load failed: {type(e).__name__}: {e}"
|
||||
|
||||
# Which reference wins, and -- just as important -- SAY SO when something the user
|
||||
# supplied is being dropped. Silently falling through to a different mode is the
|
||||
# worst outcome here: upload an inpaint reference, forget the range sliders, and a
|
||||
# bare if/elif quietly runs audio-to-audio instead.
|
||||
#
|
||||
# The engines take ONE reference (generate_eager reads a single init_audio_path and
|
||||
# derives the inpaint keep-mask from its latents), so inpaint and audio-to-audio
|
||||
# cannot use two different files the way the MLX backend can.
|
||||
kw, mode, notes = {}, "text-to-audio", []
|
||||
sec = float(seconds)
|
||||
rng = None
|
||||
if inpaint_path and float(inp_end) > float(inp_start) and float(inp_start) < sec:
|
||||
if float(inp_end) > sec:
|
||||
notes.append(f"inpaint end clamped to the clip length ({sec:g}s)")
|
||||
rng = (float(inp_start), float(min(float(inp_end), sec)))
|
||||
elif inpaint_path and float(inp_end) <= float(inp_start):
|
||||
notes.append("inpainting ignored — set the start/end range sliders")
|
||||
elif inpaint_path:
|
||||
notes.append(f"inpainting ignored — start ({float(inp_start):g}s) is past "
|
||||
f"the end of the clip ({sec:g}s)")
|
||||
elif float(inp_end) > float(inp_start):
|
||||
notes.append("inpaint range ignored — no reference audio uploaded")
|
||||
|
||||
if rng is not None:
|
||||
kw = dict(init_audio_path=inpaint_path, inpaint_range=rng,
|
||||
init_noise_level=float(sigma_max))
|
||||
mode = "inpaint"
|
||||
if a2a_path:
|
||||
notes.append("using the inpaint reference — one reference per render, and "
|
||||
"inpainting needs it to keep the untouched part")
|
||||
elif a2a_path:
|
||||
# a2a's own slider IS the schedule start (parent-repo semantics), so it
|
||||
# overrides sigma_max rather than stacking with it.
|
||||
kw = dict(init_audio_path=a2a_path, init_noise_level=float(a2a_sigma))
|
||||
mode = "audio-to-audio"
|
||||
if float(sigma_max) != 1.0:
|
||||
notes.append(f"sigma_max {float(sigma_max):g} ignored — audio-to-audio uses "
|
||||
f"its own noise level ({float(a2a_sigma):g})")
|
||||
elif float(sigma_max) != 1.0:
|
||||
kw = dict(init_noise_level=float(sigma_max))
|
||||
neg = (negative_prompt or "").strip()
|
||||
if float(cfg) != 1.0:
|
||||
kw.update(cfg=float(cfg), apg=float(apg), negative_prompt=neg or None)
|
||||
elif neg:
|
||||
notes.append("negative prompt ignored — it only applies when CFG > 1")
|
||||
try:
|
||||
pcm, t = inf.generate(prompt, seconds=float(seconds), steps=int(steps),
|
||||
seed=seed, **kw)
|
||||
except NotImplementedError as e:
|
||||
return None, f"not supported by this engine: {e}"
|
||||
except Exception as e:
|
||||
return None, f"{type(e).__name__}: {e}"
|
||||
|
||||
debug_html = ""
|
||||
if debug:
|
||||
try:
|
||||
# Reuse the seed so the second pass is the same take, then hand both timing
|
||||
# dicts to the renderer. When the render already went eager there is nothing
|
||||
# to compare against, so skip the extra pass entirely.
|
||||
if not t.get("eager") and not getattr(inf, "_debug_warmed", False):
|
||||
# First eager call on this instance builds the _SecondContext set and its
|
||||
# buffers. Timing that pass charges ~600 ms of one-off setup to "unfused"
|
||||
# and makes the fusion win look far bigger than it is. Burn one pass.
|
||||
inf.generate_eager(prompt, seconds=float(seconds), steps=int(steps),
|
||||
seed=t["seed"], **kw)
|
||||
inf._debug_warmed = True
|
||||
t_stage = t if t.get("eager") else inf.generate_eager(
|
||||
prompt, seconds=float(seconds), steps=int(steps), seed=t["seed"], **kw)[1]
|
||||
debug_html = render_debug(inf, t, t_stage)
|
||||
except Exception as e:
|
||||
debug_html = (f"<div style='color:#e8a33d; font-size:0.82em'>debug pass failed: "
|
||||
f"{html_lib.escape(type(e).__name__)}: {html_lib.escape(str(e))}</div>")
|
||||
base = verbose_basename(prompt or "unconditional", neg,
|
||||
float(cfg), float(sigma_max), t["seed"])
|
||||
out_path = OUTPUT_DIR / f"{base}.wav"
|
||||
_save_wav(pcm, out_path)
|
||||
if fmt in ("flac", "mp3"):
|
||||
conv = out_path.with_suffix("." + fmt)
|
||||
rc = subprocess.run(["ffmpeg", "-y", "-loglevel", "error",
|
||||
"-i", str(out_path), str(conv)],
|
||||
capture_output=True).returncode
|
||||
if rc == 0 and conv.exists():
|
||||
out_path = conv
|
||||
mime = {"wav": "audio/wav", "flac": "audio/flac", "mp3": "audio/mpeg"}[out_path.suffix[1:]]
|
||||
try:
|
||||
t0 = time.time()
|
||||
png = render_spectrogram_png(pcm, sample_rate=SAMPLE_RATE, width=1400, height=280)
|
||||
spec_b64 = base64.b64encode(png).decode("ascii")
|
||||
spec_ms = (time.time() - t0) * 1000
|
||||
except Exception:
|
||||
spec_b64, spec_ms = "", 0.0
|
||||
timing = (f"engine-load {t.get('graph_build_ms', 0):.0f} ms · "
|
||||
f"Inference : {t['inference_ms']:.1f} ms · "
|
||||
f"{t.get('realtime', 0):.0f}× realtime · "
|
||||
f"seed : {t['seed']} · T_lat : {t['T_lat']} · "
|
||||
f"samples : {t['samples']} · spec {spec_ms:.0f} ms")
|
||||
return {
|
||||
"notes": notes, "debug_html": debug_html,
|
||||
"key": f"k{time.time_ns()}", "ts": time.time(), "dit": dit_name,
|
||||
"neg": neg, "cfg": float(cfg), "smx": float(sigma_max),
|
||||
"path": str(out_path), "mime": mime, "name": out_path.name,
|
||||
"size_mb": out_path.stat().st_size / 1e6, "mode": mode,
|
||||
"prompt": prompt, "seed": t["seed"],
|
||||
"lora": "", # kept for the shared renderer's contract
|
||||
"spec_b64": spec_b64, "timing": timing,
|
||||
}, None
|
||||
|
||||
def _present(state, entry, opts, queued_panel, *, force_autoplay=False):
|
||||
"""(player, timing, error, history, queued, state) for the shared renderers."""
|
||||
auto = force_autoplay or ("Auto-play" in (opts or []))
|
||||
radio = "Radio" in (opts or [])
|
||||
loop = "Loop" in (opts or [])
|
||||
hist = state.get("history", [])
|
||||
player = render_player(entry, autoplay=auto, radio=radio, advance=radio, loop=loop,
|
||||
autodl=("Auto-download" in (opts or []))) if entry else ""
|
||||
note_html = "".join(
|
||||
f"<div style='color:#e8a33d;font-size:0.85em'>note: {html_lib.escape(n)}</div>"
|
||||
for n in (entry.get("notes") or [])) if entry else ""
|
||||
return (player, (entry.get("timing", "") + note_html
|
||||
+ entry.get("debug_html", "")) if entry else "", "",
|
||||
render_history(hist, advance=radio, loop=loop),
|
||||
queued_panel, state)
|
||||
|
||||
CTRL = None # bound after the widgets exist
|
||||
|
||||
def _sig(ctrl):
|
||||
"""What the queued take was rendered FROM.
|
||||
|
||||
A pre-rendered take is only reusable while the controls still describe it. Serving
|
||||
it blind means the click after any edit silently returns the OLD settings: generate
|
||||
30 s, drag Seconds to 380, hit Generate, and you get 30 s of audio back. Everything
|
||||
in CTRL is a str/float/bool/None (audio inputs are filepaths), so str() over the
|
||||
tuple is a sound identity -- and it fails CLOSED, since anything it cannot compare
|
||||
simply looks different and forces a fresh render.
|
||||
"""
|
||||
return tuple(str(c) for c in ctrl)
|
||||
|
||||
_BLANK = {"current": None, "queued": None, "queued_sig": None, "history": []}
|
||||
|
||||
def generate(*args):
|
||||
*ctrl, opts, state = args
|
||||
state = dict(state or _BLANK)
|
||||
q = state.get("queued")
|
||||
if q is not None and state.get("queued_sig") == _sig(ctrl):
|
||||
# still the same request — show the pre-rendered take instantly
|
||||
state["queued"] = state["queued_sig"] = None
|
||||
state["current"] = q
|
||||
state["history"] = ([q] + state.get("history", []))[:24]
|
||||
return _present(state, q, opts, render_queue_status(generating=True),
|
||||
force_autoplay=True)
|
||||
# a control moved since it was queued: that take answers a question no longer
|
||||
# being asked, so drop it and render what the controls say now.
|
||||
state["queued"] = state["queued_sig"] = None
|
||||
entry, err = _entry(*ctrl)
|
||||
if err:
|
||||
return ("", "", f"<span style='color:#f88'>error: {html_lib.escape(err)}</span>",
|
||||
render_history(state.get("history", [])), render_queue_status(), state)
|
||||
state["current"] = entry
|
||||
state["history"] = ([entry] + state.get("history", []))[:24]
|
||||
return _present(state, entry, opts, render_queue_status(), force_autoplay=True)
|
||||
|
||||
def pregen(*args):
|
||||
"""Render the NEXT take in the background so the following click is instant."""
|
||||
*ctrl, opts, state = args
|
||||
state = dict(state or _BLANK)
|
||||
if ctrl[-1]:
|
||||
# Debug mode already renders twice (fused + unfused). Pre-rendering on top would
|
||||
# make it four renders per click for a take whose numbers may never be looked at.
|
||||
return "", state
|
||||
sig = _sig(ctrl)
|
||||
if state.get("queued") is not None and state.get("queued_sig") == sig:
|
||||
return render_queue_status(state["queued"]), state
|
||||
entry, err = _entry(*ctrl)
|
||||
state["queued"] = entry if not err else None
|
||||
state["queued_sig"] = sig if not err else None
|
||||
return render_queue_status(state.get("queued")), state
|
||||
|
||||
with gr.Blocks(title="SA3 · TensorRT", css=_css) as demo:
|
||||
st = gr.State(dict(_BLANK))
|
||||
with gr.Row():
|
||||
with gr.Column(scale=3):
|
||||
with gr.Row():
|
||||
dit_dd = gr.Dropdown(label="DiT model", choices=DIT_OPTIONS,
|
||||
value=initial_dit, scale=1)
|
||||
dec_dd = gr.Dropdown(label="Decoder", choices=["same-s", "same-l"],
|
||||
value=initial_decoder, scale=1)
|
||||
with gr.Row():
|
||||
var_dd = gr.Dropdown(label="DiT variant", allow_custom_value=True,
|
||||
choices=[(l, str(p)) for l, p in initial_variants],
|
||||
value=initial_variant_path, scale=1)
|
||||
dec_var_dd = gr.Dropdown(label="Decoder variant", allow_custom_value=True,
|
||||
choices=[(l, str(p)) for l, p in initial_dec_variants],
|
||||
value=initial_dec_variant_path, scale=1)
|
||||
with gr.Row():
|
||||
prompt = gr.Textbox(label="Prompt", lines=2, scale=6,
|
||||
placeholder="warm analog house groove")
|
||||
seed = gr.Textbox(label="Seed (optional)", max_lines=1, value="", scale=1)
|
||||
with gr.Row():
|
||||
seconds = gr.Slider(label="Seconds", minimum=MIN_SECONDS,
|
||||
maximum=MAX_SECONDS, step=1,
|
||||
value=default_seconds, scale=2)
|
||||
steps = gr.Slider(label="Steps", minimum=1, maximum=16, step=1,
|
||||
value=default_steps, scale=1)
|
||||
cfg = gr.Slider(label="CFG", minimum=0.0, maximum=10.0, step=0.1,
|
||||
value=1.0, scale=1)
|
||||
with gr.Accordion("Advanced", open=False):
|
||||
with gr.Row():
|
||||
apg = gr.Slider(label="APG (only applies when CFG > 1)",
|
||||
minimum=0.0, maximum=1.0, step=0.05, value=1.0)
|
||||
sigma_global = gr.Slider(label="sigma_max", minimum=0.05, maximum=1.0,
|
||||
step=0.05, value=1.0)
|
||||
negative_prompt = gr.Textbox(label="Negative prompt", lines=1)
|
||||
debug = gr.Checkbox(
|
||||
label="Debug: per-stage ms + VRAM", value=False,
|
||||
info="Adds a stage table under the timings. The mega-graph cannot be "
|
||||
"instrumented from inside (a captured graph rejects event records), "
|
||||
"so the split is measured by re-running the same engines unfused — "
|
||||
"which roughly doubles the render. The fused total is shown beside "
|
||||
"it, so the gap is the fusion win.")
|
||||
chunking = gr.Checkbox(
|
||||
label="Chunked decode (low VRAM)", value=True,
|
||||
info="On: 256-latent windows — 485 MB of decoder scratch on SAME-L, 346 MB "
|
||||
"on SAME-S. Off: single-shot on the wide profile — only 2-5% faster "
|
||||
"end-to-end (the DiT dominates), but 5-7 GB more resident. The scratch "
|
||||
"is a profile-ceiling reservation, so switching it off costs the same "
|
||||
"at 30 s as at 380 s.")
|
||||
with gr.Accordion("Audio-to-audio (guide the whole clip)", open=False):
|
||||
a2a_audio = gr.Audio(label="Guide audio — generation starts from its latents",
|
||||
type="filepath")
|
||||
a2a_sigma = gr.Slider(label="sigma_max (lower = closer to the guide)",
|
||||
minimum=0.05, maximum=1.0, step=0.05, value=0.6)
|
||||
with gr.Accordion("Inpainting (regenerate a span)", open=False):
|
||||
inpaint_audio = gr.Audio(label="Reference audio", type="filepath")
|
||||
with gr.Row():
|
||||
inp_start = gr.Slider(label="Start (s)", minimum=0,
|
||||
maximum=default_seconds, step=0.5, value=0)
|
||||
inp_end = gr.Slider(label="End (s)", minimum=0,
|
||||
maximum=default_seconds, step=0.5, value=0)
|
||||
with gr.Accordion("Output", open=False):
|
||||
output_opts = gr.CheckboxGroup(
|
||||
label="Playback", choices=["Auto-play", "Radio", "Loop", "Auto-download"],
|
||||
value=["Auto-play"])
|
||||
file_format = gr.Radio(label="File format", choices=["wav", "flac", "mp3"],
|
||||
value="wav")
|
||||
generate_btn = gr.Button("Generate", variant="primary", size="lg")
|
||||
promote_btn = gr.Button("", elem_id="sa3-promote") # CSS-hidden, DOM-present
|
||||
with gr.Column(scale=4):
|
||||
output_player = gr.HTML()
|
||||
timing = gr.HTML()
|
||||
error_box = gr.HTML()
|
||||
queued_html = gr.HTML()
|
||||
history_html = gr.HTML()
|
||||
|
||||
dit_dd.change(on_dit_change, inputs=[dit_dd], outputs=[var_dd, dec_dd, dec_var_dd])
|
||||
dec_dd.change(on_decoder_change, inputs=[dec_dd], outputs=[dec_var_dd])
|
||||
seconds.change(on_seconds_change, inputs=[seconds], outputs=[inp_start, inp_end]
|
||||
).then(None, js=_JS_FIX_SLIDERS)
|
||||
|
||||
CTRL = [dit_dd, dec_dd, var_dd, dec_var_dd, prompt, negative_prompt, seconds, steps,
|
||||
seed, cfg, apg, sigma_global, chunking, file_format,
|
||||
a2a_audio, a2a_sigma, inpaint_audio, inp_start, inp_end, debug]
|
||||
main_out = [output_player, timing, error_box, history_html, queued_html, st]
|
||||
# ⚠ js= takes a JS *function expression* ("() => ...") — gradio evaluates it as one.
|
||||
# _JS_TRY_PLAY is an inline-attribute body (bare statements, `this` = the <audio>),
|
||||
# so passing it here is a syntax error that kills the whole frontend: the config
|
||||
# fails to evaluate and the page sits on "Loading..." forever with zero controls.
|
||||
# The autoplay rescue belongs where it already is — render_player's oncanplay.
|
||||
for btn in (generate_btn, promote_btn):
|
||||
btn.click(generate, inputs=CTRL + [output_opts, st], outputs=main_out
|
||||
).then(pregen, inputs=CTRL + [output_opts, st],
|
||||
outputs=[queued_html, st]
|
||||
).then(None, js=_JS_FIX_SLIDERS)
|
||||
gr.Markdown("Renders are written to `output/`. **Radio** auto-advances to the next take as "
|
||||
"each finishes; the next one is pre-rendered while you listen.")
|
||||
|
||||
demo.queue(max_size=16).launch(share=share, server_name="0.0.0.0", show_error=True,
|
||||
allowed_paths=[str(OUTPUT_DIR)])
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--dit", choices=list(DEFAULT_DECODERS.keys()),
|
||||
default="sm-music",
|
||||
help="Initial DiT bundle (switchable at runtime)")
|
||||
ap.add_argument("--decoder", choices=["same-s", "same-l"], default=None,
|
||||
help="Initial decoder. Default: pairs with --dit")
|
||||
ap.add_argument("--default-seconds", type=float, default=120.0,
|
||||
help="Length to pre-warm the initial graph at")
|
||||
ap.add_argument("--default-steps", type=int, default=8)
|
||||
ap.add_argument("--share", action=argparse.BooleanOptionalAction, default=True,
|
||||
help="Create a public gradio.live URL (default on)")
|
||||
ap.add_argument("--chunking", action=argparse.BooleanOptionalAction, default=True,
|
||||
help="SAME-L: decode/encode in windows on the engines' low profile "
|
||||
"(509 MB of decoder scratch instead of 8143). --no-chunking uses the "
|
||||
"wide profile: single-shot, faster above L=256, ~5.4 GB more resident.")
|
||||
ap.add_argument("--quiet", action="store_true")
|
||||
args = ap.parse_args()
|
||||
global _CHUNKING
|
||||
_CHUNKING = args.chunking
|
||||
|
||||
if args.decoder is None:
|
||||
args.decoder = DEFAULT_DECODERS[args.dit]
|
||||
|
||||
T_lat = max(1, math.ceil(args.default_seconds * SAMPLE_RATE / SAMPLES_PER_LATENT))
|
||||
if args.decoder == "same-s" and T_lat % 2 != 0:
|
||||
T_lat += 1
|
||||
|
||||
print(f"\n━━━ SA3 TRT — gradio (debug build) ━━━")
|
||||
print(f" initial dit: {args.dit}")
|
||||
print(f" initial decoder: {args.decoder}")
|
||||
print(f" warmup: T_lat={T_lat} steps={args.default_steps} "
|
||||
f"(~{args.default_seconds}s)")
|
||||
# Pretty-print what variants are visible
|
||||
for ditname in DEFAULT_DECODERS:
|
||||
vs = discover_variants(ditname)
|
||||
print(f" dit-variants[{ditname}]: " + (
|
||||
", ".join(lbl for lbl, _ in vs) if vs else "(none found)"))
|
||||
for dec in ("same-s", "same-l"):
|
||||
dvs = discover_decoder_variants(dec)
|
||||
print(f" dec-variants[{dec}]: " + (
|
||||
", ".join(lbl for lbl, _ in dvs) if dvs else "(none found)"))
|
||||
print()
|
||||
|
||||
build_ui(args.dit, args.decoder, share=args.share, quiet=args.quiet,
|
||||
default_T_lat=T_lat, default_steps=args.default_steps,
|
||||
default_seconds=args.default_seconds)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user