Files
SoundGame/tools/stable-audio-3/optimized/tflite/scripts/sa3_gradio.py
T
SoundGame Dev f0f0a9b50e 初始提交: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 模板
2026-08-29 13:00:34 +08:00

1299 lines
64 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""SA3 TFLite — gradio web UI (portable CPU / XNNPACK).
The TFLite sibling of optimized/mlx/scripts/sa3_gradio.py, driving the baked-I/O
varlen TFLite models (scripts/sa3_tflite.py) instead of MLX. Every generation mode
is wired, plus a precision picker the MLX/TRT UIs don't need:
- Model picker: sm-music / sm-sfx / medium (hot-swap; the large DiT interpreter
is cached — first use of a model/precision/length loads + XNNPACK-packs the
weights, subsequent runs reuse it and only re-bind the conditioning).
- Precision picker: fp32 / w16a32 / w8a32 / w8a8-dyn (one knob for DiT + codec +
encoder, exactly like the CLI's --precision). fp32 is the CPU fast-and-accurate
default; w16a32 is ≈lossless half-size; w8a32/w8a8-dyn are GPTQ int8 (¼ size).
- CFG 0-10 next to seconds/steps (0 = negative prompt takes over, 0.5 = halfway
between prompts, 1 = off, >1 = extrapolate) + negative prompt/APG under Advanced.
- Audio-to-audio: guide audio + init_noise_level (whole clip starts from its
latents); global sigma_max under Advanced for prompt-only generations.
- Inpainting: separate reference audio + start/end range sliders (kept bit-exact
outside the range). Combinable with a2a.
- LoRA: file upload + local-folder pick (loras/<model>/) + strength, per-model
memory, add/remove. Merged once into the DiT weights at load (the frozen graph
has no per-step gating — so no Min/Max-step sliders, unlike the MLX UI). Only
fp32 / w16a32 can be LoRA-merged; the panel disables itself under int8.
- Spectrogram-as-player (numpy port, white playhead, click-to-seek), history
panel, Output options (Auto-play / Auto-download / Infinite Radio / Loop /
Hotswap), MP3 (via ffmpeg) or WAV saving.
Launch:
./sa3-gradio # share=True by default, sm-music + same-s, fp32
./sa3-gradio --dit medium --precision w8a32
./sa3-gradio --no-share # local-only
"""
from __future__ import annotations
import argparse
import base64
import gc
import html as html_lib
import os
import re
import shutil
import subprocess
import sys
import time
import urllib.parse
import wave
from pathlib import Path
import numpy as np
SCRIPTS_DIR = Path(__file__).resolve().parent
REPO = SCRIPTS_DIR.parent
sys.path.insert(0, str(REPO)) # so `from models.defs.* import` resolves
sys.path.insert(0, str(SCRIPTS_DIR)) # so `from weights / lora_* / spec import` resolves
from sa3_tflite import ( # noqa: E402
BakedDiT, BakedEncoder, BakedDecoder, read_wav,
valid_T_lat, DEFAULT_DECODER, DIT_REL, DEC_REL, T5_REL,
COND_TOKENS, COND_DIM, SAMPLE_RATE, SAMPLES_PER_LATENT,
SAMEL_CHUNK, SAMEL_OVERLAP, MIN_SIGMA,
)
from models.defs import tflite_pipeline as P # noqa: E402
from lora_core import parse_lora_spec, LoraError # noqa: E402
from lora_patch import get_patched_dit # noqa: E402
from weights import ensure_local, PRECISIONS, dit_rel, dec_rel, enc_rel # noqa: E402
from spec import render_spectrogram_png # noqa: E402
OUTPUT_DIR = REPO / "output" / "gradio"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
# Local LoRA library: drop .safetensors under loras/<model>/ and they appear in
# the per-slot dropdown (scanned per selected DiT; hidden when empty).
LORAS_DIR = REPO / "loras"
LORA_DIR_NAMES = {"medium": "sa3-medium", "sm-sfx": "sa3-sm-sfx",
"sm-music": "sa3-sm-music"}
for _d in LORA_DIR_NAMES.values():
(LORAS_DIR / _d).mkdir(parents=True, exist_ok=True)
_DD_NONE = "" # dropdown placeholder value ("--- Choose a LoRA ---")
# Precisions whose weights can be LoRA-merged (fp32 consts + w16a32 fp16-behind-
# DEQUANTIZE). Quantized-int8 graphs (w8a32 / w8a8-dyn) can't — see lora_patch.
LORA_PRECISIONS = ("fp32", "w16a32")
# Trained max clip length per model (repo README model table).
MAX_SECONDS = {"sm-music": 120, "sm-sfx": 120, "medium": 380}
DEFAULT_DECODERS = dict(DEFAULT_DECODER)
# XNNPACK CPU threads (set from --threads in main()).
_THREADS = 8
def _lora_dd_choices(dit_name: str) -> list:
"""Dropdown choices for ./loras/<model>/: placeholder + (name, abspath)."""
d = LORAS_DIR / LORA_DIR_NAMES.get(dit_name, dit_name)
hits = sorted(d.glob("*.safetensors")) if d.is_dir() else []
return [("--- Choose a LoRA ---", _DD_NONE)] + [(p.name, str(p)) for p in hits]
# MP3 (V0) saving needs ffmpeg; without it we save WAV and hide the choice.
FFMPEG = shutil.which("ffmpeg") is not None
FORMAT_MP3, FORMAT_WAV = "Save to MP3 (V0)", "Save to WAV"
# The baked TFLite models are frozen FlatBuffer graphs, so unlike MLX there is no
# ThreadLocalStream hazard — but a single XNNPACK Interpreter must NOT be invoked
# from two threads at once, and gradio runs each handler on a rotating anyio worker
# thread. The categorical fix (same structure as the MLX UI): ALL model work
# (pre-warm + every generation) runs on one dedicated owner thread; handlers submit
# to it and wait. This also serializes generations.
import concurrent.futures as _cf # noqa: E402
_EXECUTOR = _cf.ThreadPoolExecutor(max_workers=1, thread_name_prefix="tfl")
def run_serial(fn, *args, **kwargs):
"""Run fn on the owner thread and return its result (re-raises errors)."""
return _EXECUTOR.submit(fn, *args, **kwargs).result()
def read_audio_any(path: str) -> np.ndarray:
"""(2, T) float32 @ 44.1 kHz from any common upload format.
Layered: read_wav handles 16-bit/44.1k natively and shells out to ffmpeg
for the rest when installed; if that fails (no ffmpeg), soundfile's bundled
libsndfile decodes mp3/flac/ogg/24-bit/48 kHz directly, followed by a
linear resample to 44.1 kHz."""
try:
return read_wav(path)
except Exception:
import soundfile as sf
data, sr = sf.read(path, dtype="float32", always_2d=True) # (T, ch)
y = data.T
y = np.stack([y[0], y[0]]) if y.shape[0] == 1 else y[:2]
if sr != SAMPLE_RATE:
in_len = y.shape[-1]
new_len = int(round(in_len * SAMPLE_RATE / sr))
scale = in_len / new_len
pos = np.clip((np.arange(new_len) + 0.5) * scale - 0.5, 0, in_len - 1)
y = np.stack([np.interp(pos, np.arange(in_len), ch) for ch in y])
return np.ascontiguousarray(y, dtype=np.float32)
def condense_prompt(prompt: str) -> str:
"""Prompt → filename fragment (the main repo's verbose-naming rule):
filesystem-special characters become hyphens, capped at 150 chars."""
prompt = re.sub(r'[\\/:*?"<>|]', '-', prompt)[:150]
return prompt or "_"
def verbose_basename(prompt, negative_prompt, cfg, sigma_max, seed, precision) -> str:
"""prompt[.neg-…][.cfg{scale}][.smx{σ}][.{precision}].{seed} — the main repo's
gradio 'verbose' file naming, plus a precision tag for non-fp32 runs so A/B
runs across --precision don't overwrite each other."""
base = condense_prompt(prompt)
if negative_prompt and negative_prompt.strip():
base += ".neg-" + condense_prompt(negative_prompt.strip())
if cfg != 1.0:
base += f".cfg{cfg:g}"
if sigma_max != 1.0:
base += f".smx{sigma_max:g}"
if precision and precision != "fp32":
base += f".{precision}"
return f"{base}.{seed}"
def _save_wav(pcm_int16, out_path):
"""pcm_int16: (T, 2) int16 interleaved."""
with wave.open(str(out_path), "wb") as w:
w.setnchannels(2)
w.setsampwidth(2)
w.setframerate(SAMPLE_RATE)
w.writeframes(pcm_int16.tobytes())
# ── Model caches ─────────────────────────────────────────────────────────────
# The T5Gemma front-end is loaded once (fixed 256 tokens). The DiT interpreter is
# large (medium fp32 ≈ 5.8 GB) and holds its conditioning as RESIDENT tensors, so
# it's cached by FILE PATH (dit family + precision, or the LoRA-patched clone) and
# re-bound per generation via set_conditioning — resize+allocate only when the
# batch (cfg) or length changes. Codec enc/dec are cached by (name, precision).
_t5 = None
_tok = None
_dit_cache: dict[str, BakedDiT] = {}
_dit_lru: list[str] = []
_DIT_CACHE_MAX = 1 # one big DiT resident at a time (medium fp32 = 5.8 GB)
_dec_cache: dict[tuple, BakedDecoder] = {}
_dec_lru: list[tuple] = []
_enc_cache: dict[tuple, BakedEncoder] = {}
_enc_lru: list[tuple] = []
_CODEC_CACHE_MAX = 2
# Encoded-latent cache: rerunning a2a/inpainting with the same input audio skips
# the encoder. Keyed on the file's CONTENT hash (gradio re-uploads get fresh temp
# paths) + codec + precision + latent length.
_latent_cache: dict[tuple, np.ndarray] = {}
_LATENT_CACHE_MAX = 32
def get_t5():
global _t5, _tok
if _t5 is None:
_t5 = P.T5GemmaTFLite(ensure_local(T5_REL), _THREADS)
_tok = P.Tokenizer()
return _t5, _tok
def get_dit(dit_path, T_lat, t5_hidden, t5_mask, seconds, cfg, apg,
null_h, null_m, local_add_cond, batched):
"""Cache the (large) baked DiT interpreter by its file path; re-bind the
per-generation conditioning via set_conditioning. Returns (model, load_ms)."""
key = str(dit_path)
if key in _dit_cache:
if key in _dit_lru:
_dit_lru.remove(key)
_dit_lru.append(key)
model = _dit_cache[key]
model.set_conditioning(T_lat, t5_hidden, t5_mask, seconds, cfg=cfg, apg=apg,
null_hidden=null_h, null_mask=null_m,
local_add_cond=local_add_cond, batched=batched)
return model, 0.0
while len(_dit_cache) >= _DIT_CACHE_MAX:
old = _dit_lru.pop(0)
print(f" ← LRU-evicting DiT {Path(old).name}")
_dit_cache.pop(old, None)
gc.collect()
t0 = time.time()
model = BakedDiT(dit_path, T_lat, t5_hidden, t5_mask, seconds, _THREADS,
cfg=cfg, apg=apg, null_hidden=null_h, null_mask=null_m,
local_add_cond=local_add_cond, batched=batched)
load_ms = (time.time() - t0) * 1000
_dit_cache[key] = model
_dit_lru.append(key)
return model, load_ms
def get_decoder(name: str, precision: str) -> BakedDecoder:
key = (name, precision)
if key in _dec_cache:
_dec_lru.remove(key)
_dec_lru.append(key)
return _dec_cache[key]
while len(_dec_cache) >= _CODEC_CACHE_MAX:
old = _dec_lru.pop(0)
_dec_cache.pop(old, None)
gc.collect()
dec = BakedDecoder(ensure_local(dec_rel(name, precision)), _THREADS,
needs_even=(name == "same-s"))
_dec_cache[key] = dec
_dec_lru.append(key)
return dec
def get_encoder(name: str, precision: str) -> BakedEncoder:
key = (name, precision)
if key in _enc_cache:
_enc_lru.remove(key)
_enc_lru.append(key)
return _enc_cache[key]
while len(_enc_cache) >= _CODEC_CACHE_MAX:
old = _enc_lru.pop(0)
_enc_cache.pop(old, None)
gc.collect()
enc = BakedEncoder(ensure_local(enc_rel(name, precision)), _THREADS,
needs_even=(name == "same-s"))
_enc_cache[key] = enc
_enc_lru.append(key)
return enc
# ── One full generation (mirrors sa3_tflite.py main(), all modes) ─────────────
def run_generation(dit_name: str, decoder_name: str, precision: str, prompt: str,
negative_prompt: str, seconds: float, steps: int,
seed: int, cfg: float, apg: float, sigma_max: float,
a2a_audio_path: str | None = None,
inpaint_audio_path: str | None = None,
inpaint_range_sec=None,
lora_specs=None, cfg_batched: bool = True):
"""Returns (audio_np (2,T) float32, timings dict).
a2a and inpainting are independent and combinable (same as the MLX UI):
- a2a_audio_path: the whole generation STARTS from this audio's latents
(x0 = lat*(1-σmax) + noise*σmax)
- inpaint_audio_path + inpaint_range_sec: kept bit-exact outside the range
(local_add_cond context + per-step paste-back)
- both: the inpainted span regenerates FROM the a2a guide instead of pure noise
"""
t = {}
dec = decoder_name
needs_even = (dec == "same-s")
T_lat = valid_T_lat(seconds)
# 1. T5Gemma (tokenize + encode; + negative prompt for CFG)
t0 = time.time()
t5, tok = get_t5()
ids, mask = tok(prompt)
t5_hidden = t5(ids, mask)
null_h = null_m = None
if cfg != 1.0:
if negative_prompt and negative_prompt.strip():
n_ids, n_mask = tok(negative_prompt.strip())
null_h = t5(n_ids, n_mask)
null_m = n_mask.astype(np.float32)
else:
# All-zero hidden+mask → in-graph conditioner emits learned padding
# embeds for every position (the standard unconditional branch).
null_h = np.zeros((1, COND_TOKENS, COND_DIM), np.float32)
null_m = np.zeros((1, COND_TOKENS), np.float32)
t["t5_ms"] = (time.time() - t0) * 1000
# 2. Encode the provided audio inputs → latents (each optional). Memoized by
# (file content, codec, precision, T_lat) so re-running with the same audio
# skips the encoder.
def encode_audio(path):
import hashlib
ck = (hashlib.sha1(Path(path).read_bytes()).hexdigest(), dec, precision, T_lat)
if ck in _latent_cache:
t["enc_cached"] = t.get("enc_cached", 0) + 1
return _latent_cache[ck]
enc = get_encoder(dec, precision)
# SAME-S encoder needs even L; round the encode grid up, trim to T_lat.
enc_L = T_lat + 1 if (needs_even and T_lat % 2 != 0) else T_lat
target = enc_L * SAMPLES_PER_LATENT
audio_np = read_audio_any(path)
if audio_np.shape[-1] >= target:
audio_np = audio_np[:, :target]
else:
audio_np = np.pad(audio_np, ((0, 0), (0, target - audio_np.shape[-1])))
lat = enc.encode(audio_np[None], T_lat) # (1,256,T_lat)
while len(_latent_cache) >= _LATENT_CACHE_MAX:
_latent_cache.pop(next(iter(_latent_cache)))
_latent_cache[ck] = lat
return lat
a2a_latents = None # start-point guide (whole clip)
ctx_latents = None # inpaint context (kept outside the range)
t["enc_ms"] = 0.0
if a2a_audio_path:
t0 = time.time()
a2a_latents = encode_audio(a2a_audio_path)
t["enc_ms"] += (time.time() - t0) * 1000
if inpaint_audio_path and inpaint_range_sec is not None:
t0 = time.time()
ctx_latents = encode_audio(inpaint_audio_path)
t["enc_ms"] += (time.time() - t0) * 1000
# 3. inpaint local_add_cond + paste-back (TRT channel layout: [1,257,L])
local_add_cond = None
paste_back = None
if ctx_latents is not None:
s0 = max(0, int(round(inpaint_range_sec[0] * SAMPLE_RATE / SAMPLES_PER_LATENT)))
s1 = min(T_lat, int(round(inpaint_range_sec[1] * SAMPLE_RATE / SAMPLES_PER_LATENT)))
keep = np.ones((1, 1, T_lat), np.float32)
keep[:, :, s0:s1] = 0.0 # 1=keep init, 0=regenerate
masked = ctx_latents.astype(np.float32) * keep
local_add_cond = np.concatenate([keep, masked], axis=1) # (1,257,T_lat)
paste_back = (ctx_latents.astype(np.float32), keep)
# 4. initial noise + a2a start-point mix (independent of inpainting)
x0, step_noise = P.make_noise(T_lat, steps, seed)
if a2a_latents is not None:
x0 = a2a_latents.astype(np.float32) * (1.0 - sigma_max) + x0 * sigma_max
# 5. DiT load (+ optional LoRA patch) + pingpong sample
dit_path = ensure_local(dit_rel(dit_name, precision))
if lora_specs:
# Raises LoraError for quantized-int8 precisions (caught upstream → error box).
dit_path = get_patched_dit(dit_path, lora_specs, family=dit_name,
precision=precision, log=lambda m: print(f" {m}"))
backend, t["dit_load_ms"] = get_dit(str(dit_path), T_lat, t5_hidden,
mask.astype(np.float32), float(seconds),
cfg, apg, null_h, null_m, local_add_cond,
cfg_batched)
sigmas = P.build_pingpong_schedule(steps, sigma_max)
t0 = time.time()
latents = P.sample(backend, x0, step_noise, sigmas, None, None, paste_back=paste_back)
t["sample_ms"] = (time.time() - t0) * 1000
t["n_fwd"] = backend.n_fwd
# 6. Decode (whole for SAME-S; chunked for SAME-L past the window)
decoder = get_decoder(dec, precision)
t0 = time.time()
if dec == "same-l" and T_lat > SAMEL_CHUNK:
audio = decoder.decode_chunked(latents, SAMEL_CHUNK, SAMEL_OVERLAP)
else:
audio = decoder.decode_whole(latents)
t["decode_ms"] = (time.time() - t0) * 1000
audio_np = audio[0] # (2, L*4096)
requested = int(round(seconds * SAMPLE_RATE))
if audio_np.shape[-1] > requested:
audio_np = audio_np[..., :requested]
t["T_lat"] = T_lat
t["samples"] = audio_np.shape[-1]
t["inference_ms"] = sum(t.get(k, 0) for k in
("t5_ms", "enc_ms", "dit_load_ms", "sample_ms", "decode_ms"))
t["realtime"] = (audio_np.shape[-1] / SAMPLE_RATE) / max(t["inference_ms"] / 1000, 1e-9)
return audio_np, t
# ── HTML rendering (inline handlers only — gradio HTML runs attributes, not
# <script> tags) ────────────────────────────────────────────────────────────
# One-at-a-time playback. Chrome does NOT stop a playing <audio> when it's
# removed from the DOM (verified via CDP: swapped-out elements keep playing as
# detached ghosts), and querySelectorAll can't see detached nodes — so every
# player registers itself in window._sa3All on play, and pause-others sweeps
# that registry (reaching ghosts) plus the document. Entries that are detached
# AND paused drop from the registry so ghosts can be garbage-collected.
# A DETACHED element's play event must never silence the living player (the
# swap can briefly leave a superseded copy whose pending play() resolves late).
# gradio's slider track fill lives in an inline CSS var (--range_progress) that
# is only written on user input / first mount — a slider REmounted by a
# visibility toggle (LoRA slots on model switch / add) comes back without it and
# renders a blank track. Recompute it client-side after any remount batch;
# idempotent for healthy sliders (same formula gradio uses).
_JS_FIX_SLIDERS = (
"() => setTimeout(() => {"
"document.querySelectorAll('input[type=range]').forEach(r => {"
"const min = +r.min || 0, max = +r.max || 100, v = +r.value;"
"if (max > min) r.style.setProperty('--range_progress',"
"((v - min) / (max - min) * 100) + '%');"
"});}, 150)"
)
_JS_PAUSE_OTHERS = (
"if(this.isConnected){"
"var t=this;var L=window._sa3All=(window._sa3All||[]);"
"if(L.indexOf(t)<0)L.push(t);"
"L.forEach(function(o){if(o!==t){try{o.pause()}catch(e){}}});"
"document.querySelectorAll('audio').forEach(function(o){if(o!==t)o.pause();});"
"window._sa3All=L.filter(function(o){return o===t||o.isConnected;});}"
"else{this.pause();}")
_JS_PLAYHEAD = ("var p=this.closest('.blk').querySelector('.ph');"
"if(p&&this.duration)p.style.left=(this.currentTime/this.duration*100)+'%';")
# Main-player position ledger (for Hotswap): every timeupdate/play/pause stamps
# the current position so a freshly swapped-in element can resume exactly there.
# Guards (all verified via CDP against the live page):
# - isConnected: teardown fires a pause AND a final timeupdate on the removed
# element — detached elements never write the ledger.
# - unresumed hotswap candidates (data-hs set, hsd not yet) don't write either:
# with autoplay, the NEW element's 'play' event fires BEFORE loadedmetadata
# and would stamp t≈0 over the old clip's position right before the resume
# handler reads it.
_JS_POS_RECORD = ("if(this.isConnected&&!(this.dataset.hs==='1'&&!this.dataset.hsd))"
"{window._sa3Pos={t:this.currentTime,playing:!this.paused,ts:Date.now()};}")
# timeupdate variant: a PAUSED element only fires timeupdate on seeks — e.g. the
# resume handler's own currentTime assignment. Stamping playing:false there
# poisons the ledger for any second render of the same clip (gradio sometimes
# renders a component update twice), which froze the handoff chain. While
# paused, only real pause events may write.
_JS_POS_RECORD_TU = ("if(this.isConnected&&!this.paused&&"
"!(this.dataset.hs==='1'&&!this.dataset.hsd))"
"{window._sa3Pos={t:this.currentTime,playing:true,ts:Date.now()};}")
# Resilient play: Chrome's autoplay policy can reject programmatic play() once
# the transient user activation from the Generate click has expired (seconds),
# leaving the clip correctly positioned but paused. Try, retry shortly after,
# and as a last resort arm a one-shot listener so the user's next click or
# keypress anywhere resumes playback.
# Every attempt re-checks isConnected: a superseded (detached) element must
# never revive itself from a queued retry and fight the current player.
_JS_TRY_PLAY = (
"var A=this;A.play().catch(function(){setTimeout(function(){"
"if(!A.isConnected)return;"
"A.play().catch(function(){console.warn('play blocked by autoplay policy — "
"will resume on next interaction');"
"var f=function(){if(A.isConnected)A.play();"
"document.removeEventListener('pointerdown',f,true);"
"document.removeEventListener('keydown',f,true);};"
"document.addEventListener('pointerdown',f,true);"
"document.addEventListener('keydown',f,true);});},150);});")
# Hotswap resume: if the previous main audio was playing when this one arrived,
# jump to its position (+ the split-second since the last stamp) and keep going;
# beyond the new clip's duration -> start at zero. Guarded (hsd) so it applies
# once, on whichever of loadedmetadata/canplay fires first.
_JS_HOTSWAP = ("if(this.isConnected&&this.dataset.hs==='1'&&!this.dataset.hsd&&this.duration){"
"this.dataset.hsd=1;var s=window._sa3Pos;"
"this.dataset.dbg=s?(s.playing?'playing@'+s.t.toFixed(2):'notplaying@'+s.t.toFixed(2)):'noledger';"
"if(s&&s.playing){var tt=s.t+(Date.now()-s.ts)/1000;"
"this.currentTime=(tt<this.duration)?tt:0;" + _JS_TRY_PLAY + "}}")
_JS_SEEK = ("var a=this.closest('.blk').querySelector('audio');"
"var r=this.getBoundingClientRect();"
"if(a&&a.duration){a.currentTime=(event.clientX-r.left)/r.width*a.duration;a.play();}")
_JS_PROMOTE = ("var b=document.getElementById('sa3-promote');"
"if(b){(b.tagName==='BUTTON'?b:b.querySelector('button')||b).click();}")
# Scroll anchoring for the history panel: onscroll continuously records which
# item sits at the viewport top (+offset); after every re-render a hidden
# bootstrap <img onerror> restores that anchor — stick-to-top when at top,
# otherwise keep hovering over the same old item as new ones prepend.
_JS_SCROLL_RECORD = (
"var c=this,k=null,off=0,ch=c.querySelectorAll('[data-key]');"
"for(var i=0;i<ch.length;i++){if(ch[i].offsetTop+ch[i].offsetHeight>c.scrollTop)"
"{k=ch[i].getAttribute('data-key');off=ch[i].offsetTop-c.scrollTop;break}}"
"window._sa3S={atTop:c.scrollTop<8,key:k,off:off};")
_JS_SCROLL_RESTORE = (
"var c=document.getElementById('sa3-hist');var s=window._sa3S||{};"
"if(c){var el=s.key?c.querySelector('[data-key=&quot;'+s.key+'&quot;]'):null;"
"if(el&&!s.atTop){c.scrollTop=el.offsetTop-(s.off||0)}else{c.scrollTop=0}}"
"this.remove();")
def _ago(ts) -> str:
if not ts:
return ""
d = max(0.0, time.time() - ts)
if d < 10:
return "just now"
if d < 60:
return f"{int(d)}s ago"
if d < 3600:
return f"{int(d // 60)}m ago"
if d < 86400:
return f"{int(d // 3600)}h ago"
return f"{int(d // 86400)}d ago"
def _lora_specs_from_ui(vals, notes):
"""Turn the 3 LoRA slots' flat values (file, dropdown, strength)×3 into
lora_core specs. A slot's adapter is the dropped file or, failing that, the
dropdown pick from ./loras/<model>/; empty slots are skipped. Unlike the MLX
UI there is no per-step range (the frozen TFLite graph merges once at load).
A bad spec gets a note and is skipped (never an error — matches the app's
permissive-by-design generation policy)."""
specs = []
for i in range(3):
path, picked, strength = vals[3 * i: 3 * i + 3]
if not path:
path = picked if picked and picked != _DD_NONE else None
if not path:
continue
try:
specs.append(parse_lora_spec([path, f"strength={float(strength)}"]))
except LoraError as e:
notes.append(f"LoRA {i + 1}: {e}")
return specs or None
def _lora_disp(specs) -> str:
"""Short human tag per adapter: 'plini×0.8, rain'."""
tags = []
for s in specs or []:
name = Path(s["path"]).stem
if len(name) > 25:
name = name[:24] + "…"
tag = name
if s["strength"] != 1.0:
tag += f"×{s['strength']:g}"
tags.append(tag)
return ", ".join(tags)
def _meta_suffix(entry) -> str:
"""' · sm-mus · w8a32 · cfg 1.5 · noise 0.92 · audio2audio · inpainting' —
only the non-defaults. The sigma tag reads 'noise' for a2a runs (it IS the
init_noise_level there), 'smx' otherwise."""
parts = []
cfg = entry.get("cfg", 1.0)
smx = entry.get("smx", 1.0)
mode = entry.get("mode", "")
is_a2a = "a2a" in mode or mode == "audio-to-audio"
dit = entry.get("dit", "")
prec = entry.get("prec", "fp32")
if dit:
parts.append({"medium": "med", "sm-music": "sm-mus", "sm-sfx": "sm-sfx"}.get(dit, dit))
if prec and prec != "fp32":
parts.append(prec)
if cfg != 1.0:
parts.append(f"cfg {cfg:g}")
if smx != 1.0:
parts.append(f"{'noise' if is_a2a else 'smx'} {smx:g}")
if is_a2a:
parts.append("audio2audio")
if "inpaint" in mode:
parts.append("inpainting")
if entry.get("lora"):
parts.append(f"lora {entry['lora']}")
return "".join(f" · {p}" for p in parts)
def _neg_disp(entry) -> str:
"""Labeled negative prompt, shown only when it actually acted (cfg != 1)."""
neg = (entry.get("neg") or "").strip()
if neg and entry.get("cfg", 1.0) != 1.0:
return f' · <span style="opacity:0.75">neg: {html_lib.escape(neg)}</span>'
return ""
def render_player(entry, *, small=False, autoplay=False, autodl=False, radio=False,
bg=None, advance=False, loop=False, hotswap=False):
"""One self-contained player block: audio + caption + seekable spectrogram
with playhead. Global one-at-a-time playback via onplay pause-others.
small: audio + spectrogram side by side (half width each), 'Xm ago' caption.
advance: on ended, hop to the next history item's audio (Auto-play).
loop: native loop attribute — finished audio restarts (suppresses ended).
hotswap: main slot only — resume at the previous clip's position on arrival."""
# assemble per-event handler bodies (a duplicated attribute name would
# silently drop one handler, so each event is emitted exactly once)
on_canplay = []
if small:
attrs = f'onplay="{_JS_PAUSE_OTHERS}" ontimeupdate="{_JS_PLAYHEAD}"'
else:
attrs = (f'onplay="{_JS_PAUSE_OTHERS}{_JS_POS_RECORD}" '
f'onpause="{_JS_POS_RECORD}" '
f'ontimeupdate="{_JS_PLAYHEAD}{_JS_POS_RECORD_TU}"')
if hotswap and not small:
attrs += f' data-hs="1" onloadedmetadata="{_JS_HOTSWAP}"'
on_canplay.append(_JS_HOTSWAP) # fallback if loadedmetadata already passed
if autoplay and not small:
# the autoplay attribute is subject to the same policy — rescue a
# policy-paused clip (radio transitions, long generations)
on_canplay.append("if(this.paused&&!(this.dataset.hs==='1'&&!this.dataset.hsd))"
"{" + _JS_TRY_PLAY + "}")
if autodl:
js_name = entry["name"].replace("\\", "").replace("'", "\\'")
on_canplay.append("if(!this.dataset.dld){this.dataset.dld=1;"
"var l=document.createElement('a');l.href=this.src;"
f"l.download='{js_name}';l.click();}}")
if on_canplay:
attrs += f' oncanplay="{"".join(on_canplay)}"'
if loop:
attrs += " loop"
elif radio:
attrs += f' onended="{_JS_PROMOTE}"'
elif advance:
attrs += (' onended="var b=this.closest(\'.blk\');'
"var n=b?b.nextElementSibling:null;"
"while(n&&!n.classList.contains('blk'))n=n.nextElementSibling;"
"if(n){var a=n.querySelector('audio');if(a)a.play();}\"")
auto = "autoplay " if autoplay else ""
# Serve audio via gradio's file route instead of a data: URI — the URL ends
# with the real (verbose) filename, so right-click "Save audio as…" offers
# it instead of download.mp3, and the page stays light as history grows.
src = "gradio_api/file=" + urllib.parse.quote(entry["path"], safe="/")
audio_el = (f'<audio controls {auto}style="width:100%" {attrs} '
f'src="{src}"></audio>')
prompt_disp = html_lib.escape(entry["prompt"]) or "<i>(no prompt)</i>"
spec_core = ""
if entry.get("spec_b64"):
height = "height:56px;" if small else ""
tip = ("3-band tinted stereo mel · red=bass / green=mid / blue=high · "
"L top, R bottom · click to seek")
spec_core = (f'<div style="position:relative; cursor:pointer" title="{tip}" '
f'onclick="{_JS_SEEK}">'
f'<img src="data:image/png;base64,{entry["spec_b64"]}" '
f'style="width:100%; {height} display:block; image-rendering:pixelated; '
f'border:1px solid #333" alt="spectrogram"/>'
f'<div class="ph" style="position:absolute; top:0; bottom:0; left:0%; '
f'width:2px; background:#fff; pointer-events:none; '
f'box-shadow:0 0 4px rgba(0,0,0,.8)"></div>'
f'</div>')
if small:
row = (f'<div style="display:flex; gap:8px; align-items:center">'
f'<div style="flex:1; min-width:0">{audio_el}</div>'
f'<div style="flex:1; min-width:0">{spec_core}</div></div>')
cap = (f'<div style="font-size:0.8em; margin:2px 0; color:#888">'
f'{_ago(entry.get("ts"))} · {prompt_disp}{_neg_disp(entry)} · seed {entry["seed"]}'
f'{_meta_suffix(entry)}</div>')
style = f"padding:6px 8px;{' background:' + bg + ';' if bg else ''}"
return (f'<div class="blk" data-key="{entry["key"]}" style="{style}">'
f'{row}{cap}</div>')
return (f'<div class="blk" data-key="{entry["key"]}">'
f'{audio_el}{spec_core}</div>')
def render_history(hist, advance=False, loop=False):
if not hist:
return ""
# zebra striping instead of separators: light grey vs medium grey rows
items = "".join(
render_player(e, small=True, advance=advance, loop=loop,
bg="rgba(127,127,127,0.24)" if i % 2 else "rgba(127,127,127,0.08)")
for i, e in enumerate(hist))
boot = f'<img src="data:," style="display:none" onerror="{_JS_SCROLL_RESTORE}"/>'
return (f'<div style="font-weight:600; margin-top:14px">Previous generations ({len(hist)})</div>'
f'<div id="sa3-hist" onscroll="{_JS_SCROLL_RECORD}" '
f'style="max-height:480px; overflow-y:auto; position:relative; margin-top:6px; '
f'padding-right:6px">{boot}{items}</div>')
def render_queue_status(entry=None, generating=False):
"""Subdued status chip for the Infinite Radio queue — no audio/spectrogram,
just Generating…/Ready (the swap uses the entry held in server state)."""
if generating:
body = "generating…"
elif entry is not None:
p = html_lib.escape(entry["prompt"]) or "<i>(no prompt)</i>"
body = f"ready — {p}{_neg_disp(entry)} · seed {entry['seed']}{_meta_suffix(entry)}"
else:
return ""
return (f'<div style="margin-top:2px; padding:6px 10px; '
f'background:rgba(127,127,127,0.12); border-radius:6px; '
f'color:#888; font-size:0.85em">'
f'<b>Queued next</b> · {body}</div>')
# ── Gradio UI ──────────────────────────────────────────────────────────────
def build_ui(initial_dit: str, initial_decoder: str, initial_precision: str, *,
share: bool, default_seconds: float, default_steps: int,
server_port: int | None = None):
import gradio as gr
import random as _random
# Pre-warm the initial pipeline so the first click is fast — ON the owner
# thread, so all model state lives where generations run.
warm_T = valid_T_lat(default_seconds)
print(f" pre-warming {initial_dit}+{initial_decoder} @ {initial_precision} (T_lat={warm_T})...")
def _warm():
t5, tok = get_t5()
ids, mask = tok("")
t5h = t5(ids, mask)
get_dit(str(ensure_local(dit_rel(initial_dit, initial_precision))), warm_T,
t5h, mask.astype(np.float32), float(default_seconds),
1.0, 1.0, None, None, None, True)
get_decoder(initial_decoder, initial_precision)
try:
run_serial(_warm)
except Exception as e:
print(f" (warm-up skipped: {type(e).__name__}: {e})")
def on_dit_change(dit_name, cur_seconds, prev, mem, vis, *slot_vals):
"""Model switch: pair the decoder, clamp seconds, and swap the LoRA
panel to the new model — the outgoing model's slots (files, picks,
strengths, visibility) are snapshotted into per-model memory and the
incoming model's snapshot is restored (or an empty panel). Dropdowns
rescan loras/<model>/ (hidden when empty)."""
max_s = MAX_SECONDS.get(dit_name, 120)
ch = _lora_dd_choices(dit_name)
valid = {v for _, v in ch}
lbl = f"…or pick from loras/{LORA_DIR_NAMES.get(dit_name, dit_name)}/"
mem = dict(mem or {})
mem[prev] = {"vis": list(vis),
"slots": [list(slot_vals[3 * i: 3 * i + 3]) for i in range(3)]}
snap = mem.get(dit_name)
if snap:
nvis = list(snap["vis"])
slots = [list(s) for s in snap["slots"]]
for s in slots:
if s[1] not in valid:
s[1] = _DD_NONE
# a slot with an adapter must never be an invisible-but-active one
nvis = [v or bool(s[0] or s[1] != _DD_NONE)
for v, s in zip(nvis, slots)]
else:
nvis = [False, False, False]
slots = [[None, _DD_NONE, 1.0] for _ in range(3)]
dd_ups = [gr.update(choices=ch, value=slots[i][1], visible=len(ch) > 1,
label=lbl) for i in range(3)]
srow_ups = [gr.update(visible=bool(slots[i][0] or slots[i][1] != _DD_NONE))
for i in range(3)]
row_ups = [gr.update(visible=v) for v in nvis]
return (gr.update(value=DEFAULT_DECODERS.get(dit_name, "same-s")),
gr.update(maximum=max_s, value=min(cur_seconds, max_s)),
*dd_ups, *srow_ups, *row_ups,
gr.update(visible=not all(nvis)), nvis, dit_name, mem,
slots[0][0], slots[1][0], slots[2][0],
slots[0][2], slots[1][2], slots[2][2])
def _generate_entry(dit_name, precision, decoder_name, prompt, negative_prompt,
seconds, steps, seed_text, cfg, apg, sigma_max, init_noise,
a2a_audio, inpaint_audio, inp_start, inp_end,
output_opts, file_format, *lora_vals):
"""Run one generation and package it as a history entry.
Returns (entry, None) or (None, error_message)."""
prompt = (prompt or "").strip()
# Permissive by design: a generation should succeed with whatever IS set.
# Half-configured features are ignored with a visible note, never an error.
notes = []
lora_specs = (_lora_specs_from_ui(lora_vals, notes) if lora_vals else None)
if lora_specs and precision not in LORA_PRECISIONS:
# UI already hides LoRA under int8; this is the belt-and-braces guard for
# API callers / stale state. get_patched_dit would raise anyway — surface
# it as a clean note + skip the LoRA rather than failing the generation.
notes.append(f"LoRA ignored — precision {precision} can't be LoRA-merged "
f"(needs fp32 or w16a32)")
lora_specs = None
# blank or -1 → random seed, kept small (1-9999) for readability
try:
seed = int(seed_text.strip()) if seed_text and seed_text.strip() else -1
except ValueError:
seed = -1
notes.append("seed wasn't an integer — used a random one")
if seed == -1:
seed = _random.randint(1, 9999)
# backstop for API callers bypassing the slider maxima
max_s = MAX_SECONDS.get(dit_name, 120)
if seconds > max_s:
notes.append(f"seconds clamped to {dit_name}'s trained max ({max_s:g}s)")
seconds = float(max_s)
# sigma_max governs every generation's schedule start; when guide audio is
# present (a2a), init_noise_level overrides it (parent-repo gradio semantics).
sigma_max = float(init_noise) if a2a_audio else float(sigma_max)
if sigma_max < MIN_SIGMA:
which = "init_noise_level" if a2a_audio else "sigma_max"
notes.append(f"{which} 0 runs at {MIN_SIGMA} (model is undefined at t≈0)"
+ (" — output ≈ the re-encoded input" if a2a_audio else ""))
sigma_max = MIN_SIGMA
inpaint_range = None
if inpaint_audio and inp_end > inp_start and inp_start < seconds:
if inp_end > seconds:
notes.append(f"inpaint end clamped to the clip length ({seconds:g}s)")
inpaint_range = (float(inp_start), float(min(inp_end, seconds)))
elif inpaint_audio:
notes.append("inpainting ignored — set the start/end range sliders")
elif inp_end > inp_start:
notes.append("inpaint range ignored — no reference audio uploaded")
opts = output_opts or []
if ("Infinite Radio" in opts and seed_text and seed_text.strip()
and seed_text.strip() != "-1"):
notes.append("Infinite Radio with a fixed seed repeats the same clip — "
"clear the seed for endless variety")
mode = ("a2a+inpaint" if (a2a_audio and inpaint_range) else
"inpaint" if inpaint_range else
"audio-to-audio" if a2a_audio else "text-to-audio")
try:
audio_np, t = run_serial(
run_generation,
dit_name, decoder_name, precision, prompt, negative_prompt or "",
float(seconds), int(steps), seed, float(cfg), float(apg),
float(sigma_max), a2a_audio or None, inpaint_audio or None,
inpaint_range, lora_specs)
except LoraError as e:
return None, f"LoRA error: {e}"
except Exception as e:
return None, f"error: {type(e).__name__}: {e}"
if not np.isfinite(audio_np).all():
return None, "error: model produced non-finite audio (try a higher σmax or different seed)"
pcm = (np.clip(audio_np, -1, 1) * 32767.0).astype(np.int16).T # (T, 2)
basename = verbose_basename(prompt, negative_prompt, cfg, sigma_max, seed, precision)
out_path = OUTPUT_DIR / f"{basename}.wav"
_save_wav(pcm, out_path)
mime = "audio/wav"
if FFMPEG and file_format == FORMAT_MP3:
mp3_path = OUTPUT_DIR / f"{basename}.mp3"
r = subprocess.run(["ffmpeg", "-y", "-v", "error", "-i", str(out_path),
"-codec:a", "libmp3lame", "-q:a", "0", str(mp3_path)],
capture_output=True)
if r.returncode == 0 and mp3_path.exists():
out_path.unlink()
out_path, mime = mp3_path, "audio/mpeg"
else:
notes.append("mp3 encode failed — saved WAV instead")
spec_b64 = None
try:
spec_png = render_spectrogram_png(pcm, sample_rate=SAMPLE_RATE,
width=1200, height=240)
spec_b64 = base64.b64encode(spec_png).decode("ascii")
except Exception as e:
notes.append(f"spectrogram failed: {type(e).__name__}: {e}")
load_note = (f"DiT-load {t['dit_load_ms']:.0f} ms ·&nbsp; "
if t.get("dit_load_ms", 0) > 100 else "")
cached_tag = " (cached)" if t.get("enc_cached") else ""
enc_note = f"encode {t['enc_ms']:.0f} ms{cached_tag} ·&nbsp; " if t.get("enc_ms") else ""
prec_note = f"<b>precision</b>: {precision} ·&nbsp; " if precision != "fp32" else ""
apg_note = f" (apg {apg:g})" if apg != 1.0 else ""
cfg_note = f"cfg {cfg:g}{apg_note} ·&nbsp; " if cfg != 1.0 else ""
lora_note = (f"<b>lora</b>: {html_lib.escape(_lora_disp(lora_specs))} ·&nbsp; "
if lora_specs else "")
prompt_disp = html_lib.escape(prompt) or "<i>(no prompt)</i>"
neg_disp = (f' · <span style="opacity:0.75">neg: {html_lib.escape(negative_prompt.strip())}</span>'
if negative_prompt and negative_prompt.strip() and cfg != 1.0 else "")
timing_html = (
f"{prompt_disp}{neg_disp} ·&nbsp; "
f"{prec_note}{load_note}{enc_note}{cfg_note}{lora_note}"
f"<b>Inference</b>: {t['inference_ms']:.0f} ms "
f"<span style='color:#888'>(t5={t['t5_ms']:.0f} · sample={t['sample_ms']:.0f} · "
f"decode={t['decode_ms']:.0f})</span> ·&nbsp; "
f"<b>{t['realtime']:.1f}× realtime</b> ·&nbsp; "
f"<b>seed</b>: <code>{seed}</code> ·&nbsp; "
f"<b>seq_len</b>: {t['T_lat']} ·&nbsp; <b>samples</b>: {t['samples']}"
)
if notes:
timing_html = ("".join(f"<div style='color:#fa3; font-size:0.85em'>note: {n}</div>"
for n in notes) + timing_html)
entry = {
"key": f"k{time.time_ns()}",
"ts": time.time(),
"dit": dit_name,
"prec": precision,
"neg": (negative_prompt or "").strip(),
"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": seed,
"lora": _lora_disp(lora_specs),
"spec_b64": spec_b64,
"timing": timing_html,
}
return entry, None
def _present(state, entry, opts, queued_panel, *, force_autoplay=False):
"""Make `entry` the current clip (previous current moves to history top)
and render all output panels."""
if state["current"] is not None:
state["history"].insert(0, state["current"])
state["current"] = entry
loop = "Loop" in opts
main = render_player(entry,
autoplay=force_autoplay or "Auto-play" in opts,
autodl="Auto-download" in opts,
radio="Infinite Radio" in opts and not loop,
loop=loop,
hotswap="Hotswap" in opts)
return (main, entry["timing"], "",
render_history(state["history"], advance="Auto-play" in opts, loop=loop),
queued_panel, state)
def generate(dit_name, precision, decoder_name, prompt, negative_prompt,
seconds, steps, seed_text, cfg, apg, sigma_max, init_noise,
a2a_audio, inpaint_audio, inp_start, inp_end,
output_opts, file_format, *lora_and_state):
*lora_vals, state = lora_and_state
entry, err_msg = _generate_entry(
dit_name, precision, decoder_name, prompt, negative_prompt, seconds, steps,
seed_text, cfg, apg, sigma_max, init_noise, a2a_audio,
inpaint_audio, inp_start, inp_end, output_opts, file_format,
*lora_vals)
if entry is None:
return (gr.update(), gr.update(),
f"<span style='color:#f88'>{err_msg}</span>",
gr.update(), gr.update(), state)
# a manual Generate makes any queued clip stale (settings may have changed);
# the chained pregen refills it when Infinite Radio is on
state["queued"] = None
opts = output_opts or []
queued_panel = render_queue_status(generating=True) if "Infinite Radio" in opts else ""
return _present(state, entry, opts, queued_panel)
def promote(dit_name, precision, decoder_name, prompt, negative_prompt,
seconds, steps, seed_text, cfg, apg, sigma_max, init_noise,
a2a_audio, inpaint_audio, inp_start, inp_end,
output_opts, file_format, *lora_and_state):
"""Infinite Radio: swap the pre-generated clip in when playback ends.
Falls back to a full generate if the queue is empty. Either way the
next track ALWAYS autoplays — Auto-play only governs whether a clip
starts when nothing was playing; radio transitions are continuations."""
*lora_vals, state = lora_and_state
q = state.get("queued")
if q is None:
entry, err_msg = _generate_entry(
dit_name, precision, decoder_name, prompt, negative_prompt, seconds, steps,
seed_text, cfg, apg, sigma_max, init_noise, a2a_audio,
inpaint_audio, inp_start, inp_end, output_opts, file_format,
*lora_vals)
if entry is None:
return (gr.update(), gr.update(),
f"<span style='color:#f88'>{err_msg}</span>",
gr.update(), gr.update(), state)
q = entry
state["queued"] = None
return _present(state, q, output_opts or [],
render_queue_status(generating=True), force_autoplay=True)
def pregen(dit_name, precision, decoder_name, prompt, negative_prompt,
seconds, steps, seed_text, cfg, apg, sigma_max, init_noise,
a2a_audio, inpaint_audio, inp_start, inp_end,
output_opts, file_format, *lora_and_state):
"""Chained after generate/promote: pre-generate the NEXT clip while the
current one plays (Infinite Radio only)."""
*lora_vals, state = lora_and_state
if "Infinite Radio" not in (output_opts or []):
state["queued"] = None
return "", state
entry, err_msg = _generate_entry(
dit_name, precision, decoder_name, prompt, negative_prompt, seconds, steps,
seed_text, cfg, apg, sigma_max, init_noise, a2a_audio,
inpaint_audio, inp_start, inp_end, output_opts, file_format,
*lora_vals)
if entry is None:
return (f"<div style='color:#f88; font-size:0.85em'>queue: {err_msg}</div>",
state)
state["queued"] = entry
return render_queue_status(entry), state
# sa3-promote must stay MOUNTED for the onended click to find it —
# visible=False would remove it from the DOM entirely, so hide via CSS.
# sa3-out: collapse the column's flex gap + wrapper padding so the
# spectrogram caption / timing line / queued status sit tightly together.
_css = ("#sa3-promote{display:none !important}"
"#sa3-out{gap:4px !important}"
"#sa3-out .html-container{padding:0 !important; margin:0 !important}")
with gr.Blocks(title="SA3 TFLite") as demo:
gr.Markdown(
"# SA3 TFLite — portable CPU (XNNPACK)\n"
"Text-to-audio, CFG + negative prompt, audio-to-audio, inpainting, LoRA. "
"Pick a precision (fp32 / w16a32 / w8a32 / w8a8-dyn). First use of a "
"model/precision loads weights; subsequent runs are cached."
)
st = gr.State({"current": None, "queued": None, "history": []})
with gr.Row():
with gr.Column(scale=3):
with gr.Row():
dit_dd = gr.Dropdown(label="DiT model", choices=list(DIT_REL.keys()),
value=initial_dit, scale=1)
precision_dd = gr.Dropdown(label="Precision", choices=list(PRECISIONS),
value=initial_precision, scale=1)
decoder_dd = gr.Dropdown(label="Decoder (codec)",
choices=list(DEC_REL.keys()),
value=initial_decoder, scale=1)
with gr.Row():
prompt = gr.Textbox(label="Prompt", lines=2, scale=6,
placeholder="e.g. 'Impending tribal, epic orchestral buildup'")
seed = gr.Textbox(label="Seed (optional)", max_lines=1, value="",
scale=1, min_width=80)
with gr.Row():
seconds = gr.Slider(label="Seconds", minimum=1,
maximum=MAX_SECONDS.get(initial_dit, 120),
value=default_seconds, step=1)
steps = gr.Slider(label="Steps", minimum=1, maximum=16,
value=default_steps, step=1)
cfg = gr.Slider(label="CFG", minimum=0.0, maximum=10.0,
value=1.0, step=0.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, value=1.0, step=0.05)
sigma_global = gr.Slider(label="sigma_max",
minimum=0.0, maximum=1.0, value=1.0, step=0.01)
negative_prompt = gr.Textbox(label="Negative prompt", lines=1)
with gr.Accordion("LoRA", open=False):
# Shown only under a quantized precision (LoRA disabled there).
lora_note = gr.Markdown(
"*LoRA needs precision **fp32** or **w16a32** — quantized-int8 "
"graphs (w8a32 / w8a8-dyn) can't be LoRA-merged.*",
visible=initial_precision not in LORA_PRECISIONS)
with gr.Group(visible=initial_precision in LORA_PRECISIONS) as lora_editor:
_dd0 = _lora_dd_choices(initial_dit)
lora_inputs, lora_rows, lora_rm_btns = [], [], []
lora_dds, lora_files, lora_srows = [], [], []
for _i in range(1, 4):
with gr.Group(visible=False) as _grp:
with gr.Row(equal_height=True):
_lf = gr.File(label=f"LoRA {_i} (.safetensors)",
file_types=[".safetensors"],
type="filepath", scale=1, height=88)
_ld = gr.Dropdown(
label=f"…or pick from loras/{LORA_DIR_NAMES[initial_dit]}/",
choices=_dd0, value=_DD_NONE, scale=1,
visible=len(_dd0) > 1)
_rm = gr.Button("✕", size="sm", scale=0,
min_width=36)
with gr.Row(visible=False) as _srow:
_ls = gr.Slider(label="strength", minimum=0.0,
maximum=10.0, value=1.0, step=0.1,
scale=2, min_width=110)
lora_rows.append(_grp)
lora_rm_btns.append(_rm)
lora_dds.append(_ld)
lora_files.append(_lf)
lora_srows.append(_srow)
lora_inputs += [_lf, _ld, _ls]
lora_add_btn = gr.Button("+ Add LoRA", size="sm")
lora_vis = gr.State([False, False, False])
# Per-model LoRA memory: switching DiT saves the outgoing
# model's slots and restores the incoming model's.
lora_mem = gr.State({})
prev_dit = gr.State(initial_dit)
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")
sigma_slider = gr.Slider(
label="init_noise_level (1.0 = prompt, ~0.92 = fusion, 0 = input)",
minimum=0.0, maximum=1.0, value=0.92, step=0.01)
with gr.Accordion("Inpainting (regenerate a span of reference audio)", open=False):
inpaint_audio = gr.Audio(label="Reference audio — kept bit-exact outside the range",
type="filepath")
with gr.Row():
inp_start = gr.Slider(label="Start (s)", minimum=0, maximum=120,
value=0, step=0.5)
inp_end = gr.Slider(label="End (s)", minimum=0, maximum=120,
value=0, step=0.5)
with gr.Accordion("Output", open=False):
output_opts = gr.CheckboxGroup(
["Auto-play", "Auto-download", "Infinite Radio", "Loop", "Hotswap"],
value=["Auto-play"], label="Playback Options")
opts_state = gr.State(["Auto-play"])
file_format = gr.Radio(
[FORMAT_MP3, FORMAT_WAV] if FFMPEG else [FORMAT_WAV],
value=FORMAT_MP3 if FFMPEG else FORMAT_WAV,
label="Format", visible=FFMPEG)
with gr.Column(scale=2, elem_id="sa3-out"):
generate_btn = gr.Button("Generate", variant="primary", size="lg",
elem_id="sa3-generate")
promote_btn = gr.Button("", elem_id="sa3-promote") # CSS-hidden, DOM-present
gr.Markdown("**Output**")
output_player = gr.HTML()
timing = gr.HTML()
error_box = gr.HTML()
queued_html = gr.HTML()
history_html = gr.HTML()
def _sync_dd_vis(dit_name):
"""Second batch after any group (re)mount: gradio 6 drops visibility
updates sent to components inside a hidden container (the component
isn't in the DOM), so dropdown show/hide must be applied AFTER the
slot groups are mounted — chained via .then() on every event that
reveals them."""
v = len(_lora_dd_choices(dit_name)) > 1
return [gr.update(visible=v) for _ in range(3)]
dit_dd.change(on_dit_change,
inputs=[dit_dd, seconds, prev_dit, lora_mem,
lora_vis] + lora_inputs,
outputs=[decoder_dd, seconds] + lora_dds + lora_srows
+ lora_rows
+ [lora_add_btn, lora_vis, prev_dit, lora_mem]
+ lora_files
+ [lora_inputs[3 * i + 2] for i in range(3)]
).then(_sync_dd_vis, inputs=[dit_dd], outputs=lora_dds
).then(None, js=_JS_FIX_SLIDERS)
def on_precision_change(prec):
"""Quantized-int8 precisions can't be LoRA-merged — hide the editor
and show the note. fp32 / w16a32 re-enable it."""
disabled = prec not in LORA_PRECISIONS
return gr.update(visible=disabled), gr.update(visible=not disabled)
precision_dd.change(on_precision_change, inputs=[precision_dd],
outputs=[lora_note, lora_editor]
).then(_sync_dd_vis, inputs=[dit_dd], outputs=lora_dds
).then(None, js=_JS_FIX_SLIDERS)
def on_seconds_change(sec):
return gr.update(maximum=sec), gr.update(maximum=sec)
seconds.change(on_seconds_change, inputs=[seconds], outputs=[inp_start, inp_end])
def on_opts_change(opts, prev):
"""Loop and Infinite Radio are mutually exclusive — the one just
ticked wins, the other unticks."""
opts = opts or []
if "Loop" in opts and "Infinite Radio" in opts:
added = set(opts) - set(prev or [])
drop = "Infinite Radio" if "Loop" in added else "Loop"
opts = [o for o in opts if o != drop]
return gr.update(value=opts), opts
return gr.update(), opts
output_opts.change(on_opts_change, inputs=[output_opts, opts_state],
outputs=[output_opts, opts_state])
def lora_add(vis, dit_name, f1, d1, f2, d2, f3, d3):
"""Reveal the first hidden LoRA slot; hide the button at 3/3.
Slider-row and dropdown visibility are recomputed from slot content
and the current model's library on every add, so a re-opened slot
can never inherit a stale visible state."""
vis = list(vis)
for i, v in enumerate(vis):
if not v:
vis[i] = True
break
srows = [gr.update(visible=bool(f or (d and d != _DD_NONE)))
for f, d in ((f1, d1), (f2, d2), (f3, d3))]
return ([gr.update(visible=v) for v in vis] + srows
+ [gr.update(visible=not all(vis)), vis])
lora_add_btn.click(lora_add,
inputs=[lora_vis, dit_dd]
+ [c for i in range(3)
for c in lora_inputs[3 * i: 3 * i + 2]],
outputs=lora_rows + lora_srows
+ [lora_add_btn, lora_vis]
).then(_sync_dd_vis, inputs=[dit_dd], outputs=lora_dds
).then(None, js=_JS_FIX_SLIDERS)
def _lora_remove(idx):
def _rm(vis):
vis = list(vis)
vis[idx] = False
return ([gr.update(visible=v) for v in vis]
+ [gr.update(visible=True), vis,
# reset the slot: file, dropdown, strength (srow hidden)
None, gr.update(value=_DD_NONE), 1.0,
gr.update(visible=False)])
_rm.__name__ = f"lora_remove_{idx + 1}"
return _rm
for _idx, _btn in enumerate(lora_rm_btns):
_btn.click(_lora_remove(_idx), inputs=[lora_vis],
outputs=lora_rows + [lora_add_btn, lora_vis]
+ lora_inputs[3 * _idx: 3 * _idx + 3]
+ [lora_srows[_idx]])
# A slot's strength slider appears once an adapter is loaded — via file
# drop OR dropdown pick — and the two sources are exclusive (loading one
# resets the other). All three events are user-driven (upload/clear/input),
# so the programmatic resets here can't loop.
def _lora_wire_slot(idx):
lf, ld, srow = lora_files[idx], lora_dds[idx], lora_srows[idx]
def _up(f):
return gr.update(visible=bool(f)), gr.update(value=_DD_NONE)
_up.__name__ = f"lora_file_{idx + 1}"
lf.upload(_up, inputs=[lf], outputs=[srow, ld]
).then(None, js=_JS_FIX_SLIDERS)
def _cl(dd):
return gr.update(visible=bool(dd and dd != _DD_NONE))
_cl.__name__ = f"lora_file_clear_{idx + 1}"
lf.clear(_cl, inputs=[ld], outputs=[srow])
def _pick(dd, f):
loaded = bool(dd and dd != _DD_NONE)
return (gr.update(visible=loaded or bool(f)),
gr.update(value=None) if loaded else gr.update())
_pick.__name__ = f"lora_pick_{idx + 1}"
ld.input(_pick, inputs=[ld, lf], outputs=[srow, lf]
).then(None, js=_JS_FIX_SLIDERS)
for _idx in range(3):
_lora_wire_slot(_idx)
ctrl_inputs = [dit_dd, precision_dd, decoder_dd, prompt, negative_prompt,
seconds, steps, seed, cfg, apg, sigma_global,
sigma_slider, a2a_audio, inpaint_audio,
inp_start, inp_end, output_opts, file_format] + lora_inputs
main_outputs = [output_player, timing, error_box, queued_html, history_html, st]
# NB: _present returns (player, timing, err, history, queued, state) —
# map onto components in that order.
def _reorder(fn):
def wrapped(*args):
player, tim, err_, hist, queued, state = fn(*args)
return player, tim, err_, queued, hist, state
wrapped.__name__ = fn.__name__
return wrapped
generate_btn.click(_reorder(generate), inputs=ctrl_inputs + [st],
outputs=main_outputs
).then(pregen, inputs=ctrl_inputs + [st],
outputs=[queued_html, st])
promote_btn.click(_reorder(promote), inputs=ctrl_inputs + [st],
outputs=main_outputs
).then(pregen, inputs=ctrl_inputs + [st],
outputs=[queued_html, st])
gr.Markdown(
"<p style='color:#888; font-size:0.85em'>"
"WAVs saved under <code>output/gradio/</code>. "
"Models run on CPU (XNNPACK). Precision applies to the DiT + codec "
"(+ encoder for a2a/inpaint); T5Gemma is fp16.</p>"
)
# gradio 6 moved `css` from the Blocks constructor to launch().
demo.queue(max_size=16).launch(share=share, server_name="0.0.0.0",
server_port=server_port, css=_css,
allowed_paths=[str(OUTPUT_DIR)],
prevent_thread_lock=False, show_error=True)
def main():
ap = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--dit", choices=list(DIT_REL.keys()), default="sm-music",
help="Initial DiT bundle (switchable at runtime)")
ap.add_argument("--decoder", choices=list(DEC_REL.keys()), default=None,
help="Initial decoder. Default: pairs with --dit")
ap.add_argument("--precision", choices=list(PRECISIONS), default="fp32",
help="Initial precision (switchable at runtime): fp32 (default, "
"CPU fast+accurate) | w16a32 (fp16, ≈lossless, half size) | "
"w8a32 / w8a8-dyn (GPTQ int8, ¼ size)")
ap.add_argument("--default-seconds", type=float, default=30.0,
help="Length to pre-warm the initial DiT at")
ap.add_argument("--default-steps", type=int, default=8)
ap.add_argument("--threads", type=int, default=8, help="XNNPACK CPU threads")
ap.add_argument("--port", type=int, default=None,
help="Server port (default: gradio picks 7860+, auto-incrementing)")
ap.add_argument("--share", action=argparse.BooleanOptionalAction, default=True,
help="Create a public gradio.live URL (default on)")
args = ap.parse_args()
if args.decoder is None:
args.decoder = DEFAULT_DECODERS[args.dit]
global _THREADS
_THREADS = args.threads
print(f"\n━━━ SA3 TFLite — gradio ━━━")
print(f" initial dit: {args.dit}")
print(f" initial decoder: {args.decoder}")
print(f" initial precision:{args.precision}")
print(f" threads: {args.threads}")
print(f" models: {', '.join(DIT_REL.keys())} (runtime-switchable)")
build_ui(args.dit, args.decoder, args.precision, share=args.share,
default_seconds=args.default_seconds, default_steps=args.default_steps,
server_port=args.port)
if __name__ == "__main__":
main()