File size: 8,989 Bytes
2de4138 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | #!/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())
|