修改保存图片节点

nodes/image.py

@@ -2,6 +2,8 @@ import os
 import uuid
 from PIL import Image
 import numpy as np
+import torch  # 添加这一行，用于类型检查
+
 class SaveImagePath:
     @classmethod
     def INPUT_TYPES(s):
@@ -10,24 +12,36 @@ class SaveImagePath:
                 "image_path": ("IMAGE", {"forceInput": True}),
             }
         }
+
     RETURN_TYPES = ("STRING",)
     FUNCTION = "load"
     CATEGORY = "不忘科技-自定义节点🚩"
+
     def load(self, image_path):
+        # 如果是torch.Tensor类型，转换为numpy数组
+        if isinstance(image_path, torch.Tensor):
+            image_path = image_path.cpu().numpy()
+
         # 确保数据类型为uint8且在0 - 255范围内
         if image_path.dtype != np.uint8:
             image_path = np.clip(image_path, 0, 255).astype(np.uint8)
+
         # 如果是单通道图像，转换为3通道
         if len(image_path.shape) == 2:
             image_path = np.stack([image_path] * 3, axis=-1)
+
         # 如果是通道优先格式 (C, H, W)，转换为通道最后格式 (H, W, C)
         elif len(image_path.shape) == 3 and image_path.shape[0] <= 4:
             image_path = np.transpose(image_path, (1, 2, 0))
+
         pil_image = Image.fromarray(image_path)
+
         output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
         if not os.path.exists(output_dir):
             os.makedirs(output_dir)
+
         file_name = "%s.jpg" % str(uuid.uuid4())
         p = os.path.join(output_dir, file_name)
         pil_image.save(p)
+
         return (p,)