"""Compare direct keyframe FL2VA packing with a matched Comfy DiT capture.""" import argparse from pathlib import Path import numpy as np from PIL import Image import torch from h3_blackwell_runtime.checkpoint import H3Checkpoint from h3_blackwell_runtime.denoiser import H3PackedDenoiser from h3_blackwell_runtime.packing import H3PromptPacker from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner from h3_blackwell_runtime.qwen3vl_vision import Qwen3VL32BVision, build_fl2va_presentation, resize_keyframe from h3_blackwell_runtime.qwen3vl_vision import mrope_freqs_cis, mrope_position_ids from h3_blackwell_runtime.sampler import _model_sigma, beta_sigmas from h3_blackwell_runtime.t2v import random_av_latents from h3_blackwell_runtime.token_refiner import H3TokenRefiner from h3_blackwell_runtime.vae_encoder import MiniMaxH3VideoVAEEncoder parser = argparse.ArgumentParser() parser.add_argument("--capture", required=True) parser.add_argument("--first", required=True) parser.add_argument("--last", required=True) parser.add_argument("--model", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors") parser.add_argument("--qwen", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors") parser.add_argument("--vae", default="/vae/minimax_h3_video_vae_fp16.safetensors") parser.add_argument("--tokenizer", default="/opt/h3-blackwell-runtime/src/h3_blackwell_runtime/qwen25_tokenizer") parser.add_argument("--reference-vision-first") parser.add_argument("--reference-vision-last") parser.add_argument("--qwen-capture-dir") parser.add_argument("--sampler-capture-dir") parser.add_argument("--oracle-qwen-input", action="store_true") parser.add_argument("--vae-dtype", choices=("float16", "bfloat16", "float32"), default="float32") parser.add_argument("--vae-no-tiling", action="store_true") parser.add_argument("--load-dit-first", action="store_true") parser.add_argument("--preview-order", action="store_true") args = parser.parse_args() prompt = "A studio time-lapse of the same pink peony bud opening into the same fully bloomed pink peony, fixed camera, cream background." width = height = 384 requested_frames = 22 seed = 440207 def load_image(path): image = Image.open(path).convert("RGB") return torch.from_numpy(np.asarray(image).copy()).unsqueeze(0).cuda().float().div(255.0) def report(name, actual, expected): expected = expected.to(actual.device) delta = (actual.float() - expected.float()).abs() print({ "stage": name, "shape": tuple(actual.shape), "mean_delta": float(delta.mean()), "max_delta": float(delta.max()), }, flush=True) capture = torch.load(args.capture, map_location="cuda", weights_only=False) checkpoint = H3Checkpoint(args.model) dit_probe = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend="sage2").eval() if args.load_dit_first else None early_refiner = H3TokenRefiner(checkpoint, attention_backend="sage2") if args.preview_order else None early_packer = H3PromptPacker(checkpoint) if args.preview_order else None conditioner = Qwen3VLPromptConditioner(args.qwen, args.tokenizer) vision = Qwen3VL32BVision(args.qwen, device="cuda", dtype=torch.float32) if args.reference_vision_first and args.reference_vision_last: class CapturedVision: def __init__(self, paths): self.outputs = [torch.load(path, map_location="cuda", weights_only=False) for path in paths] def __call__(self, flatten, grid): output = self.outputs.pop(0) return output["merged"].cuda(), [value.cuda() for value in output["deepstack"]] vision = CapturedVision([args.reference_vision_first, args.reference_vision_last]) video, audio, frame_count = random_av_latents(width, height, requested_frames, seed) if args.sampler_capture_dir: initial = torch.load(Path(args.sampler_capture_dir) / "initial.pt", map_location="cuda", weights_only=False) direct_initial = torch.cat((video.reshape(-1), audio.reshape(-1))) report("sampler_initial", direct_initial, initial["initial_x"].reshape(-1)) presentation = build_fl2va_presentation( prompt, load_image(args.first), load_image(args.last), width=width, height=height, frame_count=frame_count, tokenizer=conditioner.tokenizer, vision=vision, text_encoder=conditioner.encoder, device="cuda", ) if args.qwen_capture_dir: qwen_capture = Path(args.qwen_capture_dir) expected_ids = torch.load(qwen_capture / "qwen_input_ids.pt", map_location="cuda", weights_only=False) print({"stage": "qwen_input_ids", "equal": torch.equal(presentation.input_ids, expected_ids), "direct_shape": tuple(presentation.input_ids.shape), "comfy_shape": tuple(expected_ids.shape)}, flush=True) direct_embeds = conditioner.encoder._embed_rows(presentation.input_ids) visual_mask = torch.zeros((1, direct_embeds.shape[1]), dtype=torch.bool, device="cuda") deepstack_by_index = {} for embed in presentation.embeds_info: start = embed["index"] end = start + embed["size"] direct_embeds[0, start:end] = embed["extra"]["merged"].to(direct_embeds) visual_mask[0, start:end] = True for index, value in enumerate(embed["extra"]["deepstack"]): deepstack_by_index.setdefault(index, []).append(value) compact_ids_path = qwen_capture / "qwen_compact_token_ids.pt" if compact_ids_path.exists(): compact_ids = torch.load(compact_ids_path, map_location="cuda", weights_only=False) print({"stage": "qwen_compact_token_ids", "equal": torch.equal(presentation.input_ids[~visual_mask], compact_ids.reshape(-1)), "direct": presentation.input_ids[~visual_mask].tolist(), "comfy": compact_ids.reshape(-1).tolist()}, flush=True) direct_deepstack = [torch.cat(values, dim=0) for _, values in sorted(deepstack_by_index.items())] expected_embeds = torch.load(qwen_capture / "qwen_input_embeds.pt", map_location="cuda", weights_only=False) raw_rows = torch.nn.functional.embedding(presentation.input_ids, conditioner.encoder.embed_tokens) raw_scales = torch.nn.functional.embedding(presentation.input_ids, conditioner.encoder.embed_scale) fp16_embeds = (raw_rows.to(torch.float16) * raw_scales.to(torch.float16)).float() bf16_embeds = (raw_rows.to(torch.bfloat16) * raw_scales.to(torch.bfloat16)).float() fp32_to_fp16_embeds = (raw_rows.float() * raw_scales.float()).half().float() fp32_to_bf16_embeds = (raw_rows.float() * raw_scales.float()).bfloat16().float() report("qwen_input_text_rows_fp16", fp16_embeds[~visual_mask], expected_embeds[~visual_mask]) report("qwen_input_text_rows_bf16", bf16_embeds[~visual_mask], expected_embeds[~visual_mask]) report("qwen_input_text_rows_fp32_to_fp16", fp32_to_fp16_embeds[~visual_mask], expected_embeds[~visual_mask]) report("qwen_input_text_rows_fp32_to_bf16", fp32_to_bf16_embeds[~visual_mask], expected_embeds[~visual_mask]) report("qwen_input_embeds", direct_embeds, expected_embeds) report("qwen_input_text_rows", direct_embeds[~visual_mask], expected_embeds.to(direct_embeds.device)[~visual_mask]) report("qwen_input_visual_rows", direct_embeds[visual_mask], expected_embeds.to(direct_embeds.device)[visual_mask]) position_ids = mrope_position_ids(presentation.embeds_info, direct_embeds.shape[1], "cuda") freqs = mrope_freqs_cis(position_ids) hidden = (expected_embeds if args.oracle_qwen_input else direct_embeds).to(conditioner.encoder.dtype) for index, layer in enumerate(conditioner.encoder.layers): hidden = layer(hidden, freqs) expected_layer = torch.load(qwen_capture / "qwen_layers" / f"{index:02d}.pt", map_location="cuda", weights_only=False) report(f"qwen_layer_{index:02d}", hidden, expected_layer) if index < len(direct_deepstack): hidden[visual_mask] = hidden[visual_mask] + direct_deepstack[index].to(hidden) expected_layer50 = torch.load(qwen_capture / "qwen_layer50.pt", map_location="cuda", weights_only=False) report("qwen_layer50", presentation.text_states, expected_layer50) vae_dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}[args.vae_dtype] vae = MiniMaxH3VideoVAEEncoder.from_safetensors(args.vae, device="cuda", dtype=vae_dtype, tiling=not args.vae_no_tiling).eval() cond_latents = [] cond_images = [] for keyframe in presentation.keyframes: resized = resize_keyframe( keyframe["image"], width, height, crop="disabled" if keyframe["resolved_frame_index"] == 0 else "center", ) cond_images.append(resized) pixels = resized.movedim(-1, 1).cuda().float().mul(2.0).sub(1.0) cond_latents.append(vae.encode(pixels)) if args.qwen_capture_dir: for index, latent in enumerate(cond_latents): captured_vae = torch.load(Path(args.qwen_capture_dir) / f"vae_keyframe_{index}.pt", map_location="cuda", weights_only=False) print({"stage": f"vae_meta_{index}", **captured_vae.get("meta", {})}, flush=True) report(f"vae_image_{index}", cond_images[index], captured_vae["image"]) report(f"vae_keyframe_{index}", latent, captured_vae["latent"]) text = (early_refiner or H3TokenRefiner(checkpoint))(presentation.text_states) packer = early_packer or H3PromptPacker(checkpoint) sigma = beta_sigmas(12, device=video.device)[0] hidden, times, segments, positions, _, _ = packer( text, video, audio, _model_sigma(sigma), text_token_tags=presentation.text_token_tags, cond_latents=cond_latents, cond_frame_indices=[keyframe["resolved_frame_index"] for keyframe in presentation.keyframes], frame_count=frame_count, seed=seed, ) expected_hidden = capture["hidden"] text_length = text.shape[1] frame_rows = (video.shape[-2] // 2) * (video.shape[-1] // 2) direct_first = hidden[text_length:text_length + frame_rows] direct_last = hidden[text_length + frame_rows:text_length + 2 * frame_rows] comfy_first = expected_hidden[text_length:text_length + frame_rows] comfy_last = expected_hidden[text_length + frame_rows:text_length + 2 * frame_rows] print({"stage": "lengths", "text": text_length, "cond_each": frame_rows, "direct_total": hidden.shape[0], "comfy_total": expected_hidden.shape[0]}, flush=True) report("text_rows", hidden[:text_length], expected_hidden[:text_length]) report("cond_first_to_first", direct_first, comfy_first) report("cond_first_to_last", direct_first, comfy_last) report("cond_last_to_last", direct_last, comfy_last) report("cond_last_to_first", direct_last, comfy_first) report("timesteps", times, capture["timesteps"]) report("positions", positions, capture["position_ids"]) print({"stage": "segments", "direct": segments, "comfy": capture["segments"]}, flush=True)