79 lines
4.3 KiB
Python
79 lines
4.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, _conv3d
|
|
|
|
|
|
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", dtype=torch.float16).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 = _conv3d(z_d, direct.post_quant_conv.weight, direct.post_quant_conv.bias)
|
|
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)
|
|
|
|
patches = h_d.shape[1] - 1 - dd.num_register_tokens
|
|
out_d = dd.proj_out(dd.norm_out(h_d))[:, :patches]
|
|
out_u = du.proj_out(du.norm_out(h_u))[:, :patches]
|
|
stats("proj_out", out_d, out_u)
|
|
|
|
out_d = out_d.view(b, latent_t, latent_h, latent_w, dd.out_channels, dd.patch_size_t, dd.patch_size, dd.patch_size)
|
|
out_u = out_u.view(b, latent_t, latent_h, latent_w, du.out_channels, du.patch_size_t, du.patch_size, du.patch_size)
|
|
out_d = out_d.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(b, dd.out_channels, latent_t * dd.patch_size_t, latent_h * dd.patch_size, latent_w * dd.patch_size)
|
|
out_u = out_u.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(b, du.out_channels, latent_t * du.patch_size_t, latent_h * du.patch_size, latent_w * du.patch_size)
|
|
stats("decoded", out_d, out_u)
|
|
|
|
pix_d = out_d.float().mul_(direct.pixel_std.to(out_d)).add_(direct.pixel_mean.to(out_d)).clamp_(0, 1).mul_(2).sub_(1)
|
|
pix_u = out_u.float().mul_(upstream.pixel_std.to(out_u)).add_(upstream.pixel_mean.to(out_u)).clamp_(0, 1).mul_(2).sub_(1)
|
|
stats("pixels", pix_d, pix_u)
|