import os
import uuid
from PIL import Image
import numpy as np
import torch

class SaveImagePath:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "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)

        # 去除多余的维度，如果形状是(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)

        # 如果是通道优先格式 (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,)
# 节点类定义结束，以下是用于注册节点的字典结构（通常在实际使用中由ComfyUI等框架来解析和注册）
NODE_CLASS_MAPPINGS = {
    "SaveImagePath": SaveImagePath
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "SaveImagePath": "保存图片路径"
}