SoundGame Dev
f0f0a9b50e
- 共享引擎（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 模板
933 lines
49 KiB
Python
933 lines
49 KiB
Python
"""SA3 text-to-audio (+ audio-to-audio + inpainting + CFG) on CPU via the BAKED-I/O
|
||
varlen TFLite / LiteRT models — the portable-CPU sibling of the MLX and TensorRT
|
||
releases. CLI at feature parity with sa3_mlx.py / sa3_trt.py, sharing their flag names.
|
||
|
||
Baked I/O (ONNX / TensorRT convention — conditioner + patch/unpatch are IN-GRAPH):
|
||
DiT (6-input): x[1,256,L], t[1], t5_hidden[1,256,768], t5_mask[1,256],
|
||
seconds_total[1], local_add_cond[1,257,L] -> velocity[1,256,L]
|
||
(feed RAW T5 outputs; prompt-padding + seconds-embed happen in-graph)
|
||
Encoder (audio-in): audio[1,2,N] -> latents[1,256,N/4096] (patch-encode baked)
|
||
Decoder (audio-out): latents[1,256,L] -> audio[1,2,L*4096] (unpatch baked)
|
||
T5Gemma: t5gemma/encoder_fp16.tflite (fixed 256 text tokens; tokenizer bundled)
|
||
|
||
Modes (identical flags to sa3_mlx / sa3_trt):
|
||
text-to-audio --prompt P
|
||
audio-to-audio --prompt P --init-audio IN.wav [--init-noise-level σ]
|
||
inpainting --prompt P --init-audio IN.wav --inpaint-range START,END
|
||
negative CFG --prompt P --cfg N [--negative-prompt P_NEG] [--apg A]
|
||
|
||
CFG uses the TRT/ONNX baked-conditioner convention: the uncond branch feeds the negative
|
||
prompt's T5 output (or an all-zero hidden+mask, which the in-graph conditioner turns into
|
||
learned padding embeddings) — it does NOT zero cross_attn like MLX (our DiT bakes the
|
||
conditioner, so there's no raw cross_attn). CFG runs cond+uncond as ONE batch=2 invoke on the
|
||
variable-batch DiT (--cfg-batched, default; ~7-29%% faster on Apple-Silicon AMX) or as a
|
||
sequential batch=1 dual-pass (--no-cfg-batched). Bit-identical at fp32/w8a32; ~80 dB at
|
||
w16a32; w8a8-dyn diverges by design (activation scales shared across the batch).
|
||
|
||
All model files auto-download from HuggingFace (stabilityai/stable-audio-3-optimized)
|
||
on first use and symlink into models/tflite/ from the HF cache. See scripts/weights.py.
|
||
"""
|
||
from __future__ import annotations
|
||
import argparse, math, os, random, re, subprocess, sys, time, wave
|
||
from pathlib import Path
|
||
import numpy as np
|
||
|
||
# Windows consoles default to a legacy code page (cp1252/cp437) that can't encode
|
||
# the banners' box-drawing / emoji characters. Force UTF-8 with replacement so no
|
||
# print can crash the CLI. No-op on macOS/Linux, where stdout is UTF-8 already.
|
||
for _stream in (sys.stdout, sys.stderr):
|
||
try:
|
||
_stream.reconfigure(encoding="utf-8", errors="replace")
|
||
except AttributeError:
|
||
pass # exotic stream without reconfigure() (non-TextIOWrapper) — leave as-is
|
||
|
||
# Enable ANSI/VT escape processing in legacy Windows consoles (cmd.exe). Calling
|
||
# os.system("") flips the console into VT mode; harmless elsewhere and on
|
||
# modern Windows Terminal, which has it on already.
|
||
if os.name == "nt":
|
||
os.system("")
|
||
|
||
# The arrow-key picker below uses POSIX terminal APIs. Keep the non-interactive
|
||
# CLI usable on Windows, where those modules are not available.
|
||
try:
|
||
import termios
|
||
import tty
|
||
except ImportError:
|
||
termios = tty = None
|
||
|
||
REPO = Path(__file__).resolve().parent.parent # project root (scripts/ is one level down)
|
||
sys.path.insert(0, str(REPO)) # so `from models.defs.* import` resolves
|
||
sys.path.insert(0, str(REPO / "scripts")) # so `from weights import *` resolves
|
||
|
||
from models.defs import tflite_pipeline as P # Tokenizer, T5GemmaTFLite, build_pingpong_schedule, make_noise, sample, save_wav
|
||
from weights import ensure_local, is_present, PRECISIONS, dit_rel, dec_rel, enc_rel
|
||
|
||
SAMPLE_RATE = 44100
|
||
SAMPLES_PER_LATENT = 4096
|
||
COND_TOKENS = 256
|
||
COND_DIM = 768
|
||
MIN_SIGMA = 0.01 # rf_denoiser is undefined at t≈0 → NaN below this
|
||
|
||
# ─── Model manifest (local rel paths; resolved via weights.ensure_local at load) ───
|
||
# --precision picks the DiT + decoder variant (same flag as the TensorRT release's
|
||
# --precision; values here use the wXaY convention — weights/activations bit-widths,
|
||
# "16" = fp16, there is no int16). Encoders + T5Gemma have one precision, like TRT.
|
||
# fp32 reference models (default — on CPU fp16 is SLOWER than fp32, so unlike
|
||
# TRT the fastest-and-accurate choice is fp32)
|
||
# w16a32 fp16 weights / fp32 activations — half size, ≈lossless, 1.5-3× slower
|
||
# on an M4 Pro (legacy name: fp16; ≈ TRT's fp16 in spirit,
|
||
# storage-only here)
|
||
# w8a32 GPTQ weight-only int8 — ¼ size at fp32 speed (legacy: woint8)
|
||
# w8a8-dyn GPTQ dynamic int8 — fastest (~1.2-1.3×), lowest quality (legacy: dynint8)
|
||
# PRECISIONS + the path builders live in weights.py (single source of truth with
|
||
# the download manifest — a precision the CLI accepts is guaranteed downloadable).
|
||
DIT_REL = {d: dit_rel(d) for d in ("sm-music", "sm-sfx", "medium")} # fp32 (picker)
|
||
DEC_REL = {d: dec_rel(d) for d in ("same-s", "same-l")}
|
||
ENC_REL = {d: enc_rel(d) for d in ("same-s", "same-l")}
|
||
T5_REL = "models/tflite/t5gemma/encoder_fp16.tflite"
|
||
DEFAULT_DECODER = {"sm-music": "same-s", "sm-sfx": "same-s", "medium": "same-l"}
|
||
|
||
# SAME-L dense SWA mask is O(S^2): chunk long decodes. SAME-S has a narrow field: whole.
|
||
SAMEL_CHUNK = 64 # latent tokens/window — throughput optimum (6.5x RT) vs the O(S^2) dense SWA mask
|
||
SAMEL_OVERLAP = 8 # symmetric latent-token context each interior side (SAME-L needs >=8)
|
||
|
||
|
||
def _free():
|
||
import gc; gc.collect()
|
||
|
||
|
||
# ─── ANSI display (match sa3_mlx.py style) ─────────────────────────────────
|
||
_USE_COLOR = sys.stdout.isatty()
|
||
_RULE_W = 64
|
||
def _c(code, s): return f"\x1b[{code}m{s}\x1b[0m" if _USE_COLOR else s
|
||
def bold(s): return _c("1", s)
|
||
def dim(s): return _c("2", s)
|
||
def cyan(s): return _c("36", s)
|
||
def yellow(s): return _c("33", s)
|
||
def green(s): return _c("32", s)
|
||
def magenta(s): return _c("35", s)
|
||
def red(s): return _c("31", s)
|
||
def rule(ch="━"): print(cyan(ch * _RULE_W))
|
||
def banner(t):
|
||
rule(); print(f" {bold(t)}"); rule()
|
||
def stage(idx, label, ms=None):
|
||
head = f" {cyan(idx)} {bold(label)}"
|
||
if ms is None:
|
||
print(head, flush=True); return
|
||
visible = len(f" {idx} {label}")
|
||
dots = dim("·" * max(2, _RULE_W - visible - 9))
|
||
print(f"{head} {dots} {yellow(f'{ms:>5.0f} ms')}", flush=True)
|
||
def sub(t): print(f" {dim(t)}", flush=True)
|
||
|
||
|
||
# ─── Interactive arrow-key picker (from sa3_mlx.py) ────────────────────────
|
||
def _arrow_pick(prompt: str, options: list[str], default: str | None = None) -> str:
|
||
"""Tiny arrow-key picker — no external deps, posix termios only.
|
||
|
||
Up/Down to move, Enter to select, Ctrl-C to abort. Falls back to a
|
||
numeric prompt when stdin isn't a TTY (piped input, CI, etc.).
|
||
"""
|
||
if termios is None or tty is None or not sys.stdin.isatty():
|
||
print(prompt)
|
||
for i, o in enumerate(options):
|
||
mark = "*" if o == default else " "
|
||
print(f" {mark} [{i}] {o}")
|
||
s = input(f"Choose [0-{len(options)-1}] (Enter for default): ").strip()
|
||
if s == "":
|
||
return default or options[0]
|
||
if s.isdigit() and 0 <= int(s) < len(options):
|
||
return options[int(s)]
|
||
return s if s in options else (default or options[0])
|
||
|
||
idx = options.index(default) if default in options else 0
|
||
fd = sys.stdin.fileno()
|
||
old = termios.tcgetattr(fd)
|
||
print(prompt)
|
||
for _ in options:
|
||
print()
|
||
try:
|
||
tty.setcbreak(fd)
|
||
while True:
|
||
sys.stdout.write(f"\x1b[{len(options)}A")
|
||
for i, o in enumerate(options):
|
||
if i == idx:
|
||
sys.stdout.write(f"\x1b[2K\x1b[36m▶ {o}\x1b[0m\n")
|
||
else:
|
||
sys.stdout.write(f"\x1b[2K {o}\n")
|
||
sys.stdout.flush()
|
||
ch = sys.stdin.read(1)
|
||
if ch == "\x1b":
|
||
seq = sys.stdin.read(2)
|
||
if seq == "[A":
|
||
idx = (idx - 1) % len(options)
|
||
elif seq == "[B":
|
||
idx = (idx + 1) % len(options)
|
||
elif ch in ("\n", "\r"):
|
||
return options[idx]
|
||
elif ch == "\x03":
|
||
raise KeyboardInterrupt
|
||
finally:
|
||
termios.tcsetattr(fd, termios.TCSADRAIN, old)
|
||
|
||
|
||
def prompt_user_if_missing(args):
|
||
"""Interactive arrow-key selection when --dit / --decoder aren't supplied."""
|
||
if args.dit is None:
|
||
args.dit = _arrow_pick("Choose DiT model:", list(DIT_REL.keys()), default="sm-music")
|
||
print(f" → {args.dit}")
|
||
if args.decoder is None:
|
||
suggested = DEFAULT_DECODER[args.dit]
|
||
args.decoder = _arrow_pick("Choose audio decoder:", list(DEC_REL.keys()), default=suggested)
|
||
print(f" → {args.decoder}")
|
||
if args.seed is None:
|
||
args.seed = random.randint(0, 2**31 - 1)
|
||
return args
|
||
|
||
|
||
def _preflight_download(args, dec) -> None:
|
||
"""Resolve every model file this run needs and download any that are missing —
|
||
BEFORE the banner prints and BEFORE the wall-clock starts. Network time then
|
||
isn't charged against '×realtime' and the user sees download progress as a
|
||
clearly separate setup step."""
|
||
needed = [T5_REL, dit_rel(args.dit, args.dit_precision), dec_rel(dec, args.decoder_precision)]
|
||
if args.init_audio:
|
||
needed.append(enc_rel(dec, args.encoder_precision))
|
||
missing = [p for p in needed if not is_present(p)]
|
||
if not missing:
|
||
return
|
||
print(f" Fetching {len(missing)} missing model file(s) before starting:")
|
||
for rel in missing:
|
||
ensure_local(rel)
|
||
print()
|
||
|
||
|
||
# ─── WAV read (16-bit PCM fast path + ffmpeg fallback) ─────────────────────
|
||
def read_wav(path: str) -> np.ndarray:
|
||
"""Return (2, T) float32 in [-1, 1]. 16-bit/44.1k native; else via ffmpeg. Mono→stereo."""
|
||
try:
|
||
with wave.open(path, "rb") as w:
|
||
nch, sw, sr, n = w.getnchannels(), w.getsampwidth(), w.getframerate(), w.getnframes()
|
||
if sr == SAMPLE_RATE and sw == 2:
|
||
raw = np.frombuffer(w.readframes(n), np.int16).astype(np.float32) / 32767.0
|
||
if nch == 1:
|
||
return np.stack([raw, raw], 0)
|
||
return raw.reshape(-1, nch).T[:2]
|
||
except wave.Error:
|
||
pass
|
||
try:
|
||
r = subprocess.run(["ffmpeg", "-v", "error", "-i", path, "-f", "s16le",
|
||
"-ar", str(SAMPLE_RATE), "-ac", "2", "-"],
|
||
capture_output=True, check=True)
|
||
except FileNotFoundError:
|
||
ffmpeg_hint = {"darwin": "brew install ffmpeg",
|
||
"win32": "winget install ffmpeg (or: choco install ffmpeg)",
|
||
}.get(sys.platform, "apt install ffmpeg")
|
||
raise RuntimeError(f"{path}: unsupported WAV format. Install ffmpeg for 24/32-bit or "
|
||
f"non-44.1kHz input: {ffmpeg_hint}")
|
||
except subprocess.CalledProcessError as e:
|
||
raise RuntimeError(f"{path}: ffmpeg failed — {e.stderr.decode().strip()}")
|
||
raw = np.frombuffer(r.stdout, np.int16).astype(np.float32) / 32767.0
|
||
return raw.reshape(-1, 2).T
|
||
|
||
|
||
# ─── Baked DiT backend (6-input; batched OR sequential CFG + inpaint local_add_cond) ───
|
||
def _apply_cfg(x, t, v_cond, v_uncond, cfg, apg):
|
||
"""Combine cond/uncond velocities in denoised space (RF), optional APG. fp32.
|
||
Matches sa3_trt_core.model_fn. Returns cfg_v so `denoised = x - t*v` is the guided one."""
|
||
sigma = np.float32(t)
|
||
xf = x.astype(np.float32)
|
||
cond_d = xf - v_cond.astype(np.float32) * sigma
|
||
uncond_d = xf - v_uncond.astype(np.float32) * sigma
|
||
diff = cond_d - uncond_d
|
||
if apg <= 0.0:
|
||
cfg_diff = diff
|
||
else:
|
||
norm = np.sqrt((cond_d * cond_d).sum(axis=(-2, -1), keepdims=True))
|
||
unit = cond_d / np.maximum(norm, 1e-8)
|
||
parallel = (diff * unit).sum(axis=(-2, -1), keepdims=True) * unit
|
||
diff_orth = diff - parallel
|
||
cfg_diff = diff_orth if apg >= 1.0 else (apg * diff_orth + (1.0 - apg) * diff)
|
||
cfg_d = cond_d + (cfg - 1.0) * cfg_diff
|
||
return ((xf - cfg_d) / sigma).astype(np.float32)
|
||
|
||
|
||
class BakedDiT:
|
||
"""model_fn(x,t,cross,gcond)->v compatible with P.sample. cross/gcond are IGNORED
|
||
(conditioning is baked in-graph, driven by the raw T5 outputs held here).
|
||
|
||
cfg==1.0 -> one batch=1 forward with the (cond) T5 output.
|
||
cfg!=1.0 -> CFG combining a cond and an uncond (null_hidden/null_mask) velocity in
|
||
denoised space with optional APG. Two backends (bit-identical output at
|
||
fp32/w8a32; ~80 dB at w16a32; w8a8-dyn diverges — batch-shared activation
|
||
scales; auto-falls-back to batch=1 for cfg==1.0):
|
||
batched=True (default): ONE batch=2 invoke/step — cond=row0, uncond=row1 — on the
|
||
canonical variable-batch DiT (its batch axis is dynamic). ~7-29%% faster
|
||
on Apple-Silicon (the AMX matrix unit amortizes weight loads across the 2 rows).
|
||
batched=False: SEQUENTIAL dual-pass — two batch=1 invokes/step, like TensorRT (whose
|
||
engine is static-batch=1). Use if a backend/delegate dislikes batch=2.
|
||
Returns cfg_v so P.sample's `denoised = x - t*v` yields the guided denoised."""
|
||
def __init__(self, path, L, t5_hidden, t5_mask, seconds, threads=8,
|
||
cfg=1.0, apg=1.0, null_hidden=None, null_mask=None, local_add_cond=None,
|
||
batched=True):
|
||
from ai_edge_litert import interpreter as tfl
|
||
self.it = tfl.Interpreter(model_path=str(path), num_threads=threads)
|
||
det = self.it.get_input_details()
|
||
def pick(pred): return [d for d in det if pred(d)]
|
||
self.i_x = pick(lambda d: len(d["shape"]) == 3 and d["shape"][1] == 256 and d["shape"][2] != COND_DIM)[0]
|
||
self.i_lac = pick(lambda d: len(d["shape"]) == 3 and d["shape"][1] == 257)[0]
|
||
self.i_t5h = pick(lambda d: len(d["shape"]) == 3 and d["shape"][2] == COND_DIM)[0]
|
||
self.i_t5m = pick(lambda d: len(d["shape"]) == 2)[0]
|
||
scalars = sorted(pick(lambda d: len(d["shape"]) == 1), key=lambda d: d["name"])
|
||
self.i_t, self.i_sec = scalars[0], scalars[1] # args_1=t < args_4=seconds by name
|
||
self.out = self.it.get_output_details()[0]["index"]
|
||
# Track the currently-allocated (batch, length) so set_conditioning can skip the
|
||
# expensive resize+allocate when only the conditioning tensors change.
|
||
self._alloc_B = self._alloc_L = None
|
||
self.set_conditioning(L, t5_hidden, t5_mask, seconds, cfg=cfg, apg=apg,
|
||
null_hidden=null_hidden, null_mask=null_mask,
|
||
local_add_cond=local_add_cond, batched=batched)
|
||
|
||
def set_conditioning(self, L, t5_hidden, t5_mask, seconds, *, cfg=1.0, apg=1.0,
|
||
null_hidden=None, null_mask=None, local_add_cond=None, batched=True):
|
||
"""(Re)bind per-generation conditioning onto this (possibly reused) interpreter.
|
||
|
||
The baked DiT holds the conditioning (T5 hidden/mask, seconds, local_add_cond) as
|
||
RESIDENT input tensors — a fresh generation only changes those, not the graph. A
|
||
cached BakedDiT can therefore be re-pointed at a new prompt/length WITHOUT rebuilding
|
||
the interpreter or re-packing XNNPACK's weights: resize+allocate_tensors runs only
|
||
when the batch (cfg on/off) or length L actually changes, otherwise the residents are
|
||
overwritten in place (same reuse pattern BakedDecoder/BakedEncoder use for L). Called
|
||
from __init__ (so the CLI's one-shot path is unchanged) and by the gradio DiT cache."""
|
||
self.L = L
|
||
self.cfg = float(cfg); self.apg = float(apg)
|
||
self.n_fwd = 0
|
||
# Batched CFG only applies when there's an uncond branch to co-batch (cfg != 1.0).
|
||
self.batched = bool(batched) and self.cfg != 1.0
|
||
B = 2 if self.batched else 1
|
||
self.B = B
|
||
# Resize batch axis to B (the canonical DiT is variable-batch) + length axis to L.
|
||
if (self._alloc_B, self._alloc_L) != (B, L):
|
||
self.it.resize_tensor_input(self.i_x["index"], [B, 256, L], strict=False)
|
||
self.it.resize_tensor_input(self.i_lac["index"], [B, 257, L], strict=False)
|
||
self.it.resize_tensor_input(self.i_t5h["index"], [B, COND_TOKENS, COND_DIM], strict=False)
|
||
self.it.resize_tensor_input(self.i_t5m["index"], [B, COND_TOKENS], strict=False)
|
||
self.it.resize_tensor_input(self.i_t["index"], [B], strict=False)
|
||
self.it.resize_tensor_input(self.i_sec["index"], [B], strict=False)
|
||
self.it.allocate_tensors()
|
||
self._alloc_B, self._alloc_L = B, L
|
||
|
||
t5h = t5_hidden.astype(np.float32)
|
||
t5m = t5_mask.astype(np.float32)
|
||
sec = np.array([np.float32(seconds)], np.float32)
|
||
lac = (np.zeros((1, 257, L), np.float32) if local_add_cond is None
|
||
else local_add_cond.astype(np.float32))
|
||
self.null_h = None if null_hidden is None else null_hidden.astype(np.float32)
|
||
self.null_m = None if null_mask is None else null_mask.astype(np.float32)
|
||
if self.batched:
|
||
# row0 = cond, row1 = uncond. t5_hidden/mask differ per row; seconds + lac are
|
||
# shared → tiled. All four are constant across steps → set resident once.
|
||
self.it.set_tensor(self.i_t5h["index"], np.concatenate([t5h, self.null_h], axis=0))
|
||
self.it.set_tensor(self.i_t5m["index"], np.concatenate([t5m, self.null_m], axis=0))
|
||
self.it.set_tensor(self.i_sec["index"], np.concatenate([sec, sec], axis=0))
|
||
self.it.set_tensor(self.i_lac["index"], np.concatenate([lac, lac], axis=0))
|
||
else:
|
||
self.t5h, self.t5m = t5h, t5m
|
||
self.it.set_tensor(self.i_sec["index"], sec) # seconds + lac constant → resident
|
||
self.it.set_tensor(self.i_lac["index"], lac)
|
||
return self
|
||
|
||
def _fwd(self, x, t, t5h, t5m):
|
||
"""One batch=1 invoke (cfg==1.0 or sequential CFG)."""
|
||
self.it.set_tensor(self.i_x["index"], x.astype(np.float32))
|
||
self.it.set_tensor(self.i_t["index"], np.array([np.float32(t)], np.float32))
|
||
self.it.set_tensor(self.i_t5h["index"], t5h)
|
||
self.it.set_tensor(self.i_t5m["index"], t5m)
|
||
self.it.invoke()
|
||
self.n_fwd += 1
|
||
return self.it.get_tensor(self.out).copy()
|
||
|
||
def _fwd_batched(self, x, t):
|
||
"""One batch=2 invoke -> (v_cond, v_uncond). x is the shared batch=1 state (both CFG
|
||
branches denoise the same latent), tiled to [2,256,L]; t5h/t5m/sec/lac already resident."""
|
||
x2 = np.concatenate([x, x], axis=0).astype(np.float32)
|
||
self.it.set_tensor(self.i_x["index"], x2)
|
||
self.it.set_tensor(self.i_t["index"], np.array([np.float32(t), np.float32(t)], np.float32))
|
||
self.it.invoke()
|
||
self.n_fwd += 2
|
||
v2 = self.it.get_tensor(self.out)
|
||
return v2[0:1].copy(), v2[1:2].copy()
|
||
|
||
def __call__(self, x, t, cross=None, gcond=None):
|
||
if self.cfg == 1.0:
|
||
return self._fwd(x, t, self.t5h, self.t5m)
|
||
if self.batched:
|
||
v_cond, v_uncond = self._fwd_batched(x, t)
|
||
else:
|
||
v_cond = self._fwd(x, t, self.t5h, self.t5m)
|
||
v_uncond = self._fwd(x, t, self.null_h, self.null_m)
|
||
return _apply_cfg(x, t, v_cond, v_uncond, self.cfg, self.apg)
|
||
|
||
|
||
# ─── Baked audio-in encoder (audio -> latents; SAME-S needs even L) ────────
|
||
class BakedEncoder:
|
||
def __init__(self, path, threads=8, needs_even=False):
|
||
from ai_edge_litert import interpreter as tfl
|
||
self.it = tfl.Interpreter(model_path=str(path), num_threads=threads)
|
||
self.i = self.it.get_input_details()[0]["index"]
|
||
self.o = self.it.get_output_details()[0]["index"]
|
||
self.needs_even = needs_even
|
||
self._cur_N = None
|
||
|
||
def _resize(self, N):
|
||
if self._cur_N != N:
|
||
self.it.resize_tensor_input(self.i, [1, 2, N], strict=False)
|
||
self.it.allocate_tensors()
|
||
self._cur_N = N
|
||
|
||
def encode(self, audio, T_lat):
|
||
"""audio: (1,2,M), M a multiple of 4096 (caller pads to even L for SAME-S).
|
||
Returns latents (1,256,T_lat) — trimmed back to the natural (decoder-independent) T_lat."""
|
||
self._resize(audio.shape[-1])
|
||
self.it.set_tensor(self.i, audio.astype(np.float32))
|
||
self.it.invoke()
|
||
return self.it.get_tensor(self.o)[:, :, :T_lat].copy()
|
||
|
||
|
||
# ─── Baked audio-out decoder (whole for SAME-S; chunked for SAME-L) ─────────
|
||
class BakedDecoder:
|
||
def __init__(self, path, threads=8, needs_even=False):
|
||
from ai_edge_litert import interpreter as tfl
|
||
self.tfl = tfl
|
||
self.path = str(path)
|
||
self.threads = threads
|
||
self.needs_even = needs_even # SAME-S: T_aud=L*16 must be %32 -> pad odd L->even
|
||
self.it = tfl.Interpreter(model_path=self.path, num_threads=threads)
|
||
self.i = self.it.get_input_details()[0]["index"]
|
||
self.o = self.it.get_output_details()[0]["index"]
|
||
self.i_det = self.it.get_input_details()[0]
|
||
self._cur_L = None
|
||
|
||
def _resize(self, L):
|
||
if self._cur_L != L:
|
||
self.it.resize_tensor_input(self.i, [1, 256, L], strict=False)
|
||
self.it.allocate_tensors()
|
||
self._cur_L = L
|
||
|
||
def decode_whole(self, latents):
|
||
L = latents.shape[2]
|
||
if self.needs_even and L % 2 != 0:
|
||
# SAME-S needs even L. Pad one edge-replicated latent token, decode at L+1,
|
||
# trim the extra token's audio. Narrow receptive field => negligible boundary
|
||
# effect on the kept L*4096 samples. Keeps odd-L requests working without
|
||
# changing the DiT/noise path (which stays natural-ceil, MLX/TRT-matched).
|
||
latents = np.concatenate([latents, latents[:, :, -1:]], axis=2)
|
||
self._resize(L + 1)
|
||
self.it.set_tensor(self.i, latents.astype(np.float32))
|
||
self.it.invoke()
|
||
return self.it.get_tensor(self.o)[:, :, :L * SAMPLES_PER_LATENT].copy()
|
||
self._resize(L)
|
||
self.it.set_tensor(self.i, latents.astype(np.float32))
|
||
self.it.invoke()
|
||
return self.it.get_tensor(self.o).copy() # [1,2,L*4096]
|
||
|
||
def decode_chunked(self, latents, chunk, overlap, on_chunk=None):
|
||
"""Stitch audio directly (stride = 4096 samples per latent token)."""
|
||
B, C, L = latents.shape
|
||
if L <= chunk:
|
||
return self.decode_whole(latents)
|
||
S = SAMPLES_PER_LATENT
|
||
out = np.zeros((B, 2, L * S), np.float32)
|
||
K = chunk - 2 * overlap
|
||
assert K > 0, (chunk, overlap)
|
||
core = 0
|
||
n_windows = (L + K - 1) // K
|
||
wi = 0
|
||
while core < L:
|
||
core_end = min(core + K, L)
|
||
win_start = core - overlap
|
||
win_end = win_start + chunk
|
||
if win_start < 0:
|
||
win_start, win_end = 0, chunk
|
||
if win_end > L:
|
||
win_end, win_start = L, L - chunk
|
||
win = latents[:, :, win_start:win_end]
|
||
y = self.decode_whole(win) # [1,2,chunk*4096]
|
||
ks = (core - win_start) * S
|
||
ke = (core_end - win_start) * S
|
||
out[:, :, core * S: core_end * S] = y[:, :, ks:ke]
|
||
wi += 1
|
||
if on_chunk:
|
||
on_chunk(wi, n_windows)
|
||
core = core_end
|
||
return out
|
||
|
||
|
||
def valid_T_lat(seconds):
|
||
"""seconds -> T_lat via natural ceil, DECODER-INDEPENDENT and identical to MLX/TRT
|
||
(sa3_mlx / sa3_trt.resolve_T_lat = max(1, ceil(seconds*44100/4096))). So MLX == TRT ==
|
||
TFLite pick the SAME length -> no length-driven divergence, and true ODD-length
|
||
requests are honored. The DiT handles odd L natively; SAME-L takes any L; SAME-S
|
||
(even T_aud=L*16 % 32) pads odd->even at encode/decode and trims."""
|
||
return max(1, int(np.ceil(seconds * SAMPLE_RATE / SAMPLES_PER_LATENT)))
|
||
|
||
|
||
def _play_backend(platform: str | None = None) -> tuple[str, list[str] | None]:
|
||
"""Pick the --play playback backend for `platform` (default: sys.platform)
|
||
WITHOUT executing anything. Split out of main() so it can be import-tested.
|
||
|
||
Returns (kind, argv_prefix):
|
||
("subprocess", ["afplay"]) macOS — argv_prefix + [wav_path] is run
|
||
("subprocess", ["aplay", …]) Linux with alsa-utils installed
|
||
("winsound", None) Windows — stdlib winsound.PlaySound (blocking)
|
||
("none", None) no player available — caller prints the path
|
||
"""
|
||
plat = sys.platform if platform is None else platform
|
||
if plat == "darwin":
|
||
return ("subprocess", ["afplay"])
|
||
if plat == "win32":
|
||
return ("winsound", None)
|
||
if plat.startswith("linux"):
|
||
from shutil import which
|
||
if which("aplay") is not None:
|
||
return ("subprocess", ["aplay", "-q"])
|
||
return ("none", None)
|
||
|
||
|
||
class _HelpfulParser(argparse.ArgumentParser):
|
||
"""argparse that prints full help (not just usage) when a flag is unknown / invalid,
|
||
and tacks the shared example-commands block onto the end of -h / --help."""
|
||
def error(self, message):
|
||
sys.stderr.write(f"\nerror: {message}\n\n")
|
||
self.print_help(sys.stderr)
|
||
sys.exit(2)
|
||
def print_help(self, file=None):
|
||
super().print_help(file)
|
||
try:
|
||
from examples import print_example_commands
|
||
print_example_commands()
|
||
except Exception:
|
||
pass # never let an examples-block failure mask the actual --help
|
||
|
||
|
||
def main():
|
||
ap = _HelpfulParser(
|
||
description="SA3 text-to-audio (+ audio-to-audio + inpainting + CFG) — baked-I/O varlen TFLite / CPU",
|
||
allow_abbrev=False, # --pr/--di/--de became ambiguous once *-precision flags landed
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog=("modes\n"
|
||
" text-to-audio --prompt P\n"
|
||
" audio-to-audio --prompt P --init-audio IN.wav [--init-noise-level σ]\n"
|
||
" inpainting --prompt P --init-audio IN.wav --inpaint-range START,END\n"
|
||
" negative CFG --prompt P --cfg N [--negative-prompt P_NEG] [--apg A]\n"))
|
||
# Inputs
|
||
ap.add_argument("--prompt", default=None,
|
||
help="Text prompt. Empty string = unconditional. If omitted, asked via stdin.")
|
||
ap.add_argument("--negative-prompt", default=None,
|
||
help="Negative prompt for CFG's uncond branch. Ignored when --cfg=1.0. "
|
||
"When unset and --cfg≠1.0, the uncond branch uses all-zero T5 hidden+mask "
|
||
"(→ learned padding embeddings in-graph).")
|
||
ap.add_argument("--init-audio", default=None,
|
||
help="WAV to start from. With --init-noise-level → audio-to-audio; with "
|
||
"--inpaint-range → inpainting. Encoder loaded automatically; audio is "
|
||
"trimmed/zero-padded to --seconds. Any format via ffmpeg fallback.")
|
||
ap.add_argument("--inpaint-range", default=None,
|
||
help="Inpaint time range 'START,END' in seconds (e.g. '5,10'). Requires "
|
||
"--init-audio. Regenerates the masked span, preserves the rest (paste-back).")
|
||
# Models
|
||
ap.add_argument("--dit", choices=["sm-music", "sm-sfx", "medium"], default=None,
|
||
help="DiT model (names match sa3_mlx / sa3_trt). If omitted, prompts "
|
||
"interactively with an arrow-key picker.")
|
||
ap.add_argument("--decoder", choices=["same-s", "same-l"], default=None,
|
||
help="Audio decoder. Default: same-s for sm-music/sm-sfx, same-l for medium. "
|
||
"If omitted, prompts interactively with an arrow-key picker.")
|
||
ap.add_argument("--precision", choices=list(PRECISIONS), default="fp32",
|
||
help="DiT + decoder precision, wXaY = weight/activation bits (same flag "
|
||
"as the TensorRT CLI): 'fp32' (default — reference; on CPU also the "
|
||
"fast choice) | 'w16a32' (fp16 weights, half size, ≈lossless, 1.5-3× "
|
||
"slower) | 'w8a32' (GPTQ int8 weights, ¼ size at fp32 speed) | "
|
||
"'w8a8-dyn' (dynamic int8, fastest, lower quality). "
|
||
"Applies to the DiT, codec decoder, and (for a2a/inpaint) the "
|
||
"encoder; T5Gemma is single-precision. Per-component overrides below.")
|
||
ap.add_argument("--dit-precision", choices=list(PRECISIONS), default=None,
|
||
help="Override --precision for the DiT only (e.g. a quantized DiT "
|
||
"with an fp32 codec).")
|
||
ap.add_argument("--decoder-precision", choices=list(PRECISIONS), default=None,
|
||
help="Override --precision for the SAME codec only. The codec runs "
|
||
"once on a fixed latent, so its precision maps directly to "
|
||
"audio quality (w8a32 is transparent at 40-46 dB).")
|
||
ap.add_argument("--encoder-precision", choices=list(PRECISIONS), default=None,
|
||
help="Override --precision for the SAME encoder (audio-to-audio / "
|
||
"inpainting only). Latent PSNR vs fp32 (same-s/same-l): "
|
||
"w16a32 ≈lossless (66/71 dB), w8a32 (GPTQ) 36/46 dB, "
|
||
"w8a8-dyn 24/30 dB.")
|
||
# LoRA (merged into the DiT weights at load; mirrors the MLX CLI's --lora)
|
||
ap.add_argument("--lora", action="append", nargs="+", default=None,
|
||
metavar=("ADAPTER", "strength=S"),
|
||
help="A LoRA adapter to merge into the DiT — repeat the flag to stack "
|
||
"several. Each --lora takes the adapter path followed by an optional "
|
||
"strength=S (default --lora-strength). Example: "
|
||
"--lora plini.safetensors strength=0.8 . The adapter is a .safetensors "
|
||
"file (SA3-native train_lora.py / underfit output) or a PEFT adapter "
|
||
"directory; ONLY .safetensors is accepted. The merged weights are "
|
||
"written into a cached copy of the DiT (models/tflite/lora_cache/) and "
|
||
"reused on repeat runs. Requires --dit-precision fp32 or w16a32 "
|
||
"(quantized-int8 DiTs can't be LoRA-merged). NOTE: per-step gating "
|
||
"(steps=) is MLX-only — the frozen TFLite graph merges once at load.")
|
||
ap.add_argument("--lora-strength", type=float, default=1.0,
|
||
help="Default strength for --lora adapters without their own strength= "
|
||
"(default 1.0). 0 disables (bit-identical to no LoRA); >1 amplifies.")
|
||
# Sampling
|
||
ap.add_argument("--seconds", type=float, default=30.0,
|
||
help="Output length. T_lat = ceil(seconds*44100/4096) (natural ceil, decoder-"
|
||
"independent, matches MLX/TRT). Final WAV trimmed to exactly --seconds.")
|
||
ap.add_argument("--steps", type=int, default=8,
|
||
help="Pingpong sampling steps (≥1). rf_denoiser is distilled for 8 (default).")
|
||
ap.add_argument("--seed", type=int, default=None,
|
||
help="Random seed. If omitted, a random seed is chosen and printed at the end.")
|
||
ap.add_argument("--init-noise-level", type=float, default=1.0,
|
||
help="σmax — schedule's starting noise level. With --init-audio: 0.4-0.8 varies, "
|
||
"1.0 = full regen (init ignored). Min %.2f (model NaNs at t≈0)." % MIN_SIGMA)
|
||
ap.add_argument("--cfg", type=float, default=1.0,
|
||
help="Classifier-Free Guidance scale. 1.0 = off (single pass). >1 toward prompt, "
|
||
"<1 toward uncond. Any value ≠1.0 runs cond+uncond each step (see --cfg-batched).")
|
||
ap.add_argument("--apg", type=float, default=1.0,
|
||
help="Adaptive Projected Guidance [0..1], only when --cfg≠1.0. 1.0 = full APG "
|
||
"(orthogonal projection), 0.0 = vanilla CFG. rf_denoiser default 1.0.")
|
||
ap.add_argument("--cfg-batched", action=argparse.BooleanOptionalAction, default=True,
|
||
help="When --cfg≠1.0, run cond+uncond as ONE batch=2 invoke on the variable-batch "
|
||
"DiT (default; ~7-29%% faster on Apple-Silicon AMX). --no-cfg-batched forces the "
|
||
"sequential batch=1 dual-pass (like TensorRT). Bit-identical at "
|
||
"fp32/w8a32; ~80 dB at w16a32; w8a8-dyn diverges by design (batch-"
|
||
"shared activation scales) — use --no-cfg-batched to reproduce "
|
||
"sequential baselines with w8a8-dyn.")
|
||
# Runtime / output
|
||
ap.add_argument("--threads", type=int, default=8, help="XNNPACK CPU threads.")
|
||
ap.add_argument("--free-models", action=argparse.BooleanOptionalAction, default=True,
|
||
help="Free each model after its last use to lower peak RAM (default on).")
|
||
ap.add_argument("--out", "-o", default=None,
|
||
help="Output WAV path. Relative → output/<file>; absolute → as-is. "
|
||
"If omitted, auto-named from the prompt + seed.")
|
||
ap.add_argument("--play", action="store_true",
|
||
help="Play the WAV after writing (blocking): `afplay` on macOS, "
|
||
"winsound on Windows, `aplay` on Linux (prints the path if "
|
||
"no player is available).")
|
||
args = ap.parse_args()
|
||
if args.steps < 1:
|
||
ap.error(f"--steps must be ≥ 1 (got {args.steps})")
|
||
# per-component overrides fall back to the shared --precision
|
||
args.dit_precision = args.dit_precision or args.precision
|
||
args.decoder_precision = args.decoder_precision or args.precision
|
||
args.encoder_precision = args.encoder_precision or args.precision
|
||
|
||
# Parse --lora specs up front (fail fast, before any model work).
|
||
args.lora_specs = None
|
||
if args.lora:
|
||
from lora_core import parse_lora_spec, LoraError
|
||
try:
|
||
args.lora_specs = [parse_lora_spec(toks, args.lora_strength) for toks in args.lora]
|
||
except LoraError as e:
|
||
ap.error(str(e))
|
||
|
||
# Interactive fills (match MLX/TRT).
|
||
args = prompt_user_if_missing(args)
|
||
if args.prompt is None:
|
||
args.prompt = input("Prompt: ").strip()
|
||
if args.seed is None:
|
||
args.seed = random.randint(0, 2**31 - 1)
|
||
|
||
dec = args.decoder or DEFAULT_DECODER[args.dit]
|
||
T_lat = valid_T_lat(args.seconds)
|
||
target_dur = T_lat * SAMPLES_PER_LATENT / SAMPLE_RATE
|
||
|
||
# Resolve output path — auto-name from prompt+seed when --out is not given.
|
||
# Relative paths land in output/ (auto-created); absolute paths honored as-is.
|
||
# A relative path that already starts with "output/" is taken as-is (relative to
|
||
# cwd) so `--out output/foo.wav` doesn't become output/output/foo.wav.
|
||
if args.out is None:
|
||
slug = re.sub(r'[^a-z0-9]+', '_', args.prompt.lower()).strip('_')[:48] or "out"
|
||
# non-fp32 runs get a precision suffix so same-prompt/seed A/B runs across
|
||
# --precision values don't overwrite each other
|
||
if args.dit_precision == args.decoder_precision:
|
||
prec_tag = "" if args.precision == "fp32" else f"_{args.precision}"
|
||
else:
|
||
prec_tag = f"_dit-{args.dit_precision}_dec-{args.decoder_precision}"
|
||
lora_tag = ""
|
||
if args.lora_specs:
|
||
names = "-".join(re.sub(r'[^a-z0-9]+', '', Path(s["path"]).stem.lower())[:12]
|
||
for s in args.lora_specs)
|
||
lora_tag = f"_lora-{names}"
|
||
args.out = f"{slug}{prec_tag}{lora_tag}_{args.seed}.wav"
|
||
out_path = Path(args.out)
|
||
if not out_path.is_absolute() and out_path.parts[:1] != ("output",):
|
||
out_path = REPO / "output" / out_path
|
||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||
args.out = str(out_path)
|
||
|
||
# Inpaint validation → latent range.
|
||
inpaint_range = None
|
||
inp_start_sec = inp_end_sec = None
|
||
if args.inpaint_range is not None:
|
||
if args.init_audio is None:
|
||
sys.exit("error: --inpaint-range requires --init-audio (the audio to inpaint into)")
|
||
try:
|
||
s_str, e_str = args.inpaint_range.split(",")
|
||
inp_start_sec = float(s_str.strip()); inp_end_sec = float(e_str.strip())
|
||
except ValueError:
|
||
sys.exit(f"error: --inpaint-range must be 'START,END' in seconds; got {args.inpaint_range!r}")
|
||
if not (0 <= inp_start_sec < inp_end_sec <= args.seconds):
|
||
sys.exit(f"error: invalid inpaint range {inp_start_sec}-{inp_end_sec}s "
|
||
f"(must satisfy 0 <= start < end <= {args.seconds}s)")
|
||
inp_start_lat = max(0, int(round(inp_start_sec * SAMPLE_RATE / SAMPLES_PER_LATENT)))
|
||
inp_end_lat = min(T_lat, int(round(inp_end_sec * SAMPLE_RATE / SAMPLES_PER_LATENT)))
|
||
inpaint_range = (inp_start_lat, inp_end_lat)
|
||
|
||
sigma_max = float(args.init_noise_level)
|
||
if sigma_max < MIN_SIGMA:
|
||
sys.exit(f"error: --init-noise-level={sigma_max} too low (min {MIN_SIGMA}; model NaNs at t≈0)")
|
||
mode = ("inpaint" if inpaint_range else
|
||
"audio-to-audio" if args.init_audio else "text-to-audio")
|
||
|
||
# ── Preflight: download any missing model files BEFORE the banner/wall-clock. ──
|
||
_preflight_download(args, dec)
|
||
|
||
# ── Banner ──
|
||
print()
|
||
banner(f"SA3 → TFLite/CPU {mode}")
|
||
k = lambda s: dim(f"{s:>11}")
|
||
print(f" {k('prompt')} {bold(repr(args.prompt))}")
|
||
if args.negative_prompt:
|
||
suffix = "" if args.cfg != 1.0 else dim(" (ignored: --cfg=1.0)")
|
||
print(f" {k('neg prompt')} {bold(repr(args.negative_prompt))}{suffix}")
|
||
line = f" {k('dit')} {magenta(args.dit)} {k('decoder')} {magenta(dec)}"
|
||
if args.init_audio:
|
||
enc_note = "" if args.encoder_precision == args.dit_precision else f" ({args.encoder_precision})"
|
||
line += f" {k('encoder')} {magenta(dec + enc_note)}"
|
||
prec_disp = (args.dit_precision if args.dit_precision == args.decoder_precision
|
||
else f"dit {args.dit_precision} · codec {args.decoder_precision}")
|
||
line += f" {k('precision')} {magenta(prec_disp)} {k('threads')} {args.threads}"
|
||
print(line)
|
||
if args.init_audio:
|
||
print(f" {k('init audio')} {bold(args.init_audio)}")
|
||
if inpaint_range:
|
||
s0, s1 = inpaint_range
|
||
print(f" {k('inpaint')} {bold(f'{inp_start_sec:.2f}s..{inp_end_sec:.2f}s')} "
|
||
f"{dim(f'(latent {s0}..{s1} of {T_lat})')}")
|
||
cfg_line = f" {k('σmax')} {bold(f'{sigma_max:.2f}')} {k('cfg')} {args.cfg}"
|
||
if args.cfg != 1.0:
|
||
cfg_line += dim(f" (apg={args.apg}, {'batched' if args.cfg_batched else 'sequential'} CFG)")
|
||
print(cfg_line)
|
||
if args.lora_specs:
|
||
lora_disp = ", ".join(
|
||
os.path.basename(s["path"].rstrip("/"))
|
||
+ (f"×{s['strength']:g}" if s["strength"] != 1.0 else "")
|
||
for s in args.lora_specs)
|
||
print(f" {k('lora')} {magenta(lora_disp)}")
|
||
print(f" {k('seconds')} {args.seconds}s {k('steps')} {args.steps} {k('seed')} {args.seed}")
|
||
print(f" {k('T_lat')} {T_lat} {dim(f'({target_dur:.2f}s → trimmed to {args.seconds:.2f}s)')}")
|
||
print()
|
||
|
||
# Stage numbering (extra stage when encoding init audio).
|
||
N = 3 + (1 if args.init_audio else 0)
|
||
TAG = {"t5": f"[1/{N}]"}
|
||
if args.init_audio:
|
||
TAG.update(enc=f"[2/{N}]", dit=f"[3/{N}]", dec=f"[4/{N}]")
|
||
else:
|
||
TAG.update(dit=f"[2/{N}]", dec=f"[3/{N}]")
|
||
|
||
t_wall = time.perf_counter()
|
||
|
||
# ── T5Gemma (tokenize + encode; + negative prompt for CFG) ──
|
||
stage(TAG["t5"], "T5Gemma (tokenize + encode)")
|
||
t0 = time.perf_counter()
|
||
tok = P.Tokenizer()
|
||
ids, mask = tok(args.prompt) # (1,256) int32 each
|
||
t5 = P.T5GemmaTFLite(ensure_local(T5_REL), args.threads)
|
||
t5_hidden = t5(ids, mask) # (1,256,768) fp32
|
||
null_h = null_m = None
|
||
if args.cfg != 1.0:
|
||
if args.negative_prompt:
|
||
n_ids, n_mask = tok(args.negative_prompt)
|
||
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). No extra T5 pass.
|
||
null_h = np.zeros((1, COND_TOKENS, COND_DIM), np.float32)
|
||
null_m = np.zeros((1, COND_TOKENS), np.float32)
|
||
t5_ms = (time.perf_counter() - t0) * 1000
|
||
stage(TAG["t5"], "T5Gemma (tokenize + encode)", t5_ms)
|
||
sub(f"t5_hidden {t5_hidden.shape} mask sum={int(mask.sum())}"
|
||
+ (f" neg ({'prompt' if args.negative_prompt else 'zeros'})" if null_h is not None else ""))
|
||
if args.free_models:
|
||
del t5; _free()
|
||
|
||
# ── (audio-to-audio / inpaint) Encode init audio → init_latents ──
|
||
init_latents = None
|
||
if args.init_audio:
|
||
stage(TAG["enc"], f"Encode init audio → latents ({dec})")
|
||
t0 = time.perf_counter()
|
||
# SAME-S encoder needs even L; round the encode grid up, trim latents to T_lat.
|
||
enc_L = T_lat + 1 if (dec == "same-s" and T_lat % 2 != 0) else T_lat
|
||
target_samples = enc_L * SAMPLES_PER_LATENT
|
||
audio_in = read_wav(args.init_audio) # (2, T_in)
|
||
if audio_in.shape[-1] >= target_samples:
|
||
audio_in = audio_in[:, :target_samples]
|
||
init_action = f"trimmed to {target_samples} samples"
|
||
else:
|
||
pad = target_samples - audio_in.shape[-1]
|
||
audio_in = np.pad(audio_in, ((0, 0), (0, pad)))
|
||
init_action = f"zero-padded by {pad} samples"
|
||
audio_in = audio_in[None] # (1,2,target_samples)
|
||
enc = BakedEncoder(ensure_local(enc_rel(dec, args.encoder_precision)), args.threads, needs_even=(dec == "same-s"))
|
||
init_latents = enc.encode(audio_in, T_lat) # (1,256,T_lat)
|
||
enc_ms = (time.perf_counter() - t0) * 1000
|
||
stage(TAG["enc"], f"Encode init audio → latents ({dec})", enc_ms)
|
||
sub(f"{init_action} latents {init_latents.shape}")
|
||
if args.free_models:
|
||
del enc; _free()
|
||
|
||
# ── Build inpaint local_add_cond + paste-back, and initial noise ──
|
||
local_add_cond = None
|
||
paste_back = None
|
||
if inpaint_range is not None:
|
||
s0, s1 = inpaint_range
|
||
keep = np.ones((1, 1, T_lat), np.float32); keep[:, :, s0:s1] = 0.0 # 1=keep, 0=regen
|
||
masked = init_latents.astype(np.float32) * keep
|
||
local_add_cond = np.concatenate([keep, masked], axis=1) # (1,257,T_lat), TRT layout
|
||
paste_back = (init_latents.astype(np.float32), keep)
|
||
|
||
x0, step_noise = P.make_noise(T_lat, args.steps, args.seed) # x0 = pure noise
|
||
if init_latents is not None and inpaint_range is None:
|
||
# rf_denoiser init mix (linear interp): noise = init*(1-σmax) + pure*σmax
|
||
x0 = init_latents.astype(np.float32) * (1.0 - sigma_max) + x0 * sigma_max
|
||
|
||
# ── DiT load + pingpong sample ──
|
||
stage(TAG["dit"], f"DiT — load + sample ({args.steps} steps, σmax={sigma_max:.2f})")
|
||
t0 = time.perf_counter()
|
||
cfg_note = (("CFG batched (1× batch=2 invoke/step)" if args.cfg_batched
|
||
else "CFG sequential (2× batch=1 invokes/step)") if args.cfg != 1.0 else "")
|
||
print(f" {dim('loading baked DiT ' + args.dit + ' ...')}", flush=True)
|
||
dit_path = ensure_local(dit_rel(args.dit, args.dit_precision))
|
||
if args.lora_specs:
|
||
from lora_patch import get_patched_dit
|
||
from lora_core import LoraError
|
||
try:
|
||
dit_path = get_patched_dit(dit_path, args.lora_specs, family=args.dit,
|
||
precision=args.dit_precision, log=sub)
|
||
except LoraError as e:
|
||
sys.exit(f"error: {e}")
|
||
backend = BakedDiT(dit_path, T_lat, t5_hidden, mask.astype(np.float32),
|
||
args.seconds, args.threads, cfg=args.cfg, apg=args.apg,
|
||
null_hidden=null_h, null_mask=null_m, local_add_cond=local_add_cond,
|
||
batched=args.cfg_batched)
|
||
load_ms = (time.perf_counter() - t0) * 1000
|
||
sub(f"load {load_ms/1000:.1f}s" + (f" {cfg_note}" if cfg_note else ""))
|
||
|
||
sig = P.build_pingpong_schedule(args.steps, sigma_max)
|
||
sched_str = " · ".join(f"{float(x):.3f}" for x in sig)
|
||
sub(f"schedule {sched_str}")
|
||
|
||
step_prev = [time.perf_counter()]
|
||
def on_step(i, total):
|
||
now = time.perf_counter(); el = (now - step_prev[0]) * 1000; step_prev[0] = now
|
||
if _USE_COLOR:
|
||
bar_w = 20; filled = int(round(bar_w * i / total))
|
||
bar = cyan("█" * filled) + dim("·" * (bar_w - filled))
|
||
sys.stdout.write(f"\r\x1b[K {dim('sampling')} {bar} "
|
||
f"{bold(f'step {i}/{total}')} {yellow(f'{el:.0f} ms')}")
|
||
sys.stdout.flush()
|
||
else:
|
||
print(f" sampling step {i}/{total} {el:.0f} ms", flush=True)
|
||
|
||
t0 = time.perf_counter()
|
||
latents = P.sample(backend, x0, step_noise, sig, None, None,
|
||
on_step=on_step, paste_back=paste_back)
|
||
samp_ms = (time.perf_counter() - t0) * 1000
|
||
if _USE_COLOR:
|
||
sys.stdout.write("\r\x1b[K")
|
||
stage(TAG["dit"], f"DiT — load + sample ({args.steps} steps, σmax={sigma_max:.2f})",
|
||
load_ms + samp_ms)
|
||
sub(f"sample {samp_ms:.0f} ms ({samp_ms/max(args.steps,1):.0f} ms/step, "
|
||
f"{backend.n_fwd} forwards) latents {latents.shape}")
|
||
if args.free_models:
|
||
del backend; _free()
|
||
|
||
# ── Decode (audio-out) + WAV ──
|
||
stage(TAG["dec"], f"Decoder ({dec}, audio-out) + WAV")
|
||
t0 = time.perf_counter()
|
||
print(f" {dim('loading baked decoder ' + dec + ' ...')}", flush=True)
|
||
decoder = BakedDecoder(ensure_local(dec_rel(dec, args.decoder_precision)), args.threads, needs_even=(dec == "same-s"))
|
||
load2_ms = (time.perf_counter() - t0) * 1000
|
||
sub(f"load {load2_ms:.0f} ms")
|
||
|
||
t0 = time.perf_counter()
|
||
if dec == "same-l" and T_lat > SAMEL_CHUNK:
|
||
print(f" {dim(f'SAME-L chunked decode (chunk={SAMEL_CHUNK}, ovl={SAMEL_OVERLAP}) ...')}", flush=True)
|
||
def on_chunk(i, n):
|
||
print(f" {dim(f'decode chunk {i}/{n}')}", flush=True)
|
||
audio = decoder.decode_chunked(latents, SAMEL_CHUNK, SAMEL_OVERLAP, on_chunk=on_chunk)
|
||
dmode = f"chunked (chunk={SAMEL_CHUNK}, ovl={SAMEL_OVERLAP})"
|
||
else:
|
||
print(f" {dim('whole decode ...')}", flush=True)
|
||
audio = decoder.decode_whole(latents)
|
||
dmode = "whole"
|
||
dec_ms = (time.perf_counter() - t0) * 1000
|
||
|
||
audio_np = audio[0] # (2, L*4096)
|
||
req = int(round(args.seconds * SAMPLE_RATE))
|
||
if audio_np.shape[-1] > req:
|
||
audio_np = audio_np[:, :req]
|
||
P.save_wav(args.out, audio_np)
|
||
stage(TAG["dec"], f"Decoder ({dec}, audio-out) + WAV", load2_ms + dec_ms)
|
||
peak = float(np.abs(audio_np).max()); rms = float(np.sqrt((audio_np**2).mean()))
|
||
sub(f"decode {dmode} {dec_ms:.0f} ms audio {audio_np.shape} peak {peak:.3f} rms {rms:.3f}")
|
||
if args.free_models:
|
||
del decoder; _free()
|
||
|
||
total = time.perf_counter() - t_wall
|
||
dur = audio_np.shape[-1] / SAMPLE_RATE
|
||
print()
|
||
rule()
|
||
print(f" {bold(green('done'))} {bold(f'{total:.2f}s')} wall → {dur:.1f}s audio → "
|
||
f"{bold(yellow(f'{dur/total:.2f}× realtime'))} {dim(f'seed {args.seed}')}")
|
||
abs_out = os.path.abspath(args.out)
|
||
try:
|
||
rel_out = os.path.relpath(abs_out)
|
||
except ValueError:
|
||
rel_out = abs_out
|
||
shown = rel_out if len(rel_out) <= len(abs_out) and not rel_out.startswith("..") else abs_out
|
||
print(f" {bold(green('▸ saved'))} {bold(shown)} {dim(f'({abs_out})' if shown != abs_out else '')}".rstrip())
|
||
rule()
|
||
|
||
if args.play:
|
||
kind, argv = _play_backend()
|
||
try:
|
||
if kind == "winsound":
|
||
print(f" {bold('▶ playing')} {args.out} {dim('(Ctrl-C to stop)')}")
|
||
import winsound
|
||
# SND_ASYNC + an interruptible sleep instead of the synchronous
|
||
# call: PlaySound(..., SND_FILENAME) blocks inside native code,
|
||
# so Python can't deliver Ctrl-C until playback finishes. This
|
||
# way Ctrl-C lands in time.sleep and SND_PURGE stops the audio.
|
||
with wave.open(args.out) as _w:
|
||
_dur = _w.getnframes() / _w.getframerate()
|
||
winsound.PlaySound(args.out,
|
||
winsound.SND_FILENAME | winsound.SND_ASYNC)
|
||
try:
|
||
time.sleep(_dur + 0.25)
|
||
except KeyboardInterrupt:
|
||
winsound.PlaySound(None, winsound.SND_PURGE)
|
||
raise
|
||
elif kind == "subprocess":
|
||
print(f" {bold('▶ playing')} {args.out} {dim('(Ctrl-C to stop)')}")
|
||
subprocess.run(argv + [args.out], check=False)
|
||
else:
|
||
print(f" {bold('▶ play')} no audio player found ('aplay' missing) — "
|
||
f"file saved at {args.out}")
|
||
except KeyboardInterrupt:
|
||
print()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|