133 lines
5.5 KiB
Python
133 lines
5.5 KiB
Python
import torch
|
|
from comfy.model_patcher import ModelPatcher
|
|
import uuid
|
|
import os
|
|
|
|
class SimplifiedFlux1Merge:
|
|
CATEGORY = "Flux.1"
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "run"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"interpolate_groups": ("BOOLEAN", {"default": False}),
|
|
"merge_mode": (["standard", "add_difference", "train_difference"], {"default": "standard"}),
|
|
"model_a": ("MODEL",),
|
|
"model_b": ("MODEL",),
|
|
"model_c": ("MODEL",),
|
|
|
|
"group_00_01": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 2.0, "step": 0.1}),
|
|
"group_02_03": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 2.0, "step": 0.1}),
|
|
"group_04_05": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 2.0, "step": 0.1}),
|
|
"group_06_07": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 2.0, "step": 0.1}),
|
|
"group_08_09": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 2.0, "step": 0.1}),
|
|
"group_10_11": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 2.0, "step": 0.1}),
|
|
"group_12_13": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 2.0, "step": 0.1}),
|
|
"group_14_15": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 2.0, "step": 0.1}),
|
|
"group_16_18": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 2.0, "step": 0.1}),
|
|
}
|
|
}
|
|
|
|
def run(self, merge_mode, model_a, model_b, model_c, interpolate_groups,
|
|
group_00_01, group_02_03, group_04_05, group_06_07,
|
|
group_08_09, group_10_11, group_12_13, group_14_15, group_16_18):
|
|
|
|
weights_raw = {
|
|
"group_00_01": group_00_01,
|
|
"group_02_03": group_02_03,
|
|
"group_04_05": group_04_05,
|
|
"group_06_07": group_06_07,
|
|
"group_08_09": group_08_09,
|
|
"group_10_11": group_10_11,
|
|
"group_12_13": group_12_13,
|
|
"group_14_15": group_14_15,
|
|
"group_16_18": group_16_18,
|
|
}
|
|
|
|
group_ranges = [
|
|
range(0, 2), range(2, 4), range(4, 6), range(6, 8), range(8, 10),
|
|
range(10, 12), range(12, 14), range(14, 16), range(16, 19)
|
|
]
|
|
|
|
weights = {}
|
|
keys = list(weights_raw.keys())
|
|
if interpolate_groups:
|
|
for idx, group in enumerate(group_ranges):
|
|
w_start = weights_raw[keys[idx]]
|
|
if idx + 1 < len(keys):
|
|
w_end = weights_raw[keys[idx + 1]]
|
|
else:
|
|
w_end = w_start
|
|
steps = len(group)
|
|
for i, j in enumerate(group):
|
|
weights[j] = w_start + (w_end - w_start) * (i / (steps - 1)) if steps > 1 else w_start
|
|
else:
|
|
for idx, group in enumerate(group_ranges):
|
|
for j in group:
|
|
weights[j] = weights_raw[keys[idx]]
|
|
|
|
model_a_weights = dict(model_a.model.state_dict())
|
|
model_b_weights = dict(model_b.model.state_dict())
|
|
model_c_weights = dict(model_c.model.state_dict())
|
|
|
|
|
|
merged_model = {"model": {}}
|
|
|
|
for key in model_a_weights:
|
|
if key == "model_sampling.sigmas":
|
|
continue
|
|
|
|
device = model_a.load_device
|
|
a_tensor = model_a_weights[key].to(dtype=torch.bfloat16, device=device)
|
|
b_tensor = model_b_weights.get(key, a_tensor).to(dtype=torch.bfloat16, device=device)
|
|
c_tensor = model_c_weights.get(key, a_tensor).to(dtype=torch.bfloat16, device=device)
|
|
|
|
if isinstance(a_tensor, torch.Tensor) and isinstance(b_tensor, torch.Tensor) and a_tensor.shape == b_tensor.shape:
|
|
weight = 1.0
|
|
for i, r in enumerate(group_ranges):
|
|
for j in r:
|
|
if f"double_blocks.{j}" in key or f"single_blocks.{j * 2}" in key or f"single_blocks.{j * 2 + 1}" in key:
|
|
weight = weights[j]
|
|
break
|
|
|
|
if merge_mode == "add_difference":
|
|
diff = weight * (b_tensor - a_tensor)
|
|
merged_model["model"][key] = a_tensor + diff
|
|
elif merge_mode == "train_difference":
|
|
diff = weight * (b_tensor - c_tensor)
|
|
merged_model["model"][key] = a_tensor + diff
|
|
else:
|
|
merged_model["model"][key] = a_tensor * weight + b_tensor * (1 - weight)
|
|
else:
|
|
merged_model["model"][key] = a_tensor
|
|
|
|
del a_tensor, b_tensor, c_tensor
|
|
|
|
torch.cuda.empty_cache()
|
|
|
|
del model_a_weights, model_b_weights, model_c_weights
|
|
|
|
model_class = type(model_a.model)
|
|
model_config = getattr(model_a.model, "model_config", None)
|
|
new_model = model_class(model_config=model_config)
|
|
new_model.load_state_dict(merged_model["model"], strict=False)
|
|
|
|
patched = ModelPatcher(model=new_model, load_device=model_a.load_device, offload_device=model_a.offload_device)
|
|
patched.model_hash = str(uuid.uuid4())
|
|
|
|
return (patched,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SimplifiedFlux1Merge": SimplifiedFlux1Merge
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SimplifiedFlux1Merge": "Simplified Flux.1 Merge"
|
|
}
|
|
|
|
WEB_DIRECTORY = os.path.join(os.path.dirname(__file__), "js")
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|