fix: make reboot routes awaitable (#1045)

Refs #1045

README.ZH_CN.md

@@ -39,6 +39,7 @@
 - 支持 kolors 模型
 - 支持 flux 模型
 - 支持 惰性条件判断（ifElse）和 for循环
+- 支持 Anima 与 Krea2 diffusion 模型，可通过 `easy diffusionModelLoader` 加载（需显式选择文本编码器与 VAE），并使用 `easy XYInputs: DiffusionModel` 进行 XY 对比
 
 ## 👨🏻‍🔧 安装
 
@@ -52,6 +53,10 @@ git clone https://github.com/yolain/ComfyUI-Easy-Use
 
 ## 📜 更新日志
 
+**v1.4.1**
+
+- 修复 `easy saveText` 将文本输出限制在输出目录 #1032
+
 **v1.4.0**
 
 - 添加 `easy tableEditor` 节点 - 用于编辑和显示表格数据的节点

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:
@@ -47,6 +48,10 @@ Double-click install.bat to install the required dependencies
 
 ## 📜 Changelog
 
+**v1.4.1**
+
+- Fix `easy saveText` to confine text outputs to output directory #1032
+
 **v1.4.0**
 
 - Add `easy tableEditor` node - A node for editing and displaying table data

__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.7"
+__version__ = "1.4.1"
 
 import yaml
 import json
@@ -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"]

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": {
@@ -6271,6 +6438,25 @@
       }
     }
   },
+  "easy PassOrNone": {
+    "display_name": "Pass or None",
+    "inputs": {
+      "any": {
+        "name": "anything"
+      },
+      "default": {
+        "name": "default"
+      }
+    },
+    "outputs": {
+      "0": {
+        "name": "output"
+      },
+      "1": {
+        "name": "is_none"
+      }
+    }
+  },
   "easy isNone": {
     "display_name": "Is None",
     "inputs": {

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": {
@@ -6321,6 +6488,25 @@
       }
     }
   },
+  "easy PassOrNone": {
+    "display_name": "传递或为空",
+    "inputs": {
+      "any": {
+        "name": "任何"
+      },
+      "default": {
+        "name": "默认值"
+      }
+    },
+    "outputs": {
+      "0": {
+        "name": "输出"
+      },
+      "1": {
+        "name": "是否为空"
+      }
+    }
+  },
   "easy isNone": {
     "display_name": "是否为空",
     "inputs": {

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",
+}

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,)
 

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

py/libs/path_utils.py

@@ -0,0 +1,29 @@
+import os
+
+
+def resolve_output_file_path(output_root, output_file_path, file_name, file_extension):
+    """Resolve a workflow-provided output path beneath ``output_root``.
+
+    Relative output directories remain supported, but are interpreted relative
+    to ComfyUI's configured output directory rather than the process working
+    directory. Resolving both paths prevents ``..`` components and existing
+    symlinks from escaping the allowed root.
+    """
+    output_root = os.path.realpath(output_root)
+    requested_directory = output_file_path
+    if not os.path.isabs(requested_directory):
+        requested_directory = os.path.join(output_root, requested_directory)
+
+    candidate = os.path.realpath(
+        os.path.join(requested_directory, f"{file_name}.{file_extension}")
+    )
+    try:
+        is_within_output = os.path.commonpath((output_root, candidate)) == output_root
+    except ValueError:
+        # Different Windows drives and paths containing null bytes are unsafe.
+        is_within_output = False
+
+    if not is_within_output:
+        raise ValueError("Saving outside the ComfyUI output directory is not allowed")
+
+    return candidate

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
 

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'
 

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'],

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",

py/nodes/logic.py

@@ -8,6 +8,7 @@ from ..libs.utils import AlwaysEqualProxy, ByPassTypeTuple, cleanGPUUsedForce, c
 from ..libs.cache import cache, update_cache, remove_cache
 from ..libs.log import log_node_info, log_node_warn
 from ..libs.math import evaluate_formula
+from ..libs.path_utils import resolve_output_file_path
 import numpy as np
 import time
 import os
@@ -999,13 +1000,59 @@ class isMaskEmpty(io.ComfyNode):
         return io.NodeOutput(False)
 
 
