739 lines
29 KiB
Python
739 lines
29 KiB
Python
"""ComfyUI loader nodes for INT-Crush quantized models.
|
|
|
|
Provides five ComfyUI nodes:
|
|
- SimpleINT4UNetLoader: loads a model quantized with INT4 weights
|
|
- SimpleINT8UNetLoader: loads a model quantized with INT8 weights
|
|
- OrbitQuantLoader: loads a model quantized with OrbitQuant (RPBH + codebook)
|
|
- IntCrushLoRALoader: attaches a LoRA as a residual buffer
|
|
- IntCrushLoRAUnloader: removes LoRA residual buffers
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import torch
|
|
import uuid
|
|
|
|
from ._ops import make_intcrush_ops
|
|
from ._lora import (
|
|
clear_intcrush_lora,
|
|
parse_lora_adapters,
|
|
apply_lora_stack,
|
|
rotated_adapter_for_module,
|
|
)
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
# ── Shared constants ─────────────────────────────────────────────────────────
|
|
|
|
# Valid rotation-group sizes (0 = disable rotation, fall back to W4A16 for INT4).
|
|
# Powers of 4 (16, 64, 256, 1024, 4096) use Regular Hadamard (balanced row sums,
|
|
# preferred for quantization quality). Other powers of 2 (32, 128, 512, 2048)
|
|
# use Sylvester Hadamard fallback.
|
|
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", "6")
|
|
_INT8_FORMAT_VERSIONS: tuple[str, ...] = ("1",)
|
|
|
|
# Safetensors format_version values that use OrbitQuant quantization.
|
|
_ORBITQUANT_FORMAT_VERSIONS: tuple[str, ...] = ("4",)
|
|
|
|
|
|
# ── aimdo integration ──────────────────────────────────────────────────────
|
|
|
|
def _reset_aimdo_state():
|
|
"""Flush aimdo VBAR watermarks, PyTorch caching allocator, and OrbitQuant cache.
|
|
|
|
aimdo's watermark tracks which weights have been evicted. Without a
|
|
reset, a newly-loaded model inherits stale eviction priorities from the
|
|
previous model, causing unnecessary re-faults.
|
|
|
|
WARNING: vbars_reset_watermark_limits() is global — it resets ALL
|
|
VBARs across all devices and models. If other models are loaded
|
|
concurrently, this will cause unnecessary re-faults for their
|
|
already-resident weights. Only call this when switching models.
|
|
"""
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
try:
|
|
from ._orbitquant import clear_orbitquant_globals
|
|
clear_orbitquant_globals()
|
|
except ImportError:
|
|
pass
|
|
try:
|
|
from comfy_aimdo.model_vbar import vbars_reset_watermark_limits
|
|
vbars_reset_watermark_limits()
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
# ── Loader helpers ───────────────────────────────────────────────────────────
|
|
|
|
def _get_diffusion_model_list() -> list[str]:
|
|
try:
|
|
import folder_paths
|
|
return folder_paths.get_filename_list("diffusion_models")
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
def _get_lora_list() -> list[str]:
|
|
try:
|
|
import folder_paths
|
|
return folder_paths.get_filename_list("loras")
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
def _load_metadata(unet_path: str) -> dict:
|
|
"""Load safetensors metadata, warning on errors instead of swallowing them.
|
|
|
|
Returns ``{}`` for files with no metadata (e.g. a .pt checkpoint) so callers
|
|
can proceed with defaults. On a raised exception, logs at WARNING with the
|
|
path and error so a wrong-rot_size / wrong-format-version mismatch is
|
|
diagnosable rather than silently producing wrong output.
|
|
"""
|
|
import comfy.utils
|
|
try:
|
|
_, metadata = comfy.utils.load_torch_file(unet_path, return_metadata=True)
|
|
except Exception as exc:
|
|
log.warning(
|
|
"[INT-Crush] Failed to load metadata from '%s': %s. "
|
|
"Falling back to the node's defaults, which may produce incorrect "
|
|
"output if the model was quantized with non-default settings.",
|
|
unet_path, exc,
|
|
)
|
|
return {}
|
|
return metadata or {}
|
|
|
|
|
|
def _detect_rot_size(metadata: dict | None, default_rot_size: int, format_versions: tuple[str, ...]) -> int:
|
|
"""Read rot_size from safetensors metadata if available."""
|
|
if not metadata:
|
|
return default_rot_size
|
|
if metadata.get("int_crush.format_version") not in format_versions:
|
|
return default_rot_size
|
|
detected = metadata.get("int_crush.rot_size")
|
|
if detected is not None:
|
|
try:
|
|
detected = int(detected)
|
|
if detected in VALID_ROT_SIZES:
|
|
log.info("[INT-Crush] Auto-detected rot_size=%d from metadata", detected)
|
|
return detected
|
|
except (ValueError, TypeError):
|
|
pass
|
|
log.warning(
|
|
"[INT-Crush] Metadata confirms INT-Crush format but rot_size is missing; "
|
|
"using default rot_size=%d. This may produce incorrect results if the "
|
|
"model was quantized with a different rot_size (including Sylvester "
|
|
"sizes like 32, 128, 512, 2048).",
|
|
default_rot_size,
|
|
)
|
|
return default_rot_size
|
|
|
|
|
|
def _detect_orbitquant(metadata: dict | None) -> bool:
|
|
"""Return True if metadata indicates an OrbitQuant model."""
|
|
if not metadata:
|
|
return False
|
|
return (
|
|
metadata.get("int_crush.format_version") == "4"
|
|
and metadata.get("int_crush.orbitquant") == "true"
|
|
)
|
|
|
|
|
|
def _detect_rpbh(metadata: dict | None) -> bool:
|
|
"""Return True if metadata indicates an RPBH uniform-grid model
|
|
(non-OrbitQuant RPBH used with rtn/gptq/ldlq)."""
|
|
if not metadata:
|
|
return False
|
|
fv = metadata.get("int_crush.format_version")
|
|
return (
|
|
fv in _INT4_FORMAT_VERSIONS + _INT8_FORMAT_VERSIONS
|
|
and metadata.get("int_crush.rpbh") == "true"
|
|
)
|
|
|
|
|
|
# ── ComfyUI Node: INT4 Loader ────────────────────────────────────────────────
|
|
|
|
class SimpleINT4UNetLoader:
|
|
"""ComfyUI node: loads a INT-Crush-quantized model with INT4 weights."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"unet_name": (_get_diffusion_model_list(),),
|
|
"rot_size": ([s for s in VALID_ROT_SIZES], {"default": 256}),
|
|
},
|
|
"optional": {
|
|
"kernel_backend": (["auto", "triton", "cuda", "pytorch"],
|
|
{"default": "auto"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "loaders/INT-Crush"
|
|
|
|
def load(self, unet_name: str, rot_size: int,
|
|
kernel_backend: str = "auto") -> tuple[object]:
|
|
import folder_paths
|
|
import comfy.utils
|
|
from comfy.sd import load_diffusion_model
|
|
|
|
unet_path = folder_paths.get_full_path("diffusion_models", unet_name)
|
|
if unet_path is None:
|
|
raise FileNotFoundError(f"Model not found: {unet_name}")
|
|
|
|
metadata = _load_metadata(unet_path)
|
|
|
|
rot_size = _detect_rot_size(metadata, rot_size, _INT4_FORMAT_VERSIONS)
|
|
is_rpbh = _detect_rpbh(metadata)
|
|
|
|
use_w4a16 = False
|
|
if rot_size == 0 and not is_rpbh:
|
|
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,
|
|
mp_plan=mp_plan)
|
|
model_options = {"custom_operations": ops_cls}
|
|
_reset_aimdo_state()
|
|
|
|
from ._ops import _set_rpbh_metadata_cache
|
|
|
|
# Pre-load RPBH globals after _reset_aimdo_state cleared them.
|
|
if is_rpbh:
|
|
from ._orbitquant import load_rpbh_globals
|
|
_set_rpbh_metadata_cache(metadata)
|
|
full_sd = comfy.utils.load_torch_file(str(unet_path))
|
|
model_path = (metadata or {}).get("int_crush.rpbh_dims", "default")
|
|
load_rpbh_globals(full_sd, metadata, model_path)
|
|
del full_sd
|
|
log.info("[INT-Crush] RPBH uniform-grid detected (INT4), pre-loaded globals")
|
|
else:
|
|
_set_rpbh_metadata_cache(None)
|
|
|
|
try:
|
|
model = load_diffusion_model(unet_path, model_options=model_options)
|
|
finally:
|
|
if is_rpbh:
|
|
from ._ops import _set_rpbh_metadata_cache as _clear_meta
|
|
_clear_meta(None)
|
|
|
|
# RPBH may pad in_features — apply padded layer fixes.
|
|
if is_rpbh:
|
|
padded_str = (metadata or {}).get("int_crush.padded_layers", "")
|
|
if padded_str:
|
|
_apply_padded_layer_fixes(model, padded_str)
|
|
|
|
_register_lora_sync(model)
|
|
return (model,)
|
|
|
|
|
|
# ── ComfyUI Node: INT8 Loader ────────────────────────────────────────────────
|
|
|
|
class SimpleINT8UNetLoader:
|
|
"""ComfyUI node: loads a INT-Crush-quantized model with INT8 weights."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"unet_name": (_get_diffusion_model_list(),),
|
|
"rot_size": ([s for s in VALID_ROT_SIZES], {"default": 256}),
|
|
},
|
|
"optional": {
|
|
"kernel_backend": (["auto", "triton", "cuda", "pytorch"],
|
|
{"default": "auto"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "loaders/INT-Crush"
|
|
|
|
def load(self, unet_name: str, rot_size: int,
|
|
kernel_backend: str = "auto") -> tuple[object]:
|
|
import folder_paths
|
|
import comfy.utils
|
|
from comfy.sd import load_diffusion_model
|
|
|
|
unet_path = folder_paths.get_full_path("diffusion_models", unet_name)
|
|
if unet_path is None:
|
|
raise FileNotFoundError(f"Model not found: {unet_name}")
|
|
|
|
metadata = _load_metadata(unet_path)
|
|
|
|
rot_size = _detect_rot_size(metadata, rot_size, _INT8_FORMAT_VERSIONS)
|
|
|
|
is_rpbh = _detect_rpbh(metadata)
|
|
|
|
# Pass metadata only for RPBH — avoids triggering OrbitQuant scale-shape fallback.
|
|
ops_cls = make_intcrush_ops("int8_crush", rot_size,
|
|
metadata=metadata if is_rpbh else None,
|
|
kernel_backend=kernel_backend)
|
|
model_options = {"custom_operations": ops_cls}
|
|
_reset_aimdo_state()
|
|
|
|
from ._ops import _set_rpbh_metadata_cache
|
|
|
|
# Pre-load RPBH globals after _reset_aimdo_state cleared them.
|
|
if is_rpbh:
|
|
from ._orbitquant import load_rpbh_globals
|
|
_set_rpbh_metadata_cache(metadata)
|
|
full_sd = comfy.utils.load_torch_file(str(unet_path))
|
|
model_path = (metadata or {}).get("int_crush.rpbh_dims", "default")
|
|
load_rpbh_globals(full_sd, metadata, model_path)
|
|
del full_sd
|
|
log.info("[INT-Crush] RPBH uniform-grid detected (INT8), pre-loaded globals")
|
|
else:
|
|
_set_rpbh_metadata_cache(None)
|
|
|
|
try:
|
|
model = load_diffusion_model(unet_path, model_options=model_options)
|
|
finally:
|
|
if is_rpbh:
|
|
from ._ops import _set_rpbh_metadata_cache as _clear_meta
|
|
_clear_meta(None)
|
|
|
|
# Fix model config for padded layers (rotation padding inflates in_features).
|
|
padded_str = (metadata or {}).get("int_crush.padded_layers", "")
|
|
if padded_str:
|
|
_apply_padded_layer_fixes(model, padded_str)
|
|
|
|
_register_lora_sync(model)
|
|
return (model,)
|
|
|
|
|
|
# ── ComfyUI Node: OrbitQuant Loader ─────────────────────────────────────────
|
|
|
|
class OrbitQuantLoader:
|
|
"""ComfyUI node: loads an OrbitQuant-quantized model.
|
|
|
|
OrbitQuant uses a calibration-free RPBH rotation + Lloyd-Max codebook
|
|
for weight-activation quantization (arXiv:2607.02461). Format version 4.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"unet_name": (_get_diffusion_model_list(),),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "loaders/INT-Crush"
|
|
|
|
def load(self, unet_name: str) -> tuple[object]:
|
|
import folder_paths
|
|
import comfy.utils
|
|
from comfy.sd import load_diffusion_model
|
|
|
|
unet_path = folder_paths.get_full_path("diffusion_models", unet_name)
|
|
if unet_path is None:
|
|
raise FileNotFoundError(f"Model not found: {unet_name}")
|
|
|
|
# Load metadata for validation
|
|
metadata = _load_metadata(str(unet_path))
|
|
|
|
if not _detect_orbitquant(metadata):
|
|
raise ValueError(
|
|
"Not an OrbitQuant model (expected format_version=4 and "
|
|
"int_crush.orbitquant=true). Use SimpleINT4UNetLoader or "
|
|
"SimpleINT8UNetLoader for other INT-Crush formats."
|
|
)
|
|
|
|
_reset_aimdo_state()
|
|
|
|
ops_cls = make_intcrush_ops("orbitquant", rot_size=0, metadata=metadata)
|
|
model_options = {"custom_operations": ops_cls}
|
|
|
|
# Cache metadata for _load_from_state_dict (global tensors are
|
|
# invisible during per-module loading).
|
|
from ._ops import _set_orbitquant_metadata_cache
|
|
from ._orbitquant import load_orbitquant_globals, clear_orbitquant_globals
|
|
_set_orbitquant_metadata_cache(metadata)
|
|
|
|
# Pre-load global RPBH/codebook tensors into the cache.
|
|
full_sd = comfy.utils.load_torch_file(str(unet_path))
|
|
model_path = metadata.get("int_crush.orbitquant_dims", "default")
|
|
load_orbitquant_globals(full_sd, metadata, model_path)
|
|
del full_sd
|
|
|
|
try:
|
|
model = load_diffusion_model(str(unet_path), model_options=model_options)
|
|
finally:
|
|
from ._ops import _set_orbitquant_metadata_cache as _clear_meta
|
|
_clear_meta(None)
|
|
clear_orbitquant_globals()
|
|
|
|
_register_lora_sync(model)
|
|
return (model,)
|
|
|
|
|
|
def _apply_padded_layer_fixes(model, padded_str: str) -> None:
|
|
"""Trim weight rows and fix in_channels for layers whose in_features
|
|
were padded during quantization.
|
|
|
|
*padded_str* is a semicolon-separated ``"key=orig_in_features"`` list
|
|
stored in safetensors metadata under ``int_crush.padded_layers``.
|
|
"""
|
|
import comfy.utils
|
|
|
|
padded_map = {}
|
|
for entry in padded_str.split(";"):
|
|
if "=" in entry:
|
|
k, v = entry.split("=", 1)
|
|
padded_map[k.strip()] = int(v.strip())
|
|
|
|
m = model.model
|
|
for layer_key, orig_in in padded_map.items():
|
|
module_path = layer_key.rsplit(".weight", 1)[0]
|
|
try:
|
|
module = comfy.utils.get_attr(m, module_path)
|
|
except AttributeError:
|
|
continue
|
|
|
|
if hasattr(module, 'weight') and module.weight is not None:
|
|
weight = module.weight
|
|
if hasattr(weight, '_qdata'):
|
|
qdata = weight._qdata
|
|
params = weight._params
|
|
# Decide packing factor from the loader-set quant_format rather
|
|
# than the qdata dtype: uint8 is ambiguous (OrbitQuant INT4
|
|
# codebook indices are also uint8 but 1-per-byte, not 2-per-byte).
|
|
quant_format = getattr(module, 'quant_format', None)
|
|
if quant_format == "int4_crush":
|
|
is_int4_packed = True
|
|
elif quant_format == "int8_crush":
|
|
is_int4_packed = False
|
|
else:
|
|
is_int4_packed = qdata.dtype == torch.uint8
|
|
unpacked_in = qdata.shape[1] * (2 if is_int4_packed else 1) \
|
|
if qdata.dim() > 1 else qdata.shape[0] * (2 if is_int4_packed else 1)
|
|
if unpacked_in > orig_in:
|
|
if is_int4_packed and orig_in % 2 != 0:
|
|
log.warning(
|
|
"[INT-Crush] %s: orig_in=%d is odd; INT4 is packed "
|
|
"2-per-byte, trimming to %d (one value dropped).",
|
|
layer_key, orig_in, orig_in // 2 * 2,
|
|
)
|
|
trim_cols = orig_in // 2 if is_int4_packed else orig_in
|
|
if qdata.dim() > 1:
|
|
new_qdata = qdata[:, :trim_cols].contiguous()
|
|
else:
|
|
new_qdata = qdata[:trim_cols].contiguous()
|
|
new_params = type(params)(
|
|
scale=params.scale,
|
|
orig_dtype=params.orig_dtype,
|
|
orig_shape=(params.orig_shape[0], orig_in),
|
|
**{k: v for k, v in {
|
|
'rot_need': getattr(params, 'rot_need', None),
|
|
'rot_size': getattr(params, 'rot_size', None),
|
|
'perm': getattr(params, 'perm', None),
|
|
'zp': getattr(params, 'zp', None),
|
|
}.items() if v is not None},
|
|
)
|
|
from comfy.quant_ops import QuantizedTensor
|
|
module.weight = torch.nn.Parameter(
|
|
QuantizedTensor(new_qdata, weight._layout_type, new_params),
|
|
requires_grad=False,
|
|
)
|
|
log.info("[INT-Crush] Fixed %s in_features %d -> %d",
|
|
layer_key, unpacked_in, orig_in)
|
|
else:
|
|
padded_in = weight.shape[1]
|
|
if padded_in > orig_in:
|
|
module.weight = torch.nn.Parameter(
|
|
weight[:, :orig_in].contiguous(),
|
|
requires_grad=False,
|
|
)
|
|
log.info("[INT-Crush] Fixed %s in_features %d -> %d",
|
|
layer_key, padded_in, orig_in)
|
|
|
|
if module_path == "img_in" and hasattr(m, 'in_channels'):
|
|
ps = getattr(m, 'patch_size', 1)
|
|
m.in_channels = orig_in // (ps * ps) if ps else orig_in
|
|
log.info("[INT-Crush] INT8: Fixed model.in_channels = %d", m.in_channels)
|
|
|
|
|
|
# ── INT-Crush LoRA nodes ─────────────────────────────────────────────────────
|
|
|
|
def _on_load_apply_lora(patcher, *args):
|
|
"""ON_LOAD/ON_PRE_RUN callback: sync module adapters with the active patcher.
|
|
|
|
Every run re-enters here via load() or pre_run(). A uuid mismatch
|
|
applies the incoming patcher's stack, or clears adapters when the
|
|
patcher has no stack. An empty stack means the LoRA node left the
|
|
workflow.
|
|
"""
|
|
applied = getattr(patcher.model, '_intcrush_lora_applied_uuid', None)
|
|
if applied == patcher.patches_uuid:
|
|
return
|
|
stack = patcher.attachments.get("intcrush_lora_stack")
|
|
if stack:
|
|
apply_lora_stack(patcher, stack)
|
|
else:
|
|
clear_intcrush_lora(patcher)
|
|
patcher.model._intcrush_lora_applied_uuid = patcher.patches_uuid
|
|
|
|
|
|
def _on_load_clear_lora(patcher, *args):
|
|
"""ON_LOAD callback: clear adapters when an unloader patcher loads."""
|
|
applied = getattr(patcher.model, '_intcrush_lora_applied_uuid', None)
|
|
if applied == patcher.patches_uuid:
|
|
return
|
|
clear_intcrush_lora(patcher)
|
|
patcher.model._intcrush_lora_applied_uuid = patcher.patches_uuid
|
|
|
|
|
|
def _register_lora_sync(model: object) -> None:
|
|
"""Attach the adapter-sync callback to a loader-produced patcher so
|
|
adapter state follows patcher switches without an unloader node."""
|
|
from comfy.patcher_extension import CallbacksMP
|
|
for event in (CallbacksMP.ON_LOAD, CallbacksMP.ON_PRE_RUN):
|
|
model.remove_callbacks_with_key(event, "intcrush_lora")
|
|
model.add_callback_with_key(event, "intcrush_lora", _on_load_apply_lora)
|
|
|
|
|
|
_std_lora_bridge_installed = False
|
|
|
|
def install_std_lora_bridge() -> None:
|
|
"""Make ComfyUI's built-in LoraLoader work on INT-Crush quantized models.
|
|
|
|
ComfyUI's LoRA machinery (LowVramPatch, pinned patch buffers, VBAR
|
|
watermark accounting) merges adapters into the stored weight. INT-Crush
|
|
stores *rotated* weights, so a raw merge corrupts the adapter. This
|
|
wraps ModelPatcher.add_patches:
|
|
|
|
- Quantized layers: patches divert to the intcrush_lora_stack
|
|
attachment and apply as activation residuals. This is exact for
|
|
additive adapters and avoids dequantizing weights for a merge.
|
|
- LoRAs needing a true weight merge (DoRA, conv reshape, attention
|
|
mid) on plain group-Hadamard layers: the down-projection is
|
|
pre-rotated (rotated_adapter_for_module) and the patch passes
|
|
through as a first-class ComfyUI patch.
|
|
"""
|
|
global _std_lora_bridge_installed
|
|
import comfy.model_patcher
|
|
import comfy.utils
|
|
from comfy import weight_adapter
|
|
from comfy.patcher_extension import CallbacksMP
|
|
|
|
original_add_patches = comfy.model_patcher.ModelPatcher.add_patches
|
|
# Module-level flag guards same-process install; the marker guards
|
|
# reinstall after ComfyUI reloads this custom-node module.
|
|
if _std_lora_bridge_installed or getattr(original_add_patches, "_intcrush_bridge", False):
|
|
return
|
|
|
|
def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
|
|
rotated = 0
|
|
residual = {}
|
|
passthrough = {}
|
|
for k, v in patches.items():
|
|
op = None
|
|
if isinstance(k, str) and k.endswith(".weight") \
|
|
and isinstance(v, weight_adapter.WeightAdapterBase):
|
|
try:
|
|
op = comfy.utils.get_attr(self.model, k[: -len(".weight")])
|
|
except AttributeError:
|
|
op = None
|
|
if op is not None and getattr(op, "_intcrush_is_quantized", False):
|
|
wrapped = rotated_adapter_for_module(v, op)
|
|
if wrapped is None:
|
|
residual[k] = v
|
|
else:
|
|
if wrapped is not v:
|
|
rotated += 1
|
|
passthrough[k] = wrapped
|
|
else:
|
|
passthrough[k] = v
|
|
|
|
if residual:
|
|
stack = list(self.attachments.get("intcrush_lora_stack", []))
|
|
stack.append((residual, strength_patch))
|
|
self.attachments["intcrush_lora_stack"] = stack
|
|
self.remove_callbacks_with_key(CallbacksMP.ON_LOAD, "intcrush_lora")
|
|
self.add_callback_with_key(CallbacksMP.ON_LOAD, "intcrush_lora",
|
|
_on_load_apply_lora)
|
|
log.info("[INT-Crush LoRA] Native loader: %d layer(s) as "
|
|
"activation residuals", len(residual))
|
|
if rotated:
|
|
log.info("[INT-Crush LoRA] Native loader: %d quantized layer(s) "
|
|
"pre-rotated for weight merge", rotated)
|
|
|
|
accepted = original_add_patches(self, passthrough, strength_patch,
|
|
strength_model)
|
|
# Report diverted keys as accepted so ComfyUI doesn't warn
|
|
# "lora key not loaded".
|
|
return accepted + list(residual)
|
|
|
|
add_patches._intcrush_bridge = True
|
|
comfy.model_patcher.ModelPatcher.add_patches = add_patches
|
|
_std_lora_bridge_installed = True
|
|
|
|
|
|
install_std_lora_bridge()
|
|
|
|
|
|
class IntCrushLoRALoader:
|
|
"""Load a LoRA file and attach it as a residual on the raw (unrotated)
|
|
activation, preserving ConvRot/SmoothRot quantization and the fast
|
|
Triton INT4/INT8 GEMM kernel paths.
|
|
|
|
Standard ComfyUI LoRA patching adds dW to the *rotated* weight matrix,
|
|
which corrupts the LoRA contribution. This node instead stores the LoRA
|
|
as separate low-rank buffers (down-projection A, up-projection B) and
|
|
applies ``out += (x_raw @ A^T) @ B^T * scale`` in the forward pass,
|
|
*before* any rotation or smoothing. This is mathematically equivalent
|
|
to the true unrotated LoRA and keeps the main weight as a
|
|
QuantizedTensor so the Triton kernels stay active.
|
|
|
|
The node clones the ModelPatcher and stores the parsed adapter mapping
|
|
as a patcher attachment. An ON_LOAD callback reapplies the adapters
|
|
whenever ComfyUI loads the model, so the correct adapters are active
|
|
even when the execution cache returns a stale patcher whose shared
|
|
model was mutated by a later run.
|
|
|
|
Usage: connect MODEL output from an INT4/INT8 loader into this node,
|
|
then connect the output to KSampler (or wherever the model is used).
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"lora_name": (_get_lora_list(), {"tooltip": "The name of the LoRA."}),
|
|
"strength": ("FLOAT", {
|
|
"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01,
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "loaders/INT-Crush"
|
|
|
|
def load(self, model: object, lora_name: str, strength: float) -> tuple[object]:
|
|
import folder_paths
|
|
import comfy.utils
|
|
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
lora_sd, _metadata = comfy.utils.load_torch_file(
|
|
lora_path, safe_load=True, return_metadata=True)
|
|
|
|
if not lora_sd:
|
|
raise ValueError(f"INT-Crush LoRA: empty file '{lora_name}'")
|
|
|
|
# Parse adapters WITHOUT mutating the model.
|
|
adapters = parse_lora_adapters(model, lora_sd)
|
|
if not adapters:
|
|
log.warning("[INT-Crush LoRA] No modules matched in '%s'", lora_name)
|
|
|
|
# Clone so each LoRA node output is a distinct patcher that
|
|
# ComfyUI's cache can track independently.
|
|
n = model.clone()
|
|
|
|
# Store the adapter mapping on the clone as a stack entry.
|
|
# clone() copies attachments by reference, so copy the list
|
|
# before appending to avoid mutating the parent's attachment.
|
|
stack = list(n.attachments.get("intcrush_lora_stack", []))
|
|
stack.append((adapters, strength))
|
|
n.attachments["intcrush_lora_stack"] = stack
|
|
|
|
# Bump patches_uuid so ComfyUI's VRAM management knows the model
|
|
# state changed and triggers a reload -> ON_LOAD fires.
|
|
n.patches_uuid = uuid.uuid4()
|
|
|
|
# Re-register the sync callback. Remove first to avoid duplicates
|
|
# from clone().
|
|
_register_lora_sync(n)
|
|
|
|
# Apply immediately so the current execution has the right state
|
|
# even if ComfyUI doesn't reload the model (already loaded, not dirty).
|
|
apply_lora_stack(n, stack)
|
|
n.model._intcrush_lora_applied_uuid = n.patches_uuid
|
|
|
|
if adapters:
|
|
log.info("[INT-Crush LoRA] Attached to %d layer(s)", len(adapters))
|
|
return (n,)
|
|
|
|
|
|
class IntCrushLoRAUnloader:
|
|
"""Remove INT-Crush LoRA residual buffers from a model.
|
|
|
|
Usually optional: the model loader's sync callback clears adapters
|
|
automatically when the LoRA node leaves the workflow. Use this node
|
|
for explicit mid-workflow detaches.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"model": ("MODEL",)}}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "unload"
|
|
CATEGORY = "loaders/INT-Crush"
|
|
|
|
def unload(self, model: object) -> tuple[object]:
|
|
from comfy.patcher_extension import CallbacksMP
|
|
|
|
n = model.clone()
|
|
# Remove any inherited adapter attachment.
|
|
n.attachments.pop("intcrush_lora_stack", None)
|
|
n.patches_uuid = uuid.uuid4()
|
|
|
|
# Register ON_LOAD that clears adapters when model reloads.
|
|
n.remove_callbacks_with_key(CallbacksMP.ON_LOAD, "intcrush_lora")
|
|
n.add_callback_with_key(CallbacksMP.ON_LOAD, "intcrush_lora_clear",
|
|
_on_load_clear_lora)
|
|
|
|
# Clear immediately for the current execution.
|
|
clear_intcrush_lora(n)
|
|
n.model._intcrush_lora_applied_uuid = n.patches_uuid
|
|
|
|
log.info("[INT-Crush LoRA] Cleared LoRA adapters")
|
|
return (n,)
|
|
|
|
|
|
# ── Node registration ────────────────────────────────────────────────────────
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SimpleINT4UNetLoader": SimpleINT4UNetLoader,
|
|
"SimpleINT8UNetLoader": SimpleINT8UNetLoader,
|
|
"OrbitQuantLoader": OrbitQuantLoader,
|
|
"IntCrushLoRALoader": IntCrushLoRALoader,
|
|
"IntCrushLoRAUnloader": IntCrushLoRAUnloader,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SimpleINT4UNetLoader": "INT4 UNet Loader (INT-Crush)",
|
|
"SimpleINT8UNetLoader": "INT8 UNet Loader (INT-Crush)",
|
|
"OrbitQuantLoader": "OrbitQuant UNet Loader (INT-Crush)",
|
|
"IntCrushLoRALoader": "LoRA Loader (INT-Crush)",
|
|
"IntCrushLoRAUnloader": "LoRA Unloader (INT-Crush)",
|
|
}
|
|
|
|
__all__ = [
|
|
"SimpleINT4UNetLoader",
|
|
"SimpleINT8UNetLoader",
|
|
"OrbitQuantLoader",
|
|
"IntCrushLoRALoader",
|
|
"IntCrushLoRAUnloader",
|
|
"NODE_CLASS_MAPPINGS",
|
|
"NODE_DISPLAY_NAME_MAPPINGS",
|
|
]
|