212 lines
7.6 KiB
Python
212 lines
7.6 KiB
Python
"""INT-Crush adapter residual management.
|
|
|
|
Provides helpers to attach/detach weight adapters (LoRA, LoKr, LoHa,
|
|
OFT, BOFT, GLoRA) as residual buffers on IntCrushOps.Linear modules.
|
|
The adapter's ``h(x, out)`` computes the additive delta and ``g(out)``
|
|
applies any output transformation (identity for most; orthogonal
|
|
rotation for OFT/BOFT).
|
|
|
|
All six adapter types supported by ComfyUI's weight_adapter system
|
|
work out of the box — the adapter's own math is used directly.
|
|
"""
|
|
|
|
import logging
|
|
import torch
|
|
from typing import Any
|
|
|
|
from ._quant_utils import rotate_activations
|
|
|
|
__all__ = [
|
|
"clear_intcrush_lora",
|
|
"attach_lora_as_buffers",
|
|
"parse_lora_adapters",
|
|
"apply_lora_adapters",
|
|
"apply_lora_stack",
|
|
"rotated_adapter_for_module",
|
|
]
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
|
|
def _offload_adapter_weights(adapter: Any) -> None:
|
|
"""Move adapter tensors to CPU, preserving dtype.
|
|
|
|
Detached adapters can stay referenced by ComfyUI's output cache, so
|
|
dropping the module reference alone does not free their GPU memory.
|
|
"""
|
|
weights = getattr(adapter, "weights", None)
|
|
if not isinstance(weights, (list, tuple)):
|
|
return
|
|
moved = []
|
|
changed = False
|
|
for w in weights:
|
|
if isinstance(w, torch.Tensor) and w.device.type != "cpu":
|
|
w = w.to("cpu")
|
|
changed = True
|
|
moved.append(w)
|
|
if changed:
|
|
adapter.weights = type(weights)(moved)
|
|
|
|
|
|
def clear_intcrush_lora(model: Any) -> int:
|
|
"""Remove adapter residuals from all IntCrushOps.Linear modules."""
|
|
count = 0
|
|
for module in model.model.modules():
|
|
if hasattr(module, '_intcrush_adapters'):
|
|
for entry in module._intcrush_adapters:
|
|
_offload_adapter_weights(entry["adapter"])
|
|
module._intcrush_adapters = []
|
|
# Invalidate any cached CUDA graph — it may have been captured
|
|
# with different layer attributes (smooth/perm/device state).
|
|
# The next forward pass will rebuild it from the current state.
|
|
module._graph_runner = None
|
|
count += 1
|
|
return count
|
|
|
|
|
|
def _new_adapter_entry(adapter: Any, strength: float) -> dict:
|
|
"""Build a per-module adapter stack entry with its runtime state."""
|
|
return {
|
|
"adapter": adapter,
|
|
"strength": strength,
|
|
"ready": False,
|
|
"device": None,
|
|
"dtype": None,
|
|
"transient": False,
|
|
"lokr_cache": None,
|
|
}
|
|
|
|
|
|
def _build_weight_key_to_module(model: Any) -> dict:
|
|
"""Map '{module_name}.weight' -> module for all IntCrushOps.Linear modules."""
|
|
result = {}
|
|
for name, module in model.model.named_modules():
|
|
if hasattr(module, '_intcrush_adapters'):
|
|
result[f"{name}.weight"] = module
|
|
return result
|
|
|
|
|
|
def parse_lora_adapters(model: Any, lora_sd: dict) -> dict:
|
|
"""Parse a LoRA state-dict into a ``{weight_key: adapter}`` mapping.
|
|
|
|
Does NOT mutate the model. The returned mapping can be stored on a
|
|
ModelPatcher attachment and applied later via apply_lora_adapters /
|
|
apply_lora_stack.
|
|
"""
|
|
import comfy.lora
|
|
import comfy.lora_convert
|
|
import comfy.weight_adapter
|
|
|
|
lora_sd = comfy.lora_convert.convert_lora(lora_sd)
|
|
key_map = comfy.lora.model_lora_keys_unet(model.model, {})
|
|
patches = comfy.lora.load_lora(lora_sd, key_map)
|
|
if not patches:
|
|
return {}
|
|
|
|
weight_key_to_module = _build_weight_key_to_module(model)
|
|
adapters: dict = {}
|
|
for weight_key, patch in patches.items():
|
|
if weight_key not in weight_key_to_module:
|
|
continue
|
|
if not isinstance(patch, comfy.weight_adapter.WeightAdapterBase):
|
|
continue
|
|
adapters[weight_key] = patch
|
|
return adapters
|
|
|
|
|
|
def apply_lora_adapters(model: Any, adapters: dict, strength: float) -> int:
|
|
"""Attach parsed adapters to the model's IntCrushOps.Linear modules.
|
|
|
|
Clears any existing adapters first. Returns the number of modules
|
|
that received an adapter.
|
|
"""
|
|
clear_intcrush_lora(model)
|
|
if not adapters:
|
|
return 0
|
|
weight_key_to_module = _build_weight_key_to_module(model)
|
|
attached = 0
|
|
for weight_key, patch in adapters.items():
|
|
module = weight_key_to_module.get(weight_key)
|
|
if module is None:
|
|
continue
|
|
module._intcrush_adapters.append(_new_adapter_entry(patch, strength))
|
|
attached += 1
|
|
return attached
|
|
|
|
|
|
def apply_lora_stack(model: Any, stack: list) -> int:
|
|
"""Apply a stack of ``(adapters, strength)`` pairs to the model.
|
|
|
|
Clears any existing adapters first, then applies each entry in
|
|
order. Entries targeting the same module accumulate; each one
|
|
contributes its own scaled residual in the forward pass.
|
|
"""
|
|
clear_intcrush_lora(model)
|
|
if not stack:
|
|
return 0
|
|
weight_key_to_module = _build_weight_key_to_module(model)
|
|
attached = 0
|
|
for adapters, strength in stack:
|
|
for weight_key, patch in adapters.items():
|
|
module = weight_key_to_module.get(weight_key)
|
|
if module is None:
|
|
continue
|
|
module._intcrush_adapters.append(_new_adapter_entry(patch, strength))
|
|
attached += 1
|
|
return attached
|
|
|
|
|
|
def rotated_adapter_for_module(adapter: Any, module: Any) -> Any | None:
|
|
"""Rebuild a LoRA adapter in the module's Hadamard-rotated weight space.
|
|
|
|
INT-Crush stores ``W·H`` for rotated layers, so a standard LoRA merge
|
|
of ``W·H + s·BA`` corrupts the adapter. Pre-rotating the down-projection
|
|
is exact because Hadamard is orthogonal:
|
|
``(x·H)(W·H + s·B·(A·H))ᵀ = x·Wᵀ + s·x·(BA)ᵀ``.
|
|
|
|
Returns the adapter unchanged when the module is not quantized, or
|
|
when the LoRA needs a true weight merge on a merge-compatible layer.
|
|
Returns None otherwise, diverting the patch to the
|
|
activation-residual path. Residual application is exact for plain
|
|
additive LoRA and every non-LoRA adapter type, and avoids
|
|
dequantizing the weight for a merge.
|
|
"""
|
|
if not getattr(module, "_intcrush_is_quantized", False):
|
|
return adapter
|
|
if adapter.name != "lora":
|
|
# Non-LoRA adapters (LoKr, LoHa, OFT, ...) can't be pre-rotated.
|
|
# Divert to the activation-residual path instead of letting
|
|
# ComfyUI dequantize the full weight for a merge.
|
|
return None
|
|
up, down, alpha, mid, dora_scale, reshape = adapter.weights
|
|
if mid is None and dora_scale is None and reshape is None:
|
|
# Plain additive LoRA: residual path is exact on raw activations
|
|
# and avoids dequantizing the weight for a merge.
|
|
return None
|
|
# DoRA, conv reshape, or attention mid: only a weight merge is correct.
|
|
if (module._intcrush_perm is not None or module._intcrush_is_rpbh
|
|
or module._intcrush_smooth is not None
|
|
or module._intcrush_smoothrot_factors is not None):
|
|
return None
|
|
rot_size = module._intcrush_rot_size
|
|
if not module._intcrush_rot_need or rot_size <= 0:
|
|
return adapter
|
|
# fp32 like the converter; slice group padding back to stored width.
|
|
rot = rotate_activations(down.float(), rot_size)
|
|
if rot.shape[-1] != down.shape[-1]:
|
|
rot = rot[..., : down.shape[-1]]
|
|
rot = rot.to(down.dtype)
|
|
return type(adapter)(adapter.loaded_keys, (up, rot, alpha, mid, dora_scale, reshape))
|
|
|
|
|
|
def attach_lora_as_buffers(model: Any, lora_sd: dict, strength: float) -> tuple[int, int]:
|
|
"""Parse a LoRA state-dict and attach adapter objects as residual buffers.
|
|
|
|
Returns (attached_count, total_patches).
|
|
"""
|
|
adapters = parse_lora_adapters(model, lora_sd)
|
|
if not adapters:
|
|
return 0, 0
|
|
attached = apply_lora_adapters(model, adapters, strength)
|
|
return attached, len(adapters)
|