110 lines
4.3 KiB
Python
110 lines
4.3 KiB
Python
"""Compare direct and upstream MiniMax H3 VAE encoding on one keyframe."""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.vae_encoder import (
|
|
MiniMaxH3VideoVAEEncoder as DirectVAEEncoder,
|
|
_causal_conv3d,
|
|
_downsample,
|
|
_group_norm_3d,
|
|
_resnet,
|
|
)
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--image", required=True)
|
|
parser.add_argument("--vae", default="/vae/minimax_h3_video_vae_fp16.safetensors")
|
|
parser.add_argument("--comfy-path", default="/opt/ComfyUI")
|
|
parser.add_argument("--tiling", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
sys.path.insert(0, args.comfy_path)
|
|
from h3_blackwell_runtime.upstream_vae import MiniMaxH3VideoVAE as UpstreamVAE # noqa: E402
|
|
|
|
|
|
def load_checkpoint(path):
|
|
if os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}:
|
|
from safetensors.torch import load
|
|
|
|
with open(path, "rb") as file:
|
|
return load(file.read())
|
|
from safetensors.torch import load_file
|
|
|
|
return load_file(path, device="cuda")
|
|
|
|
|
|
def report(name, actual, expected):
|
|
actual = actual.detach()
|
|
expected = expected.detach().to(actual.device)
|
|
delta = (actual.float() - expected.float()).abs()
|
|
print({
|
|
"stage": name,
|
|
"shape": tuple(actual.shape),
|
|
"actual_min": float(actual.min()),
|
|
"actual_max": float(actual.max()),
|
|
"expected_min": float(expected.min()),
|
|
"expected_max": float(expected.max()),
|
|
"mean_delta": float(delta.mean()),
|
|
"max_delta": float(delta.max()),
|
|
}, flush=True)
|
|
|
|
|
|
image = Image.open(args.image).convert("RGB")
|
|
pixels = torch.from_numpy(np.asarray(image).copy()).permute(2, 0, 1).unsqueeze(0).cuda().float()
|
|
pixels = pixels.div(127.5).sub(1.0)
|
|
|
|
direct = DirectVAEEncoder.from_safetensors(args.vae, device="cuda", tiling=args.tiling).eval()
|
|
upstream = UpstreamVAE(tiling=args.tiling).to("cuda").eval()
|
|
upstream = upstream.to(dtype=direct.compute_dtype)
|
|
state = load_checkpoint(args.vae)
|
|
upstream.load_state_dict(state, strict=True)
|
|
del state
|
|
|
|
trace = {}
|
|
trace_names = {"encoder.conv_in", "encoder.norm_out", "encoder.conv_out", "quant_conv"}
|
|
for level in range(6):
|
|
trace_names.update(f"encoder.down.{level}.block.{block}" for block in range(2))
|
|
if level < 4:
|
|
trace_names.add(f"encoder.down.{level}.downsample")
|
|
for name, module in upstream.named_modules():
|
|
if name in trace_names:
|
|
module.register_forward_hook(lambda _module, _inputs, output, name=name: trace.__setitem__(name, output.detach().cpu()))
|
|
|
|
report("quant_conv_weight", direct.quant_conv.weight, upstream.quant_conv.weight)
|
|
report("quant_conv_bias", direct.quant_conv.bias, upstream.quant_conv.bias)
|
|
with torch.inference_mode():
|
|
direct_latent = direct.encode(pixels.clone())
|
|
upstream_latent = upstream.encode(pixels.to(direct.compute_dtype))
|
|
report("normalized_latent", direct_latent, upstream_latent)
|
|
|
|
with torch.inference_mode():
|
|
x = pixels.unsqueeze(2)
|
|
x = (x + 1.0) * 0.5
|
|
x = (x - direct.pixel_mean.to(x)) / direct.pixel_std.to(x)
|
|
params = direct.W
|
|
x = _causal_conv3d(x, params["conv_in"][0], params["conv_in"][1], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
|
|
report("encoder.conv_in", x, trace.pop("encoder.conv_in"))
|
|
for level_index, level in enumerate(params["down"]):
|
|
for block_index, block in enumerate(level["blocks"]):
|
|
x = _resnet(x, block)
|
|
name = f"encoder.down.{level_index}.block.{block_index}"
|
|
report(name, x, trace.pop(name))
|
|
if level["down"] is not None:
|
|
x = _downsample(x, level["down"])
|
|
name = f"encoder.down.{level_index}.downsample"
|
|
report(name, x, trace.pop(name))
|
|
x = _group_norm_3d(x, params["norm_out_w"], params["norm_out_b"])
|
|
report("encoder.norm_out", x, trace.pop("encoder.norm_out"))
|
|
x = torch.nn.functional.silu(x)
|
|
x = _causal_conv3d(x, params["conv_out"][0], params["conv_out"][1], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
|
|
report("encoder.conv_out", x, trace.pop("encoder.conv_out"))
|
|
x = torch.nn.functional.conv3d(x, direct.quant_conv.weight, direct.quant_conv.bias)
|
|
report("quant_conv", x, trace.pop("quant_conv"))
|
|
report("latents_mean", direct.latents_mean, upstream.latents_mean)
|
|
report("latents_std", direct.latents_std, upstream.latents_std)
|