358 lines
18 KiB
Python
358 lines
18 KiB
Python
"""PromptForge LoRA loader: a searchable list of every LoRA you have, with per-LoRA strengths.
|
|
|
|
The list is drawn by web/promptforge_lora.js; this file holds the state format, the library index
|
|
(served at /promptforge/loras) and the node. The state is one JSON string:
|
|
|
|
{"filter": "all", "items": [{"name": "sub/a.safetensors", "on": true, "model": 1.0, "clip": null,
|
|
"random": false, "min": 0.6, "max": 1.0, "group": "", "part": "all"}]}
|
|
|
|
* `clip` null means "same as model".
|
|
* `random` draws the strength between `min` and `max`, seeded by the node seed.
|
|
* Items sharing a non-empty `group` are alternatives: one of the enabled ones is picked per run.
|
|
* `part` limits a LoRA to the video layers or the audio layers of an audio+video model (MiniMax H3, LTX-2).
|
|
* With two or more LoRAs on, the report says how much they overlap and how hard they push together; `balance`
|
|
lowers all strengths by the same factor when together they push harder than the strongest one alone.
|
|
"""
|
|
import json
|
|
import os
|
|
import random
|
|
import re
|
|
|
|
import torch
|
|
from aiohttp import web
|
|
from safetensors import safe_open
|
|
from server import PromptServer
|
|
|
|
import comfy.lora
|
|
import comfy.lora_convert
|
|
import comfy.utils
|
|
import folder_paths
|
|
|
|
from . import families
|
|
|
|
_index = {} # LoRA name -> ((mtime, size), entry)
|
|
|
|
|
|
def _trigger_words(sidecar):
|
|
words = sidecar.get("trainedWords") or (sidecar.get("civitai") or {}).get("trainedWords") or []
|
|
return [word.strip() for word in words if isinstance(word, str) and word.strip()]
|
|
|
|
|
|
def _describe(name):
|
|
path = folder_paths.get_full_path("loras", name)
|
|
if path is None:
|
|
return None
|
|
stat = os.stat(path)
|
|
signature = (stat.st_mtime_ns, stat.st_size)
|
|
cached = _index.get(name)
|
|
if cached and cached[0] == signature:
|
|
return cached[1]
|
|
family = families.family_of_lora(name, path)
|
|
sidecar = {}
|
|
sidecar_path = path.rsplit(".", 1)[0] + ".metadata.json"
|
|
try:
|
|
with open(sidecar_path, encoding="utf-8") as handle:
|
|
sidecar = json.load(handle)
|
|
except (OSError, ValueError):
|
|
pass
|
|
entry = {"name": name, "family": family.key if family else "", "words": _trigger_words(sidecar), "size": stat.st_size}
|
|
_index[name] = (signature, entry)
|
|
return entry
|
|
|
|
|
|
def library():
|
|
"""Every LoRA with its family and trigger words, plus the families present (for the filter)."""
|
|
entries = [entry for entry in map(_describe, folder_paths.get_filename_list("loras")) if entry]
|
|
present = sorted({entry["family"] for entry in entries if entry["family"]})
|
|
return {
|
|
"loras": entries,
|
|
"families": [{"key": key, "label": families.BY_KEY[key].label} for key in present],
|
|
}
|
|
|
|
|
|
@PromptServer.instance.routes.get("/promptforge/loras")
|
|
async def _library_route(request):
|
|
return web.json_response(library())
|
|
|
|
|
|
PREVIEW_SUFFIXES = (".preview.jpeg", ".preview.jpg", ".preview.png", ".preview.webp", ".jpeg", ".jpg", ".png", ".webp")
|
|
|
|
|
|
def _preview_path(name):
|
|
stem = os.path.splitext(folder_paths.get_full_path_or_raise("loras", name))[0]
|
|
return next((stem + suffix for suffix in PREVIEW_SUFFIXES if os.path.isfile(stem + suffix)), None)
|
|
|
|
|
|
def _plain(html):
|
|
return re.sub(r"\n{3,}", "\n\n", re.sub(r"<[^>]+>", "", re.sub(r"<(?:br|/p|/li)\s*/?>", "\n", html or ""))).strip()
|
|
|
|
|
|
TRAINING_FIELDS = (("ss_network_module", "Network"), ("ss_network_dim", "Rank"), ("ss_network_alpha", "Alpha"), ("ss_base_model_version", "Trained on"), ("modelspec.architecture", "Architecture"),
|
|
("ss_resolution", "Resolution"), ("ss_num_epochs", "Epochs"), ("ss_max_train_steps", "Steps"), ("ss_learning_rate", "Learning rate"), ("ss_optimizer", "Optimizer"),
|
|
("ss_num_train_images", "Training images"), ("modelspec.title", "Title"), ("ss_training_comment", "Comment"))
|
|
|
|
|
|
def _training_info(path):
|
|
"""What the file itself says about its training: [(label, value)] and the most common training tags."""
|
|
try:
|
|
with safe_open(path, framework="pt") as handle:
|
|
metadata = handle.metadata() or {}
|
|
except (OSError, ValueError):
|
|
return [], []
|
|
fields = [(label, str(metadata[key])[:200]) for key, label in TRAINING_FIELDS if metadata.get(key) not in (None, "", "None")]
|
|
counts = {}
|
|
try:
|
|
for dataset in json.loads(metadata.get("ss_tag_frequency") or "{}").values():
|
|
for tag, count in dataset.items():
|
|
counts[tag.strip()] = counts.get(tag.strip(), 0) + int(count)
|
|
except (ValueError, AttributeError, TypeError):
|
|
pass
|
|
return fields, [tag for tag, _ in sorted(counts.items(), key=lambda entry: -entry[1])[:24] if tag]
|
|
|
|
|
|
@PromptServer.instance.routes.get("/promptforge/lora_info")
|
|
async def _info_route(request):
|
|
name = request.query.get("name", "")
|
|
path = folder_paths.get_full_path("loras", name)
|
|
if path is None:
|
|
return web.json_response({"error": "not found"}, status=404)
|
|
try:
|
|
with open(os.path.splitext(path)[0] + ".metadata.json", encoding="utf-8") as handle:
|
|
sidecar = json.load(handle)
|
|
except (OSError, ValueError):
|
|
sidecar = {}
|
|
civitai = sidecar.get("civitai") or {}
|
|
model_id = civitai.get("modelId")
|
|
model = civitai.get("model") or {}
|
|
training, training_tags = _training_info(path)
|
|
images = []
|
|
for image in civitai.get("images") or []:
|
|
meta = image.get("meta") or {}
|
|
if image.get("url"):
|
|
images.append({"url": image["url"], "video": image.get("type") == "video", "seed": meta.get("seed"), "steps": meta.get("steps"), "cfg": meta.get("cfgScale"),
|
|
"sampler": meta.get("sampler"), "model": meta.get("Model"), "positive": meta.get("prompt"), "negative": meta.get("negativePrompt")})
|
|
return web.json_response({
|
|
"name": name,
|
|
"title": sidecar.get("model_name") or os.path.basename(name),
|
|
"version": civitai.get("name") or "",
|
|
"type": model.get("type") or "",
|
|
"base_model": sidecar.get("base_model", ""),
|
|
"sha256": sidecar.get("sha256", ""),
|
|
"size": sidecar.get("size") or os.path.getsize(path),
|
|
"creator": (civitai.get("creator") or {}).get("username", ""),
|
|
"downloads": (civitai.get("stats") or {}).get("downloadCount"),
|
|
"likes": (civitai.get("stats") or {}).get("thumbsUpCount"),
|
|
"words": _trigger_words(sidecar),
|
|
"training": training,
|
|
"training_tags": training_tags,
|
|
"tips": sidecar.get("usage_tips") or "",
|
|
"notes": sidecar.get("notes") or "",
|
|
"tags": sidecar.get("tags") or [],
|
|
"description": _plain(sidecar.get("modelDescription") or ""),
|
|
"url": f"https://civitai.com/models/{model_id}" + (f"?modelVersionId={civitai['id']}" if civitai.get("id") else "") if model_id else "",
|
|
"preview": _preview_path(name) is not None,
|
|
"images": images,
|
|
})
|
|
|
|
|
|
@PromptServer.instance.routes.get("/promptforge/lora_preview")
|
|
async def _preview_route(request):
|
|
name = request.query.get("name", "")
|
|
if folder_paths.get_full_path("loras", name) is None or (path := _preview_path(name)) is None:
|
|
return web.Response(status=404)
|
|
return web.FileResponse(path)
|
|
|
|
|
|
def _model_family(model):
|
|
config = getattr(getattr(model, "model", None), "model_config", None)
|
|
return families.family_from_config(type(config).__name__) if config is not None else None
|
|
|
|
|
|
def choose(items, seed):
|
|
"""The (item, model strength, clip strength) list that applies for this seed."""
|
|
rng = random.Random(seed)
|
|
groups = {}
|
|
for item in items:
|
|
if item.get("on"):
|
|
groups.setdefault((item.get("group") or "").strip(), []).append(item)
|
|
chosen = []
|
|
picked = {group: rng.choice(members) for group, members in groups.items() if group}
|
|
for item in items:
|
|
if not item.get("on"):
|
|
continue
|
|
group = (item.get("group") or "").strip()
|
|
if group and picked[group] is not item:
|
|
continue
|
|
if item.get("random"):
|
|
low, high = sorted((float(item.get("min", 0.0)), float(item.get("max", 1.0))))
|
|
strength = round(rng.uniform(low, high), 2)
|
|
else:
|
|
strength = float(item.get("model", 1.0))
|
|
clip = item.get("clip")
|
|
chosen.append((item, strength, strength if clip is None or item.get("random") else float(clip)))
|
|
return chosen
|
|
|
|
|
|
AUDIO_KEY = "audio"
|
|
|
|
|
|
def in_part(key, part):
|
|
"""Whether a patched weight belongs to the chosen part of an audio+video model."""
|
|
return part == "all" or (AUDIO_KEY in str(key).lower()) == (part == "audio")
|
|
|
|
|
|
def _factors(adapter):
|
|
"""(up, down, scale) of a plain LoRA patch, or None for other kinds of patches."""
|
|
if getattr(adapter, "name", "") != "lora":
|
|
return None
|
|
up, down, alpha, mid, dora, _ = adapter.weights
|
|
if mid is not None or dora is not None:
|
|
return None
|
|
down = down.flatten(1).float()
|
|
return up.float(), down, (alpha if alpha is not None else down.shape[0]) / down.shape[0]
|
|
|
|
|
|
def stack_stats(stack, limit=160):
|
|
"""Energy of every LoRA's change and the overlap of each pair, from the low-rank factors (no full weight matrices).
|
|
|
|
stack: [(name, strength, patches)]. Returns (norms2, dots) where dots[(i, j)] is the inner product of the two changes
|
|
over the layers both touch, and norms2[i] the squared size of change i over its own layers; each is measured on at
|
|
most `limit` layers per LoRA.
|
|
"""
|
|
factors = []
|
|
for _, _, patches in stack:
|
|
layers = {key: f for key, adapter in patches.items() if (f := _factors(adapter)) is not None}
|
|
keys = sorted(layers, key=str)
|
|
step = max(1, len(keys) // limit)
|
|
factors.append({key: layers[key] for key in keys[::step]})
|
|
norms2 = [0.0] * len(stack)
|
|
dots = {}
|
|
for i, mine in enumerate(factors):
|
|
for key, (up, down, scale) in mine.items():
|
|
norms2[i] += float(scale ** 2 * ((up.T @ up) * (down @ down.T)).sum())
|
|
for j in range(i + 1, len(stack)):
|
|
total, shared = 0.0, 0
|
|
for key, (up, down, scale) in mine.items():
|
|
other = factors[j].get(key)
|
|
if other is None or other[0].shape[0] != up.shape[0] or other[1].shape[1] != down.shape[1]:
|
|
continue
|
|
total += float(scale * other[2] * ((up.T @ other[0]) * (down @ other[1].T)).sum())
|
|
shared += 1
|
|
dots[(i, j)] = (total, shared)
|
|
return norms2, dots
|
|
|
|
|
|
def balance_factor(strengths, norms2, dots, floor):
|
|
"""(factor, combined/strongest energy ratio): the uniform scale that brings the combined push down to that of the strongest LoRA alone."""
|
|
single = max(s * s * n for s, n in zip(strengths, norms2))
|
|
combined = sum(s * s * n for s, n in zip(strengths, norms2)) + 2 * sum(strengths[i] * strengths[j] * dot for (i, j), (dot, _) in dots.items())
|
|
if single <= 0 or combined <= 0:
|
|
return 1.0, 1.0
|
|
return min(1.0, max(floor, (single / combined) ** 0.5)), combined / single
|
|
|
|
|
|
def overlap_words(cosine):
|
|
return "opposing" if cosine < -0.15 else "aligned" if cosine > 0.4 else "partly aligned" if cosine > 0.15 else "independent"
|
|
|
|
|
|
class PROMPTFORGE_LoraLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"loras": ("STRING", {"default": "{}", "tooltip": "Managed by the list below; edit the list, not this text."}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 2 ** 32 - 1, "tooltip": "Seeds random strengths and random group picks."}),
|
|
"skip_mismatched": ("BOOLEAN", {"default": True, "tooltip": "Skip LoRAs trained for a different model family than the connected model."}),
|
|
},
|
|
"optional": {
|
|
"clip": ("CLIP",),
|
|
"pipe": ("PROMPTFORGE_PIPE", {"tooltip": "From a PromptForge Loader: the list switches to that model's family."}),
|
|
"balance": ("BOOLEAN", {"default": False, "tooltip": "When the LoRAs together push harder than the strongest one alone, lower all strengths by the same factor (the report shows the numbers). Image models can go down to 0.85 of the set strengths; video and audio models are left alone unless balance_floor is set."}),
|
|
"balance_floor": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "Lowest factor balance may apply. 0 = by model kind: 0.85 for image models, 1.0 (no change) for video and audio."}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "STRING", "STRING")
|
|
RETURN_NAMES = ("MODEL", "CLIP", "TRIGGER_WORDS", "REPORT")
|
|
FUNCTION = "load"
|
|
CATEGORY = "🔮 PromptForge/Load"
|
|
DESCRIPTION = "Searchable LoRA list with per-LoRA strengths, random strengths and random groups."
|
|
|
|
def load(self, model, loras, seed, skip_mismatched, clip=None, pipe=None, balance=False, balance_floor=0.0):
|
|
try:
|
|
items = json.loads(loras or "{}").get("items", [])
|
|
except ValueError as exc:
|
|
raise ValueError(f"[PromptForge] the LoRA list is not valid JSON: {exc}")
|
|
target = _model_family(model) if skip_mismatched else None
|
|
|
|
report, words = [], []
|
|
key_map = None
|
|
stack = [] # (name, model strength, clip strength, patches, metadata)
|
|
for item, strength_model, strength_clip in choose(items, seed):
|
|
name = item["name"]
|
|
entry = _describe(name)
|
|
if entry is None:
|
|
raise ValueError(f"[PromptForge] LoRA '{name}' was not found in the loras folders.")
|
|
if target is not None and entry["family"] and entry["family"] != target.key:
|
|
report.append(f"skipped {name}: made for {entry['family']}, model is {target.key}")
|
|
continue
|
|
if strength_model == 0 and strength_clip == 0:
|
|
continue
|
|
lora, metadata = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("loras", name), safe_load=True, return_metadata=True)
|
|
if key_map is None: # the key maps depend only on the base model and clip, so build them once for all LoRAs
|
|
key_map = comfy.lora.model_lora_keys_unet(model.model, {})
|
|
if clip is not None:
|
|
key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map)
|
|
patches = comfy.lora.load_lora(comfy.lora_convert.convert_lora(lora), key_map)
|
|
part = item.get("part") or "all"
|
|
if part != "all":
|
|
patches = {key: patch for key, patch in patches.items() if in_part(key, part)}
|
|
if not patches:
|
|
report.append(f"skipped {name}: it has no {part} layers")
|
|
continue
|
|
stack.append((name, strength_model, strength_clip, patches, metadata, part))
|
|
words.extend(word for word in entry["words"] if word not in words)
|
|
|
|
factor = 1.0
|
|
if stack:
|
|
norms2, dots = stack_stats([(name, strength, patches) for name, strength, _, patches, _, _ in stack])
|
|
if balance and len(stack) > 1:
|
|
floor = balance_floor or (0.85 if target is None or target.kind == "image" else 1.0)
|
|
factor, ratio = balance_factor([strength for _, strength, *_ in stack], norms2, dots, floor)
|
|
for (name, strength_model, strength_clip, patches, _, part), norm2 in zip(stack, norms2):
|
|
rank = next((a.weights[1].shape[0] for a in patches.values() if getattr(a, "name", "") == "lora"), 0)
|
|
report.append(f"{name}: model {strength_model * factor:g}, clip {strength_clip * factor:g}" + (f" ({part} layers only)" if part != "all" else "")
|
|
+ f", {len(patches)} layers" + (f", rank {rank}" if rank else "") + (f", strength-weighted size {strength_model * norm2 ** 0.5:.3g}" if norm2 else ""))
|
|
for (i, j), (dot, shared) in dots.items():
|
|
size = (norms2[i] * norms2[j]) ** 0.5
|
|
if shared and size:
|
|
cosine = dot / size
|
|
report.append(f"overlap {stack[i][0]} / {stack[j][0]}: {shared} shared layers, similarity {cosine:+.2f} ({overlap_words(cosine)})")
|
|
if len(stack) > 1:
|
|
single = max(s * s * n for (_, s, *_), n in zip(stack, norms2))
|
|
combined = sum(s * s * n for (_, s, *_), n in zip(stack, norms2)) + 2 * sum(stack[i][1] * stack[j][1] * dot for (i, j), (dot, _) in dots.items())
|
|
if single > 0:
|
|
report.append(f"together they push {max(combined, 0) / single:.2f}x as hard as the strongest alone"
|
|
+ (f"; balance lowered every strength to {factor:.2f}x" if factor < 1 else "; turn on balance to even this out" if combined / single > 1.5 else ""))
|
|
|
|
patched_model, patched_clip = model, clip
|
|
for name, strength_model, strength_clip, patches, metadata, _ in stack:
|
|
if strength_model:
|
|
if patched_model is model:
|
|
patched_model = model.clone()
|
|
patched_model.add_patches(patches, strength_model * factor)
|
|
if metadata:
|
|
patched_model.set_attachments("lora_metadata", metadata)
|
|
if clip is not None and strength_clip:
|
|
if patched_clip is clip:
|
|
patched_clip = clip.clone()
|
|
patched_clip.add_patches(patches, strength_clip * factor)
|
|
if metadata:
|
|
patched_clip.patcher.set_attachments("lora_metadata", metadata)
|
|
return (patched_model, patched_clip, ", ".join(words), "\n".join(report))
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"PROMPTFORGE_LoraLoader": PROMPTFORGE_LoraLoader}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"PROMPTFORGE_LoraLoader": "PromptForge LoRA Loader"}
|