nodes/doubley2000-ComfyUI-PS-Connector
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
doubley2000
2026-04-18 00:28:29 +08:00
committed by GitHub
parent 3b5a22879d
commit fa6b607845
3 changed files with 485 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
PSWorkflowConfig.py
+304
View File
@@ -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" # 注册新节点显示名称
}
__init__.py
+47
View File
@@ -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']
routes.py
+134
View File
@@ -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"})