remixXL / webp_test.py
Stkzzzz222's picture
Upload webp_test.py
52c37ff verified
import os
from PIL import Image
import numpy as np
import torch
class LoadWebPAnimation:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"path": ("STRING", {"default": ""}),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "load_animation"
def load_animation(self, path):
# Check if the file exists
if not os.path.exists(path):
raise FileNotFoundError(f"File not found: {path}")
# Open the .webp file
with Image.open(path) as img:
# Check if it's an animation
if hasattr(img, 'is_animated') and img.is_animated:
frames = []
# Extract each frame
for i in range(img.n_frames):
img.seek(i)
frame = img.copy().convert('RGB') # Ensure 3 channels (RGB)
frames.append(np.array(frame))
batch = np.stack(frames, axis=0) # Shape: (num_frames, height, width, 3)
else:
# Single image case
frame = img.convert('RGB')
batch = np.expand_dims(np.array(frame), axis=0) # Shape: (1, height, width, 3)
# Convert to float32 and normalize to [0, 1]
batch = batch.astype(np.float32) / 255.0
# Convert to PyTorch tensor without permuting
batch = torch.from_numpy(batch) # Shape: (num_frames, height, width, 3)
return (batch,) # Return as a tuple
# Register the node with ComfyUI
NODE_CLASS_MAPPINGS = {
"LoadWebPAnimation": LoadWebPAnimation
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LoadWebPAnimation": "Load WebP Animation"
}