Add PDD support and improve adaLN grid handling

Introduce PDD/Acc head handling and a robust adaLN grid/basis pipeline.

- Add h3lora/pdd.py: peel Acc/PDD output-head banks from LoRAs, keep banks aside, and blend rows at sample time (PddFinalLayer + attach wrapper) to avoid crashes.
- Refactor h3lora/adaln.py: generate/load silu grids from time_embedder, safetensors, or live embedder; add hashing, basis cache, sidecar support, residual acceptance threshold, and scanning of diffusion_model safetensors.
- Integrate PDD into apply.py, detect.py, nodes.py and README; add .gitignore entries.

.gitignore

@@ -0,0 +1,31 @@
+__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

README.md

@@ -80,14 +80,17 @@ trained against the checkpoint you are running.
 
 #### The silu grid
 
-Dense↔curve porting needs `h3_silu_temb_grid.safetensors`. It is searched for in
-`models/h3_adaln/`, `models/loras/`, `models/diffusion_models/` and one level
-under `custom_nodes/`. Table-to-table rebasing does not need it. Set
-`adaln_port` to `off` to disable porting entirely.
+Dense↔curve porting prefers a basis already baked into the loaded checkpoint
+(`adaln_basis` + `adaln_mean`, or `silu_t_emb_grid`). If those keys are absent
+it uses a live `time_embedder` on a dense base, then `h3_silu_temb_grid.safetensors`
+(searched in `models/h3_adaln/`, `models/loras/`, `models/diffusion_models/` and
+one level under `custom_nodes/`), then a scored scan of other H3 checkpoints.
+Table-to-table rebasing does not need a grid. Set `adaln_port` to `off` to
+disable porting entirely.
 
-With a grid from a different build the fit bottoms out at ~1.7e-3 relative,
-because the 7th and 8th curve directions are near-degenerate (σ₇ ≈ σ₈) and so
-differ between bakes. That is still ~6× below the int8 quantization floor.
+A grid from a different build bottoms out at ~1.7e-3 relative because the 7th
+and 8th curve directions are near-degenerate (σ₇ ≈ σ₈). Fits worse than 5e-3
+are rejected and the next source is tried.
 
 ### 3. Key conventions
 
@@ -105,6 +108,16 @@ two tokens:
 
 Verified against all 37 H3 LoRAs in `models/loras/h3`: **zero unmatched keys.**
 
+### 4. Acc / PDD output heads
+
+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.
+
 ## Stacking
 
 Multiple LoRAs on the same layer are fused into a single pair by concatenating

h3lora/adaln.py

@@ -39,18 +39,30 @@ its own weights.
 
 from __future__ import annotations
 
+import hashlib
+import json
 import logging
+import math
 import os
+import struct
 
 import torch
+import torch.nn.functional as F
 
 LOG = logging.getLogger(__name__)
 
 GRID_FILENAME = "h3_silu_temb_grid.safetensors"
 _GRID_TENSOR = "silu_t_emb_grid"
+BASIS_SUBDIR = "h3_adaln"
+TE_IN_W = "time_embedder.proj_in.weight"
+TE_IN_B = "time_embedder.proj_in.bias"
+TE_OUT_W = "time_embedder.proj_out.weight"
+TE_OUT_B = "time_embedder.proj_out.bias"
+MAX_RESIDUAL = 5e-3
+FREQ_DIM = 256
 
 _grid_cache: dict[str, torch.Tensor] = {}
-_basis_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}
+_basis_cache: dict[tuple, tuple] = {}
 
 
 def find_silu_grid(extra_dirs=()) -> str:
@@ -97,6 +109,148 @@ def load_silu_grid(path: str) -> torch.Tensor:
     return grid
 
 
