diff --git a/nodes/image.py b/nodes/image.py index bdd79d5..a681366 100644 --- a/nodes/image.py +++ b/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": "保存图片路径" -} +} \ No newline at end of file