added load image resize fitter

This commit is contained in:
Hamilton
2023-11-23 14:39:31 -08:00
parent 047853c9d2
commit df266b0e58
2 changed files with 219 additions and 29 deletions
+9 -5
View File
@@ -1,17 +1,21 @@
from .nodes import FitSize, FitSizeFromImage, FitResizeImage
from .nodes import FitSize, FitSizeFromImage, FitResizeImage, FitResizeLatent, LoadToFitResizeLatent
# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
"FitSizeByMax": FitSize,
"FitSizeFromInt": FitSize,
"FitSizeFromImage": FitSizeFromImage,
"FitResizeImage": FitResizeImage,
"FitSizeResizeImage": FitResizeImage,
"FitSizeResizeLatent": FitResizeLatent,
"LoadToFitResizeLatent": LoadToFitResizeLatent,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FitSizeByMax": "Fit Size",
"FitSizeFromInt": "Fit Size From Int",
"FitSizeFromImage": "Fit Size From Image",
"FitResizeImage": "Fit Resize Image",
"FitSizeResizeImage": "Fit Resize Image",
"FitSizeResizeLatent": "Fit Resize Latent",
"LoadToFitResizeLatent": "Load Image To Fit Resize Latent",
}
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
+210 -24
View File
@@ -1,6 +1,9 @@
from PIL import Image
import numpy as np
from PIL import Image, ImageOps
import torch
import os
import hashlib
import folder_paths
import numpy as np
# Tensor to PIL
def tensor2pil(image):
@@ -10,12 +13,15 @@ def tensor2pil(image):
def pil2tensor(image):
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
def get_max_size (width, height, max):
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:
@@ -28,6 +34,18 @@ def get_image_size(IMAGE) -> tuple[int, int]:
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)
resample_filters = {
'nearest': 0,
'lanczos': 1,
'bilinear': 2,
'bicubic': 3,
}
class FitSize:
def __init__(self):
pass
@@ -39,6 +57,7 @@ class FitSize:
"original_width": ("INT", {}),
"original_height": ("INT", {}),
"max_size": ("INT", {"default": 768, "step": 8}),
"upscale": (["false", "true"],)
}
}
@@ -48,8 +67,8 @@ class FitSize:
CATEGORY = "Fitsize"
def fit_to_size (self, original_width, original_height, max_size):
values = get_max_size(original_width, original_height, max_size)
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:
@@ -62,6 +81,7 @@ class FitSizeFromImage:
"required": {
"image": ("IMAGE",),
"max_size": ("INT", {"default": 768, "step": 8}),
"upscale": (["false", "true"],)
}
}
@@ -71,9 +91,9 @@ class FitSizeFromImage:
CATEGORY = "Fitsize"
def fit_to_size_from_image (self, image, max_size):
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)
values = get_max_size(size[0], size[1], max_size, upscale)
return values
class FitResizeImage:
@@ -87,6 +107,7 @@ class FitResizeImage:
"image": ("IMAGE",),
"max_size": ("INT", {"default": 768, "step": 8}),
"resampling": (["lanczos", "nearest", "bilinear", "bicubic"],),
"upscale": (["false", "true"],)
}
}
@@ -96,26 +117,191 @@ class FitResizeImage:
CATEGORY = "Fitsize"
def fit_resize_image (self, image, max_size=768, resampling="bicubic"):
def fit_resize_image (self, image, max_size=768, resampling="bicubic", upscale="false", latent=False):
size = get_image_size(image)
img = tensor2pil(image)
octalwidth = size[0] if size[0] % 8 == 0 else size[0] + (8 - size[0] % 8)
octalheight = size[1] if size[1] % 8 == 0 else size[1] + (8 - size[1] % 8)
new_width, new_height, aspect_ratio = get_max_size(octalwidth, octalheight, max_size)
print("new values",new_width, new_height, aspect_ratio)
resample_filters = {
'nearest': 0,
'bilinear': 2,
'bicubic': 3,
'lanczos': 1
}
print("Resampling:", resampling, resample_filters[resampling])
octalwidth, octalheight = octal_sizes(size[0], size[1])
new_width, new_height, aspect_ratio = get_max_size(octalwidth, octalheight, 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)
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}),
}
}
RETURN_TYPES = (
"LATENT",
"IMAGE",
"INT",
"INT",
"FLOAT",
)
RETURN_NAMES = (
"Latent",
"Image",
"Fit Width",
"Fit Height",
"Aspect Ratio",
)
FUNCTION = "fit_resize_latent"
CATEGORY = "Fitsize"
@staticmethod
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 fit_resize_latent (self, image, vae, max_size=768, resampling="bicubic", upscale="false", batch_size=1):
size = get_image_size(image)
img = tensor2pil(image)
octalwidth, octalheight = octal_sizes(size[0], size[1])
new_width, new_height, aspect_ratio = get_max_size(octalwidth, octalheight, max_size, upscale)
resized_image = img.resize((new_width, new_height), resample=Image.Resampling(resample_filters[resampling]))
tensor_img = pil2tensor(resized_image)
# vae encode the image
pixels = self.vae_encode_crop_pixels(tensor_img)
t = vae.encode(pixels[:,:,:,:3])
# batch the latent vectors
batched = t.repeat((batch_size, 1,1,1))
return (
{"samples":batched},
tensor_img,
new_width,
new_height,
aspect_ratio,
)
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}),
}
}
RETURN_TYPES = (
"LATENT",
"IMAGE",
"INT",
"INT",
"FLOAT",
)
RETURN_NAMES = (
"Latent",
"Image",
"Fit Width",
"Fit Height",
"Aspect Ratio",
)
FUNCTION = "fit_resize_latent"
CATEGORY = "Fitsize"
@staticmethod
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
@staticmethod
def load_image(image):
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):
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):
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):
got_image,mask = self.load_image(image)
size = get_image_size(got_image)
img = tensor2pil(got_image)
new_width, new_height, aspect_ratio = get_max_size(size[0], size[1], max_size, upscale)
octalwidth, octalheight = octal_sizes(new_width, new_height)
resized_image = img.resize((octalwidth, octalheight), resample=Image.Resampling(resample_filters[resampling]))
tensor_img = pil2tensor(resized_image)
# vae encode the image
pixels = self.vae_encode_crop_pixels(tensor_img)
t = vae.encode(pixels[:,:,:,:3])
# batch the latent vectors
batched = t.repeat((batch_size, 1,1,1))
return (
{"samples":batched},
tensor_img,
octalwidth,
octalheight,
aspect_ratio,
)