from PIL import Image
import numpy as np
import torch

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

    CATEGORY = "SwarmUI/images"
    RETURN_TYPES = ("IMAGE", "MASK",)
    FUNCTION = "rem"

    def rem(self, images):
        from rembg import remove

        output = []
        masks = []
        for image in images:
            i = 255.0 * image.cpu().numpy()
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
            img = img.convert("RGBA")
            img = remove(img, post_process_mask=True)
            output.append(np.array(img).astype(np.float32) / 255.0)
            if 'A' in img.getbands():
                mask = np.array(img.getchannel('A')).astype(np.float32) / 255.0
                masks.append(1. - mask)
            else:
                masks.append(np.zeros((64,64), dtype=np.float32))
        return (torch.from_numpy(np.array(output)), torch.from_numpy(np.array(masks)))

NODE_CLASS_MAPPINGS = {
    "SwarmRemBg": SwarmRemBg,
}