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

@@ -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

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.

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:

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

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,
                 )

h3lora/detect.py

@@ -24,7 +24,12 @@ _CONVENTIONS = (
 
 def read_header(path: str) -> dict:
     with open(path, "rb") as f:
-        length = struct.unpack("<Q", f.read(8))[0]
+        raw_length = f.read(8)
+        if len(raw_length) != 8:
+            raise ValueError("truncated safetensors header length")
+        length = struct.unpack("<Q", raw_length)[0]
+        if length > 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():

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)
 

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 = (

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

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,

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}"):

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

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 = ""

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()

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()

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()

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);
     };