| |
| """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] |
| 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 = [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: |
| |
| |
| |
| 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()) |
|
|
