Skip to content

Commit

Permalink
LoadImage now loads all the frames from animated images as a batch.
Browse files Browse the repository at this point in the history
  • Loading branch information
comfyanonymous committed Dec 20, 2023
1 parent 5f54614 commit a1e1c69
Showing 1 changed file with 24 additions and 11 deletions.
35 changes: 24 additions & 11 deletions nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import time
import random

from PIL import Image, ImageOps
from PIL import Image, ImageOps, ImageSequence
from PIL.PngImagePlugin import PngInfo
import numpy as np
import safetensors.torch
Expand Down Expand Up @@ -1410,17 +1410,30 @@ def INPUT_TYPES(s):
FUNCTION = "load_image"
def load_image(self, 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 'A' in i.getbands():
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
mask = 1. - torch.from_numpy(mask)
img = Image.open(image_path)
output_images = []
output_masks = []
for i in ImageSequence.Iterator(img):
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)
else:
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
output_images.append(image)
output_masks.append(mask.unsqueeze(0))

if len(output_images) > 1:
output_image = torch.cat(output_images, dim=0)
output_mask = torch.cat(output_masks, dim=0)
else:
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
return (image, mask.unsqueeze(0))
output_image = output_images[0]
output_mask = output_masks[0]

return (output_image, output_mask)

@classmethod
def IS_CHANGED(s, image):
Expand Down

0 comments on commit a1e1c69

Please sign in to comment.