#!/usr/bin/env python3 """CPU format/schema validation for native ComfyUI OmniRef quantized checkpoints. Checks, without CUDA: * the exact quantized-linear set and native Comfy marker JSON; * storage dtypes/shapes against the pinned ComfyUI layout rules; * every dense tensor is byte-identical to the BF16 source; * the pinned ComfyUI detector recovers the expected architecture; * when possible, a strict native state-dict load on CPU through ``comfy.ops``. FP8/NVFP4 kernel execution is validated separately on the GPU. """ from __future__ import annotations import argparse import datetime as dt import hashlib import json import math from pathlib import Path import sys import torch from safetensors import safe_open QUANT_LINEAR = __import__("re").compile( r"^blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj|to_gate_compress)|mlp\.(?:fc1|fc2))\.weight$" ) def sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for block in iter(lambda: handle.read(64 * 1024 * 1024), b""): digest.update(block) return digest.hexdigest() def header(path: Path) -> tuple[list[str], dict, dict]: with safe_open(str(path), framework="pt", device="cpu") as handle: keys = list(handle.keys()) metadata = handle.metadata() or {} details = {key: (str(handle.get_slice(key).get_dtype()), list(handle.get_slice(key).get_shape())) for key in keys} return keys, metadata, details def is_nvfp4(metadata: dict) -> bool: return json.loads(metadata.get("quantization", "{}")).get("format") == "nvfp4" def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--src", required=True, type=Path, help="BF16 native Comfy source") parser.add_argument("--dst", required=True, type=Path, help="Quantized native Comfy file") parser.add_argument("--format", required=True, choices=["fp8", "nvfp4"]) parser.add_argument("--receipt", required=True, type=Path) parser.add_argument("--comfy-root", default=None, help="Pinned ComfyUI checkout for native detection/load") args = parser.parse_args() errors: list[str] = [] src_keys, src_meta, src_details = header(args.src) dst_keys, dst_meta, dst_details = header(args.dst) expected_quant = {key for key in src_keys if QUANT_LINEAR.match(key)} expected_markers = {key[: -len(".weight")] + ".comfy_quant" for key in expected_quant} if len(expected_quant) != 250: errors.append(f"expected 250 quantized linears in source, found {len(expected_quant)}") actual_base = {key for key in dst_keys if QUANT_LINEAR.match(key)} if actual_base != expected_quant: errors.append("quantized weight set differs from expected block linears") for name in expected_quant: marker_key = name[: -len(".weight")] + ".comfy_quant" if marker_key not in dst_details: errors.append(f"missing marker: {marker_key}") continue with safe_open(str(args.dst), framework="pt", device="cpu") as handle: marker = json.loads(handle.get_tensor(marker_key).numpy().tobytes()) expected_format = "float8_e4m3fn" if args.format == "fp8" else "nvfp4" if marker.get("format") != expected_format: errors.append(f"bad marker for {name}: {marker}") scale_key = name + "_scale" if scale_key not in dst_details: errors.append(f"missing scale: {scale_key}") if args.format == "nvfp4": if name + "_scale_2" not in dst_details: errors.append(f"missing global scale: {name}_scale_2") logical = src_details[name][1] rows, cols = logical expected_storage = [rows, math.ceil(cols / 16) * 8] # K/2 with 16x padding storage_shape = dst_details[name][1] if storage_shape != expected_storage: errors.append(f"{name}: NVFP4 storage {storage_shape} != {expected_storage}") scale_shape = dst_details[scale_key][1] expected_scale = [rows, math.ceil(cols / 16)] if scale_shape != expected_scale: errors.append(f"{name}: block scale {scale_shape} != {expected_scale}") if dst_details[name][0] != "U8": errors.append(f"{name}: NVFP4 storage dtype {dst_details[name][0]} != U8") else: if dst_details[name][1] != src_details[name][1]: errors.append(f"{name}: FP8 shape {dst_details[name][1]} != {src_details[name][1]}") if dst_details[name][0] != "F8_E4M3": errors.append(f"{name}: FP8 dtype {dst_details[name][0]} != F8_E4M3") if dst_details[scale_key][1] not in ([], [1]): errors.append(f"{name}: FP8 weight_scale shape {dst_details[scale_key][1]} is not per-tensor") # Dense pass-through tensors must be byte-identical. dense = [key for key in src_keys if not QUANT_LINEAR.match(key)] changed = [] with safe_open(str(args.src), framework="pt", device="cpu") as src_reader, safe_open( str(args.dst), framework="pt", device="cpu" ) as dst_reader: for name in dense: left = src_reader.get_tensor(name) right = dst_reader.get_tensor(name) if left.dtype != right.dtype or left.shape != right.shape or not torch.equal(left, right): changed.append(name) if changed: errors.append(f"dense tensors changed: {changed[:8]}") detected = None native_load = {"status": "not_requested"} if args.comfy_root: sys.path.insert(0, str(Path(args.comfy_root).resolve())) try: import comfy.model_detection import comfy.ops from comfy.ldm.minimax.model import MiniMaxH3Model from safetensors.torch import load_file state = load_file(str(args.dst), device="cpu") detected = comfy.model_detection.detect_unet_config(state, "", metadata=dst_meta) expected_arch = { "image_model": "minimax_h3", "num_layers": 50, "token_refiner_num_layers": 2, "hidden_size": 5376, "latents_dim": 24, "audio_latents_dim": 32, "attention_head_dim": 128, "num_attention_heads": 56, "ffn_hidden_size": 14336, "gate_compress": True, "pdd_steps": 32, } for key, value in expected_arch.items(): if detected.get(key) != value: errors.append(f"detected {key}={detected.get(key)!r} != {value!r}") operations = comfy.ops.mixed_precision_ops( {"mixed_ops": True}, compute_dtype=torch.float32, disabled=set() ) model = MiniMaxH3Model(**detected, dtype=torch.float32, device="cpu", operations=operations) result = model.load_state_dict(state, strict=True) native_load = { "status": "passed", "missing_keys": list(result.missing_keys), "unexpected_keys": list(result.unexpected_keys), } if result.missing_keys or result.unexpected_keys: errors.append("native strict load reported missing/unexpected keys") except Exception as exc: # noqa: BLE001 - report exact environment limitation # The format/schema checks above are authoritative on CPU. Some # comfy-kitchen builds cannot construct quantized modules without a # CUDA context; the GPU end-to-end run covers that path. native_load = {"status": "unavailable", "error": f"{type(exc).__name__}: {exc}"} receipt = { "schema": "omniref-comfy-schema-validation-v1", "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(), "format": args.format, "source": {"path": str(args.src), "sha256": sha256_file(args.src), "tensor_count": len(src_keys)}, "quantized": {"path": str(args.dst), "sha256": sha256_file(args.dst), "tensor_count": len(dst_keys)}, "checks": { "quantized_linear_count": len(expected_quant), "markers_present": sorted(expected_markers) == sorted( key for key in dst_keys if key.endswith(".comfy_quant") ), "dense_tensors_checked": len(dense), "dense_tensors_changed": changed, "detected_config": detected, "native_cpu_strict_load": native_load, }, "errors": errors, "result": "passed" if not errors else "failed", } args.receipt.parent.mkdir(parents=True, exist_ok=True) args.receipt.write_text(json.dumps(receipt, indent=2) + "\n") print(json.dumps({"result": receipt["result"], "errors": errors, "native": native_load}, indent=2)) return 0 if not errors else 1 if __name__ == "__main__": raise SystemExit(main())