39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
"""Run the direct H3 core against one matched ComfyUI capture."""
|
|
|
|
import argparse
|
|
import time
|
|
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
|
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--capture-dir", default="/artifacts/capture")
|
|
parser.add_argument("--model", default="/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")
|
|
args = parser.parse_args()
|
|
|
|
capture_dir = args.capture_dir
|
|
inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
|
|
expected = torch.load(f"{capture_dir}/output.pt", map_location="cuda", weights_only=False)
|
|
checkpoint = H3Checkpoint(args.model)
|
|
|
|
start = time.perf_counter()
|
|
model = H3PackedDenoiser.from_checkpoint(checkpoint).eval()
|
|
with torch.inference_mode():
|
|
video, audio = model(
|
|
inputs["hidden"],
|
|
inputs["timesteps"],
|
|
inputs["position_ids"],
|
|
inputs["segments"],
|
|
expected["video_segment"],
|
|
expected["audio_segment"],
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
for name, actual, reference in (("video", video, expected["video"]), ("audio", audio, expected["audio"])):
|
|
reference = reference.reshape_as(actual)
|
|
delta = (actual.float() - reference.float()).abs()
|
|
print(f"{name} shape={tuple(actual.shape)} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
|
print(f"elapsed_s={time.perf_counter() - start:.3f}")
|