289 lines
9.7 KiB
Python
289 lines
9.7 KiB
Python
# iamccs_ltx2_lora_stack_segmented6.py
|
||
# ===============================================================
|
||
# Segmented LoRA stacks for workflows with 3 segments × 2 stages.
|
||
# Outputs either 6 LoRA stacks or 6 MODELs (MODEL only, no CLIP).
|
||
# ===============================================================
|
||
|
||
import logging
|
||
from typing import Any, Dict, Optional
|
||
|
||
import comfy.sd
|
||
import comfy.utils
|
||
import folder_paths
|
||
|
||
from .iamccs_ltx2_lora_stack import SuppressLTX2MissingKeysFilter, standardize_ltx2_lora_keys
|
||
|
||
|
||
def _load_lora_state_dict(name: str, cache: Dict[str, Any]) -> Optional[dict]:
|
||
if not name or name == "no":
|
||
return None
|
||
if name in cache:
|
||
return cache[name]
|
||
|
||
path = folder_paths.get_full_path_or_raise("loras", name)
|
||
sd = comfy.utils.load_torch_file(path, safe_load=True)
|
||
sd = standardize_ltx2_lora_keys(sd)
|
||
cache[name] = sd
|
||
return sd
|
||
|
||
|
||
def _append_lora(stack: list, name: str, strength: float, cache: Dict[str, Any]) -> None:
|
||
if not name or name == "no":
|
||
return
|
||
s = float(strength)
|
||
if s == 0.0:
|
||
return
|
||
|
||
sd = _load_lora_state_dict(name, cache)
|
||
if not sd:
|
||
return
|
||
|
||
stack.append({"name": name, "strength": s, "state_dict": sd})
|
||
|
||
|
||
def _build_segment_stage_stack(
|
||
*,
|
||
fixed_lora: str,
|
||
fixed_strength: float,
|
||
var_lora1: str,
|
||
var1_strength: float,
|
||
var_lora2: str,
|
||
var2_strength: float,
|
||
cache: Dict[str, Any],
|
||
) -> list:
|
||
stack: list = []
|
||
_append_lora(stack, fixed_lora, fixed_strength, cache)
|
||
_append_lora(stack, var_lora1, var1_strength, cache)
|
||
_append_lora(stack, var_lora2, var2_strength, cache)
|
||
return stack
|
||
|
||
|
||
def _apply_lora_stack_to_model(model, lora_stack: list):
|
||
if not lora_stack:
|
||
return model
|
||
|
||
model_out = model
|
||
for entry in lora_stack:
|
||
sd = entry["state_dict"]
|
||
strength = float(entry["strength"])
|
||
model_out, _ = comfy.sd.load_lora_for_models(model_out, None, sd, strength, 0)
|
||
return model_out
|
||
|
||
|
||
class IAMCCS_LTX2_LoRAStackSegmented6:
|
||
"""Builds 6 LORA stacks: 3 segments × 2 stages (MODEL-only workflows)."""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
lora_list = folder_paths.get_filename_list("loras") + ["no"]
|
||
|
||
required: Dict[str, Any] = {
|
||
"fixed_lora": (lora_list, {"default": "no"}),
|
||
}
|
||
|
||
# 3 segments (0..2), each has 2 stages and 2 variable loras
|
||
for seg in range(3):
|
||
required[f"seg{seg}_var_lora1"] = (lora_list, {"default": "no"})
|
||
required[f"seg{seg}_var_lora2"] = (lora_list, {"default": "no"})
|
||
|
||
# fixed strength per stage
|
||
required[f"seg{seg}_fixed_strength_stage1"] = (
|
||
"FLOAT",
|
||
{"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
|
||
)
|
||
required[f"seg{seg}_fixed_strength_stage2"] = (
|
||
"FLOAT",
|
||
{"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
|
||
)
|
||
|
||
# var strengths per stage
|
||
for i in (1, 2):
|
||
required[f"seg{seg}_var{i}_strength_stage1"] = (
|
||
"FLOAT",
|
||
{"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
|
||
)
|
||
required[f"seg{seg}_var{i}_strength_stage2"] = (
|
||
"FLOAT",
|
||
{"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
|
||
)
|
||
|
||
return {"required": required}
|
||
|
||
RETURN_TYPES = ("LORA", "LORA", "LORA", "LORA", "LORA", "LORA")
|
||
RETURN_NAMES = (
|
||
"seg0_stage1_lora",
|
||
"seg0_stage2_lora",
|
||
"seg1_stage1_lora",
|
||
"seg1_stage2_lora",
|
||
"seg2_stage1_lora",
|
||
"seg2_stage2_lora",
|
||
)
|
||
FUNCTION = "build"
|
||
CATEGORY = "IAMCCS/LoRA"
|
||
|
||
def build(self, fixed_lora: str, **kwargs):
|
||
cache: Dict[str, Any] = {}
|
||
|
||
out: list[list] = []
|
||
for seg in range(3):
|
||
var1 = str(kwargs.get(f"seg{seg}_var_lora1") or "no")
|
||
var2 = str(kwargs.get(f"seg{seg}_var_lora2") or "no")
|
||
|
||
fixed_s1 = kwargs.get(f"seg{seg}_fixed_strength_stage1", 0.0)
|
||
fixed_s2 = kwargs.get(f"seg{seg}_fixed_strength_stage2", 0.0)
|
||
|
||
v1s1 = kwargs.get(f"seg{seg}_var1_strength_stage1", 0.0)
|
||
v1s2 = kwargs.get(f"seg{seg}_var1_strength_stage2", 0.0)
|
||
v2s1 = kwargs.get(f"seg{seg}_var2_strength_stage1", 0.0)
|
||
v2s2 = kwargs.get(f"seg{seg}_var2_strength_stage2", 0.0)
|
||
|
||
out.append(
|
||
_build_segment_stage_stack(
|
||
fixed_lora=fixed_lora,
|
||
fixed_strength=fixed_s1,
|
||
var_lora1=var1,
|
||
var1_strength=v1s1,
|
||
var_lora2=var2,
|
||
var2_strength=v2s1,
|
||
cache=cache,
|
||
)
|
||
)
|
||
out.append(
|
||
_build_segment_stage_stack(
|
||
fixed_lora=fixed_lora,
|
||
fixed_strength=fixed_s2,
|
||
var_lora1=var1,
|
||
var1_strength=v1s2,
|
||
var_lora2=var2,
|
||
var2_strength=v2s2,
|
||
cache=cache,
|
||
)
|
||
)
|
||
|
||
# Logging summary (compact)
|
||
total = sum(len(s) for s in out)
|
||
if total == 0:
|
||
logging.warning("[IAMCCS_LTX2_LoRAStackSegmented6] ⚠ No LoRA selected")
|
||
else:
|
||
logging.info(f"[IAMCCS_LTX2_LoRAStackSegmented6] ✅ Built 6 stacks ({total} active entries)")
|
||
|
||
return tuple(out)
|
||
|
||
|
||
class IAMCCS_LTX2_ModelWithLoRA_Segmented6:
|
||
"""Applies 6 stacks (3 segments × 2 stages) to a base MODEL and outputs 6 MODELs."""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
# Mirror config of stack node, but include base model
|
||
lora_list = folder_paths.get_filename_list("loras") + ["no"]
|
||
|
||
required: Dict[str, Any] = {
|
||
"model": ("MODEL",),
|
||
"fixed_lora": (lora_list, {"default": "no"}),
|
||
}
|
||
|
||
for seg in range(3):
|
||
required[f"seg{seg}_var_lora1"] = (lora_list, {"default": "no"})
|
||
required[f"seg{seg}_var_lora2"] = (lora_list, {"default": "no"})
|
||
|
||
required[f"seg{seg}_fixed_strength_stage1"] = (
|
||
"FLOAT",
|
||
{"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
|
||
)
|
||
required[f"seg{seg}_fixed_strength_stage2"] = (
|
||
"FLOAT",
|
||
{"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
|
||
)
|
||
|
||
for i in (1, 2):
|
||
required[f"seg{seg}_var{i}_strength_stage1"] = (
|
||
"FLOAT",
|
||
{"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
|
||
)
|
||
required[f"seg{seg}_var{i}_strength_stage2"] = (
|
||
"FLOAT",
|
||
{"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
|
||
)
|
||
|
||
return {"required": required}
|
||
|
||
RETURN_TYPES = ("MODEL", "MODEL", "MODEL", "MODEL", "MODEL", "MODEL")
|
||
RETURN_NAMES = (
|
||
"seg0_stage1_model",
|
||
"seg0_stage2_model",
|
||
"seg1_stage1_model",
|
||
"seg1_stage2_model",
|
||
"seg2_stage1_model",
|
||
"seg2_stage2_model",
|
||
)
|
||
FUNCTION = "apply_segmented"
|
||
CATEGORY = "IAMCCS/LoRA"
|
||
|
||
def apply_segmented(self, model, fixed_lora: str, **kwargs):
|
||
cache: Dict[str, Any] = {}
|
||
|
||
# Build all 6 stacks
|
||
stacks: list[list] = []
|
||
for seg in range(3):
|
||
var1 = str(kwargs.get(f"seg{seg}_var_lora1") or "no")
|
||
var2 = str(kwargs.get(f"seg{seg}_var_lora2") or "no")
|
||
|
||
fixed_s1 = kwargs.get(f"seg{seg}_fixed_strength_stage1", 0.0)
|
||
fixed_s2 = kwargs.get(f"seg{seg}_fixed_strength_stage2", 0.0)
|
||
|
||
v1s1 = kwargs.get(f"seg{seg}_var1_strength_stage1", 0.0)
|
||
v1s2 = kwargs.get(f"seg{seg}_var1_strength_stage2", 0.0)
|
||
v2s1 = kwargs.get(f"seg{seg}_var2_strength_stage1", 0.0)
|
||
v2s2 = kwargs.get(f"seg{seg}_var2_strength_stage2", 0.0)
|
||
|
||
stacks.append(
|
||
_build_segment_stage_stack(
|
||
fixed_lora=fixed_lora,
|
||
fixed_strength=fixed_s1,
|
||
var_lora1=var1,
|
||
var1_strength=v1s1,
|
||
var_lora2=var2,
|
||
var2_strength=v2s1,
|
||
cache=cache,
|
||
)
|
||
)
|
||
stacks.append(
|
||
_build_segment_stage_stack(
|
||
fixed_lora=fixed_lora,
|
||
fixed_strength=fixed_s2,
|
||
var_lora1=var1,
|
||
var1_strength=v1s2,
|
||
var_lora2=var2,
|
||
var2_strength=v2s2,
|
||
cache=cache,
|
||
)
|
||
)
|
||
|
||
# Apply with log suppression
|
||
logger = logging.getLogger()
|
||
missing_keys_filter = SuppressLTX2MissingKeysFilter()
|
||
logger.addFilter(missing_keys_filter)
|
||
try:
|
||
models = []
|
||
for idx, stack in enumerate(stacks):
|
||
out_model = _apply_lora_stack_to_model(model, stack)
|
||
models.append(out_model)
|
||
if stack:
|
||
names = ", ".join(f"{e['name']}({e['strength']})" for e in stack)
|
||
logging.info(f"[IAMCCS_LTX2_ModelWithLoRA_Segmented6] segStage[{idx}] -> {names}")
|
||
return tuple(models)
|
||
finally:
|
||
logger.removeFilter(missing_keys_filter)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"IAMCCS_LTX2_LoRAStackSegmented6": IAMCCS_LTX2_LoRAStackSegmented6,
|
||
"IAMCCS_LTX2_ModelWithLoRA_Segmented6": IAMCCS_LTX2_ModelWithLoRA_Segmented6,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"IAMCCS_LTX2_LoRAStackSegmented6": "LoRA Stack (LTX-2, segmented: 3 seg × 2 stages)",
|
||
"IAMCCS_LTX2_ModelWithLoRA_Segmented6": "Apply LoRA to MODEL (LTX-2, segmented: 3 seg × 2 stages)",
|
||
}
|