From 4ef4de0a5dce1cbfad674faa6633841373636086 Mon Sep 17 00:00:00 2001 From: doggeddalle <153059103+cicalooo@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:22:53 +1000 Subject: [PATCH] Robust adaln, pdd, keymap, modality fixes Multiple improvements and hardening across the Power LoRA stack: - adaln: stronger caching (file-version keys), device-safe grid/basis computation, normalized tensor hashing for basis cache, more robust table-to-table rebasing with residual checks and fallback, better detection of adaLN weight name variants. - apply: validate numeric strengths, reject unsupported passengers, handle loading errors, detect shape mismatches, track rejected/skipped/quantized merges, improved PDD handoff and reporting. - pdd: introduce movable PddBank/PddState, register banks as additional models, thread-local PDD context, strength-scheduled blending and safer interval clamping. - keymap/detect/gain/modality/nodes/branch: add support for varied LoRA key spellings, more validation, skip incomplete adapters, and small bugfixes. - server: surface partial calibration note. - pyproject/README/.gitignore: packaging, docs, and ignore cleanup; add Python >=3.10 requirement and publisher repo fix. - web UI: fix balance button behavior, safer auto-balance error handling and factor sanitization. - tests: add unit tests for adaln, modality and pdd. Overall: better robustness to varied LoRA formats, mixed devices, malformed inputs, and improved user-facing error/reporting. --- .gitignore | 35 +------- README.md | 41 +++++++-- h3lora/adaln.py | 97 +++++++++++++------- h3lora/apply.py | 178 +++++++++++++++++++++++++++++-------- h3lora/branch.py | 3 +- h3lora/detect.py | 15 +++- h3lora/gain.py | 29 ++++-- h3lora/keymap.py | 6 +- h3lora/modality.py | 13 ++- h3lora/nodes.py | 40 +++++++-- h3lora/pdd.py | 149 ++++++++++++++++++++++++++----- h3lora/server.py | 2 + pyproject.toml | 5 +- tests/test_adaln.py | 74 +++++++++++++++ tests/test_modality.py | 23 +++++ tests/test_pdd.py | 118 ++++++++++++++++++++++++ web/h3_power_lora_stack.js | 20 +++-- 17 files changed, 689 insertions(+), 159 deletions(-) create mode 100644 tests/test_adaln.py create mode 100644 tests/test_modality.py create mode 100644 tests/test_pdd.py diff --git a/.gitignore b/.gitignore index 636b2b8..e8c7e05 100644 --- a/.gitignore +++ b/.gitignore @@ -1,31 +1,4 @@ -__pycache__/__init__.cpython-312.pyc -h3lora/__pycache__/__init__.cpython-312.pyc -h3lora/__pycache__/adaln.cpython-312.pyc -h3lora/__pycache__/apply.cpython-312.pyc -h3lora/__pycache__/branch.cpython-312.pyc -h3lora/__pycache__/detect.cpython-312.pyc -h3lora/__pycache__/gain.cpython-312.pyc -h3lora/__pycache__/keymap.cpython-312.pyc -h3lora/__pycache__/modality.cpython-312.pyc -h3lora/__pycache__/nodes.cpython-312.pyc -h3lora/__pycache__/pdd.cpython-312.pyc -h3lora/__pycache__/schedule.cpython-312.pyc -h3lora/__pycache__/server.cpython-312.pyc -tests/__pycache__/test_adaln.cpython-312.pyc -tests/__pycache__/test_branch_schedule.cpython-312.pyc -tests/__pycache__/test_pdd.cpython-312.pyc -tests/__pycache__/test_schedule.cpython-312.pyc -tests/test_adaln.py -tests/test_pdd.py -h3lora/__pycache__/server.cpython-312.pyc -h3lora/__pycache__/schedule.cpython-312.pyc -h3lora/__pycache__/nodes.cpython-312.pyc -h3lora/__pycache__/modality.cpython-312.pyc -h3lora/__pycache__/keymap.cpython-312.pyc -h3lora/__pycache__/gain.cpython-312.pyc -h3lora/__pycache__/detect.cpython-312.pyc -h3lora/__pycache__/branch.cpython-312.pyc -h3lora/__pycache__/apply.cpython-312.pyc -h3lora/__pycache__/adaln.cpython-312.pyc -h3lora/__pycache__/__init__.cpython-312.pyc -__pycache__/__init__.cpython-312.pyc +__pycache__/ +*.py[cod] +.pytest_cache/ +AUDIT_REPORT.md diff --git a/README.md b/README.md index 9dc184c..82b9d79 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,8 @@ This node routes quantized layers through an exact runtime low-rank branch `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) +- `branch` — never modify a weight, even in bf16 (fast strength A/B testing); + adapters without a runtime branch are reported as rejected - `merge` — stock behaviour; only useful for comparison One layer is special-cased: `mlp.fc2` under `TensorWiseINT8Layout` is reached @@ -106,7 +107,9 @@ two tokens: | 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.** +The mapper resolves the advertised conventions against the model's own state +dict. Unmatched keys are reported per row; the installed LoRA collection is not +part of this package and should be checked when adding new trainer formats. ### 4. Acc / PDD output heads @@ -114,9 +117,11 @@ alibaba-pai Acc LoRAs ship a 32-interval output-head bank as `final_layer.video_out.set_weight` of shape `[3072, 5376]` (`32 × 96`) plus the audio twin. Stock Comfy `copy_`s that onto the native `[96, 5376]` head and crashes. This node peels those tensors before the stock `set` path and blends -the heads the sampler step spans (same rule as Comfy PR 15908). Use `simple` at -8 steps with shifts 12/3 so the steps land on the trained grid. Later stack rows -that also carry a bank replace the earlier one. +the heads the sampler step spans (same rule as Comfy PR 15908). The row +strength interpolates between the native and blended heads; schedules multiply +that row strength. Use `simple` at 8 steps with shifts 12/3 so the steps land on +the trained grid. Later stack rows that also carry a bank replace the earlier +one. ## Stacking @@ -152,8 +157,8 @@ deltas are scheduled with their LoRA rather than being left at a fixed value. ## Auto-balance -**Strength 1.0 is not a unit.** Measured across the 27 non-distillation H3 LoRAs -in `models/loras/h3`, the perturbation produced at strength 1.0 spans **65×** — +**Strength 1.0 is not a unit.** In the reference non-distillation H3 corpus, +the perturbation produced at strength 1.0 spans **65×** — 0.054% of the base weights at one end, 5.24% at the other. Neither rank nor file size predicts it: a rank-128 adapter sits at 0.088% while a rank-16 one sits at 0.40%. So a strength that worked on one LoRA carries no information about the @@ -288,9 +293,27 @@ auto-balance would have used. - 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. +- DoRA, LoHa, LoKr and locon adapters merge in `auto`; `branch` rejects them + because only plain rank decompositions have the stack's runtime branch path. ## License Licensed under the [Apache License, Version 2.0](LICENSE). + +The package requires Python 3.10 or newer and a ComfyUI revision providing the +MiniMax H3 model, `QuantizedTensor`, weight adapters, and patcher wrappers used +by this node. + +## Development checks + +From the ComfyUI root, make the package and ComfyUI import roots visible before +running the tests: + +```powershell +$env:PYTHONPATH = "custom_nodes/ComfyUI-H3-PowerLoraStack;." +python -m pytest custom_nodes/ComfyUI-H3-PowerLoraStack/tests --import-mode=importlib -q +``` + +The package tests cover CPU paths and skip the mixed-device check when no CUDA +device is available. `node --check web/h3_power_lora_stack.js` validates the +frontend syntax. diff --git a/h3lora/adaln.py b/h3lora/adaln.py index 2178549..ad3cd70 100644 --- a/h3lora/adaln.py +++ b/h3lora/adaln.py @@ -61,7 +61,7 @@ TE_OUT_B = "time_embedder.proj_out.bias" MAX_RESIDUAL = 5e-3 FREQ_DIM = 256 -_grid_cache: dict[str, torch.Tensor] = {} +_grid_cache: dict[tuple[str, int, int], torch.Tensor] = {} _basis_cache: dict[tuple, tuple] = {} @@ -96,8 +96,10 @@ def find_silu_grid(extra_dirs=()) -> str: def load_silu_grid(path: str) -> torch.Tensor: - """Load the ``[grid, 2688]`` silu(t_emb) grid, cached by path.""" - cached = _grid_cache.get(path) + """Load the ``[grid, 2688]`` silu(t_emb) grid, cached by file version.""" + stat = os.stat(path) + cache_key = (path, stat.st_mtime_ns, stat.st_size) + cached = _grid_cache.get(cache_key) if cached is not None: return cached from safetensors.torch import load_file @@ -105,7 +107,10 @@ def load_silu_grid(path: str) -> torch.Tensor: 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 + for old_key in tuple(_grid_cache): + if old_key[0] == path and old_key != cache_key: + _grid_cache.pop(old_key, None) + _grid_cache[cache_key] = grid return grid @@ -123,9 +128,11 @@ def tensor_hash(tensor): def silu_temb_grid(proj_in_w, proj_in_b, proj_out_w, proj_out_b, rows=1025, freq_dim=FREQ_DIM): - t = torch.arange(rows, dtype=torch.float32) / float(rows - 1) + device = proj_in_w.device + t = torch.arange(rows, device=device, dtype=torch.float32) / float(rows - 1) half = freq_dim // 2 - freqs = torch.exp(-math.log(10000.0) * torch.arange(half, dtype=torch.float32) / half) + freqs = torch.exp(-math.log(10000.0) * torch.arange( + half, device=device, dtype=torch.float32) / half) args = t[:, None] * freqs[None] emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) hidden = F.silu(emb @ proj_in_w.float().T + proj_in_b.float()) @@ -262,13 +269,15 @@ def fit_basis(table: torch.Tensor, grid: torch.Tensor): 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)) + table_cpu = _cpu_f32(table) + grid_cpu = _cpu_f32(grid) + key = (tensor_hash(table_cpu), tensor_hash(grid_cpu)) cached = _basis_cache.get(key) if cached is not None: return cached - t = table.to(torch.float64) - s = grid.to(torch.float64) + t = table_cpu.to(torch.float64) + s = grid_cpu.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] @@ -310,8 +319,8 @@ def fit_table_to_table(source_table: torch.Tensor, target_table: torch.Tensor): Returns ``(M [k_src, k_tgt], a [k_src], residual)``. """ - src = source_table.to(torch.float64) - tgt = target_table.to(torch.float64) + src = _cpu_f32(source_table).to(torch.float64) + tgt = _cpu_f32(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]})" @@ -363,23 +372,22 @@ class AdalnContext: def _scan_grids(self, table): rows = int(table.shape[0]) - key = tensor_hash(table) - cached = _read_basis_cache(key) - if cached is not None and _accept(cached[2]): - v, c, residual, source = cached - return self._record(v, c, residual, "%s (cached)" % source) + table_key = tensor_hash(table) best = None for label, path in _iter_diffusion_safetensors(): grid = _grid_from_checkpoint(path, rows) if grid is None: continue - v, c, residual = fit_basis(table.detach().to("cpu"), grid) + key = "d%s%s" % (table_key[:8], tensor_hash(grid)[:8]) + cached = _read_basis_cache(key) + if cached is not None: + v, c, residual, _src = cached + else: + v, c, residual = fit_basis(table.detach().to("cpu"), grid) + _write_basis_cache(key, v, c, residual, label) if best is None or residual < best[2]: best = (v, c, residual, label) - if best is not None: - _write_basis_cache(key, best[0], best[1], best[2], best[3]) - return self._record(*best) - return None + return self._record(*best) if best is not None else None def _scan_tables(self, grid): grid_key = tensor_hash(grid) @@ -388,7 +396,7 @@ class AdalnContext: table = _table_from_checkpoint(path) if table is None or table.shape[0] != grid.shape[0]: continue - key = "d%s%s" % (grid_key[:8], tensor_hash(table)[:8]) + key = "d%s%s" % (tensor_hash(table)[:8], grid_key[:8]) cached = _read_basis_cache(key) if cached is not None: v, c, residual, _src = cached @@ -404,7 +412,7 @@ class AdalnContext: def basis(self, curve_table=None): """``(V, c, residual)`` relating the dense space to a curve basis.""" table = self.table if self.table is not None else curve_table - cache_key = None if table is None else id(table) + cache_key = None if table is None else tensor_hash(table) if table is not self.table and cache_key in self._curve_cache: return self._curve_cache[cache_key] if self.table is not None and (self._basis is not None or self._failed): @@ -481,36 +489,55 @@ def port_adaln_pairs(sd: dict, ctx: AdalnContext, source_table=None): # curve -> curve via the LoRA's own table, when it carries one table_map = None + table_map_failed = source_table is not None and ctx.table is not 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) + if _accept(residual): + table_map = (m, a_const) + stats["residual"] = residual + table_map_failed = False + LOG.info("H3 PowerLoraStack: adaLN table-to-table fit residual %.2e", residual) + else: + LOG.warning("H3 PowerLoraStack: adaLN table-to-table fit residual %.2e exceeds %.0e", + residual, MAX_RESIDUAL) except Exception as exc: LOG.warning("H3 PowerLoraStack: adaLN table rebase unavailable (%s)", exc) out = dict(sd) basis = None + basis_pinv = 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") + a_key = next((parts.get(suffix) for suffix in ( + "lora_A.weight", "lora_A.default.weight", "lora_A", + "lora_down.weight", "_lora.down.weight", "lora.down.weight", + "lora_linear_layer.down.weight", + ) if parts.get(suffix)), None) + b_key = next((parts.get(suffix) for suffix in ( + "lora_B.weight", "lora_B.default.weight", "lora_B", + "lora_up.weight", "_lora.up.weight", "lora.up.weight", + "lora_linear_layer.up.weight", + ) if parts.get(suffix)), None) if a_key is None or b_key is None: continue a = sd[a_key] - if a.ndim != 2: + b = sd[b_key] + if a.ndim != 2 or b.ndim != 2 or b.shape[1] != a.shape[0]: + LOG.warning("H3 PowerLoraStack: malformed adaLN pair %s, skipped", module) + for key in parts.values(): + out.pop(key, None) + stats["skipped"] += 1 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 + m, a_const = (x.to(device=a32.device) for x in 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: + elif source_dim == ctx.target_dim and not table_map_failed: stats["ok"] += 1 continue else: @@ -522,13 +549,17 @@ def port_adaln_pairs(sd: dict, ctx: AdalnContext, source_table=None): stats["skipped"] += 1 continue v, c, residual = basis # V: [2688, k], c: [2688] + v = v.to(device=a32.device) + c = c.to(device=a32.device) 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] + if basis_pinv is None: + basis_pinv = torch.linalg.pinv(v) + a_new = a32 @ basis_pinv # [r, 2688] const = a_new @ c # [r] sign = -1.0 else: diff --git a/h3lora/apply.py b/h3lora/apply.py index d65988b..c17c0da 100644 --- a/h3lora/apply.py +++ b/h3lora/apply.py @@ -3,8 +3,10 @@ from __future__ import annotations import logging +import math from collections import OrderedDict +import torch import comfy.lora import comfy.patcher_extension import comfy.utils @@ -34,6 +36,7 @@ class StackReport: self.merged = 0 self.branched = 0 self.skipped = 0 + self.rejected = 0 self.bank_bytes = 0 def add(self, line: str): @@ -125,8 +128,7 @@ def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="", video_dim = int(getattr(diffusion_model.final_layer.video_out, "out_features", 0) or 0) audio_dim = int(getattr(diffusion_model.final_layer.audio_out, "out_features", 0) or 0) hidden = int(getattr(diffusion_model, "hidden_size", 0) or 0) - pdd_bank = None - pdd_from = None + pdd_entry = None target_dim, table = adaln_mod.read_target(diffusion_model) if adaln_mode == "off": adaln_ctx = None @@ -166,15 +168,32 @@ def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="", for fallback_row, entry in enumerate(entries, start=1): row_index = int(entry.get("row", fallback_row)) name = entry.get("name") or entry["path"] - strength = float(entry.get("strength", 1.0)) + try: + strength = float(entry.get("strength", 1.0)) + except (TypeError, ValueError): + report.rejected += 1 + report.add(f"{name} rejected: strength must be numeric") + continue + if not math.isfinite(strength): + report.rejected += 1 + report.add(f"{name} rejected: strength must be finite") + continue row_schedule = schedule_mod.resolve(schedule, row_index) - lora_sd = comfy.utils.load_torch_file(entry["path"], safe_load=True) + try: + lora_sd = comfy.utils.load_torch_file(entry["path"], safe_load=True) + except Exception as exc: + report.rejected += 1 + report.add(f"{name} rejected: could not load LoRA ({exc})") + continue normalized, unmatched = keymap.normalize(lora_sd, index) bank, pdd_note = pdd_mod.peel(normalized, video_dim, audio_dim, hidden) - if bank is not None: - pdd_bank = bank - pdd_from = name + row_pdd_entry = ({ + "bank": bank, + "name": name, + "strength": strength, + "schedule": row_schedule, + } if bank is not None else None) measured = gain.measure_state_dict(normalized, name) # before porting: the port derives its bias delta as ``B @ const``, so @@ -183,8 +202,9 @@ def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="", normalized, mod_values, mod_geom) mod_note = modality_mod.describe(mod_values, mod_stats) adaln_note = "" + table_key = next((key for key in normalized if key.endswith("adaln_t_table")), None) + source_table = normalized.pop(table_key) if table_key is not None else None 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"]: @@ -196,10 +216,58 @@ def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="", if stats["residual"] is not None: adaln_note += f" (basis fit {stats['residual']:.1e})" - loaded = comfy.lora.load_lora(normalized, key_map) + passengers = [key for key in normalized if keymap.split_suffix(key)[1] is None] + unsupported_passengers = [key for key in passengers + if not key.endswith("adaln_t_table")] + if unsupported_passengers: + report.rejected += len(unsupported_passengers) + example = ", ".join(unsupported_passengers[:2]) + if len(unsupported_passengers) > 2: + example += ", ..." + report.add(f"{name} rejected: {len(unsupported_passengers)} unsupported " + f"passenger tensor(s) ({example}); no partial application") + continue + if row_pdd_entry is not None: + pdd_entry = row_pdd_entry + try: + loaded = comfy.lora.load_lora(normalized, key_map) + except Exception as exc: + report.rejected += 1 + report.add(f"{name} rejected: adapter load failed ({exc})") + continue merge: dict = dict(loaded) branched_here = 0 + skipped_here = 0 + branch_rejected_here = 0 + for weight_key, patch in loaded.items(): + if (not weight_key.endswith(".weight") or not isinstance(patch, tuple) + or len(patch) <= 1 or patch[0] not in {"diff", "set"} + or not patch[1]): + continue + candidate = patch[1][0] + module_path = weight_key[:-len(".weight")] + weight = _module_weight(patcher, module_path) + if (not torch.is_tensor(candidate) or weight is None + or tuple(candidate.shape) != tuple(weight.shape)): + LOG.warning("H3 PowerLoraStack: shape mismatch on %s, skipped", weight_key) + report.skipped += 1 + skipped_here += 1 + merge.pop(weight_key, None) + for bias_key, bias_patch in loaded.items(): + if (not bias_key.endswith(".bias") or not isinstance(bias_patch, tuple) + or len(bias_patch) <= 1 or bias_patch[0] != "diff" + or not bias_patch[1]): + continue + module_path = bias_key[:-len(".bias")] + candidate = bias_patch[1][0] + weight = _module_weight(patcher, module_path) + if (not torch.is_tensor(candidate) or weight is None + or candidate.ndim != 1 or candidate.shape[0] != weight.shape[0]): + LOG.warning("H3 PowerLoraStack: shape mismatch on %s, skipped", bias_key) + report.skipped += 1 + skipped_here += 1 + merge.pop(bias_key, None) for weight_key, adapter in loaded.items(): if not weight_key.endswith(".weight"): continue @@ -207,29 +275,33 @@ def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="", if module_path is None or not _is_plain_lora(adapter): continue weight = _module_weight(patcher, module_path) - effective_mode = "branch" if row_schedule is not None else mode - if not _branchable(weight, module_path, effective_mode): - continue up, down = adapter.weights[0], adapter.weights[1] - if up.shape[0] != weight.shape[0] or down.shape[1] != weight.shape[1]: + if (not torch.is_tensor(up) or not torch.is_tensor(down) + or up.ndim != 2 or down.ndim != 2 or up.shape[1] != down.shape[0] + or weight is None or 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 + skipped_here += 1 merge.pop(weight_key, None) continue + effective_mode = "branch" if row_schedule is not None else mode + if not _branchable(weight, module_path, effective_mode): + continue alpha = adapter.weights[2] alpha_scale = float(alpha) / down.shape[0] if alpha is not None else 1.0 bias = None - if row_schedule is not None: - bias_key = f"{module_path}.bias" - bias_patch = loaded.get(bias_key) - if (isinstance(bias_patch, tuple) and len(bias_patch) > 1 - and bias_patch[0] == "diff" and bias_patch[1]): - candidate = bias_patch[1][0] - if getattr(candidate, "ndim", None) == 1 and candidate.shape[0] == weight.shape[0]: - bias = candidate - merge.pop(bias_key, None) + bias_key = f"{module_path}.bias" + bias_patch = loaded.get(bias_key) + if (isinstance(bias_patch, tuple) and len(bias_patch) > 1 + and bias_patch[0] == "diff" and bias_patch[1]): + candidate = bias_patch[1][0] + if getattr(candidate, "ndim", None) == 1 and candidate.shape[0] == weight.shape[0]: + bias = candidate.to(torch.float32) * strength + merge.pop(bias_key, None) + if row_schedule is not None or bias is not None: per_module.setdefault(module_path, []).append( - (up, down, alpha_scale, row_schedule, bias)) + (up, down, strength * alpha_scale, row_schedule, bias)) else: # Preserve the original strength-folded contribution exactly. per_module.setdefault(module_path, []).append( @@ -237,17 +309,46 @@ def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="", merge.pop(weight_key, None) branched_here += 1 - if merge: - patcher.add_patches(merge, strength) - report.merged += len(merge) + if mode == "branch" and merge: + # Branch mode must not silently fall back to a weight patch for an + # adapter or passenger that has no runtime branch implementation. + branch_rejected_here = len(merge) + report.rejected += branch_rejected_here + merge.clear() + + accepted = patcher.add_patches(merge, strength) if merge else [] + rejected_here = len(merge) - len(accepted) + quantized_merge = [] + for patch_key in accepted: + weight_key = patch_key[0] if isinstance(patch_key, tuple) else patch_key + if not isinstance(weight_key, str) or not weight_key.endswith(".weight"): + continue + module_path = weight_key[:-len(".weight")] + if isinstance(_module_weight(patcher, module_path), QuantizedTensor): + quantized_merge.append(weight_key) + report.rejected += rejected_here + report.merged += len(accepted) report.branched += branched_here - detail = f"{len(merge)} merged, {branched_here} branched" + detail = f"{len(accepted)} merged, {branched_here} branched" if unmatched: detail += f", {len(unmatched)} unmatched" + if passengers: + detail += f", {len(passengers)} ignored" + if skipped_here: + detail += f", {skipped_here} skipped" + if rejected_here: + detail += f", {rejected_here} rejected" + if branch_rejected_here: + detail += f", {branch_rejected_here} rejected in branch mode" + if quantized_merge: + detail += f", {len(quantized_merge)} quantized merge" if pdd_note: detail += f", {pdd_note}" report.add(f"{name} @ {strength:g}: {detail}{adaln_note}{mod_note}") + if quantized_merge: + report.add(" ! quantized targets were merged because no exact runtime " + "adapter path was available") if row_schedule is not None: arrow = "\u2192" report.add( @@ -256,20 +357,21 @@ def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="", f"({row_schedule.domain} {row_schedule.start_percent:g}\u2013" f"{row_schedule.end_percent:g}%)" ) - if merge: + if accepted: report.add( - f" ! {len(merge)} patches on {name} cannot be scheduled " + f" ! {len(accepted)} patches on {name} cannot be scheduled " f"(merged at {strength:.2f})" ) if measured.get("rel"): # what this LoRA actually does to the weights, so a strength that is # far off the calibrated unit is visible without the UI button note = f" rel dW {measured['rel'] * 100:.3f}%" - suggested = measured["factor"] - if abs(strength - suggested) > 0.1 * max(suggested, 1e-6): - note += f" (auto-balance would use {suggested:.2f})" + factor = measured["factor"] + suggested = strength * factor + if abs(strength - suggested) > 0.1 * max(abs(suggested), 1e-6): + note += f" (auto-balance would use {suggested:.2f}, factor {factor:.2f})" report.add(note) - if not merge and not branched_here: + if not accepted and not branched_here and row_pdd_entry is None: report.add(f" ! {name} matched no layers on this model") if per_module: @@ -289,9 +391,11 @@ def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="", report.add(f"branch bank: {len(per_module)} layers, " f"{report.bank_bytes / (1024 ** 2):.0f} MB") - if pdd_bank is not None: - pdd_mod.attach(patcher, pdd_bank) - report.add(f"PDD: {pdd_bank['n']} heads from {pdd_from} " - f"(simple @ 8 steps, shifts 12/3)") + if pdd_entry is not None: + pdd_bank = pdd_entry["bank"] + pdd_mod.attach(patcher, pdd_bank, pdd_entry["strength"], pdd_entry["schedule"]) + report.add(f"PDD: {pdd_bank['n']} heads from {pdd_entry['name']} " + f"(strength {pdd_entry['strength']:g}, " + f"{pdd_mod.bank_bytes(pdd_bank) / (1024 ** 2):.0f} MB, shifts 12/3)") return patcher, report diff --git a/h3lora/branch.py b/h3lora/branch.py index 90355c3..bc88b19 100644 --- a/h3lora/branch.py +++ b/h3lora/branch.py @@ -160,7 +160,8 @@ class ScheduleController: ]) if fused.biases: bias_scales[name] = torch.tensor( - [strength(schedule) for _bias, schedule in fused.biases], + [strength(schedule) if schedule is not None else 1.0 + for _bias, schedule in fused.biases], dtype=self.dtype, device=device, ) diff --git a/h3lora/detect.py b/h3lora/detect.py index 503b7ea..6968873 100644 --- a/h3lora/detect.py +++ b/h3lora/detect.py @@ -24,7 +24,12 @@ _CONVENTIONS = ( def read_header(path: str) -> dict: with open(path, "rb") as f: - length = struct.unpack(" 64 * 1024 * 1024: + raise ValueError("safetensors header exceeds 64 MiB") return json.loads(f.read(length)) @@ -79,7 +84,11 @@ def inspect(path: str) -> dict: has_alpha = True continue shape = v.get("shape") or [] - is_down = k.endswith(("lora_A.weight", "lora_down.weight")) + is_down = k.endswith(( + "lora_A.weight", "lora_A.default.weight", "lora_A", + "lora_down.weight", "_lora.down.weight", ".lora.down.weight", + "lora_linear_layer.down.weight", + )) if is_down and len(shape) == 2: ranks.add(int(shape[0])) if "adaln" in k: @@ -97,7 +106,7 @@ def inspect(path: str) -> dict: 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) + if not re.search(r"(lora_A|lora_B|lora_down|lora_up|_lora|lora_linear_layer|reshape_weight|alpha|diff|dora|hada|lokr|oft)", k) )[:8] pdd_n = 0 for k, v in header.items(): diff --git a/h3lora/gain.py b/h3lora/gain.py index 595da01..3ce8c1b 100644 --- a/h3lora/gain.py +++ b/h3lora/gain.py @@ -55,7 +55,7 @@ _BASE_RMS = { "other": 0.1620, } -# Median ``rel`` over the 27 non-distillation H3 LoRAs in models/loras/h3. +# Median ``rel`` over the reference non-distillation H3 LoRA corpus. # This is the anchor that makes strength 1.0 mean "the usual amount of change" # rather than "whatever this trainer happened to emit". Chosen as the median # rather than a hand-picked target so the calibration agrees with trainer @@ -69,8 +69,10 @@ REFERENCE_REL = 0.00292 MIN_FACTOR = 0.05 MAX_FACTOR = 1.0 -_DOWN = (".lora_down.weight", ".lora_A.weight", ".lora_A.default.weight", ".lora_A") -_UP = (".lora_up.weight", ".lora_B.weight", ".lora_B.default.weight", ".lora_B") +_DOWN = (".lora_down.weight", ".lora_A.weight", ".lora_A.default.weight", ".lora_A", + "_lora.down.weight", ".lora.down.weight", ".lora_linear_layer.down.weight") +_UP = (".lora_up.weight", ".lora_B.weight", ".lora_B.default.weight", ".lora_B", + "_lora.up.weight", ".lora.up.weight", ".lora_linear_layer.up.weight") _cache: dict[tuple, dict] = {} @@ -117,6 +119,7 @@ class _Accum: self.layers = 0 self.ranks: set[int] = set() self.kinds: set[str] = set() + self.incomplete = False def add(self, module: str, delta2: float, out: int, inp: int, kind: str, rank=None): rms = _BASE_RMS.get(_group(module), _BASE_RMS["other"]) @@ -135,10 +138,11 @@ class _Accum: "layers": self.layers, "ranks": sorted(self.ranks), "kinds": sorted(self.kinds), + "complete": not self.incomplete, # exact-duplicate fingerprint: two files with the same layer count # and the same total delta energy are the same adapter "fingerprint": None if rel is None else f"{self.layers}:{self.delta2:.10g}", - "factor": factor_for(rel), + "factor": factor_for(rel) if not self.incomplete else 1.0, } @@ -164,7 +168,12 @@ def _pairs(keys): elif suffix == ".alpha": slot = "alpha" elif suffix in (".lokr_w1", ".lokr_w2", ".lokr_w1_a", ".lokr_w1_b", - ".lokr_w2_a", ".lokr_w2_b", ".diff"): + ".lokr_w2_a", ".lokr_w2_b", ".lokr_t2", ".diff", ".dora_scale"): + slot = suffix[1:] + elif suffix in (".lora_mid.weight", ".diff_b", ".set_weight", ".reshape_weight", + ".hada_w1_a", ".hada_w1_b", ".hada_w2_a", ".hada_w2_b", + ".hada_t1", ".hada_t2", ".oft_blocks", ".boft_blocks", + ".rescale", ".w_norm", ".b_norm"): slot = suffix[1:] if slot: mods.setdefault(body, {})[slot] = key @@ -174,9 +183,16 @@ def _pairs(keys): def _measure(mods, get, name: str) -> dict: """``mods`` from :func:`_pairs`, ``get(key) -> tensor``.""" acc = _Accum() + measurable = { + "down", "up", "alpha", "lokr_w1", "lokr_w2", "lokr_w1_a", "lokr_w1_b", + "lokr_w2_a", "lokr_w2_b", "diff", + } for module, slots in mods.items(): if _skip(module): continue + if set(slots) - measurable or "lokr_t2" in slots: + acc.incomplete = True + continue try: alpha_key = slots.get("alpha") if "down" in slots and "up" in slots: @@ -202,7 +218,10 @@ def _measure(mods, get, name: str) -> dict: elif "diff" in slots: d = _f32(get(slots["diff"])) acc.add(module, float(d.norm() ** 2), d.shape[0], d.shape[1], "diff") + else: + acc.incomplete = True except Exception as exc: # one odd layer must not + acc.incomplete = True LOG.debug("gain: skipped %s (%s)", module, exc) # sink the file return acc.result(name) diff --git a/h3lora/keymap.py b/h3lora/keymap.py index 44edc89..90b7238 100644 --- a/h3lora/keymap.py +++ b/h3lora/keymap.py @@ -24,6 +24,9 @@ import re # suffixes comfy.lora.load_lora understands, longest first so that # ``.lora_down.weight`` is not truncated to ``.weight`` _LORA_SUFFIXES = ( + ".lora_linear_layer.down.weight", ".lora_linear_layer.up.weight", + "_lora.down.weight", "_lora.up.weight", + ".lora.down.weight", ".lora.up.weight", ".lora_down.weight", ".lora_up.weight", ".lora_A.weight", ".lora_B.weight", ".lora_A.default.weight", ".lora_B.default.weight", @@ -32,7 +35,8 @@ _LORA_SUFFIXES = ( ".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", + ".lokr_t2", ".reshape_weight", ".oft_blocks", ".boft_blocks", ".rescale", + ".w_norm", ".b_norm", ) _STRIP_PREFIXES = ( diff --git a/h3lora/modality.py b/h3lora/modality.py index 0f7c93f..535b240 100644 --- a/h3lora/modality.py +++ b/h3lora/modality.py @@ -41,6 +41,7 @@ not everything it eventually reaches. from __future__ import annotations import logging +import math import torch @@ -51,7 +52,11 @@ TAGS = ("video", "text", "audio") # lora_B spellings; ``.diff`` is a full-weight replacement, which slices the same # way. ``.set_weight`` is an absolute value, not a delta, so it is never scaled. -_ROW_KEYS = (".lora_B.weight", ".lora_up.weight", ".diff") +_ROW_KEYS = ( + ".lora_B.weight", ".lora_B.default.weight", ".lora_B", + ".lora_up.weight", "_lora.up.weight", ".lora.up.weight", + ".lora_linear_layer.up.weight", ".diff", ".diff_b", +) _ADALN_SUFFIX = "adaln_proj.linear" @@ -91,6 +96,8 @@ def normalize_scales(scales) -> tuple: values = tuple(float(v) for v in scales) if len(values) != len(TAGS): return None + if not all(math.isfinite(value) for value in values): + raise ValueError("modality scales must be finite") return values @@ -132,7 +139,9 @@ def apply_to_state_dict(sd: dict, values, geom): continue stats["present"] += 1 shape = getattr(tensor, "shape", None) - if shape is None or len(shape) != 2 or shape[0] != rows: + is_bias = suffix == ".diff_b" + expected_rank = 1 if is_bias else 2 + if shape is None or len(shape) != expected_rank or shape[0] != rows: # final_layer.adaln_proj (one modality) and any future geometry land # here and are left exactly as they were stats["mismatched"] += 1 diff --git a/h3lora/nodes.py b/h3lora/nodes.py index 1a1b4db..afc0e7f 100644 --- a/h3lora/nodes.py +++ b/h3lora/nodes.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import math import folder_paths @@ -54,14 +55,22 @@ def _resolve_lora(name: str): if path: return path wanted = name.replace("\\", "/").lower() + matches = [] 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) + if normalized == wanted: + matches.append(candidate) + elif "/" not in wanted and normalized.endswith("/" + wanted): + matches.append(candidate) + if len(matches) == 1: + return folder_paths.get_full_path("loras", matches[0]) + if len(matches) > 1: + LOG.warning("H3 Power LoRA Stack: ambiguous lora name %r (%d matches)", + name, len(matches)) return None -def _collect(kwargs): +def _collect(kwargs, issues=None): """Pull enabled rows out of the dynamic ``lora_N`` inputs, in UI order.""" rows = [] for key, value in kwargs.items(): @@ -80,7 +89,20 @@ def _collect(kwargs): for order, value in rows: if not value.get("on", True): continue - strength = float(value.get("strength", 1.0)) + try: + strength = float(value.get("strength", 1.0)) + except (TypeError, ValueError): + message = f"row {order}: strength is not numeric" + LOG.warning("H3 Power LoRA Stack: %s", message) + if issues is not None: + issues.append(message) + continue + if not math.isfinite(strength): + message = f"row {order}: strength must be finite" + LOG.warning("H3 Power LoRA Stack: %s", message) + if issues is not None: + issues.append(message) + continue if strength == 0.0: continue name = value.get("lora") @@ -88,7 +110,10 @@ def _collect(kwargs): 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) + message = f"row {order}: could not resolve lora {name!r}" + LOG.warning("H3 Power LoRA Stack: %s", message) + if issues is not None: + issues.append(message) continue entries.append({"name": name, "path": path, "strength": strength, "row": order}) return entries @@ -151,9 +176,10 @@ class H3PowerLoraStack: if model is None: raise ValueError("H3 Power LoRA Stack: no model connected") - entries = _collect(kwargs) + issues = [] + entries = _collect(kwargs, issues) if not entries: - return (model, "no LoRAs enabled") + return (model, "no LoRAs enabled" + ("\n" + "\n".join(issues) if issues else "")) patcher, report = apply_mod.apply_stack( model, entries, mode=quantized_layers, adaln_mode=adaln_port, diff --git a/h3lora/pdd.py b/h3lora/pdd.py index 8a61037..f26921d 100644 --- a/h3lora/pdd.py +++ b/h3lora/pdd.py @@ -6,20 +6,26 @@ onto the native ``[96, hidden]`` head and crashes. Core PR 15908 teaches ``FinalLayer`` to blend those rows; this module does the same inside the Power LoRA Stack so those files load without a core change. -The bank is kept beside the model (object patch + DIFFUSION_MODEL wrapper). -The native head is not resized, so a run without the bank is unchanged. +The bank is kept beside the model (object patch + DIFFUSION_MODEL wrapper) and +is registered as an additional model so ComfyUI can move and account for it. +The native head is not resized, and PDD strength interpolates from the native +head to the blended bank. """ from __future__ import annotations import logging +import threading import torch import torch.nn as nn import torch.nn.functional as F +import comfy.model_management +import comfy.model_patcher from comfy.ldm.minimax.model import _mod_row, time_shift_sigma from comfy.patcher_extension import WrappersMP +from .schedule import progress_from LOG = logging.getLogger("h3.powerlorastack") @@ -83,7 +89,11 @@ def peel(sd, video_dim, audio_dim, hidden): aw, aw_key = _pop_first(sd, _AUDIO_WEIGHT) ab, ab_key = _pop_first(sd, _AUDIO_BIAS) if vw is None and aw is None: - return None, "" + if vb is not None: + sd[vb_key] = vb + if ab is not None: + sd[ab_key] = ab + return None, "PDD bias without head ignored" if vb is not None or ab is not None else "" def _restore(): if vw is not None: @@ -97,12 +107,15 @@ def peel(sd, video_dim, audio_dim, hidden): video = _as_bank(vw, video_dim) if vw is not None else None audio = _as_bank(aw, audio_dim) if aw is not None else None + notes = [] if video is None and vw is not None: LOG.warning("H3 PowerLoraStack: ignored %s shape %s (native video head is %d)", vw_key, tuple(vw.shape), video_dim) + notes.append("video weight dropped") if audio is None and aw is not None: LOG.warning("H3 PowerLoraStack: ignored %s shape %s (native audio head is %d)", aw_key, tuple(aw.shape), audio_dim) + notes.append("audio weight dropped") if video is None and audio is None: return None, "PDD heads dropped (shape mismatch)" @@ -125,11 +138,17 @@ def peel(sd, video_dim, audio_dim, hidden): video_b = None if vb is not None: packed = _as_bank(vb, video_dim) - video_b = packed[1] if packed is not None and packed[0] == n else None + video_b = (packed[1] if video is not None and packed is not None + and packed[0] == n else None) + if video_b is None: + notes.append("video bias dropped") audio_b = None if ab is not None: packed = _as_bank(ab, audio_dim) - audio_b = packed[1] if packed is not None and packed[0] == n else None + audio_b = (packed[1] if audio is not None and packed is not None + and packed[0] == n else None) + if audio_b is None: + notes.append("audio bias dropped") bank = { "n": n, @@ -140,7 +159,10 @@ def peel(sd, video_dim, audio_dim, hidden): "audio_w": audio_w, "audio_b": audio_b, } - return bank, f"PDD bank n={n}" + note = f"PDD bank n={n}" + if notes: + note += " (" + ", ".join(notes) + ")" + return bank, note def _blend(weight, bias, h, n, start, stop, flow_shift): @@ -159,15 +181,21 @@ def _blend(weight, bias, h, n, start, stop, flow_shift): def interval(sigma, sigma_next, n, video_shift): start, stop = (round(float(1.0 - time_shift_sigma(s, video_shift, 1.0)) * n) for s in (sigma, sigma_next)) - return start, max(stop, start + 1) + start = max(0, min(n - 1, start)) + stop = max(start + 1, min(n, stop)) + return start, stop class PddFinalLayer(nn.Module): def __init__(self, inner, bank, ctx): super().__init__() - self.inner = inner - self.bank = bank - self.ctx = ctx + # The bank belongs to the additional-model patcher, so keep that + # reference out of this module. Register the native head children at + # their original paths so the object patch preserves the model state + # dict and its main-patcher ownership. + object.__setattr__(self, "inner", inner) + object.__setattr__(self, "bank", bank) + object.__setattr__(self, "ctx", ctx) self.norm = inner.norm self.adaln_proj = inner.adaln_proj self.video_out = inner.video_out @@ -177,11 +205,14 @@ class PddFinalLayer(nn.Module): inner = self.inner bank = self.bank ctx = self.ctx - n = bank["n"] + n = bank.n sample_sigmas = ctx.get("sample_sigmas") sigma = ctx.get("sigma") if sample_sigmas is None or sigma is None: raise ValueError("MiniMax H3 PDD heads need the sampler's sigma schedule") + sample_sigmas = torch.as_tensor(sample_sigmas, device=sigma.device).flatten() + if sample_sigmas.numel() == 0: + raise ValueError("MiniMax H3 PDD heads need a non-empty sigma schedule") shift, scale = inner.adaln_proj(t_emb) @@ -193,33 +224,105 @@ class PddFinalLayer(nn.Module): i = int((sample_sigmas - sigma).abs().argmin()) sigma_next = sample_sigmas[min(i + 1, sample_sigmas.shape[0] - 1)] shifts = ctx.get("shifts") or (12.0, 3.0) + strength = float(ctx.get("strength") or 0.0) start, stop = interval(sigma, sigma_next, n, shifts[0]) vh = mod(video_seg) ah = mod(audio_seg) - video = _blend(bank["video_w"], bank["video_b"], vh, n, start, stop, shifts[0]) - if bank["audio_w"] is None: - audio = inner.audio_out(ah) + native_video = inner.video_out(vh) + if bank.video_w is None: + video = native_video else: - audio = _blend(bank["audio_w"], bank["audio_b"], ah, n, start, stop, shifts[1]) + pdd_video = _blend(bank.video_w, bank.video_b, vh, n, start, stop, shifts[0]) + video = native_video + strength * (pdd_video - native_video) + native_audio = inner.audio_out(ah) + if bank.audio_w is None: + audio = native_audio + else: + pdd_audio = _blend(bank.audio_w, bank.audio_b, ah, n, start, stop, shifts[1]) + audio = native_audio + strength * (pdd_audio - native_audio) return video, audio -def attach(patcher, bank): +class PddBank(nn.Module): + """Movable, accountable storage for a PDD head bank.""" + + def __init__(self, bank): + super().__init__() + self.n = bank["n"] + self.video_w = nn.Parameter(bank["video_w"], requires_grad=False) if bank["video_w"] is not None else None + self.video_b = nn.Parameter(bank["video_b"], requires_grad=False) if bank["video_b"] is not None else None + self.audio_w = nn.Parameter(bank["audio_w"], requires_grad=False) if bank["audio_w"] is not None else None + self.audio_b = nn.Parameter(bank["audio_b"], requires_grad=False) if bank["audio_b"] is not None else None + + +class PddState: + def __init__(self, strength, schedule): + self.default_strength = float(strength) + self.schedule = schedule + self.local = threading.local() + + def set(self, sigma, sample_sigmas, shifts, strength): + self.local.values = { + "sigma": sigma, + "sample_sigmas": sample_sigmas, + "shifts": shifts, + "strength": strength, + } + + def get(self, key): + return getattr(self.local, "values", {}).get(key) + + def clear(self): + self.local.values = {} + + +def bank_bytes(bank): + return sum( + tensor.numel() * tensor.element_size() + for key in ("video_w", "video_b", "audio_w", "audio_b") + if (tensor := bank.get(key)) is not None + ) + + +def attach(patcher, bank, strength=1.0, schedule=None): fl = patcher.get_model_object("diffusion_model.final_layer") - ctx = {"sigma": None, "sample_sigmas": None, "shifts": (12.0, 3.0)} - wrapped = PddFinalLayer(fl, bank, ctx) + bank_model = PddBank(bank) + bank_patcher = comfy.model_patcher.ModelPatcher( + bank_model, + load_device=comfy.model_management.get_torch_device(), + offload_device=comfy.model_management.unet_offload_device(), + ) + n = 0 + while patcher.get_additional_models_with_key(f"h3_power_lora_pdd_bank_{n}"): + n += 1 + patcher.set_additional_models(f"h3_power_lora_pdd_bank_{n}", [bank_patcher]) + + ctx = PddState(strength, schedule) + wrapped = PddFinalLayer(fl, bank_model, ctx) patcher.add_object_patch("diffusion_model.final_layer", wrapped) def wrapper(executor, x, timestep, context, transformer_options=None, **kwargs): transformer_options = transformer_options or {} sigma = (timestep.flatten()[0] / 1000.0).float().clamp(min=1e-6) - ctx["sigma"] = sigma - ctx["sample_sigmas"] = transformer_options.get("sample_sigmas") - ctx["shifts"] = ( + sample_sigmas = transformer_options.get("sample_sigmas") + if sample_sigmas is None: + sample_sigmas = transformer_options.get("sigmas") + effective_strength = ctx.default_strength + if ctx.schedule is not None: + schedule_options = transformer_options + if "sample_sigmas" not in schedule_options and sample_sigmas is not None: + schedule_options = dict(transformer_options) + schedule_options["sample_sigmas"] = sample_sigmas + progress = progress_from(timestep, schedule_options, ctx.schedule.domain) + effective_strength *= ctx.schedule.evaluate(progress) + ctx.set(sigma, sample_sigmas, ( float(transformer_options.get("minimax_h3_sigma_shift_video", 12.0)), float(transformer_options.get("minimax_h3_sigma_shift_audio", 3.0)), - ) - return executor(x, timestep, context, transformer_options, **kwargs) + ), effective_strength) + try: + return executor(x, timestep, context, transformer_options, **kwargs) + finally: + ctx.clear() n = 0 while patcher.get_wrappers(WrappersMP.DIFFUSION_MODEL, f"h3_power_lora_pdd_{n}"): diff --git a/h3lora/server.py b/h3lora/server.py index 3b39ac5..eba2ce9 100644 --- a/h3lora/server.py +++ b/h3lora/server.py @@ -47,6 +47,8 @@ def measure_names(names): entry["note"] = "unreadable" elif not info.get("rel"): entry["note"] = "no measurable layers" + elif not info.get("complete", True): + entry["note"] = "calibration is partial for unsupported adapter tensors" else: fp = info.get("fingerprint") first = seen.get(fp) if fp else None diff --git a/pyproject.toml b/pyproject.toml index 113abfd..231e74a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,13 +2,14 @@ name = "comfyui-h3-powerlorastack" description = "Stacked multi-LoRA loader for MiniMax H3 across bf16, fp8, int8 and w4a8 weights." version = "1.0.0" +requires-python = ">=3.10" license = { text = "Apache-2.0" } dependencies = [] [project.urls] -Repository = "https://github.com/doggeddalle/ComfyUI-H3-PowerLoraStack" +Repository = "https://github.com/cicalooo/ComfyUI-H3-PowerLoraStack" [tool.comfy] -PublisherId = "doggeddalle" +PublisherId = "cicalooo" DisplayName = "MiniMax H3 Power LoRA Stack" Icon = "" diff --git a/tests/test_adaln.py b/tests/test_adaln.py new file mode 100644 index 0000000..f99c1af --- /dev/null +++ b/tests/test_adaln.py @@ -0,0 +1,74 @@ +import unittest + +import torch + +from h3lora.adaln import AdalnContext, fit_basis, port_adaln_pairs, sidecar_basis + + +class AdalnSidecarTests(unittest.TestCase): + def test_fit_basis_recovers_affine(self): + g, k, d = 32, 4, 16 + table = torch.randn(g, k) + v = torch.randn(d, k) + c = torch.randn(d) + grid = c.unsqueeze(0) + table @ v.T + v2, c2, residual = fit_basis(table, grid) + self.assertLess(residual, 1e-5) + self.assertTrue(torch.allclose(v2, v, atol=1e-4)) + self.assertTrue(torch.allclose(c2, c, atol=1e-4)) + + def test_sidecar_v_c_preferred(self): + v = torch.randn(16, 4) + c = torch.randn(16) + got = sidecar_basis({"adaln_basis": v, "adaln_mean": c}, torch.randn(8, 4)) + self.assertIsNotNone(got) + self.assertEqual(got[3], "checkpoint sidecars") + self.assertEqual(got[2], 0.0) + self.assertTrue(torch.equal(got[0], v.float())) + + def test_context_uses_sidecars_without_grid_file(self): + g, k, d = 16, 3, 8 + table = torch.randn(g, k) + v = torch.randn(d, k) + c = torch.randn(d) + ctx = AdalnContext(k, table, grid_path="", sidecars={"adaln_basis": v, "adaln_mean": c}) + basis = ctx.basis() + self.assertIsNotNone(basis) + self.assertEqual(ctx.source, "checkpoint sidecars") + a = torch.randn(2, d) + b = torch.randn(6, 2) + sd = { + "blocks.0.adaln_proj.linear.lora_A.weight": a, + "blocks.0.adaln_proj.linear.lora_B.weight": b, + } + out, stats = port_adaln_pairs(sd, ctx) + self.assertEqual(stats["ported"], 1) + self.assertEqual(out["blocks.0.adaln_proj.linear.lora_A.weight"].shape[1], k) + + def test_fit_basis_normalizes_mixed_devices(self): + if not torch.cuda.is_available(): + self.skipTest("CUDA is not available") + table = torch.randn(16, 3, device="cuda") + grid = torch.randn(16, 8) + _v, _c, residual = fit_basis(table, grid) + self.assertTrue(torch.isfinite(torch.tensor(residual))) + + def test_default_suffix_is_ported(self): + table = torch.randn(16, 3) + ctx = AdalnContext( + 3, + table, + sidecars={"adaln_basis": torch.randn(8, 3), "adaln_mean": torch.randn(8)}, + ) + sd = { + "blocks.0.adaln_proj.linear.lora_A.default.weight": torch.randn(2, 8), + "blocks.0.adaln_proj.linear.lora_B.default.weight": torch.randn(6, 2), + } + out, stats = port_adaln_pairs(sd, ctx) + self.assertEqual(stats["ported"], 1) + self.assertIn("blocks.0.adaln_proj.linear.lora_A.weight", out) + self.assertIn("blocks.0.adaln_proj.linear.lora_B.weight", out) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_modality.py b/tests/test_modality.py new file mode 100644 index 0000000..2767d3d --- /dev/null +++ b/tests/test_modality.py @@ -0,0 +1,23 @@ +import unittest + +import torch + +from h3lora.modality import apply_to_state_dict + + +class ModalityTests(unittest.TestCase): + def test_default_b_and_diff_b_are_scaled(self): + expand, hidden = 2, 4 + rows = 3 * expand * hidden + sd = { + "diffusion_model.blocks.0.adaln_proj.linear.lora_B.default.weight": torch.ones(rows, 2), + "diffusion_model.blocks.0.adaln_proj.linear.diff_b": torch.ones(rows), + } + out, stats = apply_to_state_dict(sd, (0.0, 1.0, 1.0), (expand, 3, hidden)) + self.assertEqual(stats["scaled"], 2) + self.assertTrue(torch.equal(out[next(iter(sd))][: expand * hidden], torch.zeros(expand * hidden, 2))) + self.assertTrue(torch.equal(out[next(reversed(sd))][: expand * hidden], torch.zeros(expand * hidden))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pdd.py b/tests/test_pdd.py new file mode 100644 index 0000000..e9ee2b4 --- /dev/null +++ b/tests/test_pdd.py @@ -0,0 +1,118 @@ +import unittest +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from h3lora.pdd import PddBank, PddFinalLayer, PddState, _as_bank, _blend, interval, peel + + +class PeelTests(unittest.TestCase): + def test_stacked_set_weight(self): + sd = { + "diffusion_model.final_layer.video_out.set_weight": torch.randn(3072, 5376), + "diffusion_model.final_layer.video_out.set_bias": torch.randn(3072), + "diffusion_model.final_layer.audio_out.set_weight": torch.randn(1024, 5376), + "diffusion_model.blocks.0.attn.qkv_proj.lora_A.weight": torch.randn(8, 5376), + } + bank, note = peel(sd, 96, 32, 5376) + self.assertIsNotNone(bank) + self.assertEqual(bank["n"], 32) + self.assertEqual(bank["video_w"].shape, (3072, 5376)) + self.assertEqual(bank["audio_w"].shape, (1024, 5376)) + self.assertIn("n=32", note) + self.assertNotIn("set_weight", sd) + self.assertIn("diffusion_model.blocks.0.attn.qkv_proj.lora_A.weight", sd) + + def test_proj_out_3d(self): + sd = { + "proj_out.weight": torch.randn(32, 96, 5376), + "proj_out.bias": torch.randn(32, 96), + "audio_proj_out.weight": torch.randn(32, 32, 5376), + "audio_proj_out.bias": torch.randn(32, 32), + } + bank, note = peel(sd, 96, 32, 5376) + self.assertEqual(bank["n"], 32) + self.assertEqual(bank["video_w"].shape, (3072, 5376)) + self.assertEqual(bank["video_b"].shape, (3072,)) + self.assertEqual(bank["audio_w"].shape, (1024, 5376)) + self.assertFalse(sd) + + def test_native_head_left_alone(self): + w = torch.randn(96, 5376) + sd = {"diffusion_model.final_layer.video_out.set_weight": w} + bank, note = peel(sd, 96, 32, 5376) + self.assertIsNone(bank) + self.assertEqual(note, "") + self.assertIs(sd["diffusion_model.final_layer.video_out.set_weight"], w) + + def test_garbage_dropped(self): + sd = {"diffusion_model.final_layer.video_out.set_weight": torch.randn(100, 5376)} + bank, note = peel(sd, 96, 32, 5376) + self.assertIsNone(bank) + self.assertIn("dropped", note) + self.assertFalse(sd) + + +class BlendTests(unittest.TestCase): + def test_as_bank_3d_and_2d(self): + t = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3) + n, flat = _as_bank(t, 4) + self.assertEqual(n, 2) + self.assertEqual(flat.shape, (8, 3)) + n2, flat2 = _as_bank(flat, 4) + self.assertEqual(n2, 2) + self.assertTrue(torch.equal(flat, flat2)) + + def test_single_interval_is_that_head(self): + n, o, hidden = 4, 2, 3 + weight = torch.zeros(n * o, hidden) + for i in range(n): + weight[i * o:(i + 1) * o] = float(i + 1) + h = torch.ones(5, hidden) + out = _blend(weight, None, h, n, 2, 3, 1.0) + self.assertEqual(out.shape, (5, o)) + self.assertTrue(torch.allclose(out, torch.full((5, o), 3.0 * hidden))) + + def test_interval_monotonic(self): + start, stop = interval(torch.tensor(1.0), torch.tensor(0.0), 32, 12.0) + self.assertGreaterEqual(stop, start + 1) + self.assertGreaterEqual(start, 0) + self.assertLessEqual(stop, 32) + + def test_partial_audio_bank_falls_back_to_native_video(self): + sd = { + "final_layer.audio_out.set_weight": torch.randn(2 * 2, 3), + } + bank, note = peel(sd, 4, 2, 3) + self.assertIsNotNone(bank) + self.assertIsNone(bank["video_w"]) + self.assertIn("n=2", note) + + def test_strength_interpolates_native_and_pdd_heads(self): + class StaticAdaln(nn.Module): + def forward(self, _t): + return torch.zeros(2, 3), torch.zeros(2, 3) + + inner = SimpleNamespace( + norm=nn.Identity(), + adaln_proj=StaticAdaln(), + video_out=nn.Linear(3, 2), + audio_out=nn.Linear(3, 1), + ) + bank, _ = peel({ + "final_layer.video_out.set_weight": torch.ones(2 * 2, 3), + }, 2, 1, 3) + state = PddState(0.5, None) + state.set(torch.tensor(1.0), torch.tensor([1.0, 0.0]), (12.0, 3.0), 0.5) + layer = PddFinalLayer(inner, PddBank(bank), state) + video, audio = layer( + torch.ones(4, 3), torch.zeros(1, 3), (0, 2, 0), (2, 4, 1)) + native = inner.video_out(torch.ones(2, 3)) + pdd = torch.full_like(native, 3.0) + self.assertTrue(torch.allclose(video, native + 0.5 * (pdd - native))) + self.assertEqual(audio.shape, (2, 1)) + + +if __name__ == "__main__": + unittest.main() diff --git a/web/h3_power_lora_stack.js b/web/h3_power_lora_stack.js index feda02a..87a7843 100644 --- a/web/h3_power_lora_stack.js +++ b/web/h3_power_lora_stack.js @@ -366,6 +366,7 @@ async function showLoraMenu(node, widget, event, { removeOnCancel = false } = {} if (removeOnCancel && !committed) removeLoraWidget(node, widget); }; const commit = (value) => { + const wasBalanced = isBalanced(node); committed = true; const next = { ...widget.value, lora: value }; // the balance factor belonged to the file that was here before, so drop the @@ -382,7 +383,7 @@ async function showLoraMenu(node, widget, event, { removeOnCancel = false } = {} widget.value = next; node.setDirtyCanvas(true, true); close(); - if (isBalanced(node)) applyBalance(node); + if (wasBalanced) applyBalance(node); }; const onOutside = (e) => { if (!root.contains(e.target)) close(); @@ -543,6 +544,12 @@ async function applyBalance(node, { force = false } = {}) { data = await fetchBalance([...new Set(rows.map((w) => w.value.lora))]); } catch (err) { console.error("[H3PowerLoraStack] auto-balance failed", err); + const button = (node.widgets || []).find((w) => w.h3Role === "balance"); + if (button) { + button.name = "⚖ Balance failed"; + button.tooltip = err?.message || "Could not measure the selected LoRAs"; + node.setDirtyCanvas(true, true); + } return; } const results = data?.results ?? {}; @@ -550,12 +557,16 @@ async function applyBalance(node, { force = false } = {}) { const result = results[widget.value.lora]; if (!result) continue; const manual = - widget.value.manual !== undefined ? widget.value.manual : widget.value.strength; - const factor = result.factor ?? 1; + widget.value.autoApplied && widget.value.manual !== undefined + ? widget.value.manual + : widget.value.strength; + const rawFactor = Number(result.factor ?? 1); + if (!Number.isFinite(rawFactor) || rawFactor < 0) continue; + const factor = Math.min(1, rawFactor); widget.value = { ...widget.value, manual, - strength: round2(manual * factor), + strength: manual * factor, factor, rel: result.rel ?? null, note: result.note ?? "", @@ -666,7 +677,6 @@ app.registerExtension({ // serialization drops them from both halves of that round trip. for (const button of [add, balance, restore]) button.serialize = false; - fetchLoraList(); // warm the cache so the first picker opens instantly resize(this); };