58 lines
2.2 KiB
Python
58 lines
2.2 KiB
Python
"""Compare direct and upstream full VAE decode on a saved latent."""
|
|
|
|
import argparse
|
|
import gc
|
|
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("--frames", type=int)
|
|
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 pixelize(decoded: torch.Tensor, frames: int) -> torch.Tensor:
|
|
return ((decoded[:, :, :frames].clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu()
|
|
|
|
|
|
def summarize(a: torch.Tensor, b: torch.Tensor) -> None:
|
|
diff = (a.float() - b.float()).abs()
|
|
frame_rows = []
|
|
for index in range(diff.shape[2]):
|
|
frame = diff[:, :, index]
|
|
frame_rows.append((index, float(frame.max()), float(frame.mean()), float(torch.quantile(frame.flatten(), 0.99))))
|
|
print({
|
|
"shape": tuple(a.shape),
|
|
"max": float(diff.max()),
|
|
"mean": float(diff.mean()),
|
|
"top_max": sorted(frame_rows, key=lambda row: row[1], reverse=True)[:12],
|
|
"top_mean": sorted(frame_rows, key=lambda row: row[2], reverse=True)[:12],
|
|
}, 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")
|
|
frames = args.frames or int(state.get("frames", latent.shape[2] * 4))
|
|
|
|
with torch.inference_mode():
|
|
direct = DirectVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval()
|
|
direct_pixels = pixelize(direct.decode(latent.to(next(direct.parameters()).dtype)), frames)
|
|
del direct
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
upstream = UpstreamVAE().to("cuda", dtype=torch.float16).eval()
|
|
upstream.load_state_dict(load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda"), strict=True)
|
|
upstream_pixels = pixelize(upstream.decode(latent.to(next(upstream.parameters()).dtype)), frames)
|
|
|
|
summarize(direct_pixels, upstream_pixels)
|