diff --git a/_loaders.py b/_loaders.py index e08a0a1..27de81b 100644 --- a/_loaders.py +++ b/_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() diff --git a/_ops.py b/_ops.py index b07edc8..2e3a860 100644 --- a/_ops.py +++ b/_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"] diff --git a/pyproject.toml b/pyproject.toml index 0d57d8f..638efc7 100644 --- a/pyproject.toml +++ b/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",