|
|
| """Compare MiniMax H3 AdaLN modulation *outputs*, not curve-basis weights.
|
|
|
| No torch/safetensors imports; reads .safetensors arrays via NumPy memmap.
|
| This avoids PyTorch native crashes and only materializes one AdaLN layer
|
| at a time. Does not modify source files.
|
| """
|
|
|
| import argparse
|
| import csv
|
| import json
|
| import math
|
| import os
|
| import struct
|
| import numpy as np
|
|
|
| PREFIXES = ("model.diffusion_model.", "diffusion_model.")
|
| DTYPES = {"F16": np.dtype("<f2"), "BF16": np.dtype("<u2"), "F32": np.dtype("<f4")}
|
|
|
|
|
| def normalize_key(k):
|
| for p in PREFIXES:
|
| if k.startswith(p):
|
| return k[len(p) :]
|
| return k
|
|
|
|
|
| class TensorFile:
|
| def __init__(self, path):
|
| self.path = path
|
| with open(path, "rb") as f:
|
| hraw = f.read(8)
|
| if len(hraw) != 8:
|
| raise ValueError("Not a Safetensors file: " + path)
|
| length = struct.unpack("<Q", hraw)[0]
|
| if length > 500_000_000:
|
| raise ValueError("Unreasonably large Safetensors header")
|
| raw = f.read(length)
|
| if len(raw) != length:
|
| raise ValueError("Truncated Safetensors header")
|
| header = json.loads(raw)
|
| self.offset = 8 + length
|
| self.info = {}
|
| for key, info in header.items():
|
| if key == "__metadata__":
|
| continue
|
| nk = normalize_key(key)
|
| if nk in self.info:
|
| raise ValueError(f"Duplicate normalized key in {path}: {nk}")
|
| self.info[nk] = info
|
|
|
| def get(self, key):
|
| info = self.info[key]
|
| code = info["dtype"]
|
| if code not in DTYPES:
|
| raise ValueError(f"Unsupported tensor dtype {code}: {key}")
|
| dtype = DTYPES[code]
|
| shape = tuple(info["shape"])
|
| a, b = info["data_offsets"]
|
| expected_bytes = math.prod(shape) * dtype.itemsize
|
| if b - a != expected_bytes:
|
| raise ValueError(f"Invalid byte size for {key}")
|
| mmap = np.memmap(
|
| self.path, dtype=dtype, mode="r", offset=self.offset + a, shape=shape
|
| )
|
| if code == "BF16":
|
| result = (np.asarray(mmap, dtype=np.uint32) << 16).view(np.float32)
|
| else:
|
| result = np.asarray(mmap, dtype=np.float32)
|
| return result
|
|
|
|
|
| def metrics(base, target):
|
| if base.shape != target.shape:
|
| raise ValueError(f"Shape mismatch: {base.shape} and {target.shape}")
|
| a = base.astype(np.float64, copy=False).ravel()
|
| b = target.astype(np.float64, copy=False).ravel()
|
| delta = b - a
|
| d2 = float(np.dot(delta, delta))
|
| a2 = float(np.dot(a, a))
|
| b2 = float(np.dot(b, b))
|
| cross = float(np.dot(a, b))
|
| rel = math.sqrt(d2 / a2) if a2 > 0 else (0.0 if d2 == 0 else float("inf"))
|
| denom = math.sqrt(a2 * b2)
|
| cos = max(-1.0, min(1.0, cross / denom)) if denom else (1.0 if d2 == 0 else 0.0)
|
| return rel * 100.0, cos, float(np.max(np.abs(delta)))
|
|
|
|
|
| def evaluate(base_file, target_file, num_samples=9, num_blocks=50, include_final=True):
|
| bt = base_file.get("adaln_t_table")
|
| tt = target_file.get("adaln_t_table")
|
| if bt.shape != tt.shape or bt.ndim != 2:
|
| raise ValueError(f"Incompatible time tables: {bt.shape} / {tt.shape}")
|
| indices = np.unique(np.linspace(0, bt.shape[0] - 1, num_samples, dtype=np.int64))
|
| bv, tv = bt[indices], tt[indices]
|
| print(f"Timestep table shape: {bt.shape}", flush=True)
|
| print(f"Comparing rows: {list(map(int, indices))}", flush=True)
|
| table_rel, table_cos, table_max = metrics(bt, tt)
|
| print(
|
| f"RAW time table difference: {table_rel:.4f}% (cos {table_cos:.6f})", flush=True
|
| )
|
| print()
|
|
|
| names = [f"blocks.{i}.adaln_proj.linear" for i in range(num_blocks)]
|
| if include_final:
|
| names.append("final_layer.adaln_proj.linear")
|
| rows = []
|
| for i, stem in enumerate(names, 1):
|
| bw = base_file.get(stem + ".weight")
|
| tw = target_file.get(stem + ".weight")
|
| bb = base_file.get(stem + ".bias")
|
| tb = target_file.get(stem + ".bias")
|
| if bw.shape != tw.shape or bb.shape != tb.shape:
|
| raise ValueError(f"Unexpected AdaLN shape mismatch: {stem}")
|
| if bw.ndim != 2 or bw.shape[1] != bt.shape[1] or bb.shape != (bw.shape[0],):
|
| raise ValueError(f"Invalid AdaLN projection dimensions: {stem}")
|
|
|
|
|
| base_out = bv @ bw.T + bb
|
| target_out = tv @ tw.T + tb
|
| rel, cos, max_abs = metrics(base_out, target_out)
|
| wrel, wcos, _ = metrics(bw, tw)
|
| row = {
|
| "key": stem,
|
| "modulation_relerr_percent": rel,
|
| "modulation_cosine": cos,
|
| "modulation_max_abs": max_abs,
|
| "raw_weight_relerr_percent": wrel,
|
| "raw_weight_cosine": wcos,
|
| }
|
| rows.append(row)
|
| print(
|
| f"[{i:2}/{len(names)}] output rel={rel:9.4f}% cos={cos:.8f} | "
|
| f"raw W rel={wrel:9.3f}% | {stem}",
|
| flush=True,
|
| )
|
| return rows
|
|
|
|
|
| def main():
|
| ap = argparse.ArgumentParser(description=__doc__)
|
| ap.add_argument("base", help="Original Ref2VA checkpoint")
|
| ap.add_argument("target", help="FastH3 Omni checkpoint")
|
| ap.add_argument("--samples", type=int, default=9)
|
| ap.add_argument("--blocks", type=int, default=50)
|
| ap.add_argument("--no-final", action="store_true")
|
| ap.add_argument("--report", default="h3_omni_adaln_function.tsv")
|
| args = ap.parse_args()
|
| if args.samples < 2 or args.blocks < 0:
|
| ap.error("--samples must be >= 2 and --blocks must be >= 0")
|
| print("Using NumPy, no torch/CUDA. No checkpoint changes.\n", flush=True)
|
| rows = evaluate(
|
| TensorFile(args.base),
|
| TensorFile(args.target),
|
| args.samples,
|
| args.blocks,
|
| not args.no_final,
|
| )
|
| with open(args.report, "w", newline="", encoding="utf-8") as f:
|
| writer = csv.DictWriter(
|
| f, fieldnames=list(rows[0]) if rows else ["key"], delimiter="\t"
|
| )
|
| writer.writeheader()
|
| writer.writerows(rows)
|
| print("\nSUMMARY", flush=True)
|
| if rows:
|
| ordered = sorted(
|
| rows, key=lambda x: x["modulation_relerr_percent"], reverse=True
|
| )
|
| print(f"Compared layers: {len(rows)}")
|
| print(
|
| f"Mean output relative error: {sum(x['modulation_relerr_percent'] for x in rows)/len(rows):.5f}%"
|
| )
|
| print("Worst 10 output mismatches:")
|
| for x in ordered[:10]:
|
| print(
|
| f" {x['modulation_relerr_percent']:9.5f}% cos {x['modulation_cosine']:.7f} {x['key']}"
|
| )
|
| print(f"Full report: {args.report}", flush=True)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
