32 lines
1.4 KiB
Python
32 lines
1.4 KiB
Python
"""Compare every direct H3 block output with one ComfyUI per-block capture."""
|
|
|
|
import argparse
|
|
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
|
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
|
|
from h3_blackwell_runtime.rope import h3_rope_rotation
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--capture-dir", default="/artifacts/capture")
|
|
parser.add_argument("--model", default="/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")
|
|
parser.add_argument("--reference-input", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
capture_dir = args.capture_dir
|
|
inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
|
|
checkpoint = H3Checkpoint(args.model)
|
|
model = H3PackedDenoiser.from_checkpoint(checkpoint).eval()
|
|
|
|
hidden = inputs["hidden"]
|
|
rotation = h3_rope_rotation(inputs["position_ids"], model.backbone.inv_freq, hidden.dtype)
|
|
with torch.inference_mode():
|
|
for index, (block, adaln) in enumerate(zip(model.backbone.blocks, model.backbone.adaln, strict=True)):
|
|
hidden = block(hidden, rotation, *adaln(inputs["timesteps"]), inputs["segments"])
|
|
expected = torch.load(f"{capture_dir}/blocks/{index:02d}.pt", map_location="cuda", weights_only=False)
|
|
delta = (hidden.float() - expected.float()).abs()
|
|
print(f"block={index:02d} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
|
if args.reference_input:
|
|
hidden = expected
|