Mixed precision

_loaders.py

@@ -8,6 +8,7 @@ Provides five ComfyUI nodes:
   - IntCrushLoRAUnloader: removes LoRA residual buffers
 """
 
+import json
 import logging
 import torch
 
@@ -25,7 +26,7 @@ log = logging.getLogger(__name__)
 VALID_ROT_SIZES: tuple[int, ...] = (0, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096)
 
 # Safetensors format_version values that use INT-Crush quantization.
-_INT4_FORMAT_VERSIONS: tuple[str, ...] = ("2", "3")
+_INT4_FORMAT_VERSIONS: tuple[str, ...] = ("2", "3", "6")
 _INT8_FORMAT_VERSIONS: tuple[str, ...] = ("1",)
 
 # Safetensors format_version values that use OrbitQuant quantization.
@@ -169,11 +170,16 @@ class SimpleINT4UNetLoader:
             log.info("[INT-Crush] W4A8 requires rotation — falling back to W4A16")
             use_w4a16 = True
 
+        # Parse mixed precision plan if present (format_version "6")
+        mp_plan_json = metadata.get("int_crush.mp_plan") if metadata else None
+        mp_plan = json.loads(mp_plan_json) if mp_plan_json else None
+
         # Pass metadata only for RPBH — avoids triggering OrbitQuant scale-shape fallback.
         ops_cls = make_intcrush_ops("int4_crush", rot_size,
                                     use_w4a16=use_w4a16,
                                     metadata=metadata if is_rpbh else None,
-                                    kernel_backend=kernel_backend)
+                                    kernel_backend=kernel_backend,
+                                    mp_plan=mp_plan)
         model_options = {"custom_operations": ops_cls}
         _reset_aimdo_state()
 

_ops.py

@@ -8,6 +8,7 @@ per-input QuantizedTensor.from_float / cast_bias_weight round-trip.
 
 from __future__ import annotations
 
+import json
 import logging
 import os
 from typing import Any, Callable
@@ -293,6 +294,7 @@ def make_intcrush_ops(
     use_w4a16: bool = False,
     metadata: dict | None = None,
     kernel_backend: str = "auto",
+    mp_plan: dict | None = None,
 ) -> type:
     """Build a manual_cast ops class with lean INT-Crush forward.
 
@@ -303,10 +305,13 @@ def make_intcrush_ops(
         use_w4a16: disable W4A8 path, use W4A16 unpack+GEMM instead
         metadata: safetensors file metadata (used for OrbitQuant detection)
         kernel_backend: "auto" (CUDA→Triton→PyTorch), "triton", "cuda", or "pytorch"
+        mp_plan: per-layer mixed precision plan {name: {"w": 4|8, "a": 4|8}} or None
     """
     # Use a hash of the orbitquant_dims as cache key differentiator
     _oq_key = metadata.get("int_crush.orbitquant_dims", "") if metadata else ""
-    cache_key = (quant_format, rot_size, use_pytorch, use_w4a16, _oq_key, kernel_backend)
+    _mp_key = json.dumps(mp_plan, sort_keys=True) if mp_plan else ""
+    cache_key = (quant_format, rot_size, use_pytorch, use_w4a16,
+                 _oq_key, kernel_backend, _mp_key)
     if cache_key in _ops_cache:
         return _ops_cache[cache_key]
     from comfy.ops import manual_cast
@@ -315,6 +320,7 @@ def make_intcrush_ops(
         _orbitquant_metadata: dict | None = metadata  # class-level attribute
         _orbitquant_model_path: str = ""  # set during first _load_from_state_dict
         _rpbh_metadata: dict | None = None  # class-level RPBH metadata cache
+        _mp_plan: dict | None = mp_plan  # mixed precision plan (per-layer W/A bit-widths)
 
         class Linear(manual_cast.Linear):
             """INT-Crush Linear: QuantizedTensor storage with lean forward()."""
@@ -364,6 +370,9 @@ def make_intcrush_ops(
                 self._intcrush_rpbh_h: int | None = None               # Hadamard block size
                 self._intcrush_rpbh_d: int | None = None               # Original input dim
                 self._intcrush_rpbh_d_padded: int | None = None        # Padded input dim
+                # ── Mixed precision state ──
+                self._intcrush_mp_weight_bits: int | None = None
+                self._intcrush_mp_act_bits: int | None = None
 
 
             def _apply_adapter(self, x_raw: torch.Tensor, out: torch.Tensor) -> torch.Tensor:
@@ -1363,6 +1372,27 @@ def make_intcrush_ops(
                 self._intcrush_rot_size = rot_size
                 self._intcrush_is_quantized = True
 
+                # ── Mixed precision: read per-layer bit-widths ──
+                mp_plan = getattr(type(self), '_mp_plan', None)
+                if mp_plan:
+                    bare_name = prefix.rstrip(".")
+                    mp_entry = mp_plan.get(bare_name)
+                    if mp_entry is None:
+                        # Fallback: strip common prefixes
+                        for strip_prefix in ("model.", "model.diffusion_model."):
+                            if bare_name.startswith(strip_prefix):
+                                mp_entry = mp_plan.get(
+                                    bare_name[len(strip_prefix):])
+                                if mp_entry is not None:
+                                    break
+                    if mp_entry is not None:
+                        self._intcrush_mp_weight_bits = mp_entry["w"]
+                        self._intcrush_mp_act_bits = mp_entry["a"]
+                        _log.debug("mp_plan: %s -> W%dA%d",
+                                   bare_name, mp_entry["w"], mp_entry["a"])
+                    else:
+                        _log.debug("mp_plan: no entry for '%s'", bare_name)
+
                 perm = tensors["perm"]
                 zp = tensors["zp"]
 

pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "setuptools.backends._legacy:_Backend"
 
 [project]
 name = "ComfyUI-INT-Crush"
-version = "0.1.8"
-description = "INT4/INT8 quantized inference with ConvRot, PermuQuant, SmoothQuant, SmoothRot and OrbitQuant support"
+version = "0.1.9"
+description = "INT4/INT8 quantized inference with ConvRot, PermuQuant, SmoothQuant, SmoothRot, OrbitQuant and mixed precision support"
 requires-python = ">=3.10"
 dependencies = [
     "torch>=2.1",