From 36ca672cfb20b05bf06ee04328fd5f433b0a751c Mon Sep 17 00:00:00 2001
From: yp <yp@bowong.ai>
Date: Wed, 11 Jun 2025 18:17:53 +0800
Subject: [PATCH] =?UTF-8?q?fix=20=E4=BF=9D=E5=AD=98=E5=9B=BE=E7=89=87?=
 =?UTF-8?q?=E8=87=AA=E5=AE=9A=E4=B9=89=E8=8A=82=E7=82=B9?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 nodes/image.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/nodes/image.py b/nodes/image.py
index 08bd9bb..bdd79d5 100644
--- a/nodes/image.py
+++ b/nodes/image.py
@@ -2,8 +2,7 @@ import os
 import uuid
 from PIL import Image
 import numpy as np
-import torch  # 添加这一行，用于类型检查
-
+import torch
 
 class SaveImagePath:
     @classmethod
@@ -27,6 +26,10 @@ class SaveImagePath:
         if image_path.dtype != np.uint8:
             image_path = np.clip(image_path, 0, 255).astype(np.uint8)
 
+        # 去除多余的维度，如果形状是(1, 1, height, width, channels)或(1, height, width, channels)等情况
+        while len(image_path.shape) > 3:
+            image_path = image_path.squeeze(0)
+
         # 如果是单通道图像，转换为3通道
         if len(image_path.shape) == 2:
             image_path = np.stack([image_path] * 3, axis=-1)
@@ -46,3 +49,10 @@ class SaveImagePath:
         pil_image.save(p)
 
         return (p,)
+# 节点类定义结束，以下是用于注册节点的字典结构（通常在实际使用中由ComfyUI等框架来解析和注册）
+NODE_CLASS_MAPPINGS = {
+    "SaveImagePath": SaveImagePath
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SaveImagePath": "保存图片路径"
+}