""" H3 Model Inspector (detect the base precision / quant format) ============================================================== Reads the loaded MODEL and reports which precision/quant format the H3 DiT is stored in: BF16, FP8 (e4m3 / e5m2), INT8 (+ convrot), NVFP4, MXFP8, ConvRot-W4A4, W4A8, or a mix. Report-only — a manual hint you read and act on. WHY IT'S FUTURE-PROOF (incl. MXFP8 "once one comes out") -------------------------------------------------------- It doesn't sniff dtypes and guess. ComfyUI tags every quantized layer at load with module.quant_format, using fixed strings it already recognizes: nvfp4, mxfp8, float8_e4m3fn, float8_e5m2, int8_tensorwise, convrot_w4a4, asym_w4a8_int8 (see comfy/ops.py). This node reads that tag. MXFP8 is already a recognized format in ComfyUI (comfy/ops.py + comfy/float.py + model_management.supports_mxfp8_compute), so the day someone ships an MXFP8 H3 checkpoint, this node labels it correctly with no change. Any brand-new tag lands under "other: " instead of crashing, so it degrades gracefully. It also reports whether YOUR card can run NVFP4 / MXFP8 natively (model_management.supports_nvfp4_compute / supports_mxfp8_compute). NOT detected: pruned vs full (that's an architecture axis — factorized AdaLN — not a quant format). Reported as a caveat, not guessed. INSTALL: drop into ComfyUI/custom_nodes/, restart. Node: MiniMax-H3 -> Model Inspector. """ # ---- pure helpers (no torch; unit-testable) ------------------------------- _FRIENDLY = { "float8_e4m3fn": "FP8 (e4m3)", "float8_e5m2": "FP8 (e5m2)", "mxfp8": "MXFP8", "nvfp4": "NVFP4", "int8_tensorwise": "INT8", "int8_tensorwise+convrot": "INT8 convrot", "convrot_w4a4": "ConvRot W4A4 (int4)", "asym_w4a8_int8": "W4A8 (int4/int8)", "bf16": "BF16", "fp16": "FP16", } # implication note per format, tied to a Blackwell 16GB context _IMPLICATION = { "NVFP4": "native on Blackwell (sm_120); half the size of INT8.", "MXFP8": "needs Blackwell + torch >= 2.10 for native compute.", "FP8 (e4m3)": "fp8 storage; runs on Ada/Blackwell.", "FP8 (e5m2)": "fp8 storage; runs on Ada/Blackwell.", "INT8": "int8 storage.", "INT8 convrot": "int8+ConvRot — needs working sm_120 kernels (absent on some 50-series setups).", "ConvRot W4A4 (int4)": "4-bit ConvRot; requires the matching custom nodes/branch.", "W4A8 (int4/int8)": "4-bit weight / 8-bit activation.", "BF16": "full precision; largest footprint, cleanest LoRA apply.", "FP16": "half precision.", } def friendly(fmt): return _FRIENDLY.get(fmt, f"other: {fmt}") def summarize(counts): """counts: {raw_format: n}. Returns (label, per_format_summary_lines). Label = the dominant NON-bf16/fp16 quant format if any (the main blocks), else the dominant plain dtype.""" quant = {k: v for k, v in counts.items() if k not in ("bf16", "fp16")} lines = [] for raw, n in sorted(counts.items(), key=lambda kv: -kv[1]): lines.append(f" {friendly(raw)}: {n} layer(s)") if quant: top = max(quant.items(), key=lambda kv: kv[1])[0] label = friendly(top) elif counts: top = max(counts.items(), key=lambda kv: kv[1])[0] label = friendly(top) else: label = "unknown" return label, lines # ---- ComfyUI node --------------------------------------------------------- def _detect(model): """Walk the DiT modules, tally quant_format tags (and dtype for the rest). Returns (label, counts, report_lines). Imports torch/mm lazily so the pure helpers above stay importable without a ComfyUI runtime.""" import torch import comfy.model_management as mm # locate the diffusion model inside the ModelPatcher dm = getattr(getattr(model, "model", None), "diffusion_model", None) if dm is None: dm = getattr(model, "model", None) or model def dtype_label(dt): return { torch.bfloat16: "bf16", torch.float16: "fp16", torch.float8_e4m3fn: "float8_e4m3fn", torch.float8_e5m2: "float8_e5m2", torch.int8: "int8_tensorwise", }.get(dt, str(dt).replace("torch.", "")) counts = {} if hasattr(dm, "modules"): for m in dm.modules(): fmt = getattr(m, "quant_format", None) if fmt is not None: # distinguish int8 convrot via the packed weight's params if fmt == "int8_tensorwise": params = getattr(getattr(m, "weight", None), "_params", None) if getattr(params, "convrot", False): fmt = "int8_tensorwise+convrot" counts[fmt] = counts.get(fmt, 0) + 1 continue w = getattr(m, "weight", None) if w is not None and hasattr(w, "dtype"): counts[dtype_label(w.dtype)] = counts.get(dtype_label(w.dtype), 0) + 1 label, lines = summarize(counts) # hardware capability for the relevant 4-bit/8-bit formats try: nv = mm.supports_nvfp4_compute() except Exception: nv = None try: mx = mm.supports_mxfp8_compute() except Exception: mx = None report = [f"Detected base precision: {label}"] report += lines impl = _IMPLICATION.get(label) if impl: report.append(f" -> {impl}") report.append(f" card supports NVFP4 compute: {nv}; MXFP8 compute: {mx}") if label == "MXFP8" and mx is False: report.append(" WARNING: MXFP8 file but this card/torch can't run it natively.") if label == "NVFP4" and nv is False: report.append(" WARNING: NVFP4 file but this card can't run it natively.") report.append(" (pruned-vs-full is a separate architecture axis; not detected here.)") return label, counts, "\n".join(report) class H3ModelInspector: CATEGORY = "MiniMax-H3/multishot" FUNCTION = "inspect" RETURN_TYPES = ("STRING", "STRING") RETURN_NAMES = ("format", "report") OUTPUT_NODE = True @classmethod def INPUT_TYPES(cls): return {"required": {"model": ("MODEL",)}} def inspect(self, model): label, _counts, report = _detect(model) return (label, report) NODE_CLASS_MAPPINGS = {"H3ModelInspector": H3ModelInspector} NODE_DISPLAY_NAME_MAPPINGS = {"H3ModelInspector": "H3 Model Inspector"} __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] if __name__ == "__main__": # exercise the pure aggregation logic with mocked layer tallies cases = { "NVFP4 file (200 main + bf16 rest)": {"nvfp4": 200, "bf16": 132}, "INT8 convrot": {"int8_tensorwise+convrot": 170, "bf16": 30}, "plain bf16": {"bf16": 340}, "FP8": {"float8_e4m3fn": 200, "bf16": 140}, "MXFP8 (future file)": {"mxfp8": 200, "bf16": 132}, "some unknown new tag": {"fp6_e3m2": 200, "bf16": 132}, } for name, counts in cases.items(): label, lines = summarize(counts) print(f"{name:38s} -> {label}")