import os.path
import uuid

import torch
import torchvision


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):
        image_path = image_path.float()
        # 假设数据范围在 [0, 255]，进行归一化
        image_path = image_path / 255.0
        # 检查并调整形状（这里只是示例，具体调整需根据实际情况）
        if len(image_path.shape) == 3:
            image_path = image_path.unsqueeze(0)  # 添加批次维度
        u = uuid.uuid4()
        p = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output", "%s.jpg" % str(u))
        torchvision.utils.save_image(image_path, p)
        # u = uuid.uuid4()
        # p = os.path.join(os.path.dirname(os.path.abspath(__file__)),"output","%s.jpg" % str(u))
        # torchvision.utils.save_image(image_path, p)
        return (p,)