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