import random from PIL import Image, ImageOps import torch import os import hashlib import folder_paths import numpy as np # Tensor to PIL def tensor2pil(image): return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) # PIL to Tensor def pil2tensor(image): return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) def get_max_size (width, height, max, upscale="false"): aspect_ratio = width / height fit_width = max fit_height = max if upscale == "false" and width <= max and height <= max: return (width, height, aspect_ratio) if aspect_ratio > 1: fit_height = int(max / aspect_ratio) else: fit_width = int(max * aspect_ratio) new_width, new_height = octal_sizes(fit_width, fit_height) return (new_width, new_height, aspect_ratio) def get_image_size(IMAGE) -> tuple[int, int]: samples = IMAGE.movedim(-1, 1) size = samples.shape[3], samples.shape[2] return size def octal_sizes (width, height): octalwidth = width if width % 8 == 0 else width + (8 - width % 8) octalheight = height if height % 8 == 0 else height + (8 - height % 8) return (octalwidth, octalheight) def vae_encode_crop_pixels(pixels): x = (pixels.shape[1] // 8) * 8 y = (pixels.shape[2] // 8) * 8 if pixels.shape[1] != x or pixels.shape[2] != y: x_offset = (pixels.shape[1] % 8) // 2 y_offset = (pixels.shape[2] % 8) // 2 pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :] return pixels def blend_latents(latent, noised_latent, alpha): return latent * alpha + noised_latent * (1. - alpha) def fit_and_resize_image (image, vae, max_size=768, resampling="bicubic", upscale="false", batch_size=1, add_noise=0.0): size = get_image_size(image) new_width, new_height, aspect_ratio = get_max_size(size[0], size[1], max_size, upscale) img = tensor2pil(image) resized_image = img.resize((new_width, new_height), resample=Image.Resampling(resample_filters[resampling])) tensor_img = pil2tensor(resized_image) pixels = vae_encode_crop_pixels(tensor_img) if add_noise > 0.0: noise = torch.randn_like(vae.encode(pixels[:,:,:,:3])) noised_latent = blend_latents(noise, vae.encode(pixels[:,:,:,:3]), add_noise) noised_latent = noised_latent.repeat((batch_size, 1,1,1)) # vae encode the image t = vae.encode(pixels[:,:,:,:3]) # batch the latent vectors batched = t.repeat((batch_size, 1,1,1)) return ( {"samples": noised_latent if add_noise > 0.0 else batched}, tensor_img, new_width, new_height, aspect_ratio, ) resample_filters = { 'nearest': 0, 'lanczos': 1, 'bilinear': 2, 'bicubic': 3, } class FitSize: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "original_width": ("INT", {}), "original_height": ("INT", {}), "max_size": ("INT", {"default": 768, "step": 8}), "upscale": (["false", "true"],) } } RETURN_TYPES = ("INT", "INT", "FLOAT") RETURN_NAMES = ("Fit Width", "Fit Height", "Aspect Ratio") FUNCTION = "fit_to_size" CATEGORY = "Fitsize/Numbers" def fit_to_size (self, original_width, original_height, max_size, upscale="false"): values = get_max_size(original_width, original_height, max_size, upscale) return values class FitSizeFromImage: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "max_size": ("INT", {"default": 768, "step": 8}), "upscale": (["false", "true"],) } } RETURN_TYPES = ("INT", "INT", "FLOAT") RETURN_NAMES = ("Fit Width", "Fit Height", "Aspect Ratio") FUNCTION = "fit_to_size_from_image" CATEGORY = "Fitsize/Numbers" def fit_to_size_from_image (self, image, max_size, upscale="false"): size = get_image_size(image) values = get_max_size(size[0], size[1], max_size, upscale) return values class FitResizeImage: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "max_size": ("INT", {"default": 768, "step": 8}), "resampling": (["lanczos", "nearest", "bilinear", "bicubic"],), "upscale": (["false", "true"],) } } RETURN_TYPES = ("IMAGE","INT","INT","FLOAT") RETURN_NAMES = ("Image","Fit Width", "Fit Height", "Aspect Ratio") FUNCTION = "fit_resize_image" CATEGORY = "Fitsize/Image" def fit_resize_image (self, image, max_size=768, resampling="bicubic", upscale="false", latent=False): size = get_image_size(image) img = tensor2pil(image) new_width, new_height, aspect_ratio = get_max_size(size[0], size[1], max_size, upscale) resized_image = img.resize((new_width, new_height), resample=Image.Resampling(resample_filters[resampling])) return (pil2tensor(resized_image),new_width,new_height,aspect_ratio) class FitResizeLatent(): def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "vae": ("VAE",), "max_size": ("INT", {"default": 768, "step": 8}), "resampling": (["lanczos", "nearest", "bilinear", "bicubic"],), "upscale": (["false", "true"],), "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}), "add_noise": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.01}), } } RETURN_TYPES = ( "LATENT", "IMAGE", "INT", "INT", "FLOAT", ) RETURN_NAMES = ( "Latent", "Image", "Fit Width", "Fit Height", "Aspect Ratio", ) FUNCTION = "fit_resize_latent" CATEGORY = "Fitsize/Image" def fit_resize_latent (self, image, vae, max_size=768, resampling="bicubic", upscale="false", batch_size=1, add_noise=0.0): return fit_and_resize_image(image, vae, max_size, resampling, upscale, batch_size, add_noise) class LoadToFitResizeLatent(): def __init__(self): pass @classmethod def INPUT_TYPES(s): input_dir = folder_paths.get_input_directory() files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] return { "required": { "vae": ("VAE",), "image": (sorted(files), {"image_upload": True}), "max_size": ("INT", {"default": 768, "step": 8}), "resampling": (["lanczos", "nearest", "bilinear", "bicubic"],), "upscale": (["false", "true"],), "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}), "add_noise": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.01}), } } RETURN_TYPES = ( "LATENT", "IMAGE", "INT", "INT", "FLOAT", "MASK", ) RETURN_NAMES = ( "Latent", "Image", "Width", "Height", "Aspect Ratio", "Mask", ) FUNCTION = "fit_resize_latent" CATEGORY = "Fitsize/Image" @staticmethod def load_image(image): if (type(image) == str): image_path = folder_paths.get_annotated_filepath(image) i = Image.open(image_path) 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)) @classmethod def IS_CHANGED(s, vae, image, max_size=768, resampling="bicubic", upscale="false", batch_size=1, add_noise=0.0): image_path = folder_paths.get_annotated_filepath(image) m = hashlib.sha256() with open(image_path, 'rb') as f: m.update(f.read()) return m.digest().hex() @classmethod def VALIDATE_INPUTS(s, vae, image, max_size=768, resampling="bicubic", upscale="false", batch_size=1, add_noise=0.0): if not folder_paths.exists_annotated_filepath(image): return "Invalid image file: {}".format(image) return True def fit_resize_latent (self, vae, image, max_size=768, resampling="bicubic", upscale="false", batch_size=1, add_noise=0.0): got_image,mask = self.load_image(image) latent,img,new_width,new_height,aspect_ratio = fit_and_resize_image(got_image, vae, max_size, resampling, upscale, batch_size, add_noise) return ( latent, img, new_width, new_height, aspect_ratio, mask, ) class CropImageIntoEvenPieces: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "rows": ("INT", {"default": 3, "min": 1, "max": 32, "step": 1,}), "columns": ("INT", {"default": 1, "min": 1, "max": 32, "step": 1,}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "run" CATEGORY = "Fitsize/Image" def run(self, image, rows, columns): if rows < 1: rows = 1 if columns < 1: columns = 1 w = image.shape[2] # width h = image.shape[1] # height crop_width = int(w / columns) crop_height = int(h / rows) image = image.numpy() pieces = [] for i in range(rows): for j in range(columns): y = i * crop_height x = j * crop_width crop = image[: , y : y + crop_height , x : x + crop_width , :] pieces.append(torch.from_numpy(crop)) return (torch.cat(pieces, dim=0), ) class ImageRegionMask: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "rows": ("INT", {"default": 3, "min": 1, "max": 32, "step": 1,}), "columns": ("INT", {"default": 1, "min": 1, "max": 32, "step": 1,}), "chosen_row": ("INT", {"default": 0, "min": 0, "max": 32, "step": 1,}), "chosen_column": ("INT", {"default": 0, "min": 0, "max": 32, "step": 1,}), }, } RETURN_TYPES = ("MASK",) FUNCTION = "run" CATEGORY = "Fitsize/Mask" def run(self, image, rows, columns, chosen_row, chosen_column): if rows < 1: rows = 1 if columns < 1: columns = 1 w = image.shape[2] # width h = image.shape[1] # height crop_width = int(w / columns) crop_height = int(h / rows) mask = torch.zeros((h, w)) min_y = crop_height * chosen_row max_y = min_y + crop_height min_x = crop_width * chosen_column max_x = min_x + crop_width mask[int(min_y):int(max_y), int(min_x):int(max_x)] = 1 return (mask.unsqueeze(0), ) class RandomImageFromBatch: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "images": ("IMAGE", ), "seed": ("INT", {"default": 0}), "start_index": ("INT", {"default": -1, "min": -1, "max": 32, "step": 1,}), "select_amount": ("INT", {"default": 1, "min": 1, "max": 32, "step": 1,}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "run" CATEGORY = "Fitsize/Image" def run(self, images, seed, start_index, select_amount): # if type(images) == torch.Tensor: # images = images.numpy() if start_index == -1: start_index = np.random.randint(0, images.shape[0]) if start_index >= images.shape[0]: start_index = images.shape[0]-1 if select_amount > images.shape[0]: select_amount = images.shape[0] if select_amount < 1: select_amount = 1 selected = images[start_index:start_index + select_amount] return (selected, ) class RandomImageFromList: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "list": ("IMAGE", ), "seed": ("INT", {"default": 0}), "start_index": ("INT", {"default": -1, "min": -1, "max": 32, "step": 1,}), "select_amount": ("INT", {"default": 1, "min": 1, "max": 32, "step": 1,}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "run" CATEGORY = "Fitsize/Image" def run(self, list, seed, start_index, select_amount): print(f'type of list: {type(list)}, length: {len(list)}') list_length = len(list) if start_index == -1: start_index = np.random.randint(0, list_length) # return random.choice(list, select_amount) if start_index >= list_length: start_index = list_length-1 if select_amount > list_length: select_amount = list_length if select_amount < 1: select_amount = 1 selected = list[start_index:start_index + select_amount] print(f'selected: {start_index} to {start_index + select_amount} found {len(selected)}') return (selected, ) class RandomImageFromBatches: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "seed": ("INT",{"default": 0}), "start_index": ("INT", {"default": -1, "min": -1, "max": 32, "step": 1,}), "select_amount": ("INT", {"default": 1, "min": 1, "max": 32, "step": 1,}), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "run" CATEGORY = "Fitsize/Image" def run(self, seed=0, start_index=0, select_amount=1, **kwargs): batches = kwargs.values() selected = [] print(f'len(batches): {len(batches)}') for img in batches: # if type(images) == torch.Tensor: # images = images.numpy() if start_index == -1: start_index = np.random.randint(0, img.shape[0]) if start_index >= img.shape[0]: start_index = img.shape[0]-1 if select_amount > img.shape[0]: select_amount = img.shape[0] if select_amount < 1: select_amount = 1 # add images to selected selected.append(img[start_index:start_index + select_amount]) # try to return a tensor of images if all widths and heights match try: selected = torch.cat(selected, dim=0) except: pass return (selected, )