updated images, added mask output

README.md

@@ -30,7 +30,7 @@ This node accepts any image input and will extract the width and height automati
 
 ![Fit Size From Image](assets/fitsizefromimage.png)
 
-### Fit Resize Image
+### Fit And Resize Image
 
 Now this is where things get interesting. This node accepts a vae so that we can skip right to outputting a rescaled image. It will output both an image and a latent batch. This makes it a very useful tool for img2img workflows. 
 
@@ -48,7 +48,7 @@ This workflow uses a number of other custom node sets to showcase that this node
 
 ![Fit Resize Image](assets/loadtofitresizeimage.png)
 
-Could this workflow example be simpler? There are a lot of nodes like the Efficiency Nodes that can combine a bunch of functionality, and this node can fit gently in the center of them.
+Here is a simpler workflow example using the Efficiency nodes. The Fit nodes should fit right into the middle of a lot of other workflows.
 
 ![Fit Resize Image](assets/loadtofitresizeimageb.png)
 

__init__.py

@@ -3,17 +3,17 @@ from .nodes import FitSize, FitSizeFromImage, FitResizeImage, FitResizeLatent, L
 # A dictionary that contains all nodes you want to export with their names
 # NOTE: names should be globally unique
 NODE_CLASS_MAPPINGS = {
-    "FitSizeFromInt": FitSize,
-    "FitSizeFromImage": FitSizeFromImage,
-    "FitSizeResizeImage": FitResizeLatent,
-    "LoadToFitResizeImage": LoadToFitResizeLatent,
+    "Fit Size From Int": FitSize,
+    "Fit Size From Image": FitSizeFromImage,
+    "Fit Image And Resize": FitResizeLatent,
+    "Load Image And Resize To Fit": LoadToFitResizeLatent,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "FitSizeFromInt": "Fit Size From Int",
-    "FitSizeFromImage": "Fit Size From Image",
-    "FitSizeResizeImage": "Fit Resize Image",
-    "LoadToFitResizeImage": "Load Image And Resize To Fit",
+    "Fit Size From Int": "Fit Size From Int",
+    "Fit Size From Image": "Fit Size From Image",
+    "Fit Image And Resize": "Fit Image And Resize",
+    "Load Image And Resize To Fit": "Load Image And Resize To Fit",
 }
 
 __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

nodes.py

@@ -41,6 +41,49 @@ def octal_sizes (width, height):
     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,
@@ -145,6 +188,7 @@ class FitResizeLatent():
                 "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}),
             }
         }
 
@@ -166,40 +210,9 @@ class FitResizeLatent():
 
     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, add_noise=0.0):
 
-    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)
-
-        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]))
-        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,
-            )
+        return fit_and_resize_image(image, vae, max_size, resampling, upscale, batch_size, add_noise)
     
 class LoadToFitResizeLatent():
     def __init__(self):
@@ -217,6 +230,7 @@ class LoadToFitResizeLatent():
                 "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}),
             }
         }
 
@@ -226,36 +240,31 @@ class LoadToFitResizeLatent():
         "INT",
         "INT",
         "FLOAT",
+        "MASK",
         )
     RETURN_NAMES = (
         "Latent",
         "Image",
-        "Fit Width",
-        "Fit Height",
+        "Width",
+        "Height",
         "Aspect Ratio",
+        "Mask",
         )
     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 (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)
@@ -264,7 +273,7 @@ class LoadToFitResizeLatent():
         return (image, mask.unsqueeze(0))
 
     @classmethod
-    def IS_CHANGED(s, vae, image, max_size=768, resampling="bicubic", upscale="false", batch_size=1):
+    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:
@@ -272,35 +281,22 @@ class LoadToFitResizeLatent():
         return m.digest().hex()
 
     @classmethod
-    def VALIDATE_INPUTS(s, vae, image, max_size=768, resampling="bicubic", upscale="false", batch_size=1):
+    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):
+    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)
 
-        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)
-        
-        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))
+        latent,img,new_width,new_height,aspect_ratio = fit_and_resize_image(got_image, vae, max_size, resampling, upscale, batch_size, add_noise)
 
         return (
-            {"samples":batched},
-            tensor_img,
+            latent,
+            img,
             new_width,
             new_height,
             aspect_ratio,
+            mask,
             )