XLOmniV2.4 Merge Recipe

XLOmniV2.4.py

import sd_mecha
from sd_mecha.hypers import Hyper
from sd_mecha.merge_methods import SameMergeSpace
from sd_mecha.extensions.merge_method import LiftFlag, convert_to_recipe
from sd_mecha.merge_space import MergeSpace
import torch
from torch import Tensor
from typing import Optional, Dict, Mapping, TypeVar
import functools
import math
import operator
sd_mecha.set_log_level()

RecipeNodeOrPath = sd_mecha.extensions.merge_method.RecipeNodeOrPath
path_to_node = sd_mecha.extensions.merge_method.path_to_node

DeltaMergeSpace = TypeVar("DeltaMergeSpace", bound=LiftFlag[MergeSpace.DELTA])

subtract = sd_mecha.merge_methods.subtract
perpendicular_component = sd_mecha.merge_methods.perpendicular_component
add_perpendicular = sd_mecha.add_perpendicular
train_difference = sd_mecha.merge_methods.train_difference
add_difference_ties = sd_mecha.add_difference_ties
add_difference = sd_mecha.merge_methods.add_difference
clamp = sd_mecha.merge_methods.clamp
cosine_add_b = sd_mecha.cosine_add_b
cosine_add_a = sd_mecha.cosine_add_a
geometric_sum = sd_mecha.geometric_sum
multiply_quotient = sd_mecha.merge_methods.multiply_quotient
rotate = sd_mecha.merge_methods.rotate
ties_sum = sd_mecha.merge_methods.ties_sum
filter_top_k = sd_mecha.merge_methods.filter_top_k
dropout = sd_mecha.merge_methods.dropout
weighted_sum = sd_mecha.merge_methods.weighted_sum
slerp = sd_mecha.merge_methods.slerp
ties_with_dare = sd_mecha.ties_with_dare
ties_sum_with_dropout = sd_mecha.merge_methods.ties_sum_with_dropout

device = "cuda"
dtype = torch.float32

from sd_mecha.hypers import Hyper

@convert_to_recipe
def train_difference_return_delta(
    a: Tensor | SameMergeSpace,
    b: Tensor | SameMergeSpace,
    c: Tensor | SameMergeSpace,
    *,
    alpha: Hyper = 2.0,
    **kwargs,
) -> Tensor | DeltaMergeSpace:
    threshold = torch.maximum(torch.abs(a - c), torch.abs(b - c))
    dissimilarity = torch.clamp(torch.nan_to_num((c - a) * (b - c) / threshold**2, nan=0), 0)

    return ((b - c) * alpha * dissimilarity)

merger = sd_mecha.RecipeMerger(
    models_dir="/mnt/tsar/ComfyUI/models/checkpoints/SDXL/",
)

grouped_diffs_compass = [
    subtract(mobius, compass, device=device, dtype=dtype),
    subtract(colorful, compass, device=device, dtype=dtype),
    subtract(animagine, compass, device=device, dtype=dtype),
    subtract(juggernaut, compass, device=device, dtype=dtype),
    subtract(epicrealismxl, compass, device=device, dtype=dtype),
]

grouped_diffs_base = [
    subtract(mobius, base, device=device, dtype=dtype),
    subtract(colorful, base, device=device, dtype=dtype),
    subtract(animagine, base, device=device, dtype=dtype),
    subtract(juggernaut, base, device=device, dtype=dtype),
    subtract(epicrealismxl, base, device=device, dtype=dtype),
]

grouped_diffs_trained = [
    train_difference_return_delta(compass, mobius, base, device=device, dtype=dtype),
    train_difference_return_delta(compass, colorful, base, device=device, dtype=dtype),
    train_difference_return_delta(compass, animagine, base, device=device, dtype=dtype),
    train_difference_return_delta(compass, juggernaut, base, device=device, dtype=dtype),
    train_difference_return_delta(compass, epicrealismxl, base, device=device, dtype=dtype),
]


recipe = ties_with_dare(compass, *grouped_diffs_compass, alpha=1.0, k=1.0, probability=0.666667, apply_median=1.0, maxiter=100, no_rescale=1.0, vote_sgn=1.0, seed=80085, device="cpu", dtype=torch.float32)
recipe2 = ties_with_dare(recipe, *grouped_diffs_base, alpha=1.0, k=1.0, probability=0.666667, apply_median=1.0, maxiter=100, no_rescale=1.0, vote_sgn=1.0, seed=69, device="cpu", dtype=torch.float32)
recipe3 = ties_with_dare(recipe2, *grouped_diffs_trained, alpha=1.0, k=1.0, probability=0.666667, apply_median=1.0, maxiter=100, no_rescale=1.0, vote_sgn=1.0, seed=21, device="cpu", dtype=torch.float32)
recipe4 = slerp(recipe3, compass, alpha=(
        sd_mecha.default("sdxl", "txt", 0.7) |
        sd_mecha.default("sdxl", "txt2", 0.7) |
        sd_mecha.default("sdxl", "unet", 0.0)
        ), device=device, dtype=dtype)

merger.merge_and_save(recipe4, output="XLOmniV2.4.safetensors")