50 lines
1.3 KiB
Python
50 lines
1.3 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageOps
|
|
import base64
|
|
from io import BytesIO
|
|
import folder_paths
|
|
import os
|
|
|
|
|
|
class ImagePin:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_dir = folder_paths.get_input_directory()
|
|
files = [
|
|
f
|
|
for f in os.listdir(input_dir)
|
|
if os.path.isfile(os.path.join(input_dir, f))
|
|
]
|
|
return {
|
|
"required": {
|
|
"image": (sorted(files) + ["#DECODE_FROM_BASE64"], {"image_upload": True}),
|
|
"image_data": ("STRING", {"multiline": False}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "utils"
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "load_image"
|
|
|
|
def load_image(self, image, image_data):
|
|
if image_data.startswith('data:image/webp;base64,'):
|
|
image_data = base64.b64decode(image_data.split(",")[1])
|
|
else:
|
|
image_data = base64.b64decode(image_data.split(",")[1])
|
|
|
|
i = Image.open(BytesIO(image_data))
|
|
i = ImageOps.exif_transpose(i)
|
|
image = i.convert("RGB")
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
return (image,)
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, image, image_data):
|
|
if image_data is None:
|
|
return "image_data is required"
|
|
|
|
return True
|