import json import logging import math import os import re from collections import OrderedDict, defaultdict from hashlib import md5 import comfy.lora import comfy.lora_convert import comfy.sd import comfy.utils import folder_paths import torch from aiohttp import web from server import PromptServer LORA_TYPES = ["general", "action", "style", "quality", "character"] LORA_STACK_TYPE = "FUNPACK_LORA_STACK" TRANSFORMER_BLOCK_PATTERN = re.compile(r"(?:^|\.)transformer_blocks\.(\d+)\.") LTX_IMAGE_MODELS = {"ltxv", "ltxav"} LORA_RAW_CACHE_SIZE = 12 LORA_PATCH_CACHE_SIZE = 24 LORA_PROFILE_CACHE_SIZE = 24 LORA_BLOCK_TYPE_PROFILES = { "character": {"priority": 1.18, "yield": 0.48}, "action": {"priority": 1.12, "yield": 0.62}, "concept": {"priority": 1.12, "yield": 0.62}, "quality": {"priority": 1.04, "yield": 0.72}, "style": {"priority": 0.96, "yield": 0.96}, "general": {"priority": 0.90, "yield": 1.12}, } @PromptServer.instance.routes.get("/funpack/loras") async def funpack_loras(_): return web.json_response( ["None"] + folder_paths.get_filename_list("loras"), headers={"Cache-Control": "no-store, max-age=0"}, ) 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 prompt_key_for_v2(prompt): return re.sub(r"\s+", " ", str(prompt or "").strip()) def lora_state_id(lora_name, lora_type): return md5(f"{lora_name}::{lora_type}".encode("utf-8")).hexdigest()[:16] def normalize_lora_type(lora_type): lora_type = str(lora_type or "general").strip().lower() if lora_type == "concept": return "action" return lora_type if lora_type in LORA_TYPES else "general" 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") def refiner_v2_state_path(refinement_key): base_dir = os.path.dirname(os.path.abspath(__file__)) refinements_dir = os.path.join(base_dir, "refinements") safe_key = md5(f"clip::{refinement_key}".encode("utf-8")).hexdigest() return os.path.join(refinements_dir, f"refine_v2_{safe_key}.json") def coerce_bool(value): if isinstance(value, str): return value.strip().lower() in {"1", "true", "yes", "on"} return bool(value) def safe_float(value, fallback=1.0): try: result = float(value) except (TypeError, ValueError): return fallback return result if math.isfinite(result) else fallback def patch_target_key(patch_key): if isinstance(patch_key, tuple) and patch_key: return patch_key[0] return patch_key def transformer_block_index(patch_key): target_key = patch_target_key(patch_key) if not isinstance(target_key, str): return None match = TRANSFORMER_BLOCK_PATTERN.search(target_key) if not match: return None return int(match.group(1)) def patch_energy(value): if isinstance(value, torch.Tensor): return float(value.abs().mean().item()) weights = getattr(value, "weights", None) if weights is not None: return patch_energy(weights) if isinstance(value, dict): return sum(patch_energy(item) for item in value.values()) if isinstance(value, (list, tuple)): return sum(patch_energy(item) for item in value) return 0.0 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_list": ("STRING", {"default": "[]", "multiline": False}), "refinement_key_input": ("STRING", { "default": "", "multiline": False, "forceInput": True, "tooltip": "Optional linked refinement key, for example from FunPack Refinement Key Loader. Overrides the refinement_key widget when connected.", }), "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"}), "per_block": ( "BOOLEAN", { "default": False, "tooltip": "For LTX-mode stacks, analyze LoRA block deltas and balance competing block strengths automatically.", }, ), }, "optional": optional, } @classmethod def IS_CHANGED(cls, **kwargs): return float("nan") @classmethod def VALIDATE_INPUTS(cls, lora_list=None, lora_0=None): return True def _load_suggestions(self, refinement_key, mode, prompt_key, v2_prompt_key): for path, key, label in ( (refiner_v2_state_path(refinement_key), v2_prompt_key, "V2 refiner suggestions applied"), (refiner_state_path(refinement_key, mode), prompt_key, "legacy refiner suggestions applied"), ): if not os.path.exists(path): continue try: with open(path, "r", encoding="utf-8") as file: data = json.load(file) except (json.JSONDecodeError, OSError, ValueError): continue prompt_history = data.get("prompt_histories", {}).get(key) if not prompt_history: continue suggestions = prompt_history.get("lora_weight_suggestions", {}) if suggestions: return suggestions, label return {}, "base weights: prompt suggestions not available yet" def _entry_from_row(self, index, row): if not isinstance(row, dict): return None if not row.get("on", True): return None name = row.get("lora", row.get("name", "None")) if not name or name == "None": return None lora_type = normalize_lora_type(row.get("type", row.get("lora_type", "general"))) return { "slot": index, "name": name, "type": lora_type, "id": lora_state_id(name, lora_type), "base_model_weight": safe_float(row.get("strength", row.get("base_weight", 1.0))), } def _iter_lora_list(self, value): if value in (None, "", "[]"): return None if isinstance(value, str): try: value = json.loads(value) except (json.JSONDecodeError, TypeError, ValueError): return None if not isinstance(value, list): return None entries = [] for index, row in enumerate(value): entry = self._entry_from_row(index, row) if entry: entries.append(entry) return entries def _iter_slots(self, kwargs): listed_entries = self._iter_lora_list(kwargs.get("lora_list")) if listed_entries is not None: yield from listed_entries return 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): entry = self._entry_from_row(index, lora_name) if entry: yield entry continue shifted_base_weight = None lora_type = kwargs.get(f"lora_{index}_type", "general") if not isinstance(lora_name, str) and isinstance(lora_type, str) and lora_type not in LORA_TYPES and lora_type != "concept": shifted_base_weight = safe_float(lora_name, 1.0) lora_name = lora_type lora_type = "general" if not lora_name or lora_name == "None": continue lora_type = normalize_lora_type(lora_type) base_weight = safe_float(kwargs.get(f"lora_{index}_base_weight", 1.0), shifted_base_weight or 1.0) yield { "slot": index, "name": lora_name, "type": lora_type, "id": lora_state_id(lora_name, lora_type), "base_model_weight": base_weight, } def _get_suggestion(self, suggestions, entry): suggestion = suggestions.get(entry["id"]) if self._suggestion_matches_base(suggestion, entry): return suggestion for legacy in suggestions.values(): if ( isinstance(legacy, dict) and legacy.get("name") == entry["name"] and legacy.get("type", "general") == entry["type"] and self._suggestion_matches_base(legacy, entry) ): return legacy return {} def _suggestion_matches_base(self, suggestion, entry): if not isinstance(suggestion, dict): return False saved_base = suggestion.get("base_model_weight") if saved_base is None: return True return abs(float(saved_base) - float(entry["base_model_weight"])) <= 1e-6 def apply_lora_weights(self, positive_prompt, refinement_key, mode, per_block=False, **kwargs): linked_refinement_key = str(kwargs.pop("refinement_key_input", "") or "").strip() if linked_refinement_key: refinement_key = linked_refinement_key mode = (mode or "ltx2").lower() per_block = coerce_bool(per_block) prompt_key = prompt_key_for_mode(positive_prompt, mode) v2_prompt_key = prompt_key_for_v2(positive_prompt) suggestions, source_message = self._load_suggestions(refinement_key, mode, prompt_key, v2_prompt_key) loras = [] lines = [f"FunPack Apply LoRA Weights | {source_message}"] lines.append(f"Per-block application: {'enabled' if per_block else 'disabled'}") 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": 2, "refinement_key": refinement_key, "mode": mode, "per_block": per_block, "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.raw_lora_cache = OrderedDict() self.model_patch_cache = OrderedDict() self.block_profile_cache = OrderedDict() def _cache_get(self, cache, key): if key not in cache: return None value = cache.pop(key) cache[key] = value return value def _cache_put(self, cache, key, value, max_items): cache[key] = value cache.move_to_end(key) while len(cache) > max_items: cache.popitem(last=False) def _lora_file_cache_key(self, lora_name): lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) stat = os.stat(lora_path) mtime_ns = getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1_000_000_000)) return (lora_path, mtime_ns, stat.st_size) @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("MODEL",), "lora_stack": (LORA_STACK_TYPE,), }, "optional": { "clip": ("CLIP",), }, } def _load_lora_file(self, lora_name): cache_key = self._lora_file_cache_key(lora_name) lora = self._cache_get(self.raw_lora_cache, cache_key) if lora is None: lora = comfy.utils.load_torch_file(cache_key[0], safe_load=True) self._cache_put(self.raw_lora_cache, cache_key, lora, LORA_RAW_CACHE_SIZE) return lora, cache_key def _model_image_model(self, model): model_wrapper = getattr(model, "model", None) model_config = getattr(model_wrapper, "model_config", None) unet_config = getattr(model_config, "unet_config", None) if isinstance(unet_config, dict): return unet_config.get("image_model") return None def _per_block_requested(self, entry, lora_stack): return coerce_bool(entry.get("per_block", lora_stack.get("per_block", False))) def _per_block_supported(self, model, lora_stack, entry): if not self._per_block_requested(entry, lora_stack): return False if (lora_stack.get("mode") or "ltx2").lower() != "ltx2": return False return self._model_image_model(model) in LTX_IMAGE_MODELS def _model_cache_key(self, model): model_wrapper = getattr(model, "model", None) model_config = getattr(model_wrapper, "model_config", None) unet_config = getattr(model_config, "unet_config", None) image_model = unet_config.get("image_model") if isinstance(unet_config, dict) else None sampling = getattr(model_wrapper, "model_sampling", None) return ( id(model_wrapper), type(model_wrapper).__module__, type(model_wrapper).__name__, type(model_config).__module__ if model_config is not None else None, type(model_config).__name__ if model_config is not None else None, type(sampling).__module__ if sampling is not None else None, type(sampling).__name__ if sampling is not None else None, image_model, ) def _load_model_lora_patches(self, model, lora, lora_cache_key, model_cache_key=None): if lora_cache_key is None: key_map = comfy.lora.model_lora_keys_unet(model.model, {}) converted_lora = comfy.lora_convert.convert_lora(lora) return comfy.lora.load_lora(converted_lora, key_map) cache_key = (lora_cache_key, model_cache_key or self._model_cache_key(model)) loaded = self._cache_get(self.model_patch_cache, cache_key) if loaded is not None: return loaded key_map = comfy.lora.model_lora_keys_unet(model.model, {}) converted_lora = comfy.lora_convert.convert_lora(lora) loaded = comfy.lora.load_lora(converted_lora, key_map) self._cache_put(self.model_patch_cache, cache_key, loaded, LORA_PATCH_CACHE_SIZE) return loaded def _split_model_patches_by_block(self, loaded): global_patches = {} block_patches = defaultdict(dict) for patch_key, patch_value in loaded.items(): block_index = transformer_block_index(patch_key) if block_index is None: global_patches[patch_key] = patch_value continue block_patches[block_index][patch_key] = patch_value return global_patches, dict(block_patches) def _block_scores_from_patches(self, block_patches): block_scores = {} for block_index, patches in block_patches.items(): score = 0.0 for patch_value in patches.values(): score += patch_energy(patch_value) if score > 0.0: block_scores[block_index] = score return block_scores def _block_scales_from_scores(self, block_scores): if len(block_scores) < 2: return {} mean_score = sum(block_scores.values()) / len(block_scores) if mean_score <= 0.0: return {} scales = {} for block_index, score in block_scores.items(): ratio = max(0.0, score / mean_score) scales[block_index] = max(0.25, min(1.75, ratio ** 0.5)) return scales def _block_scales_from_patches(self, block_patches): return self._block_scales_from_scores(self._block_scores_from_patches(block_patches)) def _normalized_block_scores(self, block_scores): total = sum(block_scores.values()) if total <= 0.0: return {} return {block_index: score / total for block_index, score in block_scores.items()} def _top_block_summary(self, normalized_scores, limit=4): top_blocks = sorted(normalized_scores.items(), key=lambda item: item[1], reverse=True)[:limit] return ",".join(str(block_index) for block_index, _ in top_blocks) if top_blocks else "none" def _block_type_profile(self, entry): return LORA_BLOCK_TYPE_PROFILES.get(entry.get("type", "general"), LORA_BLOCK_TYPE_PROFILES["general"]) def _block_profile_template(self, loaded, patch_cache_key=None): if patch_cache_key is not None: cached = self._cache_get(self.block_profile_cache, patch_cache_key) if cached is not None: return cached global_patches, block_patches = self._split_model_patches_by_block(loaded) block_scores = self._block_scores_from_patches(block_patches) base_scales = self._block_scales_from_scores(block_scores) normalized_scores = self._normalized_block_scores(block_scores) if not base_scales or not normalized_scores: return None template = { "global_count": len(global_patches), "block_count": len(block_patches), "block_scores": block_scores, "base_scales": base_scales, "normalized_scores": normalized_scores, "top_blocks": self._top_block_summary(normalized_scores), "concentration": max(normalized_scores.values()) if normalized_scores else 0.0, } if patch_cache_key is not None: self._cache_put(self.block_profile_cache, patch_cache_key, template, LORA_PROFILE_CACHE_SIZE) return template def _lora_block_profile(self, entry, loaded, model_weight, patch_cache_key=None): template = self._block_profile_template(loaded, patch_cache_key) if template is None: return None type_profile = self._block_type_profile(entry) return { "entry": entry, "loaded": loaded, "model_weight": model_weight, "global_count": template["global_count"], "block_count": template["block_count"], "block_scores": template["block_scores"], "base_scales": template["base_scales"], "stack_scales": dict(template["base_scales"]), "normalized_scores": template["normalized_scores"], "top_blocks": template["top_blocks"], "concentration": template["concentration"], "priority": type_profile["priority"], "yield": type_profile["yield"], "overlap_score": 0.0, } def _pair_overlap_factor(self, entry, other_entry): lora_type = entry.get("type", "general") other_type = other_entry.get("type", "general") if "quality" in {lora_type, other_type}: return 0.72 lora_type = normalize_lora_type(lora_type) other_type = normalize_lora_type(other_type) if lora_type in {"style", "general"} and other_type in {"action", "character"}: return 1.14 if lora_type in {"action", "character"} and other_type in {"action", "character"}: return 1.08 return 1.0 def _block_presence_threshold(self, profile): block_count = max(1, len(profile["normalized_scores"])) return max(0.003, 0.40 / block_count) def _stack_block_scales(self, profiles): if len(profiles) < 2: return block_indices = sorted({block_index for profile in profiles for block_index in profile["normalized_scores"]}) for block_index in block_indices: contributors = [ profile for profile in profiles if profile["normalized_scores"].get(block_index, 0.0) >= self._block_presence_threshold(profile) ] if len(contributors) < 2: continue for profile in contributors: own_presence = profile["normalized_scores"].get(block_index, 0.0) own_weight = max(0.05, abs(profile["model_weight"])) own_signal = own_presence * own_weight * profile["priority"] other_signals = [] other_presence = 0.0 for other in contributors: if other is profile: continue factor = self._pair_overlap_factor(profile["entry"], other["entry"]) presence = other["normalized_scores"].get(block_index, 0.0) * factor other_presence += presence other_signals.append( presence * max(0.05, abs(other["model_weight"])) * other["priority"] ) strongest_other = max(other_signals) if other_signals else 0.0 if strongest_other <= 0.0: continue overlap_ratio = other_presence / max(own_presence + other_presence, 1e-9) advantage = (own_signal - strongest_other) / max(own_signal + strongest_other, 1e-9) if advantage >= 0.0: multiplier = 1.0 + min(0.18, advantage * 0.14) * min(1.0, overlap_ratio * 1.25) else: pressure = min(1.0, overlap_ratio * 1.35) damp = min(0.35, (-advantage) * 0.22 * profile["yield"] * pressure) multiplier = 1.0 - damp profile["stack_scales"][block_index] = max( 0.18, min(1.90, profile["stack_scales"].get(block_index, 1.0) * multiplier), ) profile["overlap_score"] = max(profile["overlap_score"], overlap_ratio) def _apply_model_patches(self, model, loaded, model_weight, block_scales=None): new_model = model.clone() applied = set() global_patches, block_patches = self._split_model_patches_by_block(loaded) if global_patches: applied.update(new_model.add_patches(global_patches, model_weight)) if block_scales: for block_index in sorted(block_patches): block_strength = model_weight * block_scales.get(block_index, 1.0) applied.update(new_model.add_patches(block_patches[block_index], block_strength)) else: for block_index in sorted(block_patches): applied.update(new_model.add_patches(block_patches[block_index], model_weight)) for patch_key in loaded: if patch_key not in applied: logging.warning("NOT LOADED %s", patch_key) return new_model, len(global_patches), len(block_patches) def _per_block_status(self, profile): scales = profile["stack_scales"] min_scale = min(scales.values()) max_scale = max(scales.values()) mode = "smart-per-block" if profile["overlap_score"] > 0.0 else "per-block" return ( f"{mode} blocks={len(scales)} non_block={profile['global_count']} " f"range={min_scale:.2f}..{max_scale:.2f} " f"top={profile['top_blocks']} overlap={profile['overlap_score']:.2f}" ) def _load_lora_per_block(self, model, lora, model_weight): lora_cache_key = None loaded = self._load_model_lora_patches(model, lora, lora_cache_key) profile = self._lora_block_profile({}, loaded, model_weight) if not profile: return None, "per-block fallback=global" new_model, non_block_count, _ = self._apply_model_patches( model, loaded, model_weight, block_scales=profile["stack_scales"], ) min_scale = min(profile["stack_scales"].values()) max_scale = max(profile["stack_scales"].values()) status = ( f"per-block blocks={len(profile['stack_scales'])} non_block={non_block_count} " f"range={min_scale:.2f}..{max_scale:.2f}" ) return new_model, status def load_loras(self, model, lora_stack, clip=None): loras = lora_stack.get("loras", []) if isinstance(lora_stack, dict) else [] per_block = coerce_bool(lora_stack.get("per_block", False)) if isinstance(lora_stack, dict) else False lines = [f"FunPack LoRA Loader | loading {len(loras)} LoRA(s)"] lines.append(f"Per-block application: {'enabled' if per_block else 'disabled'}") loaded_count = 0 model_cache_key = self._model_cache_key(model) prepared = [] per_block_profiles = [] 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, lora_cache_key = self._load_lora_file(entry["name"]) if self._per_block_supported(model, lora_stack, entry): patch_cache_key = (lora_cache_key, model_cache_key) loaded = self._load_model_lora_patches(model, lora, lora_cache_key, model_cache_key) profile = self._lora_block_profile(entry, loaded, model_weight, patch_cache_key) if profile is not None: item = { "entry": entry, "lora": lora, "mode": "per_block", "model_weight": model_weight, "profile": profile, } prepared.append(item) per_block_profiles.append(profile) else: prepared.append( { "entry": entry, "loaded": loaded, "mode": "patch_global", "model_weight": model_weight, "status": "per-block fallback=global", } ) elif self._per_block_requested(entry, lora_stack): prepared.append( { "entry": entry, "lora": lora, "mode": "global", "model_weight": model_weight, "status": "per-block unsupported -> global", } ) else: prepared.append( { "entry": entry, "lora": lora, "mode": "global", "model_weight": model_weight, "status": "global", } ) self._stack_block_scales(per_block_profiles) for item in prepared: entry = item["entry"] model_weight = item["model_weight"] if item["mode"] == "per_block": profile = item["profile"] model, _, _ = self._apply_model_patches( model, profile["loaded"], model_weight, block_scales=profile["stack_scales"], ) apply_status = self._per_block_status(profile) elif item["mode"] == "patch_global": apply_status = item["status"] model, _, _ = self._apply_model_patches(model, item["loaded"], model_weight) else: apply_status = item["status"] model, clip = comfy.sd.load_lora_for_models(model, clip, item["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')} mode={apply_status}" ) if loaded_count == 0: lines.append("No LoRAs were applied.") return (model, clip, lora_stack, "\n".join(lines))