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": "保存图片路径" +}