diff --git a/tools/compare_vae_decoder_clip.py b/tools/compare_vae_decoder_clip.py new file mode 100644 index 0000000..986c799 --- /dev/null +++ b/tools/compare_vae_decoder_clip.py @@ -0,0 +1,66 @@ +"""Compare direct and upstream H3 VAE decoder internals on one latent clip.""" + +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("--start", type=int, default=0) +parser.add_argument("--tokens", type=int, default=8) +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 stats(name: str, a: torch.Tensor, b: torch.Tensor) -> None: + diff = (a.float() - b.float()).abs() + print({"stage": name, "max": float(diff.max()), "mean": float(diff.mean()), "shape": tuple(a.shape)}, flush=True) + + +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=False).eval() +upstream = UpstreamVAE(tiling=False).to("cuda").eval() +upstream.load_state_dict(load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda"), strict=True) + +with torch.inference_mode(): + z_d = latent[:, :, args.start:args.start + args.tokens].to(next(direct.parameters()).dtype) + z_u = z_d.clone().to(next(upstream.parameters()).dtype) + z_d = z_d * direct.latents_std.view(1, -1, 1, 1, 1).to(z_d) + direct.latents_mean.view(1, -1, 1, 1, 1).to(z_d) + z_u = z_u * upstream.latents_std.view(1, -1, 1, 1, 1).to(z_u) + upstream.latents_mean.view(1, -1, 1, 1, 1).to(z_u) + z_d = direct.post_quant_conv(z_d) + z_u = upstream.post_quant_conv(z_u) + stats("post_quant_conv", z_d, z_u) + + dd, du = direct.decoder, upstream.decoder + h_d = dd.x_embedder(z_d.flatten(2).transpose(1, 2)) + h_u = du.x_embedder(z_u.flatten(2).transpose(1, 2)) + stats("x_embedder", h_d, h_u) + + b, _, latent_t, latent_h, latent_w = z_d.shape + h_d = torch.cat((h_d, dd.register_tokens.to(h_d).expand(b, -1, -1), torch.zeros_like(h_d[:, :1])), dim=1) + h_u = torch.cat((h_u, upstream.decoder.register_tokens.to(h_u).expand(b, -1, -1), torch.zeros_like(h_u[:, :1])), dim=1) + ids_d = __import__("h3_blackwell_runtime.vae_decoder", fromlist=["create_token_ids"]).create_token_ids((latent_t, latent_h, latent_w), z_d.device, z_d.dtype).expand(b, -1, -1) + ids_d = torch.cat((ids_d, torch.zeros(b, 1 + dd.num_register_tokens, 3, device=z_d.device, dtype=z_d.dtype)), dim=1) + ids_u = __import__("h3_blackwell_runtime.upstream_vae", fromlist=["create_token_ids"]).create_token_ids((latent_t, latent_h, latent_w), z_u.device, z_u.dtype).expand(b, -1, -1) + ids_u = torch.cat((ids_u, torch.zeros(b, 1 + du.num_register_tokens, 3, device=z_u.device, dtype=z_u.dtype)), dim=1) + rope_d = dd.pos_embed(ids_d) + rope_u = du.pos_embed(ids_u) + stats("rope", rope_d, rope_u) + + for index, (block_d, block_u) in enumerate(zip(dd.transformer_blocks, du.transformer_blocks)): + h_d = block_d(h_d, rope_d) + h_u = block_u(h_u, rope_u) + stats(f"block_{index:02d}", h_d, h_u) + if index >= 5 and (h_d.float() - h_u.float()).abs().mean() > 1e-3: + break