e4cfa9d1e9
Fork of dgrauet/Hunyuan3D-2.1-mlx + our generate_e2e.py CLI, env-tunable remesh (HY3D_REMESH_FACES), and HARDWARE.md. Upstream training data (mini_trainset) and demo images stripped — inference needs none of it. Full upstream history: github.com/dgrauet/Hunyuan3D-2.1-mlx
798 lines
28 KiB
Python
798 lines
28 KiB
Python
"""Compare PyTorch and MLX UNet/VAE outputs numerically.
|
|
|
|
Loads both models, runs the same input through each, and compares
|
|
intermediate activations to find where divergence occurs.
|
|
|
|
Usage:
|
|
cd /Users/dgrauet/Work/Hunyuan3D-2.1-mlx/hy3dpaint
|
|
python tests/compare_mlx_pytorch.py
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
|
|
|
import numpy as np
|
|
import gc
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# PyTorch source weights for ground-truth comparison. Override via env var
|
|
# or pass --pt-weights on CLI if you keep them elsewhere.
|
|
PT_WEIGHTS_ROOT = os.environ.get(
|
|
"HUNYUAN3D_PT_WEIGHTS_ROOT",
|
|
os.path.expanduser(
|
|
"~/Work/mlx-forge/downloads/hunyuan3d-2.1/hunyuan3d-paintpbr-v2-1"
|
|
),
|
|
)
|
|
PT_UNET_DIR = os.path.join(PT_WEIGHTS_ROOT, "unet")
|
|
PT_VAE_DIR = os.path.join(PT_WEIGHTS_ROOT, "vae")
|
|
|
|
# MLX weights: resolve HF repo ID (cached under ~/.cache/huggingface/hub).
|
|
# Override with HUNYUAN3D_MLX_WEIGHTS_DIR to use a local checkout.
|
|
_env_mlx = os.environ.get("HUNYUAN3D_MLX_WEIGHTS_DIR")
|
|
if _env_mlx:
|
|
MLX_WEIGHTS_DIR = _env_mlx
|
|
else:
|
|
from huggingface_hub import snapshot_download
|
|
MLX_WEIGHTS_DIR = snapshot_download(
|
|
repo_id="dgrauet/hunyuan3d-2.1-mlx",
|
|
allow_patterns=[
|
|
"paint_unet.safetensors",
|
|
"paint_vae.safetensors",
|
|
"paint_dino.safetensors",
|
|
],
|
|
)
|
|
|
|
# Latent size (small so both models fit comfortably)
|
|
LATENT_H, LATENT_W = 8, 8
|
|
IN_CHANNELS = 12 # 4 latent + 4 normal + 4 position
|
|
TEXT_LEN = 77
|
|
CROSS_DIM = 1024
|
|
SEED = 42
|
|
TOL = 0.01
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def separator(title: str):
|
|
print(f"\n{'=' * 70}")
|
|
print(f" {title}")
|
|
print(f"{'=' * 70}\n")
|
|
|
|
|
|
def compare_tensors(name: str, pt_np: np.ndarray, mlx_np: np.ndarray,
|
|
tol: float = TOL, detail: bool = True) -> dict:
|
|
"""Compare two numpy arrays, print a MATCH/MISMATCH verdict, return stats."""
|
|
result = {"name": name, "pt_shape": pt_np.shape, "mlx_shape": mlx_np.shape}
|
|
if pt_np.shape != mlx_np.shape:
|
|
print(f" [{name}] SHAPE MISMATCH: PT {pt_np.shape} vs MLX {mlx_np.shape}")
|
|
result["status"] = "SHAPE_MISMATCH"
|
|
return result
|
|
|
|
diff = np.abs(pt_np - mlx_np)
|
|
max_diff = float(diff.max())
|
|
mean_diff = float(diff.mean())
|
|
denom = np.maximum(np.abs(pt_np), 1e-8)
|
|
rel = diff / denom
|
|
max_rel = float(rel.max())
|
|
pt_norm = float(np.linalg.norm(pt_np.ravel()))
|
|
mlx_norm = float(np.linalg.norm(mlx_np.ravel()))
|
|
|
|
status = "MATCH" if max_diff < tol else "MISMATCH"
|
|
print(f" [{name}] shape={pt_np.shape} max_abs={max_diff:.6e} "
|
|
f"max_rel={max_rel:.6e} mean_abs={mean_diff:.6e} "
|
|
f"pt_norm={pt_norm:.4f} mlx_norm={mlx_norm:.4f} [{status}]")
|
|
if detail and max_diff > tol:
|
|
flat_idx = int(np.argmax(diff.ravel()))
|
|
coords = np.unravel_index(flat_idx, diff.shape)
|
|
print(f" worst at {coords}: PT={pt_np[coords]:.6f} "
|
|
f"MLX={mlx_np[coords]:.6f}")
|
|
print(f" first 6 PT : {pt_np.ravel()[:6]}")
|
|
print(f" first 6 MLX: {mlx_np.ravel()[:6]}")
|
|
|
|
result.update({
|
|
"status": status,
|
|
"max_abs": max_diff,
|
|
"max_rel": max_rel,
|
|
"mean_abs": mean_diff,
|
|
"pt_norm": pt_norm,
|
|
"mlx_norm": mlx_norm,
|
|
})
|
|
return result
|
|
|
|
|
|
def pt_to_nhwc(t):
|
|
"""Convert PyTorch NCHW tensor to NHWC numpy."""
|
|
import torch # noqa: F401
|
|
arr = t.detach().cpu().float().numpy()
|
|
if arr.ndim == 4:
|
|
return arr.transpose(0, 2, 3, 1)
|
|
return arr
|
|
|
|
|
|
def nhwc_to_nchw(arr):
|
|
return arr.transpose(0, 3, 1, 2)
|
|
|
|
|
|
def mx_sync(x):
|
|
import mlx.core as mx
|
|
mx.eval(x)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# PyTorch UNet loader: force in_channels=12 despite config saying 4
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def load_pt_unet():
|
|
"""Build a fresh UNet2DConditionModel (in_channels=12) and load matching
|
|
weights from the bundled checkpoint.
|
|
|
|
The checkpoint stores the full 2.5D model with keys prefixed ``unet.``
|
|
(plus an extra ``unet_dual.`` copy). We strip the prefix, drop keys that
|
|
don't correspond to standard diffusers UNet submodules, and load only the
|
|
matching subset. Everything extra (``attn_multiview``, ``attn_refview``,
|
|
``.transformer.``, ``_mr``, etc.) is simply ignored.
|
|
"""
|
|
import torch
|
|
from diffusers import UNet2DConditionModel
|
|
|
|
# Read config and override in_channels=12
|
|
import json
|
|
with open(os.path.join(PT_UNET_DIR, "config.json")) as f:
|
|
cfg = json.load(f)
|
|
cfg["in_channels"] = IN_CHANNELS
|
|
# Remove diffusers-internal keys that from_config dislikes
|
|
cfg_clean = {k: v for k, v in cfg.items() if not k.startswith("_")}
|
|
|
|
unet_pt = UNet2DConditionModel.from_config(cfg_clean)
|
|
unet_pt = unet_pt.to(torch.float32)
|
|
unet_pt.eval()
|
|
assert unet_pt.conv_in.weight.shape[1] == IN_CHANNELS
|
|
|
|
# Load checkpoint
|
|
sd = torch.load(
|
|
os.path.join(PT_UNET_DIR, "diffusion_pytorch_model.bin"),
|
|
map_location="cpu",
|
|
weights_only=False,
|
|
)
|
|
|
|
model_keys = set(unet_pt.state_dict().keys())
|
|
stripped: dict = {}
|
|
# In the 2.5D checkpoint, the vanilla BasicTransformerBlock lives one
|
|
# level deeper: ``transformer_blocks.0.transformer.<subkey>`` instead of
|
|
# ``transformer_blocks.0.<subkey>``. We rewrite those keys so they map
|
|
# onto the standard diffusers UNet state dict.
|
|
for k, v in sd.items():
|
|
if not k.startswith("unet."):
|
|
continue
|
|
nk = k[len("unet."):]
|
|
# Collapse the 2.5D wrapper: "...transformer_blocks.0.transformer.X"
|
|
# -> "...transformer_blocks.0.X"
|
|
nk2 = nk.replace(".transformer_blocks.0.transformer.",
|
|
".transformer_blocks.0.")
|
|
# Skip 2.5D-only extras: attn_multiview / attn_refview / attn_dino /
|
|
# attn1.processor / *_mr weights — these don't exist in a plain
|
|
# UNet2DConditionModel state dict.
|
|
if any(tag in nk2 for tag in (
|
|
".attn_multiview.", ".attn_refview.", ".attn_dino.",
|
|
".processor.",
|
|
)):
|
|
continue
|
|
if nk2 in model_keys:
|
|
stripped[nk2] = v
|
|
|
|
missing, unexpected = unet_pt.load_state_dict(stripped, strict=False)
|
|
print(f" PT UNet: loaded {len(stripped)} weights "
|
|
f"(missing={len(missing)}, unexpected={len(unexpected)})")
|
|
if missing:
|
|
print(f" first missing: {missing[:3]}")
|
|
return unet_pt
|
|
|
|
|
|
def load_mlx_unet_for_comparison():
|
|
"""Create MLX UNet at in_channels=12 and load ONLY stock-UNet weights.
|
|
|
|
We do NOT call _enhance_unet here so that the MLX model matches the
|
|
plain diffusers UNet2DConditionModel topology (no 2.5D modules).
|
|
"""
|
|
import mlx.core as mx
|
|
from mlx.utils import tree_flatten
|
|
from hunyuanpaintpbr_mlx.unet.unet_mlx import UNet2DConditionModelMLX
|
|
|
|
unet_mlx = UNet2DConditionModelMLX(
|
|
in_channels=IN_CHANNELS,
|
|
out_channels=4,
|
|
block_out_channels=(320, 640, 1280, 1280),
|
|
cross_attention_dim=CROSS_DIM,
|
|
attention_head_dim=(5, 10, 20, 20), # matches PT config
|
|
)
|
|
|
|
unet_w = dict(mx.load(os.path.join(MLX_WEIGHTS_DIR, "paint_unet.safetensors")))
|
|
stripped = {}
|
|
for k, v in unet_w.items():
|
|
if k.startswith("unet."):
|
|
stripped[k[5:]] = v
|
|
|
|
model_keys = set(k for k, _ in tree_flatten(unet_mlx.parameters()))
|
|
matched = [(k, v) for k, v in stripped.items() if k in model_keys]
|
|
unet_mlx.load_weights(matched)
|
|
|
|
# cast to fp32 for fair comparison
|
|
float32_params = [(k, v.astype(mx.float32))
|
|
for k, v in tree_flatten(unet_mlx.parameters())]
|
|
unet_mlx.load_weights(float32_params)
|
|
|
|
n_total = len(stripped)
|
|
n_loaded = len(matched)
|
|
print(f" Loaded {n_loaded}/{n_total} MLX UNet weights "
|
|
f"(ignored {n_total - n_loaded} 2.5D/extra tensors)")
|
|
return unet_mlx
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Part 1: Weight spot-check
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def compare_weights(results: list):
|
|
separator("PART 1: WEIGHT SPOT CHECK")
|
|
|
|
import mlx.core as mx
|
|
|
|
unet_pt = load_pt_unet()
|
|
pt_state = unet_pt.state_dict()
|
|
|
|
print(f" PT UNet in_channels (weight): {unet_pt.conv_in.weight.shape[1]}")
|
|
print(f" PT attention_head_dim (config): {unet_pt.config.attention_head_dim}")
|
|
print(f" PT use_linear_projection: {unet_pt.config.use_linear_projection}")
|
|
|
|
mlx_raw = dict(mx.load(os.path.join(MLX_WEIGHTS_DIR, "paint_unet.safetensors")))
|
|
mlx_weights = {}
|
|
for k, v in mlx_raw.items():
|
|
if k.startswith("unet."):
|
|
mlx_weights[k[5:]] = v
|
|
del mlx_raw
|
|
|
|
keys = [
|
|
"conv_in.weight",
|
|
"conv_in.bias",
|
|
"time_embedding.linear_1.weight",
|
|
"time_embedding.linear_1.bias",
|
|
"time_embedding.linear_2.weight",
|
|
"conv_out.weight",
|
|
"conv_out.bias",
|
|
"conv_norm_out.weight",
|
|
"down_blocks.0.resnets.0.norm1.weight",
|
|
"down_blocks.0.resnets.0.conv1.weight",
|
|
"down_blocks.0.resnets.0.time_emb_proj.weight",
|
|
"down_blocks.0.attentions.0.norm.weight",
|
|
"down_blocks.0.attentions.0.proj_in.weight",
|
|
"down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q.weight",
|
|
"down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_k.weight",
|
|
"down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_v.weight",
|
|
"down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_out.0.weight",
|
|
"mid_block.attentions.0.transformer_blocks.0.attn1.to_q.weight",
|
|
]
|
|
|
|
for key in keys:
|
|
if key not in pt_state:
|
|
print(f" [{key}] NOT IN PT state_dict")
|
|
continue
|
|
if key not in mlx_weights:
|
|
print(f" [{key}] NOT IN MLX weights")
|
|
continue
|
|
|
|
pt_w = pt_state[key].detach().cpu().float().numpy()
|
|
mlx_w = np.array(mlx_weights[key].astype(mx.float32))
|
|
|
|
if pt_w.ndim == 4:
|
|
# PT (O, I, kH, kW) -> MLX (O, kH, kW, I)
|
|
pt_w_n = pt_w.transpose(0, 2, 3, 1)
|
|
results.append(compare_tensors(f"W:{key}", pt_w_n, mlx_w))
|
|
elif pt_w.ndim == 2:
|
|
if pt_w.shape == mlx_w.shape:
|
|
results.append(compare_tensors(f"W:{key}", pt_w, mlx_w))
|
|
elif pt_w.T.shape == mlx_w.shape:
|
|
results.append(compare_tensors(f"W:{key}(T)", pt_w.T, mlx_w))
|
|
else:
|
|
print(f" [W:{key}] SHAPE MISMATCH: PT {pt_w.shape} MLX {mlx_w.shape}")
|
|
else:
|
|
results.append(compare_tensors(f"W:{key}", pt_w, mlx_w))
|
|
|
|
del unet_pt, pt_state, mlx_weights
|
|
gc.collect()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Part 2: Attention head dim analysis
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def check_attention_head_dims():
|
|
separator("PART 2: ATTENTION HEAD DIM ANALYSIS")
|
|
|
|
unet_pt = load_pt_unet()
|
|
|
|
print(f" PT config.attention_head_dim: {unet_pt.config.attention_head_dim}")
|
|
for i, block in enumerate(unet_pt.down_blocks):
|
|
if hasattr(block, "attentions") and len(block.attentions) > 0:
|
|
tb = block.attentions[0].transformer_blocks[0]
|
|
print(f" down_blocks[{i}]: channels={unet_pt.config.block_out_channels[i]}, "
|
|
f"heads={tb.attn1.heads}, to_q={tuple(tb.attn1.to_q.weight.shape)}")
|
|
|
|
if hasattr(unet_pt.mid_block, "attentions"):
|
|
tb = unet_pt.mid_block.attentions[0].transformer_blocks[0]
|
|
print(f" mid_block: heads={tb.attn1.heads}, "
|
|
f"to_q={tuple(tb.attn1.to_q.weight.shape)}")
|
|
|
|
del unet_pt
|
|
gc.collect()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Part 3: Timestep embedding detail
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def compare_timestep_embedding_detail(results: list):
|
|
separator("PART 3: TIMESTEP EMBEDDING")
|
|
import torch
|
|
import mlx.core as mx
|
|
from diffusers.models.embeddings import get_timestep_embedding as pt_get
|
|
from hunyuanpaintpbr_mlx.unet.blocks_mlx import get_timestep_embedding as mlx_get
|
|
|
|
t = 500.0
|
|
dim = 320
|
|
|
|
emb_pt = pt_get(torch.tensor([t]), dim,
|
|
flip_sin_to_cos=True, downscale_freq_shift=0).numpy()[0]
|
|
emb_mlx = np.array(mlx_get(mx.array([t]), dim))[0]
|
|
|
|
results.append(compare_tensors("timestep_sinusoidal", emb_pt, emb_mlx))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Part 4: UNet forward pass comparison
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def compare_unet_forward(results: list) -> dict:
|
|
separator("PART 4: UNET FORWARD PASS")
|
|
|
|
import torch
|
|
import mlx.core as mx
|
|
|
|
# --- Deterministic inputs ---
|
|
rng = np.random.RandomState(SEED)
|
|
sample_nhwc = rng.randn(1, LATENT_H, LATENT_W, IN_CHANNELS).astype(np.float32) * 0.1
|
|
text_np = np.zeros((1, TEXT_LEN, CROSS_DIM), dtype=np.float32)
|
|
timestep_val = 500
|
|
|
|
sample_pt = torch.tensor(nhwc_to_nchw(sample_nhwc))
|
|
text_pt = torch.tensor(text_np)
|
|
timestep_pt = torch.tensor([timestep_val])
|
|
|
|
# --- PyTorch ---
|
|
print("Loading PyTorch UNet...")
|
|
unet_pt = load_pt_unet()
|
|
|
|
pt_intermediates: dict = {}
|
|
|
|
def mk_hook(name, is_block=False):
|
|
def _h(module, inp, output):
|
|
val = output[0] if (is_block and isinstance(output, tuple)) else output
|
|
pt_intermediates[name] = val.detach()
|
|
return _h
|
|
|
|
handles = []
|
|
handles.append(unet_pt.conv_in.register_forward_hook(mk_hook("conv_in")))
|
|
handles.append(unet_pt.down_blocks[0].register_forward_hook(
|
|
mk_hook("down_block_0", is_block=True)))
|
|
handles.append(unet_pt.mid_block.register_forward_hook(mk_hook("mid_block")))
|
|
handles.append(unet_pt.up_blocks[0].register_forward_hook(mk_hook("up_block_0")))
|
|
handles.append(unet_pt.conv_out.register_forward_hook(mk_hook("conv_out")))
|
|
|
|
print("Running PyTorch forward pass...")
|
|
with torch.no_grad():
|
|
out_pt = unet_pt(sample_pt, timestep_pt, encoder_hidden_states=text_pt).sample
|
|
for h in handles:
|
|
h.remove()
|
|
|
|
pt_intermediates["final"] = out_pt.detach()
|
|
pt_results = {k: pt_to_nhwc(v) for k, v in pt_intermediates.items()}
|
|
|
|
print(f" PT final (NHWC): {pt_results['final'].shape} "
|
|
f"min={pt_results['final'].min():.4f} max={pt_results['final'].max():.4f}")
|
|
|
|
del unet_pt, out_pt
|
|
gc.collect()
|
|
if torch.backends.mps.is_available():
|
|
torch.mps.empty_cache()
|
|
|
|
# --- MLX ---
|
|
print("\nLoading MLX UNet...")
|
|
unet_mlx = load_mlx_unet_for_comparison()
|
|
|
|
from hunyuanpaintpbr_mlx.unet.blocks_mlx import get_timestep_embedding
|
|
|
|
sample_mlx = mx.array(sample_nhwc)
|
|
text_mlx = mx.array(text_np)
|
|
timestep_mlx = mx.array([timestep_val])
|
|
|
|
mlx_results: dict = {}
|
|
|
|
# Timestep embedding
|
|
t_emb = get_timestep_embedding(timestep_mlx, unet_mlx.time_proj_dim)
|
|
emb = unet_mlx.time_embedding(t_emb)
|
|
mx_sync(emb)
|
|
|
|
# Step 1: conv_in
|
|
x = unet_mlx.conv_in(sample_mlx)
|
|
mx_sync(x)
|
|
mlx_results["conv_in"] = np.array(x)
|
|
|
|
# Step 2: down blocks (capture first)
|
|
down_res = [x]
|
|
block = unet_mlx.down_blocks[0]
|
|
if getattr(block, "has_cross_attention", False):
|
|
x, res = block(x, emb, text_mlx)
|
|
else:
|
|
x, res = block(x, emb)
|
|
down_res.extend(res)
|
|
mx_sync(x)
|
|
mlx_results["down_block_0"] = np.array(x)
|
|
|
|
for block in unet_mlx.down_blocks[1:]:
|
|
if getattr(block, "has_cross_attention", False):
|
|
x, res = block(x, emb, text_mlx)
|
|
else:
|
|
x, res = block(x, emb)
|
|
down_res.extend(res)
|
|
|
|
# Step 3: mid
|
|
x = unet_mlx.mid_block(x, emb, text_mlx)
|
|
mx_sync(x)
|
|
mlx_results["mid_block"] = np.array(x)
|
|
|
|
# Step 4: up blocks (capture first)
|
|
first_up = True
|
|
for block in unet_mlx.up_blocks:
|
|
n_res = len(block.resnets)
|
|
res_samples = down_res[-n_res:]
|
|
down_res = down_res[:-n_res]
|
|
if getattr(block, "has_cross_attention", False):
|
|
x = block(x, emb, res_samples, text_mlx)
|
|
else:
|
|
x = block(x, emb, res_samples)
|
|
if first_up:
|
|
mx_sync(x)
|
|
mlx_results["up_block_0"] = np.array(x)
|
|
first_up = False
|
|
|
|
# Step 5: out
|
|
h = unet_mlx.conv_act(unet_mlx.conv_norm_out(x))
|
|
out = unet_mlx.conv_out(h)
|
|
mx_sync(out)
|
|
mlx_results["conv_out"] = np.array(out)
|
|
mlx_results["final"] = np.array(out)
|
|
|
|
print(f" MLX final: {mlx_results['final'].shape} "
|
|
f"min={mlx_results['final'].min():.4f} max={mlx_results['final'].max():.4f}")
|
|
|
|
# --- Compare layer by layer ---
|
|
separator("LAYER-BY-LAYER COMPARISON")
|
|
order = ["conv_in", "down_block_0", "mid_block", "up_block_0", "conv_out", "final"]
|
|
layer_results = {}
|
|
divergence = None
|
|
for name in order:
|
|
pt_arr = pt_results.get(name)
|
|
mlx_arr = mlx_results.get(name)
|
|
if pt_arr is None or mlx_arr is None:
|
|
print(f" [{name}] missing (PT={pt_arr is not None}, MLX={mlx_arr is not None})")
|
|
continue
|
|
r = compare_tensors(name, pt_arr, mlx_arr)
|
|
results.append(r)
|
|
layer_results[name] = r
|
|
if divergence is None and r.get("status") == "MISMATCH":
|
|
divergence = name
|
|
|
|
del unet_mlx
|
|
gc.collect()
|
|
|
|
return {"layers": layer_results, "divergence": divergence,
|
|
"pt": pt_results, "mlx": mlx_results}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Part 5: Dig into the divergence point
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def drill_down_first_layer(results: list):
|
|
"""If conv_in already diverges, compare conv_in weights directly.
|
|
|
|
Also manually applies the exact PyTorch conv_in via numpy math, and
|
|
compares both MLX and the reference numpy conv to PyTorch output.
|
|
This isolates weight-layout vs numerical-op bugs.
|
|
"""
|
|
separator("PART 5: DRILL-DOWN (CONV_IN ISOLATION)")
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import mlx.core as mx
|
|
|
|
rng = np.random.RandomState(SEED)
|
|
sample_nhwc = rng.randn(1, LATENT_H, LATENT_W, IN_CHANNELS).astype(np.float32) * 0.1
|
|
sample_nchw = nhwc_to_nchw(sample_nhwc)
|
|
|
|
unet_pt = load_pt_unet()
|
|
w_pt = unet_pt.conv_in.weight.detach().cpu().float().numpy() # (O, I, kH, kW)
|
|
b_pt = unet_pt.conv_in.bias.detach().cpu().float().numpy()
|
|
|
|
# PT reference
|
|
with torch.no_grad():
|
|
ref = F.conv2d(torch.tensor(sample_nchw),
|
|
torch.tensor(w_pt), torch.tensor(b_pt), padding=1)
|
|
ref_nhwc = pt_to_nhwc(ref)
|
|
|
|
# MLX with weights loaded from safetensors
|
|
unet_mlx = load_mlx_unet_for_comparison()
|
|
x_mlx = unet_mlx.conv_in(mx.array(sample_nhwc))
|
|
mx_sync(x_mlx)
|
|
mlx_out = np.array(x_mlx)
|
|
|
|
# Direct MLX conv with transposed PT weights (sanity check)
|
|
import mlx.nn as nn
|
|
conv = nn.Conv2d(IN_CHANNELS, 320, 3, padding=1)
|
|
conv.weight = mx.array(w_pt.transpose(0, 2, 3, 1)) # (O, kH, kW, I)
|
|
conv.bias = mx.array(b_pt)
|
|
direct = conv(mx.array(sample_nhwc))
|
|
mx_sync(direct)
|
|
direct_np = np.array(direct)
|
|
|
|
# MLX conv_in weight stored
|
|
mlx_w = np.array(unet_mlx.conv_in.weight)
|
|
# Compare MLX stored weight with PT transposed weight
|
|
print(" Weight tensor comparison (MLX stored vs PT transposed to OHWI):")
|
|
compare_tensors("conv_in.weight", w_pt.transpose(0, 2, 3, 1), mlx_w)
|
|
|
|
print("\n Activation comparisons:")
|
|
results.append(compare_tensors("conv_in PT-vs-MLX(loaded)", ref_nhwc, mlx_out))
|
|
results.append(compare_tensors("conv_in PT-vs-MLX(direct weight set)", ref_nhwc, direct_np))
|
|
|
|
del unet_pt, unet_mlx
|
|
gc.collect()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Part 6: Drill-down into down_block_0
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def drill_down_block0(results: list):
|
|
"""Capture PT outputs at resnet-0, attention-0, resnet-1, attention-1,
|
|
downsample inside down_blocks[0] and compare to MLX step-by-step.
|
|
|
|
Also tries a 'corrected timestep' MLX run that uses the PyTorch temb
|
|
instead of the MLX timestep embedding, to confirm/deny the timestep
|
|
embedding bug as the root cause of downstream divergence.
|
|
"""
|
|
separator("PART 6: DRILL-DOWN INTO down_blocks[0]")
|
|
|
|
import torch
|
|
import mlx.core as mx
|
|
from diffusers.models.embeddings import get_timestep_embedding as pt_get
|
|
|
|
rng = np.random.RandomState(SEED)
|
|
sample_nhwc = rng.randn(1, LATENT_H, LATENT_W, IN_CHANNELS).astype(np.float32) * 0.1
|
|
text_np = np.zeros((1, TEXT_LEN, CROSS_DIM), dtype=np.float32)
|
|
timestep_val = 500
|
|
sample_pt = torch.tensor(nhwc_to_nchw(sample_nhwc))
|
|
|
|
# --- PT forward, hooking each sub-module of down_blocks[0] ---
|
|
unet_pt = load_pt_unet()
|
|
block = unet_pt.down_blocks[0]
|
|
captured: dict = {}
|
|
|
|
def hook(key):
|
|
def _h(module, inp, output):
|
|
if isinstance(output, tuple):
|
|
output = output[0]
|
|
captured[key] = output.detach()
|
|
return _h
|
|
|
|
handles = [
|
|
block.resnets[0].register_forward_hook(hook("resnet0")),
|
|
block.attentions[0].register_forward_hook(hook("attn0")),
|
|
block.resnets[1].register_forward_hook(hook("resnet1")),
|
|
block.attentions[1].register_forward_hook(hook("attn1")),
|
|
]
|
|
if block.downsamplers:
|
|
handles.append(
|
|
block.downsamplers[0].register_forward_hook(hook("downsample"))
|
|
)
|
|
|
|
with torch.no_grad():
|
|
unet_pt(sample_pt, torch.tensor([timestep_val]),
|
|
encoder_hidden_states=torch.tensor(text_np))
|
|
for h in handles:
|
|
h.remove()
|
|
|
|
pt_sub = {k: pt_to_nhwc(v) for k, v in captured.items()}
|
|
|
|
# Also grab the PT timestep embedding for MLX "corrected" run
|
|
pt_temb_sinus = pt_get(
|
|
torch.tensor([timestep_val]).float(), 320,
|
|
flip_sin_to_cos=True, downscale_freq_shift=0,
|
|
)
|
|
pt_emb = unet_pt.time_embedding(pt_temb_sinus).detach().cpu().float().numpy()
|
|
|
|
del unet_pt
|
|
gc.collect()
|
|
|
|
# --- MLX forward: standard (buggy timestep) ---
|
|
unet_mlx = load_mlx_unet_for_comparison()
|
|
from hunyuanpaintpbr_mlx.unet.blocks_mlx import get_timestep_embedding as mlx_get
|
|
|
|
sample_mlx = mx.array(sample_nhwc)
|
|
text_mlx = mx.array(text_np)
|
|
|
|
t_emb = mlx_get(mx.array([timestep_val]), unet_mlx.time_proj_dim)
|
|
emb_buggy = unet_mlx.time_embedding(t_emb)
|
|
|
|
x0 = unet_mlx.conv_in(sample_mlx)
|
|
b = unet_mlx.down_blocks[0]
|
|
# Inspect sub-modules individually
|
|
r0 = b.resnets[0](x0, emb_buggy); mx_sync(r0)
|
|
a0 = b.attentions[0](r0, text_mlx); mx_sync(a0)
|
|
r1 = b.resnets[1](a0, emb_buggy); mx_sync(r1)
|
|
a1 = b.attentions[1](r1, text_mlx); mx_sync(a1)
|
|
|
|
mlx_sub = {
|
|
"resnet0": np.array(r0),
|
|
"attn0": np.array(a0),
|
|
"resnet1": np.array(r1),
|
|
"attn1": np.array(a1),
|
|
}
|
|
|
|
print(" [MLX using its own (buggy?) timestep embedding]")
|
|
for k in ("resnet0", "attn0", "resnet1", "attn1"):
|
|
if k in pt_sub and k in mlx_sub:
|
|
results.append(compare_tensors(f"block0.{k}", pt_sub[k], mlx_sub[k]))
|
|
|
|
# --- MLX with pytorch_compatible GroupNorm patched everywhere ---
|
|
print("\n [MLX with GroupNorm.pytorch_compatible=True applied to all GNs]")
|
|
import mlx.nn as mnn
|
|
|
|
def patch_gn(module):
|
|
for _, m in module.named_modules():
|
|
if isinstance(m, mnn.GroupNorm):
|
|
m.pytorch_compatible = True
|
|
patch_gn(unet_mlx)
|
|
|
|
x0p = unet_mlx.conv_in(sample_mlx)
|
|
r0p = b.resnets[0](x0p, emb_buggy); mx_sync(r0p)
|
|
a0p = b.attentions[0](r0p, text_mlx); mx_sync(a0p)
|
|
r1p = b.resnets[1](a0p, emb_buggy); mx_sync(r1p)
|
|
a1p = b.attentions[1](r1p, text_mlx); mx_sync(a1p)
|
|
for k, arr in [("resnet0", r0p), ("attn0", a0p),
|
|
("resnet1", r1p), ("attn1", a1p)]:
|
|
results.append(compare_tensors(
|
|
f"block0.{k} (GN fix)", pt_sub[k], np.array(arr)))
|
|
|
|
# Undo patch for the next experiment
|
|
def unpatch_gn(module):
|
|
for _, m in module.named_modules():
|
|
if isinstance(m, mnn.GroupNorm):
|
|
m.pytorch_compatible = False
|
|
unpatch_gn(unet_mlx)
|
|
|
|
# --- MLX forward: use PyTorch's temb instead ---
|
|
print("\n [MLX using PyTorch's timestep embedding]")
|
|
emb_fixed = mx.array(pt_emb)
|
|
r0f = b.resnets[0](x0, emb_fixed); mx_sync(r0f)
|
|
a0f = b.attentions[0](r0f, text_mlx); mx_sync(a0f)
|
|
r1f = b.resnets[1](a0f, emb_fixed); mx_sync(r1f)
|
|
a1f = b.attentions[1](r1f, text_mlx); mx_sync(a1f)
|
|
|
|
for k, arr in [("resnet0", r0f), ("attn0", a0f),
|
|
("resnet1", r1f), ("attn1", a1f)]:
|
|
if k in pt_sub:
|
|
results.append(compare_tensors(
|
|
f"block0.{k} (PT temb)", pt_sub[k], np.array(arr)))
|
|
|
|
del unet_mlx
|
|
gc.collect()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Write markdown report
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def write_report(results: list, forward_summary: dict, path: str):
|
|
lines = [
|
|
"# MLX vs PyTorch Numerical Comparison Report",
|
|
"",
|
|
f"- latent size: {LATENT_H}x{LATENT_W}",
|
|
f"- in_channels: {IN_CHANNELS}",
|
|
f"- tolerance: {TOL}",
|
|
"",
|
|
"## Layer-by-layer divergence",
|
|
"",
|
|
]
|
|
layers = forward_summary.get("layers", {}) if forward_summary else {}
|
|
if layers:
|
|
lines += [
|
|
"| layer | status | max_abs | max_rel | mean_abs | pt_norm | mlx_norm |",
|
|
"| --- | --- | --- | --- | --- | --- | --- |",
|
|
]
|
|
for name, r in layers.items():
|
|
lines.append(
|
|
f"| {name} | {r.get('status','?')} | "
|
|
f"{r.get('max_abs',float('nan')):.3e} | "
|
|
f"{r.get('max_rel',float('nan')):.3e} | "
|
|
f"{r.get('mean_abs',float('nan')):.3e} | "
|
|
f"{r.get('pt_norm',float('nan')):.3f} | "
|
|
f"{r.get('mlx_norm',float('nan')):.3f} |"
|
|
)
|
|
div = forward_summary.get("divergence")
|
|
lines += ["", f"**First divergent layer:** {div or 'none — all layers MATCH'}", ""]
|
|
|
|
lines += ["## All comparisons", "",
|
|
"| comparison | status | max_abs | max_rel |",
|
|
"| --- | --- | --- | --- |"]
|
|
for r in results:
|
|
if "status" not in r:
|
|
continue
|
|
lines.append(
|
|
f"| {r['name']} | {r['status']} | "
|
|
f"{r.get('max_abs', float('nan')):.3e} | "
|
|
f"{r.get('max_rel', float('nan')):.3e} |"
|
|
)
|
|
|
|
with open(path, "w") as f:
|
|
f.write("\n".join(lines) + "\n")
|
|
print(f"\nReport written to {path}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main
|
|
# ---------------------------------------------------------------------------
|
|
|
|
if __name__ == "__main__":
|
|
print("=" * 70)
|
|
print(" MLX vs PyTorch UNet Numerical Comparison")
|
|
print(f" Latent size: {LATENT_H}x{LATENT_W} in_channels={IN_CHANNELS}")
|
|
print("=" * 70)
|
|
|
|
results: list = []
|
|
forward_summary: dict = {}
|
|
|
|
sections = [
|
|
("Weight spot check", lambda: compare_weights(results)),
|
|
("Attention head dims", check_attention_head_dims),
|
|
("Timestep embedding", lambda: compare_timestep_embedding_detail(results)),
|
|
("UNet forward pass",
|
|
lambda: forward_summary.update(compare_unet_forward(results) or {})),
|
|
("Conv_in drill down", lambda: drill_down_first_layer(results)),
|
|
("down_block_0 drill down", lambda: drill_down_block0(results)),
|
|
]
|
|
|
|
for name, func in sections:
|
|
try:
|
|
func()
|
|
except Exception as e:
|
|
print(f"\n ERROR in {name}: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
separator("SUMMARY")
|
|
print(f" First divergent layer: "
|
|
f"{forward_summary.get('divergence') or 'none'}")
|
|
|
|
report_path = os.path.join(os.path.dirname(__file__), "comparison_report_auto.md")
|
|
try:
|
|
write_report(results, forward_summary, report_path)
|
|
except Exception as e:
|
|
print(f" Failed to write report: {e}")
|