|
|
|
@ -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 |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|