+def _cpu_f32(t):
+    return t.detach().to(device="cpu", dtype=torch.float32).contiguous()
+
+
+def tensor_hash(tensor):
+    t = _cpu_f32(tensor)
+    h = hashlib.sha256()
+    h.update(str(tuple(t.shape)).encode())
+    h.update(t.numpy().tobytes())
+    return h.hexdigest()[:16]
+
+
+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)
+    half = freq_dim // 2
+    freqs = torch.exp(-math.log(10000.0) * torch.arange(half, 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())
+    temb = hidden @ proj_out_w.float().T + proj_out_b.float()
+    return F.silu(temb)
+
+
+def grid_from_time_embedder(time_embedder, rows=1025):
+    try:
+        proj_in, proj_out = time_embedder.proj_in, time_embedder.proj_out
+        tensors = (proj_in.weight, proj_in.bias, proj_out.weight, proj_out.bias)
+        if any(t is None or t.device.type == "meta" for t in tensors):
+            return None
+        freq_dim = int(getattr(time_embedder, "freq_dim", FREQ_DIM))
+        return silu_temb_grid(*tensors, rows=rows, freq_dim=freq_dim).clone()
+    except Exception:
+        return None
+
+
+def _safetensors_header(path):
+    with open(path, "rb") as fh:
+        raw = fh.read(8)
+        if len(raw) < 8:
+            return None
+        size = struct.unpack("<Q", raw)[0]
+        if size <= 0 or size > (64 << 20):
+            return None
+        return json.loads(fh.read(size))
+
+
+def _time_embedder_prefix(header):
+    if not header:
+        return None
+    if any(k.endswith("adaln_t_table") for k in header):
+        return None
+    for key in header:
+        if key.endswith(TE_OUT_W):
+            prefix = key[:-len(TE_OUT_W)]
+            if all(prefix + n in header for n in (TE_IN_W, TE_IN_B, TE_OUT_B)):
+                return prefix
+    return None
+
+
+def _load_named(path, names):
+    from safetensors import safe_open
+    with safe_open(path, framework="pt") as f:
+        return [_cpu_f32(f.get_tensor(n)) for n in names]
+
+
+def _basis_dir():
+    try:
+        import folder_paths
+        return os.path.join(folder_paths.models_dir, BASIS_SUBDIR)
+    except Exception:
+        return ""
+
+
+def _cache_path(key):
+    d = _basis_dir()
+    return os.path.join(d, "basis_%s.safetensors" % key) if d else ""
+
+
+def _read_basis_cache(key):
+    path = _cache_path(key)
+    if not path or not os.path.isfile(path):
+        return None
+    try:
+        from safetensors import safe_open
+        with safe_open(path, framework="pt") as f:
+            meta = f.metadata() or {}
+            return (_cpu_f32(f.get_tensor("V")), _cpu_f32(f.get_tensor("c")),
+                    float(meta.get("residual", "nan")),
+                    meta.get("source", "cache"))
+    except Exception:
+        return None
+
+
+def _write_basis_cache(key, v, c, residual, source):
+    path = _cache_path(key)
+    if not path:
+        return
+    try:
+        from safetensors.torch import save_file
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        save_file({"V": v.contiguous(), "c": c.contiguous()}, path,
+                  metadata={"residual": repr(residual), "source": source, "table_hash": key})
+    except Exception:
+        pass
+
+
+def _iter_diffusion_safetensors():
+    try:
+        import folder_paths
+    except Exception:
+        return
+    for name in folder_paths.get_filename_list("diffusion_models"):
+        if not name.lower().endswith(".safetensors"):
+            continue
+        path = folder_paths.get_full_path("diffusion_models", name)
+        if path:
+            yield os.path.basename(name), path
+
+
+def _grid_from_checkpoint(path, rows):
+    try:
+        prefix = _time_embedder_prefix(_safetensors_header(path))
+        if prefix is None:
+            return None
+        tensors = _load_named(path, [prefix + n for n in (TE_IN_W, TE_IN_B, TE_OUT_W, TE_OUT_B)])
+        return silu_temb_grid(*tensors, rows=rows)
+    except Exception:
+        return None
+
+
+def _table_from_checkpoint(path):
+    try:
+        header = _safetensors_header(path)
+        key = next((k for k in (header or ()) if k.endswith("adaln_t_table")), None)
+        if key is None:
+            return None
+        return _load_named(path, [key])[0]
+    except Exception:
+        return None
+
+
 def fit_basis(table: torch.Tensor, grid: torch.Tensor):
     """Least-squares fit of ``grid ~= 1 c^T + table V^T``.
 
@@ -121,25 +275,31 @@ def fit_basis(table: torch.Tensor, grid: torch.Tensor):
     v = solution[1:].transpose(0, 1).contiguous()                                    # [2688, k]
 
     residual = float((design @ solution - s).norm() / s.norm().clamp(min=1e-12))
-    # The last two curve directions are near-degenerate (sigma_7 ~ sigma_8), so
-    # they differ between bakes and a grid from a different build only agrees to
-    # ~2e-3.  That is still ~6x below the int8 quantization floor, and vastly
-    # better than dropping adaLN entirely, so only a much larger residual is
-    # worth complaining about.
-    if residual > 2e-2:
-        LOG.warning(
-            "H3 PowerLoraStack: adaLN basis fit residual %.2e is high - the silu grid "
-            "and the checkpoint's adaln_t_table probably come from different bakes",
-            residual,
-        )
-    else:
-        LOG.info("H3 PowerLoraStack: adaLN basis fit residual %.2e", residual)
-
     out = (v.to(torch.float32), c.to(torch.float32), residual)
     _basis_cache[key] = out
     return out
 
 
+def sidecar_basis(sidecars, table=None):
+    if not sidecars:
+        return None
+    v, c = sidecars.get("adaln_basis"), sidecars.get("adaln_mean")
+    if torch.is_tensor(v) and torch.is_tensor(c):
+        v, c = _cpu_f32(v), _cpu_f32(c)
+        if v.ndim == 2 and c.ndim == 1 and v.shape[0] == c.shape[0]:
+            if table is None or v.shape[1] == table.shape[1]:
+                return v, c, 0.0, "checkpoint sidecars"
+    grid = sidecars.get("silu_t_emb_grid")
+    if torch.is_tensor(grid) and table is not None:
+        v, c, residual = fit_basis(table, grid)
+        return v, c, residual, "checkpoint silu_t_emb_grid"
+    return None
+
+
+def _accept(residual):
+    return residual == residual and residual <= MAX_RESIDUAL
+
+
 def fit_table_to_table(source_table: torch.Tensor, target_table: torch.Tensor):
     """Map a source curve basis onto a target curve basis.
 
@@ -167,52 +327,134 @@ def fit_table_to_table(source_table: torch.Tensor, target_table: torch.Tensor):
 class AdalnContext:
     """Everything needed to move adaLN LoRA pairs between the two bases."""
 
-    def __init__(self, target_dim: int, table=None, grid_path: str = ""):
-        self.target_dim = int(target_dim)
+    def __init__(self, target_dim: int, table=None, grid_path: str = "",
+                 sidecars=None, time_embedder=None):
+        self.target_dim = int(target_dim) if target_dim is not None else 0
         self.table = table
         self.grid_path = grid_path
+        self.sidecars = sidecars or {}
+        self.time_embedder = time_embedder
         self.residual = None
+        self.source = ""
         self._basis = None
         self._failed = False
+        self._curve_cache = {}
 
     @property
     def is_curve(self) -> bool:
         return self.table is not None
 
-    def basis(self, curve_table=None):
-        """``(V, c, residual)`` relating the dense space to a curve basis.
+    def _record(self, v, c, residual, source):
+        if not _accept(residual):
+            LOG.warning(
+                "H3 PowerLoraStack: adaLN basis from %s residual %.2e exceeds %.0e",
+                source, residual, MAX_RESIDUAL,
+            )
+            return None
+        LOG.info("H3 PowerLoraStack: adaLN basis from %s, residual %.2e", source, residual)
+        self._basis = (v, c, residual)
+        self.residual = residual
+        self.source = source
+        return self._basis
 
-        The fit needs whichever side of the conversion is the *curve* side.  On
-        a pruned target that is the checkpoint's own table; on a dense target
-        converting a curve-trained LoRA it has to be the table the LoRA shipped,
-        because nothing else records which bake it was trained against.
-        """
+    def _fit_grid(self, table, grid, source):
+        v, c, residual = fit_basis(table.detach().to("cpu"), grid)
+        return self._record(v, c, residual, source)
+
+    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)
+        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)
+            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
+
+    def _scan_tables(self, grid):
+        grid_key = tensor_hash(grid)
+        best = None
+        for label, path in _iter_diffusion_safetensors():
+            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])
+            cached = _read_basis_cache(key)
+            if cached is not None:
+                v, c, residual, _src = cached
+            else:
+                v, c, residual = fit_basis(table, 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 None:
+            return None
+        return self._record(*best)
+
+    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
-        if table is None:
+        cache_key = None if table is None else id(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):
+            return self._basis
+        if self._failed and table is None:
+            return None
+
+        got = None
+        if table is not None and table is self.table:
+            got = sidecar_basis(self.sidecars, table)
+            if got is not None:
+                got = self._record(*got)
+
+        if got is None and table is not None:
+            live = grid_from_time_embedder(self.time_embedder, rows=int(table.shape[0])) if self.time_embedder is not None else None
+            if live is not None:
+                got = self._fit_grid(table, live, "live time_embedder")
+
+        if got is None and table is not None and self.grid_path:
+            try:
+                got = self._fit_grid(table, load_silu_grid(self.grid_path), self.grid_path)
+            except Exception as exc:
+                LOG.warning("H3 PowerLoraStack: could not load silu grid (%s)", exc)
+
+        if got is None and table is not None:
+            got = self._scan_grids(table)
+
+        if got is None and table is None:
+            live = grid_from_time_embedder(self.time_embedder) if self.time_embedder is not None else None
+            if live is not None:
+                got = sidecar_basis(self.sidecars, None)
+                if got is not None:
+                    got = self._record(*got)
+                if got is None:
+                    got = self._scan_tables(live)
+
+        if got is None:
             if not self._failed:
                 LOG.warning(
-                    "H3 PowerLoraStack: cannot rebase adaLN onto a dense checkpoint - "
-                    "the LoRA does not carry the adaln_t_table of the curve bake it "
-                    "was trained against"
+                    "H3 PowerLoraStack: adaLN porting unavailable - no baked basis, "
+                    "silu grid, live time embedder, or matching diffusion_models bake"
                 )
                 self._failed = True
+            if table is not self.table and table is not None:
+                self._curve_cache[cache_key] = None
             return None
-        if self._basis is not None or self._failed:
-            return self._basis
-        try:
-            if not self.grid_path:
-                raise ValueError(
-                    f"no {GRID_FILENAME} found; put it in models/h3_adaln/ to enable "
-                    "adaLN porting"
-                )
-            grid = load_silu_grid(self.grid_path)
-            self._basis = fit_basis(table.detach().to("cpu"), grid)
-            self.residual = self._basis[2]
-        except Exception as exc:
-            LOG.warning("H3 PowerLoraStack: adaLN porting unavailable (%s)", exc)
-            self._failed = True
-            self._basis = None
-        return self._basis
+
+        if table is not self.table and table is not None:
+            self._curve_cache[cache_key] = got
+        return got
 
 
 def port_adaln_pairs(sd: dict, ctx: AdalnContext, source_table=None):

h3lora/apply.py

@@ -16,6 +16,7 @@ from . import branch as branch_mod
 from . import gain
 from . import keymap
 from . import modality as modality_mod
+from . import pdd as pdd_mod
 from . import schedule as schedule_mod
 
 LOG = logging.getLogger("h3.powerlorastack")
@@ -121,19 +122,37 @@ def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="",
         diffusion_model = patcher.get_model_object("diffusion_model")
     except Exception as exc:
         raise ValueError("H3 Power LoRA Stack requires a MiniMax H3 model") from exc
+    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
     target_dim, table = adaln_mod.read_target(diffusion_model)
     if adaln_mode == "off":
         adaln_ctx = None
     else:
         if not grid_path:
             grid_path = adaln_mod.find_silu_grid()
-        adaln_ctx = adaln_mod.AdalnContext(target_dim, table, grid_path)
+        cfg = getattr(getattr(patcher, "model", None), "model_config", None)
+        sidecars = getattr(cfg, "adaln_sidecars", None) if cfg is not None else None
+        adaln_ctx = adaln_mod.AdalnContext(
+            target_dim, table, grid_path,
+            sidecars=sidecars,
+            time_embedder=getattr(diffusion_model, "time_embedder", None),
+        )
 
     mod_values = modality_mod.normalize_scales(modality)
     mod_geom = None if modality_mod.is_identity(mod_values) else modality_mod.geometry(diffusion_model)
 
     report.add(f"base: {detect_quantization(patcher)}")
     report.add(f"adaLN: {'curve' if table is not None else 'dense'} (input dim {target_dim})")
+    if adaln_ctx is not None and table is not None:
+        adaln_ctx.basis()
+        if adaln_ctx.source:
+            line = f"  basis: {adaln_ctx.source}"
+            if adaln_ctx.residual is not None:
+                line += f" (residual {adaln_ctx.residual:.1e})"
+            report.add(line)
     if not modality_mod.is_identity(mod_values) and mod_geom is None:
         report.add("  ! adaLN modality control requested but this model's adaLN "
                    "does not split into the expected modalities - ignored")
@@ -152,6 +171,10 @@ def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="",
         lora_sd = comfy.utils.load_torch_file(entry["path"], safe_load=True)
 
         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
         measured = gain.measure_state_dict(normalized, name)
 
         # before porting: the port derives its bias delta as ``B @ const``, so
@@ -222,6 +245,8 @@ def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="",
         detail = f"{len(merge)} merged, {branched_here} branched"
         if unmatched:
             detail += f", {len(unmatched)} unmatched"
+        if pdd_note:
+            detail += f", {pdd_note}"
         report.add(f"{name} @ {strength:g}: {detail}{adaln_note}{mod_note}")
         if row_schedule is not None:
             arrow = "\u2192"
@@ -264,4 +289,9 @@ 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)")
+
     return patcher, report

