From 369cb3c5cee597b3d9932a5b6433933babcf44e4 Mon Sep 17 00:00:00 2001 From: jueming Date: Thu, 8 May 2025 14:50:48 +0800 Subject: [PATCH] fix image preview --- nodes/image_downloader.py | 147 +++++++++++++++++++----- web/js/image_downloader_preview.js | 177 +++++++++++------------------ 2 files changed, 185 insertions(+), 139 deletions(-) diff --git a/nodes/image_downloader.py b/nodes/image_downloader.py index b10ef2c..9c8fbf6 100644 --- a/nodes/image_downloader.py +++ b/nodes/image_downloader.py @@ -17,12 +17,15 @@ class KLingAIImageDownloader: """ KLingAI Image Downloader Node Downloads image from URL and saves to local directory + 具有图片预览功能 """ def __init__(self): self.default_output_dir = folder_paths.get_output_directory() self.type = "image" - self.preview_image = None + self.output_dir = self.default_output_dir + self.prefix = "KLingAI" + self.last_downloaded_image = None @classmethod def INPUT_TYPES(s): @@ -34,7 +37,7 @@ class KLingAIImageDownloader: "multiline": False, "placeholder": "Base filename for downloaded image" }), - "preview_image": ("STRING", {"default": "#NULL", "multiline": False}), + "image": (["#DATA"], {"image_upload": True}), # 添加支持上传图片的控件 "image_data": ("STRING", {"default": "", "multiline": False}), }, "optional": { @@ -43,19 +46,15 @@ class KLingAIImageDownloader: "multiline": False, "placeholder": "Optional: Custom output directory path" }), - }, - "hidden": { - "image_data_hidden": ("STRING", {"default": "", "multiline": False}), - "node_id": ("STRING", {"default": ""}), } } - RETURN_TYPES = ("STRING", "STRING", "IMAGE", "STRING") - RETURN_NAMES = ("image_path", "image_url", "image", "image_data") - FUNCTION = "download_image" + RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING") + RETURN_NAMES = ("IMAGE", "MASK", "image_path", "image_url") + FUNCTION = "process_image" CATEGORY = "JM-KLingAI-API" - OUTPUT_NODE = True # 标记为可作为终端节点的节点 - OUTPUT_IS_LIST = (False, False, False, False) # 指示输出不是列表 + OUTPUT_NODE = True + OUTPUT_IS_LIST = (False, False, False, False) def get_next_sequence_number(self, directory, filename_prefix): """ @@ -91,7 +90,7 @@ class KLingAIImageDownloader: """ Path(directory).mkdir(parents=True, exist_ok=True) return directory - + def image_to_base64(self, pil_image): """ 将PIL图像转换为base64字符串 @@ -100,11 +99,11 @@ class KLingAIImageDownloader: return "" buffered = BytesIO() - pil_image.save(buffered, format="JPEG") + pil_image.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") - return f"data:image/jpeg;base64,{img_str}" - - def download_image(self, image_url, filename_prefix="KLingAI", preview_image="#NULL", image_data="", custom_output_dir="", image_data_hidden="", node_id=""): + return f"data:image/png;base64,{img_str}" + + def download_image(self, image_url, filename_prefix="KLingAI", custom_output_dir=""): """ Download image from URL and save to local directory """ @@ -113,7 +112,9 @@ class KLingAIImageDownloader: if image_url is None or (isinstance(image_url, str) and not image_url.strip()): error_msg = "图片URL不能为空" print(error_msg) - return (error_msg, image_url, None, "") + # 创建小的空白图像用于显示错误信息 + empty_tensor = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + return (empty_tensor, error_msg, image_url) if not filename_prefix: filename_prefix = "KLingAI" @@ -161,27 +162,117 @@ class KLingAIImageDownloader: # 转换为PyTorch张量 image_tensor = torch.from_numpy(image_array)[None,] - # 生成base64编码的图像数据用于预览 - image_data_base64 = self.image_to_base64(pil_image) + # 存储最后下载的图像以便在UI中显示 + self.last_downloaded_image = pil_image print(f"成功加载图片，尺寸: {pil_image.width}x{pil_image.height}") - return (filepath, image_url, image_tensor, image_data_base64) + return (image_tensor, filepath, image_url) except Exception as e: print(f"图片加载错误: {str(e)}") - return (filepath, image_url, None, "") + # 创建小的空白图像 + empty_tensor = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + return (empty_tensor, filepath, image_url) except ValueError as ve: error_msg = f"参数错误: {str(ve)}" print(error_msg) - return (error_msg, image_url, None, "") # 返回错误消息而不是None + # 创建小的空白图像 + empty_tensor = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + return (empty_tensor, error_msg, image_url) except requests.exceptions.RequestException as re: error_msg = f"下载错误: {str(re)}" print(error_msg) - return (error_msg, image_url, None, "") # 返回错误消息而不是None + # 创建小的空白图像 + empty_tensor = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + return (empty_tensor, error_msg, image_url) except Exception as e: error_msg = f"图片下载错误: {str(e)}" print(error_msg) - return (error_msg, image_url, None, "") # 返回错误消息而不是None + # 创建小的空白图像 + empty_tensor = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + return (empty_tensor, error_msg, image_url) + + def process_image(self, image_url, filename_prefix="KLingAI", image="#DATA", image_data="", custom_output_dir=""): + """ + 处理图片：如果提供了图片数据，直接处理；否则尝试下载URL中的图片 + """ + try: + # 设置输出目录 + output_dir = custom_output_dir if custom_output_dir else self.default_output_dir + output_dir = self.ensure_directory(output_dir) + self.output_dir = output_dir + self.prefix = filename_prefix + + # 默认掩码（全透明） + mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu").unsqueeze(0) + + # 如果提供了图片数据（通过上传） + if image_data and image_data != "" and image_data.startswith("data:"): + try: + # 解码base64图片数据 + image_data_binary = base64.b64decode(image_data.split(",")[1]) + i = Image.open(BytesIO(image_data_binary)) + + # 转换图像格式 + img_rgb = i.convert("RGB") + image_array = np.array(img_rgb).astype(np.float32) / 255.0 + image_tensor = torch.from_numpy(image_array)[None,] + + # 检查是否有alpha通道，创建掩码 + if 'A' in i.getbands(): + alpha = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(alpha).unsqueeze(0) + + # 保存上传的图片 + seq_num = self.get_next_sequence_number(output_dir, filename_prefix) + filepath = os.path.join(output_dir, f"{filename_prefix}_{seq_num:04d}.png") + img_rgb.save(filepath, "PNG") + print(f"上传的图片已保存至: {filepath}") + self.save_image_preview_info(filepath) + + return (image_tensor, mask, filepath, "本地上传") + + except Exception as e: + print(f"处理上传的图片时出错: {str(e)}") + # 如果处理上传失败但有URL，则继续尝试下载 + if not image_url or not image_url.strip(): + raise ValueError(f"图片处理失败: {str(e)}") + + # 如果没有上传图片或处理失败，尝试从URL下载 + if image_url and image_url.strip(): + try: + filepath = self.download_image(image_url, output_dir, filename_prefix) + + # 加载下载的图片 + pil_image = Image.open(filepath).convert("RGB") + image_array = np.array(pil_image).astype(np.float32) / 255.0 + image_tensor = torch.from_numpy(image_array)[None,] + + # 转换为base64用于预览 + image_data_base64 = self.image_to_base64(pil_image) + + print(f"成功加载图片，尺寸: {pil_image.width}x{pil_image.height}") + + return (image_tensor, mask, filepath, image_url) + + except Exception as e: + print(f"从URL下载图片失败: {str(e)}") + if not image_data or image_data == "" or not image_data.startswith("data:"): + raise ValueError(f"无法下载图片: {str(e)}") + + # 如果没有提供有效的URL或图片数据 + if (not image_url or not image_url.strip()) and (not image_data or not image_data.startswith("data:")): + raise ValueError("请提供有效的图片URL或上传图片") + + # 如果代码执行到这里，可能是前面的某个步骤成功但没有正确返回 + raise ValueError("未知错误，无法处理图片") + + except Exception as e: + print(f"图片处理错误: {str(e)}") + # 创建空白图像和掩码 + empty_tensor = torch.zeros((1, 64, 64, 3), dtype=torch.float32) + empty_mask = torch.zeros((1, 64, 64), dtype=torch.float32) + return (empty_tensor, empty_mask, f"错误: {str(e)}", image_url) def save_image_preview_info(self, filepath): """ @@ -228,6 +319,10 @@ class KLingAIImageDownloader: # 确保节点可以在没有连接的情况下运行 @classmethod - def IS_CHANGED(cls, image_url, filename_prefix="KLingAI", preview_image="#NULL", image_data="", custom_output_dir="", image_data_hidden="", node_id=""): - # 返回当前时间，确保节点总是执行 + def IS_CHANGED(cls, image_url, filename_prefix="KLingAI", image="#DATA", image_data="", custom_output_dir=""): + # 如果有图片数据上传，返回图片数据的哈希值 + if image_data and image_data.startswith("data:"): + import hashlib + return hashlib.md5(image_data.encode()).hexdigest() + # 否则使用当前时间，确保节点总是执行 return time.time() \ No newline at end of file diff --git a/web/js/image_downloader_preview.js b/web/js/image_downloader_preview.js index f711f3e..534c467 100644 --- a/web/js/image_downloader_preview.js +++ b/web/js/image_downloader_preview.js @@ -6,135 +6,86 @@ app.registerExtension({ nodeCreated(node, app) { // 检查节点类型是否为KLingAIImageDownloader if (node.comfyClass === "KLingAIImageDownloader") { + // 获取节点原始的大小 + const origSize = node.size || [200, 100]; + // 设置节点的大小，增加足够的空间来显示预览 + node.size = [Math.max(origSize[0], 300), Math.max(origSize[1], 400)]; - // 把节点的image_data隐藏起来 - for (const w of node.widgets) { - if (w.name === "preview_image" || w.name === "image_data") { - w.hide = true; + // 保存原始的onDrawBackground方法 + const origDrawBackground = node.onDrawBackground; + + // 重写onDrawBackground，增加图片预览功能 + node.onDrawBackground = function(ctx) { + // 调用原始的绘制方法 + if (origDrawBackground) { + origDrawBackground.call(this, ctx); } - } - - // 获取节点ID并设置到隐藏参数 - const node_id = node.id; - const node_id_widget = node.widgets.find(w => w.name === "node_id"); - if (node_id_widget) { - node_id_widget.value = node_id; - } - - // 找到image_data widget - const image_data_widget = node.widgets.find(w => w.name === "image_data"); - - if (!image_data_widget) return; - - // 设置节点尺寸以便有足够空间显示预览 - node.setSize([300, 400]); - - // 添加预览控件 - node.addCustomWidget({ - name: "preview_widget", - type: "image", - draw(ctx, node, widget_width, widget_height) { - // 获取节点的标题栏高度和控件所占区域高度 - const headerHeight = 30; // 标题栏近似高度 - const widgetsHeight = node.widgets.filter(w => !w.hide).length * 32; // 假设每个可见控件高32像素 + // 检查是否有图像需要预览 + if (this.imgs && this.imgs.length > 0) { + const img = this.imgs[0]; + if (!img || !img.width) return; - // 计算预览区域的大小和位置 - const previewX = 10; - const previewY = headerHeight + widgetsHeight + 5; - const maxWidth = node.size[0] - 20; - const maxHeight = node.size[1] - previewY - 10; + // 计算预览区域 + const headerHeight = 30; // 标题栏高度 + const widgetsHeight = this.widgets.length * 28; // 控件区域高度 + const margin = 10; - // 如果没有图像数据，就不显示预览 - if (!image_data_widget.value) { - // 显示提示文字 - ctx.fillStyle = "#888"; - ctx.font = "13px Arial"; - ctx.textAlign = "center"; - ctx.fillText("下载图片后在此处显示预览", - node.size[0] / 2, - previewY + maxHeight / 2); - return; - } + // 计算可用于预览的区域 + const previewAreaX = margin; + const previewAreaY = headerHeight + widgetsHeight + margin; + const previewAreaWidth = this.size[0] - margin * 2; + const previewAreaHeight = this.size[1] - previewAreaY - margin; - // 尝试创建图像对象 - if (!node._preview_img && image_data_widget.value) { - node._preview_img = new Image(); - node._preview_img.src = image_data_widget.value; - node._preview_img.onload = () => { - app.canvas.setDirty(true); - }; - } + // 计算图像显示大小，保持比例 + const imgRatio = img.width / img.height; + let drawWidth, drawHeight; - // 如果图像已加载，绘制它 - if (node._preview_img && node._preview_img.complete) { - const img = node._preview_img; - const imgRatio = img.width / img.height; - - let w, h; - if (imgRatio > maxWidth / maxHeight) { - // 图像较宽，以宽度为限制 - w = maxWidth; - h = w / imgRatio; - } else { - // 图像较高，以高度为限制 - h = maxHeight; - w = h * imgRatio; - } - - // 计算居中位置 - const x = previewX + (maxWidth - w) / 2; - const y = previewY + (maxHeight - h) / 2; - - // 绘制图像 - ctx.drawImage(img, x, y, w, h); - } - }, - - // 当widget值改变时 - async value_changed(value, node) { - if (value) { - node._preview_img = new Image(); - node._preview_img.src = value; - app.canvas.setDirty(true); - } - } - }); - - // 监听图像数据变化 - image_data_widget.callback = function(value) { - if (value && value !== node._last_image_data) { - node._last_image_data = value; - - if (node._preview_img) { - node._preview_img.src = value; + if (imgRatio > previewAreaWidth / previewAreaHeight) { + // 图像较宽，以宽度为限 + drawWidth = previewAreaWidth; + drawHeight = drawWidth / imgRatio; } else { - node._preview_img = new Image(); - node._preview_img.src = value; + // 图像较高，以高度为限 + drawHeight = previewAreaHeight; + drawWidth = drawHeight * imgRatio; } - app.canvas.setDirty(true); + // 计算居中位置 + const drawX = previewAreaX + (previewAreaWidth - drawWidth) / 2; + const drawY = previewAreaY + (previewAreaHeight - drawHeight) / 2; + + // 绘制预览图像 + ctx.drawImage(img, drawX, drawY, drawWidth, drawHeight); + + // 绘制图像尺寸信息 + ctx.fillStyle = "#CCC"; + ctx.font = "12px Arial"; + ctx.fillText(`${img.naturalWidth}x${img.naturalHeight}`, + drawX + drawWidth - 70, drawY + drawHeight - 10); + } else { + // 没有图像时显示提示文本 + ctx.fillStyle = "#999"; + ctx.font = "14px Arial"; + ctx.textAlign = "center"; + ctx.fillText("下载图片后将在此处显示预览", + this.size[0] / 2, + this.size[1] / 2); + ctx.textAlign = "left"; } - return value; }; - // 因为自定义控件是在节点创建后添加的，我们需要手动更新节点的大小 - // 这样能确保有足够的空间显示预览 - const orig_onExecuted = node.onExecuted; + // 重写onExecuted方法，在节点执行后更新图像 + const origOnExecuted = node.onExecuted; node.onExecuted = function(message) { - if (orig_onExecuted) { - orig_onExecuted.apply(this, arguments); + if (origOnExecuted) { + origOnExecuted.call(this, message); } - // 检查输出中是否有image_data - if (message && message.outputs && message.outputs.image_data) { - const image_data = message.outputs.image_data; - if (image_data && image_data_widget) { - image_data_widget.value = image_data; - - // 强制重绘 - app.canvas.setDirty(true); - } + // 检查是否有图像输出 + if (message && message.outputs && message.outputs.image) { + // 触发预览图像的更新 + app.canvas.setDirty(true); } }; }