Files
SoundGame/tools/stable-audio-3/optimized/mlx/scripts/parity_forward_mlx.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

108 lines
4.8 KiB
Python

"""§9 forward-parity harness — MLX side + diff.
Reads shared_inputs.npz + torch_out.npz (from parity_forward_torch.py) and runs
the SAME controlled forward through the MLX base DiT in fp32, then reports three
PSNRs that LOCALIZE any mismatch:
(1) DiT-only — feed torch's cross/global + the shared noised into the MLX DiT.
Isolates the DiT forward convention (weights are the same fp16 source→fp32,
so ~79 dB is the cross-framework fp32 floor; a convention bug reads ≪40 dB).
(2) conditioning — MLX-native cross/global vs torch's (T5 + seconds; fp16-T5 bound).
(3) end-to-end — full MLX-native prediction + loss vs torch.
Also asserts the torch model's pretransform.scale == 1.0 (the softnorm no-scale
convention the MLX trainer relies on). Run in the MLX venv:
python parity_forward_mlx.py --base-npz <dit_sm-music-base_f16.npz> --workdir <dir>
"""
import argparse
import importlib
import sys
import numpy as np
def psnr(a, b):
a = np.asarray(a, np.float64); b = np.asarray(b, np.float64)
mse = np.mean((a - b) ** 2)
if mse == 0:
return float("inf")
return 20 * np.log10(max(np.abs(a).max(), np.abs(b).max()) / np.sqrt(mse))
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--base-npz", required=True, help="MLX base DiT weights (dit_<model>-base_f16.npz)")
ap.add_argument("--workdir", required=True)
ap.add_argument("--dit-module", default="models.defs.dit_mlx")
ap.add_argument("--prompt", default="techno")
ap.add_argument("--seconds", type=float, default=30.0)
args = ap.parse_args()
import mlx.core as mx
from sa3_mlx import T5GEMMA_NPZ_REL
from weights import ensure_local
from models.defs.sa3_pipeline import apply_prompt_padding, load_conditioner_from_npz
from models.defs.t5gemma_mlx import T5Gemma
tt = np.load(f"{args.workdir}/torch_out.npz")
inp = np.load(f"{args.workdir}/shared_inputs.npz", allow_pickle=True)
assert abs(float(tt["scale"]) - 1.0) < 1e-9, \
f"pretransform.scale={float(tt['scale'])} != 1.0 — the MLX trainer's no-scale " \
"softnorm convention is WRONG for this model (real bug)."
latents = mx.array(inp["latents"]).astype(mx.float32)
noise = mx.array(inp["noise"]).astype(mx.float32)
t = mx.array(inp["t"]).astype(mx.float32)
B, C, T = latents.shape
dit = importlib.import_module(args.dit_module).load_dit(
args.base_npz, T_lat=T, dtype=mx.float32, compile_=False)
noised = latents * (1.0 - t)[:, None, None] + noise * t[:, None, None]
target = noise - latents
lac = mx.concatenate([mx.ones((B, 1, T)), mx.zeros((B, C, T))], axis=1).transpose(0, 2, 1).astype(mx.float32)
# (1) DiT-only: inject torch's cross/global
pred1 = dit(noised, t, mx.array(tt["cross"]).astype(mx.float32),
mx.array(tt["gcond"]).astype(mx.float32), local_add_cond=lac)
mx.eval(pred1)
# (2) MLX-native conditioning
t5 = T5Gemma.from_npz(str(ensure_local(T5GEMMA_NPZ_REL)))
padding_emb, secs = load_conditioner_from_npz(args.base_npz, prefix="cond.")
embeds, mask = t5.encode([args.prompt] * B, max_len=256)
prompt_cond = apply_prompt_padding(embeds.astype(mx.float32), mask, padding_emb.astype(mx.float32))
sec_tok = mx.stop_gradient(secs([args.seconds] * B).astype(mx.float32))
cross_m = mx.concatenate([prompt_cond, sec_tok], axis=1)
gcond_m = sec_tok[:, 0, :]
mx.eval(cross_m, gcond_m)
# (3) end-to-end
pred3 = dit(noised, t, cross_m, gcond_m, local_add_cond=lac)
mx.eval(pred3)
loss_m = float(mx.mean((pred3 - target) ** 2))
# (4) BACKWARD: ∂loss/∂noised through the DiT (with torch's cross/global so
# the backward is compared on identical inputs — isolates the DiT backward).
cross_t = mx.array(tt["cross"]).astype(mx.float32)
gcond_t = mx.array(tt["gcond"]).astype(mx.float32)
def loss_of_noised(nz):
pred = dit(nz, t, cross_t, gcond_t, local_add_cond=lac)
return mx.mean((pred - target) ** 2)
grad_m = mx.grad(loss_of_noised)(noised)
mx.eval(grad_m)
p_grad = psnr(np.array(grad_m), tt["grad_noised"]) if "grad_noised" in tt else float("nan")
print(f"noised sanity PSNR = {psnr(np.array(noised), tt['noised']):.1f} dB (inf = bit-identical)")
print(f"(1) DiT-only forward = {psnr(np.array(pred1), tt['prediction']):6.2f} dB [conv bug ≪40; fp32 floor ~79]")
print(f"(2) conditioning = cross {psnr(np.array(cross_m), tt['cross']):.2f} dB / global {psnr(np.array(gcond_m), tt['gcond']):.2f} dB")
print(f"(3) end-to-end pred = {psnr(np.array(pred3), tt['prediction']):6.2f} dB")
print(f"(4) BACKWARD d loss/d noised PSNR = {p_grad:6.2f} dB [full DiT backward; conv bug ≪40]")
print(f"loss: torch {float(tt['loss']):.6f} MLX {loss_m:.6f} |Δ| {abs(float(tt['loss'])-loss_m):.2e}")
if __name__ == "__main__":
sys.exit(main())