67 lines
3.3 KiB
Python
67 lines
3.3 KiB
Python
|
"""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
|