Add VAE temporal assembly tracer

tools/compare_vae_temporal_assembly.py

@@ -0,0 +1,81 @@
+"""Trace direct/upstream temporal VAE assembly from matching tiled clips."""
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+
+from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE as DirectVAE
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--latent", type=Path, required=True)
+parser.add_argument("--comfy-path", default="/opt/ComfyUI")
+args = parser.parse_args()
+
+sys.path.insert(0, args.comfy_path)
+from h3_blackwell_runtime.upstream_vae import MiniMaxH3VideoVAE as UpstreamVAE  # noqa: E402
+
+
+def diff_stats(label: str, a: torch.Tensor, b: torch.Tensor) -> None:
+    diff = (a.float() - b.float()).abs()
+    print({"label": label, "shape": tuple(a.shape), "max": float(diff.max()), "mean": float(diff.mean())}, flush=True)
+
+
+def plan(vae, z_len: int):
+    pseudo = z_len + vae.token_drop
+    pad = (-pseudo) % vae.tokens_chunk_size
+    pseudo += pad
+    chunks = pseudo // vae.tokens_chunk_size - int(vae.token_drop > 0)
+    if chunks < 1:
+        pad += vae.tokens_chunk_size
+        chunks += 1
+    return chunks, pad, vae._decode_temporal_frame_plan(z_len + pad, chunks, pad)
+
+
+state = torch.load(args.latent, map_location="cuda", weights_only=False)
+latent = state["latent"].to("cuda") if isinstance(state, dict) else state.to("cuda")
+
+direct = DirectVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda", tiling=True).eval()
+upstream = UpstreamVAE(tiling=True).to("cuda").eval()
+upstream.load_state_dict(load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda"), strict=True)
+
+with torch.inference_mode():
+    zd = latent.to(next(direct.parameters()).dtype)
+    zu = latent.to(next(upstream.parameters()).dtype)
+    zd = zd * direct.latents_std.view(1, -1, 1, 1, 1).to(zd) + direct.latents_mean.view(1, -1, 1, 1, 1).to(zd)
+    zu = zu * upstream.latents_std.view(1, -1, 1, 1, 1).to(zu) + upstream.latents_mean.view(1, -1, 1, 1, 1).to(zu)
+    chunks, pad, frames = plan(direct, zd.shape[2])
+    if pad:
+        zd = torch.cat((zd, zd[:, :, -1:].repeat(1, 1, pad, 1, 1)), dim=2)
+        zu = torch.cat((zu, zu[:, :, -1:].repeat(1, 1, pad, 1, 1)), dim=2)
+    print({"chunks": chunks, "pad": pad, "frames": frames}, flush=True)
+
+    chunk_dec = direct.tokens_chunk_size * direct.vae_ratio_t
+    split_count = int(direct.token_drop > 0) + 1
+    prev_d = prev_u = None
+    write_pos = 0
+    for chunk_index in range(chunks):
+        start = chunk_index * direct.tokens_chunk_size
+        end = start + direct.tokens_chunk_size + direct.token_overlap
+        clip_d = direct.tiled_decode(zd[:, :, start:end])
+        clip_u = upstream.tiled_decode(zu[:, :, start:end])
+        diff_stats(f"clip_{chunk_index}", clip_d, clip_u)
+        for split in range(split_count):
+            frame_start = split * chunk_dec
+            frame_end = min(frame_start + chunk_dec, clip_d.shape[2])
+            part_d = clip_d[:, :, frame_start:frame_end][:, :, direct.frame_pre_padding:]
+            part_u = clip_u[:, :, frame_start:frame_end][:, :, upstream.frame_pre_padding:]
+            if split == 0 and prev_d is not None:
+                part_d = direct.blend(prev_d, part_d, direct.frame_overlap, -3)
+                part_u = upstream.blend(prev_u, part_u, upstream.frame_overlap, -3)
+                diff_stats(f"chunk_{chunk_index}_blended_at_{write_pos}", part_d, part_u)
+            elif split == 0:
+                diff_stats(f"chunk_{chunk_index}_write_at_{write_pos}", part_d, part_u)
+            else:
+                prev_d, prev_u = part_d.contiguous(), part_u.contiguous()
+                diff_stats(f"chunk_{chunk_index}_overlap", prev_d, prev_u)
+                continue
+            write_pos += part_d.shape[2]