x-stride-ComfyUI_XS_DARE_TIES/dare_ties_stack.py

import folder_paths

from .dare_ties_profile import PROFILE_TYPE

STACK_TYPE = "DARE_TIES_MODEL_STACK"


def checkpoint_list_with_none():
    return ["None"] + folder_paths.get_filename_list("checkpoints")


class DARETIESModelStack:
    """
    Chainable 3-slot checkpoint stack.

    This node does not load or merge anything.
    It only builds a list of child checkpoints plus model/clip weights.

    Chain multiple stack nodes for 3 / 6 / 9 / 12... models.
    """

    DESCRIPTION = (
        "Builds a chainable list of SDXL/Illustrious-derived child checkpoints with per-slot "
        "model and CLIP weights and optional per-slot bezier UNet profiles.\n"
        "\n"
        "Profile mapping along the SDXL UNet (left to right): 9 IN (composition) -> 3 MID "
        "(global coherence) -> 9 OUT (style/details). A profile multiplies each child's delta "
        "per block.\n"
        "\n"
        "No checkpoints are loaded here; loading and merging happen in DARE TIES Merge Models. "
        "Chain multiple stack nodes for 3/6/9/12+ children. Disabled, 'None', or zero-zero slots "
        "are dropped.\n"
        "\n"
        "Stack order is preserved in the list sent to Apply: slot index i is used in "
        "stable_key_seed(seed, key, i) for DARE and affects TIES sign election, so reordering "
        "slots changes merge output when using dare_delta, ties_delta, or dare_ties_delta."
    )

    @classmethod
    def INPUT_TYPES(cls):
        ckpts = checkpoint_list_with_none()
        weight_cfg = {
            "default": 1.0,
            "min": 0.0,
            "max": 10.0,
            "step": 0.01,
            "tooltip": (
                "Pre-normalization weight for this child. With normalize_*_weights=True in the "
                "Apply node, weights are rescaled so they sum to 1 across all stack entries."
            ),
        }
        clip_weight_cfg = dict(weight_cfg)
        clip_weight_cfg["tooltip"] = (
            "Pre-normalization weight for this child's CLIP/text-encoder delta. Only used when "
            "merge_clip is not 'off' in the Apply node."
        )
        ckpt_tooltip = (
            "Child checkpoint to mix into the reference. Set to 'None' to skip this slot."
        )
        enabled_tooltip = (
            "Disable to skip this slot without changing its values. Disabled slots contribute "
            "nothing to the merge."
        )

        return {
            "required": {
                "ckpt_name_1": (ckpts, {"tooltip": ckpt_tooltip}),
                "model_weight_1": ("FLOAT", weight_cfg),
                "clip_weight_1": ("FLOAT", clip_weight_cfg),
                "enabled_1": ("BOOLEAN", {"default": True, "tooltip": enabled_tooltip}),

                "ckpt_name_2": (ckpts, {"tooltip": ckpt_tooltip}),
                "model_weight_2": ("FLOAT", weight_cfg),
                "clip_weight_2": ("FLOAT", clip_weight_cfg),
                "enabled_2": ("BOOLEAN", {"default": True, "tooltip": enabled_tooltip}),

                "ckpt_name_3": (ckpts, {"tooltip": ckpt_tooltip}),
                "model_weight_3": ("FLOAT", weight_cfg),
                "clip_weight_3": ("FLOAT", clip_weight_cfg),
                "enabled_3": ("BOOLEAN", {"default": True, "tooltip": enabled_tooltip}),
            },
            "optional": {
                "model_stack": (
                    STACK_TYPE,
                    {"tooltip": "Optional upstream stack to extend. Entries from this input come first in the output list."},
                ),
                "profile_1": (
                    PROFILE_TYPE,
                    {"tooltip": "Optional UNet block curve for slot 1 (21 samples). Multiplies that child's delta per block."},
                ),
                "profile_2": (
                    PROFILE_TYPE,
                    {"tooltip": "Optional UNet block curve for slot 2."},
                ),
                "profile_3": (
                    PROFILE_TYPE,
                    {"tooltip": "Optional UNet block curve for slot 3."},
                ),
            },
        }

    RETURN_TYPES = (STACK_TYPE,)
    RETURN_NAMES = ("model_stack",)
    OUTPUT_TOOLTIPS = (
        "List of child checkpoints with model_weight, clip_weight, and optional profile for UNet shaping.",
    )
    FUNCTION = "make_stack"
    CATEGORY = "advanced/model_merging/dare_ties"

    def make_stack(
        self,
        ckpt_name_1,
        model_weight_1,
        clip_weight_1,
        enabled_1,
        ckpt_name_2,
        model_weight_2,
        clip_weight_2,
        enabled_2,
        ckpt_name_3,
        model_weight_3,
        clip_weight_3,
        enabled_3,
        model_stack=None,
        profile_1=None,
        profile_2=None,
        profile_3=None,
    ):
        stack = []

        if model_stack is not None:
            stack.extend(list(model_stack))

        slots = [
            (ckpt_name_1, model_weight_1, clip_weight_1, enabled_1, profile_1),
            (ckpt_name_2, model_weight_2, clip_weight_2, enabled_2, profile_2),
            (ckpt_name_3, model_weight_3, clip_weight_3, enabled_3, profile_3),
        ]

        for ckpt_name, model_weight, clip_weight, enabled, profile in slots:
            if not enabled:
                continue

            if ckpt_name is None or ckpt_name == "None":
                continue

            if float(model_weight) == 0.0 and float(clip_weight) == 0.0:
                continue

            entry = {
                "ckpt_name": ckpt_name,
                "model_weight": float(model_weight),
                "clip_weight": float(clip_weight),
            }
            if profile is not None:
                entry["profile"] = profile
            stack.append(entry)

        return (stack,)