540 lines
15 KiB
Python
540 lines
15 KiB
Python
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, )
|