From 9fab4e953b4fb54f371b4fd92102db8029c07dae Mon Sep 17 00:00:00 2001 From: LAOGOU-666 <441059767@qq.com> Date: Fri, 6 Jun 2025 02:43:39 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=E6=A1=A5=E6=8E=A5=E9=A2=84?= =?UTF-8?q?=E8=A7=88V2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- py/bridge_preview.py | 272 ++++++++++++++++++++++++++++++++++++++++++ web/bridge_preview.js | 212 ++++++++++++++++++++++++++++++++ 2 files changed, 484 insertions(+) create mode 100644 py/bridge_preview.py create mode 100644 web/bridge_preview.js diff --git a/py/bridge_preview.py b/py/bridge_preview.py new file mode 100644 index 0000000..00dbf9c --- /dev/null +++ b/py/bridge_preview.py @@ -0,0 +1,272 @@ +from .md import * +import torch +import json +import os +import random +import numpy as np +import hashlib +from PIL import Image +from PIL.PngImagePlugin import PngInfo +import folder_paths +import comfy.utils +from threading import Event +import threading +from server import PromptServer +from aiohttp import web +def get_bridge_storage(): + if not hasattr(PromptServer.instance, '_bridge_node_data'): + PromptServer.instance._bridge_node_data = {} + return PromptServer.instance._bridge_node_data +def get_bridge_cache(): + """获取桥接节点的缓存存储""" + if not hasattr(PromptServer.instance, '_bridge_node_cache'): + PromptServer.instance._bridge_node_cache = {} + return PromptServer.instance._bridge_node_cache +class BridgePreviewNode(PreviewImage): + """桥接预览节点，等待前端遮罩操作完成后输出图片""" + def __init__(self): + super().__init__() + self.prefix_append = "_bridge_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE", ), + "file_info": ("STRING", {"default": "", "readonly": True, "multiline": False}), + "skip": ("BOOLEAN", {"default": False, "label_on": "Skip", "label_off": "Open dialog"}), + }, + "hidden": { + "prompt": "PROMPT", + "extra_pnginfo": "EXTRA_PNGINFO", + "unique_id": "UNIQUE_ID" + }, + } + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("处理后图像", "遮罩") + FUNCTION = "process_image" + OUTPUT_NODE = True + CATEGORY = "🎈LAOGOU/Image" + def calculate_image_hash(self, images): + """计算图片的哈希值用于检测是否改变""" + try: + np_images = images.cpu().numpy() + image_bytes = np_images.tobytes() + return hashlib.md5(image_bytes).hexdigest() + except: + return None + def save_output_images(self, images, mask, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): + """保存输出图片（带遮罩的RGBA图片）用于前端显示""" + try: + if images.shape[0] != mask.shape[0]: + return None + batch_size = images.shape[0] + saved_images = [] + for i in range(batch_size): + single_image = images[i:i+1] + single_mask = mask[i] + np_image = (single_image.squeeze(0).cpu().numpy() * 255).astype(np.uint8) + np_mask = (single_mask.cpu().numpy() * 255).astype(np.uint8) + h, w, c = np_image.shape + rgba_image = np.zeros((h, w, 4), dtype=np.uint8) + rgba_image[:, :, :3] = np_image + rgba_image[:, :, 3] = np_mask + pil_image = Image.fromarray(rgba_image, 'RGBA') + counter = 1 + while True: + filename = f"{filename_prefix}_{counter:05d}_.png" + full_path = os.path.join(folder_paths.get_output_directory(), filename) + if not os.path.exists(full_path): + break + counter += 1 + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for key, value in extra_pnginfo.items(): + metadata.add_text(key, json.dumps(value)) + pil_image.save(full_path, pnginfo=metadata, compress_level=4) + saved_images.append({ + "filename": filename, + "subfolder": "", + "type": "output" + }) + return {"ui": {"images": saved_images}} + except Exception as e: + return None + def process_image(self, images, file_info="", skip=False, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, unique_id=None): + """处理图像，等待前端遮罩操作完成""" + try: + node_id = str(unique_id) + current_hash = self.calculate_image_hash(images) + cache = get_bridge_cache() + + # 如果启用跳过预览，直接返回缓存结果或原图 + if skip: + if node_id in cache and current_hash: + cached_data = cache[node_id] + if (cached_data.get("input_hash") == current_hash and + cached_data.get("final_result")): + # 返回缓存的最终结果 + cached_images, cached_mask = cached_data["final_result"] + return (cached_images, cached_mask) + + # 没有缓存或输入改变，返回原图和全白遮罩 + batch_size, height, width, channels = images.shape + default_mask = torch.ones((batch_size, height, width), dtype=torch.float32) + return (images, default_mask) + + # 正常预览遮罩编辑流程 + event = Event() + preview_urls = [] + should_send_to_frontend = True + + if node_id in cache and current_hash: + cached_data = cache[node_id] + if cached_data.get("input_hash") == current_hash and cached_data.get("output_urls"): + preview_urls = cached_data["output_urls"] + should_send_to_frontend = True + else: + preview_result = self.save_images(images, filename_prefix, prompt, extra_pnginfo) + preview_urls = preview_result["ui"]["images"] if preview_result and "ui" in preview_result else [] + else: + preview_result = self.save_images(images, filename_prefix, prompt, extra_pnginfo) + preview_urls = preview_result["ui"]["images"] if preview_result and "ui" in preview_result else [] + get_bridge_storage()[node_id] = { + "event": event, + "result": None, + "original_images": images, + "processing_complete": False, + "input_hash": current_hash + } + + try: + PromptServer.instance.send_sync("bridge_preview_update", { + "node_id": node_id, + "urls": preview_urls + }) + + if not event.wait(timeout=30): + if node_id in get_bridge_storage(): + del get_bridge_storage()[node_id] + batch_size, height, width, channels = images.shape + default_mask = torch.ones((batch_size, height, width), dtype=torch.float32) + return (images, default_mask) + + result_data = None + if node_id in get_bridge_storage(): + result_data = get_bridge_storage()[node_id]["result"] + del get_bridge_storage()[node_id] + + batch_size, height, width, channels = images.shape + default_mask = torch.ones((batch_size, height, width), dtype=torch.float32) + + if result_data is not None and isinstance(result_data, tuple) and len(result_data) == 2: + final_images, final_mask = result_data + if current_hash: + output_result = self.save_output_images(final_images, final_mask, filename_prefix + "_output", prompt, extra_pnginfo) + output_urls = output_result["ui"]["images"] if output_result and "ui" in output_result else [] + cache[node_id] = { + "input_hash": current_hash, + "output_urls": output_urls, + "final_result": (final_images, final_mask) + } + return (final_images, final_mask) + else: + return (images, default_mask) + + except Exception as e: + if node_id in get_bridge_storage(): + del get_bridge_storage()[node_id] + batch_size, height, width, channels = images.shape + default_mask = torch.ones((batch_size, height, width), dtype=torch.float32) + return (images, default_mask) + + except Exception as e: + batch_size, height, width, channels = images.shape + default_mask = torch.ones((batch_size, height, width), dtype=torch.float32) + return (images, default_mask) +@PromptServer.instance.routes.post("/bridge_preview/confirm") +async def confirm_bridge_preview(request): + """处理前端确认遮罩操作完成""" + try: + data = await request.json() + node_id = str(data.get("node_id")) + file_info = data.get("file_info") + if node_id not in get_bridge_storage(): + return web.json_response({"success": False, "error": "节点未找到或已超时"}) + try: + node_info = get_bridge_storage()[node_id] + if file_info: + processed_image, mask_image = load_processed_image(file_info) + if processed_image is not None and mask_image is not None: + node_info["result"] = (processed_image, mask_image) + node_info["processing_complete"] = True + node_info["event"].set() + return web.json_response({"success": True}) + except Exception as e: + if node_id in get_bridge_storage() and "event" in get_bridge_storage()[node_id]: + get_bridge_storage()[node_id]["event"].set() + return web.json_response({"success": False, "error": str(e)}) + except Exception as e: + return web.json_response({"success": False, "error": str(e)}) +@PromptServer.instance.routes.post("/bridge_preview/cancel") +async def cancel_bridge_preview(request): + """取消桥接预览操作""" + try: + data = await request.json() + node_id = str(data.get("node_id")) + if node_id in get_bridge_storage(): + get_bridge_storage()[node_id]["event"].set() + return web.json_response({"success": True, "message": f"节点 {node_id} 已取消"}) + else: + return web.json_response({"success": False, "error": f"节点 {node_id} 未找到或已超时"}) + except Exception as e: + return web.json_response({"success": False, "error": str(e)}) +def load_processed_image(file_info): + """从文件信息加载处理后的图片，返回图像和遮罩""" + try: + if isinstance(file_info, dict): + filename = file_info.get("filename") + subfolder = file_info.get("subfolder", "") + file_type = file_info.get("type", "output") + elif isinstance(file_info, str): + filename = file_info + subfolder = "" + file_type = "output" + else: + return None, None + if not filename: + return None, None + if file_type == "input": + base_dir = folder_paths.get_input_directory() + elif file_type == "output": + base_dir = folder_paths.get_output_directory() + elif file_type == "temp": + base_dir = folder_paths.get_temp_directory() + else: + base_dir = folder_paths.get_output_directory() + if subfolder: + file_path = os.path.join(base_dir, subfolder, filename) + else: + file_path = os.path.join(base_dir, filename) + if os.path.exists(file_path): + image = Image.open(file_path) + if image.mode != 'RGBA': + image = image.convert('RGBA') + np_image = np.array(image) + rgb_image = np_image[:, :, :3] + alpha_channel = np_image[:, :, 3] + tensor_image = torch.from_numpy(rgb_image / 255.0).float().unsqueeze(0) + mask_tensor = torch.from_numpy(alpha_channel / 255.0).float().unsqueeze(0) + return tensor_image, mask_tensor + else: + return None, None + except Exception as e: + pass + return None, None +NODE_CLASS_MAPPINGS = { + "BridgePreviewNode": BridgePreviewNode +} +NODE_DISPLAY_NAME_MAPPINGS = { + "BridgePreviewNode": "🎈LG_PreviewBridge_V2" +} \ No newline at end of file diff --git a/web/bridge_preview.js b/web/bridge_preview.js new file mode 100644 index 0000000..d058fb0 --- /dev/null +++ b/web/bridge_preview.js @@ -0,0 +1,212 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; +const waitingNodes = new Map(); +const processedNodes = new Set(); +function updateNodePreview(nodeId, fileInfo) { + const node = app.graph.getNodeById(parseInt(nodeId)); + if (!node || !fileInfo) return; + const imageUrl = fileInfo.subfolder + ? `/view?filename=${encodeURIComponent(fileInfo.filename)}&subfolder=${encodeURIComponent(fileInfo.subfolder)}&type=${fileInfo.type || 'input'}` + : `/view?filename=${encodeURIComponent(fileInfo.filename)}&type=${fileInfo.type || 'input'}`; + node.images = [{ + filename: fileInfo.filename, + subfolder: fileInfo.subfolder || "", + type: fileInfo.type || "input", + url: imageUrl + }]; + if (node.onDrawBackground) { + const img = new Image(); + img.onload = () => { + node.imgs = [img]; + app.graph.setDirtyCanvas(true); + }; + img.src = imageUrl; + } + updateStatusWidget(node, `${nodeId}_${fileInfo.filename}`); +} +function updateStatusWidget(node, statusText) { + const statusWidget = node.widgets?.find(w => w.name === "file_info"); + if (statusWidget) { + statusWidget.value = statusText; + app.graph.setDirtyCanvas(true); + app.graph.change(); + console.log(`[BridgePreview] 节点 ${node.id} 状态已更新: ${statusText}`); + } +} +api.addEventListener("bridge_preview_update", (event) => { + const { node_id, urls } = event.detail; + const node = app.graph._nodes_by_id[node_id]; + if (!node || !urls?.length) return; + waitingNodes.set(node_id, { urls, timestamp: Date.now(), node }); + const imageData = urls.map((url, index) => ({ + index, + filename: url.filename, + subfolder: url.subfolder, + type: url.type + })); + node.imageData = imageData; + node.imgs = []; + imageData.forEach((imgData, i) => { + const img = new Image(); + img.onload = () => { + app.graph.setDirtyCanvas(true); + if (i === imageData.length - 1) { + setupClipspace(node_id, urls); + } + }; + img.src = `/view?filename=${encodeURIComponent(imgData.filename)}&type=${imgData.type}&subfolder=${imgData.subfolder || ''}&${app.getPreviewFormatParam()}`; + node.imgs.push(img); + }); + node.setSizeForImage?.(); + node.update?.(); +}); +function setupClipspace(nodeId, urls) { + const ComfyApp = app.constructor; + if (!ComfyApp.clipspace) ComfyApp.clipspace = {}; + if (!app.clipspace) app.clipspace = {}; + const images = urls.map(url => ({ + filename: url.filename, + subfolder: url.subfolder || "", + type: url.type || "output" + })); + const imgs = urls.map(url => ({ + src: `${window.location.origin}/view?filename=${encodeURIComponent(url.filename)}&type=${url.type}&subfolder=${encodeURIComponent(url.subfolder || '')}`, + filename: url.filename, + subfolder: url.subfolder || "", + type: url.type || "output" + })); + [ComfyApp.clipspace, app.clipspace].forEach(clipspace => { + clipspace.images = images; + clipspace.imgs = imgs; + clipspace.selectedIndex = 0; + }); + setTimeout(() => { + const openMaskEditor = ComfyApp.open_maskeditor || app.open_maskeditor; + openMaskEditor?.(); + bindCancelButton(); + }, 100); +} +function bindCancelButton() { + const checkInterval = setInterval(() => { + const maskEditor = findMaskEditor(); + if (!maskEditor) return; + const cancelButtons = Array.from(maskEditor.querySelectorAll('button')).filter(btn => { + const text = btn.textContent.trim().toLowerCase(); + return text === 'cancel' || text === '取消'; + }); + if (cancelButtons.length > 0) { + cancelButtons.forEach(button => { + if (!button.hasAttribute('data-bridge-bound')) { + button.setAttribute('data-bridge-bound', 'true'); + button.addEventListener('click', () => { + setTimeout(handleMaskEditorCancel, 50); + }); + } + }); + clearInterval(checkInterval); + } + }, 300); + setTimeout(() => clearInterval(checkInterval), 10000); +} +function findMaskEditor() { + const modals = document.querySelectorAll('div.comfy-modal, .comfy-modal, [class*="modal"]'); + for (const modal of modals) { + if (modal.querySelector('canvas') && modal.style.display !== 'none') { + return modal; + } + } + const elements = document.querySelectorAll('*'); + for (const element of elements) { + const buttons = element.querySelectorAll('button'); + if (buttons.length >= 2 && element.querySelector('canvas')) { + const buttonTexts = Array.from(buttons).map(btn => btn.textContent.trim().toLowerCase()); + if (buttonTexts.some(text => text.includes('cancel') || text.includes('取消')) && + buttonTexts.some(text => text.includes('save') || text.includes('保存'))) { + return element; + } + } + } + return null; +} +function handleMaskEditorCancel() { + waitingNodes.forEach((nodeInfo, nodeId) => { + if (!processedNodes.has(String(nodeId))) { + sendCancelSignal(nodeId); + } + }); + waitingNodes.clear(); + processedNodes.clear(); +} +async function sendCancelSignal(nodeId) { + try { + const response = await api.fetchApi("/bridge_preview/cancel", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ node_id: String(nodeId) }) + }); + if (!response.ok) { + console.error(`[BridgePreview] 取消节点 ${nodeId} 失败:`, response.status); + } + } catch (error) { + console.error(`[BridgePreview] 发送取消信号失败:`, error); + } +} +const originalFetch = api.fetchApi; +api.fetchApi = async function(url, options) { + const result = await originalFetch.call(this, url, options); + if (url === "/upload/mask" && result.ok) { + await handleMaskUpload(result); + } + return result; +}; +async function handleMaskUpload(result) { + try { + const responseData = await result.clone().json(); + const fileInfo = responseData?.name ? { + filename: responseData.name, + subfolder: responseData.subfolder || "clipspace", + type: responseData.type || "input" + } : null; + if (!fileInfo) return; + let latestNodeId = null; + let latestTimestamp = 0; + for (const [nodeId, nodeInfo] of waitingNodes) { + if (nodeInfo.timestamp > latestTimestamp) { + latestTimestamp = nodeInfo.timestamp; + latestNodeId = nodeId; + } + } + if (!latestNodeId) return; + try { + updateNodePreview(latestNodeId, fileInfo); + const confirmResponse = await originalFetch.call(api, "/bridge_preview/confirm", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + node_id: String(latestNodeId), + file_info: fileInfo + }) + }); + if (confirmResponse.ok) { + processedNodes.add(String(latestNodeId)); + setTimeout(() => { + waitingNodes.delete(latestNodeId); + }, 1000); + } + } catch (error) { + console.error(`[BridgePreview] 处理节点 ${latestNodeId} 失败:`, error); + } + } catch (error) { + console.error(`[BridgePreview] 处理mask结果失败:`, error); + } +} +setInterval(() => { + const now = Date.now(); + for (const [nodeId, nodeInfo] of waitingNodes) { + if (now - nodeInfo.timestamp > 30000) { + waitingNodes.delete(nodeId); + processedNodes.delete(String(nodeId)); + } + } +}, 5000); +console.log("[BridgePreview] 桥接预览模块已加载"); \ No newline at end of file