543 lines
27 KiB
Python
543 lines
27 KiB
Python
"""
|
||
Easy LoRA Merger Node (with Smart‑Pro Defaults)
|
||
Merges three LoRAs using advanced triple‑merge methods and auto‑formats output.
|
||
"""
|
||
|
||
import torch
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Dict, Any, List, Optional
|
||
|
||
import folder_paths
|
||
import comfy.sd
|
||
import comfy.utils
|
||
from .config import MergeConfig, PRECISION_STANDARD, DEVICE_OPTIONS, DevicePrecisionConfig, ACTIVE_THRESHOLD_DEFAULT
|
||
from .utils import (
|
||
load_lora_with_metadata,
|
||
save_safetensors_file,
|
||
get_user_output_path,
|
||
get_experiment_temp_path,
|
||
categorize_key,
|
||
NodeCache,
|
||
)
|
||
from .engine.triple_forensics import build_triple_forensic_report
|
||
from .engine.klein_normalizer import (
|
||
universal_normalize,
|
||
detect_lora_format,
|
||
infer_separator_style,
|
||
infer_naming_style,
|
||
)
|
||
from .engine.identity_normalizer import identity_normalize
|
||
from .engine.serialization_factory import finalize_for_save
|
||
from .engine.triple_methods import merge_triple_method
|
||
from .engine.methods import resolve_blend_mode_triple
|
||
from .validation import MetadataMerger
|
||
from .engine.lora_studio_converter import MusubiLoraConverter
|
||
from .engine.metadata_factory import finalize_metadata
|
||
from .engine.merge_engine_v2 import IdentityMergeEngine
|
||
|
||
|
||
class EasyLoRATripleMerger:
|
||
"""Merge THREE LoRAs at once and apply to model (experimental)"""
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
loras = folder_paths.get_filename_list("loras")
|
||
default_folder = ""
|
||
lora_folders = folder_paths.get_folder_paths("loras")
|
||
if lora_folders:
|
||
default_folder = str(lora_folders[0])
|
||
|
||
return {
|
||
"required": {
|
||
"model": ("MODEL",),
|
||
"clip": ("CLIP",),
|
||
"method": (["linear", "ties_strict", "ties_gentle", "dare_lite",
|
||
"dare_rescale", "subtract", "magnitude", "feature_mix",
|
||
"svd_preserve", "noise_aware", "gradient_alignment",
|
||
"slerp", "cross", "ties_contrast", "block_swap"], {
|
||
"default": "linear"
|
||
}),
|
||
"density": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 1.0, "step": 0.01}),
|
||
},
|
||
"optional": {
|
||
# ── SECTION 2: SOURCES ──────────────────────────────────
|
||
"lora_a": (["None"] + loras,),
|
||
"lora_b": (["None"] + loras,),
|
||
"lora_c": (["None"] + loras,),
|
||
"lora_data_a": ("LORA",),
|
||
"lora_data_b": ("LORA",),
|
||
"lora_data_c": ("LORA",),
|
||
"weight_a": ("FLOAT", {"default": 1.0, "min": -5.0, "max": 5.0, "step": 0.05,
|
||
"tooltip": "Strength of first LoRA"}),
|
||
"weight_b": ("FLOAT", {"default": 1.0, "min": -5.0, "max": 5.0, "step": 0.05,
|
||
"tooltip": "Strength of second LoRA"}),
|
||
"weight_c": ("FLOAT", {"default": 1.0, "min": -5.0, "max": 5.0, "step": 0.05,
|
||
"tooltip": "Strength of third LoRA"}),
|
||
|
||
# ── SECTION 3: SETTINGS ─────────────────────────────────
|
||
"blend_mode": (["auto", "active", "dense"], {
|
||
"default": "auto",
|
||
"tooltip": "auto: Smart choice based on trainer metadata (match → dense, mismatch → active) | dense: Traditional weighted sum"
|
||
}),
|
||
|
||
# ── SECTION 4: ADVANCED ─────────────────────────────────
|
||
"energy_preservation": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "Preserve energy distribution during merge (recommended). Disable for raw weighted sum."
|
||
}),
|
||
"balancing_mode": (["safe", "creative", "disabled", "intensity", "impact"], {
|
||
"default": "safe",
|
||
"tooltip": "Auto-weight-balancing: safe (hard-match energy on shared layers only, rank-independent), creative (preserve ratio with reduced magnitude), intensity (peak-energy detection for cross-concept merges), impact (intensity + sparsity correction for sparse-vs-dense merges like Anima), disabled (no adjustment)."
|
||
}),
|
||
"magnitude_scaling": (["none", "rms", "top_5%", "top_10%", "top_20%", "top_30%"], {
|
||
"default": "none",
|
||
"tooltip": "Signal magnitude scaling before merging – scales LoRA B and C to match LoRA A's energy using RMS or top‑X% percentile."
|
||
}),
|
||
"active_threshold": ("BOOLEAN", {"default": True,
|
||
"tooltip": "Enable active region detection (threshold value from config.py). Disable for dense-style behavior."}),
|
||
"uniqueness": ("FLOAT", {"default": 0.7, "min": 0.1, "max": 1.0, "step": 0.01,
|
||
"tooltip": "For feature_mix: higher = preserve more unique features"}),
|
||
"threshold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01,
|
||
"tooltip": "For subtract: minimum magnitude to subtract"}),
|
||
"blend": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
|
||
"tooltip": "For magnitude: 0=strict, 1=blended"}),
|
||
|
||
# ── SECTION 5: HARDWARE ─────────────────────────────────
|
||
"device": (DEVICE_OPTIONS, {"default": "auto"}),
|
||
"precision": (PRECISION_STANDARD, {"default": "auto"}),
|
||
"batch_size": ("INT", {"default": 32, "min": 1, "max": 256, "step": 8,
|
||
"tooltip": "Number of keys to process per batch. "
|
||
"DeviceManager.suggest_batch_size() can auto-tune based on VRAM."}),
|
||
"streaming": ("BOOLEAN", {"default": True,
|
||
"tooltip": "Stream tensors to save VRAM"}),
|
||
|
||
# ── SECTION 6: OUTPUT ───────────────────────────────────
|
||
"save_trigger": ("BOOLEAN", {"default": False}),
|
||
"filename": ("STRING", {"default": "triple_merged", "multiline": False}),
|
||
"save_folder": ("STRING", {"default": default_folder, "multiline": False}),
|
||
"metadata_mode": (["none", "preserve_a", "preserve_b", "merge_basic"], {
|
||
"default": "merge_basic",
|
||
"tooltip": "How to handle metadata from source LoRAs"
|
||
}),
|
||
}
|
||
}
|
||
|
||
@classmethod
|
||
def IS_CHANGED(cls, **kwargs):
|
||
"""Cache implementation for ComfyUI change detection."""
|
||
# Inline data guard: wired lora_data inputs can't be cheaply hashed
|
||
for guard in ['lora_data_a', 'lora_data_b', 'lora_data_c']:
|
||
if kwargs.get(guard) is not None:
|
||
return float("nan")
|
||
|
||
return NodeCache.is_changed(
|
||
cls.__name__,
|
||
lora_a=kwargs.get('lora_a', 'None'),
|
||
lora_b=kwargs.get('lora_b', 'None'),
|
||
lora_c=kwargs.get('lora_c', 'None'),
|
||
weight_a=kwargs.get('weight_a', 1.0),
|
||
weight_b=kwargs.get('weight_b', 1.0),
|
||
weight_c=kwargs.get('weight_c', 1.0),
|
||
method=kwargs.get('method', 'linear'),
|
||
density=kwargs.get('density', 1.0),
|
||
save_trigger=kwargs.get('save_trigger', False),
|
||
uniqueness=kwargs.get('uniqueness', 0.7),
|
||
threshold=kwargs.get('threshold', 0.0),
|
||
blend=kwargs.get('blend', 0.5),
|
||
active_threshold=kwargs.get('active_threshold', 1e-8),
|
||
magnitude_scaling=kwargs.get('magnitude_scaling', 'none'),
|
||
blend_mode=kwargs.get('blend_mode', 'auto'),
|
||
precision=kwargs.get('precision', 'auto'),
|
||
device=kwargs.get('device', 'auto'),
|
||
batch_size=kwargs.get('batch_size', 32),
|
||
streaming=kwargs.get('streaming', True),
|
||
balancing_mode=kwargs.get('balancing_mode', 'safe'),
|
||
energy_preservation=kwargs.get('energy_preservation', True),
|
||
metadata_mode=kwargs.get('metadata_mode', 'merge_basic'),
|
||
)
|
||
|
||
|
||
RETURN_TYPES = ("LORA", "MODEL", "CLIP", "STRING", "STRING")
|
||
RETURN_NAMES = ("lora", "model", "clip", "output_path", "forensic_report")
|
||
FUNCTION = "merge_triple"
|
||
CATEGORY = "LoRA/Experimental"
|
||
|
||
def merge_triple(self, model, clip, method="linear", density=1.0,
|
||
lora_a="None", lora_b="None", lora_c="None",
|
||
lora_data_a=None, lora_data_b=None, lora_data_c=None,
|
||
weight_a=1.0, weight_b=1.0, weight_c=1.0,
|
||
uniqueness=0.7, threshold=0.0, blend=0.5, blend_mode="auto",
|
||
active_threshold=True,
|
||
magnitude_scaling="none",
|
||
save_trigger=False, save_folder="", filename="triple_merged",
|
||
device="auto", precision="auto", batch_size=32,
|
||
streaming=True, energy_preservation=True, balancing_mode="safe", metadata_mode="merge_basic"):
|
||
|
||
print("\n" + "="*50)
|
||
print("🎨 Easy LoRA Merger (with Smart‑Pro Defaults)")
|
||
print("="*50)
|
||
|
||
# Hidden studio knobs – smart‑pro defaults
|
||
target_format = "auto"
|
||
keep_alphas = False
|
||
bake_custom_scale = 1.0
|
||
|
||
# ══ Runtime precision guard — protect against stale workflow JSONs ══
|
||
if precision not in PRECISION_STANDARD:
|
||
print(f" ⚠️ Precision '{precision}' is no longer available, falling back to 'auto'")
|
||
precision = "auto"
|
||
|
||
# Resolve active_threshold bool to actual float threshold value
|
||
# True → use the default threshold from config.py
|
||
# False → use 0.0 (all non-zero values treated as active, effectively disabling)
|
||
active_threshold_float = ACTIVE_THRESHOLD_DEFAULT if active_threshold else 0.0
|
||
|
||
# Create config with only the variables we have
|
||
config = MergeConfig(
|
||
method=method,
|
||
density=density,
|
||
weight_a=weight_a,
|
||
weight_b=weight_b,
|
||
uniqueness=uniqueness,
|
||
threshold=threshold,
|
||
blend=blend,
|
||
blend_mode=blend_mode,
|
||
device_type=device,
|
||
precision=precision,
|
||
batch_size=batch_size,
|
||
streaming=streaming,
|
||
energy_preservation=energy_preservation,
|
||
balancing_mode=balancing_mode,
|
||
active_threshold=active_threshold_float,
|
||
magnitude_scaling=magnitude_scaling,
|
||
metadata_mode=metadata_mode
|
||
)
|
||
|
||
# Get paths for all three LoRAs
|
||
paths = []
|
||
loras_data = [lora_data_a, lora_data_b, lora_data_c]
|
||
loras_dropdown = [lora_a, lora_b, lora_c]
|
||
names = ["A", "B", "C"]
|
||
|
||
two_way_mode = False
|
||
missing_allowed_for_c = False
|
||
sds = []
|
||
metas = []
|
||
for i, (data, dropdown, name) in enumerate(zip(loras_data, loras_dropdown, names)):
|
||
if data is not None:
|
||
# LORA data input: unpack state dict directly — no temp file needed
|
||
lora_dict = data[0] if isinstance(data, (tuple, list)) else data
|
||
sds.append(lora_dict)
|
||
metas.append({})
|
||
print(f"📄 {name}: LORA data (in‑memory, {len(lora_dict)} tensors)")
|
||
elif dropdown != "None" and dropdown:
|
||
path = folder_paths.get_full_path("loras", dropdown)
|
||
print(f"📄 {name}: {dropdown}")
|
||
sd, meta = load_lora_with_metadata(Path(path))
|
||
sds.append(sd)
|
||
metas.append(meta)
|
||
else:
|
||
# Missing input
|
||
if name == "C":
|
||
# Allow missing C, treat as two‑way merge
|
||
print(f"⚠️ {name}: Missing input – falling back to two‑way merge (A + B)")
|
||
two_way_mode = True
|
||
missing_allowed_for_c = True
|
||
continue
|
||
else:
|
||
# Missing A or B is fatal
|
||
print(f"❌ {name}: Missing input")
|
||
return (None, model, clip, "",
|
||
"❌ Aborted: Missing LoRA A or B — cannot proceed")
|
||
|
||
# Output path logic
|
||
if save_trigger:
|
||
output_path = get_user_output_path(save_folder, filename)
|
||
else:
|
||
output_path = get_experiment_temp_path("triple")
|
||
|
||
print("📥 Loaded LoRAs in memory")
|
||
|
||
# Extract trigger words and warn about conflicts
|
||
converter = MusubiLoraConverter()
|
||
trigger_words = []
|
||
for i, meta in enumerate(metas):
|
||
info = converter.analyze_lora_structure(sds[i], meta)
|
||
concept = info.get('concept_analysis', {})
|
||
trigger = concept.get('trigger_word')
|
||
trigger_words.append(trigger)
|
||
# Compare pairs
|
||
names = ['A', 'B', 'C']
|
||
for i in range(len(trigger_words)):
|
||
for j in range(i+1, len(trigger_words)):
|
||
if trigger_words[i] and trigger_words[j] and trigger_words[i] == trigger_words[j]:
|
||
print(f"⚠️ Warning: Both LoRA {names[i]} and LoRA {names[j]} use the same trigger word '{trigger_words[i]}'. Results may be unpredictable.")
|
||
|
||
# Check if they're all the same type
|
||
print("🔍 Detecting formats...")
|
||
formats = [detect_lora_format(sd) for sd in sds]
|
||
print(f" Formats: {formats}")
|
||
|
||
# Normalize all to common format with identity mapping
|
||
# SD1.5 Diffusers LoRA detection: bridge preserves block structure (avoids ~175/396 key loss)
|
||
from .engine.diffusers_bridge import detect_diffusers_sd15_lora, normalize_diffusers_preserving
|
||
|
||
print("🔄 Normalizing with identity mapping...")
|
||
normalized_sds = []
|
||
mappings = []
|
||
for sd, meta in zip(sds, metas):
|
||
if detect_diffusers_sd15_lora(sd):
|
||
print(" 🔄 SD1.5 Diffusers LoRA detected – using Diffusers-preserving normalization")
|
||
norm, mapping = normalize_diffusers_preserving(sd)
|
||
else:
|
||
norm, mapping = identity_normalize(sd, meta)
|
||
normalized_sds.append(norm)
|
||
mappings.append(mapping)
|
||
|
||
# Forensics data collectors
|
||
energy_by_component = None
|
||
adjusted_weights = None
|
||
warnings_list: List[str] = []
|
||
|
||
# Decide between triple and two‑way merge
|
||
if len(sds) == 2:
|
||
print("🔄 Two‑way merge (LoRA C missing) – using IdentityMergeEngine")
|
||
warnings_list.append("Two-way fallback: LoRA C was missing")
|
||
# Use IdentityMergeEngine with the same config
|
||
engine = IdentityMergeEngine(config)
|
||
merged_dict, master_map = engine.merge_with_mapping(sds[0], sds[1], metas[0], metas[1])
|
||
print(f"✅ Two‑way merge completed ({len(merged_dict)} keys)")
|
||
resolved_blend_mode = resolve_blend_mode_triple(blend_mode, metas)
|
||
else:
|
||
# Triple merge
|
||
print(f"🎯 Method: {method}")
|
||
if method == "slerp":
|
||
print(" 🔮 SLERP selected — uses |weight| for interpolation factor t; weight sign is discarded")
|
||
print(" ℹ️ SLERP acts on the 2 strongest LoRAs (|weight| > 0); falls back to dense sum for 1 or 3 active")
|
||
# Resolve UI blend_mode to internal value
|
||
resolved_blend_mode = resolve_blend_mode_triple(blend_mode, metas)
|
||
if resolved_blend_mode != blend_mode:
|
||
print(f" 🔀 blend_mode '{blend_mode}' → '{resolved_blend_mode}'")
|
||
|
||
merged_dict, master_map, energy_by_component, adjusted_weights = merge_triple_method(
|
||
normalized_sds,
|
||
[weight_a, weight_b, weight_c],
|
||
method, density, uniqueness, threshold, blend, blend_mode=resolved_blend_mode,
|
||
device=config.device_type,
|
||
precision=config.precision,
|
||
magnitude_scaling=config.magnitude_scaling,
|
||
max_scaling_factor=config.max_scaling_factor,
|
||
batch_size=config.batch_size,
|
||
streaming=config.streaming,
|
||
energy_preservation=config.energy_preservation,
|
||
balancing_mode=config.balancing_mode,
|
||
mappings=mappings,
|
||
original_sds=sds,
|
||
metas=metas
|
||
)
|
||
|
||
# ==================== SMART‑PRO DEFAULTS ====================
|
||
# Apply up‑weight scaling (if requested)
|
||
effective_scaling = abs(bake_custom_scale - 1.0) > 1e-6
|
||
if effective_scaling:
|
||
print(f"⚖️ Applying up‑weight scaling: {bake_custom_scale}x")
|
||
scaled_keys = 0
|
||
for key in list(merged_dict.keys()):
|
||
if key.endswith(".weight") and ("lora_B" in key or "lora_up" in key):
|
||
merged_dict[key] = merged_dict[key] * bake_custom_scale
|
||
scaled_keys += 1
|
||
print(f" Scaled {scaled_keys} up‑weight tensors")
|
||
|
||
# Auto‑detect target format if "auto"
|
||
effective_target_format = target_format
|
||
if target_format == "auto":
|
||
# Infer separator and naming styles from normalized keys
|
||
separator = infer_separator_style(list(merged_dict.keys()))
|
||
naming = infer_naming_style(list(merged_dict.keys()))
|
||
# Determine ecosystem (use metadata from first LoRA as proxy)
|
||
target_ecosystem = metas[0].get('target_ecosystem', 'Unknown')
|
||
primary_structure = metas[0].get('primary_structure', '')
|
||
# Decision logic (same as Musubi converter)
|
||
if target_ecosystem in ("Flux.1-Dev/S", "Z-Image"):
|
||
# Flux/Z‑Image models benefit from forge_optimized (prefix mapping)
|
||
if separator == "underscore" and naming == "lora_down_up":
|
||
effective_target_format = "forge_optimized"
|
||
else:
|
||
# Fallback to comfy_native for dot separator
|
||
effective_target_format = "comfy_native"
|
||
elif separator == "dot" and naming == "lora_a_b":
|
||
effective_target_format = "comfy_native"
|
||
else:
|
||
# Default to standard_webui (underscore + lora_down_up)
|
||
effective_target_format = "standard_webui"
|
||
print(f"🔍 Auto‑selected target format: {effective_target_format} (separator={separator}, naming={naming}, ecosystem={target_ecosystem})")
|
||
|
||
# Legacy‑aware alpha baking
|
||
bake_alphas_flag = not keep_alphas
|
||
if not keep_alphas:
|
||
# User left checkbox unchecked → apply smart default based on model type
|
||
target_ecosystem = metas[0].get('target_ecosystem', 'Unknown')
|
||
if target_ecosystem in ("Flux.1-Dev/S", "Z-Image"):
|
||
bake_alphas_flag = True # bake alphas for Flux/Z‑Image
|
||
else:
|
||
bake_alphas_flag = False # keep alphas for SDXL/SD1.5
|
||
print(f"📝 Legacy‑aware alpha baking: alphas will be {'baked' if bake_alphas_flag else 'kept'} (detected {target_ecosystem})")
|
||
|
||
# Handle alpha baking BEFORE finalize (separate from format conversion)
|
||
if bake_alphas_flag:
|
||
from .utils import bake_alphas as bake_alphas_func
|
||
# infer_naming_style already imported at module level (line 28)
|
||
baking_naming_style = infer_naming_style(list(merged_dict.keys()))
|
||
merged_dict = bake_alphas_func(merged_dict, naming_style=baking_naming_style)
|
||
# Remove baked alpha keys from master_map
|
||
master_map = {k: v for k, v in master_map.items() if not k.endswith('.alpha')}
|
||
print(f"🔥 Alpha keys baked into up‑weights ({len(merged_dict)} tensors)")
|
||
|
||
# Use identity mapping to restore original keys (avoids fragile pattern matching)
|
||
restored_dict, _ = finalize_for_save(
|
||
merged_dict, master_map,
|
||
target_format=effective_target_format,
|
||
meta_a=metas[0] if metas else None,
|
||
meta_b=metas[1] if metas and len(metas) > 1 else None
|
||
)
|
||
|
||
# Apply forge_optimized prefix mapping (Flux DiT blocks: diffusion_model. → transformer.)
|
||
if effective_target_format == 'forge_optimized':
|
||
forge_count = 0
|
||
for key in list(restored_dict.keys()):
|
||
new_key = key
|
||
if ('double_blocks' in new_key or 'single_blocks' in new_key):
|
||
if new_key.startswith('diffusion_model.'):
|
||
new_key = 'transformer.' + new_key[len('diffusion_model.'):]
|
||
forge_count += 1
|
||
if new_key != key:
|
||
restored_dict[new_key] = restored_dict.pop(key)
|
||
if forge_count:
|
||
print(f" 🔧 Applied forge prefix mapping to {forge_count} Flux keys")
|
||
|
||
converted_dict = restored_dict
|
||
print(f"🎯 Converted to {effective_target_format}: {len(converted_dict)} keys")
|
||
|
||
# Merge metadata using unified metadata_factory
|
||
merged_metadata = {}
|
||
if metas and len(metas) >= 1:
|
||
metas_nonnull = [m or {} for m in metas]
|
||
|
||
if metadata_mode == "merge_basic":
|
||
# Chain merge using MetadataMerger to combine multiple LoRA sources,
|
||
# then finalize_metadata handles scrub + sign.
|
||
merged = MetadataMerger.merge(metas_nonnull[0], metas_nonnull[1] if len(metas_nonnull) > 1 else {}, mode="merge_basic")
|
||
if len(metas_nonnull) > 2:
|
||
merged = MetadataMerger.merge(merged, metas_nonnull[2], mode="merge_basic")
|
||
print(f"📝 Metadata mode 'merge_basic' – merged metadata from {len(metas_nonnull)} LoRAs ({len(merged)} entries)")
|
||
|
||
# Extract trigger words
|
||
converter = MusubiLoraConverter()
|
||
trigger_words = []
|
||
for meta in metas_nonnull:
|
||
flat = converter._parse_metadata(meta)
|
||
triggers = converter._extract_trigger_words(flat)
|
||
if triggers:
|
||
trigger_words.append(triggers)
|
||
if trigger_words:
|
||
all_triggers = set()
|
||
for tw in trigger_words:
|
||
if isinstance(tw, tuple):
|
||
words = tw[0]
|
||
elif isinstance(tw, list):
|
||
words = tw
|
||
else:
|
||
words = [p.strip() for p in tw.split(',') if p.strip()]
|
||
all_triggers.update(words)
|
||
if all_triggers:
|
||
merged["trigger_words"] = ', '.join(sorted(all_triggers))
|
||
|
||
merged_metadata = finalize_metadata(
|
||
metadata=merged,
|
||
mode=metadata_mode,
|
||
component="merger",
|
||
extra_fields={"merged_key_count": str(len(merged_dict))}
|
||
)
|
||
else:
|
||
# none / preserve_a / preserve_b — handled directly by finalize_metadata
|
||
merged_metadata = finalize_metadata(
|
||
metadata=metas_nonnull[0],
|
||
mode=metadata_mode,
|
||
component="merger",
|
||
extra_fields={"merged_key_count": str(len(merged_dict))}
|
||
)
|
||
|
||
# Convert output LoRA tensors to user-chosen dtype (merge computation
|
||
# stays in bfloat16 for stability, but the output honors the precision parameter)
|
||
target_dtype = config.device_precision.dtype
|
||
if target_dtype is not None:
|
||
converted_count = 0
|
||
for key in list(converted_dict.keys()):
|
||
t = converted_dict[key]
|
||
if t.dtype != target_dtype:
|
||
converted_dict[key] = t.to(target_dtype)
|
||
converted_count += 1
|
||
if converted_count:
|
||
print(f"🎯 Converted {converted_count} output tensors to {target_dtype}")
|
||
|
||
# Finalize and save
|
||
if save_trigger:
|
||
output_path, save_path_str = save_safetensors_file(
|
||
converted_dict, save_folder, filename, metadata=merged_metadata,
|
||
folder_type="loras", default_name="triple_merged",
|
||
)
|
||
if output_path is None:
|
||
print(f" ⚠️ Save failed: {save_path_str}")
|
||
else:
|
||
print("🧠 Preview mode: LoRA kept in RAM, no file written")
|
||
|
||
# Load into model (always from RAM - converted_dict is still in memory)
|
||
try:
|
||
lora_data = converted_dict
|
||
print("🧠 Using merged LoRA from RAM")
|
||
model_lora, clip_lora = comfy.sd.load_lora_for_models(
|
||
model, clip, lora_data, 1.0, 1.0
|
||
)
|
||
lora_tuple = (lora_data, 1.0, 1.0)
|
||
except Exception as e:
|
||
print(f"❌ Load failed: {e}")
|
||
error_report = f"❌ Load failed after merge: {e}"
|
||
return (None, model, clip, "", "❌ Load failed after merge: " + str(e))
|
||
|
||
# Determine save path for return value
|
||
if save_trigger and output_path is not None:
|
||
final_save_path = str(output_path)
|
||
else:
|
||
final_save_path = ""
|
||
|
||
# Collect density warning
|
||
if density < 1.0:
|
||
warnings_list.append(f"Density < 1.0 ({density}): sparse selection applied — individual key magnitudes may differ")
|
||
|
||
# Build forensic report for success
|
||
active_sds = [normalized_sds[i] for i in range(len(normalized_sds))]
|
||
lora_labels = ["A", "B", "C"][:len(sds)]
|
||
active_weights = [weight_a, weight_b, weight_c][:len(sds)]
|
||
|
||
forensic_report = build_triple_forensic_report(
|
||
method=method,
|
||
density=density,
|
||
blend_mode=blend_mode,
|
||
resolved_blend_mode=resolved_blend_mode,
|
||
balancing_mode=balancing_mode,
|
||
weights=active_weights,
|
||
lora_names=lora_labels,
|
||
two_way_mode=(len(sds) == 2),
|
||
active_lora_count=len(sds),
|
||
merged_dict=merged_dict,
|
||
original_sds=active_sds,
|
||
energy_by_component=energy_by_component,
|
||
adjusted_weights=adjusted_weights,
|
||
warnings=warnings_list if warnings_list else None,
|
||
)
|
||
|
||
print("="*50)
|
||
print("🎉 Triple Merge Complete!")
|
||
print("="*50)
|
||
|
||
return (lora_tuple, model_lora, clip_lora, final_save_path, forensic_report) |