30 lines
1.1 KiB
Python
30 lines
1.1 KiB
Python
"""Compare direct block-0 AdaLN modulation and gates with Comfy."""
|
|
|
|
import argparse
|
|
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
|
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--capture-dir", required=True)
|
|
parser.add_argument("--model", required=True)
|
|
args = parser.parse_args()
|
|
|
|
inputs = torch.load(f"{args.capture_dir}/input.pt", map_location="cuda", weights_only=False)
|
|
expected = torch.load(f"{args.capture_dir}/block0_norm1_adaln.pt", map_location="cuda", weights_only=False)
|
|
model = H3PackedDenoiser.from_checkpoint(H3Checkpoint(args.model)).eval()
|
|
actual = model.backbone.adaln[0](inputs["timesteps"])
|
|
|
|
for name, value, reference in zip(
|
|
("shift", "scale", "gate_msa", "shift_mlp", "scale_mlp", "gate_mlp"),
|
|
actual,
|
|
(expected["shift"], expected["scale"], expected["gate_msa"], None, None, expected["gate_mlp"]),
|
|
strict=True,
|
|
):
|
|
if reference is None:
|
|
continue
|
|
delta = (value.float() - reference.float()).abs()
|
|
print(f"{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|