From 1fc69b650bd5f1cf12ea0d9ff5ca60baf350dcee Mon Sep 17 00:00:00 2001 From: jueming Date: Thu, 8 May 2025 15:02:29 +0800 Subject: [PATCH] fix image download --- nodes/image_downloader.py | 138 +++++++---------------------- web/js/image_downloader_preview.js | 72 ++++++++++----- 2 files changed, 85 insertions(+), 125 deletions(-) diff --git a/nodes/image_downloader.py b/nodes/image_downloader.py index 9c8fbf6..37e8e8e 100644 --- a/nodes/image_downloader.py +++ b/nodes/image_downloader.py @@ -21,6 +21,7 @@ class KLingAIImageDownloader: """ def __init__(self): + # 确保使用ComfyUI的output目录 self.default_output_dir = folder_paths.get_output_directory() self.type = "image" self.output_dir = self.default_output_dir @@ -37,28 +38,27 @@ class KLingAIImageDownloader: "multiline": False, "placeholder": "Base filename for downloaded image" }), - "image": (["#DATA"], {"image_upload": True}), # 添加支持上传图片的控件 - "image_data": ("STRING", {"default": "", "multiline": False}), }, "optional": { "custom_output_dir": ("STRING", { "default": "", "multiline": False, - "placeholder": "Optional: Custom output directory path" + "placeholder": "Optional: Custom output directory path within outputs folder" }), } } - RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING") - RETURN_NAMES = ("IMAGE", "MASK", "image_path", "image_url") - FUNCTION = "process_image" + RETURN_TYPES = ("IMAGE", "STRING", "STRING") + RETURN_NAMES = ("IMAGE", "image_path", "image_url") + FUNCTION = "download_image" CATEGORY = "JM-KLingAI-API" - OUTPUT_NODE = True - OUTPUT_IS_LIST = (False, False, False, False) + OUTPUT_NODE = True # 标记为可作为终端节点的节点 + OUTPUT_IS_LIST = (False, False, False) # 指示输出不是列表 def get_next_sequence_number(self, directory, filename_prefix): """ Get the next sequence number for the image file + 始终从0001开始，如果文件已存在则自动递增 """ # 支持常见的图片格式 patterns = [ @@ -73,11 +73,11 @@ class KLingAIImageDownloader: existing_files.extend(glob.glob(os.path.join(directory, pattern))) if not existing_files: - return 1 + return 1 # 始终从1开始，对应_0001格式 numbers = [] for f in existing_files: - # 匹配不同格式的文件名 + # 匹配不同格式的文件名，提取序号部分 match = re.search(f"{filename_prefix}_(\d+)\.(png|jpg|jpeg|webp)$", f) if match: numbers.append(int(match.group(1))) @@ -87,6 +87,7 @@ class KLingAIImageDownloader: def ensure_directory(self, directory): """ Ensure the directory exists, create if it doesn't + 对于自定义目录，确保是在output目录下的子目录 """ Path(directory).mkdir(parents=True, exist_ok=True) return directory @@ -119,11 +120,26 @@ class KLingAIImageDownloader: if not filename_prefix: filename_prefix = "KLingAI" - # Determine output directory - output_dir = custom_output_dir if custom_output_dir else self.default_output_dir + # 确定输出目录，如果指定了自定义目录，确保是output目录的子目录 + output_dir = self.default_output_dir + if custom_output_dir: + # 确保自定义路径是一个相对路径，或将绝对路径转为相对output目录的路径 + if os.path.isabs(custom_output_dir): + # 如果是绝对路径，尝试将其转为相对于output目录的路径，或直接使用output目录 + rel_path = os.path.relpath(custom_output_dir, start=self.default_output_dir) + if rel_path.startswith('..'): # 路径在output目录之外 + print(f"警告: 自定义路径 {custom_output_dir} 不在output目录内，将使用默认output目录") + else: + output_dir = os.path.join(self.default_output_dir, rel_path) + else: + # 相对路径直接拼接到output目录 + output_dir = os.path.join(self.default_output_dir, custom_output_dir) + + # 确保目录存在 output_dir = self.ensure_directory(output_dir) + print(f"图片将保存至: {output_dir}") - # Get next sequence number + # 获取下一个序列号 seq_num = self.get_next_sequence_number(output_dir, filename_prefix) # 检测图片格式 @@ -133,16 +149,16 @@ class KLingAIImageDownloader: elif ".webp" in image_url.lower(): image_format = "webp" - # Create filename + # 创建文件名，格式为 prefix_0001.png filename = f"{filename_prefix}_{seq_num:04d}.{image_format}" filepath = os.path.join(output_dir, filename) - # Download image + # 下载图片 print(f"正在从 {image_url} 下载图片") response = requests.get(image_url, stream=True) response.raise_for_status() - # Save image + # 保存图片 with open(filepath, 'wb') as f: for chunk in response.iter_content(chunk_size=8192): if chunk: @@ -192,88 +208,6 @@ class KLingAIImageDownloader: 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): """ 保存图片预览信息以确保在历史记录中显示 @@ -319,10 +253,6 @@ class KLingAIImageDownloader: # 确保节点可以在没有连接的情况下运行 @classmethod - 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() - # 否则使用当前时间，确保节点总是执行 + def IS_CHANGED(cls, image_url, filename_prefix="KLingAI", custom_output_dir=""): + # 返回当前时间，确保节点总是执行 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 534c467..648de24 100644 --- a/web/js/image_downloader_preview.js +++ b/web/js/image_downloader_preview.js @@ -9,7 +9,10 @@ app.registerExtension({ // 获取节点原始的大小 const origSize = node.size || [200, 100]; // 设置节点的大小，增加足够的空间来显示预览 - node.size = [Math.max(origSize[0], 300), Math.max(origSize[1], 400)]; + node.size = [Math.max(origSize[0], 320), Math.max(origSize[1], 450)]; + + // 创建一个图像对象用于预览 + node._previewImage = null; // 保存原始的onDrawBackground方法 const origDrawBackground = node.onDrawBackground; @@ -21,21 +24,24 @@ app.registerExtension({ origDrawBackground.call(this, ctx); } + // 计算预览区域 + const headerHeight = 30; // 标题栏高度 + const widgetsHeight = this.widgets.length * 28; // 控件区域高度 + const margin = 10; + + // 计算可用于预览的区域 + const previewAreaX = margin; + const previewAreaY = headerHeight + widgetsHeight + margin; + const previewAreaWidth = this.size[0] - margin * 2; + const previewAreaHeight = this.size[1] - previewAreaY - margin; + + // 绘制预览区域的边框 + ctx.fillStyle = "#2a2a2a"; + ctx.fillRect(previewAreaX, previewAreaY, previewAreaWidth, previewAreaHeight); + // 检查是否有图像需要预览 - if (this.imgs && this.imgs.length > 0) { - const img = this.imgs[0]; - if (!img || !img.width) return; - - // 计算预览区域 - const headerHeight = 30; // 标题栏高度 - const widgetsHeight = this.widgets.length * 28; // 控件区域高度 - const margin = 10; - - // 计算可用于预览的区域 - const previewAreaX = margin; - const previewAreaY = headerHeight + widgetsHeight + margin; - const previewAreaWidth = this.size[0] - margin * 2; - const previewAreaHeight = this.size[1] - previewAreaY - margin; + if (this._previewImage && this._previewImage.width) { + const img = this._previewImage; // 计算图像显示大小，保持比例 const imgRatio = img.width / img.height; @@ -70,22 +76,46 @@ app.registerExtension({ ctx.textAlign = "center"; ctx.fillText("下载图片后将在此处显示预览", this.size[0] / 2, - this.size[1] / 2); + this.size[1] / 2 + 20); ctx.textAlign = "left"; } }; - // 重写onExecuted方法，在节点执行后更新图像 + // 重写onExecuted方法，在节点执行后加载和显示图像 const origOnExecuted = node.onExecuted; node.onExecuted = function(message) { if (origOnExecuted) { origOnExecuted.call(this, message); } - // 检查是否有图像输出 - if (message && message.outputs && message.outputs.image) { - // 触发预览图像的更新 - app.canvas.setDirty(true); + // 检查是否有图像输出和路径输出 + if (message && message.outputs) { + // 获取图像路径 + const imagePath = message.outputs.image_path; + if (imagePath && typeof imagePath === "string" && !imagePath.startsWith("错误:")) { + // 创建图像预览URL + let previewUrl; + if (imagePath.startsWith("/")) { + // 绝对路径 + previewUrl = `/view?filename=${encodeURIComponent(imagePath)}`; + } else { + // 相对路径 + previewUrl = `/view?filename=${encodeURIComponent(imagePath)}`; + } + + // 加载预览图像 + if (this._previewImage) { + this._previewImage.src = previewUrl; + } else { + this._previewImage = new Image(); + this._previewImage.onload = () => { + app.canvas.setDirty(true); + }; + this._previewImage.src = previewUrl; + } + + console.log(`正在加载预览图片: ${previewUrl}`); + } } }; }