doggeddalle
4ef4de0a5d
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.
133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
"""Cheap safetensors introspection for MiniMax H3 LoRA files.
|
|
|
|
Reads only the header, so scanning a 2.5 GB LoRA costs a few kilobytes of IO.
|
|
Used both for the node's diagnostic report and for deciding, before anything is
|
|
loaded, whether a file needs adaLN porting.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import struct
|
|
|
|
_cache: dict[tuple, dict] = {}
|
|
|
|
_CONVENTIONS = (
|
|
("kohya", re.compile(r"^lora_unet_")),
|
|
("lycoris", re.compile(r"^lycoris_")),
|
|
("peft", re.compile(r"^(base_model\.model\.|transformer\.)")),
|
|
("comfy", re.compile(r"^diffusion_model\.")),
|
|
)
|
|
|
|
|
|
def read_header(path: str) -> dict:
|
|
with open(path, "rb") as f:
|
|
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))
|
|
|
|
|
|
def _convention(keys) -> str:
|
|
for name, pattern in _CONVENTIONS:
|
|
if any(pattern.match(k) for k in keys):
|
|
return name
|
|
if any(k.startswith(("blocks.", "token_refiner.", "final_layer.")) for k in keys):
|
|
return "bare"
|
|
return "unknown"
|
|
|
|
|
|
def inspect(path: str) -> dict:
|
|
"""Summarise one LoRA file. Cached on (path, mtime, size)."""
|
|
try:
|
|
st = os.stat(path)
|
|
except OSError as exc:
|
|
return {"error": str(exc), "name": os.path.basename(path)}
|
|
key = (path, st.st_mtime_ns, st.st_size)
|
|
hit = _cache.get(key)
|
|
if hit is not None:
|
|
return hit
|
|
|
|
info: dict = {"name": os.path.basename(path), "path": path, "size": st.st_size}
|
|
try:
|
|
header = read_header(path)
|
|
except Exception as exc:
|
|
info["error"] = f"unreadable header: {exc}"
|
|
_cache[key] = info
|
|
return info
|
|
|
|
metadata = header.pop("__metadata__", {}) or {}
|
|
keys = list(header)
|
|
info["tensors"] = len(keys)
|
|
info["convention"] = _convention(keys)
|
|
info["dtypes"] = sorted({v.get("dtype", "?") for v in header.values()})
|
|
info["base_model"] = metadata.get("ss_base_model_version", "")
|
|
info["trainer"] = ""
|
|
try:
|
|
info["trainer"] = json.loads(metadata.get("software", "{}")).get("name", "")
|
|
except Exception:
|
|
pass
|
|
if not info["trainer"] and metadata.get("ss_network_module"):
|
|
info["trainer"] = metadata["ss_network_module"]
|
|
|
|
ranks: set[int] = set()
|
|
adaln_dims: set[int] = set()
|
|
modules: set[str] = set()
|
|
has_alpha = False
|
|
for k, v in header.items():
|
|
if k.endswith(".alpha"):
|
|
has_alpha = True
|
|
continue
|
|
shape = v.get("shape") or []
|
|
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:
|
|
adaln_dims.add(int(shape[1]))
|
|
for tag in ("adaln_proj", "attn.qkv_proj", "attn_qkv_proj", "attn.out_proj",
|
|
"attn_out_proj", "mlp.fc1", "mlp_fc1", "mlp.fc2", "mlp_fc2",
|
|
"token_refiner", "final_layer"):
|
|
if tag in k:
|
|
modules.add(tag.replace("_proj", "").replace(".", "_").replace("attn_qkv", "qkv")
|
|
.replace("attn_out", "out"))
|
|
|
|
info["ranks"] = sorted(ranks)
|
|
info["alpha"] = has_alpha
|
|
info["adaln_dims"] = sorted(adaln_dims)
|
|
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|_lora|lora_linear_layer|reshape_weight|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
|
|
|
|
|
|
def summary_line(info: dict) -> str:
|
|
if info.get("error"):
|
|
return f"{info['name']}: {info['error']}"
|
|
ranks = "/".join(str(r) for r in info.get("ranks", [])) or "?"
|
|
adaln = info.get("adaln_dims")
|
|
adaln_s = f", adaLN dim {'/'.join(str(d) for d in adaln)}" if adaln else ", no adaLN"
|
|
return (f"{info['name']}: {info.get('tensors', 0)} tensors, {info.get('convention', '?')} "
|
|
f"format, rank {ranks}{adaln_s}")
|