diff --git a/README.md b/README.md index ffd6f0c..e322504 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,6 @@ The classifier models have been taken from the sdweb-auto-MBW repo. - the resulting model will contain the text encoder and VAE sent to the node ### Bugs -- filename box doesn't use the standard comfy "prefix" method - merging process doesn't use the comfy ModelPatcher method and takes hundreds of milliseconds - - as a result, --highvram flag recommended. both models will be kept in VRAM and the process is much faster - the unet will (probably) be fp16 and the rest fp32. that's how they're sent to the node diff --git a/__init__.py b/__init__.py index 9879bdb..c8af03f 100644 --- a/__init__.py +++ b/__init__.py @@ -15,11 +15,12 @@ import nodes sys.path.append(str(pathlib.Path(__file__).parent)) import classifiers +from nodes_model_merging import CheckpointSave BLOCK_ORDER = [12, 11, 13, 10, 14, 9, 15, 8, 16, 7, 17, 6, 18, 5, 19, 4, 20, 3, 21, 2, 22, 1, 23, 0, 24] -class AutoMBW: +class AutoMBW(CheckpointSave): def __init__(self): self.type = "output" @@ -42,7 +43,7 @@ class AutoMBW: "search_depth": ("INT", {"default": 4, "min": 2}), "sample_count": ("INT", {"default": 1, "min": 1}), "classifier": (classifiers.__all__,), - "filename": ("STRING", { "multiline": False, "default": "ambw" }), + "filename_prefix": ("STRING", { "multiline": False, "default": "ambw" }), }} RETURN_TYPES = () @@ -103,8 +104,9 @@ class AutoMBW: return maximum def ambw(self, model1, model2, clip, vae, prompt, negative, search_depth, - sample_count, classifier, filename): + sample_count, classifier, filename_prefix): # python setup + self.output_dir = folder_paths.get_output_directory() self.model1 = model1 self.model2 = model2 self.vae = vae @@ -142,34 +144,7 @@ class AutoMBW: self.merge(block, self.ratios[block]) print(self.ratios) - sd1 = self.model1.model.state_dict() - vae = vae.first_stage_model.state_dict() - for key in vae: - sd1[f"first_stage_model.{key}"] = vae[key] - clip = clip.cond_stage_model.state_dict() - for key in clip: - sd1[f"cond_stage_model.{key}"] = clip[key] - - create_ckpt = False - if filename.endswith(".safetensors"): - filename = filename[0:-12] - elif filename.endswith(".ckpt"): - filename = filename[0:-5] - create_ckpt = True - - filename = pathlib.Path(folder_paths.folder_names_and_paths[ - "checkpoints"][0][0]).joinpath(f"{filename}") - print(f"saving as {filename}", end="") - if not create_ckpt: - try: - import safetensors.torch - print(".safetensors") - safetensors.torch.save_file(sd1, f"{filename}.safetensors") - except ModuleNotFoundError: - create_ckpt = True - if create_ckpt: - print(".ckpt") - torch.save(sd1, f"{filename}.ckpt") + self.save(self.model1, clip, vae, filename_prefix) return () diff --git a/nodes_model_merging.py b/nodes_model_merging.py new file mode 100644 index 0000000..bce4b3d --- /dev/null +++ b/nodes_model_merging.py @@ -0,0 +1,149 @@ +import comfy.sd +import comfy.utils +import comfy.model_base + +import folder_paths +import json +import os + +from comfy.cli_args import args + +class ModelMergeSimple: + @classmethod + def INPUT_TYPES(s): + return {"required": { "model1": ("MODEL",), + "model2": ("MODEL",), + "ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "merge" + + CATEGORY = "advanced/model_merging" + + def merge(self, model1, model2, ratio): + m = model1.clone() + kp = model2.get_key_patches("diffusion_model.") + for k in kp: + m.add_patches({k: kp[k]}, 1.0 - ratio, ratio) + return (m, ) + +class CLIPMergeSimple: + @classmethod + def INPUT_TYPES(s): + return {"required": { "clip1": ("CLIP",), + "clip2": ("CLIP",), + "ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + }} + RETURN_TYPES = ("CLIP",) + FUNCTION = "merge" + + CATEGORY = "advanced/model_merging" + + def merge(self, clip1, clip2, ratio): + m = clip1.clone() + kp = clip2.get_key_patches() + for k in kp: + if k.endswith(".position_ids") or k.endswith(".logit_scale"): + continue + m.add_patches({k: kp[k]}, 1.0 - ratio, ratio) + return (m, ) + +class ModelMergeBlocks: + @classmethod + def INPUT_TYPES(s): + return {"required": { "model1": ("MODEL",), + "model2": ("MODEL",), + "input": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "middle": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "out": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}) + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "merge" + + CATEGORY = "advanced/model_merging" + + def merge(self, model1, model2, **kwargs): + m = model1.clone() + kp = model2.get_key_patches("diffusion_model.") + default_ratio = next(iter(kwargs.values())) + + for k in kp: + ratio = default_ratio + k_unet = k[len("diffusion_model."):] + + last_arg_size = 0 + for arg in kwargs: + if k_unet.startswith(arg) and last_arg_size < len(arg): + ratio = kwargs[arg] + last_arg_size = len(arg) + + m.add_patches({k: kp[k]}, 1.0 - ratio, ratio) + return (m, ) + +class CheckpointSave: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + + @classmethod + def INPUT_TYPES(s): + return {"required": { "model": ("MODEL",), + "clip": ("CLIP",), + "vae": ("VAE",), + "filename_prefix": ("STRING", {"default": "checkpoints/ComfyUI"}),}, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},} + RETURN_TYPES = () + FUNCTION = "save" + OUTPUT_NODE = True + + CATEGORY = "advanced/model_merging" + + def save(self, model, clip, vae, filename_prefix, prompt=None, extra_pnginfo=None): + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) + prompt_info = "" + if prompt is not None: + prompt_info = json.dumps(prompt) + + metadata = {} + + enable_modelspec = True + if isinstance(model.model, comfy.model_base.SDXL): + metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-base" + elif isinstance(model.model, comfy.model_base.SDXLRefiner): + metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-refiner" + else: + enable_modelspec = False + + if enable_modelspec: + metadata["modelspec.sai_model_spec"] = "1.0.0" + metadata["modelspec.implementation"] = "sgm" + metadata["modelspec.title"] = "{} {}".format(filename, counter) + + #TODO: + # "stable-diffusion-v1", "stable-diffusion-v1-inpainting", "stable-diffusion-v2-512", + # "stable-diffusion-v2-768-v", "stable-diffusion-v2-unclip-l", "stable-diffusion-v2-unclip-h", + # "v2-inpainting" + + if model.model.model_type == comfy.model_base.ModelType.EPS: + metadata["modelspec.predict_key"] = "epsilon" + elif model.model.model_type == comfy.model_base.ModelType.V_PREDICTION: + metadata["modelspec.predict_key"] = "v" + + if not args.disable_metadata: + metadata["prompt"] = prompt_info + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata[x] = json.dumps(extra_pnginfo[x]) + + output_checkpoint = f"{filename}_{counter:05}_.safetensors" + output_checkpoint = os.path.join(full_output_folder, output_checkpoint) + + comfy.sd.save_checkpoint(output_checkpoint, model, clip, vae, metadata=metadata) + return {} + + +NODE_CLASS_MAPPINGS = { + "ModelMergeSimple": ModelMergeSimple, + "ModelMergeBlocks": ModelMergeBlocks, + "CheckpointSave": CheckpointSave, + "CLIPMergeSimple": CLIPMergeSimple, +}