h3lora/detect.py

@@ -99,6 +99,16 @@ def inspect(path: str) -> dict:
         k for k in keys
         if not re.search(r"(lora_A|lora_B|lora_down|lora_up|alpha|diff|dora|hada|lokr|oft)", k)
     )[:8]
+    pdd_n = 0
+    for k, v in header.items():
+        if not k.endswith(("video_out.set_weight", "proj_out.weight")):
+            continue
+        shape = v.get("shape") or []
+        if len(shape) == 3 and shape[1] == 96:
+            pdd_n = int(shape[0])
+        elif len(shape) >= 1 and shape[0] % 96 == 0 and shape[0] > 96:
+            pdd_n = int(shape[0] // 96)
+    info["pdd_heads"] = pdd_n
     _cache[key] = info
     return info
 

h3lora/nodes.py

@@ -104,6 +104,8 @@ class H3PowerLoraStack:
       branch instead of a lossy dequantize/requantize merge;
     * adaLN basis mismatch -- LoRAs trained against a dense checkpoint are
       rebased onto a pruned checkpoint's curve, instead of being skipped;
+    * Acc/PDD head banks -- ``video_out.set_weight`` of shape ``[N*96, hidden]``
+      is peeled off before the stock ``set`` path and blended per sampler step;
     * key conventions -- ai-toolkit, kohya, lycoris, peft and bare-prefix
       layouts all resolve against the model's own key set.
     """
@@ -300,6 +302,8 @@ class H3LoraInspector:
             f"adaLN     : {'/'.join(str(d) for d in info['adaln_dims']) or 'none'}",
             f"modules   : {', '.join(info['modules'])}",
         ]
+        if info.get("pdd_heads"):
+            lines.append(f"PDD       : {info['pdd_heads']} output heads")
         if info.get("base_model"):
             lines.append(f"base      : {info['base_model']}")
         if info.get("passenger"):

h3lora/pdd.py

@@ -0,0 +1,228 @@
+"""Peel Acc/PDD output-head banks off H3 LoRAs and blend them at sample time.
+
+alibaba-pai Acc LoRAs store a 32-interval bank as ``final_layer.video_out.set_weight``
+of shape ``[N * 96, hidden]`` (and the audio twin). Stock Comfy ``copy_``s that
+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.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from comfy.ldm.minimax.model import _mod_row, time_shift_sigma
+from comfy.patcher_extension import WrappersMP
+
+LOG = logging.getLogger("h3.powerlorastack")
+
+_VIDEO_WEIGHT = (
+    "diffusion_model.final_layer.video_out.set_weight",
+    "final_layer.video_out.set_weight",
+    "proj_out.weight",
+)
+_VIDEO_BIAS = (
+    "diffusion_model.final_layer.video_out.set_bias",
+    "final_layer.video_out.set_bias",
+    "proj_out.bias",
+)
+_AUDIO_WEIGHT = (
+    "diffusion_model.final_layer.audio_out.set_weight",
+    "final_layer.audio_out.set_weight",
+    "audio_proj_out.weight",
+)
+_AUDIO_BIAS = (
+    "diffusion_model.final_layer.audio_out.set_bias",
+    "final_layer.audio_out.set_bias",
+    "audio_proj_out.bias",
+)
+
+
+def _pop_first(sd, names):
+    for name in names:
+        if name in sd:
+            return sd.pop(name), name
+    return None, None
+
+
+def _as_bank(tensor, out_features):
+    """``[N, O, I]`` or ``[N*O, I]`` / ``[N*O]`` -> ``(N, flat)`` or None."""
+    if tensor is None or out_features <= 0:
+        return None
+    t = tensor.detach().float().contiguous()
+    if t.ndim == 3 and t.shape[1] == out_features:
+        n = t.shape[0]
+        return n, t.reshape(n * out_features, *t.shape[2:])
+    if t.ndim == 2 and t.shape[1] == out_features:
+        n = t.shape[0]
+        return n, t.reshape(n * out_features)
+    if t.ndim == 2 and t.shape[0] % out_features == 0:
+        n = t.shape[0] // out_features
+        return n, t
+    if t.ndim == 1 and t.shape[0] % out_features == 0:
+        n = t.shape[0] // out_features
+        return n, t
+    return None
+
+
+def peel(sd, video_dim, audio_dim, hidden):
+    """Pull a PDD bank out of a LoRA state dict. Returns ``(bank, note)``.
+
+    ``bank`` is None when the file has no Acc heads. Matching keys that are
+    not a whole number of native heads are dropped so stock ``set`` cannot crash.
+    """
+    vw, vw_key = _pop_first(sd, _VIDEO_WEIGHT)
+    vb, vb_key = _pop_first(sd, _VIDEO_BIAS)
+    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, ""
+
+    def _restore():
+        if vw is not None:
+            sd[vw_key] = vw
+        if vb is not None:
+            sd[vb_key] = vb
+        if aw is not None:
+            sd[aw_key] = aw
+        if ab is not None:
+            sd[ab_key] = ab
+
+    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
+    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)
+    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)
+    if video is None and audio is None:
+        return None, "PDD heads dropped (shape mismatch)"
+
+    n = video[0] if video is not None else audio[0]
+    if video is not None and video[0] != n:
+        return None, f"PDD video n={video[0]} != audio n={n}"
+    if audio is not None and audio[0] != n:
+        return None, f"PDD audio n={audio[0]} != video n={n}"
+    if n <= 1:
+        _restore()
+        return None, ""
+
+    video_w = video[1] if video is not None else None
+    audio_w = audio[1] if audio is not None else None
+    if video_w is not None and video_w.shape[-1] != hidden:
+        return None, f"PDD video in_features {video_w.shape[-1]} != {hidden}"
+    if audio_w is not None and audio_w.shape[-1] != hidden:
+        return None, f"PDD audio in_features {audio_w.shape[-1]} != {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
+    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
+
+    bank = {
+        "n": n,
+        "video_dim": video_dim,
+        "audio_dim": audio_dim,
+        "video_w": video_w,
+        "video_b": video_b,
+        "audio_w": audio_w,
+        "audio_b": audio_b,
+    }
+    return bank, f"PDD bank n={n}"
+
+
+def _blend(weight, bias, h, n, start, stop, flow_shift):
+    grid = torch.linspace(1.0, 0.0, n + 1, dtype=torch.float64)
+    dt = (1.0 - flow_shift * grid / (1.0 + (flow_shift - 1.0) * grid)).diff()[start:stop]
+    w = (dt / dt.sum()).to(device=h.device, dtype=h.dtype)
+    rows = weight.to(device=h.device, dtype=h.dtype).reshape(n, -1, weight.shape[-1])[start:stop]
+    out_w = torch.einsum("n,noi->oi", w, rows)
+    out_b = None
+    if bias is not None:
+        b_rows = bias.to(device=h.device, dtype=h.dtype).reshape(n, -1)[start:stop]
+        out_b = torch.einsum("n,no->o", w, b_rows)
+    return F.linear(h, out_w, out_b)
+
+
+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)
+
+
+class PddFinalLayer(nn.Module):
+    def __init__(self, inner, bank, ctx):
+        super().__init__()
+        self.inner = inner
+        self.bank = bank
+        self.ctx = ctx
+        self.norm = inner.norm
+        self.adaln_proj = inner.adaln_proj
+        self.video_out = inner.video_out
+        self.audio_out = inner.audio_out
+
+    def forward(self, x, t_emb, video_seg, audio_seg):
+        inner = self.inner
+        bank = self.bank
+        ctx = self.ctx
+        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")
+
+        shift, scale = inner.adaln_proj(t_emb)
+
+        def mod(seg):
+            a, b, row = seg
+            return (inner.norm(x[a:b]) * (1.0 + _mod_row(scale, row, scale.dtype))
+                    + _mod_row(shift, row, shift.dtype)).to(torch.float32)
+
+        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)
+        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)
+        else:
+            audio = _blend(bank["audio_w"], bank["audio_b"], ah, n, start, stop, shifts[1])
+        return video, audio
+
+
+def attach(patcher, bank):
+    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)
+    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"] = (
+            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)
+
+    n = 0
+    while patcher.get_wrappers(WrappersMP.DIFFUSION_MODEL, f"h3_power_lora_pdd_{n}"):
+        n += 1
+    patcher.add_wrapper_with_key(WrappersMP.DIFFUSION_MODEL, f"h3_power_lora_pdd_{n}", wrapper)
+    return n