Files
SoundGame/tools/stable-audio-3/optimized/tensorRT/build/build_autoencoders.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

298 lines
16 KiB
Python

#!/usr/bin/env python3
"""Build the SAME-L / SAME-S autoencoder engines for the local GPU.
Consumer flow: download build-ready ONNX from HuggingFace, compile to TensorRT, verify. No model
checkpoints, no calibration data, no stable-audio-tools -- just tensorrt + torch + huggingface-hub.
Everything that needs a checkpoint has already been baked into the published ONNX: the decoders
arrive deterministic, with the limiter grafted in and its ceiling baked as a constant, and the fp8
encoders arrive with amax-calibrated activation scales. So this script only has to compile.
python build/build_autoencoders.py # all 8, auto-detect arch
python build/build_autoencoders.py --model same-l # one model
python build/build_autoencoders.py --precision fp8 # one precision
python build/build_autoencoders.py --no-verify # skip the post-build checks
python build/build_autoencoders.py --list # show targets and exit
WHAT GETS BUILT (8 engines). Each carries TWO optimization profiles:
band 0 L=1..256 the chunked / low-VRAM mode decoder ~509 MB scratch (SAME-L)
band 1 L=1..4096 single-shot decoder ~8143 MB
TensorRT commits a context's scratch at create_execution_context(), sized from the profile ceiling
and *before any shape is bound* -- so a single 4096-ceiling engine reserves the full amount whether
it decodes five seconds or six minutes. Two bands let one file serve both a low-memory deployment
and a fast single-shot one. ⚠ The saving only materialises with USER_MANAGED contexts sized by
get_device_memory_size_for_profile_v2(band): device_memory_size_v2 is the MAX across profiles, so a
DEFAULT context reserves the wide band and the benefit silently disappears. The runtime's
TRTRunner(..., profile=N) does this correctly.
`opt` is a separate lever from the ceiling and costs no memory -- TRT tunes tactic selection around
it. Measured across L=32..4096, the naive default of min(1292, ceiling) costs 17-23% at short L on
the wide band while buying nothing at long L. The values below are the measured optima.
ARCH. TensorRT bakes the compute capability into the engine, so run this ON the target GPU. The
arch you build on is the arch the engine runs on.
"""
import argparse
import os
import sys
import time
from pathlib import Path
BUILD_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(BUILD_DIR))
from _arch import detect_arch, arch_dir as _arch_dir # noqa: E402
HF_REPO = "stabilityai/stable-audio-3-optimized"
SPL = 4096 # samples per latent
# ── targets ────────────────────────────────────────────────────────────────────────────────────
# recipe "samel": STRONGLY_TYPED + the diff_attn_swa plugin. fp8 rides in the graph's own QDQ
# nodes, so the FP8 builder flag must NOT be set (strongly-typed forbids it).
# recipe "sames": weakly-typed EXPLICIT_BATCH + BF16. Here the FP8 flag IS required for the fp8
# tiers -- it is what lets TRT honour the in-graph QDQ.
TARGETS = {
"same-l/dec_fp16": dict(model="same-l", kind="dec", recipe="samel", fp8=False,
onnx="onnx/same-l/dec_fp16_chunkable_limiter.onnx",
out="same-l/dec_fp16_chunkable_limiter.trt",
bands=[256, 4096], opts=[256, 1024]),
"same-l/enc_fp16": dict(model="same-l", kind="enc", recipe="samel", fp8=False,
onnx="onnx/same-l/enc_dynamic_triton_swa.onnx",
out="same-l/enc_fp16_chunkable.trt",
bands=[256, 4096], opts=[256, 256]),
"same-l/dec_fp8": dict(model="same-l", kind="dec", recipe="samel", fp8=False,
onnx="onnx/same-l/dec_fp8_chunkable_limiter.onnx",
out="same-l/dec_fp8_chunkable_limiter.trt",
bands=[256, 4096], opts=[256, 1024]),
"same-l/enc_fp8": dict(model="same-l", kind="enc", recipe="samel", fp8=False,
onnx="onnx/same-l/enc_fp8.onnx",
out="same-l/enc_fp8_chunkable.trt",
bands=[256, 4096], opts=[256, 256]),
"same-s/dec_bf16": dict(model="same-s", kind="dec", recipe="sames", fp8=False,
onnx="onnx/same-s/dec_bf16_chunkable_limiter.onnx",
out="same-s/dec_bf16_chunkable_limiter.trt",
bands=[256, 4096], opts=[256, 512]),
"same-s/enc_bf16": dict(model="same-s", kind="enc", recipe="sames", fp8=False,
onnx="onnx/same-s/enc_dynamic_bf16.onnx",
out="same-s/enc_bf16_chunkable.trt",
bands=[256, 4096], opts=[256, 256]),
"same-s/dec_fp8": dict(model="same-s", kind="dec", recipe="sames", fp8=True,
onnx="onnx/same-s/dec_fp8_chunkable_limiter.onnx",
out="same-s/dec_fp8_chunkable_limiter.trt",
bands=[256, 4096], opts=[256, 512]),
"same-s/enc_fp8": dict(model="same-s", kind="enc", recipe="sames", fp8=True,
onnx="onnx/same-s/enc_fp8.onnx",
out="same-s/enc_fp8_chunkable.trt",
bands=[256, 4096], opts=[256, 256]),
}
LIMITER_CEILING = 0.977 # -0.2021 dBFS, baked into every published decoder
def _fetch(rel: str) -> str:
from huggingface_hub import hf_hub_download
return hf_hub_download(HF_REPO, rel)
def build_one(key: str, spec: dict, out_root: Path, workspace_gb: int = 16) -> Path:
import tensorrt as trt
out = out_root / spec["out"]
out.parent.mkdir(parents=True, exist_ok=True)
print(f"\n━━━ {key} → {spec['out']} ━━━", flush=True)
print(f" onnx: {spec['onnx']}", flush=True)
src = _fetch(spec["onnx"])
logger = trt.Logger(trt.Logger.ERROR)
trt.init_libnvinfer_plugins(logger, "")
flags = 0
if spec["recipe"] == "samel":
# AOT, not JIT: a JIT plugin runs Python inside the enqueue path, is not
# stream-capturable (so the engine cannot go inside a CUDA graph), and was measured to
# corrupt the heap on repeated enqueues in a multi-profile engine where AOT was bit-exact.
os.environ.setdefault("SA3_SWA_PLUGIN", "aot")
sys.path.insert(0, str(BUILD_DIR.parent / "scripts"))
import diff_attn_nocast_plugin # noqa: F401 registers samel::diff_attn_swa
flags |= 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
flags |= 1 << int(trt.NetworkDefinitionCreationFlag.PREFER_AOT_PYTHON_PLUGINS)
builder = trt.Builder(logger)
net = builder.create_network(flags)
parser = trt.OnnxParser(net, logger)
if not parser.parse_from_file(src):
for i in range(parser.num_errors):
print(f" {parser.get_error(i)}", file=sys.stderr)
raise SystemExit(f"ONNX parse failed for {key}")
cfg = builder.create_builder_config()
cfg.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_gb << 30)
if spec["recipe"] == "sames":
cfg.set_flag(trt.BuilderFlag.BF16)
if spec["fp8"]:
cfg.set_flag(trt.BuilderFlag.FP8)
tname = "audio" if spec["kind"] == "enc" else "latent"
for hi, opt in zip(spec["bands"], spec["opts"]):
prof = builder.create_optimization_profile()
if spec["kind"] == "enc":
prof.set_shape(tname, (1, 2, SPL), (1, 2, opt * SPL), (1, 2, hi * SPL))
else:
prof.set_shape(tname, (1, 256, 1), (1, 256, opt), (1, 256, hi))
cfg.add_optimization_profile(prof)
print(" profiles: " + " | ".join(f"L=1..{hi} opt={o}"
for hi, o in zip(spec["bands"], spec["opts"])), flush=True)
print(f" recipe: {'STRONGLY_TYPED + AOT plugin' if spec['recipe']=='samel' else 'EXPLICIT_BATCH + BF16'}"
f"{' + FP8' if spec['fp8'] else ''}", flush=True)
t0 = time.time()
ser = builder.build_serialized_network(net, cfg)
if ser is None:
raise SystemExit(f"BUILD FAILED for {key}")
out.write_bytes(memoryview(ser))
eng = trt.Runtime(logger).deserialize_cuda_engine(out.read_bytes())
print(f" built {out.stat().st_size/1e6:.0f} MB in {time.time()-t0:.0f}s", flush=True)
div = SPL if spec["kind"] == "enc" else 1
for i in range(eng.num_optimization_profiles):
lo, o, hi = eng.get_tensor_profile_shape(tname, i)
print(f" band {i}: L={lo[-1]//div}..{hi[-1]//div} opt={o[-1]//div} "
f"scratch {eng.get_device_memory_size_for_profile_v2(i)/1e6:>7.1f} MB", flush=True)
return out
def _test_signal(n_lat: int):
"""A harmonic, envelope-shaped stereo signal at a realistic level.
⚠ This is a STRUCTURAL smoke test -- does the engine produce finite, correctly-shaped,
non-silent audio and hold the limiter -- NOT a quality measurement. Autoencoder quality must
be measured on real music; synthetic noise in particular decodes past the int16 rail and
fabricates plausible-looking numbers.
"""
import numpy as np
t = np.arange(n_lat * SPL, dtype=np.float64) / 44100.0
sig = sum(np.sin(2*np.pi*f*t) / (i + 1) for i, f in enumerate((110., 220., 330., 550., 880.)))
env = 0.5 + 0.5 * np.sin(2*np.pi*0.7*t)
x = (sig / np.abs(sig).max() * env * 0.7).astype(np.float32)
return np.stack([x, np.roll(x, 977)]).copy()
def verify(built: dict, out_root: Path) -> bool:
"""Per-engine structural checks, plus the one cross-check that needs no reference: the two
bands of the same engine must agree with each other."""
import numpy as np, torch, tensorrt as trt
sys.path.insert(0, str(BUILD_DIR.parent / "scripts"))
import sa3_trt_core as canon
canon._import_heavy()
print("\n━━━ verify ━━━", flush=True)
ok = True
def run(path, prof, tensor, data):
"""Returns (array, is_pcm). is_pcm marks int32 output scaled to int16 range -- these
decoders bake the clip+scale+transpose tail, so the caller must divide by 32767. Deciding
that from the returned array's dtype does not work: it has already been cast to float."""
r = canon.TRTRunner(str(path), None, True, prof)
need = r.engine.get_device_memory_size_for_profile_v2(prof)
buf = torch.empty(need, dtype=torch.uint8, device="cuda")
r.context.set_device_memory(buf.data_ptr(), need)
if tensor == "audio":
return canon.encoder_encode(r, data).float().cpu().numpy(), False
out = canon.decoder_decode(r, data.to(r.in_dtype["latent"]))
return out.float().cpu().numpy(), ("pcm" in r.out_dtype)
for key, spec in built.items():
p = out_root / spec["out"]
L = 128 # fits band 0 and band 1 alike
try:
if spec["kind"] == "enc":
x = torch.from_numpy(_test_signal(L)).unsqueeze(0).cuda()
a, _ = run(p, 0, "audio", x)
b, _ = run(p, 1, "audio", x)
fin = np.isfinite(a).all() and np.isfinite(b).all()
shape_ok = a.shape[-1] == L
cs = float((a.ravel() @ b.ravel()) /
(np.linalg.norm(a) * np.linalg.norm(b) + 1e-30))
good = fin and shape_ok and cs > 0.99 and np.abs(a).max() > 1e-3
print(f" {spec['out']:>40s} latents{a.shape[-1]:>5d} band0/band1 cos {cs:.6f}"
f" {'OK' if good else 'FAIL'}", flush=True)
else:
enc_key = key.replace("dec", "enc")
ep = out_root / TARGETS[enc_key]["out"]
if not ep.exists():
# The decoder check is a round trip, so it needs the matching encoder.
# Building with --kind dec alone is legitimate (the encoder may already
# be deployed elsewhere), so say what is missing instead of reporting a
# failure the engine is not responsible for. No synthetic latent stands
# in here: a fabricated signal would not exercise the decoder the way
# real content does, and a green tick from one would be worse than none.
print(f" {spec['out']:>40s} SKIP needs {TARGETS[enc_key]['out']} for "
f"the round trip — build without --kind, or --kind enc first",
flush=True)
continue
x = torch.from_numpy(_test_signal(L)).unsqueeze(0).cuda()
z, _ = run(ep, 0, "audio", x)
zt = torch.from_numpy(z).cuda()
a, pcm_a = run(p, 0, "latent", zt)
b, pcm_b = run(p, 1, "latent", zt)
fa = np.clip(a[0] / 32767.0, -1, 1) if pcm_a else a[0]
fb = np.clip(b[0] / 32767.0, -1, 1) if pcm_b else b[0]
pk = float(np.abs(fa).max())
fin = np.isfinite(fa).all() and np.isfinite(fb).all()
cs = float((fa.ravel() @ fb.ravel()) /
(np.linalg.norm(fa) * np.linalg.norm(fb) + 1e-30))
# the limiter is baked: output must never exceed the ceiling (+1 LSB slack)
lim_ok = pk <= LIMITER_CEILING + 1.0 / 32767
good = fin and cs > 0.99 and pk > 1e-3 and lim_ok
print(f" {spec['out']:>40s} peak {20*np.log10(max(pk,1e-9)):>+7.2f} dBFS"
f" band0/band1 cos {cs:.6f} limiter {'held' if lim_ok else 'EXCEEDED'}"
f" {'OK' if good else 'FAIL'}", flush=True)
ok &= bool(good)
except Exception as e:
print(f" {spec['out']:>40s} FAIL {type(e).__name__}: {str(e)[:90]}", flush=True)
ok = False
torch.cuda.empty_cache()
return ok
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--model", choices=("same-l", "same-s"), help="build only this model")
ap.add_argument("--precision", choices=("fp16", "bf16", "fp8"), help="build only this precision")
ap.add_argument("--kind", choices=("enc", "dec"), help="build only encoders or only decoders")
ap.add_argument("--out", default=None, help="output root (default ../models/<arch>/)")
ap.add_argument("--workspace-gb", type=int, default=16)
ap.add_argument("--verify", action=argparse.BooleanOptionalAction, default=True)
ap.add_argument("--force", action="store_true", help="rebuild engines that already exist")
ap.add_argument("--list", action="store_true", help="list targets and exit")
a = ap.parse_args()
sel = {k: v for k, v in TARGETS.items()
if (not a.model or v["model"] == a.model)
and (not a.kind or v["kind"] == a.kind)
and (not a.precision or a.precision in k)}
if not sel:
print("no targets match those filters", file=sys.stderr); return 2
arch = detect_arch()
out_root = Path(a.out) if a.out else Path(_arch_dir(arch))
print(f" arch {arch} → {out_root}")
if a.list:
for k, v in sel.items():
print(f" {k:>16s} {v['onnx']:>46s} → {v['out']}")
return 0
todo = {}
for k, v in sel.items():
p = out_root / v["out"]
if p.exists() and not a.force:
print(f" skip {v['out']} (exists; --force to rebuild)")
else:
todo[k] = v
for k, v in todo.items():
build_one(k, v, out_root, a.workspace_gb)
if a.verify:
# verify everything selected, including engines that were already present
if not verify(sel, out_root):
print("\n VERIFICATION FAILED", file=sys.stderr); return 1
print("\n all checks passed")
return 0
if __name__ == "__main__":
raise SystemExit(main())