File size: 1,676 Bytes
ca2a3d8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
from PIL import Image, ImageOps
import numpy as np
import torch, base64, io
def b64_to_img_and_mask(image_base64):
imageData = base64.b64decode(image_base64)
i = Image.open(io.BytesIO(imageData))
if hasattr(i, 'is_animated') and i.is_animated:
images = []
for frame in range(i.n_frames):
i.seek(frame)
images.append(i.convert("RGB"))
i.seek(0)
image = np.array(images).astype(np.float32) / 255.0
image = torch.from_numpy(image)
else:
i = ImageOps.exif_transpose(i)
image = i.convert("RGB")
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
if 'A' in i.getbands():
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
mask = 1. - torch.from_numpy(mask)
else:
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
return (image, mask.unsqueeze(0))
class SwarmLoadImageB64:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image_base64": ("STRING", {"multiline": True})
}
}
CATEGORY = "SwarmUI/images"
RETURN_TYPES = ("IMAGE", "MASK")
FUNCTION = "load_image_b64"
DESCRIPTION = "Loads an image from a base64 string. Works like a regular LoadImage node, but with input format designed to be easier to use through automated calls, including SwarmUI with custom workflows."
def load_image_b64(self, image_base64):
return b64_to_img_and_mask(image_base64)
NODE_CLASS_MAPPINGS = {
"SwarmLoadImageB64": SwarmLoadImageB64,
}
|