diff --git a/input/example.png b/input/example.png new file mode 100644 index 00000000..7b7f3c9c Binary files /dev/null and b/input/example.png differ diff --git a/nodes.py b/nodes.py index a297fc3b..20ae2838 100644 --- a/nodes.py +++ b/nodes.py @@ -58,6 +58,10 @@ class VAEEncode: FUNCTION = "encode" def encode(self, vae, pixels): + x = (pixels.shape[1] // 64) * 64 + y = (pixels.shape[2] // 64) * 64 + if pixels.shape[1] != x or pixels.shape[2] != y: + pixels = pixels[:,:x,:y,:] return (vae.encode(pixels), ) class CheckpointLoader: @@ -205,6 +209,24 @@ class SaveImage: img.save(f"output/ComfyUI_{self.counter:05}_.png", pnginfo=metadata, optimize=True) self.counter += 1 +class LoadImage: + input_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input") + @classmethod + def INPUT_TYPES(s): + return {"required": + {"image": (os.listdir(s.input_dir), )}, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "load_image" + def load_image(self, image): + image_path = os.path.join(self.input_dir, image) + image = Image.open(image_path).convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image[None])[None,] + return image + + NODE_CLASS_MAPPINGS = { "KSampler": KSampler, @@ -216,6 +238,7 @@ NODE_CLASS_MAPPINGS = { "EmptyLatentImage": EmptyLatentImage, "LatentUpscale": LatentUpscale, "SaveImage": SaveImage, + "LoadImage": LoadImage }