ComfyUI-CustomNode/nodes/image.py

37 lines
1.1 KiB
Python

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,)