+class PassOrNone(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="easy PassOrNone",
+            description="Passes the input through, or a default value when input is None/not provided, or outputs None when both inputs are None/not provided.",
+            category="EasyUse/Logic",
+            search_aliases=[
+                "null",
+                "nothing",
+                "empty",
+                "blank",
+            ],
+            inputs=[
+                io.AnyType.Input(
+                    "anything",
+                    tooltip="Passes the input through, or a default value when input is None/not provided, or outputs None when both inputs are None/not provided.",
+                    optional=True,
+                ),
+                io.AnyType.Input(
+                    "default",
+                    tooltip="Fallback value to use when the input is None/not provided.",
+                    optional=True,
+                ),
+            ],
+            outputs=[
+                io.AnyType.Output("output"),
+                io.Boolean.Output("is_none"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, anything=None, default=None):
+        if anything is None:
+            return io.NodeOutput(
+                default,
+                default is None,
+            )
+
+        return io.NodeOutput(
+            anything,
+            False,
+        )
+
+
 class isNone(io.ComfyNode):
     @classmethod
     def define_schema(cls):
         return io.Schema(
             node_id="easy isNone",
+            description="Returns true if the input is None, an empty string, or zero.",
             category="EasyUse/Logic",
-            inputs=[io.AnyType.Input("any")],
+            inputs=[io.AnyType.Input("any", tooltip="Returns true if the input is None, an empty string, or zero.")],
             outputs=[io.Boolean.Output("boolean")],
         )
 
@@ -1585,9 +1632,14 @@ class saveText(io.ComfyNode):
             log_node_warn("Save Text", "No file details found. No file output.")
             return io.NodeOutput(text, None)
 
-        filepath = os.path.join(output_file_path, file_name) + "." + file_extension
-        if not os.path.exists(output_file_path):
-            os.makedirs(output_file_path)
+        if file_extension not in ("txt", "csv"):
+            raise ValueError("Unsupported text file extension")
+
+        output_dir = folder_paths.get_output_directory()
+        filepath = resolve_output_file_path(
+            output_dir, output_file_path, file_name, file_extension
+        )
+        os.makedirs(os.path.dirname(filepath), exist_ok=True)
 
         file_mode = "w" if overwrite else "a"
         log_node_info("Save Text", f"Saving to {filepath}")
@@ -1604,27 +1656,24 @@ class saveText(io.ComfyNode):
 
         result_image = None
         if image is not None:
-            imagepath = os.path.join(output_file_path, file_name)
-            index = 1
+            imagepath = resolve_output_file_path(
+                output_dir, output_file_path, file_name, "png"
+            )
             if not overwrite:
-                while os.path.exists(filepath):
-                    imagepath = os.path.join(output_file_path, file_name) + "_" + str(index)
+                index = 1
+                while os.path.exists(imagepath):
+                    imagepath = resolve_output_file_path(
+                        output_dir, output_file_path, f"{file_name}_{index}", "png"
+                    )
                     index += 1
 
-            output_dir = folder_paths.output_directory
-            output_path_val = "" if output_file_path in [None, "", "none", "."] else output_file_path
-            if not os.path.isabs(output_file_path):
-                output_path_val = os.path.join(output_dir, output_path_val)
-            if output_path_val.strip():
-                if not os.path.isabs(output_path_val):
-                    output_path_val = os.path.join(folder_paths.output_directory, output_path_val)
-                if not os.path.exists(output_path_val.strip()):
-                    print(f"The path `{output_path_val.strip()}` does not exist! Creating directory.")
-                    os.makedirs(output_path_val, exist_ok=True)
+            image_output_path = os.path.dirname(imagepath)
+            os.makedirs(image_output_path, exist_ok=True)
 
             images_tensor = torch.cat([image], dim=0)
-            cls.save_image(images_tensor, imagepath, "png", 100, None, None,
-                           filename_number_start="true", output_path=output_path_val,
+            image_name = os.path.splitext(os.path.basename(imagepath))[0]
+            cls.save_image(images_tensor, image_name, "png", 100, None, None,
+                           filename_number_start="true", output_path=image_output_path,
                            delimiter="_", number_padding=4, lossless_webp=False)
             log_node_info("Save Text", f"Saving Image to {imagepath}")
             result_image = image
@@ -1679,6 +1728,7 @@ NODE_CLASS_MAPPINGS = {
     "easy blocker": Blocker,
     "easy ifElse": IfElse,
     "easy isMaskEmpty": isMaskEmpty,
+    "easy PassOrNone": PassOrNone,
     "easy isNone": isNone,
     "easy isSDXL": isSDXL,
     "easy isFileExist": isFileExist,
@@ -1728,6 +1778,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "easy ifElse": "If else",
     "easy blocker": "Blocker",
     "easy isMaskEmpty": "Is Mask Empty",
+    "easy PassOrNone": "Pass or None",
     "easy isNone": "Is None",
     "easy isSDXL": "Is SDXL",
     "easy isFileExist": "Is File Exist",

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

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)

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':

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",

py/routes.py

@@ -1,10 +1,16 @@
 import os
 import hashlib
+import hmac
 import sys
 import json
 import shutil
+import secrets
+import tempfile
+from functools import lru_cache
+from urllib.parse import urlsplit
 import folder_paths
 from aiohttp import web
+from PIL import Image, UnidentifiedImageError
 from server import PromptServer
 from .config import RESOURCES_DIR, FOOOCUS_STYLES_DIR, FOOOCUS_STYLES_SAMPLES
 from .libs.model import easyModelManager
@@ -50,8 +56,30 @@ async def translate(request):
     else:
         return web.json_response({"text": text})
 
-@PromptServer.instance.routes.get("/easyuse/reboot")
-def reboot(request):
+_reboot_token = secrets.token_urlsafe(32)
+
+
+def _same_origin_request(request):
+    fetch_site = request.headers.get("Sec-Fetch-Site")
+    if fetch_site and fetch_site not in ("same-origin", "none"):
+        return False
+    origin = request.headers.get("Origin")
+    return not origin or urlsplit(origin).netloc == request.host
+
+
+@PromptServer.instance.routes.get("/easyuse/reboot-token")
+async def get_reboot_token(request):
+    if not _same_origin_request(request):
+        return web.Response(status=403)
+    return web.json_response({"token": _reboot_token}, headers={"Cache-Control": "no-store"})
+
+
+@PromptServer.instance.routes.post("/easyuse/reboot")
+async def reboot(request):
+    token = request.headers.get("X-EasyUse-Reboot-Token", "")
+    if not _same_origin_request(request) or not hmac.compare_digest(token, _reboot_token):
+        return web.Response(status=403)
+
     try:
         sys.stdout.close_log()
     except Exception as e:
@@ -169,9 +197,9 @@ async def getModelsList(request):
 @PromptServer.instance.routes.post("/easyuse/metadata/notes/{name}")
 async def save_notes(request):
     name = request.match_info["name"]
-    pos = name.index("/")
-    type = name[0:pos]
-    name = name[pos+1:]
+    type, separator, name = name.partition("/")
+    if not separator or type not in ("checkpoints", "loras", "embeddings"):
+        return web.Response(status=400)
 
     file_path = None
     if type == "embeddings" or type == "loras":
@@ -189,24 +217,34 @@ async def save_notes(request):
             if file_path is not None:
                 break
     else:
-        file_path = folder_paths.get_full_path(
-            type, name)
+        if name in folder_paths.get_filename_list(type):
+            file_path = folder_paths.get_full_path(type, name)
     if not file_path:
         return web.Response(status=404)
 
     file_no_ext = os.path.splitext(file_path)[0]
     info_file = file_no_ext + ".txt"
-    with open(info_file, "w") as f:
-        f.write(await request.text())
+    staged_path = None
+    try:
+        with tempfile.NamedTemporaryFile(
+                mode="w", encoding="utf-8", dir=os.path.dirname(info_file),
+                prefix=".easyuse-notes-", delete=False
+        ) as staged:
+            staged_path = staged.name
+            staged.write(await request.text())
+        os.replace(staged_path, info_file)
+    finally:
+        if staged_path and os.path.exists(staged_path):
+            os.unlink(staged_path)
 
     return web.Response(status=200)
 
 @PromptServer.instance.routes.get("/easyuse/metadata/{name}")
 async def load_metadata(request):
     name = request.match_info["name"]
-    pos = name.index("/")
-    type = name[0:pos]
-    name = name[pos+1:]
+    type, separator, name = name.partition("/")
+    if not separator or type not in ("checkpoints", "loras", "embeddings"):
+        return web.Response(status=400)
 
     file_path = None
     if type == "embeddings":
@@ -224,7 +262,8 @@ async def load_metadata(request):
             if file_path is not None:
                 break
     else:
-        file_path = folder_paths.get_full_path(type, name)
+        if name in folder_paths.get_filename_list(type):
+            file_path = folder_paths.get_full_path(type, name)
     if not file_path:
         return web.Response(status=404)
 
@@ -241,47 +280,93 @@ async def load_metadata(request):
     file_no_ext = os.path.splitext(file_path)[0]
 
     info_file = file_no_ext + ".txt"
-    if os.path.isfile(info_file):
+    if os.path.isfile(info_file) and not os.path.islink(info_file):
         with open(info_file, "r") as f:
             meta["easyuse.notes"] = f.read()
 
-    hash_file = file_no_ext + ".sha256"
-    if os.path.isfile(hash_file):
-        with open(hash_file, "rt") as f:
-            meta["easyuse.sha256"] = f.read()
-    else:
-        with open(file_path, "rb") as f:
-            meta["easyuse.sha256"] = hashlib.sha256(f.read()).hexdigest()
-        with open(hash_file, "wt") as f:
-            f.write(meta["easyuse.sha256"])
+    # Sidecar hashes are user-controlled; never use them as proof of the model's hash.
+    stat = os.stat(file_path)
+    meta["easyuse.sha256"] = _model_sha256(
+        file_path, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns
+    )
 
     return web.json_response(meta)
 
+
+@lru_cache(maxsize=128)
+def _model_sha256(path, size, mtime_ns, ctime_ns):
+    digest = hashlib.sha256()
+    with open(path, "rb") as model:
+        for chunk in iter(lambda: model.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+_PREVIEW_FORMATS = {
+    ".png": "PNG",
+    ".jpg": "JPEG",
+    ".jpeg": "JPEG",
+    ".webp": "WEBP",
+    ".gif": "GIF",
+}
+
 @PromptServer.instance.routes.post("/easyuse/save/{name}")
 async def save_preview(request):
     name = request.match_info["name"]
-    pos = name.index("/")
-    type = name[0:pos]
-    name = name[pos+1:]
+    model_type, separator, model_name = name.partition("/")
+    if not separator or model_type not in ("checkpoints", "loras"):
+        return web.Response(status=400)
+    if model_name not in folder_paths.get_filename_list(model_type):
+        return web.Response(status=404)
+
+    model_path = folder_paths.get_full_path(model_type, model_name)
+    if not model_path:
+        return web.Response(status=404)
 
     body = await request.json()
-
-    dir = folder_paths.get_directory_by_type(body.get("type", "output"))
-    subfolder = body.get("subfolder", "")
-    full_output_folder = os.path.join(dir, os.path.normpath(subfolder))
-
-    if os.path.commonpath((dir, os.path.abspath(full_output_folder))) != dir:
+    filename = body.get("filename")
+    if (body.get("type") != "temp" or body.get("subfolder", "") != ""
+            or not isinstance(filename, str) or not filename
+            or os.path.basename(filename) != filename or filename in (".", "..")):
         return web.Response(status=400)
 
-    filepath = os.path.join(full_output_folder, body.get("filename", ""))
-    image_path = folder_paths.get_full_path(type, name)
-    image_path = os.path.splitext(
-        image_path)[0] + os.path.splitext(filepath)[1]
+    extension = os.path.splitext(filename)[1].lower()
+    if extension not in _PREVIEW_FORMATS:
+        return web.Response(status=400)
 
-    shutil.copyfile(filepath, image_path)
+    temp_dir = folder_paths.get_directory_by_type("temp")
+    filepath = os.path.join(temp_dir, filename)
+    if (os.path.commonpath((os.path.realpath(temp_dir), os.path.realpath(filepath)))
+            != os.path.realpath(temp_dir) or not os.path.isfile(filepath)):
+        return web.Response(status=400)
+
+    image_path = os.path.splitext(model_path)[0] + extension
+    if (os.path.islink(image_path)
+            or os.path.realpath(os.path.dirname(image_path))
+            != os.path.realpath(os.path.dirname(model_path))):
+        return web.Response(status=400)
+
+    staged_path = None
+    try:
+        with tempfile.NamedTemporaryFile(
+                dir=os.path.dirname(image_path), prefix=".easyuse-preview-", delete=False
+        ) as staged:
+            staged_path = staged.name
+            with open(filepath, "rb") as source:
+                shutil.copyfileobj(source, staged)
+        with Image.open(staged_path) as image:
+            if image.format != _PREVIEW_FORMATS[extension]:
+                return web.Response(status=400)
+            image.verify()
+        os.replace(staged_path, image_path)
+    except (OSError, ValueError, UnidentifiedImageError):
+        return web.Response(status=400)
+    finally:
+        if staged_path and os.path.exists(staged_path):
+            os.unlink(staged_path)
 
     return web.json_response({
-        "image":  type + "/" + os.path.basename(image_path)
+        "image":  model_type + "/" + os.path.basename(image_path)
     })
 
 @PromptServer.instance.routes.post("/easyuse/model/download")

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-easy-use"
 description = "To enhance the usability of ComfyUI, optimizations and integrations have been implemented for several commonly used nodes."
-version = "1.4.0"
+version = "1.4.1"
 license = { file = "LICENSE" }
 dependencies = ["diffusers", "accelerate", "clip_interrogator", "sentencepiece", "lark", "onnxruntime", "spandrel", "opencv-python-headless", "matplotlib", "peft"]
 

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()

tests/test_path_utils.py

@@ -0,0 +1,71 @@
+import importlib.util
+import os
+from pathlib import Path
+import tempfile
+import unittest
+
+
+MODULE_PATH = Path(__file__).parents[1] / "py" / "libs" / "path_utils.py"
+SPEC = importlib.util.spec_from_file_location("easyuse_path_utils", MODULE_PATH)
+path_utils = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(path_utils)
+
+
+class ResolveOutputFilePathTests(unittest.TestCase):
+    def test_relative_subdirectory_is_resolved_under_output_root(self):
+        with tempfile.TemporaryDirectory() as output_root:
+            result = path_utils.resolve_output_file_path(
+                output_root, "metadata", "prompt", "txt"
+            )
+
+            self.assertEqual(
+                result,
+                os.path.join(os.path.realpath(output_root), "metadata", "prompt.txt"),
+            )
+
+    def test_absolute_directory_inside_output_root_is_allowed(self):
+        with tempfile.TemporaryDirectory() as output_root:
+            inside = os.path.join(output_root, "metadata")
+
+            result = path_utils.resolve_output_file_path(
+                output_root, inside, "prompt", "txt"
+            )
+
+            self.assertEqual(result, os.path.join(inside, "prompt.txt"))
+
+    def test_absolute_directory_outside_output_root_is_rejected(self):
+        with tempfile.TemporaryDirectory() as output_root:
+            with tempfile.TemporaryDirectory() as outside:
+                with self.assertRaises(ValueError):
+                    path_utils.resolve_output_file_path(
+                        output_root, outside, "marker", "txt"
+                    )
+
+    def test_output_directory_traversal_is_rejected(self):
+        with tempfile.TemporaryDirectory() as output_root:
+            with self.assertRaises(ValueError):
+                path_utils.resolve_output_file_path(
+                    output_root, "../outside", "marker", "txt"
+                )
+
+    def test_file_name_traversal_is_rejected(self):
+        with tempfile.TemporaryDirectory() as output_root:
+            with self.assertRaises(ValueError):
+                path_utils.resolve_output_file_path(
+                    output_root, ".", "../../marker", "txt"
+                )
+
+    @unittest.skipUnless(hasattr(os, "symlink"), "symlinks are unavailable")
+    def test_symlink_escape_is_rejected(self):
+        with tempfile.TemporaryDirectory() as output_root:
+            with tempfile.TemporaryDirectory() as outside:
+                os.symlink(outside, os.path.join(output_root, "linked"))
+
+                with self.assertRaises(ValueError):
+                    path_utils.resolve_output_file_path(
+                        output_root, "linked", "marker", "txt"
+                    )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_routes_security.py

@@ -0,0 +1,174 @@
+import ast
+import asyncio
+import hashlib
+import inspect
+import json
+import os
+from pathlib import Path
+import shutil
+import sys
+import tempfile
+import unittest
+from functools import lru_cache
+from types import SimpleNamespace
+from urllib.parse import urlsplit
+from unittest.mock import patch
+
+from aiohttp import web
+from aiohttp.test_utils import TestClient, TestServer
+from PIL import Image, UnidentifiedImageError
+
+
+ROUTES_PATH = Path(__file__).parents[1] / "py" / "routes.py"
+
+
+def load_handlers(folder_paths, get_metadata):
+    """Load the actual handlers without importing ComfyUI's GPU dependencies."""
+    names = {
+        "_same_origin_request", "get_reboot_token", "reboot", "_model_sha256",
+        "load_metadata", "save_notes", "save_preview",
+    }
+    tree = ast.parse(ROUTES_PATH.read_text())
+    functions = []
+    for node in tree.body:
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in names:
+            node.decorator_list = []
+            functions.append(node)
+    namespace = {
+        "os": os, "sys": sys, "hashlib": hashlib, "hmac": __import__("hmac"),
+        "json": json, "shutil": shutil, "tempfile": tempfile,
+        "lru_cache": lru_cache, "urlsplit": urlsplit, "web": web,
+        "Image": Image, "UnidentifiedImageError": UnidentifiedImageError,
+        "folder_paths": folder_paths, "getMetadata": get_metadata,
+        "_reboot_token": "test-reboot-token",
+        "_PREVIEW_FORMATS": {
+            ".png": "PNG", ".jpg": "JPEG", ".jpeg": "JPEG",
+            ".webp": "WEBP", ".gif": "GIF",
+        },
+    }
+    exec(compile(ast.Module(body=functions, type_ignores=[]), str(ROUTES_PATH), "exec"), namespace)
+    return SimpleNamespace(**namespace)
+
+
+class SecurityRouteTests(unittest.TestCase):
+    def setUp(self):
+        self.workspace = tempfile.TemporaryDirectory()
+        self.addCleanup(self.workspace.cleanup)
+        self.root = Path(self.workspace.name)
+        self.model_dir = self.root / "models"
+        self.temp_dir = self.root / "temp"
+        self.model_dir.mkdir()
+        self.temp_dir.mkdir()
+        self.model_path = self.model_dir / "sample.safetensors"
+        self.model_path.write_bytes(b"model data")
+        paths = SimpleNamespace(
+            get_filename_list=lambda kind: [self.model_path.name],
+            get_full_path=lambda kind, name: str(self.model_path),
+            get_directory_by_type=lambda kind: str(self.temp_dir),
+        )
+        self.handlers = load_handlers(
+            paths,
+            lambda path: json.dumps({"__metadata__": {"easyuse.notes": "<img onerror=alert(1)>"}}),
+        )
+
+    def request(self, name="loras/sample.safetensors", filename="preview.png", **body):
+        payload = {"type": "temp", "filename": filename, **body}
+        return SimpleNamespace(
+            match_info={"name": name},
+            json=lambda: asyncio.sleep(0, result=payload),
+            headers={}, host="localhost:8188",
+        )
+
+    def test_save_rejects_script_and_custom_node_target(self):
+        (self.temp_dir / "payload.py").write_text("print('sentinel')")
+        response = asyncio.run(self.handlers.save_preview(self.request(filename="payload.py")))
+        self.assertEqual(response.status, 400)
+        response = asyncio.run(self.handlers.save_preview(
+            self.request(name="custom_nodes/package/__init__.py", filename="payload.py")
+        ))
+        self.assertEqual(response.status, 400)
+
+    def test_save_accepts_real_image_and_rejects_disguised_script(self):
+        Image.new("RGB", (1, 1)).save(self.temp_dir / "preview.png")
+        response = asyncio.run(self.handlers.save_preview(self.request()))
+        self.assertEqual(response.status, 200)
+        with Image.open(self.model_dir / "sample.png") as saved:
+            self.assertEqual(saved.format, "PNG")
+
+        (self.temp_dir / "preview.png").write_text("print('sentinel')")
+        response = asyncio.run(self.handlers.save_preview(self.request()))
+        self.assertEqual(response.status, 400)
+        with Image.open(self.model_dir / "sample.png") as saved:
+            self.assertEqual(saved.format, "PNG")
+
+    @unittest.skipUnless(hasattr(os, "symlink"), "symlinks are unavailable")
+    def test_save_does_not_follow_preview_symlink(self):
+        Image.new("RGB", (1, 1)).save(self.temp_dir / "preview.png")
+        protected = self.root / "protected.txt"
+        protected.write_text("untouched")
+        os.symlink(protected, self.model_dir / "sample.png")
+        response = asyncio.run(self.handlers.save_preview(self.request()))
+        self.assertEqual(response.status, 400)
+        self.assertEqual(protected.read_text(), "untouched")
+
+    def test_metadata_ignores_forged_hash_sidecar(self):
+        (self.model_dir / "sample.sha256").write_text("0" * 64)
+        response = asyncio.run(self.handlers.load_metadata(self.request()))
+        self.assertEqual(
+            json.loads(response.text)["easyuse.sha256"],
+            hashlib.sha256(self.model_path.read_bytes()).hexdigest(),
+        )
+
+    @unittest.skipUnless(hasattr(os, "symlink"), "symlinks are unavailable")
+    def test_notes_reject_custom_nodes_and_do_not_follow_symlinks(self):
+        request = self.request(name="custom_nodes/package/__init__.py")
+        request.text = lambda: asyncio.sleep(0, result="new notes")
+        self.assertEqual(asyncio.run(self.handlers.save_notes(request)).status, 400)
+
+        protected = self.root / "protected.txt"
+        protected.write_text("untouched")
+        os.symlink(protected, self.model_dir / "sample.txt")
+        request = self.request()
+        request.text = lambda: asyncio.sleep(0, result="new notes")
+        self.assertEqual(asyncio.run(self.handlers.save_notes(request)).status, 200)
+        self.assertEqual(protected.read_text(), "untouched")
+        self.assertEqual((self.model_dir / "sample.txt").read_text(), "new notes")
+
+    def test_reboot_requires_token_and_same_origin(self):
+        self.assertTrue(inspect.iscoroutinefunction(self.handlers.get_reboot_token))
+        self.assertTrue(inspect.iscoroutinefunction(self.handlers.reboot))
+        request = self.request()
+        request.headers = {"Sec-Fetch-Site": "cross-site"}
+        self.assertEqual(asyncio.run(self.handlers.get_reboot_token(request)).status, 403)
+        request.headers = {"Sec-Fetch-Site": "same-origin"}
+        self.assertEqual(json.loads(asyncio.run(self.handlers.get_reboot_token(request)).text)["token"], "test-reboot-token")
+        request.headers = {}
+        with patch.object(self.handlers.os, "execv", return_value="restarted") as restart:
+            self.assertEqual(asyncio.run(self.handlers.reboot(request)).status, 403)
+            request.headers = {"X-EasyUse-Reboot-Token": "test-reboot-token", "Origin": "http://other.test"}
+            self.assertEqual(asyncio.run(self.handlers.reboot(request)).status, 403)
+            restart.assert_not_called()
+            request.headers["Origin"] = "http://localhost:8188"
+            request.headers["Sec-Fetch-Site"] = "same-site"
+            self.assertEqual(asyncio.run(self.handlers.reboot(request)).status, 403)
+            request.headers["Sec-Fetch-Site"] = "same-origin"
+            self.assertEqual(asyncio.run(self.handlers.reboot(request)), "restarted")
+            restart.assert_called_once()
+
+    def test_reboot_routes_return_http_responses(self):
+        async def exercise_routes():
+            app = web.Application()
+            app.router.add_get("/easyuse/reboot-token", self.handlers.get_reboot_token)
+            app.router.add_post("/easyuse/reboot", self.handlers.reboot)
+            async with TestClient(TestServer(app)) as client:
+                token_response = await client.get("/easyuse/reboot-token")
+                self.assertEqual(token_response.status, 200)
+                self.assertEqual((await token_response.json())["token"], "test-reboot-token")
+                reboot_response = await client.post("/easyuse/reboot")
+                self.assertEqual(reboot_response.status, 403)
+
+        asyncio.run(exercise_routes())
+
+
+if __name__ == "__main__":
+    unittest.main()

web_version/v1/js/common/model.js

@@ -114,10 +114,10 @@ export class ModelInfoDialog extends ComfyDialog {
 
 			let pre = this.customNotes.substring(end, pos);
 			if (pre) {
-				pre = pre.replaceAll("\n", "<br>");
 				notes.push(
 					$el("span", {
-						innerHTML: pre,
+						textContent: pre,
+						style: { whiteSpace: "pre-line" },
 					})
 				);
 			}
@@ -127,6 +127,7 @@ export class ModelInfoDialog extends ComfyDialog {
 						href: m[0],
 						textContent: m[0],
 						target: "_blank",
+						rel: "noopener noreferrer",
 					})
 				);
 			}
@@ -335,7 +336,9 @@ export class ModelInfoDialog extends ComfyDialog {
 											const blob = await (await fetch(cate.url)).blob();
 
 											// Store it in temp
-											const name = "temp_preview." + new URL(cate.url).pathname.split(".")[1];
+											const extension = ({"image/png": "png", "image/jpeg": "jpg", "image/webp": "webp", "image/gif": "gif"})[blob.type]
+												|| new URL(cate.url).pathname.split(".").pop().toLowerCase();
+											const name = "temp_preview." + extension;
 											const body = new FormData();
 											body.append("image", new File([blob], name));
 											body.append("overwrite", "true");
@@ -365,10 +368,13 @@ export class ModelInfoDialog extends ComfyDialog {
 												headers: {
 													"content-type": "application/json",
 												},
-											}).then(_=>{
-												toast.success($t('Saving Succeed'))
-												toast.hideLoading()
-											});
+											}).then(response => {
+												if (!response.ok) throw new Error(`Error saving preview (${response.status})`);
+												toast.success($t('Saving Succeed'));
+											}).catch(error => {
+												console.error(error);
+												toast.error($t('Saving Failed'));
+											}).finally(() => toast.hideLoading());
 											this.isSaving = false
 											app.refreshComboInNodes();
 										},
