From f38d0842327e56915bf5ba9abd462094d2787029 Mon Sep 17 00:00:00 2001 From: doggeddalle <153059103+cicalooo@users.noreply.github.com> Date: Thu, 27 Aug 2026 19:43:07 +1000 Subject: [PATCH] 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 | 31 +++++ README.md | 27 +++- h3lora/adaln.py | 330 ++++++++++++++++++++++++++++++++++++++++------- h3lora/apply.py | 32 ++++- h3lora/detect.py | 10 ++ h3lora/nodes.py | 4 + h3lora/pdd.py | 228 ++++++++++++++++++++++++++++++++ 7 files changed, 610 insertions(+), 52 deletions(-) create mode 100644 .gitignore create mode 100644 h3lora/pdd.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..636b2b8 --- /dev/null +++ b/.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 diff --git a/README.md b/README.md index 378ad49..9dc184c 100644 --- a/README.md +++ b/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 diff --git a/h3lora/adaln.py b/h3lora/adaln.py index 4aa4f3a..2178549 100644 --- a/h3lora/adaln.py +++ b/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(" (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): diff --git a/h3lora/apply.py b/h3lora/apply.py index 115ff7b..d65988b 100644 --- a/h3lora/apply.py +++ b/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 diff --git a/h3lora/detect.py b/h3lora/detect.py index d47fc81..503b7ea 100644 --- a/h3lora/detect.py +++ b/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 diff --git a/h3lora/nodes.py b/h3lora/nodes.py index 2672138..1a1b4db 100644 --- a/h3lora/nodes.py +++ b/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"): diff --git a/h3lora/pdd.py b/h3lora/pdd.py new file mode 100644 index 0000000..8a61037 --- /dev/null +++ b/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