Add a LoadImage node to load images for img2img.
This commit is contained in:
parent
220afe3310
commit
15f8da2849
Binary file not shown.
After Width: | Height: | Size: 8.4 KiB |
23
nodes.py
23
nodes.py
|
@ -58,6 +58,10 @@ class VAEEncode:
|
||||||
FUNCTION = "encode"
|
FUNCTION = "encode"
|
||||||
|
|
||||||
def encode(self, vae, pixels):
|
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), )
|
return (vae.encode(pixels), )
|
||||||
|
|
||||||
class CheckpointLoader:
|
class CheckpointLoader:
|
||||||
|
@ -205,6 +209,24 @@ class SaveImage:
|
||||||
img.save(f"output/ComfyUI_{self.counter:05}_.png", pnginfo=metadata, optimize=True)
|
img.save(f"output/ComfyUI_{self.counter:05}_.png", pnginfo=metadata, optimize=True)
|
||||||
self.counter += 1
|
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 = {
|
NODE_CLASS_MAPPINGS = {
|
||||||
"KSampler": KSampler,
|
"KSampler": KSampler,
|
||||||
|
@ -216,6 +238,7 @@ NODE_CLASS_MAPPINGS = {
|
||||||
"EmptyLatentImage": EmptyLatentImage,
|
"EmptyLatentImage": EmptyLatentImage,
|
||||||
"LatentUpscale": LatentUpscale,
|
"LatentUpscale": LatentUpscale,
|
||||||
"SaveImage": SaveImage,
|
"SaveImage": SaveImage,
|
||||||
|
"LoadImage": LoadImage
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue