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

import base64
import io

from PIL import Image, ImageOps

from nodes import ImageScale, VAEEncode
import numpy as np
from perlin_noise import perlin_noise
import torch
from torchvision.transforms import ToPILImage
from ..utils import VyroParams
from .perlin import perlin_power_fractal_batch
from PIL import Image


def pil2tensor(image):
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)

class VyroPipeInputV2:
    def __init__(self):
        self.vae_encoder = VAEEncode()
        self.image_scale = ImageScale()
        pass
    
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "user_prompt": ("STRING", {"multiline": True}),
                "mode": (VyroParams.MODE, {"default": VyroParams.MODE[0]}),
                "vae": ("VAE",),
                "init_noise_mode": (["perlin1","perlin2","zeros"], {"default": "perlin1"}),
            },
            "optional": {
                "user_neg_prompt": ("STRING",{"multiline": True}),
                "batch_size": ("INT", {"default": 1, "max": 4, "min": 1, "step": 1}),
                "cfg": ("FLOAT", {"default": 7.5, "min":1.0, "max":30.0, "step":0.1}),
                "steps": ("INT", {"default": 20, "min":10, "max":150, "step":1}),
                "width": ("INT", {"default": 1024, "min":64, "max":4096, "step":8}),
                "height": ("INT", {"default": 1024, "min":64, "max":4096, "step":8}),
                "seed": ("INT", {"default": 1, "min": 1, "max": 2**32 - 1}),
                "init_img": ("STRING", {"multiline": True}),
                "denoise": ("FLOAT", {"default" : 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "stage1_strength": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}),
                "stage2_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "efficiency_multiplier": (VyroParams.MULTIPLIER, {"default": VyroParams.MULTIPLIER[0]}),
                
                
            }
        }

    RETURN_TYPES = ("VYRO_PARAMS",)
    RETURN_NAMES = ("vyro_params",)
    

    FUNCTION = "prep_input"

    #OUTPUT_NODE = False

    CATEGORY = "Vyro"

    
    
    def blank_image(self, width, height, red, green, blue):
        # Ensure multiples
        width = (width // 8) * 8
        height = (height // 8) * 8

        # Blend image
        blank = Image.new(mode="RGB", size=(width, height),
                          color=(red, green, blue))

        return pil2tensor(blank)
    
    def image_to_mask(self, image, channel):
        channels = ["red", "green", "blue", "alpha"]
        mask = image[:, :, :, channels.index(channel)]
        return (mask,)
    
    def mask_to_image(self, mask):
        result = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
        return (result,)

   



    def upscale_image(self,image, upscale_method='nearest', width=512, height=512, crop='disabled'):
        if width == 0 and height == 0:
            return image
        
        original_aspect_ratio = image.width / image.height
        target_aspect_ratio = width / height
        
        pil_methods = {
            'nearest': Image.NEAREST,
            'bilinear': Image.BILINEAR,
            'area': Image.HAMMING,
            'bicubic': Image.BICUBIC,
            'lanczos': Image.LANCZOS
        }
        
        if crop == 'center':
            if original_aspect_ratio > target_aspect_ratio:
                new_height = height
                new_width = int(round(height * original_aspect_ratio))
            else:
                new_width = width
                new_height = int(round(width / original_aspect_ratio))
            
            resized_image = image.resize((new_width, new_height), pil_methods[upscale_method])
            
            left = (resized_image.width - width) / 2
            top = (resized_image.height - height) / 2
            right = (resized_image.width + width) / 2
            bottom = (resized_image.height + height) / 2
            resized_image = resized_image.crop((left, top, right, bottom))
            
        else:
            if width == 0:
                width = int(round(height * original_aspect_ratio))
            elif height == 0:
                height = int(round(width / original_aspect_ratio))
            
            resized_image = image.resize((width, height), pil_methods[upscale_method])
        
        return resized_image


    def prep_input(self, user_prompt, mode, vae, init_noise_mode, user_neg_prompt='', batch_size=1, cfg=7.5, steps=20, width=1024, height=1024, seed=1, init_img='', denoise=1.0,stage1_strength=0.25, stage2_strength=1.0, efficiency_multiplier=1.0):

        width = int(width / efficiency_multiplier)
        height = int(height / efficiency_multiplier)
        
        # width = (width // 8) * 8
        # height = (height // 8) * 8
        
        if init_img == "" or init_img is None or init_img == 'undefined':
            generator = torch.cuda.manual_seed(seed)
            if 'perlin2' in init_noise_mode:
                init_img = perlin_noise(grid_shape=(2, 8), out_shape=(width // 8, height // 8), batch_size=batch_size*4, generator=generator)
                init_img = init_img.reshape(batch_size, 4, height // 8, width // 8)
            elif 'perlin1' in init_noise_mode:
                init_img = perlin_power_fractal_batch(
                    batch_size=batch_size,
                    width=width // 8,
                    height=height // 8,
                    X=0,
                    Y=0,
                    Z=0,
                    frame=0,
                    seed=seed
                )[0]
                init_img = init_img.reshape(batch_size, 4, height // 8, width // 8)
            else:
                init_img = torch.zeros([batch_size, 4, height // 8, width // 8])
            # init_img = self.blank_image(width, height, 0, 0, 0)
            # init_img = init_img.unsqueeze(0)
            # init_img = latent = torch.zeros([batch_size, 4, height // 8, width // 8])
        else:
            # Convert init_img from byte string to PIL image
            # Get bytes from string
            init_img_bytes = base64.b64decode(init_img)
            img = Image.open(io.BytesIO(init_img_bytes))
            # Center crop to width/height
            
            init_img = pil2tensor(img)
            init_img = self.image_scale.upscale(init_img, "nearest-exact", int(width), int(height), "center")[0]
            init_img = init_img.squeeze(0)
            init_img = init_img.repeat(batch_size, 1, 1, 1)
            
            init_img = self.vae_encoder.encode(vae, init_img)[0]['samples']
        

        
        params = VyroParams(latents={"samples":init_img}, user_prompt=user_prompt, mode=mode, cfg=cfg, batch_size=batch_size, steps=steps, width=width, height=height, seed=seed, denoise=denoise, user_neg_prompt=user_neg_prompt, stage1_strength=stage1_strength, stage2_strength=stage2_strength, efficiency_multiplier=efficiency_multiplier)

         
        return (params,)

NODE_CLASS_MAPPINGS = {
    "Vyro Pipe Input V2": VyroPipeInputV2,
}

# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
    "VyroPipeInputV2": "Vyro Pipe Input V2",
}