LoadImage now loads all the frames from animated images as a batch.

This commit is contained in:
comfyanonymous 2023-12-20 16:39:09 -05:00
parent 5f54614e7f
commit a1e1c69f7d

View File

@ -9,7 +9,7 @@ import math
import time import time
import random import random
from PIL import Image, ImageOps from PIL import Image, ImageOps, ImageSequence
from PIL.PngImagePlugin import PngInfo from PIL.PngImagePlugin import PngInfo
import numpy as np import numpy as np
import safetensors.torch import safetensors.torch
@ -1410,17 +1410,30 @@ class LoadImage:
FUNCTION = "load_image" FUNCTION = "load_image"
def load_image(self, image): def load_image(self, image):
image_path = folder_paths.get_annotated_filepath(image) image_path = folder_paths.get_annotated_filepath(image)
i = Image.open(image_path) img = Image.open(image_path)
i = ImageOps.exif_transpose(i) output_images = []
image = i.convert("RGB") output_masks = []
image = np.array(image).astype(np.float32) / 255.0 for i in ImageSequence.Iterator(img):
image = torch.from_numpy(image)[None,] i = ImageOps.exif_transpose(i)
if 'A' in i.getbands(): image = i.convert("RGB")
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 image = np.array(image).astype(np.float32) / 255.0
mask = 1. - torch.from_numpy(mask) 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: else:
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") output_image = output_images[0]
return (image, mask.unsqueeze(0)) output_mask = output_masks[0]
return (output_image, output_mask)
@classmethod @classmethod
def IS_CHANGED(s, image): def IS_CHANGED(s, image):