50 lines
2.1 KiB
Python
50 lines
2.1 KiB
Python
|
"""Compare direct prompt-only FL2VA packing with an immutable Comfy DiT input."""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.nn.functional as functional
|
||
|
|
|
||
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||
|
|
from h3_blackwell_runtime.packing import H3PromptPacker, pack_audio, patchify_video
|
||
|
|
from h3_blackwell_runtime.sampler import _audio_sigma
|
||
|
|
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--capture", type=Path, required=True)
|
||
|
|
parser.add_argument("--refiner-trace", type=Path, required=True)
|
||
|
|
parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
captured = torch.load(args.capture, map_location="cuda", weights_only=False)
|
||
|
|
text = torch.load(args.refiner_trace / "refiner_output.pt", map_location="cuda", weights_only=False)
|
||
|
|
if text.ndim == 2:
|
||
|
|
text = text.unsqueeze(0)
|
||
|
|
checkpoint = H3Checkpoint(args.checkpoint)
|
||
|
|
packer = H3PromptPacker(checkpoint)
|
||
|
|
|
||
|
|
video_shape = (1, 24, 7, 12, 20)
|
||
|
|
audio_shape = (1, 32, 2, 37)
|
||
|
|
video = captured["video_x"].to("cuda")
|
||
|
|
audio = captured["audio_x"].to("cuda")
|
||
|
|
sigma = 1 - captured["timesteps"].min()
|
||
|
|
|
||
|
|
with torch.inference_mode():
|
||
|
|
hidden, times, segments, positions, _, _ = packer(text, video, audio, float(sigma))
|
||
|
|
text_rows = text[0].to(torch.bfloat16)
|
||
|
|
audio_rows = functional.linear(pack_audio(audio).float(), packer.audio_weight, packer.audio_bias).to(torch.bfloat16)
|
||
|
|
video_rows = functional.linear(patchify_video(video).float(), packer.video_weight, packer.video_bias).to(torch.bfloat16)
|
||
|
|
|
||
|
|
for name, actual, expected in (
|
||
|
|
("text", text_rows, captured["hidden"][:17]),
|
||
|
|
("audio", audio_rows, captured["hidden"][17:91]),
|
||
|
|
("video", video_rows, captured["hidden"][91:]),
|
||
|
|
("hidden", hidden, captured["hidden"]),
|
||
|
|
("times", times, captured["timesteps"]),
|
||
|
|
("positions", positions, captured["position_ids"]),
|
||
|
|
):
|
||
|
|
delta = (actual.float() - expected.to(actual.device).float()).abs()
|
||
|
|
print(f"{name} shape={tuple(actual.shape)} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
||
|
|
print(f"segments direct={segments} comfy={captured['segments']}")
|