73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
"""VAE parity: MLX decoder vs comfy torch VAE on the SAME packed latent.
|
|
|
|
Run on your Mac (needs torch+comfy+mlx). Confirms the MLX AutoencoderKL decoder
|
|
matches comfy's decode of the SeFi VAE (which does NOT apply `bn`, so bn_denorm
|
|
stays OFF here):
|
|
|
|
cd ComfyUI
|
|
python custom_nodes/ComfyUI-SeFi-MLX/sefi_mlx/parity_vae.py \
|
|
models/vae/sefi-image-5b-turbo_vae.safetensors
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
|
|
|
|
def main(vae_path=None):
|
|
here = os.path.abspath(__file__)
|
|
custom_nodes = os.path.dirname(os.path.dirname(os.path.dirname(here)))
|
|
comfy_root = os.path.dirname(custom_nodes)
|
|
sys.path.insert(0, comfy_root)
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(here)))
|
|
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
import torch
|
|
|
|
from sefi_mlx.vae_mlx import load_vae_mlx_from_torch
|
|
|
|
if vae_path is None:
|
|
vae_path = os.path.join(comfy_root, "models", "vae",
|
|
"sefi-image-5b-turbo_vae.safetensors")
|
|
print("VAE:", vae_path)
|
|
|
|
# packed texture latent (B,128,h,w)
|
|
torch.manual_seed(0)
|
|
h = w = 16
|
|
packed = torch.randn(1, 128, h, w) * 0.5
|
|
|
|
# --- comfy torch decode (unshuffle to 32-ch happens inside / via fallback) --
|
|
tvae = comfy.sd.VAE(sd=comfy.utils.load_torch_file(vae_path))
|
|
try:
|
|
timg = tvae.decode(packed)
|
|
except RuntimeError:
|
|
from einops import rearrange
|
|
up = rearrange(packed, "b (c ph pw) h w -> b c (h ph) (w pw)", ph=2, pw=2)
|
|
timg = tvae.decode(up)
|
|
timg = timg.float().cpu().numpy() # (B,H,W,3) in comfy
|
|
if timg.shape[1] in (1, 3): # (B,C,H,W) -> (B,H,W,C)
|
|
timg = np.moveaxis(timg, 1, -1)
|
|
|
|
# --- MLX decode ---------------------------------------------------------
|
|
mvae = load_vae_mlx_from_torch(vae_path)
|
|
mimg = np.array(mvae.decode_packed(mx.array(packed.numpy()), bn_denorm=False)
|
|
.astype(mx.float32))
|
|
|
|
print("torch img:", timg.shape, "mlx img:", mimg.shape)
|
|
n = min(timg.shape[1], mimg.shape[1])
|
|
a, b = timg[:, :n, :n], mimg[:, :n, :n]
|
|
diff = np.abs(a - b)
|
|
print(f"max abs diff : {diff.max():.3e}")
|
|
print(f"mean abs diff: {diff.mean():.3e}")
|
|
ok = diff.max() < 5e-2
|
|
print("VAE PARITY OK" if ok else "VAE PARITY MISMATCH — check norm/scale/order")
|
|
return 0 if ok else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main(sys.argv[1] if len(sys.argv) > 1 else None))
|