508 lines
17 KiB
Python
508 lines
17 KiB
Python
import json
|
|
import torch
|
|
from typing import Dict
|
|
|
|
import folder_paths
|
|
import comfy.sd
|
|
from nodes import LoraLoader, VAELoader
|
|
from comfy_extras.nodes_model_merging import CheckpointSave
|
|
|
|
from .logger import logger
|
|
|
|
from .utility_nodes import (
|
|
NODE_CLASS_MAPPINGS as UTIL_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as UTIL_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
from .sdxl_prompt_styler import (
|
|
NODE_CLASS_MAPPINGS as SDXL_STYLER_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as SDXL_STYLER_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
from .interrogate import (
|
|
NODE_CLASS_MAPPINGS as INTERROGATE_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as INTERROGATE_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
from .fooocus import (
|
|
NODE_CLASS_MAPPINGS as FOOOCUS_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as FOOOCUS_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
from .postprocessing import (
|
|
NODE_CLASS_MAPPINGS as PP_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as PP_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
from .controlnet import (
|
|
NODE_CLASS_MAPPINGS as CONTROLNET_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as CONTROLNET_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
from .animatediff import (
|
|
NODE_CLASS_MAPPINGS as ANIMATEDIFF_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as ANIMATEDIFF_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
from .ip_adapter_nodes import (
|
|
NODE_CLASS_MAPPINGS as IP_ADAPTER_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as IP_ADAPTER_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
from .isnet import (
|
|
NODE_CLASS_MAPPINGS as ISNET_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as ISNET_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
from .inpaint import (
|
|
NODE_CLASS_MAPPINGS as INPAINT_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as INPAINT_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
from .video import (
|
|
NODE_CLASS_MAPPINGS as VIDEO_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as VIDEO_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
from .impact import (
|
|
NODE_CLASS_MAPPINGS as IMPACT_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as IMPACT_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
from .llm import (
|
|
NODE_CLASS_MAPPINGS as LLM_NODE_CLASS_MAPPINGS,
|
|
NODE_DISPLAY_NAME_MAPPINGS as LLM_NODE_DISPLAY_NAME_MAPPINGS,
|
|
)
|
|
|
|
|
|
class AVVAELoader(VAELoader):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
inputs = VAELoader.INPUT_TYPES()
|
|
inputs["optional"] = {"vae_override": ("STRING", {"default": "None"})}
|
|
return inputs
|
|
|
|
CATEGORY = "Art Venture/Loaders"
|
|
|
|
def load_vae(self, vae_name, vae_override="None"):
|
|
if vae_override != "None":
|
|
if vae_override not in folder_paths.get_filename_list("vae"):
|
|
print(f"Warning: Not found VAE model {vae_override}. Use {vae_name} instead.")
|
|
else:
|
|
vae_name = vae_override
|
|
|
|
return super().load_vae(vae_name)
|
|
|
|
|
|
class AVLoraLoader(LoraLoader):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
inputs = LoraLoader.INPUT_TYPES()
|
|
inputs["optional"] = {
|
|
"lora_override": ("STRING", {"default": "None"}),
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
}
|
|
return inputs
|
|
|
|
CATEGORY = "Art Venture/Loaders"
|
|
|
|
def load_lora(self, model, clip, lora_name, *args, lora_override="None", enabled=True, **kwargs):
|
|
if not enabled:
|
|
return (model, clip)
|
|
|
|
if lora_override != "None":
|
|
if lora_override not in folder_paths.get_filename_list("loras"):
|
|
print(f"Warning: Not found Lora model {lora_override}. Use {lora_name} instead.")
|
|
else:
|
|
lora_name = lora_override
|
|
|
|
return super().load_lora(model, clip, lora_name, *args, **kwargs)
|
|
|
|
|
|
class AVLoraListStacker:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"data": ("STRING", {"default": "", "multiline": True, "dynamicPrompts": False}),
|
|
},
|
|
"optional": {"lora_stack": ("LORA_STACK",)},
|
|
}
|
|
|
|
RETURN_TYPES = ("LORA_STACK",)
|
|
FUNCTION = "load_list_lora"
|
|
CATEGORY = "Art Venture/Loaders"
|
|
|
|
def parse_lora_list(self, data: str):
|
|
# data is a list of lora model (lora_name, strength_model, strength_clip, url) in json format
|
|
# trim data
|
|
data = data.strip()
|
|
if data == "" or data == "[]" or data is None:
|
|
return []
|
|
|
|
print(f"Loading lora list: {data}")
|
|
|
|
lora_list = json.loads(data)
|
|
if len(lora_list) == 0:
|
|
return []
|
|
|
|
available_loras = folder_paths.get_filename_list("loras")
|
|
|
|
lora_params = []
|
|
for lora in lora_list:
|
|
lora_name = lora["name"]
|
|
strength_model = lora["strength"]
|
|
strength_clip = lora["strength"]
|
|
|
|
if strength_model == 0 and strength_clip == 0:
|
|
continue
|
|
|
|
if lora_name not in available_loras:
|
|
print(f"Not found lora {lora_name}, skipping")
|
|
continue
|
|
|
|
lora_params.append((lora_name, strength_model, strength_clip))
|
|
|
|
return lora_params
|
|
|
|
def load_list_lora(self, data, lora_stack=None):
|
|
loras = self.parse_lora_list(data)
|
|
|
|
if lora_stack is not None:
|
|
loras.extend([l for l in lora_stack if l[0] != "None"])
|
|
|
|
return (loras,)
|
|
|
|
|
|
class AVLoraListLoader(AVLoraListStacker):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"clip": ("CLIP",),
|
|
"data": ("STRING", {"default": "", "multiline": True, "dynamicPrompts": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP")
|
|
|
|
def load_list_lora(self, model, clip, data):
|
|
lora_params = self.parse_lora_list(data)
|
|
|
|
if len(lora_params) == 0:
|
|
return (model, clip)
|
|
|
|
def load_loras(lora_params, model, clip):
|
|
for lora_name, strength_model, strength_clip in lora_params:
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
lora_file = comfy.utils.load_torch_file(lora_path)
|
|
model, clip = comfy.sd.load_lora_for_models(model, clip, lora_file, strength_model, strength_clip)
|
|
return model, clip
|
|
|
|
lora_model, lora_clip = load_loras(lora_params, model, clip)
|
|
|
|
return (lora_model, lora_clip)
|
|
|
|
|
|
class AVCheckpointModelsToParametersPipe:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
|
|
},
|
|
"optional": {
|
|
"pipe": ("PIPE",),
|
|
"secondary_ckpt_name": (["None"] + folder_paths.get_filename_list("checkpoints"),),
|
|
"vae_name": (["None"] + folder_paths.get_filename_list("vae"),),
|
|
"upscaler_name": (["None"] + folder_paths.get_filename_list("upscale_models"),),
|
|
"secondary_upscaler_name": (["None"] + folder_paths.get_filename_list("upscale_models"),),
|
|
"lora_1_name": (["None"] + folder_paths.get_filename_list("loras"),),
|
|
"lora_2_name": (["None"] + folder_paths.get_filename_list("loras"),),
|
|
"lora_3_name": (["None"] + folder_paths.get_filename_list("loras"),),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("PIPE",)
|
|
CATEGORY = "Art Venture/Parameters"
|
|
FUNCTION = "checkpoint_models_to_parameter_pipe"
|
|
|
|
def checkpoint_models_to_parameter_pipe(
|
|
self,
|
|
ckpt_name,
|
|
pipe: Dict = {},
|
|
secondary_ckpt_name="None",
|
|
vae_name="None",
|
|
upscaler_name="None",
|
|
secondary_upscaler_name="None",
|
|
lora_1_name="None",
|
|
lora_2_name="None",
|
|
lora_3_name="None",
|
|
):
|
|
pipe["ckpt_name"] = ckpt_name if ckpt_name != "None" else None
|
|
pipe["secondary_ckpt_name"] = secondary_ckpt_name if secondary_ckpt_name != "None" else None
|
|
pipe["vae_name"] = vae_name if vae_name != "None" else None
|
|
pipe["upscaler_name"] = upscaler_name if upscaler_name != "None" else None
|
|
pipe["secondary_upscaler_name"] = secondary_upscaler_name if secondary_upscaler_name != "None" else None
|
|
pipe["lora_1_name"] = lora_1_name if lora_1_name != "None" else None
|
|
pipe["lora_2_name"] = lora_2_name if lora_2_name != "None" else None
|
|
pipe["lora_3_name"] = lora_3_name if lora_3_name != "None" else None
|
|
return (pipe,)
|
|
|
|
|
|
class AVPromptsToParametersPipe:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"positive": ("STRING", {"multiline": True, "default": "Positive"}),
|
|
"negative": ("STRING", {"multiline": True, "default": "Negative"}),
|
|
},
|
|
"optional": {
|
|
"pipe": ("PIPE",),
|
|
"image": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("PIPE",)
|
|
CATEGORY = "Art Venture/Parameters"
|
|
FUNCTION = "prompt_to_parameter_pipe"
|
|
|
|
def prompt_to_parameter_pipe(self, positive, negative, pipe: Dict = {}, image=None, mask=None):
|
|
pipe["positive"] = positive
|
|
pipe["negative"] = negative
|
|
pipe["image"] = image
|
|
pipe["mask"] = mask
|
|
return (pipe,)
|
|
|
|
|
|
class AVParametersPipeToCheckpointModels:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipe": ("PIPE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"PIPE",
|
|
"CHECKPOINT_NAME",
|
|
"CHECKPOINT_NAME",
|
|
"VAE_NAME",
|
|
"UPSCALER_NAME",
|
|
"UPSCALER_NAME",
|
|
"LORA_NAME",
|
|
"LORA_NAME",
|
|
"LORA_NAME",
|
|
)
|
|
RETURN_NAMES = (
|
|
"pipe",
|
|
"ckpt_name",
|
|
"secondary_ckpt_name",
|
|
"vae_name",
|
|
"upscaler_name",
|
|
"secondary_upscaler_name",
|
|
"lora_1_name",
|
|
"lora_2_name",
|
|
"lora_3_name",
|
|
)
|
|
CATEGORY = "Art Venture/Parameters"
|
|
FUNCTION = "parameter_pipe_to_checkpoint_models"
|
|
|
|
def parameter_pipe_to_checkpoint_models(self, pipe: Dict = {}):
|
|
ckpt_name = pipe.get("ckpt_name", None)
|
|
secondary_ckpt_name = pipe.get("secondary_ckpt_name", None)
|
|
vae_name = pipe.get("vae_name", None)
|
|
upscaler_name = pipe.get("upscaler_name", None)
|
|
secondary_upscaler_name = pipe.get("secondary_upscaler_name", None)
|
|
lora_1_name = pipe.get("lora_1_name", None)
|
|
lora_2_name = pipe.get("lora_2_name", None)
|
|
lora_3_name = pipe.get("lora_3_name", None)
|
|
|
|
return (
|
|
pipe,
|
|
ckpt_name,
|
|
secondary_ckpt_name,
|
|
vae_name,
|
|
upscaler_name,
|
|
secondary_upscaler_name,
|
|
lora_1_name,
|
|
lora_2_name,
|
|
lora_3_name,
|
|
)
|
|
|
|
|
|
class AVParametersPipeToPrompts:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipe": ("PIPE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"PIPE",
|
|
"STRING",
|
|
"STRING",
|
|
"IMAGE",
|
|
"MASK",
|
|
)
|
|
RETURN_NAMES = (
|
|
"pipe",
|
|
"positive",
|
|
"negative",
|
|
"image",
|
|
"mask",
|
|
)
|
|
CATEGORY = "Art Venture/Parameters"
|
|
FUNCTION = "parameter_pipe_to_prompt"
|
|
|
|
def parameter_pipe_to_prompt(self, pipe: Dict = {}):
|
|
positive = pipe.get("positive", None)
|
|
negative = pipe.get("negative", None)
|
|
image = pipe.get("image", None)
|
|
mask = pipe.get("mask", None)
|
|
|
|
return (
|
|
pipe,
|
|
positive,
|
|
negative,
|
|
image,
|
|
mask,
|
|
)
|
|
|
|
|
|
class AVCheckpointMerge:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model1": ("MODEL",),
|
|
"model2": ("MODEL",),
|
|
"model1_weight": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 1.0, "step": 0.01}),
|
|
"model2_weight": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "merge"
|
|
|
|
CATEGORY = "Art Venture/Model Merging"
|
|
|
|
def merge(self, model1, model2, model1_weight, model2_weight):
|
|
m = model1.clone()
|
|
k1 = model1.get_key_patches("diffusion_model.")
|
|
k2 = model2.get_key_patches("diffusion_model.")
|
|
for k in k1:
|
|
if k in k2:
|
|
a = k1[k][0]
|
|
b = k2[k][0]
|
|
|
|
if a.shape != b.shape and a.shape[0:1] + a.shape[2:] == b.shape[0:1] + b.shape[2:]:
|
|
if a.shape[1] == 4 and b.shape[1] == 9:
|
|
raise RuntimeError(
|
|
"When merging inpainting model with a normal one, model1 must be the inpainting model."
|
|
)
|
|
if a.shape[1] == 4 and b.shape[1] == 8:
|
|
raise RuntimeError(
|
|
"When merging instruct-pix2pix model with a normal one, model1 must be the instruct-pix2pix model."
|
|
)
|
|
|
|
c = torch.zeros_like(a)
|
|
c[:, 0:4, :, :] = b
|
|
b = c
|
|
|
|
m.add_patches({k: (b,)}, model2_weight, model1_weight)
|
|
else:
|
|
logger.warn(f"Key {k} not found in model2")
|
|
m.add_patches({k: k1[k]}, -1.0, 1.0) # zero out
|
|
|
|
return (m,)
|
|
|
|
|
|
class AVCheckpointSave(CheckpointSave):
|
|
CATEGORY = "Art Venture/Model Merging"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
inputs = CheckpointSave.INPUT_TYPES()
|
|
inputs["optional"] = {
|
|
"dtype": (["float16", "float32"], {"default": "float16"}),
|
|
}
|
|
|
|
return inputs
|
|
|
|
def save(self, *args, dtype="float16", **kwargs):
|
|
comfy_save_checkpoint = comfy.sd.save_checkpoint
|
|
|
|
if dtype == "float16":
|
|
|
|
def save_checkpoint(output_path, model, clip, vae, metadata=None):
|
|
model.model.half()
|
|
return comfy_save_checkpoint(output_path, model, clip, vae, metadata)
|
|
|
|
comfy.sd.save_checkpoint = save_checkpoint
|
|
|
|
try:
|
|
return super().save(*args, **kwargs)
|
|
finally:
|
|
comfy.sd.save_checkpoint = comfy_save_checkpoint
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AV_CheckpointModelsToParametersPipe": AVCheckpointModelsToParametersPipe,
|
|
"AV_PromptsToParametersPipe": AVPromptsToParametersPipe,
|
|
"AV_ParametersPipeToCheckpointModels": AVParametersPipeToCheckpointModels,
|
|
"AV_ParametersPipeToPrompts": AVParametersPipeToPrompts,
|
|
"AV_VAELoader": AVVAELoader,
|
|
"AV_LoraLoader": AVLoraLoader,
|
|
"AV_LoraListLoader": AVLoraListLoader,
|
|
"AV_LoraListStacker": AVLoraListStacker,
|
|
"AV_CheckpointMerge": AVCheckpointMerge,
|
|
"AV_CheckpointSave": AVCheckpointSave,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AV_CheckpointModelsToParametersPipe": "Checkpoint Models to Pipe",
|
|
"AV_PromptsToParametersPipe": "Prompts to Pipe",
|
|
"AV_ParametersPipeToCheckpointModels": "Pipe to Checkpoint Models",
|
|
"AV_ParametersPipeToPrompts": "Pipe to Prompts",
|
|
"AV_VAELoader": "VAE Loader",
|
|
"AV_LoraLoader": "Lora Loader",
|
|
"AV_LoraListLoader": "Lora List Loader",
|
|
"AV_LoraListStacker": "Lora List Stacker",
|
|
"AV_CheckpointMerge": "Checkpoint Merge",
|
|
"AV_CheckpointSave": "Checkpoint Save",
|
|
}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS.update(UTIL_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(UTIL_NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
NODE_CLASS_MAPPINGS.update(SDXL_STYLER_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(SDXL_STYLER_NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
NODE_CLASS_MAPPINGS.update(INTERROGATE_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(INTERROGATE_NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
NODE_CLASS_MAPPINGS.update(FOOOCUS_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(FOOOCUS_NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
NODE_CLASS_MAPPINGS.update(PP_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(PP_NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
NODE_CLASS_MAPPINGS.update(CONTROLNET_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(CONTROLNET_NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
NODE_CLASS_MAPPINGS.update(ANIMATEDIFF_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(ANIMATEDIFF_NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
NODE_CLASS_MAPPINGS.update(IP_ADAPTER_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(IP_ADAPTER_NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
NODE_CLASS_MAPPINGS.update(ISNET_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(ISNET_NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
NODE_CLASS_MAPPINGS.update(INPAINT_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(INPAINT_NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
NODE_CLASS_MAPPINGS.update(VIDEO_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(VIDEO_NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
NODE_CLASS_MAPPINGS.update(IMPACT_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(IMPACT_NODE_DISPLAY_NAME_MAPPINGS)
|
|
|
|
NODE_CLASS_MAPPINGS.update(LLM_NODE_CLASS_MAPPINGS)
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(LLM_NODE_DISPLAY_NAME_MAPPINGS)
|