301 lines
9.7 KiB
Python
301 lines
9.7 KiB
Python
import json
|
|
import os
|
|
import re
|
|
from hashlib import md5
|
|
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import folder_paths
|
|
|
|
|
|
LORA_TYPES = ["general", "concept", "style", "quality", "character"]
|
|
LORA_STACK_TYPE = "FUNPACK_LORA_STACK"
|
|
|
|
|
|
class AnyType(str):
|
|
def __ne__(self, _):
|
|
return False
|
|
|
|
|
|
any_type = AnyType("*")
|
|
|
|
|
|
class FlexibleOptionalInputType(dict):
|
|
def __init__(self, input_type, data=None):
|
|
super().__init__(data or {})
|
|
self.input_type = input_type
|
|
self.data = data or {}
|
|
|
|
def __contains__(self, _):
|
|
return True
|
|
|
|
def __getitem__(self, key):
|
|
return self.data.get(key, (self.input_type,))
|
|
|
|
def get(self, key, default=None):
|
|
return self.data.get(key, default)
|
|
|
|
def items(self):
|
|
return self.data.items()
|
|
|
|
def keys(self):
|
|
return self.data.keys()
|
|
|
|
def values(self):
|
|
return self.data.values()
|
|
|
|
def __iter__(self):
|
|
return iter(self.data)
|
|
|
|
def __len__(self):
|
|
return len(self.data)
|
|
|
|
|
|
def normalize_prompt_for_mode(prompt, mode):
|
|
prompt = (prompt or "").strip()
|
|
if (mode or "ltx2").lower() == "wan":
|
|
return re.sub(r"\s+", " ", prompt)
|
|
return prompt
|
|
|
|
|
|
def prompt_key_for_mode(prompt, mode):
|
|
if (mode or "ltx2").lower() == "wan":
|
|
return normalize_prompt_for_mode(prompt, mode)
|
|
return prompt or ""
|
|
|
|
|
|
def lora_state_id(lora_name, lora_type):
|
|
return md5(f"{lora_name}::{lora_type}".encode("utf-8")).hexdigest()[:16]
|
|
|
|
|
|
def refiner_state_path(refinement_key, mode):
|
|
base_dir = os.path.dirname(os.path.abspath(__file__))
|
|
refinements_dir = os.path.join(base_dir, "refinements")
|
|
safe_key = md5(f"{(mode or 'ltx2').lower()}::{refinement_key}".encode("utf-8")).hexdigest()
|
|
return os.path.join(refinements_dir, f"refine_{safe_key}.json")
|
|
|
|
|
|
class FunPackApplyLoraWeights:
|
|
"""
|
|
Builds a LoRA stack from user base weights, then applies prompt-specific
|
|
suggested weights previously saved by FunPack Gemma Embedding Refiner.
|
|
"""
|
|
|
|
CATEGORY = "FunPack/Model Management"
|
|
RETURN_TYPES = (LORA_STACK_TYPE, "STRING")
|
|
RETURN_NAMES = ("lora_stack", "status")
|
|
FUNCTION = "apply_lora_weights"
|
|
DESCRIPTION = "Reads Gemma Refiner prompt-specific LoRA suggestions and prepares a LoRA stack for loading."
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
loras = ["None"] + folder_paths.get_filename_list("loras")
|
|
optional = FlexibleOptionalInputType(
|
|
any_type,
|
|
{
|
|
"lora_0": (loras, {"default": "None"}),
|
|
"lora_0_type": (LORA_TYPES, {"default": "general"}),
|
|
"lora_0_base_weight": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": -10.0,
|
|
"max": 10.0,
|
|
"step": 0.01,
|
|
"tooltip": "Trainer-recommended model anchor weight.",
|
|
},
|
|
),
|
|
},
|
|
)
|
|
|
|
return {
|
|
"required": {
|
|
"positive_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"refinement_key": ("STRING", {"default": "my_style_v1", "multiline": False}),
|
|
"mode": (["ltx2", "wan"], {"default": "ltx2"}),
|
|
},
|
|
"optional": optional,
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("nan")
|
|
|
|
def _load_suggestions(self, refinement_key, mode, prompt_key):
|
|
path = refiner_state_path(refinement_key, mode)
|
|
if not os.path.exists(path):
|
|
return {}, "base weights: no refiner state file"
|
|
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as file:
|
|
data = json.load(file)
|
|
except (json.JSONDecodeError, OSError, ValueError):
|
|
return {}, "base weights: refiner state unreadable"
|
|
|
|
prompt_history = data.get("prompt_histories", {}).get(prompt_key)
|
|
if not prompt_history:
|
|
return {}, "base weights: prompt has no exact-match suggestions"
|
|
|
|
suggestions = prompt_history.get("lora_weight_suggestions", {})
|
|
if not suggestions:
|
|
return {}, "base weights: prompt suggestions not available yet"
|
|
|
|
return suggestions, "refiner suggestions applied"
|
|
|
|
def _iter_slots(self, kwargs):
|
|
indexed_slots = set()
|
|
for key in kwargs:
|
|
match = re.fullmatch(r"lora_(\d+)", key)
|
|
if match:
|
|
indexed_slots.add(int(match.group(1)))
|
|
|
|
for index in sorted(indexed_slots):
|
|
lora_name = kwargs.get(f"lora_{index}", "None")
|
|
if isinstance(lora_name, dict):
|
|
if not lora_name.get("on", True):
|
|
continue
|
|
name = lora_name.get("lora", "None")
|
|
if not name or name == "None":
|
|
continue
|
|
lora_type = lora_name.get("type", lora_name.get("lora_type", "general"))
|
|
if lora_type not in LORA_TYPES:
|
|
lora_type = "general"
|
|
yield {
|
|
"slot": index,
|
|
"name": name,
|
|
"type": lora_type,
|
|
"id": lora_state_id(name, lora_type),
|
|
"base_model_weight": float(lora_name.get("strength", lora_name.get("base_weight", 1.0))),
|
|
}
|
|
continue
|
|
|
|
if not lora_name or lora_name == "None":
|
|
continue
|
|
|
|
lora_type = kwargs.get(f"lora_{index}_type", "general")
|
|
if lora_type not in LORA_TYPES:
|
|
lora_type = "general"
|
|
|
|
yield {
|
|
"slot": index,
|
|
"name": lora_name,
|
|
"type": lora_type,
|
|
"id": lora_state_id(lora_name, lora_type),
|
|
"base_model_weight": float(kwargs.get(f"lora_{index}_base_weight", 1.0)),
|
|
}
|
|
|
|
def _get_suggestion(self, suggestions, entry):
|
|
suggestion = suggestions.get(entry["id"])
|
|
if suggestion:
|
|
return suggestion
|
|
|
|
for legacy in suggestions.values():
|
|
if (
|
|
isinstance(legacy, dict)
|
|
and legacy.get("name") == entry["name"]
|
|
and legacy.get("type", "general") == entry["type"]
|
|
):
|
|
return legacy
|
|
|
|
return {}
|
|
|
|
def apply_lora_weights(self, positive_prompt, refinement_key, mode, **kwargs):
|
|
mode = (mode or "ltx2").lower()
|
|
prompt_key = prompt_key_for_mode(positive_prompt, mode)
|
|
suggestions, source_message = self._load_suggestions(refinement_key, mode, prompt_key)
|
|
|
|
loras = []
|
|
lines = [f"FunPack Apply LoRA Weights | {source_message}"]
|
|
for entry in self._iter_slots(kwargs):
|
|
suggestion = self._get_suggestion(suggestions, entry)
|
|
model_weight = float(suggestion.get("model_weight", entry["base_model_weight"]))
|
|
source = "suggested" if suggestion else "base"
|
|
|
|
stack_entry = dict(entry)
|
|
stack_entry["model_weight"] = model_weight
|
|
stack_entry["source"] = source
|
|
loras.append(stack_entry)
|
|
|
|
lines.append(
|
|
f"lora_{entry['slot']}: {entry['name']} [{entry['type']}] "
|
|
f"{source}={model_weight:+.3f} base={entry['base_model_weight']:+.3f}"
|
|
)
|
|
|
|
stack = {
|
|
"version": 1,
|
|
"refinement_key": refinement_key,
|
|
"mode": mode,
|
|
"positive_prompt": positive_prompt,
|
|
"prompt_key": prompt_key,
|
|
"loras": loras,
|
|
}
|
|
|
|
if not loras:
|
|
lines.append("No LoRAs selected.")
|
|
|
|
return (stack, "\n".join(lines))
|
|
|
|
|
|
class FunPackLoraLoader:
|
|
"""Loads the LoRA stack prepared by FunPack Apply LoRA Weights."""
|
|
|
|
CATEGORY = "FunPack/Model Management"
|
|
RETURN_TYPES = ("MODEL", "CLIP", LORA_STACK_TYPE, "STRING")
|
|
RETURN_NAMES = ("MODEL", "CLIP", "lora_stack", "status")
|
|
FUNCTION = "load_loras"
|
|
DESCRIPTION = "Loads LoRAs from a FunPack LoRA stack without doing any learning."
|
|
|
|
def __init__(self):
|
|
self.loaded_lora = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"lora_stack": (LORA_STACK_TYPE,),
|
|
},
|
|
"optional": {
|
|
"clip": ("CLIP",),
|
|
},
|
|
}
|
|
|
|
def _load_lora_file(self, lora_name):
|
|
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
|
|
lora = None
|
|
if self.loaded_lora is not None:
|
|
if self.loaded_lora[0] == lora_path:
|
|
lora = self.loaded_lora[1]
|
|
else:
|
|
self.loaded_lora = None
|
|
|
|
if lora is None:
|
|
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
self.loaded_lora = (lora_path, lora)
|
|
|
|
return lora
|
|
|
|
def load_loras(self, model, lora_stack, clip=None):
|
|
loras = lora_stack.get("loras", []) if isinstance(lora_stack, dict) else []
|
|
lines = [f"FunPack LoRA Loader | loading {len(loras)} LoRA(s)"]
|
|
loaded_count = 0
|
|
|
|
for entry in loras:
|
|
model_weight = float(entry.get("model_weight", 0.0))
|
|
if model_weight == 0:
|
|
lines.append(f"lora_{entry.get('slot', '?')}: {entry.get('name', '?')} skipped at zero weight")
|
|
continue
|
|
|
|
lora = self._load_lora_file(entry["name"])
|
|
model, clip = comfy.sd.load_lora_for_models(model, clip, lora, model_weight, 0.0)
|
|
loaded_count += 1
|
|
lines.append(
|
|
f"lora_{entry.get('slot', '?')}: {entry['name']} "
|
|
f"applied={model_weight:+.3f} source={entry.get('source', 'base')}"
|
|
)
|
|
|
|
if loaded_count == 0:
|
|
lines.append("No LoRAs were applied.")
|
|
|
|
return (model, clip, lora_stack, "\n".join(lines))
|