@@ -680,4 +686,4 @@ export class LoraInfoDialog extends ModelInfoDialog {
 
 		return btns;
 	}
-}
+}

web_version/v1/js/easy/easy.js

@@ -486,10 +486,16 @@ app.registerExtension({
                 // Only show the reboot option if the server is running on a local network 仅在本地或局域网环境可重启服务
                 isLocalNetwork(window.location.host) ? {
                     content: rebootIcon.replace('currentColor','var(--error-color)') + ' '+ $t('Reboot ComfyUI') + ' (EasyUse)',
-                    callback: _ =>{
+                    callback: async _ =>{
                         if (confirm($t("Are you sure you'd like to reboot the server?"))){
                             try {
-                                api.fetchApi("/easyuse/reboot");
+                                const tokenResponse = await api.fetchApi("/easyuse/reboot-token");
+                                if (!tokenResponse.ok) throw new Error("Could not get reboot token");
+                                const {token} = await tokenResponse.json();
+                                await api.fetchApi("/easyuse/reboot", {
+                                    method: "POST",
+                                    headers: {"X-EasyUse-Reboot-Token": token},
+                                });
                             } catch (exception) {}
                         }
                     }
@@ -607,4 +613,4 @@ app.registerExtension({
             };
         }
     },
-});
+});