commit 28ac439a26f69d053f6319922dd0255a65e4ed7d Author: doggeddalle <153059103+cicalooo@users.noreply.github.com> Date: Sat Aug 8 19:27:06 2026 +1000 Add MiniMax H3 Power LoRA Stack 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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..928a9d9 --- /dev/null +++ b/README.md @@ -0,0 +1,139 @@ +# ComfyUI-H3-PowerLoraStack + +A stacked multi-LoRA loader for **MiniMax H3**, in the spirit of rgthree's Power +Lora Loader but built around the three things that actually break H3 LoRAs. + +## Nodes + +| Node | Purpose | +|---|---| +| **MiniMax H3 Power LoRA Stack** | Any number of LoRAs on one node, each with a toggle and strength | +| **MiniMax H3 LoRA Inspector** | Reports a LoRA's format, rank and adaLN basis without loading it | + +## Why not just use a normal LoRA loader + +### 1. Quantized weights: merging destroys the LoRA + +ComfyUI's stock path is dequantize → add delta → `requantize_from_float(scale="recalculate")`. +That round trip is **not idempotent**: re-fitting the codebook and re-rounding to +int4 injects ~1.5% relative weight noise, while a typical H3 LoRA delta is +0.01–0.08% of the weight. The merge therefore replaces the adapter with noise. +Measured on a w4a8 checkpoint, the stock merge recovers under a third of the +LoRA magnitude at cos 0.12–0.14 against the correct result, and takes 128 s +versus 6.5 s. + +This node routes quantized layers through an exact runtime low-rank branch +(`y = W_q(x) + B @ A @ x`) instead, keeping the quantized kernel and costing +~1.5% extra FLOPs at rank 64. + +`quantized_layers` controls this: + +- `auto` (default) — branch quantized layers, merge unquantized ones +- `branch` — never modify a weight, even in bf16 (fast strength A/B testing) +- `merge` — stock behaviour; only useful for comparison + +One layer is special-cased: `mlp.fc2` under `TensorWiseINT8Layout` is reached +through `comfy.ops.linear_input_act`, which fuses the activation into the INT8 +kernel and never calls `fc2.forward`. A branch there would be silently dropped, +so those layers merge even in `branch` mode. + +### 2. adaLN basis mismatch + +H3 ships in two forms. A *dense* checkpoint feeds `silu(time_embedder(t))`, a +2688-dim vector, into every `adaln_proj.linear`. A *curve* (pruned) checkpoint +drops the time embedder and stores an `adaln_t_table` of shape `[grid, 8]`. + +A LoRA trained against one form has the wrong `lora_A` width for the other, so +ComfyUI logs `shape '[96768, 8]' is invalid for input of size 260112384` and +skips the layer. **Dropping those pairs is not an acceptable fix** — on the +turbo distillation LoRAs the constant term alone is ~100% of the magnitude of +`dW @ S(t)`, so it discards essentially the whole adapter. + +This node changes basis instead, which preserves rank: + +``` +dense -> curve A' = A @ V, bias delta B @ (A @ c) +curve -> dense A' = A @ pinv(V), bias delta -B @ (A' @ c) +``` + +where `S(t) = c + V @ table(t)`, recovered by least squares of `[1 | table]` +against the silu grid. The bias delta is emitted as `.diff_b`, which comfy +applies as a `("diff",)` patch on the sibling `.bias`; without it the port is +nearly worthless. Measured end to end, the ported adapter reproduces the dense +contribution at **cos 0.999998**. + +The fit uses the *target checkpoint's own* table, because bakes differ — locally, +`PinkCherry` ships an uncentered basis (column norms 22.98, 2.67, 1.66, …) while +`10Eros_Max` ships comfy's mean-centered one (7.08, 2.09, 0.75, …). + +**This also catches a silent failure the stock loader cannot see.** Two curve +bakes have the same adaLN width, so a LoRA trained against one loads without +complaint on the other and is simply wrong — measured at **cos −0.375**, worse +than not applying it at all. When a LoRA ships its own `adaln_t_table`, this +node rebases table-to-table (exactly, cos 1.000000, no grid needed). When it +does not, the width matches and the mismatch is undetectable — so prefer LoRAs +trained against the checkpoint you are running. + +#### The silu grid + +Dense↔curve porting needs `h3_silu_temb_grid.safetensors`. It is searched for in +`models/h3_adaln/`, `models/loras/`, `models/diffusion_models/` and one level +under `custom_nodes/`. Table-to-table rebasing does not need it. Set +`adaln_port` to `off` to disable porting entirely. + +With a grid from a different build the fit bottoms out at ~1.7e-3 relative, +because the 7th and 8th curve directions are near-degenerate (σ₇ ≈ σ₈) and so +differ between bakes. That is still ~6× below the int8 quantization floor. + +### 3. Key conventions + +Every H3 LoRA naming convention resolves against the model's own key set rather +than by guessing where underscores split, so `qkv_proj` is never mistaken for +two tokens: + +| Convention | Example | +|---|---| +| ai-toolkit / diffusers | `diffusion_model.blocks.0.attn.qkv_proj.lora_A.weight` | +| bare (no prefix) | `blocks.0.attn.qkv_proj.lora_A.weight` | +| kohya / musubi | `lora_unet_blocks_0_attn_qkv_proj.lora_down.weight` + `.alpha` | +| lycoris | `lycoris_blocks_0_...` | +| peft / diffusers trainer | `base_model.model.blocks.0...`, `transformer.blocks.0...` | + +Verified against all 37 H3 LoRAs in `models/loras/h3`: **zero unmatched keys.** + +## Stacking + +Multiple LoRAs on the same layer are fused into a single pair by concatenating +along the rank axis: + +``` +sum_i s_i * B_i @ A_i @ x == [s_1 B_1 | ... | s_N B_N] @ [A_1; ...; A_N] @ x +``` + +so a ten-LoRA stack costs one extra matmul pair per layer, not ten. The factors +live in a `_LoraBank` registered via `set_additional_models`, so its VRAM is +accounted for and comfy's weakref bookkeeping stays quiet. + +## Output + +The `report` string output accounts for every LoRA (wire it to a preview/show-text +node to read it; it is also written to the console log): + +``` +base: ConvRotW4A4 x300, INT8 x50 +adaLN: curve (input dim 8) +turbo dense-adaLN @ 1: 26 merged, 275 branched, adaLN ported x51 (basis fit 1.7e-03) +GalaxyAce curve-adaLN @ 0.8: 0 merged, 258 branched +branch bank: 300 layers, 412 MB +``` + +## Limitations + +- A curve-trained LoRA on a **dense** checkpoint can only be ported if the LoRA + carries its own `adaln_t_table`; otherwise there is no record of which bake it + was trained against and the adaLN pairs are dropped with a warning. +- Runtime branches apply to `MODEL` from ComfyUI's native H3 loader. The + streaming loader in `minimaxh3chinkloader` uses its own `MINIMAX_H3_MODEL` + handle and its own LoRA path. +- DoRA, LoHa, LoKr and locon adapters always merge — only plain rank + decompositions can be branched. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ccdf11c --- /dev/null +++ b/__init__.py @@ -0,0 +1,7 @@ +"""ComfyUI-H3-PowerLoraStack: stacked multi-LoRA loading for MiniMax H3.""" + +from .h3lora.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +WEB_DIRECTORY = "./web" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/__pycache__/__init__.cpython-312.pyc b/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..ee877b2 Binary files /dev/null and b/__pycache__/__init__.cpython-312.pyc differ diff --git a/h3lora/__init__.py b/h3lora/__init__.py new file mode 100644 index 0000000..e421f41 --- /dev/null +++ b/h3lora/__init__.py @@ -0,0 +1 @@ +"""MiniMax H3 Power LoRA Stack.""" diff --git a/h3lora/__pycache__/__init__.cpython-312.pyc b/h3lora/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..d5be8bc Binary files /dev/null and b/h3lora/__pycache__/__init__.cpython-312.pyc differ diff --git a/h3lora/__pycache__/adaln.cpython-312.pyc b/h3lora/__pycache__/adaln.cpython-312.pyc new file mode 100644 index 0000000..66993d1 Binary files /dev/null and b/h3lora/__pycache__/adaln.cpython-312.pyc differ diff --git a/h3lora/__pycache__/apply.cpython-312.pyc b/h3lora/__pycache__/apply.cpython-312.pyc new file mode 100644 index 0000000..81d4493 Binary files /dev/null and b/h3lora/__pycache__/apply.cpython-312.pyc differ diff --git a/h3lora/__pycache__/branch.cpython-312.pyc b/h3lora/__pycache__/branch.cpython-312.pyc new file mode 100644 index 0000000..b88bcc8 Binary files /dev/null and b/h3lora/__pycache__/branch.cpython-312.pyc differ diff --git a/h3lora/__pycache__/detect.cpython-312.pyc b/h3lora/__pycache__/detect.cpython-312.pyc new file mode 100644 index 0000000..98ada7f Binary files /dev/null and b/h3lora/__pycache__/detect.cpython-312.pyc differ diff --git a/h3lora/__pycache__/keymap.cpython-312.pyc b/h3lora/__pycache__/keymap.cpython-312.pyc new file mode 100644 index 0000000..b7f7d45 Binary files /dev/null and b/h3lora/__pycache__/keymap.cpython-312.pyc differ diff --git a/h3lora/__pycache__/nodes.cpython-312.pyc b/h3lora/__pycache__/nodes.cpython-312.pyc new file mode 100644 index 0000000..9344738 Binary files /dev/null and b/h3lora/__pycache__/nodes.cpython-312.pyc differ diff --git a/h3lora/adaln.py b/h3lora/adaln.py new file mode 100644 index 0000000..4aa4f3a --- /dev/null +++ b/h3lora/adaln.py @@ -0,0 +1,337 @@ +"""AdaLN basis conversion between dense and curve MiniMax H3 checkpoints. + +H3 ships in two forms. A *dense* checkpoint feeds ``silu(time_embedder(t))``, +a 2688-dim vector, into every ``adaln_proj.linear``. A *curve* (pruned) +checkpoint drops the time embedder and instead stores a precomputed +``adaln_t_table`` of shape ``[grid, k]`` (k = 8 locally); the linears then take +a k-dim input. ComfyUI's tables are mean-centered, i.e. the two spaces are +related by + + S(t) = c + V @ table(t) V: [2688, k], c: [2688] + +A LoRA trained on one form has an ``adaln_proj.linear.lora_A`` of the wrong +width for the other, which is why ComfyUI logs + + shape '[96768, 8]' is invalid for input of size 260112384 + +and silently skips the layer. Stripping those pairs is not an acceptable fix: +on the turbo distillation LoRAs the constant term alone is ~100% of the +magnitude of ``dW @ S(t)``, so dropping adaLN discards essentially the whole +adapter. Instead we change basis, which preserves rank: + + dense -> curve A' = A @ V, plus bias delta B @ (A @ c) + curve -> dense A' = A @ pinv(V), plus bias delta -B @ (A' @ c) + +The reverse direction needs the pseudo-inverse rather than the transpose: the +least-squares ``V`` spans the right subspace but its columns are not +orthonormal, so ``V.T`` would not satisfy ``M @ V == I``. + +The bias delta is emitted as ``.diff_b``; ``comfy.lora.load_lora`` turns that +into a ``("diff",)`` patch on the sibling ``.bias``. Without it the port is +nearly worthless. + +``V`` and ``c`` are recovered by least squares of ``[1 | table]`` against the +silu grid. The fit must use the *target model's own* table: comfy's tables are +mean-centered while some third-party bakes are not, so tables from different +checkpoints are not interchangeable even though each is self-consistent with +its own weights. +""" + +from __future__ import annotations + +import logging +import os + +import torch + +LOG = logging.getLogger(__name__) + +GRID_FILENAME = "h3_silu_temb_grid.safetensors" +_GRID_TENSOR = "silu_t_emb_grid" + +_grid_cache: dict[str, torch.Tensor] = {} +_basis_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {} + + +def find_silu_grid(extra_dirs=()) -> str: + """Locate ``h3_silu_temb_grid.safetensors`` in the usual places.""" + candidates = list(extra_dirs) + try: + import folder_paths + candidates.append(os.path.join(folder_paths.models_dir, "h3_adaln")) + candidates.append(folder_paths.get_folder_paths("loras")[0]) + candidates.append(os.path.join(folder_paths.models_dir, "diffusion_models")) + candidates.append(folder_paths.get_folder_paths("custom_nodes")[0] + if "custom_nodes" in getattr(folder_paths, "folder_names_and_paths", {}) + else os.path.join(os.path.dirname(folder_paths.models_dir), "custom_nodes")) + except Exception: + pass + for base in candidates: + if not base or not os.path.isdir(base): + continue + direct = os.path.join(base, GRID_FILENAME) + if os.path.isfile(direct): + return direct + # one level down covers custom_nodes//h3_silu_temb_grid.safetensors + try: + for entry in os.listdir(base): + nested = os.path.join(base, entry, GRID_FILENAME) + if os.path.isfile(nested): + return nested + except OSError: + continue + return "" + + +def load_silu_grid(path: str) -> torch.Tensor: + """Load the ``[grid, 2688]`` silu(t_emb) grid, cached by path.""" + cached = _grid_cache.get(path) + if cached is not None: + return cached + from safetensors.torch import load_file + sd = load_file(path) + if _GRID_TENSOR not in sd: + raise ValueError(f"{path} does not contain '{_GRID_TENSOR}'") + grid = sd[_GRID_TENSOR].to(torch.float32) + _grid_cache[path] = grid + return grid + + +def fit_basis(table: torch.Tensor, grid: torch.Tensor): + """Least-squares fit of ``grid ~= 1 c^T + table V^T``. + + ``table`` is the target model's ``adaln_t_table`` ``[G, k]``; ``grid`` is + ``[G, 2688]``. Returns ``(V [2688, k], c [2688])``. + """ + if table.shape[0] != grid.shape[0]: + raise ValueError( + f"adaLN table grid ({table.shape[0]}) and silu grid ({grid.shape[0]}) " + "have different resolutions; they must come from the same bake" + ) + key = (id(table), table.shape, float(table.float().sum()), id(grid)) + cached = _basis_cache.get(key) + if cached is not None: + return cached + + t = table.to(torch.float64) + s = grid.to(torch.float64) + design = torch.cat([torch.ones(t.shape[0], 1, dtype=torch.float64), t], dim=1) # [G, 1+k] + solution = torch.linalg.lstsq(design, s).solution # [1+k, 2688] + c = solution[0].contiguous() # [2688] + v = solution[1:].transpose(0, 1).contiguous() # [2688, k] + + residual = float((design @ solution - s).norm() / s.norm().clamp(min=1e-12)) + # The last two curve directions are near-degenerate (sigma_7 ~ sigma_8), so + # they differ between bakes and a grid from a different build only agrees to + # ~2e-3. That is still ~6x below the int8 quantization floor, and vastly + # better than dropping adaLN entirely, so only a much larger residual is + # worth complaining about. + if residual > 2e-2: + LOG.warning( + "H3 PowerLoraStack: adaLN basis fit residual %.2e is high - the silu grid " + "and the checkpoint's adaln_t_table probably come from different bakes", + residual, + ) + else: + LOG.info("H3 PowerLoraStack: adaLN basis fit residual %.2e", residual) + + out = (v.to(torch.float32), c.to(torch.float32), residual) + _basis_cache[key] = out + return out + + +def fit_table_to_table(source_table: torch.Tensor, target_table: torch.Tensor): + """Map a source curve basis onto a target curve basis. + + Used when a LoRA ships its own ``adaln_t_table``: the two bakes span nearly + the same subspace but with different rotations and centering, so a LoRA + trained against one is wrong on the other even though the shapes match. + Fitting ``source ~= a + target M`` needs no silu grid at all. + + Returns ``(M [k_src, k_tgt], a [k_src], residual)``. + """ + src = source_table.to(torch.float64) + tgt = target_table.to(torch.float64) + if src.shape[0] != tgt.shape[0]: + raise ValueError( + f"adaLN tables have different grid resolutions ({src.shape[0]} vs {tgt.shape[0]})" + ) + design = torch.cat([torch.ones(tgt.shape[0], 1, dtype=torch.float64), tgt], dim=1) + solution = torch.linalg.lstsq(design, src).solution # [1+k_tgt, k_src] + residual = float((design @ solution - src).norm() / src.norm().clamp(min=1e-12)) + a = solution[0].contiguous() # [k_src] + m = solution[1:].transpose(0, 1).contiguous() # [k_src, k_tgt] + return m.to(torch.float32), a.to(torch.float32), residual + + +class AdalnContext: + """Everything needed to move adaLN LoRA pairs between the two bases.""" + + def __init__(self, target_dim: int, table=None, grid_path: str = ""): + self.target_dim = int(target_dim) + self.table = table + self.grid_path = grid_path + self.residual = None + self._basis = None + self._failed = False + + @property + def is_curve(self) -> bool: + return self.table is not None + + def basis(self, curve_table=None): + """``(V, c, residual)`` relating the dense space to a curve basis. + + The fit needs whichever side of the conversion is the *curve* side. On + a pruned target that is the checkpoint's own table; on a dense target + converting a curve-trained LoRA it has to be the table the LoRA shipped, + because nothing else records which bake it was trained against. + """ + table = self.table if self.table is not None else curve_table + if table is None: + if not self._failed: + LOG.warning( + "H3 PowerLoraStack: cannot rebase adaLN onto a dense checkpoint - " + "the LoRA does not carry the adaln_t_table of the curve bake it " + "was trained against" + ) + self._failed = True + return None + if self._basis is not None or self._failed: + return self._basis + try: + if not self.grid_path: + raise ValueError( + f"no {GRID_FILENAME} found; put it in models/h3_adaln/ to enable " + "adaLN porting" + ) + grid = load_silu_grid(self.grid_path) + self._basis = fit_basis(table.detach().to("cpu"), grid) + self.residual = self._basis[2] + except Exception as exc: + LOG.warning("H3 PowerLoraStack: adaLN porting unavailable (%s)", exc) + self._failed = True + self._basis = None + return self._basis + + +def port_adaln_pairs(sd: dict, ctx: AdalnContext, source_table=None): + """Rebase every ``adaln_proj.linear`` LoRA pair onto the target's basis. + + ``source_table`` is the LoRA's own ``adaln_t_table`` if it shipped one. + That is the better route when available: two curve bakes can have the same + width but different rotations and centering, so a LoRA whose adaLN width + already matches may still be wrong on this checkpoint, and only the table + reveals it. It also needs no silu grid. + + Mutates nothing; returns ``(new_sd, stats)``. + """ + stats = {"ported": 0, "skipped": 0, "ok": 0, "rebased": 0, "residual": None} + modules: dict[str, dict] = {} + for key in sd: + if ".adaln_proj.linear." not in key: + continue + body, _, suffix = key.rpartition(".adaln_proj.linear.") + modules.setdefault(body + ".adaln_proj.linear", {})[suffix] = key + + if not modules: + return sd, stats + + # curve -> curve via the LoRA's own table, when it carries one + table_map = None + if source_table is not None and ctx.table is not None: + try: + m, a_const, residual = fit_table_to_table(source_table, ctx.table) + table_map = (m, a_const) + stats["residual"] = residual + LOG.info("H3 PowerLoraStack: adaLN table-to-table fit residual %.2e", residual) + except Exception as exc: + LOG.warning("H3 PowerLoraStack: adaLN table rebase unavailable (%s)", exc) + + out = dict(sd) + basis = None + for module, parts in modules.items(): + a_key = parts.get("lora_A.weight") or parts.get("lora_down.weight") + b_key = parts.get("lora_B.weight") or parts.get("lora_up.weight") + if a_key is None or b_key is None: + continue + a = sd[a_key] + if a.ndim != 2: + continue + source_dim = a.shape[1] + b = sd[b_key] + a32 = a.to(torch.float32) + + if table_map is not None and source_dim == table_map[0].shape[0]: + m, a_const = table_map + a_new = a32 @ m # [r, k_tgt] + const = a32 @ a_const # [r] + sign = 1.0 + kind = "rebased" + elif source_dim == ctx.target_dim: + stats["ok"] += 1 + continue + else: + if basis is None: + basis = ctx.basis(curve_table=source_table) + if basis is None: + for key in parts.values(): + out.pop(key, None) + stats["skipped"] += 1 + continue + v, c, residual = basis # V: [2688, k], c: [2688] + stats["residual"] = residual + if source_dim == v.shape[0] and ctx.target_dim == v.shape[1]: + a_new = a32 @ v # [r, k] + const = a32 @ c # [r] + sign = 1.0 + elif source_dim == v.shape[1] and ctx.target_dim == v.shape[0]: + a_new = a32 @ torch.linalg.pinv(v) # [r, 2688] + const = a_new @ c # [r] + sign = -1.0 + else: + LOG.warning( + "H3 PowerLoraStack: adaLN pair %s has width %d, expected %d or %d - skipped", + module, source_dim, v.shape[0], v.shape[1], + ) + for key in parts.values(): + out.pop(key, None) + stats["skipped"] += 1 + continue + kind = "ported" + + # comfy scales the A/B pair by alpha/rank inside the adapter but applies + # a diff_b patch verbatim, so the factor has to be folded in here + alpha_key = parts.get("alpha") + scale = 1.0 + if alpha_key is not None: + try: + alpha = sd[alpha_key] + scale = float(alpha.item() if hasattr(alpha, "item") else alpha) / a_new.shape[0] + except Exception: + scale = 1.0 + bias_delta = sign * scale * (b.to(torch.float32) @ const) + + out.pop(a_key, None) + out.pop(b_key, None) + out[module + ".lora_A.weight"] = a_new.to(a.dtype) + out[module + ".lora_B.weight"] = b + if alpha_key is not None and alpha_key != module + ".alpha": + out.pop(alpha_key, None) + out[module + ".alpha"] = sd[alpha_key] + out[module + ".diff_b"] = bias_delta.to(b.dtype) + stats[kind] += 1 + + return out, stats + + +def read_target(diffusion_model): + """Inspect a loaded H3 DiT for its adaLN input width and curve table.""" + table = getattr(diffusion_model, "adaln_t_table", None) + dim = None + try: + dim = diffusion_model.blocks[0].adaln_proj.linear.weight.shape[1] + except Exception: + if table is not None: + dim = int(table.shape[1]) + return dim, table diff --git a/h3lora/apply.py b/h3lora/apply.py new file mode 100644 index 0000000..975a0a0 --- /dev/null +++ b/h3lora/apply.py @@ -0,0 +1,197 @@ +"""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 diff --git a/h3lora/branch.py b/h3lora/branch.py new file mode 100644 index 0000000..e76cf44 --- /dev/null +++ b/h3lora/branch.py @@ -0,0 +1,125 @@ +"""Runtime low-rank branches for quantized MiniMax H3 weights. + +Merging a LoRA into a w4a8 / int4 weight destroys it. ComfyUI's stock path is +dequantize -> add delta -> ``requantize_from_float(scale="recalculate")``, and +that round trip is not idempotent: re-fitting the codebook and re-rounding to +int4 injects ~1.5% relative weight noise, while these deltas are 0.01-0.08% of +the weight. The merge therefore replaces the adapter with noise -- measured +end to end it recovers under a third of the LoRA magnitude at cos 0.12-0.14 +against the correct result. + +Keeping the branch separate is exact, keeps the quantized kernel, and costs +~1.5% extra FLOPs at rank 64. + +Stacking is fused: N LoRAs on one Linear become a single pair by concatenating +along the rank axis, + + sum_i s_i * B_i @ A_i @ x == [s_1 B_1 | ... | s_N B_N] @ [A_1; ...; A_N] @ x + +so a ten-LoRA stack still costs one extra matmul pair per layer rather than +ten, and only one set of factors has to be moved to the GPU. +""" + +from __future__ import annotations + +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import comfy.model_management + + +class LoraBank(nn.Module): + """Holds every runtime branch's factors as one movable, accountable module.""" + + def __init__(self, pairs: "OrderedDict[str, tuple[torch.Tensor, torch.Tensor]]"): + super().__init__() + self.index: dict[str, int] = {} + for i, (name, (up, down)) in enumerate(pairs.items()): + self.register_parameter(f"up{i}", nn.Parameter(up, requires_grad=False)) + self.register_parameter(f"down{i}", nn.Parameter(down, requires_grad=False)) + self.index[name] = i + + def get(self, name: str): + i = self.index[name] + return getattr(self, f"up{i}"), getattr(self, f"down{i}") + + +class LoraBranch: + """Object patch for a Linear's ``forward``: ``y = W_q(x) + up @ down @ x``. + + Patches ``.forward`` rather than the module itself -- replacing the module + would shift the memory manager's ``named_modules`` keys. Holds the + *patcher* rather than the bare module so comfy's weakref bookkeeping does + not report a leaked model. + """ + + def __init__(self, bank_patcher, name: str, original): + self.bank_patcher = bank_patcher + self.bank: LoraBank = bank_patcher.model + self.name = name + self.original = original + + def __call__(self, input: torch.Tensor, *args, **kwargs): + out = self.original(input, *args, **kwargs) + up, down = self.bank.get(self.name) + h = F.linear(input, comfy.model_management.cast_to_device(down, input.device, input.dtype)) + delta = F.linear(h, comfy.model_management.cast_to_device(up, input.device, input.dtype)) + return out + delta.to(out.dtype) + + +def fuse(contributions, compute_dtype): + """Concatenate per-LoRA ``(up, down, scale)`` triples into one pair. + + The strength/alpha scale is folded into ``up`` so the forward path stays a + bare pair of matmuls. + """ + ups, downs = [], [] + for up, down, scale in contributions: + ups.append((up.to(torch.float32) * float(scale)).to(compute_dtype)) + downs.append(down.to(compute_dtype)) + if len(ups) == 1: + return ups[0], downs[0] + return torch.cat(ups, dim=1).contiguous(), torch.cat(downs, dim=0).contiguous() + + +def attach(model_patcher, per_module, compute_dtype, tag: str): + """Build the bank, register it for VRAM accounting, patch every forward. + + ``per_module`` maps a module path to a list of ``(up, down, scale)``. + Returns the number of branched layers. + """ + if not per_module: + return 0 + + import comfy.model_patcher + + pairs: "OrderedDict[str, tuple[torch.Tensor, torch.Tensor]]" = OrderedDict() + for module_path, contributions in per_module.items(): + pairs[module_path] = fuse(contributions, compute_dtype) + + bank = LoraBank(pairs) + bank_patcher = comfy.model_patcher.ModelPatcher( + bank, + load_device=comfy.model_management.get_torch_device(), + offload_device=comfy.model_management.unet_offload_device(), + ) + model_patcher.set_additional_models(tag, [bank_patcher]) + + for module_path in pairs: + forward_key = f"{module_path}.forward" + original = model_patcher.get_model_object(forward_key) + model_patcher.add_object_patch(forward_key, LoraBranch(bank_patcher, module_path, original)) + return len(pairs) + + +def bank_bytes(per_module, compute_dtype) -> int: + """Approximate VRAM the bank will occupy, for reporting.""" + itemsize = torch.empty((), dtype=compute_dtype).element_size() + total = 0 + for contributions in per_module.values(): + for up, down, _ in contributions: + total += (up.numel() + down.numel()) * itemsize + return total diff --git a/h3lora/detect.py b/h3lora/detect.py new file mode 100644 index 0000000..d47fc81 --- /dev/null +++ b/h3lora/detect.py @@ -0,0 +1,113 @@ +"""Cheap safetensors introspection for MiniMax H3 LoRA files. + +Reads only the header, so scanning a 2.5 GB LoRA costs a few kilobytes of IO. +Used both for the node's diagnostic report and for deciding, before anything is +loaded, whether a file needs adaLN porting. +""" + +from __future__ import annotations + +import json +import os +import re +import struct + +_cache: dict[tuple, dict] = {} + +_CONVENTIONS = ( + ("kohya", re.compile(r"^lora_unet_")), + ("lycoris", re.compile(r"^lycoris_")), + ("peft", re.compile(r"^(base_model\.model\.|transformer\.)")), + ("comfy", re.compile(r"^diffusion_model\.")), +) + + +def read_header(path: str) -> dict: + with open(path, "rb") as f: + length = struct.unpack(" str: + for name, pattern in _CONVENTIONS: + if any(pattern.match(k) for k in keys): + return name + if any(k.startswith(("blocks.", "token_refiner.", "final_layer.")) for k in keys): + return "bare" + return "unknown" + + +def inspect(path: str) -> dict: + """Summarise one LoRA file. Cached on (path, mtime, size).""" + try: + st = os.stat(path) + except OSError as exc: + return {"error": str(exc), "name": os.path.basename(path)} + key = (path, st.st_mtime_ns, st.st_size) + hit = _cache.get(key) + if hit is not None: + return hit + + info: dict = {"name": os.path.basename(path), "path": path, "size": st.st_size} + try: + header = read_header(path) + except Exception as exc: + info["error"] = f"unreadable header: {exc}" + _cache[key] = info + return info + + metadata = header.pop("__metadata__", {}) or {} + keys = list(header) + info["tensors"] = len(keys) + info["convention"] = _convention(keys) + info["dtypes"] = sorted({v.get("dtype", "?") for v in header.values()}) + info["base_model"] = metadata.get("ss_base_model_version", "") + info["trainer"] = "" + try: + info["trainer"] = json.loads(metadata.get("software", "{}")).get("name", "") + except Exception: + pass + if not info["trainer"] and metadata.get("ss_network_module"): + info["trainer"] = metadata["ss_network_module"] + + ranks: set[int] = set() + adaln_dims: set[int] = set() + modules: set[str] = set() + has_alpha = False + for k, v in header.items(): + if k.endswith(".alpha"): + has_alpha = True + continue + shape = v.get("shape") or [] + is_down = k.endswith(("lora_A.weight", "lora_down.weight")) + if is_down and len(shape) == 2: + ranks.add(int(shape[0])) + if "adaln" in k: + adaln_dims.add(int(shape[1])) + for tag in ("adaln_proj", "attn.qkv_proj", "attn_qkv_proj", "attn.out_proj", + "attn_out_proj", "mlp.fc1", "mlp_fc1", "mlp.fc2", "mlp_fc2", + "token_refiner", "final_layer"): + if tag in k: + modules.add(tag.replace("_proj", "").replace(".", "_").replace("attn_qkv", "qkv") + .replace("attn_out", "out")) + + info["ranks"] = sorted(ranks) + info["alpha"] = has_alpha + info["adaln_dims"] = sorted(adaln_dims) + info["modules"] = sorted(modules) + info["passenger"] = sorted( + k for k in keys + if not re.search(r"(lora_A|lora_B|lora_down|lora_up|alpha|diff|dora|hada|lokr|oft)", k) + )[:8] + _cache[key] = info + return info + + +def summary_line(info: dict) -> str: + if info.get("error"): + return f"{info['name']}: {info['error']}" + ranks = "/".join(str(r) for r in info.get("ranks", [])) or "?" + adaln = info.get("adaln_dims") + adaln_s = f", adaLN dim {'/'.join(str(d) for d in adaln)}" if adaln else ", no adaLN" + return (f"{info['name']}: {info.get('tensors', 0)} tensors, {info.get('convention', '?')} " + f"format, rank {ranks}{adaln_s}") diff --git a/h3lora/keymap.py b/h3lora/keymap.py new file mode 100644 index 0000000..44edc89 --- /dev/null +++ b/h3lora/keymap.py @@ -0,0 +1,126 @@ +"""Key normalisation for MiniMax H3 LoRAs. + +Every H3 LoRA in the wild is one of a handful of naming conventions. Observed +across the local collection: + +* ai-toolkit / diffusers ``diffusion_model.blocks.0.attn.qkv_proj.lora_A.weight`` +* same, no prefix ``blocks.0.attn.qkv_proj.lora_A.weight`` +* kohya / musubi ``lora_unet_blocks_0_attn_qkv_proj.lora_down.weight`` + ``.alpha`` +* lycoris ``lycoris_blocks_0_attn_qkv_proj.*`` +* peft ``base_model.model.blocks.0...`` / ``transformer.blocks.0...`` + +``comfy.lora.model_lora_keys_unet`` already covers the ``diffusion_model.``- +prefixed and ``lora_unet_`` forms, but not the bare form, and it gives us no +handle for rewriting adaLN tensors before they are matched. So we normalise +everything ourselves onto the canonical ``diffusion_model.`` module key, +resolving underscore forms through the model's own state dict rather than +guessing where the word boundaries are. +""" + +from __future__ import annotations + +import re + +# suffixes comfy.lora.load_lora understands, longest first so that +# ``.lora_down.weight`` is not truncated to ``.weight`` +_LORA_SUFFIXES = ( + ".lora_down.weight", ".lora_up.weight", + ".lora_A.weight", ".lora_B.weight", + ".lora_A.default.weight", ".lora_B.default.weight", + ".lora_A", ".lora_B", + ".lora_mid.weight", ".dora_scale", + ".diff_b", ".diff", ".set_weight", ".alpha", + ".hada_w1_a", ".hada_w1_b", ".hada_w2_a", ".hada_w2_b", ".hada_t1", ".hada_t2", + ".lokr_w1", ".lokr_w2", ".lokr_w1_a", ".lokr_w1_b", ".lokr_w2_a", ".lokr_w2_b", + ".lokr_t2", ".oft_blocks", ".boft_blocks", ".rescale", ".w_norm", ".b_norm", +) + +_STRIP_PREFIXES = ( + "diffusion_model.", "transformer.", "base_model.model.", "pipe.dit.", + "model.diffusion_model.", "unet.", +) + +_UNDERSCORE_PREFIXES = ("lora_unet_", "lycoris_", "lora_transformer_") + + +def split_suffix(key: str): + """``blocks.0.mlp.fc1.lora_A.weight`` -> ``('blocks.0.mlp.fc1', '.lora_A.weight')``.""" + for suffix in sorted(_LORA_SUFFIXES, key=len, reverse=True): + if key.endswith(suffix): + return key[: -len(suffix)], suffix + return None, None + + +def build_module_index(model_sd_keys) -> dict: + """Map every spelling of a DiT module path to its canonical form. + + Canonical form is the state-dict key minus ``.weight``, e.g. + ``diffusion_model.blocks.0.attn.qkv_proj``. Both the dotted path and the + underscored path (which is what kohya trainers emit) resolve to it, and the + underscore mapping is derived from the real keys so no heuristic is needed + to decide whether ``qkv_proj`` is one token or two. + """ + index: dict[str, str] = {} + for k in model_sd_keys: + if not k.startswith("diffusion_model.") or not k.endswith(".weight"): + continue + canonical = k[: -len(".weight")] + bare = canonical[len("diffusion_model."):] + index[canonical] = canonical + index[bare] = canonical + index[bare.replace(".", "_")] = canonical + return index + + +def _strip_prefixes(path: str) -> str: + changed = True + while changed: + changed = False + for pre in _STRIP_PREFIXES: + if path.startswith(pre): + path = path[len(pre):] + changed = True + break + for pre in _UNDERSCORE_PREFIXES: + if path.startswith(pre): + path = path[len(pre):] + break + return path + + +def canonical_module(key_body: str, index: dict): + """Resolve one LoRA key body to a canonical module path, or None.""" + body = _strip_prefixes(key_body) + hit = index.get(body) + if hit is not None: + return hit + # kohya writes ``lora_unet__blocks_0_...`` (double underscore) and some + # trainers keep a dotted tail after the underscored head + squashed = body.replace(".", "_") + hit = index.get(squashed) or index.get(re.sub(r"_{2,}", "_", squashed)) + if hit is not None: + return hit + return None + + +def normalize(lora_sd: dict, index: dict): + """Rewrite a LoRA state dict onto canonical module keys. + + Returns ``(normalized_sd, unmatched_keys)``. Non-LoRA passenger tensors + (``adaln_t_table`` and full ``adaln_proj.linear.weight`` replacements found + in "complete pruned" LoRA files) are passed through untouched so the caller + can deal with them. + """ + out: dict = {} + unmatched: list[str] = [] + for key, value in lora_sd.items(): + body, suffix = split_suffix(key) + if body is None: + out[key] = value # passenger tensor, handled upstream + continue + module = canonical_module(body, index) + if module is None: + unmatched.append(key) + continue + out[module + suffix] = value + return out, unmatched diff --git a/h3lora/nodes.py b/h3lora/nodes.py new file mode 100644 index 0000000..b938922 --- /dev/null +++ b/h3lora/nodes.py @@ -0,0 +1,202 @@ +"""Node definitions for the MiniMax H3 Power LoRA Stack.""" + +from __future__ import annotations + +import logging + +import folder_paths + +from . import apply as apply_mod +from . import detect + +LOG = logging.getLogger("h3.powerlorastack") +CATEGORY = "MiniMax-H3/lora" + + +class AnyType(str): + def __ne__(self, other): + return False + + +ANY = AnyType("*") + + +class FlexibleOptionalInputType(dict): + """Accepts the arbitrary ``lora_N`` inputs the frontend adds at runtime. + + ComfyUI validates a prompt against INPUT_TYPES; a stack whose row count is + decided in the browser has no fixed schema, so this dict reports that it + contains every key and hands back a permissive type for the ones it does + not know about. + """ + + def __init__(self, type, data=None): + super().__init__() + self.type = type + self.data = data or {} + self.update(self.data) + + def __getitem__(self, key): + if key in self.data: + return self.data[key] + return (self.type,) + + def __contains__(self, key): + return True + + +def _resolve_lora(name: str): + """Map a stored lora name onto a real file, tolerating separator drift.""" + if not name or name == "None": + return None + path = folder_paths.get_full_path("loras", name) + if path: + return path + wanted = name.replace("\\", "/").lower() + for candidate in folder_paths.get_filename_list("loras"): + normalized = candidate.replace("\\", "/").lower() + if normalized == wanted or normalized.endswith("/" + wanted.rsplit("/", 1)[-1]): + return folder_paths.get_full_path("loras", candidate) + return None + + +def _collect(kwargs): + """Pull enabled rows out of the dynamic ``lora_N`` inputs, in UI order.""" + rows = [] + for key, value in kwargs.items(): + if not key.lower().startswith("lora_") or not isinstance(value, dict): + continue + if "lora" not in value: + continue + try: + order = int(key.split("_", 1)[1]) + except (IndexError, ValueError): + order = 1 << 30 + rows.append((order, value)) + rows.sort(key=lambda item: item[0]) + + entries = [] + for _, value in rows: + if not value.get("on", True): + continue + strength = float(value.get("strength", 1.0)) + if strength == 0.0: + continue + name = value.get("lora") + if not name or name == "None": + continue # a row the user added but has not filled in yet + path = _resolve_lora(name) + if path is None: + LOG.warning("H3 Power LoRA Stack: could not find lora %r, skipping", name) + continue + entries.append({"name": name, "path": path, "strength": strength}) + return entries + + +class H3PowerLoraStack: + """Stacked multi-LoRA loader for MiniMax H3 across every weight format. + + Handles the three things the stock loader and rgthree's Power Lora Loader + get wrong on H3: + + * quantized bases -- w4a8/int4 weights are patched with an exact runtime + branch instead of a lossy dequantize/requantize merge; + * adaLN basis mismatch -- LoRAs trained against a dense checkpoint are + rebased onto a pruned checkpoint's curve, instead of being skipped; + * key conventions -- ai-toolkit, kohya, lycoris, peft and bare-prefix + layouts all resolve against the model's own key set. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": {}, + "optional": FlexibleOptionalInputType(ANY, { + "model": ("MODEL",), + "quantized_layers": (["auto", "branch", "merge"], { + "default": "auto", + "tooltip": "auto: runtime branch for quantized layers, merge for the rest. " + "branch: never modify a weight. merge: stock behaviour " + "(destroys LoRAs on w4a8/int4).", + }), + "adaln_port": (["auto", "off"], { + "default": "auto", + "tooltip": "Rebase adaLN LoRA pairs between dense (2688) and curve (8) " + "checkpoints. Needs h3_silu_temb_grid.safetensors.", + }), + }), + "hidden": {}, + } + + RETURN_TYPES = ("MODEL", "STRING") + RETURN_NAMES = ("MODEL", "report") + OUTPUT_TOOLTIPS = ("Patched model", "Per-LoRA account of what was applied") + FUNCTION = "apply" + CATEGORY = CATEGORY + DESCRIPTION = __doc__ + + def apply(self, model=None, quantized_layers="auto", adaln_port="auto", **kwargs): + if model is None: + raise ValueError("H3 Power LoRA Stack: no model connected") + + entries = _collect(kwargs) + if not entries: + return (model, "no LoRAs enabled") + + patcher, report = apply_mod.apply_stack( + model, entries, mode=quantized_layers, adaln_mode=adaln_port, + ) + text = report.text() + LOG.info("H3 Power LoRA Stack:\n%s", text) + return (patcher, text) + + +class H3LoraInspector: + """Report the format, rank and adaLN basis of a LoRA without loading it.""" + + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "lora_name": (folder_paths.get_filename_list("loras"),), + }} + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("info",) + FUNCTION = "inspect" + CATEGORY = CATEGORY + DESCRIPTION = __doc__ + + def inspect(self, lora_name): + path = _resolve_lora(lora_name) + if path is None: + return (f"{lora_name}: not found",) + info = detect.inspect(path) + if info.get("error"): + return (f"{info['name']}: {info['error']}",) + lines = [ + f"file : {info['name']}", + f"size : {info['size'] / (1024 ** 2):.0f} MB, {info['tensors']} tensors", + f"format : {info['convention']}" + + (f" ({info['trainer']})" if info.get("trainer") else ""), + f"dtype : {', '.join(info['dtypes'])}", + f"rank : {'/'.join(str(r) for r in info['ranks']) or '?'}" + + (" (alpha present)" if info["alpha"] else " (no alpha)"), + f"adaLN : {'/'.join(str(d) for d in info['adaln_dims']) or 'none'}", + f"modules : {', '.join(info['modules'])}", + ] + if info.get("base_model"): + lines.append(f"base : {info['base_model']}") + if info.get("passenger"): + lines.append(f"passenger : {', '.join(info['passenger'])}") + return ("\n".join(lines),) + + +NODE_CLASS_MAPPINGS = { + "H3PowerLoraStack": H3PowerLoraStack, + "H3LoraInspector": H3LoraInspector, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "H3PowerLoraStack": "MiniMax H3 Power LoRA Stack", + "H3LoraInspector": "MiniMax H3 LoRA Inspector", +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..5805a2f --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "comfyui-h3-powerlorastack" +description = "Stacked multi-LoRA loader for MiniMax H3 across bf16, fp8, int8 and w4a8 weights." +version = "1.0.0" +license = { text = "MIT" } +dependencies = [] + +[project.urls] +Repository = "https://github.com/doggeddalle/ComfyUI-H3-PowerLoraStack" + +[tool.comfy] +PublisherId = "doggeddalle" +DisplayName = "MiniMax H3 Power LoRA Stack" +Icon = "" diff --git a/web/h3_power_lora_stack.js b/web/h3_power_lora_stack.js new file mode 100644 index 0000000..14eee19 --- /dev/null +++ b/web/h3_power_lora_stack.js @@ -0,0 +1,502 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +const NODE_TYPE = "H3PowerLoraStack"; +const ROW_HEIGHT = 22; +const MARGIN = 12; + +// Session cache for the loras folder listing. Deliberately a module global +// rather than graph.extra: anything hung off the graph is serialized into the +// saved workflow, which would bake a stale copy of the whole folder into every +// .json. The promise itself is cached so N rows opening at once share one +// request. +let loraListPromise = null; + +/** + * The stack node declares no ``lora_name`` input of its own -- rows are added + * in the browser -- so the folder listing is borrowed from a node that does. + */ +function fetchLoraList() { + if (!loraListPromise) { + loraListPromise = api + .fetchApi("/object_info/LoraLoaderModelOnly") + .then((res) => res.json()) + .then( + (info) => info?.LoraLoaderModelOnly?.input?.required?.lora_name?.[0] ?? [] + ) + .catch((err) => { + console.warn("[H3PowerLoraStack] could not fetch lora list", err); + loraListPromise = null; // let the next picker retry + return []; + }); + } + return loraListPromise; +} + +function drawToggle(ctx, x, y, size, on) { + ctx.save(); + ctx.beginPath(); + ctx.roundRect(x, y + 3, size * 1.7, size - 6, (size - 6) / 2); + ctx.fillStyle = on ? "#3d8c5a" : "#3a3a3a"; + ctx.fill(); + ctx.beginPath(); + const knob = (size - 8) / 2; + ctx.arc(x + (on ? size * 1.7 - knob - 3 : knob + 3), y + size / 2, knob, 0, Math.PI * 2); + ctx.fillStyle = on ? "#d8f5e4" : "#8a8a8a"; + ctx.fill(); + ctx.restore(); +} + +function shortName(name) { + if (!name || name === "None") return "click to choose"; + const base = String(name).replace(/\\/g, "/").split("/").pop(); + return base.replace(/\.(safetensors|ckpt|pt|bin)$/i, ""); +} + +/** + * One stack row: enable toggle, lora name, strength, remove button. + * Values serialize as {on, lora, strength}, which is what the Python side's + * flexible `lora_N` inputs expect. + */ +function makeLoraWidget(node, name, value) { + const widget = { + name, + type: "H3_LORA", + value: Object.assign({ on: true, lora: "None", strength: 1.0 }, value || {}), + options: { serialize: true }, + hitAreas: {}, + + computeSize() { + return [node.size[0], ROW_HEIGHT]; + }, + + draw(ctx, node_, widgetWidth, y, height) { + const margin = MARGIN; + const left = margin; + const right = widgetWidth - margin; + const midY = y + height / 2; + ctx.save(); + + ctx.beginPath(); + ctx.roundRect(left, y + 1, right - left, height - 2, 6); + ctx.fillStyle = this.value.on ? "#2b2b2b" : "#232323"; + ctx.fill(); + + let cursor = left + 8; + drawToggle(ctx, cursor, y, height, this.value.on); + this.hitAreas.toggle = [cursor, cursor + height * 1.7]; + cursor += height * 1.7 + 10; + + // remove button, laid out from the right edge inward + const removeX = right - 20; + ctx.fillStyle = "#8a8a8a"; + ctx.font = "14px sans-serif"; + ctx.textAlign = "center"; + ctx.textBaseline = "middle"; + ctx.fillText("✕", removeX + 6, midY); + this.hitAreas.remove = [removeX, removeX + 16]; + + const strengthRight = removeX - 10; + const strengthLeft = strengthRight - 74; + ctx.fillStyle = "#1e1e1e"; + ctx.beginPath(); + ctx.roundRect(strengthLeft, y + 3, strengthRight - strengthLeft, height - 6, 4); + ctx.fill(); + ctx.fillStyle = "#b8b8b8"; + ctx.font = "11px sans-serif"; + ctx.fillText("◀", strengthLeft + 9, midY); + ctx.fillText("▶", strengthRight - 9, midY); + ctx.fillStyle = this.value.on ? "#e8e8e8" : "#777"; + ctx.font = "12px sans-serif"; + ctx.fillText( + Number(this.value.strength).toFixed(2), + (strengthLeft + strengthRight) / 2, + midY + ); + this.hitAreas.strengthDown = [strengthLeft, strengthLeft + 18]; + this.hitAreas.strengthUp = [strengthRight - 18, strengthRight]; + this.hitAreas.strengthValue = [strengthLeft + 18, strengthRight - 18]; + + const nameLeft = cursor; + const nameRight = strengthLeft - 10; + ctx.textAlign = "left"; + ctx.fillStyle = this.value.on ? "#e8e8e8" : "#777"; + ctx.font = "12px sans-serif"; + const label = shortName(this.value.lora); + let text = label; + const maxWidth = nameRight - nameLeft; + if (ctx.measureText(text).width > maxWidth) { + while (text.length > 4 && ctx.measureText(text + "…").width > maxWidth) { + text = text.slice(0, -1); + } + text += "…"; + } + ctx.fillText(text, nameLeft, midY); + this.hitAreas.name = [nameLeft, nameRight]; + + ctx.restore(); + }, + + mouse(event, pos, node_) { + if (event.type !== "pointerdown") return false; + const x = pos[0]; + const hit = (area) => area && x >= area[0] && x <= area[1]; + + if (hit(this.hitAreas.toggle)) { + this.value = { ...this.value, on: !this.value.on }; + node_.setDirtyCanvas(true, true); + return true; + } + if (hit(this.hitAreas.remove)) { + removeLoraWidget(node_, this); + return true; + } + if (hit(this.hitAreas.strengthDown)) { + this.value = { ...this.value, strength: round2(this.value.strength - 0.05) }; + node_.setDirtyCanvas(true, true); + return true; + } + if (hit(this.hitAreas.strengthUp)) { + this.value = { ...this.value, strength: round2(this.value.strength + 0.05) }; + node_.setDirtyCanvas(true, true); + return true; + } + if (hit(this.hitAreas.strengthValue)) { + app.canvas.prompt( + "Strength", + this.value.strength, + (v) => { + const parsed = parseFloat(v); + if (!Number.isNaN(parsed)) { + this.value = { ...this.value, strength: parsed }; + node_.setDirtyCanvas(true, true); + } + }, + event + ); + return true; + } + if (hit(this.hitAreas.name)) { + showLoraMenu(node_, this, event); + return true; + } + return false; + }, + + serializeValue() { + return { ...this.value }; + }, + }; + return widget; +} + +function round2(v) { + return Math.round(v * 100) / 100; +} + +/* ---------------------------------------------------------------- picker -- */ + +const PICKER_CSS = ` +.h3pls-picker{position:fixed;z-index:10000;display:flex;flex-direction:column; + width:460px;max-width:92vw;background:#1e1e1e;border:1px solid #4a4a4a; + border-radius:8px;box-shadow:0 12px 32px rgba(0,0,0,.6);overflow:hidden; + font-family:system-ui,-apple-system,sans-serif;} +.h3pls-picker input{margin:8px;padding:7px 9px;background:#111;color:#eee; + border:1px solid #555;border-radius:5px;font-size:13px;outline:none;} +.h3pls-picker input:focus{border-color:#3d8c5a;} +.h3pls-count{padding:0 12px 6px;font-size:11px;color:#777;} +.h3pls-list{overflow-y:auto;max-height:46vh;} +.h3pls-item{padding:6px 12px;font-size:12.5px;color:#ddd;cursor:pointer; + white-space:nowrap;overflow:hidden;text-overflow:ellipsis;} +.h3pls-item:hover{background:#2c2c2c;} +.h3pls-item.sel{background:#31543f;} +.h3pls-item .dir{color:#7a7a7a;} +.h3pls-item .hit{color:#7fd6a0;font-weight:600;} +.h3pls-item.cur::after{content:"current";float:right;color:#666;font-size:10px;} +.h3pls-empty{padding:14px 12px;font-size:12px;color:#888;} +`; + +function ensurePickerCss() { + if (document.getElementById("h3pls-css")) return; + const style = document.createElement("style"); + style.id = "h3pls-css"; + style.textContent = PICKER_CSS; + document.head.appendChild(style); +} + +/** Case-insensitive AND-match of every whitespace-separated term. */ +function matchTerms(name, terms) { + const hay = name.toLowerCase(); + const spans = []; + for (const term of terms) { + const at = hay.indexOf(term); + if (at < 0) return null; + spans.push([at, at + term.length]); + } + return spans; +} + +/** Basename hits and early hits rank first; then shortest name. */ +function score(name, terms, spans) { + if (!terms.length) return 0; + const slash = name.lastIndexOf("/") + 1; + let s = 0; + for (const [start] of spans) { + if (start >= slash) s -= 1000; + s += start; + } + return s + name.length * 0.01; +} + +function highlight(name, spans) { + const slash = name.lastIndexOf("/") + 1; + const merged = []; + for (const [a, b] of [...spans].sort((p, q) => p[0] - q[0])) { + const last = merged[merged.length - 1]; + if (last && a <= last[1]) last[1] = Math.max(last[1], b); + else merged.push([a, b]); + } + const esc = (t) => + t.replace(/&/g, "&").replace(//g, ">"); + let html = ""; + let at = 0; + for (const [a, b] of merged) { + html += wrapDir(esc(name.slice(at, a)), at, slash); + html += `${esc(name.slice(a, b))}`; + at = b; + } + html += wrapDir(esc(name.slice(at)), at, slash); + return html; +} + +function wrapDir(text, offset, slash) { + if (!text) return ""; + if (offset + text.length <= slash) return `${text}`; + if (offset >= slash) return text; + const cut = slash - offset; + return `${text.slice(0, cut)}${text.slice(cut)}`; +} + +/** + * Filterable LoRA picker. LiteGraph's ContextMenu has no search, which is + * unusable once the loras folder holds more than a screenful, so this is a + * self-contained overlay: type to narrow, arrows to move, Enter to take it. + */ +async function showLoraMenu(node, widget, event, { removeOnCancel = false } = {}) { + const names = await fetchLoraList(); + if (!names.length) { + if (removeOnCancel) removeLoraWidget(node, widget); + alert("No LoRAs found in models/loras."); + return; + } + ensurePickerCss(); + + const current = widget.value.lora; + const root = document.createElement("div"); + root.className = "h3pls-picker"; + const input = document.createElement("input"); + input.type = "text"; + input.placeholder = "Type to filter… (space-separated terms all must match)"; + input.spellcheck = false; + const count = document.createElement("div"); + count.className = "h3pls-count"; + const list = document.createElement("div"); + list.className = "h3pls-list"; + root.append(input, count, list); + document.body.appendChild(root); + + // place near the click, clamped into the viewport + const px = event?.clientX ?? window.innerWidth / 2; + const py = event?.clientY ?? window.innerHeight / 2; + root.style.left = `${Math.min(px, window.innerWidth - 470)}px`; + root.style.top = `${Math.min(py, window.innerHeight - 320)}px`; + + let rows = []; + let sel = 0; + + let committed = false; + const close = () => { + document.removeEventListener("pointerdown", onOutside, true); + root.remove(); + // a row opened straight from "Add LoRA" and then dismissed was never + // wanted, so take it back out rather than leaving an empty slot behind + if (removeOnCancel && !committed) removeLoraWidget(node, widget); + }; + const commit = (value) => { + committed = true; + widget.value = { ...widget.value, lora: value }; + node.setDirtyCanvas(true, true); + close(); + }; + const onOutside = (e) => { + if (!root.contains(e.target)) close(); + }; + + const render = () => { + const query = input.value.trim().toLowerCase(); + const terms = query ? query.split(/\s+/) : []; + rows = []; + for (const name of names) { + const spans = matchTerms(name, terms); + if (spans === null) continue; + rows.push({ name, spans, s: score(name, terms, spans) }); + } + if (terms.length) rows.sort((a, b) => a.s - b.s); + sel = 0; + if (!terms.length) { + const at = rows.findIndex((r) => r.name === current); + if (at > 0) sel = at; + } + + count.textContent = `${rows.length} of ${names.length}`; + list.replaceChildren(); + if (!rows.length) { + const empty = document.createElement("div"); + empty.className = "h3pls-empty"; + empty.textContent = "No LoRA matches that filter."; + list.appendChild(empty); + return; + } + rows.forEach((row, i) => { + const item = document.createElement("div"); + item.className = + "h3pls-item" + (i === sel ? " sel" : "") + (row.name === current ? " cur" : ""); + item.innerHTML = highlight(row.name, row.spans); + item.title = row.name; + item.addEventListener("pointerdown", (e) => { + e.preventDefault(); + e.stopPropagation(); + commit(row.name); + }); + list.appendChild(item); + }); + scrollToSel(); + }; + + const scrollToSel = () => { + const el = list.children[sel]; + if (el && el.scrollIntoView) el.scrollIntoView({ block: "nearest" }); + }; + + const move = (delta) => { + if (!rows.length) return; + const prev = list.children[sel]; + if (prev) prev.classList.remove("sel"); + sel = (sel + delta + rows.length) % rows.length; + const next = list.children[sel]; + if (next) next.classList.add("sel"); + scrollToSel(); + }; + + input.addEventListener("input", render); + input.addEventListener("keydown", (e) => { + if (e.key === "ArrowDown") { e.preventDefault(); move(1); } + else if (e.key === "ArrowUp") { e.preventDefault(); move(-1); } + else if (e.key === "PageDown") { e.preventDefault(); move(10); } + else if (e.key === "PageUp") { e.preventDefault(); move(-10); } + else if (e.key === "Enter") { + e.preventDefault(); + if (rows[sel]) commit(rows[sel].name); + } else if (e.key === "Escape") { + e.preventDefault(); + close(); + } + e.stopPropagation(); // keep ComfyUI's global hotkeys out of the box + }); + + render(); + const grabFocus = () => { + if (root.isConnected && document.activeElement !== input) input.focus(); + }; + // Defer by a full task, not a microtask: the pointerdown that opened this is + // still being dispatched, and a capture listener added now would see it and + // close the picker immediately. + setTimeout(() => { + document.addEventListener("pointerdown", onOutside, true); + grabFocus(); + }, 0); + // LiteGraph pulls focus back to the canvas on a delay in some paths; re-assert + // once so the caret really is in the filter box and typing goes there. + setTimeout(grabFocus, 60); +} + +function loraWidgets(node) { + return (node.widgets || []).filter((w) => w.type === "H3_LORA"); +} + +function renumber(node) { + loraWidgets(node).forEach((w, i) => { + w.name = `lora_${i + 1}`; + }); +} + +function addLoraWidget(node, value) { + const widget = makeLoraWidget(node, `lora_${loraWidgets(node).length + 1}`, value); + node.widgets = node.widgets || []; + // keep the "Add LoRA" button last + const buttonIndex = node.widgets.findIndex((w) => w.name === "➕ Add LoRA"); + if (buttonIndex >= 0) node.widgets.splice(buttonIndex, 0, widget); + else node.widgets.push(widget); + renumber(node); + resize(node); + return widget; +} + +function removeLoraWidget(node, widget) { + const i = node.widgets.indexOf(widget); + if (i >= 0) node.widgets.splice(i, 1); + renumber(node); + resize(node); + node.setDirtyCanvas(true, true); +} + +function resize(node) { + const [minWidth, minHeight] = node.computeSize(); + // Width stays the user's choice; height tracks the row count so removing a + // row hands the space back instead of leaving a hole. + node.setSize([Math.max(node.size[0], minWidth), minHeight]); + node.setDirtyCanvas(true, true); +} + +app.registerExtension({ + name: "h3.power.lora.stack", + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== NODE_TYPE) return; + + const onCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + onCreated?.apply(this, arguments); + this.serialize_widgets = true; + + // (value, canvas, node, pos, event) -- the event is the pointerdown that + // started the click, so the picker opens right under the button. + this.addWidget("button", "➕ Add LoRA", null, (_v, _canvas, _node, _pos, event) => { + const widget = addLoraWidget(this); + // open the picker immediately: the row is added to be filled in, so the + // filter box takes the keyboard straight away + showLoraMenu(this, widget, event, { removeOnCancel: true }); + }); + + fetchLoraList(); // warm the cache so the first picker opens instantly + resize(this); + }; + + const onConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function (info) { + onConfigure?.apply(this, arguments); + // ComfyUI restores widgets_values positionally against the widgets that + // existed before this ran, which is only the fixed combos and the button. + // The stack rows are recovered by shape instead of position, so a slot + // shifting cannot corrupt them. + const values = info?.widgets_values; + if (!Array.isArray(values)) return; + const rows = values.filter( + (v) => v && typeof v === "object" && !Array.isArray(v) && "lora" in v + ); + for (const w of loraWidgets(this)) removeLoraWidget(this, w); + for (const row of rows) addLoraWidget(this, row); + resize(this); + }; + }, +});