435 lines
19 KiB
Python
435 lines
19 KiB
Python
"""Stack orchestration: normalise, port, partition, apply."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import math
|
|
from collections import OrderedDict
|
|
|
|
import torch
|
|
import comfy.lora
|
|
import comfy.patcher_extension
|
|
import comfy.utils
|
|
from comfy.quant_ops import QuantizedTensor
|
|
from comfy.weight_adapter import LoRAAdapter
|
|
|
|
from . import adaln as adaln_mod
|
|
from . import branch as branch_mod
|
|
from . import gain
|
|
from . import keymap
|
|
from . import modality as modality_mod
|
|
from . import pdd as pdd_mod
|
|
from . import schedule as schedule_mod
|
|
from . import sidecars as sidecar_mod
|
|
|
|
LOG = logging.getLogger("h3.powerlorastack")
|
|
sidecar_mod.ensure_load_hook()
|
|
|
|
# Layers whose forward is bypassed by a fused kernel, per weight layout. A
|
|
# runtime branch on these would be silently dropped, so they must merge.
|
|
_FORWARD_BYPASSED = {"TensorWiseINT8Layout": ("mlp.fc2",)}
|
|
|
|
|
|
class StackReport:
|
|
"""Human-readable account of what happened, surfaced on the node."""
|
|
|
|
def __init__(self):
|
|
self.lines: list[str] = []
|
|
self.merged = 0
|
|
self.branched = 0
|
|
self.skipped = 0
|
|
self.rejected = 0
|
|
self.bank_bytes = 0
|
|
|
|
def add(self, line: str):
|
|
self.lines.append(line)
|
|
|
|
def text(self) -> str:
|
|
return "\n".join(self.lines)
|
|
|
|
|
|
def _module_weight(model_patcher, module_path: str):
|
|
try:
|
|
module = model_patcher.get_model_object(module_path)
|
|
except Exception:
|
|
return None
|
|
return getattr(module, "weight", None)
|
|
|
|
|
|
def _is_plain_lora(adapter) -> bool:
|
|
"""True for a bare rank-decomposition with no locon/dora/reshape extras."""
|
|
return (
|
|
isinstance(adapter, LoRAAdapter)
|
|
and adapter.weights[3] is None # locon mid
|
|
and adapter.weights[4] is None # dora scale
|
|
and adapter.weights[5] is None # reshape / pad
|
|
and adapter.weights[0].ndim == 2
|
|
)
|
|
|
|
|
|
def _branchable(weight, module_path: str, mode: str) -> bool:
|
|
if weight is None or mode == "merge":
|
|
return False
|
|
if mode == "auto" and not isinstance(weight, QuantizedTensor):
|
|
return False # unquantized merges are exact and free at runtime
|
|
layout = getattr(weight, "_layout_cls", None)
|
|
for suffix in _FORWARD_BYPASSED.get(layout, ()):
|
|
if module_path.endswith(suffix):
|
|
return False
|
|
return True
|
|
|
|
|
|
def detect_quantization(model_patcher) -> str:
|
|
"""Report the DiT's dominant weight format for the node's status line."""
|
|
layouts: dict[str, int] = {}
|
|
plain = 0
|
|
try:
|
|
diffusion_model = model_patcher.get_model_object("diffusion_model")
|
|
except Exception:
|
|
return "unknown"
|
|
for module in diffusion_model.modules():
|
|
weight = getattr(module, "weight", None)
|
|
if weight is None:
|
|
continue
|
|
if isinstance(weight, QuantizedTensor):
|
|
layout = getattr(weight, "_layout_cls", "quantized")
|
|
layouts[layout] = layouts.get(layout, 0) + 1
|
|
else:
|
|
plain += 1
|
|
if not layouts:
|
|
return "unquantized"
|
|
parts = [f"{k.replace('TensorCore', '').replace('TensorWise', '').replace('Layout', '')}"
|
|
f" x{v}" for k, v in sorted(layouts.items(), key=lambda kv: -kv[1])]
|
|
if plain:
|
|
parts.append(f"plain x{plain}")
|
|
return ", ".join(parts)
|
|
|
|
|
|
def apply_stack(model, entries, mode="auto", adaln_mode="auto", grid_path="",
|
|
report: StackReport | None = None, modality=None, schedule=None):
|
|
"""Apply a list of ``{'path', 'name', 'strength'}`` LoRAs to an H3 model.
|
|
|
|
``mode`` is ``auto`` (branch quantized layers, merge the rest), ``merge``
|
|
(stock behaviour) or ``branch`` (never touch a weight).
|
|
|
|
``adaln_mode`` is ``auto`` (rebase this stack and any adaLN adapters already
|
|
on ``model``), ``strip`` (drop mismatched pairs) or ``off``.
|
|
|
|
``modality`` optionally scales each LoRA's adaLN modulation per modality;
|
|
see :mod:`h3lora.modality`.
|
|
"""
|
|
report = report or StackReport()
|
|
patcher = model.clone()
|
|
|
|
model_sd = patcher.model.state_dict()
|
|
index = keymap.build_module_index(model_sd.keys())
|
|
key_map = {k[: -len(".weight")]: k for k in model_sd
|
|
if k.startswith("diffusion_model.") and k.endswith(".weight")}
|
|
|
|
try:
|
|
diffusion_model = patcher.get_model_object("diffusion_model")
|
|
except Exception as exc:
|
|
raise ValueError("H3 Power LoRA Stack requires a MiniMax H3 model") from exc
|
|
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_entry = 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,
|
|
sidecars=sidecar_mod.collect(patcher, diffusion_model),
|
|
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_mode != "off":
|
|
upstream = adaln_mod.upstream_adaln_fix_warning(model)
|
|
if upstream:
|
|
report.add(upstream)
|
|
LOG.warning("H3 PowerLoraStack:\n%s", upstream)
|
|
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 adaln_ctx is not None:
|
|
try:
|
|
patcher, attached = adaln_mod.port_attached_patches(
|
|
patcher, adaln_ctx, mode=adaln_mode)
|
|
except Exception as exc:
|
|
LOG.exception("H3 PowerLoraStack: incoming adaLN port failed")
|
|
report.add(f" ! incoming adaLN port failed ({exc}); left upstream patches")
|
|
else:
|
|
if attached["ported"] or attached["stripped"] or attached["unportable"]:
|
|
line = (f" incoming adaLN: ported {attached['ported']}, "
|
|
f"stripped {attached['stripped']}, "
|
|
f"unportable {attached['unportable']} "
|
|
f"across {attached['keys']} key(s)")
|
|
if attached["residual"] is not None:
|
|
line += f" (basis fit {attached['residual']:.1e})"
|
|
report.add(line)
|
|
for note in attached["notes"][:8]:
|
|
report.add(f" ! {note}")
|
|
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")
|
|
# the header carries the reason; keep the per-LoRA notes from blaming it
|
|
# on the LoRAs, which are not at fault here
|
|
mod_values = None
|
|
|
|
per_module: "OrderedDict[str, list]" = OrderedDict()
|
|
compute_dtype = model.model_dtype()
|
|
|
|
for fallback_row, entry in enumerate(entries, start=1):
|
|
row_index = int(entry.get("row", fallback_row))
|
|
name = entry.get("name") or entry["path"]
|
|
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)
|
|
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)
|
|
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
|
|
# scaling B's rows here carries through to the emitted .diff_b
|
|
normalized, mod_stats = modality_mod.apply_to_state_dict(
|
|
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:
|
|
try:
|
|
normalized, stats = adaln_mod.port_adaln_pairs(
|
|
normalized, adaln_ctx, source_table=source_table,
|
|
mode=adaln_mode)
|
|
except Exception as exc:
|
|
LOG.exception("H3 PowerLoraStack: adaLN port failed for %s", name)
|
|
report.add(f"{name}: adaLN port failed ({exc}); pairs left as-is")
|
|
stats = {"ported": 0, "skipped": 0, "ok": 0, "rebased": 0, "residual": None}
|
|
if stats["ported"]:
|
|
adaln_note = f", adaLN ported x{stats['ported']}"
|
|
if stats["rebased"]:
|
|
adaln_note += f", adaLN rebased x{stats['rebased']}"
|
|
if stats["skipped"]:
|
|
adaln_note += f", adaLN dropped x{stats['skipped']}"
|
|
if stats["residual"] is not None:
|
|
adaln_note += f" (basis fit {stats['residual']:.1e})"
|
|
|
|
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
|
|
module_path = weight_key[: -len(".weight")] if weight_key.endswith(".weight") else None
|
|
if module_path is None or not _is_plain_lora(adapter):
|
|
continue
|
|
weight = _module_weight(patcher, module_path)
|
|
up, down = adapter.weights[0], adapter.weights[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
|
|
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, strength * alpha_scale, row_schedule, bias))
|
|
else:
|
|
# Preserve the original strength-folded contribution exactly.
|
|
per_module.setdefault(module_path, []).append(
|
|
(up, down, strength * alpha_scale))
|
|
merge.pop(weight_key, None)
|
|
branched_here += 1
|
|
|
|
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(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(
|
|
f" sched: {row_schedule.start_strength:.2f} {arrow} "
|
|
f"{row_schedule.end_strength:.2f} {row_schedule.curve} "
|
|
f"({row_schedule.domain} {row_schedule.start_percent:g}\u2013"
|
|
f"{row_schedule.end_percent:g}%)"
|
|
)
|
|
if accepted:
|
|
report.add(
|
|
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}%"
|
|
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 accepted and not branched_here and row_pdd_entry is None:
|
|
report.add(f" ! {name} matched no layers on this model")
|
|
|
|
if per_module:
|
|
report.bank_bytes = branch_mod.bank_bytes(per_module, compute_dtype)
|
|
# a chain of stack nodes must not overwrite each other's bank
|
|
n = 0
|
|
while patcher.get_additional_models_with_key(f"h3_power_lora_bank_{n}"):
|
|
n += 1
|
|
tag = f"h3_power_lora_bank_{n}"
|
|
_count, controller = branch_mod.attach(patcher, per_module, compute_dtype, tag)
|
|
if controller is not None:
|
|
patcher.add_wrapper_with_key(
|
|
comfy.patcher_extension.WrappersMP.APPLY_MODEL,
|
|
f"h3_lora_schedule_{n}",
|
|
controller,
|
|
)
|
|
report.add(f"branch bank: {len(per_module)} layers, "
|
|
f"{report.bank_bytes / (1024 ** 2):.0f} MB")
|
|
|
|
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)")
|
|
|
|
adaln_mod.mark_stack_adaln(patcher)
|
|
return patcher, report
|