保存图片自定义节点

nodes/image.py

@@ -22,21 +22,30 @@ class SaveImagePath:
         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)
-
         # 去除多余的维度，如果形状是(1, 1, height, width, channels)或(1, height, width, channels)等情况
         while len(image_path.shape) > 3:
             image_path = image_path.squeeze(0)
 
+        # 如果是通道优先格式 (C, H, W)，转换为通道最后格式 (H, W, C)
+        if len(image_path.shape) == 3 and image_path.shape[0] <= 4:
+            image_path = np.transpose(image_path, (1, 2, 0))
+
         # 如果是单通道图像，转换为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))
+        # 数据范围和类型转换 - 这是关键修复
+        if image_path.dtype == np.float32 or image_path.dtype == np.float64:
+            # ComfyUI图像数据通常是0-1范围的浮点数
+            if image_path.max() <= 1.0:
+                # 从0-1范围转换到0-255范围
+                image_path = (image_path * 255.0).astype(np.uint8)
+            else:
+                # 如果已经是0-255范围，直接转换类型
+                image_path = np.clip(image_path, 0, 255).astype(np.uint8)
+        elif image_path.dtype != np.uint8:
+            # 其他数据类型，确保在0-255范围内
+            image_path = np.clip(image_path, 0, 255).astype(np.uint8)
 
         pil_image = Image.fromarray(image_path)
 
@@ -49,10 +58,11 @@ class SaveImagePath:
         pil_image.save(p)
 
         return (p,)
+
 # 节点类定义结束，以下是用于注册节点的字典结构（通常在实际使用中由ComfyUI等框架来解析和注册）
 NODE_CLASS_MAPPINGS = {
     "SaveImagePath": SaveImagePath
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "SaveImagePath": "保存图片路径"
-}
+}