| |
| """Quantize a native ComfyUI single-file OmniRef checkpoint to native FP8 or NVFP4. |
| |
| The input is the native ComfyUI layout produced by ``convert_omniref_comfy.py``: |
| Q/K/V are already fused into ``attn.qkv_proj`` and the PDD output heads use the |
| native first-head-plus-offset representation. The quantizer therefore sees one |
| fused QKV tensor and applies one Comfy weight scale to it, exactly as the |
| pinned ComfyUI runtime expects. |
| |
| Only the per-block attention/MLP projections are quantized. The token refiner, |
| AdaLN curves and projections, patch/condition/output projections, norms, RoPE, |
| and FP32 PDD heads stay dense. This matches the native ComfyUI MiniMax-H3 |
| quantized-layer policy used by the released V2/Trim Comfy checkpoints. |
| |
| Writes a JSON receipt with source/output hashes, the exact format markers, |
| per-linear reconstruction metrics, and runtime versions. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import datetime as dt |
| import hashlib |
| import importlib.metadata |
| import json |
| from pathlib import Path |
| import re |
| import sys |
| import time |
|
|
| import torch |
|
|
| |
| QUANT_LINEAR = re.compile( |
| r"^blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj|to_gate_compress)|mlp\.(?:fc1|fc2))\.weight$" |
| ) |
|
|
| FORMATS = { |
| "fp8": { |
| "layout": "TensorCoreFP8Layout", |
| "tag": {"format": "float8_e4m3fn"}, |
| "documented_name": "FP8 E4M3, per-tensor weight scale, dynamic per-tensor activation scale", |
| "activation": "dynamic-per-tensor", |
| }, |
| "nvfp4": { |
| "layout": "TensorCoreNVFP4Layout", |
| "tag": {"format": "nvfp4"}, |
| "documented_name": "NVFP4 E2M1, per-tensor global scale plus FP8 block scales (group 16), dynamic per-tensor activation scale", |
| "activation": "dynamic-per-tensor", |
| }, |
| } |
|
|
|
|
| def sha256_file(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| while True: |
| block = handle.read(64 * 1024 * 1024) |
| if not block: |
| break |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def tensor_digest(tensor: torch.Tensor) -> str: |
| return hashlib.sha256(tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()).hexdigest() |
|
|
|
|
| def safetensors_header(path: Path) -> dict: |
| from safetensors import safe_open |
|
|
| with safe_open(str(path), framework="pt", device="cpu") as handle: |
| keys = list(handle.keys()) |
| metadata = handle.metadata() or {} |
| return {"tensor_count": len(keys), "metadata": metadata, "keys": keys} |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--src", required=True, type=Path, help="Native ComfyUI BF16 single-file checkpoint") |
| parser.add_argument("--dst", required=True, type=Path, help="Output quantized single-file checkpoint") |
| parser.add_argument("--format", required=True, choices=sorted(FORMATS)) |
| parser.add_argument("--receipt", required=True, type=Path) |
| parser.add_argument("--comfyui-commit", default="b1a247c9cf9dbfda4b2c0a0c1d50950667eeaa0c") |
| parser.add_argument("--expect-tensors", type=int, default=None) |
| args = parser.parse_args() |
|
|
| if args.src == args.dst: |
| raise SystemExit("source and destination must differ") |
| if args.dst.exists(): |
| raise SystemExit(f"refusing to overwrite existing output: {args.dst}") |
| if not torch.cuda.is_available(): |
| raise SystemExit("CUDA is required for comfy-kitchen quantization") |
|
|
| import comfy_kitchen |
| import comfy_kitchen.tensor as ckt |
| from safetensors import safe_open |
| from safetensors.torch import load_file, save_file |
|
|
| started = time.time() |
| spec = FORMATS[args.format] |
| layout_cls = getattr(ckt, spec["layout"]) |
| tag_bytes = json.dumps(spec["tag"], separators=(",", ":")).encode() |
|
|
| print(f"[quantize] format={args.format} src={args.src} dst={args.dst}", flush=True) |
| src_header = safetensors_header(args.src) |
| if args.expect_tensors is not None and src_header["tensor_count"] != args.expect_tensors: |
| raise SystemExit(f"expected {args.expect_tensors} tensors, found {src_header['tensor_count']}") |
| print(f"[quantize] source tensors={src_header['tensor_count']} hashing source", flush=True) |
| src_sha256 = sha256_file(args.src) |
|
|
| |
| |
| source = load_file(str(args.src), device="cpu") |
| metadata = dict(src_header["metadata"]) |
| metadata["quantization"] = json.dumps( |
| { |
| "format": spec["tag"]["format"], |
| "layout": spec["layout"], |
| "weights": "per-tensor global scale" if args.format == "fp8" else "per-tensor global scale plus FP8 block scales, group 16", |
| "activations": spec["activation"], |
| "quantized_layers": "all 50 DiT blocks: fused qkv_proj, out_proj, to_gate_compress, mlp.fc1, mlp.fc2", |
| "dense_layers": "token_refiner, AdaLN, norms, embeddings, patch/condition projections, FP32 PDD heads", |
| "qkv_treatment": "fused Q/K/V quantized as one tensor under one weight scale", |
| "comfyui_commit": args.comfyui_commit, |
| }, |
| separators=(",", ":"), |
| ) |
| metadata["quantization_source_sha256"] = src_sha256 |
| metadata["quantization_recipe"] = "omniref-comfy-ptq-v1" |
|
|
| output: dict[str, torch.Tensor] = {} |
| linears = [] |
| for index, (name, tensor) in enumerate(source.items(), 1): |
| if not QUANT_LINEAR.match(name): |
| output[name] = tensor |
| continue |
|
|
| weight = tensor.to(device="cuda", dtype=torch.bfloat16).contiguous() |
| qdata, params = layout_cls.quantize(weight) |
| serialized = layout_cls.state_dict_tensors(qdata, params) |
| for suffix, value in serialized.items(): |
| output[name + suffix] = value.detach().cpu().contiguous() |
| output[name[: -len(".weight")] + ".comfy_quant"] = torch.tensor(list(tag_bytes), dtype=torch.uint8) |
|
|
| dequantized = layout_cls.dequantize(qdata, params).float() |
| reference = weight.float() |
| diff = dequantized - reference |
| rel_l2 = float(diff.norm() / reference.norm().clamp(min=1e-30)) |
| cosine = float(torch.nn.functional.cosine_similarity(dequantized.flatten(), reference.flatten(), dim=0)) |
| entry = { |
| "name": name, |
| "logical_shape": list(tensor.shape), |
| "weight_storage_shape": list(serialized[""].shape), |
| "weight_storage_dtype": str(serialized[""].dtype).replace("torch.", ""), |
| "reconstruction_relative_l2": rel_l2, |
| "reconstruction_cosine": cosine, |
| "reconstruction_max_abs": float(diff.abs().max()), |
| "source_amax": float(reference.abs().max()), |
| } |
| if args.format == "fp8": |
| entry["weight_scale"] = float(params.scale.reshape(-1)[0]) |
| else: |
| entry["weight_scale_2"] = float(params.scale.reshape(-1)[0]) |
| entry["weight_scale_shape"] = list(serialized["_scale"].shape) |
| entry["weight_scale_dtype"] = str(serialized["_scale"].dtype).replace("torch.", "") |
| entry["block_scale_min"] = float(serialized["_scale"].to(torch.float32).min()) |
| entry["block_scale_max"] = float(serialized["_scale"].to(torch.float32).max()) |
| linears.append(entry) |
| print( |
| f"[quantize] {index:4d}/{len(source)} {name} relL2={rel_l2:.6f} cos={cosine:.8f}", |
| flush=True, |
| ) |
| del weight, qdata, params, dequantized, reference, diff |
| torch.cuda.empty_cache() |
|
|
| expected_linears = 50 * 5 |
| if len(linears) != expected_linears: |
| raise SystemExit(f"expected {expected_linears} quantized linears, got {len(linears)}") |
|
|
| print(f"[quantize] writing {len(output)} tensors to {args.dst}", flush=True) |
| args.dst.parent.mkdir(parents=True, exist_ok=True) |
| save_file(output, str(args.dst), metadata=metadata) |
| out_sha256 = sha256_file(args.dst) |
| out_header = safetensors_header(args.dst) |
|
|
| |
| |
| pass_through = [name for name in source if not QUANT_LINEAR.match(name)] |
| identical = 0 |
| 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 pass_through: |
| 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): |
| raise SystemExit(f"pass-through tensor changed: {name}") |
| identical += 1 |
|
|
| receipt = { |
| "schema": "omniref-comfy-quant-receipt-v1", |
| "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(), |
| "format": args.format, |
| "format_description": spec["documented_name"], |
| "source": { |
| "path": str(args.src), |
| "bytes": args.src.stat().st_size, |
| "sha256": src_sha256, |
| "tensor_count": src_header["tensor_count"], |
| "metadata": src_header["metadata"], |
| }, |
| "output": { |
| "path": str(args.dst), |
| "bytes": args.dst.stat().st_size, |
| "sha256": out_sha256, |
| "tensor_count": out_header["tensor_count"], |
| }, |
| "runtime": { |
| "torch": torch.__version__, |
| "cuda": torch.version.cuda, |
| "gpu": torch.cuda.get_device_name(0), |
| "comfy_kitchen": importlib.metadata.version("comfy-kitchen"), |
| "comfyui_commit": args.comfyui_commit, |
| }, |
| "policy": { |
| "quantized_linears": "all 50 DiT blocks: fused qkv_proj, out_proj, to_gate_compress, mlp.fc1, mlp.fc2", |
| "dense_linears": "token refiner, AdaLN, norms, embeddings, patch/condition projections, FP32 PDD heads", |
| "qkv": "fused by the native Comfy conversion, then quantized as one tensor with one weight scale", |
| "activation_quantization": spec["activation"], |
| "native_comfy_layout": spec["layout"], |
| }, |
| "quantized_linear_count": len(linears), |
| "quantized_linears": linears, |
| "pass_through_tensor_count": identical, |
| "pass_through_byte_identical": True, |
| "elapsed_seconds": round(time.time() - started, 1), |
| } |
| args.receipt.parent.mkdir(parents=True, exist_ok=True) |
| args.receipt.write_text(json.dumps(receipt, indent=2) + "\n") |
| print(json.dumps({k: receipt[k] for k in ("format", "output", "quantized_linear_count", "elapsed_seconds")}, indent=2)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
