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"
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
|
||||
|
|
Loading…
Reference in New Issue