From fa6b60784586172d7cc86ce29e78ecb2b6cbc331 Mon Sep 17 00:00:00 2001 From: doubley2000 <111101395+doubley2000@users.noreply.github.com> Date: Sat, 18 Apr 2026 00:28:29 +0800 Subject: [PATCH] Add files via upload --- PSWorkflowConfig.py | 304 ++++++++++++++++++++++++++++++++++++++++++++ __init__.py | 47 +++++++ routes.py | 134 +++++++++++++++++++ 3 files changed, 485 insertions(+) create mode 100644 PSWorkflowConfig.py create mode 100644 __init__.py create mode 100644 routes.py diff --git a/PSWorkflowConfig.py b/PSWorkflowConfig.py new file mode 100644 index 0000000..ffbac6e --- /dev/null +++ b/PSWorkflowConfig.py @@ -0,0 +1,304 @@ +# PSWorkflowConfig.py +import os +import sys +import json +import torch +import torch.nn.functional as F +from PIL import Image +import numpy as np +import folder_paths +import random + +class AnyType(str): + def __ne__(self, __value: object) -> bool: + return False +any_type = AnyType("*") + +class PSGetNode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + # 节点本身的单行文本，供 PS UI 做 Label 备用 + "ui_label": ("STRING", {"default": ""}), + # 新增：布尔项，用于定义图像或mask的必要性（前端 JS 会根据连线类型动态显隐） + "is_required": ("BOOLEAN", {"default": True}), + }, + "optional": { + # 接受任意类型的输入（Image, Mask, Int, Float, String 等） + "get": (any_type,), + } + } + + RETURN_TYPES = (any_type,) + RETURN_NAMES = ("get",) + FUNCTION = "passthrough" + CATEGORY = "PS_Connector" + + def passthrough(self, ui_label="", is_required=True, get=None): + # 运行时作为一个纯粹的直通节点，不修改任何数据 + return (get,) + +class PSTextReceiver: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "text": ("STRING", {"forceInput": True}), + } + } + + RETURN_TYPES = () + OUTPUT_NODE = True + FUNCTION = "receive_text" + CATEGORY = "PS_Connector" + + def receive_text(self, text): + # 抛出 UI 字典，前端 JS websocket 会自动将其截获并识别 + return {"ui": {"text": [text]}} +# ========================================== + +class PSWorkflowConfig: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "lora": ("MODEL",), + "prompt": ("STRING", {"forceInput": True}), + "size_logic": ("BOOLEAN", {"default": True}), + "width": ("INT", {"default": 512}), + "height": ("INT", {"default": 512}), + "long_side_pixels": ("INT", {"default": 1024, "min": 256, "max": 8192, "step": 8}), + "img1_req": ("BOOLEAN", {"default": True}), + "msk1_req": ("BOOLEAN", {"default": True}), + "img2_req": ("BOOLEAN", {"default": True}), + "msk2_req": ("BOOLEAN", {"default": True}), + "img3_req": ("BOOLEAN", {"default": True}), + "msk3_req": ("BOOLEAN", {"default": True}), + "img4_req": ("BOOLEAN", {"default": True}), + "msk4_req": ("BOOLEAN", {"default": True}), + }, + "optional": { + "image1": ("IMAGE",), + "mask1": ("MASK",), + "image2": ("IMAGE",), + "mask2": ("MASK",), + "image3": ("IMAGE",), + "mask3": ("MASK",), + "image4": ("IMAGE",), + "mask4": ("MASK",), + "neg_prompt": ("STRING", {"forceInput": True}), + } + } + RETURN_TYPES = () + FUNCTION = "do_nothing" + CATEGORY = "PS_Connector" + def do_nothing(self, **kwargs): + return () + +class PSImageAndMaskScaler: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "width": ("INT", {"default": 512}), + "height": ("INT", {"default": 512}) + }, + "optional": { + "image": ("IMAGE",), + "mask": ("MASK",) + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "scale" + CATEGORY = "PS_Connector" + def scale(self, width, height, image=None, mask=None): + out_image = image + out_mask = mask + + if image is not None: + img = image.permute(0, 3, 1, 2) + img = F.interpolate(img, size=(height, width), mode='bicubic', align_corners=False) + out_image = img.permute(0, 2, 3, 1).clamp(0, 1) + + if mask is not None: + if len(mask.shape) == 2: + msk = mask.unsqueeze(0).unsqueeze(0) + elif len(mask.shape) == 3: + msk = mask.unsqueeze(1) + else: + msk = mask + msk = F.interpolate(msk, size=(height, width), mode='bicubic', align_corners=False) + out_mask = msk.squeeze(1).clamp(0, 1) + + if out_image is None: + out_image = torch.zeros((1, height, width, 3), dtype=torch.float32) + if out_mask is None: + out_mask = torch.zeros((1, height, width), dtype=torch.float32) + return (out_image, out_mask) + + +# ======= 新增/修改：并列预览节点 ======= +class PSImagePreview: + def __init__(self): + self.output_dir = folder_paths.get_temp_directory() + self.type = "temp" + self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) + + @classmethod + def INPUT_TYPES(s): + return { + "required": {}, + "optional": { + "image1": ("IMAGE",), + "image2": ("IMAGE",), + "mask": ("MASK",), + } + } + + RETURN_TYPES = () + OUTPUT_NODE = True + FUNCTION = "preview" + CATEGORY = "PS_Connector" + + def scale_tensor(self, tensor_img, target_width, target_height, is_mask=False): + if is_mask: + if len(tensor_img.shape) == 2: + msk = tensor_img.unsqueeze(0).unsqueeze(0) + elif len(tensor_img.shape) == 3: + msk = tensor_img.unsqueeze(1) + else: + msk = tensor_img + msk = F.interpolate(msk, size=(target_height, target_width), mode='bicubic', align_corners=False) + return msk.squeeze(1).clamp(0, 1) + else: + img = tensor_img.permute(0, 3, 1, 2) + img = F.interpolate(img, size=(target_height, target_width), mode='bicubic', align_corners=False) + return img.permute(0, 2, 3, 1).clamp(0, 1) + + def preview(self, image1=None, image2=None, mask=None): + if image1 is None and image2 is None and mask is None: + return {"ui": {"images": []}} + + target_width = None + target_height = None + + # 尺寸逻辑判定 + if image2 is not None: + target_height = image2.shape[1] + target_width = image2.shape[2] + elif image1 is not None: + target_height = image1.shape[1] + target_width = image1.shape[2] + elif mask is not None: + if len(mask.shape) == 2: + target_height, target_width = mask.shape + else: + target_height, target_width = mask.shape[1], mask.shape[2] + + processed_images = [] + + # 处理 image1 + if image1 is not None: + if image1.shape[1] != target_height or image1.shape[2] != target_width: + img1_scaled = self.scale_tensor(image1, target_width, target_height, is_mask=False) + else: + img1_scaled = image1 + processed_images.append(("img1", img1_scaled)) + + # 处理 image2 + if image2 is not None: + processed_images.append(("img2", image2)) + + # 处理 mask + if mask is not None: + if len(mask.shape) == 2: + h, w = mask.shape + else: + h, w = mask.shape[1], mask.shape[2] + + if h != target_height or w != target_width: + mask_scaled = self.scale_tensor(mask, target_width, target_height, is_mask=True) + else: + mask_scaled = mask + + # 将 mask 转为 3 通道图以便正常显示 + if len(mask_scaled.shape) == 2: + mask_scaled = mask_scaled.unsqueeze(0).unsqueeze(-1).repeat(1, 1, 1, 3) + elif len(mask_scaled.shape) == 3: + mask_scaled = mask_scaled.unsqueeze(-1).repeat(1, 1, 1, 3) + processed_images.append(("mask", mask_scaled)) + + if not processed_images: + return {"ui": {"images": []}} + + results = [] + + # 分别保存每张图片，而不是拼合成一张。返回的多个图片会在 UI 中并列展示。 + for img_idx, (name, img_tensor) in enumerate(processed_images): + for batch_number, image in enumerate(img_tensor): + i = 255. * image.cpu().numpy() + img_pil = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + + # 文件名包含原输入标识(img1/img2/mask)保证不会相互覆盖 + filename = f"_{name}.png" + full_path = os.path.join(self.output_dir, filename) + img_pil.save(full_path, compress_level=4) + + results.append({ + "filename": filename, + "subfolder": "", + "type": self.type + }) + + return {"ui": {"images": results}} + +# ======= 新增：获取图像尺寸及逻辑控制节点 ======= +class PSGetImageSize: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "width": ("INT", {"default": 512, "min": 1, "max": 16384, "step": 1}), + "height": ("INT", {"default": 512, "min": 1, "max": 16384, "step": 1}), + # 逻辑判断项：默认开启，开启时输出图像尺寸，关闭时输出输入框尺寸 + "use_image_size": ("BOOLEAN", {"default": True, "label": "Use Image Size"}), + } + } + + RETURN_TYPES = ("INT", "INT") + RETURN_NAMES = ("width", "height") + FUNCTION = "get_size" + CATEGORY = "PS_Connector" + + def get_size(self, image, width, height, use_image_size): + if use_image_size: + # ComfyUI 的图像 tensor 形状通常为 (batch, height, width, channels) + img_height = image.shape[1] + img_width = image.shape[2] + return (int(img_width), int(img_height)) + else: + return (int(width), int(height)) + +# =============================================== +# 注册字典追加修改 +NODE_CLASS_MAPPINGS = { + "PSWorkflowConfig": PSWorkflowConfig, + "PSImageAndMaskScaler": PSImageAndMaskScaler, + "PSTextReceiver": PSTextReceiver, + "PSGetNode": PSGetNode, + "PSImagePreview": PSImagePreview, + "PSGetImageSize": PSGetImageSize # 注册新节点 +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PSWorkflowConfig": "PS Workflow Config (PS)", + "PSImageAndMaskScaler": "PS Auto Scaler (Hidden)", + "PSTextReceiver": "PS Text Receiver (LLM)", + "PSGetNode": "PS Get (Custom UI)", + "PSImagePreview": "PS Image Preview", + "PSGetImageSize": "PS Get Image Size" # 注册新节点显示名称 +} diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..b9e906b --- /dev/null +++ b/__init__.py @@ -0,0 +1,47 @@ +# __init__.py +import os +import sys +import subprocess +from server import PromptServer +from .routes import get_workflows, read_txt, refresh_models +from .PSWorkflowConfig import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +# ====================================================================== +# 自动安装前置依赖：ComfyUI-Lora-Manager +# ====================================================================== +def check_and_install_lora_manager(): + current_dir = os.path.dirname(os.path.abspath(__file__)) + custom_nodes_dir = os.path.dirname(current_dir) + target_folder = "ComfyUI-Lora-Manager" + target_path = os.path.join(custom_nodes_dir, target_folder) + repo_url = "https://github.com/willmiao/ComfyUI-Lora-Manager.git" + + if not os.path.exists(target_path): + print(f"\n[PS_Connector] 检测到缺失核心依赖: {target_folder}。正在自动为您安装...") + try: + result = subprocess.run( + ["git", "clone", repo_url, target_path], + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True + ) + print(f"[PS_Connector] ✅ 依赖插件 {target_folder} 安装成功！") + print(f"[PS_Connector] ⚠️ 注意：由于是首次安装依赖，建议重启 ComfyUI 以确保其完全加载。\n") + except Exception as e: + print(f"\n[PS_Connector] ❌ 自动安装失败！请手动前往 custom_nodes 目录执行:") + print(f"git clone {repo_url}") + print(f"错误详情: {e}\n") + +check_and_install_lora_manager() +# ====================================================================== + +server = PromptServer.instance +server.routes.get("/ps_helper/workflows")(get_workflows) +server.routes.post("/ps_helper/refresh")(refresh_models) +server.routes.get("/ps_helper/read_txt")(read_txt) + +# 告诉 ComfyUI 加载当前目录下的 js 文件夹作为前端插件 +WEB_DIRECTORY = "./js" + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] diff --git a/routes.py b/routes.py new file mode 100644 index 0000000..d78b23a --- /dev/null +++ b/routes.py @@ -0,0 +1,134 @@ +import os +import json +import folder_paths +from aiohttp import web + +def scan_folder(folder_name): + base_paths = folder_paths.get_folder_paths(folder_name) + if not base_paths: return [] + valid_exts = {".safetensors", ".ckpt", ".pt", ".pth", ".gguf", ".bin"} + results = [] + for bp in base_paths: + if not os.path.exists(bp): continue + for root, dirs, files in os.walk(bp): + for file in files: + ext = os.path.splitext(file)[1].lower() + if ext not in valid_exts: continue + full_path = os.path.join(root, file) + full_path_norm = os.path.normpath(full_path).replace("\\", "/") + rel_path = os.path.relpath(full_path, bp).replace("\\", "/") + base_path_no_ext, _ = os.path.splitext(full_path) + json_path = base_path_no_ext + ".metadata.json" if os.path.exists(base_path_no_ext + ".metadata.json") else (base_path_no_ext + ".json" if os.path.exists(base_path_no_ext + ".json") else None) + model_name, base_model, trained_words = "-", "-", "-" + if json_path: + try: + with open(json_path, "r", encoding="utf-8") as f: + meta = json.load(f) + base_model = meta.get("base_model") or (meta.get("civitai", {}).get("baseModel") if "civitai" in meta else None) or "-" + model_name = meta.get("model_name") or (meta.get("civitai", {}).get("model", {}).get("name") if "civitai" in meta else None) or "-" + tw = meta.get("trainedWords") or (meta.get("civitai", {}).get("trainedWords") if "civitai" in meta else None) + trained_words = ",".join([str(w) for w in tw]) if isinstance(tw, list) else (str(tw).strip() if isinstance(tw, str) and tw.strip() else "-") + except: pass + results.append({"file_name": rel_path, "model_name": model_name, "file_path": full_path_norm, "base_model": base_model, "trained_words": trained_words}) + return results + +def update_txt(path, new_items, is_lora=False): + # 1. 读取旧 TXT 数据，以相对路径(file_name)为 Key 建立索引 + existing_map = {} + if os.path.exists(path): + with open(path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: continue + + # 提取干扰符号，但要记住用户的特殊标记 + has_star = line.startswith("*") + has_amp = line.startswith("&") + + clean_line = line.lstrip("*&") + parts = clean_line.replace("｜", "|").split("|") + + if len(parts) >= 3: + rel_path = parts[0].strip() + old_base = parts[3].strip() if len(parts) > 3 else "-" + existing_map[rel_path] = { + "raw_line": line, # 核心：完整保留用户原来的那一行！ + "has_star": has_star, + "has_amp": has_amp, + "old_base": old_base + } + + # 2. 核心修复：基于绝对物理路径(file_path)对新扫描的数据进行去重！ + # 解决 ComfyUI scan_folder 扫描 overlapping 文件夹导致的天生重复问题 + unique_new_items = {} + for item in new_items: + unique_new_items[item["file_path"]] = item + + final_lines = [] + + # 3. 构建最终写入的数据 + for item in unique_new_items.values(): + rel_path = item["file_name"] + new_base = str(item["base_model"]).strip() + is_new_unknown = (new_base == "-" or new_base.lower() == "unknown") + + if rel_path in existing_map: + # 如果是老文件，检查是否需要“智能更新底模信息” + old_data = existing_map[rel_path] + old_raw_line = old_data["raw_line"] + old_base = old_data["old_base"] + old_is_unknown = (old_base == "-" or old_base.lower() == "unknown" or old_data["has_amp"]) + + if old_is_unknown and not is_new_unknown: + # 触发修复逻辑：以前不知道底模，现在扫到了同名 JSON 知道了底模 + prefix = "*" if old_data["has_star"] else "" # 继承用户的隐藏标记 + line = f"{prefix}{item['file_name']}｜{item['model_name']}｜{item['file_path']}｜{item['base_model']}" + if is_lora: + tw = item['trained_words'] if item['trained_words'] else "-" + line += f"｜{tw}" + final_lines.append(line) + else: + # 正常情况：原封不动地保留老数据！（完美保护你在 TXT 里手动写的内容和标记） + final_lines.append(old_raw_line) + else: + # 这是一个全新下载的文件 + prefix = "&" if is_new_unknown else "" + line = f"{prefix}{item['file_name']}｜{item['model_name']}｜{item['file_path']}｜{item['base_model']}" + if is_lora: + tw = item['trained_words'] if item['trained_words'] else "-" + line += f"｜{tw}" + final_lines.append(line) + + # 4. 覆写到磁盘 + with open(path, "w", encoding="utf-8") as f: + f.write("\n".join(final_lines) + "\n") + +async def get_workflows(request): + base_dir = os.path.join(folder_paths.base_path, "user", "default", "workflows", "PSflows") + workflows = [] + if os.path.exists(base_dir): + for root, _, files in os.walk(base_dir): + for file in files: + if file.endswith(".json") and (file.startswith("(") or file.startswith("（")): + try: + with open(os.path.join(root, file), "r", encoding="utf-8") as f: + workflows.append({"name": file[:-5], "json": json.load(f)}) + except: pass + return web.json_response({"workflows": sorted(workflows, key=lambda x: x["name"])}) + +async def read_txt(request): + psdate_dir = os.path.join(folder_paths.base_path, "user", "default", "PSdate") + res = {"model_txt": "", "lora_txt": ""} + for k in ["model", "lora"]: + p = os.path.join(psdate_dir, f"{k}.txt") + if os.path.exists(p): + with open(p, "r", encoding="utf-8") as f: res[f"{k}_txt"] = f.read() + return web.json_response(res) + +async def refresh_models(request): + psdate_dir = os.path.join(folder_paths.base_path, "user", "default", "PSdate") + os.makedirs(psdate_dir, exist_ok=True) + models = scan_folder("checkpoints") + scan_folder("diffusion_models") + scan_folder("unet") + update_txt(os.path.join(psdate_dir, "model.txt"), models, False) + update_txt(os.path.join(psdate_dir, "lora.txt"), scan_folder("loras"), True) + return web.json_response({"status": "ok"}) \ No newline at end of file