doggeddalle
28ac439a26
Add a ComfyUI plugin implementing a stacked multi-LoRA loader for MiniMax H3. Introduces H3PowerLoraStack and H3LoraInspector node classes, plus modules: adaln (adaLN dense↔curve basis conversion), branch (runtime low-rank branches for quantized weights), apply (stack orchestration), detect (safetensors header introspection), keymap (key normalization), nodes, and a web UI (h3_power_lora_stack.js). Includes pyproject metadata and README. Preserves quantized kernels via branching, rebases adaLN pairs, handles many naming conventions, and provides a browser picker and per-LoRA reporting.
198 lines
7.3 KiB
Python
198 lines
7.3 KiB
Python
"""Stack orchestration: normalise, port, partition, apply."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from collections import OrderedDict
|
|
|
|
import comfy.lora
|
|
import comfy.utils
|
|
from comfy.quant_ops import QuantizedTensor
|
|
from comfy.weight_adapter import LoRAAdapter
|
|
|
|
from . import adaln as adaln_mod
|
|
from . import branch as branch_mod
|
|
from . import keymap
|
|
|
|
LOG = logging.getLogger("h3.powerlorastack")
|
|
|
|
# Layers whose forward is bypassed by a fused kernel, per weight layout. A
|
|
# runtime branch on these would be silently dropped, so they must merge.
|
|
_FORWARD_BYPASSED = {"TensorWiseINT8Layout": ("mlp.fc2",)}
|
|
|
|
|
|
class StackReport:
|
|
"""Human-readable account of what happened, surfaced on the node."""
|
|
|
|
def __init__(self):
|
|
self.lines: list[str] = []
|
|
self.merged = 0
|
|
self.branched = 0
|
|
self.skipped = 0
|
|
self.bank_bytes = 0
|
|
|
|
def add(self, line: str):
|
|
self.lines.append(line)
|
|
|
|
def text(self) -> str:
|
|
return "\n".join(self.lines)
|
|
|
|
|
|
def _module_weight(model_patcher, module_path: str):
|
|
try:
|
|
module = model_patcher.get_model_object(module_path)
|
|
except Exception:
|
|
return None
|
|
return getattr(module, "weight", None)
|
|
|
|
|
|
def _is_plain_lora(adapter) -> bool:
|
|
"""True for a bare rank-decomposition with no locon/dora/reshape extras."""
|
|
return (
|
|
isinstance(adapter, LoRAAdapter)
|
|
and adapter.weights[3] is None # locon mid
|
|
and adapter.weights[4] is None # dora scale
|
|
and adapter.weights[5] is None # reshape / pad
|
|
and adapter.weights[0].ndim == 2
|
|
)
|
|
|
|
|
|
def _branchable(weight, module_path: str, mode: str) -> bool:
|
|
if weight is None or mode == "merge":
|
|
return False
|
|
if mode == "auto" and not isinstance(weight, QuantizedTensor):
|
|
return False # unquantized merges are exact and free at runtime
|
|
layout = getattr(weight, "_layout_cls", None)
|
|
for suffix in _FORWARD_BYPASSED.get(layout, ()):
|
|
if module_path.endswith(suffix):
|
|
return False
|
|
return True
|
|
|
|
|
|
def detect_quantization(model_patcher) -> str:
|
|
"""Report the DiT's dominant weight format for the node's status line."""
|
|
layouts: dict[str, int] = {}
|
|
plain = 0
|
|
try:
|
|
diffusion_model = model_patcher.get_model_object("diffusion_model")
|
|
except Exception:
|
|
return "unknown"
|
|
for module in diffusion_model.modules():
|
|
weight = getattr(module, "weight", None)
|
|
if weight is None:
|
|
continue
|
|
if isinstance(weight, QuantizedTensor):
|
|
layout = getattr(weight, "_layout_cls", "quantized")
|
|
layouts[layout] = layouts.get(layout, 0) + 1
|
|
else:
|
|
plain += 1
|
|
if not layouts:
|
|
return "unquantized"
|
|
parts = [f"{k.replace('TensorCore', '').replace('TensorWise', '').replace('Layout', '')}"
|
|
f" x{v}" for k, v in sorted(layouts.items(), key=lambda kv: -kv[1])]
|
|
if plain:
|
|
parts.append(f"plain x{plain}")
|
|
return ", ".join(parts)
|
|
|
|
|
|
def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="",
|
|
report: StackReport | None = None):
|
|
"""Apply a list of ``{'path', 'name', 'strength'}`` LoRAs to an H3 model.
|
|
|
|
``mode`` is ``auto`` (branch quantized layers, merge the rest), ``merge``
|
|
(stock behaviour) or ``branch`` (never touch a weight).
|
|
"""
|
|
report = report or StackReport()
|
|
patcher = model.clone()
|
|
|
|
model_sd = patcher.model.state_dict()
|
|
index = keymap.build_module_index(model_sd.keys())
|
|
key_map = {k[: -len(".weight")]: k for k in model_sd
|
|
if k.startswith("diffusion_model.") and k.endswith(".weight")}
|
|
|
|
try:
|
|
diffusion_model = patcher.get_model_object("diffusion_model")
|
|
except Exception as exc:
|
|
raise ValueError("H3 Power LoRA Stack requires a MiniMax H3 model") from exc
|
|
target_dim, table = adaln_mod.read_target(diffusion_model)
|
|
if adaln_mode == "off":
|
|
adaln_ctx = None
|
|
else:
|
|
if not grid_path:
|
|
grid_path = adaln_mod.find_silu_grid()
|
|
adaln_ctx = adaln_mod.AdalnContext(target_dim, table, grid_path)
|
|
|
|
report.add(f"base: {detect_quantization(patcher)}")
|
|
report.add(f"adaLN: {'curve' if table is not None else 'dense'} (input dim {target_dim})")
|
|
|
|
per_module: "OrderedDict[str, list]" = OrderedDict()
|
|
compute_dtype = model.model_dtype()
|
|
|
|
for entry in entries:
|
|
name = entry.get("name") or entry["path"]
|
|
strength = float(entry.get("strength", 1.0))
|
|
lora_sd = comfy.utils.load_torch_file(entry["path"], safe_load=True)
|
|
|
|
normalized, unmatched = keymap.normalize(lora_sd, index)
|
|
adaln_note = ""
|
|
if adaln_ctx is not None and target_dim:
|
|
source_table = normalized.pop("adaln_t_table", None)
|
|
normalized, stats = adaln_mod.port_adaln_pairs(
|
|
normalized, adaln_ctx, source_table=source_table)
|
|
if stats["ported"]:
|
|
adaln_note = f", adaLN ported x{stats['ported']}"
|
|
if stats["rebased"]:
|
|
adaln_note += f", adaLN rebased x{stats['rebased']}"
|
|
if stats["skipped"]:
|
|
adaln_note += f", adaLN dropped x{stats['skipped']}"
|
|
if stats["residual"] is not None:
|
|
adaln_note += f" (basis fit {stats['residual']:.1e})"
|
|
|
|
loaded = comfy.lora.load_lora(normalized, key_map)
|
|
|
|
merge: dict = {}
|
|
branched_here = 0
|
|
for weight_key, adapter in loaded.items():
|
|
module_path = weight_key[: -len(".weight")] if weight_key.endswith(".weight") else None
|
|
if module_path is None or not _is_plain_lora(adapter):
|
|
merge[weight_key] = adapter
|
|
continue
|
|
weight = _module_weight(patcher, module_path)
|
|
if not _branchable(weight, module_path, mode):
|
|
merge[weight_key] = adapter
|
|
continue
|
|
up, down = adapter.weights[0], adapter.weights[1]
|
|
if up.shape[0] != weight.shape[0] or down.shape[1] != weight.shape[1]:
|
|
LOG.warning("H3 PowerLoraStack: shape mismatch on %s, skipped", weight_key)
|
|
report.skipped += 1
|
|
continue
|
|
alpha = adapter.weights[2]
|
|
scale = strength * (float(alpha) / down.shape[0] if alpha is not None else 1.0)
|
|
per_module.setdefault(module_path, []).append((up, down, scale))
|
|
branched_here += 1
|
|
|
|
if merge:
|
|
patcher.add_patches(merge, strength)
|
|
report.merged += len(merge)
|
|
report.branched += branched_here
|
|
|
|
detail = f"{len(merge)} merged, {branched_here} branched"
|
|
if unmatched:
|
|
detail += f", {len(unmatched)} unmatched"
|
|
report.add(f"{name} @ {strength:g}: {detail}{adaln_note}")
|
|
if not merge and not branched_here:
|
|
report.add(f" ! {name} matched no layers on this model")
|
|
|
|
if per_module:
|
|
report.bank_bytes = branch_mod.bank_bytes(per_module, compute_dtype)
|
|
# a chain of stack nodes must not overwrite each other's bank
|
|
n = 0
|
|
while patcher.get_additional_models_with_key(f"h3_power_lora_bank_{n}"):
|
|
n += 1
|
|
tag = f"h3_power_lora_bank_{n}"
|
|
branch_mod.attach(patcher, per_module, compute_dtype, tag)
|
|
report.add(f"branch bank: {len(per_module)} layers, "
|
|
f"{report.bank_bytes / (1024 ** 2):.0f} MB")
|
|
|
|
return patcher, report
|