from PIL import Image
import io
import base64
import numpy as np
import torch

class imagetob64():
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"images": ("IMAGE", ), },}

    CATEGORY = "image"
    RETURN_TYPES = ()
    OUTPUT_NODE = True
    FUNCTION = "images_tob64"
    def images_tob64(self, images):
        imges = []
        for image in images:
            i = 255. * image.cpu().numpy()
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
            buffer = io.BytesIO()
            img.save(buffer, format="JPEG")
            buffer.seek(0)
            imgb64 = base64.b64encode(buffer.getvalue()).decode()
            imges.append(imgb64)
        return { "ui": { "images": imges } }

class loadimageb64:

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        """
            Return a dictionary which contains config for all input fields.
            Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT".
            Input types "INT", "STRING" or "FLOAT" are special values for fields on the node.
            The type can be a list for selection.

            Returns: `dict`:
                - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required`
                - Value input_fields (`dict`): Contains input fields config:
                    * Key field_name (`string`): Name of a entry-point method's argument
                    * Value field_config (`tuple`):
                        + First value is a string indicate the type of field or a list for selection.
                        + Secound value is a config for type "INT", "STRING" or "FLOAT".
        """
        return {"required": {"b64img": ("STRING", {"multiline": False}),}}

    CATEGORY = "image"
    RETURN_TYPES = ("IMAGE",)

    FUNCTION = "load_imageb64"
    def load_imageb64(self, b64img):
        pilim = Image.open(io.BytesIO(base64.b64decode(b64img)))
        tensor_bw = pilim.convert("RGB")
        tensor_bw = np.array(tensor_bw).astype(np.float32) / 255.0
        tensor_bw = torch.from_numpy(tensor_bw)[None,]

        return (tensor_bw,)


NODE_CLASS_MAPPINGS = {
    "loadimageb64": loadimageb64,
    "imagetob64": imagetob64,
}