simplified logic and split into multiple files
This commit is contained in:
FugitiveExpert01
+109
-709
@@ -1,218 +1,48 @@
|
||||
"""
|
||||
FEnodes — LoRA utility nodes for VFX production pipelines.
|
||||
FEnodes -- LoRA nodes
|
||||
Author: FugitiveExpert01
|
||||
|
||||
FELoraLoad — multi-LoRA browser UI → FE_LORA_STACK
|
||||
FEApplyLora — MODEL + FE_LORA_STACK → MODEL / CLIP
|
||||
FELoraTriggerAnalysis — FE_LORA_STACK + CLIP → STRING (candidate trigger words)
|
||||
|
||||
FELoraLoad has a custom JS UI (web/js/fe_power_lora.js) providing:
|
||||
• Folder-tree browser with search
|
||||
• Per-row on/off toggle and strength
|
||||
• "+ Add LoRA" button
|
||||
• Optional separate CLIP strength (right-click node)
|
||||
|
||||
FELoraLoad -- multi-LoRA browser UI -> FE_LORA_STACK
|
||||
FEApplyLora -- MODEL + FE_LORA_STACK -> MODEL / CLIP
|
||||
FELoraTriggerAnalysis -- FE_LORA_STACK + CLIP -> STRING
|
||||
|
||||
FELoraLoad has a custom JS UI (web/js/fe_power_lora.js).
|
||||
"""
|
||||
|
||||
|
||||
__version__ = "0.1.2"
|
||||
|
||||
|
||||
import json
|
||||
import logging
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
|
||||
|
||||
from .lora_utils import (
|
||||
FE_LORA_STACK,
|
||||
_lora_weight_cache,
|
||||
_cached_file_hash,
|
||||
normalize_lora_name,
|
||||
merge_lora_weights,
|
||||
discover_encoders,
|
||||
analyse_encoder,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("FEnodes")
|
||||
|
||||
FE_LORA_STACK = "FE_LORA_STACK"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Module-level weight cache
|
||||
# Keyed on absolute lora_path so identical files shared across multiple
|
||||
# FELoraLoad nodes or tile streams are only read from disk once.
|
||||
# ---------------------------------------------------------------------------
|
||||
_lora_weight_cache: dict[str, dict] = {}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# IS_CHANGED hash cache
|
||||
# Maps lora_path -> (mtime, sha256_hex) so IS_CHANGED never reads a file
|
||||
# twice unless it has actually been modified on disk.
|
||||
# ---------------------------------------------------------------------------
|
||||
_hash_cache: dict[str, tuple[float, str]] = {}
|
||||
|
||||
def _cached_file_hash(path: str) -> str:
|
||||
"""Return SHA256 hex for path, recomputing only when mtime changes."""
|
||||
import os, hashlib
|
||||
try:
|
||||
mtime = os.path.getmtime(path)
|
||||
except OSError:
|
||||
return path # file missing — use path as sentinel
|
||||
|
||||
cached = _hash_cache.get(path)
|
||||
if cached is not None and cached[0] == mtime:
|
||||
return cached[1]
|
||||
|
||||
h = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(131072), b""):
|
||||
h.update(chunk)
|
||||
digest = h.hexdigest()
|
||||
_hash_cache[path] = (mtime, digest)
|
||||
return digest
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# API routes
|
||||
# ---------------------------------------------------------------------------
|
||||
try:
|
||||
import hashlib
|
||||
import os
|
||||
import requests as _requests
|
||||
from server import PromptServer
|
||||
from aiohttp import web as _web
|
||||
|
||||
# ── /fenodes/loras — flat file list for the JS browser ────────────────
|
||||
@PromptServer.instance.routes.get("/fenodes/loras")
|
||||
async def _fenodes_get_loras(request):
|
||||
loras = folder_paths.get_filename_list("loras")
|
||||
return _web.json_response({"loras": loras})
|
||||
|
||||
# ── /fenodes/lora_info — CivitAI + metadata info for one LoRA ────────
|
||||
def _sha256(file_path: str) -> str:
|
||||
h = hashlib.sha256()
|
||||
with open(file_path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(131072), b""):
|
||||
h.update(chunk)
|
||||
return h.hexdigest()
|
||||
|
||||
def _read_safetensors_metadata(file_path: str) -> dict:
|
||||
"""Read the __metadata__ block from a .safetensors header."""
|
||||
try:
|
||||
with open(file_path, "rb") as f:
|
||||
header_size = int.from_bytes(f.read(8), "little", signed=False)
|
||||
if header_size <= 0 or header_size > 100_000_000:
|
||||
return {}
|
||||
header = json.loads(f.read(header_size))
|
||||
meta = header.get("__metadata__", {})
|
||||
# Parse any JSON-string values
|
||||
for k, v in meta.items():
|
||||
if isinstance(v, str) and v.startswith("{"):
|
||||
try:
|
||||
meta[k] = json.loads(v)
|
||||
except Exception:
|
||||
pass
|
||||
return meta
|
||||
except Exception as e:
|
||||
logger.warning(f"[FEnodes/lora_info] Failed to read metadata: {e}")
|
||||
return {}
|
||||
|
||||
@PromptServer.instance.routes.get("/fenodes/lora_info")
|
||||
async def _fenodes_lora_info(request):
|
||||
lora_name = request.rel_url.query.get("name", "")
|
||||
refresh = request.rel_url.query.get("refresh", "false").lower() == "true"
|
||||
|
||||
if not lora_name:
|
||||
return _web.json_response({"error": "missing name"}, status=400)
|
||||
|
||||
lora_path = folder_paths.get_full_path("loras", lora_name)
|
||||
if not lora_path or not os.path.isfile(lora_path):
|
||||
return _web.json_response({"error": "file not found"}, status=404)
|
||||
|
||||
cache_path = os.path.splitext(lora_path)[0] + ".fe-info.json"
|
||||
|
||||
# Load cache if it exists and refresh not requested
|
||||
cached = {}
|
||||
if os.path.isfile(cache_path) and not refresh:
|
||||
try:
|
||||
with open(cache_path, "r", encoding="utf-8") as f:
|
||||
cached = json.load(f)
|
||||
logger.info(f"[FEnodes/lora_info] Cache hit: {lora_name}")
|
||||
return _web.json_response(cached)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger.info(f"[FEnodes/lora_info] Building info for: {lora_name}")
|
||||
|
||||
info = {"file": lora_name}
|
||||
|
||||
# Safetensors header metadata
|
||||
if lora_path.endswith(".safetensors"):
|
||||
meta = _read_safetensors_metadata(lora_path)
|
||||
if meta:
|
||||
info["metadata"] = meta
|
||||
# Pull out common fields
|
||||
for field in ("ss_sd_model_name", "ss_base_model_version",
|
||||
"ss_network_module", "ss_num_train_images",
|
||||
"modelspec.title", "modelspec.architecture"):
|
||||
if field in meta:
|
||||
info[field] = meta[field]
|
||||
|
||||
# CivitAI lookup via SHA256
|
||||
try:
|
||||
file_hash = _sha256(lora_path)
|
||||
info["sha256"] = file_hash
|
||||
civitai_url = f"https://civitai.com/api/v1/model-versions/by-hash/{file_hash}"
|
||||
logger.info(f"[FEnodes/lora_info] Querying CivitAI: {civitai_url}")
|
||||
resp = _requests.get(civitai_url, timeout=10)
|
||||
if resp.status_code == 200:
|
||||
cdata = resp.json()
|
||||
info["civitai"] = {
|
||||
"name": cdata.get("model", {}).get("name", ""),
|
||||
"versionName": cdata.get("name", ""),
|
||||
"baseModel": cdata.get("baseModel", ""),
|
||||
"modelId": cdata.get("modelId"),
|
||||
"versionId": cdata.get("id"),
|
||||
"trainedWords": cdata.get("trainedWords", []),
|
||||
"description": cdata.get("description", ""),
|
||||
"url": (
|
||||
f"https://civitai.com/models/{cdata.get('modelId')}"
|
||||
f"?modelVersionId={cdata.get('id')}"
|
||||
if cdata.get("modelId") else ""
|
||||
),
|
||||
"images": [
|
||||
{"url": img.get("url"), "nsfw": img.get("nsfwLevel", 0)}
|
||||
for img in cdata.get("images", [])[:3]
|
||||
],
|
||||
}
|
||||
elif resp.status_code == 404:
|
||||
info["civitai"] = {"error": "not found on CivitAI"}
|
||||
else:
|
||||
info["civitai"] = {"error": f"CivitAI returned {resp.status_code}"}
|
||||
except Exception as e:
|
||||
logger.warning(f"[FEnodes/lora_info] CivitAI fetch failed: {e}")
|
||||
info["civitai"] = {"error": str(e)}
|
||||
|
||||
# Save cache
|
||||
try:
|
||||
with open(cache_path, "w", encoding="utf-8") as f:
|
||||
json.dump(info, f, indent=2)
|
||||
except Exception as e:
|
||||
logger.warning(f"[FEnodes/lora_info] Could not write cache: {e}")
|
||||
|
||||
return _web.json_response(info)
|
||||
|
||||
logger.info("[FEnodes] Registered /fenodes/loras and /fenodes/lora_info API routes.")
|
||||
except Exception as _e:
|
||||
logger.warning(f"[FEnodes] Could not register API routes: {_e}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FELoraLoad — multi-LoRA browser → FE_LORA_STACK
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class FELoraLoad:
|
||||
"""
|
||||
Multi-LoRA loader with a custom folder-tree browser UI (see fe_power_lora.js).
|
||||
|
||||
The JS widget serialises the full LoRA list as JSON into the hidden
|
||||
`loras_json` STRING widget. Each entry:
|
||||
{ "enabled": bool, "lora": "path/to/lora.safetensors",
|
||||
Multi-LoRA loader with a custom folder-tree browser UI (fe_power_lora.js).
|
||||
|
||||
The JS widget serialises the LoRA list as JSON into the hidden loras_json
|
||||
STRING widget. Each entry:
|
||||
{ "enabled": bool, "lora": "path/to/file.safetensors",
|
||||
"strength_model": float, "strength_clip": float }
|
||||
|
||||
Disabled entries and empty-lora rows are skipped at execute time.
|
||||
Outputs a FE_LORA_STACK (list of dicts) for FEApplyLora.
|
||||
|
||||
Disabled and empty rows are skipped. Outputs FE_LORA_STACK for FEApplyLora.
|
||||
"""
|
||||
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
@@ -220,63 +50,62 @@ class FELoraLoad:
|
||||
"loras_json": ("STRING", {"default": "[]", "multiline": False}),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
RETURN_TYPES = (FE_LORA_STACK,)
|
||||
RETURN_NAMES = ("lora_stack",)
|
||||
FUNCTION = "load"
|
||||
CATEGORY = "FEnodes"
|
||||
DESCRIPTION = (
|
||||
"Multi-LoRA loader with folder-tree browser UI. "
|
||||
"Add any number of LoRAs, toggle each on/off, set per-LoRA strength. "
|
||||
"Add LoRAs, toggle on/off, set per-LoRA strength. "
|
||||
"Connect lora_stack to FEApplyLora."
|
||||
)
|
||||
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, loras_json):
|
||||
try:
|
||||
entries = json.loads(loras_json)
|
||||
except Exception:
|
||||
return loras_json
|
||||
|
||||
|
||||
parts = []
|
||||
for entry in entries:
|
||||
if not entry.get("enabled", True) or not entry.get("lora"):
|
||||
continue
|
||||
lora_name = entry["lora"].replace("\\", "/")
|
||||
path = folder_paths.get_full_path("loras", lora_name)
|
||||
lora_name = normalize_lora_name(entry["lora"])
|
||||
path = folder_paths.get_full_path("loras", lora_name)
|
||||
file_hash = _cached_file_hash(path) if path else lora_name
|
||||
# Normalise strengths to 4dp to avoid float serialisation jitter
|
||||
sm = round(float(entry.get("strength_model", 1.0)), 4)
|
||||
sc = round(float(entry.get("strength_clip", sm)), 4)
|
||||
parts.append(f"{lora_name}:{file_hash}:{sm}:{sc}")
|
||||
return "|".join(parts)
|
||||
|
||||
|
||||
def load(self, loras_json):
|
||||
try:
|
||||
entries = json.loads(loras_json)
|
||||
except Exception as e:
|
||||
logger.warning(f"[FEnodes/FELoraLoad] Failed to parse loras_json: {e}")
|
||||
return ([],)
|
||||
|
||||
|
||||
stack = []
|
||||
for entry in entries:
|
||||
if not entry.get("enabled", True):
|
||||
continue
|
||||
lora_name = entry.get("lora", "").replace("\\", "/")
|
||||
|
||||
lora_name = normalize_lora_name(entry.get("lora", ""))
|
||||
if not lora_name:
|
||||
continue
|
||||
|
||||
|
||||
lora_path = folder_paths.get_full_path("loras", lora_name)
|
||||
if lora_path is None:
|
||||
logger.warning(f"[FEnodes/FELoraLoad] LoRA not found: {lora_name}")
|
||||
continue
|
||||
|
||||
|
||||
strength_model = float(entry.get("strength_model", 1.0))
|
||||
# strength_clip is always written explicitly by the JS widget.
|
||||
# The fallback to strength_model is intentional: if no separate
|
||||
# CLIP strength was set, treat them as equal.
|
||||
strength_clip = float(entry.get("strength_clip", strength_model))
|
||||
|
||||
# strength_clip falls back to strength_model if not explicitly set.
|
||||
# The JS widget always writes it; fallback covers older workflows.
|
||||
strength_clip = float(entry.get("strength_clip", strength_model))
|
||||
|
||||
if lora_path in _lora_weight_cache:
|
||||
logger.info(f"[FEnodes/FELoraLoad] Cache hit: '{lora_name}'")
|
||||
weights = _lora_weight_cache[lora_path]
|
||||
@@ -287,142 +116,31 @@ class FELoraLoad:
|
||||
)
|
||||
weights = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
||||
_lora_weight_cache[lora_path] = weights
|
||||
|
||||
stack.append({
|
||||
"name": lora_name,
|
||||
"weights": weights,
|
||||
"strength_model": strength_model,
|
||||
"strength_clip": strength_clip,
|
||||
})
|
||||
|
||||
logger.info(f"[FEnodes/FELoraLoad] Stack ready: {len(stack)} LoRA(s) loaded.")
|
||||
|
||||
logger.info(f"[FEnodes/FELoraLoad] Stack ready: {len(stack)} LoRA(s).")
|
||||
return (stack,)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Merge helper — pre-scales and combines all LoRA weight dicts into one,
|
||||
# so load_lora_for_models is called once rather than once per LoRA.
|
||||
#
|
||||
# Keys containing "lora_te" are treated as text-encoder (CLIP) keys and
|
||||
# scaled by that entry's strength_clip; all others are scaled by strength_model.
|
||||
# This heuristic covers every major training convention (Kohya SD1/SDXL/Flux,
|
||||
# SimpleTuner, etc.) without requiring model-specific knowledge.
|
||||
# ---------------------------------------------------------------------------
|
||||
def _merge_lora_weights(lora_stack: list) -> dict:
|
||||
"""
|
||||
Merge all LoRAs in the stack into a single synthetic weight dict by
|
||||
concatenating lora_down / lora_up tensors along the rank dimension,
|
||||
with per-LoRA strength and alpha baked in via sqrt-scaling.
|
||||
|
||||
Mathematical guarantee:
|
||||
combined_up @ combined_down = Σ strength_i * alpha_scale_i * (up_i @ down_i)
|
||||
|
||||
Handles both naming conventions transparently:
|
||||
Kohya-style: .lora_down.weight / .lora_up.weight
|
||||
Diffusers: .lora_A.weight / .lora_B.weight (e.g. SD Turbo, SDXL Turbo)
|
||||
|
||||
Both are normalised internally and output as lora_down/up, which
|
||||
comfy.sd.load_lora_for_models accepts for all architectures.
|
||||
|
||||
Layers that appear in only one LoRA pass through unchanged (no concat overhead).
|
||||
Layers where ranks differ (shape mismatch) fall back gracefully with a warning.
|
||||
"""
|
||||
staging: dict = {} # base_key → {"downs": [], "ups": []}
|
||||
|
||||
# Supported (down_suffix, up_suffix) pairs in priority order.
|
||||
# We try lora_down first; if absent we try lora_A.
|
||||
_PAIRS = [
|
||||
(".lora_down.weight", ".lora_up.weight"),
|
||||
(".lora_A.weight", ".lora_B.weight"),
|
||||
]
|
||||
|
||||
for entry in lora_stack:
|
||||
weights = entry["weights"]
|
||||
strength_model = entry["strength_model"]
|
||||
strength_clip = entry["strength_clip"]
|
||||
|
||||
# Build a set of base keys present in this LoRA, regardless of convention.
|
||||
base_keys = set()
|
||||
for down_sfx, _ in _PAIRS:
|
||||
for k in weights:
|
||||
if k.endswith(down_sfx):
|
||||
base_keys.add(k[: -len(down_sfx)])
|
||||
|
||||
for base_key in base_keys:
|
||||
# Find which convention this layer uses
|
||||
down, up = None, None
|
||||
for down_sfx, up_sfx in _PAIRS:
|
||||
dk = base_key + down_sfx
|
||||
uk = base_key + up_sfx
|
||||
if dk in weights and uk in weights:
|
||||
down = weights[dk].float()
|
||||
up = weights[uk].float()
|
||||
break
|
||||
if down is None:
|
||||
continue
|
||||
|
||||
# Alpha scale: alpha/rank if present, else 1.0
|
||||
alpha_key = base_key + ".alpha"
|
||||
rank = down.shape[0]
|
||||
alpha_scale = float(weights[alpha_key]) / rank if alpha_key in weights else 1.0
|
||||
|
||||
# Text-encoder keys use strength_clip; everything else strength_model.
|
||||
is_te = any(f in base_key for f in ("lora_te", "lora_te1", "lora_te2"))
|
||||
strength = strength_clip if is_te else strength_model
|
||||
|
||||
# Bake strength + alpha into the matrices via sqrt-scaling so that
|
||||
# (scale*up) @ (scale*down) = strength * alpha_scale * up @ down.
|
||||
scale = (abs(strength) * alpha_scale) ** 0.5
|
||||
sign = 1.0 if strength >= 0 else -1.0
|
||||
|
||||
if base_key not in staging:
|
||||
staging[base_key] = {"downs": [], "ups": []}
|
||||
staging[base_key]["downs"].append(down * scale)
|
||||
staging[base_key]["ups"].append(up * scale * sign)
|
||||
|
||||
combined: dict = {}
|
||||
for base_key, parts in staging.items():
|
||||
downs, ups = parts["downs"], parts["ups"]
|
||||
if len(downs) == 1:
|
||||
combined[base_key + ".lora_down.weight"] = downs[0]
|
||||
combined[base_key + ".lora_up.weight"] = ups[0]
|
||||
else:
|
||||
try:
|
||||
combined[base_key + ".lora_down.weight"] = torch.cat(downs, dim=0)
|
||||
combined[base_key + ".lora_up.weight"] = torch.cat(ups, dim=1)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"[FEnodes/FEApplyLora] Merge: shape mismatch on '{base_key}' -- "
|
||||
f"falling back to first LoRA only. ({e})"
|
||||
)
|
||||
combined[base_key + ".lora_down.weight"] = downs[0]
|
||||
combined[base_key + ".lora_up.weight"] = ups[0]
|
||||
|
||||
return combined
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FEApplyLora — MODEL + FE_LORA_STACK → patched MODEL (+ optional CLIP)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class FEApplyLora:
|
||||
"""
|
||||
Applies a FE_LORA_STACK (from FELoraLoad) to a MODEL in order.
|
||||
Applies a FE_LORA_STACK to a MODEL in the chosen application mode.
|
||||
Architecture-agnostic: SD1, SDXL, Flux, WAN 2.1/2.2, HunyuanVideo, etc.
|
||||
|
||||
Optional inputs (hidden by default, enabled via right-click menu):
|
||||
clip — CLIP to patch alongside the model
|
||||
strength_scale — global multiplier applied to all per-LoRA strengths
|
||||
|
||||
Application modes
|
||||
-----------------
|
||||
Stack — each LoRA applied as a sequential patch via the model patcher.
|
||||
Safe with any combination of LoRAs. Default.
|
||||
Merge — all LoRA weight deltas are pre-scaled and summed into a single
|
||||
combined dict, then one patch is applied. Slightly cleaner at
|
||||
inference time; best when LoRAs share many of the same target layers.
|
||||
|
||||
Stack -- sequential patching, safe with any combination.
|
||||
Merge -- pre-combines all deltas into one patch, best when LoRAs share layers.
|
||||
strength_scale multiplies all per-LoRA strengths uniformly.
|
||||
lora_stack is optional -- model passes through if unconnected.
|
||||
"""
|
||||
|
||||
|
||||
APPLICATION_MODES = ["Stack", "Merge"]
|
||||
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
@@ -435,11 +153,11 @@ class FEApplyLora:
|
||||
"clip": ("CLIP",),
|
||||
"strength_scale": ("FLOAT", {
|
||||
"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01,
|
||||
"tooltip": "Global multiplier applied to all per-LoRA model and clip strengths."
|
||||
"tooltip": "Multiplies all per-LoRA model and clip strengths uniformly.",
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
RETURN_TYPES = ("MODEL", "CLIP")
|
||||
RETURN_NAMES = ("model", "clip")
|
||||
FUNCTION = "apply"
|
||||
@@ -449,367 +167,64 @@ class FEApplyLora:
|
||||
"strength_scale multiplies all per-LoRA strengths uniformly. "
|
||||
"Stack: sequential patching. Merge: pre-combines all deltas into one patch."
|
||||
)
|
||||
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, model, application_mode, lora_stack=None,
|
||||
clip=None, strength_scale=1.0):
|
||||
if not lora_stack:
|
||||
return f"empty:{application_mode}:{strength_scale}"
|
||||
return f"empty:{application_mode}:{round(float(strength_scale), 4)}"
|
||||
parts = [application_mode, str(round(float(strength_scale), 4))]
|
||||
for entry in lora_stack:
|
||||
sm = round(float(entry.get("strength_model", 1.0)), 4)
|
||||
sc = round(float(entry.get("strength_clip", sm)), 4)
|
||||
parts.append(f"{entry.get('name','')}:{sm}:{sc}")
|
||||
parts.append(f"{entry.get('name', '')}:{sm}:{sc}")
|
||||
return "|".join(parts)
|
||||
|
||||
|
||||
def apply(self, model, application_mode, lora_stack=None,
|
||||
clip=None, strength_scale=1.0):
|
||||
if not lora_stack:
|
||||
logger.info("[FEnodes/FEApplyLora] Empty stack — model unchanged.")
|
||||
logger.info("[FEnodes/FEApplyLora] No lora_stack connected -- model unchanged.")
|
||||
return (model, clip)
|
||||
|
||||
|
||||
scale = float(strength_scale)
|
||||
if scale != 1.0:
|
||||
logger.info(f"[FEnodes/FEApplyLora] Applying global strength_scale={scale}")
|
||||
|
||||
logger.info(f"[FEnodes/FEApplyLora] strength_scale={scale}")
|
||||
|
||||
if application_mode == "Merge":
|
||||
logger.info(
|
||||
f"[FEnodes/FEApplyLora] Merge mode — combining "
|
||||
f"{len(lora_stack)} LoRA(s) into one patch "
|
||||
f"(strength_scale={scale})."
|
||||
f"[FEnodes/FEApplyLora] Merge: combining {len(lora_stack)} LoRA(s) "
|
||||
f"(strength_scale={scale})"
|
||||
)
|
||||
# Scale a copy of the stack so the cached originals are untouched
|
||||
scaled_stack = [
|
||||
{**e,
|
||||
"strength_model": e["strength_model"] * scale,
|
||||
"strength_clip": e["strength_clip"] * scale}
|
||||
for e in lora_stack
|
||||
]
|
||||
merged_weights = _merge_lora_weights(scaled_stack)
|
||||
model, clip = comfy.sd.load_lora_for_models(
|
||||
model, clip, merged_weights, 1.0, 1.0
|
||||
)
|
||||
merged = merge_lora_weights(scaled_stack)
|
||||
model, clip = comfy.sd.load_lora_for_models(model, clip, merged, 1.0, 1.0)
|
||||
else:
|
||||
# Stack mode — sequential application
|
||||
for entry in lora_stack:
|
||||
sm = entry["strength_model"] * scale
|
||||
sc = entry["strength_clip"] * scale
|
||||
logger.info(
|
||||
f"[FEnodes/FEApplyLora] Stack: applying '{entry['name']}' "
|
||||
f"[FEnodes/FEApplyLora] Stack: '{entry['name']}' "
|
||||
f"(model_str={sm:.4f}, clip_str={sc:.4f})"
|
||||
)
|
||||
model, clip = comfy.sd.load_lora_for_models(
|
||||
model, clip,
|
||||
entry["weights"],
|
||||
sm, sc,
|
||||
model, clip, entry["weights"], sm, sc
|
||||
)
|
||||
|
||||
|
||||
return (model, clip)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FELoraTriggerAnalysis — encoder-agnostic trigger word analysis
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# BPE / SentencePiece tokens that score high universally and are not useful triggers.
|
||||
_STOP_TOKENS = {
|
||||
"!", '"', "#", "$", "%", "&", "'", "(", ")", "*", "+", ",", "-", ".", "/",
|
||||
":", ";", "<", "=", ">", "?", "@", "[", "\\", "]", "^", "_", "`", "{", "|",
|
||||
"}", "~", "a", "an", "the", "of", "in", "is", "it", "to", "and", "or",
|
||||
"with", "on", "at", "by", "for", "as", "be", "this", "that", "are", "was",
|
||||
"from", "its", "not", "but", "have", "he", "she", "they", "we", "you", "i",
|
||||
"</w>", "▁", "<pad>", "</s>", "<s>", "<unk>",
|
||||
"<|startoftext|>", "<|endoftext|>", "<0x00>",
|
||||
}
|
||||
|
||||
# LoRA key suffix patterns for the input-projection (down/A) half of each pair.
|
||||
_DOWN_SUFFIXES = (".lora_down.weight", ".lora_A.weight")
|
||||
|
||||
|
||||
def _discover_encoders(clip):
|
||||
"""
|
||||
Dynamically discover every text encoder sub-module present in a ComfyUI
|
||||
CLIP object, along with its token embedding table and vocabulary.
|
||||
|
||||
Returns a list of dicts:
|
||||
{
|
||||
"name": str, # attribute name, e.g. "clip_l", "t5xxl"
|
||||
"emb_table": Tensor, # (vocab_size, embed_dim) float32 CPU
|
||||
"vocab": dict[int, str], # token_id → raw token string
|
||||
"key_frags": list[str], # LoRA key fragments that target this encoder
|
||||
}
|
||||
|
||||
Strategy
|
||||
--------
|
||||
Walk every attribute of `clip.cond_stage_model`. For each nn.Module child,
|
||||
recursively search for an nn.Embedding or bare weight tensor whose shape is
|
||||
(N, D) with N > 1000 — this is the token embedding table regardless of
|
||||
whether it is called `token_embedding`, `shared`, or `embed_tokens`.
|
||||
Match the same attribute name in `clip.tokenizer` to get the vocab.
|
||||
|
||||
This approach is architecture-agnostic: CLIP-L/G (BPE, 768/1280-dim),
|
||||
T5-XXL (SentencePiece, 32128 or 256384 vocab, 4096-dim),
|
||||
LLaMA/Gemma/Qwen (SentencePiece, 32000-256384 vocab, 4096-dim) all expose
|
||||
their embedding tables as the same structural pattern.
|
||||
"""
|
||||
cond = getattr(clip, "cond_stage_model", None)
|
||||
tokenizer_root = getattr(clip, "tokenizer", None)
|
||||
if cond is None:
|
||||
return []
|
||||
|
||||
encoders = []
|
||||
|
||||
def find_embedding_weight(module):
|
||||
"""DFS through an nn.Module to find the first token embedding weight."""
|
||||
for name, child in module.named_modules():
|
||||
if isinstance(child, nn.Embedding):
|
||||
w = child.weight
|
||||
if w.dim() == 2 and w.shape[0] > 1000:
|
||||
return w.detach().float().cpu()
|
||||
# Some encoders store the embedding as a raw Parameter, not nn.Embedding
|
||||
for pname, param in child.named_parameters(recurse=False):
|
||||
if param.dim() == 2 and param.shape[0] > 1000 and pname in (
|
||||
"weight", "shared_weight", "embed_weight"
|
||||
):
|
||||
return param.detach().float().cpu()
|
||||
return None
|
||||
|
||||
def get_vocab_for(tokenizer_child):
|
||||
"""
|
||||
Extract a {token_id: token_string} vocab from whatever tokenizer object
|
||||
ComfyUI has wired up for this encoder. Four strategies tried in order:
|
||||
|
||||
1. HuggingFace get_vocab() — covers CLIP BPE, Qwen2/2.5/3, LLaMA/Gemma,
|
||||
T5 via HF wrapper, and most modern tokenizers.
|
||||
2. Wrapped HF tokenizer at .tokenizer — ComfyUI sometimes double-wraps.
|
||||
3. SentencePiece sp_model.id_to_piece() — raw T5, some LLaMA variants.
|
||||
4. tiktoken _mergeable_ranks — legacy Qwen v1 only; tiktoken Encoding
|
||||
objects are identified by their unique signature attributes rather than
|
||||
a hardcoded attribute path, so this works regardless of how deeply
|
||||
the Encoding is nested inside ComfyUI's wrapper.
|
||||
"""
|
||||
if tokenizer_child is None:
|
||||
return {}
|
||||
|
||||
# ── Strategy 1: HuggingFace get_vocab() ──────────────────────────
|
||||
if hasattr(tokenizer_child, "get_vocab"):
|
||||
try:
|
||||
v = tokenizer_child.get_vocab()
|
||||
return {idx: tok for tok, idx in v.items()}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ── Strategy 2: HF tokenizer wrapped at .tokenizer ───────────────
|
||||
inner = getattr(tokenizer_child, "tokenizer", None)
|
||||
if inner is not None and hasattr(inner, "get_vocab"):
|
||||
try:
|
||||
v = inner.get_vocab()
|
||||
return {idx: tok for tok, idx in v.items()}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ── Strategy 3: SentencePiece sp_model ───────────────────────────
|
||||
for candidate in (tokenizer_child, inner):
|
||||
if candidate is None:
|
||||
continue
|
||||
sp = getattr(candidate, "sp_model", None)
|
||||
if sp is not None:
|
||||
try:
|
||||
return {i: sp.id_to_piece(i) for i in range(sp.get_piece_size())}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ── Strategy 4: tiktoken (legacy Qwen v1) ────────────────────────
|
||||
# tiktoken Encoding objects have both `_mergeable_ranks` ({bytes: int})
|
||||
# and `_special_tokens` ({str: int}). We search recursively by these
|
||||
# signature attributes rather than a fixed attribute path.
|
||||
def _find_tiktoken(obj, _seen=None, _depth=0):
|
||||
if _depth > 5:
|
||||
return None
|
||||
if _seen is None:
|
||||
_seen = set()
|
||||
obj_id = id(obj)
|
||||
if obj_id in _seen:
|
||||
return None
|
||||
_seen.add(obj_id)
|
||||
if (hasattr(obj, "_mergeable_ranks") and
|
||||
hasattr(obj, "_special_tokens") and
|
||||
isinstance(getattr(obj, "_mergeable_ranks", None), dict)):
|
||||
return obj
|
||||
for attr in ("tokenizer", "enc", "encoding", "bpe", "inner", "_tokenizer"):
|
||||
child = getattr(obj, attr, None)
|
||||
if child is not None:
|
||||
result = _find_tiktoken(child, _seen, _depth + 1)
|
||||
if result is not None:
|
||||
return result
|
||||
return None
|
||||
|
||||
tik_enc = _find_tiktoken(tokenizer_child)
|
||||
if tik_enc is not None:
|
||||
try:
|
||||
vocab = {}
|
||||
# _mergeable_ranks: {token_bytes: token_id}
|
||||
for tok_bytes, tok_id in tik_enc._mergeable_ranks.items():
|
||||
vocab[tok_id] = tok_bytes.decode("utf-8", errors="replace")
|
||||
# _special_tokens: {token_str: token_id}
|
||||
for tok_str, tok_id in tik_enc._special_tokens.items():
|
||||
vocab[tok_id] = tok_str
|
||||
logger.info(
|
||||
f"[FEnodes/TriggerAnalysis] tiktoken vocab extracted: "
|
||||
f"{len(vocab)} tokens"
|
||||
)
|
||||
return vocab
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"[FEnodes/TriggerAnalysis] tiktoken extraction failed: {e}"
|
||||
)
|
||||
|
||||
logger.warning(
|
||||
"[FEnodes/TriggerAnalysis] Could not extract vocab for a discovered "
|
||||
"encoder — token IDs will appear as <token_N> placeholders."
|
||||
)
|
||||
return {}
|
||||
|
||||
# Walk the top-level children of cond_stage_model
|
||||
for enc_name, enc_module in cond.named_children():
|
||||
emb = find_embedding_weight(enc_module)
|
||||
if emb is None:
|
||||
continue # This child has no usable embedding table
|
||||
|
||||
# Match tokenizer sub-object by the same attribute name
|
||||
tok_child = getattr(tokenizer_root, enc_name, None)
|
||||
vocab = get_vocab_for(tok_child)
|
||||
|
||||
# Build LoRA key fragments: keys in the weights dict that target this encoder.
|
||||
# Common conventions: lora_te1_, lora_te2_, lora_te_, lora_{enc_name}_
|
||||
key_frags = [
|
||||
f"lora_te_{enc_name}_", # e.g. lora_te_clip_l_
|
||||
f"lora_{enc_name}_", # e.g. lora_clip_l_
|
||||
f"lora_te1_", # Kohya SD1 / SDXL first encoder
|
||||
f"lora_te2_", # Kohya SDXL second encoder
|
||||
f"lora_te_", # generic fallback
|
||||
enc_name, # bare name fragment
|
||||
]
|
||||
|
||||
logger.info(
|
||||
f"[FEnodes/TriggerAnalysis] Discovered encoder '{enc_name}': "
|
||||
f"vocab={emb.shape[0]}, embed_dim={emb.shape[1]}"
|
||||
)
|
||||
encoders.append({
|
||||
"name": enc_name,
|
||||
"emb_table": emb,
|
||||
"vocab": vocab,
|
||||
"key_frags": key_frags,
|
||||
})
|
||||
|
||||
return encoders
|
||||
|
||||
|
||||
def _clean_token(tok):
|
||||
"""Strip BPE/SentencePiece boundary markers and normalise to lowercase."""
|
||||
return (
|
||||
tok.replace("</w>", "")
|
||||
.replace("▁", " ")
|
||||
.replace("Ġ", " ")
|
||||
.strip()
|
||||
.lower()
|
||||
)
|
||||
|
||||
|
||||
def _analyse_encoder(enc_info, all_lora_weights, top_k):
|
||||
"""
|
||||
Run trigger analysis for one encoder against all LoRA weights.
|
||||
|
||||
For each lora_down key whose in_features matches this encoder's embed_dim,
|
||||
project every token embedding through the layer's input subspace and
|
||||
accumulate L2 activation norms. Returns a ranked list of (token, score)
|
||||
tuples, filtered for stop-tokens.
|
||||
"""
|
||||
emb_table = enc_info["emb_table"] # (vocab_size, embed_dim)
|
||||
vocab = enc_info["vocab"] # {id: raw_string}
|
||||
embed_dim = emb_table.shape[1]
|
||||
vocab_size = emb_table.shape[0]
|
||||
key_frags = enc_info["key_frags"]
|
||||
|
||||
scores = torch.zeros(vocab_size, dtype=torch.float32)
|
||||
layers_used = 0
|
||||
|
||||
for lora_entry in all_lora_weights:
|
||||
weights = lora_entry["weights"]
|
||||
|
||||
for key, tensor in weights.items():
|
||||
# Must be a lora_down / lora_A weight
|
||||
if not any(key.endswith(s) for s in _DOWN_SUFFIXES):
|
||||
continue
|
||||
# Must plausibly target this encoder
|
||||
if not any(frag in key for frag in key_frags):
|
||||
continue
|
||||
|
||||
t = tensor.detach().float().cpu()
|
||||
# Shape must be (rank, embed_dim) — 2D and inner dim matches
|
||||
if t.dim() != 2 or t.shape[1] != embed_dim:
|
||||
continue
|
||||
|
||||
# (vocab_size, embed_dim) @ (embed_dim, rank) → (vocab_size, rank)
|
||||
# Row-wise L2 norm → (vocab_size,) activation score
|
||||
scores += (emb_table @ t.T).norm(dim=1)
|
||||
layers_used += 1
|
||||
|
||||
if layers_used == 0:
|
||||
return [], 0
|
||||
|
||||
# Decode top tokens, skip stop-tokens and single chars
|
||||
sorted_ids = scores.argsort(descending=True).tolist()
|
||||
results = []
|
||||
for idx in sorted_ids:
|
||||
if len(results) >= top_k:
|
||||
break
|
||||
raw = vocab.get(idx, f"<token_{idx}>")
|
||||
cleaned = _clean_token(raw)
|
||||
if not cleaned or cleaned in _STOP_TOKENS or len(cleaned) < 2:
|
||||
continue
|
||||
results.append((cleaned, float(scores[idx])))
|
||||
|
||||
return results, layers_used
|
||||
|
||||
|
||||
|
||||
|
||||
class FELoraTriggerAnalysis:
|
||||
"""
|
||||
Analyses LoRA weight deltas against every text encoder present in the
|
||||
wired CLIP object to surface vocabulary tokens most likely to function
|
||||
as trigger words.
|
||||
|
||||
Architecture-agnostic
|
||||
---------------------
|
||||
Encoders are discovered dynamically by walking cond_stage_model's children
|
||||
looking for token embedding tables (any weight shaped vocab_size × embed_dim
|
||||
with vocab_size > 1000). Tokenizers are matched by the same attribute name.
|
||||
This covers, without any hardcoded paths:
|
||||
|
||||
CLIP-L (SD1.5, SDXL, Flux, WAN) — BPE, 49408 × 768
|
||||
CLIP-G (SDXL, SD3) — BPE, 49408 × 1280
|
||||
T5-XXL (SD3, Flux, WAN, PixArt) — SentencePiece, 32128/256384 × 4096
|
||||
LLaMA/Gemma (HunyuanVideo, LTX, Lumina) — SentencePiece, 32000+ × 4096
|
||||
Dual/triple encoders handled automatically — results labelled per encoder.
|
||||
|
||||
Method
|
||||
------
|
||||
For each text-encoder layer in the LoRA whose in_features matches a
|
||||
discovered encoder's embed_dim, project the full token embedding table
|
||||
through the lora_down input subspace and accumulate L2 activation norms.
|
||||
High-scoring tokens are the ones most aligned with what the LoRA was
|
||||
trained to respond to.
|
||||
|
||||
Output
|
||||
------
|
||||
A single STRING listing candidates from all encoders, formatted as:
|
||||
[clip_l] token_a, token_b, token_c ...
|
||||
[t5xxl] token_d, token_e, token_f ...
|
||||
|
||||
If only one encoder is present the label is omitted.
|
||||
Treat the output as a ranked shortlist, not a guaranteed match.
|
||||
Analyses LoRA weight deltas against every text encoder in the wired CLIP
|
||||
to surface candidate trigger words.
|
||||
Architecture-agnostic: CLIP-L/G, T5-XXL, LLaMA/Gemma and combinations.
|
||||
"""
|
||||
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
@@ -819,7 +234,7 @@ class FELoraTriggerAnalysis:
|
||||
"top_k": ("INT", {"default": 10, "min": 1, "max": 50}),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("candidate_triggers",)
|
||||
OUTPUT_NODE = True
|
||||
@@ -830,71 +245,56 @@ class FELoraTriggerAnalysis:
|
||||
"to surface candidate trigger words. Works with CLIP-L/G, T5-XXL, LLaMA/Gemma "
|
||||
"and any combination of dual/triple encoders."
|
||||
)
|
||||
|
||||
|
||||
def analyse(self, lora_stack, clip, top_k):
|
||||
encoders = _discover_encoders(clip)
|
||||
|
||||
encoders = discover_encoders(clip)
|
||||
|
||||
if not encoders:
|
||||
logger.warning(
|
||||
"[FEnodes/TriggerAnalysis] No token embedding tables found in this CLIP object. "
|
||||
"Cannot perform analysis."
|
||||
)
|
||||
logger.warning("[FEnodes/TriggerAnalysis] No token embedding tables found.")
|
||||
return ("",)
|
||||
|
||||
|
||||
logger.info(
|
||||
f"[FEnodes/TriggerAnalysis] Found {len(encoders)} encoder(s): "
|
||||
f"{[e['name'] for e in encoders]}"
|
||||
)
|
||||
|
||||
output_sections = []
|
||||
|
||||
sections = []
|
||||
multi = len(encoders) > 1
|
||||
|
||||
|
||||
for enc in encoders:
|
||||
results, layers_used = _analyse_encoder(enc, lora_stack, top_k)
|
||||
|
||||
results, layers_used = analyse_encoder(enc, lora_stack, top_k)
|
||||
|
||||
if layers_used == 0:
|
||||
logger.info(
|
||||
f"[FEnodes/TriggerAnalysis] '{enc['name']}': "
|
||||
"no matching LoRA layers found for this encoder — skipping."
|
||||
)
|
||||
logger.info(f"[FEnodes/TriggerAnalysis] '{enc['name']}': no matching layers.")
|
||||
continue
|
||||
|
||||
|
||||
logger.info(
|
||||
f"[FEnodes/TriggerAnalysis] '{enc['name']}': "
|
||||
f"aggregated {layers_used} layer(s), top candidates: "
|
||||
f"{[t for t, _ in results[:5]]}"
|
||||
f"{layers_used} layer(s), top: {[t for t, _ in results[:5]]}"
|
||||
)
|
||||
|
||||
|
||||
if not results:
|
||||
continue
|
||||
|
||||
|
||||
tokens_str = ", ".join(t for t, _ in results)
|
||||
if multi:
|
||||
output_sections.append(f"[{enc['name']}] {tokens_str}")
|
||||
else:
|
||||
output_sections.append(tokens_str)
|
||||
|
||||
if not output_sections:
|
||||
sections.append(f"[{enc['name']}] {tokens_str}" if multi else tokens_str)
|
||||
|
||||
if not sections:
|
||||
logger.warning(
|
||||
"[FEnodes/TriggerAnalysis] No LoRA layers matched any discovered encoder. "
|
||||
"The LoRA may only target the diffusion model (UNet/DiT) rather than the "
|
||||
"text encoder, in which case trigger analysis does not apply."
|
||||
"[FEnodes/TriggerAnalysis] No LoRA layers matched any encoder. "
|
||||
"LoRA likely targets the diffusion model only."
|
||||
)
|
||||
return ("",)
|
||||
|
||||
output = "\n".join(output_sections)
|
||||
return (output,)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
return ("\n".join(sections),)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FELoraLoad": FELoraLoad,
|
||||
"FEApplyLora": FEApplyLora,
|
||||
"FELoraTriggerAnalysis": FELoraTriggerAnalysis,
|
||||
}
|
||||
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FELoraLoad": "LoRA Load ⚡",
|
||||
"FEApplyLora": "Apply LoRA ⚡",
|
||||
|
||||
Reference in New Issue
Block a user