Imagine/imagine-v5/vyro_workflows/nodes/util.py

from typing import Tuple

from PIL import Image, ImageOps

import numpy as np
import torch

# from ..utils.face_swap import pil2mask, pil2tensor, tensor2pil

def get_latent_size(LATENT, ORIGINAL_VALUES=False) -> Tuple[int, int]:
    lc = LATENT.copy()
    size = lc["samples"].shape[3], lc["samples"].shape[2]
    if ORIGINAL_VALUES == False:
        size = size[0] * 8, size[1] * 8
    return size

class GetLatentSize:
    def __init__(self) -> None:
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "latent": ("LATENT",),
                "original": ([False, True],),
            }
        }

    RETURN_TYPES = ("INT", "INT", "TUPLE",)
    CATEGORY = 'Vyro/Utils'

    FUNCTION = 'get_size'

    def get_size(self, latent, original):
        size = get_latent_size(latent, original)
        return (size[0], size[1], size,)


class MultilineStringNode:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "Text": ("STRING", {
                    "default": "",
                    "multiline": True,
                }),
            }
        }

    RETURN_TYPES = ("STRING",)
    FUNCTION = "get_value"
    CATEGORY = "Vyro/Utils"

    def get_value(self, Text):
        return (Text,)
    
class WAS_Images_To_RGB:
    def __init__(self):
        pass

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

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("images",)
    FUNCTION = "images_to_rgb"

    CATEGORY = "Vyro/Image"

    def images_to_rgb(self, images):

        tensors = []
        for image in images:
            tensors.append(pil2tensor(tensor2pil(image).convert("RGB")))
        tensors = torch.cat(tensors, dim=0)

        return (tensors,)
        
def smooth_region(image, tolerance):
            from scipy.ndimage import gaussian_filter
            image = image.convert("L")
            mask_array = np.array(image)
            smoothed_array = gaussian_filter(mask_array, sigma=tolerance)
            threshold = np.max(smoothed_array) / 2
            smoothed_mask = np.where(smoothed_array >= threshold, 255, 0).astype(np.uint8)
            smoothed_image = Image.fromarray(smoothed_mask, mode="L")
            return ImageOps.invert(smoothed_image.convert("RGB"))   
       
class WAS_Mask_Smooth_Region:


    @classmethod
    def INPUT_TYPES(cls):
        return {
                    "required": {
                        "masks": ("MASK",),
                        "sigma": ("FLOAT", {"default":5.0, "min":0.0, "max":128.0, "step":0.1}),
                    }
                }

    CATEGORY = "Vyro/Image/Masking"

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("MASKS",)

    FUNCTION = "smooth"

    def smooth(self, masks, sigma=128):

        if masks is None:
            masks = torch.zeros((1, 1, 10, 10))  # Replace 10x10 with the desired dimensions
            masks[0, 0, 5, 5] = 1  # One white pixel at the center


        if masks.ndim > 3:
            regions = []
            for mask in masks:
                mask_np = np.clip(255. * mask.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
                pil_image = Image.fromarray(mask_np, mode="L")
                region_mask = self.WT.Masking.smooth_region(pil_image, sigma)
                region_tensor = pil2mask(region_mask).unsqueeze(0).unsqueeze(1)
                regions.append(region_tensor)
            regions_tensor = torch.cat(regions, dim=0)
            return (regions_tensor,)
        else:
            mask_np = np.clip(255. * masks.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
            pil_image = Image.fromarray(mask_np, mode="L")
            region_mask = smooth_region(pil_image, sigma)
            region_tensor = pil2mask(region_mask).unsqueeze(0).unsqueeze(1)
            return (region_tensor,)


class WAS_Latent_Upscale:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"samples": ("LATENT",), "mode": (["area", "bicubic", "bilinear", "nearest"],),
                             "factor": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.01}),
                             "align": (["true", "false"], )}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "latent_upscale"

    CATEGORY = "Vyro/Latent/Transform"

    def latent_upscale(self, samples, mode, factor, align):
        valid_modes = ["area", "bicubic", "bilinear", "nearest"]
        if mode not in valid_modes:
            print(f"Invalid interpolation mode `{mode}` selected. Valid modes are: {', '.join(valid_modes)}")
            return (s, )
        align = True if align == 'true' else False
        if not isinstance(factor, float) or factor <= 0:
            print(f"The input `factor` is `{factor}`, but should be a positive or negative float.")
            return (s, )
        s = samples.copy()
        shape = s['samples'].shape
        size = tuple(int(round(dim * factor)) for dim in shape[-2:])
        if mode in ['linear', 'bilinear', 'bicubic', 'trilinear']:
            s["samples"] = torch.nn.functional.interpolate(
                s['samples'], size=size, mode=mode, align_corners=align)
        else:
            s["samples"] = torch.nn.functional.interpolate(s['samples'], size=size, mode=mode)
        return (s,)

