diff --git a/README.ZH_CN.md b/README.ZH_CN.md index 82fe689..575e50c 100644 --- a/README.ZH_CN.md +++ b/README.ZH_CN.md @@ -39,6 +39,7 @@ - 支持 kolors 模型 - 支持 flux 模型 - 支持 惰性条件判断（ifElse）和 for循环 +- 支持 Anima 与 Krea2 diffusion 模型，可通过 `easy diffusionModelLoader` 加载（需显式选择文本编码器与 VAE），并使用 `easy XYInputs: DiffusionModel` 进行 XY 对比 ## 👨🏻‍🔧 安装 diff --git a/README.md b/README.md index 6b4ff50..aa89cb4 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,7 @@ - Support Kolors‘s model. - Support Flux's model. - Support lazy if else and for loops. +- Support Anima and Krea2 diffusion models with `easy diffusionModelLoader` and `easy XYInputs: DiffusionModel`. The loader requires an explicit text encoder and VAE selection. ## 👨🏻‍🔧 Installation Clone the repo into the **custom_nodes** directory and install the requirements: diff --git a/__init__.py b/__init__.py index e2419f1..7459342 100644 --- a/__init__.py +++ b/__init__.py @@ -12,6 +12,18 @@ comfy_path = folder_paths.base_path NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} +try: + import comfy.supported_models as _supported_models + _HAS_DIFFUSION_XY_SUPPORT = ( + hasattr(_supported_models, "Anima") + and hasattr(_supported_models, "Krea2") + ) +except Exception: + _HAS_DIFFUSION_XY_SUPPORT = False + +if not _HAS_DIFFUSION_XY_SUPPORT: + print("[ComfyUI-Easy-Use] Anima/Krea2 XY nodes need comfy.supported_models.Anima and Krea2") + importlib.import_module('.py.routes', __name__) importlib.import_module('.py.server', __name__) nodes_list = ["util", "seed", "prompt", "loaders", "adapter", "inpaint", "preSampling", "samplers", "fix", "pipe", "xyplot", "image", "logic", "api", "deprecated"] diff --git a/locales/en/nodeDefs.json b/locales/en/nodeDefs.json index 46b6fec..c4ca793 100644 --- a/locales/en/nodeDefs.json +++ b/locales/en/nodeDefs.json @@ -1202,6 +1202,173 @@ } } }, + "easy diffusionModelLoader": { + "display_name": "EasyDiffusionModelLoader", + "inputs": { + "model_name": { + "name": "model_name" + }, + "vae_name": { + "name": "vae_name" + }, + "clip_name": { + "name": "clip_name" + }, + "resolution": { + "name": "resolution" + }, + "empty_latent_width": { + "name": "empty_latent_width" + }, + "empty_latent_height": { + "name": "empty_latent_height" + }, + "positive": { + "name": "positive" + }, + "negative": { + "name": "negative" + }, + "batch_size": { + "name": "batch_size" + }, + "model_override": { + "name": "model_override" + }, + "clip_override": { + "name": "clip_override" + }, + "vae_override": { + "name": "vae_override" + } + }, + "outputs": { + "0": { + "name": "pipe" + }, + "1": { + "name": "model" + }, + "2": { + "name": "vae" + }, + "3": { + "name": "clip" + }, + "4": { + "name": "positive" + }, + "5": { + "name": "negative" + }, + "6": { + "name": "latent" + } + } + }, + "easy XYInputs: DiffusionModel": { + "display_name": "XY Inputs: Diffusion Model //EasyUse", + "inputs": { + "model_count": { + "name": "model_count" + }, + "model_name_1": { + "name": "model_name_1" + }, + "clip_name_1": { + "name": "clip_name_1" + }, + "vae_name_1": { + "name": "vae_name_1" + }, + "model_name_2": { + "name": "model_name_2" + }, + "clip_name_2": { + "name": "clip_name_2" + }, + "vae_name_2": { + "name": "vae_name_2" + }, + "model_name_3": { + "name": "model_name_3" + }, + "clip_name_3": { + "name": "clip_name_3" + }, + "vae_name_3": { + "name": "vae_name_3" + }, + "model_name_4": { + "name": "model_name_4" + }, + "clip_name_4": { + "name": "clip_name_4" + }, + "vae_name_4": { + "name": "vae_name_4" + }, + "model_name_5": { + "name": "model_name_5" + }, + "clip_name_5": { + "name": "clip_name_5" + }, + "vae_name_5": { + "name": "vae_name_5" + }, + "model_name_6": { + "name": "model_name_6" + }, + "clip_name_6": { + "name": "clip_name_6" + }, + "vae_name_6": { + "name": "vae_name_6" + }, + "model_name_7": { + "name": "model_name_7" + }, + "clip_name_7": { + "name": "clip_name_7" + }, + "vae_name_7": { + "name": "vae_name_7" + }, + "model_name_8": { + "name": "model_name_8" + }, + "clip_name_8": { + "name": "clip_name_8" + }, + "vae_name_8": { + "name": "vae_name_8" + }, + "model_name_9": { + "name": "model_name_9" + }, + "clip_name_9": { + "name": "clip_name_9" + }, + "vae_name_9": { + "name": "vae_name_9" + }, + "model_name_10": { + "name": "model_name_10" + }, + "clip_name_10": { + "name": "clip_name_10" + }, + "vae_name_10": { + "name": "vae_name_10" + } + }, + "outputs": { + "0": { + "name": "X or Y" + } + } + }, "easy loraStack": { "display_name": "EasyLoraStack", "inputs": { diff --git a/locales/zh/nodeDefs.json b/locales/zh/nodeDefs.json index a65f448..8cbe7fe 100644 --- a/locales/zh/nodeDefs.json +++ b/locales/zh/nodeDefs.json @@ -904,6 +904,173 @@ } } }, + "easy diffusionModelLoader": { + "display_name": "简易加载器（扩散模型）", + "inputs": { + "model_name": { + "name": "扩散模型" + }, + "vae_name": { + "name": "VAE" + }, + "clip_name": { + "name": "文本编码器" + }, + "resolution": { + "name": "分辨率" + }, + "empty_latent_width": { + "name": "宽度" + }, + "empty_latent_height": { + "name": "高度" + }, + "positive": { + "name": "正面提示词" + }, + "negative": { + "name": "负面提示词" + }, + "batch_size": { + "name": "批次大小" + }, + "model_override": { + "name": "模型(可选)" + }, + "clip_override": { + "name": "CLIP(可选)" + }, + "vae_override": { + "name": "VAE(可选)" + } + }, + "outputs": { + "0": { + "name": "节点束" + }, + "1": { + "name": "模型" + }, + "2": { + "name": "VAE" + }, + "3": { + "name": "CLIP" + }, + "4": { + "name": "正面提示词" + }, + "5": { + "name": "负面提示词" + }, + "6": { + "name": "潜空间" + } + } + }, + "easy XYInputs: DiffusionModel": { + "display_name": "XY输入: Diffusion Model", + "inputs": { + "model_count": { + "name": "模型数量" + }, + "model_name_1": { + "name": "扩散模型1" + }, + "clip_name_1": { + "name": "文本编码器1" + }, + "vae_name_1": { + "name": "VAE1" + }, + "model_name_2": { + "name": "扩散模型2" + }, + "clip_name_2": { + "name": "文本编码器2" + }, + "vae_name_2": { + "name": "VAE2" + }, + "model_name_3": { + "name": "扩散模型3" + }, + "clip_name_3": { + "name": "文本编码器3" + }, + "vae_name_3": { + "name": "VAE3" + }, + "model_name_4": { + "name": "扩散模型4" + }, + "clip_name_4": { + "name": "文本编码器4" + }, + "vae_name_4": { + "name": "VAE4" + }, + "model_name_5": { + "name": "扩散模型5" + }, + "clip_name_5": { + "name": "文本编码器5" + }, + "vae_name_5": { + "name": "VAE5" + }, + "model_name_6": { + "name": "扩散模型6" + }, + "clip_name_6": { + "name": "文本编码器6" + }, + "vae_name_6": { + "name": "VAE6" + }, + "model_name_7": { + "name": "扩散模型7" + }, + "clip_name_7": { + "name": "文本编码器7" + }, + "vae_name_7": { + "name": "VAE7" + }, + "model_name_8": { + "name": "扩散模型8" + }, + "clip_name_8": { + "name": "文本编码器8" + }, + "vae_name_8": { + "name": "VAE8" + }, + "model_name_9": { + "name": "扩散模型9" + }, + "clip_name_9": { + "name": "文本编码器9" + }, + "vae_name_9": { + "name": "VAE9" + }, + "model_name_10": { + "name": "扩散模型10" + }, + "clip_name_10": { + "name": "文本编码器10" + }, + "vae_name_10": { + "name": "VAE10" + } + }, + "outputs": { + "0": { + "name": "X或Y" + } + } + }, "easy zero123Loader": { "display_name": "简易加载器（Zero123）", "inputs": { diff --git a/py/config.py b/py/config.py index dee0a4b..dedc64c 100644 --- a/py/config.py +++ b/py/config.py @@ -405,3 +405,21 @@ PROMPT_TEMPLATE = { } NEW_SCHEDULERS = ['align_your_steps', 'gits'] + +DIFFUSION_MODEL_XY_DEFAULTS = { + "anima": { + "clip_name": "qwen_3_06b_base.safetensors", + "clip_type": "anima", + "vae_name": "qwen_image_vae.safetensors", + }, + "krea2": { + "clip_name": "Huihui-Qwen3-VL-4B-Instruct-abliterated-fp8_scaled.safetensors", + "clip_type": "krea2", + "vae_name": "qwen_image_vae.safetensors", + }, +} + +DIFFUSION_MODEL_CLIP_TYPES = { + "anima": "anima", + "krea2": "krea2", +} diff --git a/py/libs/conditioning.py b/py/libs/conditioning.py index b9d9733..862dae2 100644 --- a/py/libs/conditioning.py +++ b/py/libs/conditioning.py @@ -14,7 +14,7 @@ def prompt_to_cond(type, model, clip, clip_skip, lora_stack, text, prompt_token_ if model_type not in ['hydit'] and text is not None and has_chinese(text): text = zh_to_en([text])[0] - if model_type in ['hydit', 'flux', 'mochi']: + if model_type in ['hydit', 'flux', 'mochi', 'anima', 'krea2']: log_node_warn(title + "...") embeddings_final, = CLIPTextEncode().encode(clip, text) if text is not None else (None,) diff --git a/py/libs/loader.py b/py/libs/loader.py index 6702ed5..023d6b0 100644 --- a/py/libs/loader.py +++ b/py/libs/loader.py @@ -8,6 +8,8 @@ from comfy.model_patcher import ModelPatcher from nodes import NODE_CLASS_MAPPINGS from collections import defaultdict from .log import log_node_info, log_node_error +from .utils import get_sd_version +from ..config import DIFFUSION_MODEL_XY_DEFAULTS, DIFFUSION_MODEL_CLIP_TYPES from ..modules.dit.pixArt.loader import load_pixart diffusion_loaders = ["easy fullLoader", "easy a1111Loader", "easy fluxLoader", "easy comfyLoader", "easy hunyuanDiTLoader", "easy zero123Loader", "easy svdLoader"] @@ -145,6 +147,28 @@ class easyLoader: scale_soft_weights = self.get_input_value(entry, "cn_soft_weights") desired_controlnet_names.add(f'{control_net_name};{scale_soft_weights}') + elif class_type == "easy diffusionModelLoader": + desired_unet_names.add(self.get_input_value(entry, "model_name", prompt)) + clip_name = self.get_input_value(entry, "clip_name", prompt) + vae_name = self.get_input_value(entry, "vae_name", prompt) + if clip_name not in ("None", "Auto"): + desired_clip_names.add(clip_name) + if vae_name not in ("None", "Auto"): + desired_vae_names.add(vae_name) + + elif class_type == "easy XYInputs: DiffusionModel": + model_count = int(self.get_input_value(entry, "model_count", prompt) or 0) + for i in range(1, model_count + 1): + model_name = self.get_input_value(entry, f"model_name_{i}", prompt) + if model_name and model_name != "None": + desired_unet_names.add(model_name) + clip_name = self.get_input_value(entry, f"clip_name_{i}", prompt) + if clip_name not in ("None", "Auto"): + desired_clip_names.add(clip_name) + vae_name = self.get_input_value(entry, f"vae_name_{i}", prompt) + if vae_name not in ("None", "Auto"): + desired_vae_names.add(vae_name) + elif class_type in model_merge_node: desired_ckpt_names.add(self.get_input_value(entry, "ckpt_name_1")) desired_ckpt_names.add(self.get_input_value(entry, "ckpt_name_2")) @@ -282,6 +306,57 @@ class easyLoader: return model + def load_diffusion_model(self, model_name): + if model_name in self.loaded_objects["unet"]: + log_node_info("Load Diffusion Model", f"{model_name} cached") + return self.loaded_objects["unet"][model_name][0] + + model_path = folder_paths.get_full_path("diffusion_models", model_name) + if not model_path: + raise FileNotFoundError(f"[EasyUse] diffusion model not found: {model_name}") + + model = comfy.sd.load_diffusion_model(model_path) + self.add_to_cache("unet", model_name, model) + self.eviction_based_on_memory() + + return model + + def load_diffusion_xy_model(self, model_name, clip_name, vae_name): + model = self.load_diffusion_model(model_name) + family = get_sd_version(model) + + defaults = DIFFUSION_MODEL_XY_DEFAULTS.get(family) + if defaults is None: + raise RuntimeError(f"[EasyUse] unsupported diffusion model family: {family}") + + if clip_name in ("Auto", None): + clip_name = defaults["clip_name"] + if vae_name in ("Auto", None): + vae_name = defaults["vae_name"] + + clip = self.load_clip(clip_name, type=defaults["clip_type"]) + vae = self.load_vae(vae_name) + + return model, clip, vae, family + + def load_diffusion_model_required(self, model_name, clip_name, vae_name): + if clip_name in ("None", None): + raise RuntimeError("[EasyUse] clip_name is required: please select a text encoder") + if vae_name in ("None", None): + raise RuntimeError("[EasyUse] vae_name is required: please select a VAE") + + model = self.load_diffusion_model(model_name) + family = get_sd_version(model) + + clip_type = DIFFUSION_MODEL_CLIP_TYPES.get(family) + if clip_type is None: + raise RuntimeError(f"[EasyUse] unsupported diffusion model family: {family}") + + clip = self.load_clip(clip_name, type=clip_type) + vae = self.load_vae(vae_name) + + return model, clip, vae, family + def load_controlnet(self, control_net_name, scale_soft_weights=1, use_cache=True): unique_id = f'{control_net_name};{str(scale_soft_weights)}' if use_cache and unique_id in self.loaded_objects["controlnet"]: @@ -303,8 +378,9 @@ class easyLoader: return control_net def load_clip(self, clip_name, type='stable_diffusion', load_clip=None): - if clip_name in self.loaded_objects["clip"]: - return self.loaded_objects["clip"][clip_name][0] + cache_key = f"{clip_name}::{type}" + if cache_key in self.loaded_objects["clip"]: + return self.loaded_objects["clip"][cache_key][0] if type == 'stable_diffusion': clip_type = comfy.sd.CLIPType.STABLE_DIFFUSION @@ -316,9 +392,13 @@ class easyLoader: clip_type = comfy.sd.CLIPType.FLUX elif type == 'stable_audio': clip_type = comfy.sd.CLIPType.STABLE_AUDIO + elif type == 'krea2': + clip_type = comfy.sd.CLIPType.KREA2 + elif type == 'anima': + clip_type = comfy.sd.CLIPType.STABLE_DIFFUSION clip_path = folder_paths.get_full_path("clip", clip_name) load_clip = comfy.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type) - self.add_to_cache("clip", clip_name, load_clip) + self.add_to_cache("clip", cache_key, load_clip) self.eviction_based_on_memory() return load_clip diff --git a/py/libs/sampler.py b/py/libs/sampler.py index 70716ef..65c6128 100644 --- a/py/libs/sampler.py +++ b/py/libs/sampler.py @@ -65,6 +65,9 @@ class easySampler: elif model_type == 'mochi': latent = torch.zeros([batch_size, 12, ((video_length - 1) // 6) + 1, empty_latent_height // 8, empty_latent_width // 8], device=self.device) samples = {"samples": latent} + elif model_type in ("anima", "krea2"): + latent = torch.zeros([batch_size, 16, 1, empty_latent_height // 8, empty_latent_width // 8], device=self.device) + samples = {"samples": latent} elif compression == 0: latent = torch.zeros([batch_size, 4, empty_latent_height // 8, empty_latent_width // 8], device=self.device) samples = {"samples": latent} @@ -84,7 +87,7 @@ class easySampler: """ latent_size = latent_image.size() - latent_size_1batch = [1, latent_size[1], latent_size[2], latent_size[3]] + latent_size_1batch = [1] + list(latent_size[1:]) if variation_strength is not None and variation_strength > 0 or incremental_seed_mode.startswith( "variation str inc"): @@ -108,7 +111,7 @@ class easySampler: if strength_up is not None: strength += strength_up - variation_noise = variation_latent.expand(input_latent.size()[0], -1, -1, -1) + variation_noise = variation_latent.expand(input_latent.size()[0], *([-1] * (variation_latent.dim() - 1))) result = (1 - strength) * input_latent + strength * variation_noise return result diff --git a/py/libs/utils.py b/py/libs/utils.py index 0fc716a..12edd1e 100644 --- a/py/libs/utils.py +++ b/py/libs/utils.py @@ -128,6 +128,10 @@ def get_sd_version(model): return 'flux' elif isinstance(model_config, comfy.supported_models.GenmoMochi): return 'mochi' + elif isinstance(model_config, comfy.supported_models.Anima): + return 'anima' + elif isinstance(model_config, comfy.supported_models.Krea2): + return 'krea2' else: return 'unknown' diff --git a/py/libs/xyplot.py b/py/libs/xyplot.py index 37c8a6f..7f05f47 100644 --- a/py/libs/xyplot.py +++ b/py/libs/xyplot.py @@ -68,6 +68,10 @@ class easyXYPlot(): lora_weight_desc = f" w:{lora_weight:.2f}" if value_type == 'Lora' and lora_weight != 1.0 else '' value_label = f"{model_name[:25]}{lora_weight_desc}{trigger_words}" + if value_type == "DiffusionModel": + model_name = os.path.basename(os.path.splitext(value.split(",")[0])[0]) + value_label = model_name[:25] + if value_type in ["ModelMergeBlocks"]: if ":" in value: line = value.split(':') @@ -98,6 +102,32 @@ class easyXYPlot(): return plot_image_vars, value_label + @staticmethod + def _ensure_latent_for_model(model, vae, samples, plot_image_vars): + fmt = model.model.latent_format + x = samples["samples"] + expected_ndim = 2 + fmt.latent_dimensions + + if x.ndim == expected_ndim and x.shape[1] == fmt.latent_channels: + return samples + + if fmt.latent_dimensions == 3 and x.ndim == 4: + if x.count_nonzero() == 0: + x = torch.zeros( + [x.shape[0], fmt.latent_channels, 1, x.shape[2], x.shape[3]], + dtype=x.dtype, device=x.device) + elif plot_image_vars.get("images") is not None: + x = vae.encode(plot_image_vars["images"][..., :3]) + else: + raise RuntimeError( + "Switching to a 3D-latent model requires an input image " + "or an empty latent" + ) + + return {**samples, "samples": x} + + return samples + @staticmethod def get_font(font_size, font_path=None): if font_path is None: @@ -362,6 +392,28 @@ class easyXYPlot(): if "negative_cond" in plot_image_vars: negative = negative + plot_image_vars["negative_cond"] + # DiffusionModel + if self.x_type == "DiffusionModel" or self.y_type == "DiffusionModel": + xy_values = x_value if self.x_type == "DiffusionModel" else y_value + model_name, clip_name, vae_name = xy_values.split(",") + model, clip, vae, family = self.easyCache.load_diffusion_xy_model( + model_name.replace("*", ","), + clip_name.replace("*", ","), + vae_name.replace("*", ","), + ) + sd_version = family + + positive = plot_image_vars["positive"] + negative = plot_image_vars["negative"] + if positive is not None: + positive, = CLIPTextEncode().encode(clip, positive) + if negative is not None: + negative, = CLIPTextEncode().encode(clip, negative) + + samples = self._ensure_latent_for_model( + model, vae, samples, plot_image_vars + ) + # Lora if self.x_type == "Lora" or self.y_type == "Lora": # print(f"Lora: {x_value} {y_value}") @@ -399,7 +451,7 @@ class easyXYPlot(): if self.x_type == 'Positive Prompt S/R' or self.y_type == 'Positive Prompt S/R': positive = x_value if self.x_type == "Positive Prompt S/R" else y_value - if sd_version == 'flux': + if sd_version in ("flux", "anima", "krea2"): positive, = CLIPTextEncode().encode(clip, positive) else: positive = advanced_encode(clip, positive, @@ -415,7 +467,7 @@ class easyXYPlot(): if self.x_type == 'Negative Prompt S/R' or self.y_type == 'Negative Prompt S/R': negative = x_value if self.x_type == "Negative Prompt S/R" else y_value - if sd_version == 'flux': + if sd_version in ("flux", "anima", "krea2"): negative, = CLIPTextEncode().encode(clip, negative) else: negative = advanced_encode(clip, negative, @@ -483,7 +535,7 @@ class easyXYPlot(): clip = clip.clone() clip.clip_layer(plot_image_vars['clip_skip']) - if sd_version == 'flux': + if sd_version in ("flux", "anima", "krea2"): positive, = CLIPTextEncode().encode(clip, positive) else: positive = advanced_encode(clip, plot_image_vars['positive'], @@ -491,7 +543,7 @@ class easyXYPlot(): plot_image_vars['positive_weight_interpretation'], w_max=1.0, apply_to_pooled="enable",a1111_prompt_style=a1111_prompt_style, steps=steps) - if sd_version == 'flux': + if sd_version in ("flux", "anima", "krea2"): negative, = CLIPTextEncode().encode(clip, negative) else: negative = advanced_encode(clip, plot_image_vars['negative'], diff --git a/py/nodes/loaders.py b/py/nodes/loaders.py index 2ad24d0..6a254f8 100644 --- a/py/nodes/loaders.py +++ b/py/nodes/loaders.py @@ -1146,6 +1146,90 @@ class mochiLoader(fullLoader): batch_size, model_override, clip_override, vae_override, a1111_prompt_style=False, video_length=length, prompt=prompt, my_unique_id=my_unique_id ) +# Diffusion model loader +class diffusionModelLoader: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_name": (folder_paths.get_filename_list("diffusion_models"),), + "vae_name": (["None"] + folder_paths.get_filename_list("vae"), {"default": "None"}), + "clip_name": (["None"] + folder_paths.get_filename_list("text_encoders"), {"default": "None"}), + "resolution": (resolution_strings, {"default": "1024 x 1024"}), + "empty_latent_width": ("INT", {"default": 1024, "min": 64, "max": MAX_RESOLUTION, "step": 8}), + "empty_latent_height": ("INT", {"default": 1024, "min": 64, "max": MAX_RESOLUTION, "step": 8}), + "positive": ("STRING", {"default": "", "multiline": True}), + "negative": ("STRING", {"default": "", "multiline": True}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}), + }, + "optional": { + "model_override": ("MODEL",), + "clip_override": ("CLIP",), + "vae_override": ("VAE",), + }, + "hidden": {"prompt": "PROMPT", "my_unique_id": "UNIQUE_ID"}, + } + + RETURN_TYPES = ("PIPE_LINE", "MODEL", "VAE", "CLIP", "CONDITIONING", "CONDITIONING", "LATENT") + RETURN_NAMES = ("pipe", "model", "vae", "clip", "positive", "negative", "latent") + FUNCTION = "adv_pipeloader" + CATEGORY = "EasyUse/Loaders" + + def adv_pipeloader(self, model_name, vae_name, clip_name, resolution, + empty_latent_width, empty_latent_height, positive, negative, + batch_size, model_override=None, clip_override=None, + vae_override=None, prompt=None, my_unique_id=None): + easyCache.update_loaded_objects(prompt) + model, clip, vae, family = easyCache.load_diffusion_model_required( + model_name, clip_name, vae_name + ) + + if model_override is not None: + model = model_override + if clip_override is not None: + clip = clip_override + if vae_override is not None: + vae = vae_override + + samples = sampler.emptyLatent(resolution, empty_latent_width, + empty_latent_height, batch_size, + model_type=family) + + positive_cond, positive_wildcard, model, clip = prompt_to_cond( + "positive", model, clip, 0, [], positive, "none", "comfy", + False, my_unique_id, prompt, easyCache, model_type=family) + negative_cond, negative_wildcard, model, clip = prompt_to_cond( + "negative", model, clip, 0, [], negative, "none", "comfy", + False, my_unique_id, prompt, easyCache, model_type=family) + + if negative_cond is None: + negative_cond, = ConditioningZeroOut().zero_out(positive_cond) + + pipe = { + "model": model, + "positive": positive_cond, + "negative": negative_cond, + "vae": vae, + "clip": clip, + "samples": samples, + "images": None, + "loader_settings": { + "model_name": model_name, + "clip_name": clip_name, + "vae_name": vae_name, + "model_type": family, + "positive": positive, + "negative": negative, + "resolution": resolution, + "empty_latent_width": empty_latent_width, + "empty_latent_height": empty_latent_height, + "batch_size": batch_size, + }, + } + + return pipe, model, vae, clip, positive_cond, negative_cond, samples + # lora class loraSwitcher: @classmethod @@ -1542,6 +1626,7 @@ NODE_CLASS_MAPPINGS = { "easy hunyuanDiTLoader": hunyuanDiTLoader, "easy pixArtLoader": pixArtLoader, "easy mochiLoader": mochiLoader, + "easy diffusionModelLoader": diffusionModelLoader, "easy loraSwitcher": loraSwitcher, "easy loraStack": loraStack, "easy controlnetStack": controlnetStack, @@ -1564,6 +1649,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "easy hunyuanDiTLoader": "EasyLoader (HunyuanDiT)", "easy pixArtLoader": "EasyLoader (PixArt)", "easy mochiLoader": "EasyLoader (Mochi)", + "easy diffusionModelLoader": "EasyDiffusionModelLoader", "easy loraSwitcher": "EasyLoraSwitcher", "easy loraStack": "EasyLoraStack", "easy controlnetStack": "EasyControlnetStack", diff --git a/py/nodes/pipe.py b/py/nodes/pipe.py index 65225eb..5b38561 100644 --- a/py/nodes/pipe.py +++ b/py/nodes/pipe.py @@ -629,6 +629,12 @@ class pipeXYPlotAdvanced: "lora_stack": lora_stack, } + if x_axis == "advanced: DiffusionModel": + x_values = "; ".join(x_values) + + if y_axis == "advanced: DiffusionModel": + y_values = "; ".join(y_values) + if x_axis == 'advanced: Seeds++ Batch': seed = new_pipe.get('seed') or 0 value = x_values diff --git a/py/nodes/preSampling.py b/py/nodes/preSampling.py index c9069c3..03a764f 100644 --- a/py/nodes/preSampling.py +++ b/py/nodes/preSampling.py @@ -262,15 +262,16 @@ class samplerSettingsNoiseIn: model = pipe["model"] # generate base noise - batch_size, _, height, width = latent["samples"].shape + sample_shape = latent["samples"].shape + batch_size = sample_shape[0] generator = torch.manual_seed(seed) - base_noise = torch.randn((1, 4, height, width), dtype=torch.float32, device="cpu", generator=generator).repeat(batch_size, 1, 1, 1).cpu() + base_noise = torch.randn((1, *sample_shape[1:]), dtype=torch.float32, device="cpu", generator=generator).repeat(batch_size, *([1] * (len(sample_shape) - 1))).cpu() # generate variation noise if optional_noise_seed is None or optional_noise_seed == seed: optional_noise_seed = seed+1 generator = torch.manual_seed(optional_noise_seed) - variation_noise = torch.randn((batch_size, 4, height, width), dtype=torch.float32, device="cpu", + variation_noise = torch.randn(sample_shape, dtype=torch.float32, device="cpu", generator=generator).cpu() slerp_noise = self.slerp(factor, base_noise, variation_noise) diff --git a/py/nodes/samplers.py b/py/nodes/samplers.py index d00dd44..b666bdc 100644 --- a/py/nodes/samplers.py +++ b/py/nodes/samplers.py @@ -138,6 +138,14 @@ class samplerFull: to["model_patch"] = {} return to + def get_align_your_steps_sigmas(self, model, steps, denoise): + model_type = get_sd_version(model) + # Anima/Krea2 have no dedicated AYS table; keep the SDXL table they used before these families were recognized. + if model_type in ("anima", "krea2", "unknown"): + model_type = "sdxl" + sigmas, = alignYourStepsScheduler().get_sigmas(model_type.upper(), steps, denoise) + return sigmas + def get_sampler_custom(self, model, positive, negative, loader_settings): _guider = None middle = loader_settings['middle'] if "middle" in loader_settings else negative @@ -174,10 +182,7 @@ class samplerFull: elif scheduler == 'sdturbo': sigmas, = self.get_custom_cls('SDTurboScheduler').execute(model, steps, denoise) elif scheduler == 'alignYourSteps': - model_type = get_sd_version(model) - if model_type == 'unknown': - model_type = 'sdxl' - sigmas, = alignYourStepsScheduler().get_sigmas(model_type.upper(), steps, denoise) + sigmas = self.get_align_your_steps_sigmas(model, steps, denoise) elif scheduler == 'gits': sigmas, = gitsScheduler().get_sigmas(coeff, steps, denoise) else: @@ -340,10 +345,7 @@ class samplerFull: _guider, _sampler, sigmas = self.get_sampler_custom(samp_model, samp_positive, samp_negative, samp_custom) samp_samples, samp_blend_samples = sampler.custom_advanced_ksampler(_guider, _sampler, sigmas, samp_samples, add_noise, samp_seed, preview_latent=preview_latent) elif scheduler == 'align_your_steps': - model_type = get_sd_version(samp_model) - if model_type == 'unknown': - model_type = 'sdxl' - sigmas, = alignYourStepsScheduler().get_sigmas(model_type.upper(), steps, denoise) + sigmas = self.get_align_your_steps_sigmas(samp_model, steps, denoise) _sampler = comfy.samplers.sampler_object(sampler_name) samp_samples = sampler.custom_ksampler(samp_model, samp_seed, steps, cfg, _sampler, sigmas, samp_positive, samp_negative, samp_samples, disable_noise=disable_noise, preview_latent=preview_latent, noise_device=noise_device) elif scheduler == 'gits': diff --git a/py/nodes/xyplot.py b/py/nodes/xyplot.py index 5840a09..15d1fd6 100644 --- a/py/nodes/xyplot.py +++ b/py/nodes/xyplot.py @@ -528,6 +528,45 @@ class XYplot_Checkpoint: xy_values = {"axis": axis, "values": values, "lora_stack": optional_lora_stack} return (xy_values,) +# Diffusion Models +class XYplot_DiffusionModel: + @classmethod + def INPUT_TYPES(cls): + models = ["None"] + folder_paths.get_filename_list("diffusion_models") + clips = ["Auto"] + folder_paths.get_filename_list("text_encoders") + vaes = ["Auto"] + folder_paths.get_filename_list("vae") + + inputs = { + "required": { + "model_count": ("INT", {"default": 3, "min": 0, "max": 10, "step": 1}), + } + } + for i in range(1, 11): + inputs["required"][f"model_name_{i}"] = (models,) + inputs["required"][f"clip_name_{i}"] = (clips, {"default": "Auto"}) + inputs["required"][f"vae_name_{i}"] = (vaes, {"default": "Auto"}) + return inputs + + RETURN_TYPES = ("X_Y",) + RETURN_NAMES = ("X or Y",) + FUNCTION = "xy_value" + CATEGORY = "EasyUse/XY Inputs" + + def xy_value(self, model_count, **kwargs): + values = [] + for i in range(1, model_count + 1): + model_name = kwargs.get(f"model_name_{i}") + if not model_name or model_name == "None": + continue + clip_name = kwargs.get(f"clip_name_{i}", "Auto") + vae_name = kwargs.get(f"vae_name_{i}", "Auto") + values.append( + model_name.replace(",", "*") + "," + + clip_name.replace(",", "*") + "," + + vae_name.replace(",", "*") + ) + return ({"axis": "advanced: DiffusionModel", "values": values},) + #Loras class XYplot_Lora: @@ -670,6 +709,7 @@ NODE_CLASS_MAPPINGS = { "easy XYInputs: Sampler/Scheduler": XYplot_Sampler_Scheduler, "easy XYInputs: Denoise": XYplot_Denoise, "easy XYInputs: Checkpoint": XYplot_Checkpoint, + "easy XYInputs: DiffusionModel": XYplot_DiffusionModel, "easy XYInputs: Lora": XYplot_Lora, "easy XYInputs: ModelMergeBlocks": XYplot_ModelMergeBlocks, "easy XYInputs: PromptSR": XYplot_PromptSR, @@ -688,6 +728,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "easy XYInputs: Sampler/Scheduler": "XY Inputs: Sampler/Scheduler //EasyUse", "easy XYInputs: Denoise": "XY Inputs: Denoise //EasyUse", "easy XYInputs: Checkpoint": "XY Inputs: Checkpoint //EasyUse", + "easy XYInputs: DiffusionModel": "XY Inputs: Diffusion Model //EasyUse", "easy XYInputs: Lora": "XY Inputs: Lora //EasyUse", "easy XYInputs: ModelMergeBlocks": "XY Inputs: ModelMergeBlocks //EasyUse", "easy XYInputs: PromptSR": "XY Inputs: PromptSR //EasyUse", diff --git a/tests/test_diffusion_xy_helpers.py b/tests/test_diffusion_xy_helpers.py new file mode 100644 index 0000000..b5b74f8 --- /dev/null +++ b/tests/test_diffusion_xy_helpers.py @@ -0,0 +1,275 @@ +import importlib.util +import os +import sys +import tempfile +import types +import unittest +from contextlib import contextmanager +from pathlib import Path + +import torch + + +PLUGIN_ROOT = Path(__file__).parents[1] +UTILS_PATH = PLUGIN_ROOT / "py" / "libs" / "utils.py" +CONFIG_PATH = PLUGIN_ROOT / "py" / "config.py" +LOADER_PATH = PLUGIN_ROOT / "py" / "libs" / "loader.py" +XYPLOT_NODE_PATH = PLUGIN_ROOT / "py" / "nodes" / "xyplot.py" +XYPLOT_LIB_PATH = PLUGIN_ROOT / "py" / "libs" / "xyplot.py" + + +@contextmanager +def installed_modules(modules): + added = [] + for name, module in modules.items(): + if name not in sys.modules: + added.append(name) + sys.modules[name] = module + try: + yield + finally: + for name in added: + sys.modules.pop(name, None) + + +def load_module(name, path, package=None): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + if package is not None: + module.__package__ = package + spec.loader.exec_module(module) + return module + + +def make_package(name, **attrs): + package = types.ModuleType(name) + package.__path__ = [] + for key, value in attrs.items(): + setattr(package, key, value) + return package + + +def comfy_stubs(): + model_management = types.ModuleType("comfy.model_management") + model_base = types.ModuleType("comfy.model_base") + model_base.BaseModel = object + supported_models_base = types.ModuleType("comfy.supported_models_base") + supported_models_base.BASE = object + supported_models = types.ModuleType("comfy.supported_models") + supported_models.supported_models_base = supported_models_base + comfy = make_package( + "comfy", + model_management=model_management, + model_base=model_base, + supported_models_base=supported_models_base, + supported_models=supported_models, + ) + for name in ( + "SDXL", "SDXLRefiner", "SD15", "SD20", "SVD_img2vid", "SD3", + "HunyuanDiT", "Flux", "GenmoMochi", "Anima", "Krea2", + ): + setattr(supported_models, name, type(name, (), {})) + server = types.ModuleType("server") + server.PromptServer = object + return { + "comfy": comfy, + "comfy.model_management": model_management, + "comfy.model_base": model_base, + "comfy.supported_models_base": supported_models_base, + "comfy.supported_models": supported_models, + "server": server, + } + + +def xyplot_node_stubs(): + folder_paths = types.ModuleType("folder_paths") + folder_paths.get_filename_list = lambda folder: [] + return { + "comfy": make_package("comfy"), + "folder_paths": folder_paths, + "fake_py": make_package("fake_py"), + "fake_py.nodes": make_package("fake_py.nodes"), + "fake_py.libs": make_package("fake_py.libs"), + "fake_py.config": make_package("fake_py.config", RESOURCES_DIR="resources"), + "fake_py.libs.utils": make_package("fake_py.libs.utils", getMetadata=lambda *args, **kwargs: None), + } + + +def xyplot_lib_stubs(): + fake_utils = make_package("fake_py.utils", easySave=object, get_sd_version=lambda model: "unknown") + fake_adv_encode = make_package("fake_py.libs.adv_encode", advanced_encode=object) + fake_controlnet = make_package("fake_py.libs.controlnet", easyControlnet=object) + fake_log = make_package("fake_py.libs.log", log_node_warn=lambda *args, **kwargs: None) + return { + "nodes": make_package("nodes", CLIPTextEncode=object), + "fake_py": make_package("fake_py"), + "fake_py.libs": make_package("fake_py.libs"), + "fake_py.modules": make_package("fake_py.modules"), + "fake_py.utils": fake_utils, + "fake_py.libs.utils": fake_utils, + "fake_py.libs.adv_encode": fake_adv_encode, + "fake_py.libs.controlnet": fake_controlnet, + "fake_py.libs.log": fake_log, + "fake_py.modules.layer_diffuse": make_package("fake_py.modules.layer_diffuse", LayerDiffuse=object), + "fake_py.config": make_package("fake_py.config", RESOURCES_DIR="resources"), + } + + +def loader_stubs(): + comfy = make_package("comfy") + comfy.utils = make_package("comfy.utils") + comfy.sd = make_package("comfy.sd") + comfy.controlnet = make_package("comfy.controlnet") + comfy.model_patcher = make_package("comfy.model_patcher", ModelPatcher=type("ModelPatcher", (), {})) + folder_paths = make_package("folder_paths") + folder_paths.get_full_path = lambda folder, name: None + folder_paths.get_folder_paths = lambda folder: [] + folder_paths.get_filename_list = lambda folder: [] + fake_log = make_package("fake_py.libs.log", log_node_info=lambda *args, **kwargs: None, log_node_error=lambda *args, **kwargs: None) + fake_utils = make_package("fake_py.libs.utils", get_sd_version=lambda model: "unknown") + fake_config = make_package( + "fake_py.config", + DIFFUSION_MODEL_XY_DEFAULTS={}, + DIFFUSION_MODEL_CLIP_TYPES={"anima": "anima", "krea2": "krea2"}, + ) + fake_pixart = make_package("fake_py.modules.dit.pixArt.loader", load_pixart=object) + return { + "comfy": comfy, + "comfy.utils": comfy.utils, + "comfy.sd": comfy.sd, + "comfy.controlnet": comfy.controlnet, + "comfy.model_patcher": comfy.model_patcher, + "folder_paths": folder_paths, + "nodes": make_package("nodes", NODE_CLASS_MAPPINGS={}), + "fake_py": make_package("fake_py"), + "fake_py.libs": make_package("fake_py.libs"), + "fake_py.modules": make_package("fake_py.modules"), + "fake_py.modules.dit": make_package("fake_py.modules.dit"), + "fake_py.modules.dit.pixArt": make_package("fake_py.modules.dit.pixArt"), + "fake_py.libs.log": fake_log, + "fake_py.libs.utils": fake_utils, + "fake_py.config": fake_config, + "fake_py.modules.dit.pixArt.loader": fake_pixart, + } + + +class FakeModelPatcher: + def __init__(self, model_config=None, latent_format=None): + self.model = types.SimpleNamespace(model_config=model_config, latent_format=latent_format) + + +class FakeLatentFormat: + latent_dimensions = 3 + latent_channels = 16 + + +class DiffusionXYHelperTests(unittest.TestCase): + def test_get_sd_version_anima_and_krea2(self): + with installed_modules(comfy_stubs()): + utils = load_module("diffusion_xy_test_utils", UTILS_PATH) + + anima_config = utils.comfy.supported_models.Anima() + self.assertEqual(utils.get_sd_version(FakeModelPatcher(anima_config)), "anima") + + krea2_config = utils.comfy.supported_models.Krea2() + self.assertEqual(utils.get_sd_version(FakeModelPatcher(krea2_config)), "krea2") + + def test_diffusion_model_xy_defaults_are_complete(self): + with tempfile.TemporaryDirectory() as models_dir: + folder_paths = types.ModuleType("folder_paths") + folder_paths.models_dir = models_dir + with installed_modules({"folder_paths": folder_paths}): + config = load_module("diffusion_xy_test_config", CONFIG_PATH) + + for family in ("anima", "krea2"): + defaults = config.DIFFUSION_MODEL_XY_DEFAULTS[family] + self.assertTrue(defaults["clip_name"]) + self.assertTrue(defaults["clip_type"]) + self.assertTrue(defaults["vae_name"]) + self.assertEqual(config.DIFFUSION_MODEL_CLIP_TYPES[family], family) + + def test_load_diffusion_model_required_rejects_missing_clip_and_vae(self): + with installed_modules(loader_stubs()): + loader_module = load_module("fake_py.libs.loader", LOADER_PATH, "fake_py.libs") + loader = loader_module.easyLoader.__new__(loader_module.easyLoader) + loader.load_diffusion_model = lambda model_name: ("model", model_name) + loader.load_clip = lambda clip_name, type='stable_diffusion': ("clip", clip_name, type) + loader.load_vae = lambda vae_name: ("vae", vae_name) + loader_module.get_sd_version = lambda model: "krea2" + + with self.assertRaisesRegex(RuntimeError, "clip_name is required"): + loader.load_diffusion_model_required("model.safetensors", "None", "vae.safetensors") + + with self.assertRaisesRegex(RuntimeError, "vae_name is required"): + loader.load_diffusion_model_required("model.safetensors", "clip.safetensors", None) + + model, clip, vae, family = loader.load_diffusion_model_required( + "model.safetensors", "clip.safetensors", "vae.safetensors" + ) + self.assertEqual(family, "krea2") + self.assertEqual(clip[2], "krea2") + + loader_module.get_sd_version = lambda model: "flux" + with self.assertRaisesRegex(RuntimeError, "unsupported diffusion model family: flux"): + loader.load_diffusion_model_required("model.safetensors", "clip.safetensors", "vae.safetensors") + + def test_xyplot_diffusion_model_value_format(self): + with installed_modules(xyplot_node_stubs()): + node_module = load_module("fake_py.nodes.xyplot", XYPLOT_NODE_PATH, "fake_py.nodes") + node = node_module.XYplot_DiffusionModel() + + result = node.xy_value( + 2, + model_name_1="waiANIMA_v10Base10.safetensors", + clip_name_1="qwen_3_06b_base.safetensors", + vae_name_1="qwen_image_vae.safetensors", + model_name_2="moodyKrea2Mix,v70.safetensors", + clip_name_2="Auto", + vae_name_2="Auto", + ) + + self.assertEqual(result[0]["axis"], "advanced: DiffusionModel") + self.assertEqual( + result[0]["values"], + [ + "waiANIMA_v10Base10.safetensors,qwen_3_06b_base.safetensors,qwen_image_vae.safetensors", + "moodyKrea2Mix*v70.safetensors,Auto,Auto", + ], + ) + + model_name, clip_name, vae_name = result[0]["values"][0].split(",") + self.assertEqual(model_name.replace("*", ","), "waiANIMA_v10Base10.safetensors") + self.assertEqual(clip_name.replace("*", ","), "qwen_3_06b_base.safetensors") + self.assertEqual(vae_name.replace("*", ","), "qwen_image_vae.safetensors") + + model_name, clip_name, vae_name = result[0]["values"][1].split(",") + self.assertEqual(model_name.replace("*", ","), "moodyKrea2Mix,v70.safetensors") + self.assertEqual(clip_name, "Auto") + self.assertEqual(vae_name, "Auto") + + def test_ensure_latent_raises_for_nonempty_4d_latent_without_image(self): + with installed_modules(xyplot_lib_stubs()): + xyplot_module = load_module("fake_py.libs.xyplot", XYPLOT_LIB_PATH, "fake_py.libs") + + model = FakeModelPatcher(latent_format=FakeLatentFormat()) + vae = types.SimpleNamespace(encode=lambda pixels: torch.zeros([pixels.shape[0], 16, 1, pixels.shape[2], pixels.shape[3]])) + samples = {"samples": torch.ones([1, 4, 64, 64])} + + with self.assertRaisesRegex(RuntimeError, "requires an input image"): + xyplot_module.easyXYPlot._ensure_latent_for_model(model, vae, samples, {}) + + def test_ensure_latent_expands_empty_4d_latent_to_5d(self): + with installed_modules(xyplot_lib_stubs()): + xyplot_module = load_module("fake_py.libs.xyplot", XYPLOT_LIB_PATH, "fake_py.libs") + + model = FakeModelPatcher(latent_format=FakeLatentFormat()) + vae = types.SimpleNamespace(encode=lambda pixels: torch.zeros([pixels.shape[0], 16, 1, pixels.shape[2], pixels.shape[3]])) + samples = {"samples": torch.zeros([1, 4, 64, 64])} + + result = xyplot_module.easyXYPlot._ensure_latent_for_model(model, vae, samples, {}) + + self.assertEqual(result["samples"].shape, torch.Size([1, 16, 1, 64, 64])) + + +if __name__ == "__main__": + unittest.main()