- 共享引擎(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 模板
298 lines
16 KiB
Python
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())
|