TEXT_TYPE = "STRING"

class WAS_Text_Concatenate:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text_a": (TEXT_TYPE, {"forceInput": (True if TEXT_TYPE == 'STRING' else False)}),
                "text_b": (TEXT_TYPE, {"forceInput": (True if TEXT_TYPE == 'STRING' else False)}),
                "linebreak_addition": (['false','true'], ),
            },
            "optional": {
                "text_c": (TEXT_TYPE, {"forceInput": (True if TEXT_TYPE == 'STRING' else False)}),
                "text_d": (TEXT_TYPE, {"forceInput": (True if TEXT_TYPE == 'STRING' else False)}),
            }
        }

    RETURN_TYPES = (TEXT_TYPE,)
    FUNCTION = "text_concatenate"

    CATEGORY = "Vyro/Text"

    def text_concatenate(self, text_a, text_b, text_c=None, text_d=None, linebreak_addition='false'):
        return_text = text_a + ("\n" if linebreak_addition == 'true' else '') + text_b
        if text_c:
            return_text = return_text + ("\n" if linebreak_addition == 'true' else '') + text_c
        if text_d:
            return_text = return_text + ("\n" if linebreak_addition == 'true' else '') + text_d
        return (return_text, )


class WAS_Image_Threshold:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "image_threshold"

    CATEGORY = "WAS Suite/Image/Process"

    def image_threshold(self, image, threshold=0.5):
        return (pil2tensor(self.apply_threshold(tensor2pil(image), threshold)), )

    def apply_threshold(self, input_image, threshold=0.5):
        # Convert the input image to grayscale
        grayscale_image = input_image.convert('L')

        # Apply the threshold to the grayscale image
        threshold_value = int(threshold * 255)
        thresholded_image = grayscale_image.point(
            lambda x: 255 if x >= threshold_value else 0, mode='L')

        return thresholded_image
    
    

class VyroEmptyLatentImage:
    def __init__(self, device="cpu"):
        self.device = device

    @classmethod
    def INPUT_TYPES(s):
        return {"required": { "vyro_params": ("VYRO_PARAMS",),
                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})}}
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "generate"

    CATEGORY = "latent"

    def generate(self, vyro_params, batch_size=1):
        height = vyro_params.height if vyro_params.height else 1024
        width = vyro_params.width if vyro_params.width else 1024
        
        print(f"LATENT IMAGE SIZE: {height, width}")
        
        latent = torch.zeros([batch_size, 4, height // 8, width // 8])
        return ({"samples":latent}, )


  
NODE_CLASS_MAPPINGS = {
    "Get latent size": GetLatentSize,
    "Images to RGB": WAS_Images_To_RGB,
    "Mask Smooth Region": WAS_Mask_Smooth_Region,
    "Latent Upscale by Factor (WAS)": WAS_Latent_Upscale,
    "Text box": MultilineStringNode,
    "Text Concatenate": WAS_Text_Concatenate,
    "Image Threshold": WAS_Image_Threshold,
    "Vyro Empty Latent Image": VyroEmptyLatentImage
}