100 lines
3.4 KiB
Python
100 lines
3.4 KiB
Python
"""Parity harness: torch SeFiTransformer vs MLX SeFiDiTMLX on the SAME random
|
|
weights and input. Run this ON YOUR MAC (needs torch+comfy+mlx) to validate the
|
|
MLX port before trusting full-size outputs:
|
|
|
|
cd ComfyUI && python custom_nodes/ComfyUI-SeFi-MLX/sefi_mlx/parity.py
|
|
|
|
It builds a tiny model, converts torch -> MLX through the real converter/loader,
|
|
runs both forwards, and reports max/mean abs difference. Small differences
|
|
(~1e-2 in bf16-ish ranges, much less in fp32) are expected; large or structural
|
|
differences point at a bug in the port. A successful run also proves the key
|
|
mapping (no extra/missing keys) end to end.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
|
|
|
|
def main():
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
# repo root (ComfyUI) so `import comfy` works when run from there
|
|
here = os.path.abspath(__file__)
|
|
custom_nodes = os.path.dirname(os.path.dirname(os.path.dirname(here)))
|
|
comfy_root = os.path.dirname(custom_nodes)
|
|
sys.path.insert(0, comfy_root)
|
|
sys.path.insert(0, os.path.join(custom_nodes, "ComfyUI-SeFi-Image"))
|
|
|
|
import comfy.ops
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
import torch
|
|
|
|
import sefi_model # sibling torch implementation
|
|
|
|
from sefi_mlx.convert import convert_state_dict_to_mlx, infer_config_mlx
|
|
from sefi_mlx.loader import load_sefi_dit
|
|
|
|
torch.manual_seed(0)
|
|
hidden = 256 # num_heads = 2, head_dim = 128
|
|
sem, tex = 8, 16
|
|
cfg = {
|
|
"hidden_size": hidden,
|
|
"packed_in_channels": sem + tex,
|
|
"packed_out_channels": sem + tex,
|
|
"context_in_dim": 64,
|
|
"sin_dim": 256,
|
|
"half_hidden": hidden // 2,
|
|
"depth": 2,
|
|
"depth_single_blocks": 2,
|
|
"num_heads": hidden // 128,
|
|
"semantic_channels": sem,
|
|
}
|
|
|
|
# Build torch model with RANDOM (but real) weights.
|
|
model = sefi_model.SeFiTransformer(
|
|
cfg, dtype=torch.float32, device="cpu",
|
|
operations=comfy.ops.disable_weight_init)
|
|
sd = model.state_dict()
|
|
for k, v in sd.items():
|
|
sd[k] = torch.randn_like(v.float()) * 0.05
|
|
model.load_state_dict(sd)
|
|
model.eval()
|
|
|
|
# torch -> MLX through the real converter/loader (fp32, no quant).
|
|
cfg_mlx = infer_config_mlx(sd, semantic_channels=sem)
|
|
out_dir = tempfile.mkdtemp(prefix="sefi_parity_")
|
|
convert_state_dict_to_mlx(sd, cfg_mlx, out_dir, dtype="fp32")
|
|
dit, _ = load_sefi_dit(out_dir)
|
|
|
|
# Shared input.
|
|
B, h, w, L = 1, 8, 8, 5
|
|
x = torch.randn(B, sem + tex, h, w)
|
|
ctx = torch.randn(B, L, cfg["context_in_dim"])
|
|
t_sem = torch.full((B,), 0.7)
|
|
t_tex = torch.full((B,), 0.8)
|
|
|
|
with torch.no_grad():
|
|
y_t = model(x, ctx, t_sem, t_tex).float().numpy()
|
|
|
|
x_m = mx.array(x.numpy()).astype(mx.float32)
|
|
ctx_m = mx.array(ctx.numpy()).astype(mx.float32)
|
|
ts_m = mx.array(t_sem.numpy()).astype(mx.float32)
|
|
tt_m = mx.array(t_tex.numpy()).astype(mx.float32)
|
|
y_m = np.array(dit(x_m, ctx_m, ts_m, tt_m).astype(mx.float32))
|
|
|
|
diff = np.abs(y_t - y_m)
|
|
print(f"shapes: torch {y_t.shape} mlx {y_m.shape}")
|
|
print(f"max abs diff : {diff.max():.3e}")
|
|
print(f"mean abs diff: {diff.mean():.3e}")
|
|
print(f"torch range : [{y_t.min():.3f}, {y_t.max():.3f}]")
|
|
ok = diff.max() < 1e-3
|
|
print("PARITY OK" if ok else "PARITY MISMATCH — investigate the port")
|
|
return 0 if ok else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|