From 4e17e866a028a250f487068d2d2d0bfc583ad8d6 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Fri, 14 Aug 2026 01:33:33 +0700 Subject: [PATCH] Add VAE temporal assembly tracer --- tools/compare_vae_temporal_assembly.py | 81 ++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) create mode 100644 tools/compare_vae_temporal_assembly.py diff --git a/tools/compare_vae_temporal_assembly.py b/tools/compare_vae_temporal_assembly.py new file mode 100644 index 0000000..2643d7c --- /dev/null +++ b/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]