193 lines
10 KiB
Python
193 lines
10 KiB
Python
"""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)
|