423 lines
14 KiB
Python
423 lines
14 KiB
Python
import os
|
|
import json
|
|
import torch
|
|
import numpy as np
|
|
import hashlib
|
|
from typing import List, Dict, Tuple
|
|
from torch import Tensor
|
|
from PIL import Image, ImageSequence
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
import folder_paths
|
|
|
|
from .motion_module import MotionWrapper
|
|
from .model_utils import get_available_models, load_motion_module, get_available_loras, load_lora
|
|
from .utils import pil2tensor, ensure_opencv
|
|
from .sampler import AnimateDiffSampler, AnimateDiffSlidingWindowOptions
|
|
from .logger import logger
|
|
|
|
|
|
SLIDING_CONTEXT_LENGTH = 16
|
|
|
|
video_formats_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "video_formats")
|
|
video_formats = ["video/" + x[:-5] for x in os.listdir(video_formats_dir)]
|
|
|
|
|
|
class AnimateDiffModuleLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_name": (get_available_models(),),
|
|
},
|
|
"optional": {
|
|
"lora_stack": ("MOTION_LORA_STACK",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MOTION_MODULE",)
|
|
CATEGORY = "Animate Diff"
|
|
FUNCTION = "load_motion_module"
|
|
|
|
def inject_loras(self, motion_module: MotionWrapper, lora_stack: List[Tuple[Dict[str, Tensor], float]]):
|
|
for lora in lora_stack:
|
|
(state_dict, alpha) = lora
|
|
|
|
for key in state_dict:
|
|
layer_infos = key.split(".")
|
|
|
|
curr_layer = motion_module
|
|
while len(layer_infos) > 0:
|
|
temp_name = layer_infos.pop(0)
|
|
curr_layer = curr_layer.__getattr__(temp_name)
|
|
|
|
curr_layer.weight.data += alpha * state_dict[key].to(curr_layer.weight.data.device)
|
|
|
|
def eject_loras(self, motion_module: MotionWrapper, lora_stack: List[Tuple[float, Dict[str, Tensor]]]):
|
|
lora_stack.reverse() # should not matter but just in case
|
|
for lora in lora_stack:
|
|
(state_dict, alpha) = lora
|
|
|
|
for key in state_dict:
|
|
layer_infos = key.split(".")
|
|
|
|
curr_layer = motion_module
|
|
while len(layer_infos) > 0:
|
|
temp_name = layer_infos.pop(0)
|
|
curr_layer = curr_layer.__getattr__(temp_name)
|
|
|
|
curr_layer.weight.data -= alpha * state_dict[key].to(curr_layer.weight.data.device)
|
|
|
|
def load_motion_module(
|
|
self,
|
|
model_name: str,
|
|
lora_stack: List = None,
|
|
):
|
|
motion_module = load_motion_module(model_name)
|
|
|
|
# inject loras
|
|
if motion_module.is_v2:
|
|
if hasattr(motion_module, "lora_stack") and isinstance(motion_module.lora_stack, list):
|
|
self.eject_loras(motion_module, motion_module.lora_stack)
|
|
delattr(motion_module, "lora_stack")
|
|
|
|
if isinstance(lora_stack, list):
|
|
self.inject_loras(motion_module, lora_stack)
|
|
setattr(motion_module, "lora_stack", lora_stack)
|
|
|
|
elif isinstance(lora_stack, list):
|
|
logger.warning("LoRA is provided but only motion module v2 is supported.")
|
|
|
|
return (motion_module,)
|
|
|
|
|
|
class AnimateDiffLoraLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"lora_name": (get_available_loras(),),
|
|
"alpha": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}),
|
|
},
|
|
"optional": {
|
|
"lora_stack": ("MOTION_LORA_STACK",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MOTION_LORA_STACK",)
|
|
CATEGORY = "Animate Diff"
|
|
FUNCTION = "load_lora"
|
|
|
|
def load_lora(
|
|
self,
|
|
lora_name: str,
|
|
alpha: float,
|
|
lora_stack: List = None,
|
|
):
|
|
if not lora_stack:
|
|
lora_stack = []
|
|
|
|
lora = load_lora(lora_name)
|
|
lora_stack.append((lora, alpha))
|
|
|
|
return (lora_stack,)
|
|
|
|
|
|
class AnimateDiffCombine:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"frame_rate": (
|
|
"INT",
|
|
{"default": 8, "min": 1, "max": 24, "step": 1},
|
|
),
|
|
"loop_count": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
|
|
"save_image": ("BOOLEAN", {"default": True}),
|
|
"filename_prefix": ("STRING", {"default": "animate_diff"}),
|
|
"format": (["image/gif", "image/webp"] + video_formats,),
|
|
"pingpong": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "Animate Diff"
|
|
FUNCTION = "generate_gif"
|
|
|
|
def generate_gif(
|
|
self,
|
|
images,
|
|
frame_rate: int,
|
|
loop_count: int,
|
|
save_image=True,
|
|
filename_prefix="AnimateDiff",
|
|
format="image/gif",
|
|
pingpong=False,
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
# convert images to numpy
|
|
frames: List[Image.Image] = []
|
|
for image in images:
|
|
img = 255.0 * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
|
|
frames.append(img)
|
|
|
|
# save image
|
|
output_dir = folder_paths.get_output_directory() if save_image else folder_paths.get_temp_directory()
|
|
(
|
|
full_output_folder,
|
|
filename,
|
|
counter,
|
|
subfolder,
|
|
_,
|
|
) = folder_paths.get_save_image_path(filename_prefix, output_dir)
|
|
|
|
metadata = PngInfo()
|
|
if prompt is not None:
|
|
metadata.add_text("prompt", json.dumps(prompt))
|
|
if extra_pnginfo is not None:
|
|
for x in extra_pnginfo:
|
|
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
|
|
|
# save first frame as png to keep metadata
|
|
file = f"{filename}_{counter:05}_.png"
|
|
file_path = os.path.join(full_output_folder, file)
|
|
frames[0].save(
|
|
file_path,
|
|
pnginfo=metadata,
|
|
compress_level=4,
|
|
)
|
|
if pingpong:
|
|
frames = frames + frames[-2:0:-1]
|
|
|
|
format_type, format_ext = format.split("/")
|
|
|
|
if format_type == "image":
|
|
file = f"{filename}_{counter:05}_.{format_ext}"
|
|
file_path = os.path.join(full_output_folder, file)
|
|
frames[0].save(
|
|
file_path,
|
|
format=format_ext.upper(),
|
|
save_all=True,
|
|
append_images=frames[1:],
|
|
duration=round(1000 / frame_rate),
|
|
loop=loop_count,
|
|
compress_level=4,
|
|
)
|
|
else:
|
|
# save webm
|
|
import shutil
|
|
import subprocess
|
|
|
|
ffmpeg_path = shutil.which("ffmpeg")
|
|
if ffmpeg_path is None:
|
|
raise ProcessLookupError("Could not find ffmpeg")
|
|
video_format_path = os.path.join(video_formats_dir, format_ext + ".json")
|
|
with open(video_format_path, "r") as stream:
|
|
video_format = json.load(stream)
|
|
file = f"{filename}_{counter:05}_.{video_format['extension']}"
|
|
file_path = os.path.join(full_output_folder, file)
|
|
dimensions = f"{frames[0].width}x{frames[0].height}"
|
|
args = (
|
|
[
|
|
ffmpeg_path,
|
|
"-v",
|
|
"error",
|
|
"-f",
|
|
"rawvideo",
|
|
"-pix_fmt",
|
|
"rgb24",
|
|
"-s",
|
|
dimensions,
|
|
"-r",
|
|
str(frame_rate),
|
|
"-i",
|
|
"-",
|
|
]
|
|
+ video_format["main_pass"]
|
|
+ [file_path]
|
|
)
|
|
|
|
env = os.environ
|
|
if "environment" in video_format:
|
|
env.update(video_format["environment"])
|
|
with subprocess.Popen(args, stdin=subprocess.PIPE, env=env) as proc:
|
|
for frame in frames:
|
|
proc.stdin.write(frame.tobytes())
|
|
|
|
previews = [
|
|
{
|
|
"filename": file,
|
|
"subfolder": subfolder,
|
|
"type": "output" if save_image else "temp",
|
|
"format": format,
|
|
}
|
|
]
|
|
return {"ui": {"videos": previews}}
|
|
|
|
|
|
class LoadVideo:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_dir = os.path.join(folder_paths.get_input_directory(), "video")
|
|
if not os.path.exists(input_dir):
|
|
os.makedirs(input_dir, exist_ok=True)
|
|
|
|
files = [f"video/{f}" for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
|
|
|
|
return {
|
|
"required": {
|
|
"video": (sorted(files), {"video_upload": True}),
|
|
},
|
|
"optional": {
|
|
"frame_start": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF, "step": 1}),
|
|
"frame_limit": ("INT", {"default": 16, "min": 1, "max": 10240, "step": 1}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "Animate Diff/Utils"
|
|
RETURN_TYPES = ("IMAGE", "INT")
|
|
RETURN_NAMES = ("frames", "frame_count")
|
|
FUNCTION = "load"
|
|
|
|
def load_gif(self, gif_path: str, frame_start: int, frame_limit: int):
|
|
image = Image.open(gif_path)
|
|
frames = []
|
|
|
|
for i, frame in enumerate(ImageSequence.Iterator(image)):
|
|
if i < frame_start:
|
|
continue
|
|
elif i >= frame_start + frame_limit:
|
|
break
|
|
else:
|
|
frames.append(pil2tensor(frame.copy().convert("RGB")))
|
|
|
|
return frames
|
|
|
|
def load_video(self, video_path, frame_start: int, frame_limit: int):
|
|
ensure_opencv()
|
|
import cv2
|
|
|
|
video = cv2.VideoCapture(video_path)
|
|
video.set(cv2.CAP_PROP_POS_FRAMES, frame_start)
|
|
|
|
frames = []
|
|
for i in range(frame_limit):
|
|
# Read the next frame
|
|
ret, frame = video.read()
|
|
if ret:
|
|
# Convert the frame to RGB (OpenCV uses BGR)
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
# Convert the NumPy array to a PIL image and append to list
|
|
frames.append(pil2tensor(Image.fromarray(frame)))
|
|
else:
|
|
break
|
|
|
|
video.release()
|
|
|
|
return frames
|
|
|
|
def load(self, video: str, frame_start=0, frame_limit=16):
|
|
video_path = folder_paths.get_annotated_filepath(video)
|
|
(_, ext) = os.path.splitext(video_path)
|
|
|
|
if ext.lower() in {".gif", ".webp"}:
|
|
frames = self.load_gif(video_path, frame_start, frame_limit)
|
|
elif ext.lower() in {".webp", ".mp4", ".mov", ".avi", ".webm"}:
|
|
frames = self.load_video(video_path, frame_start, frame_limit)
|
|
else:
|
|
raise ValueError(f"Unsupported video format: {ext}")
|
|
|
|
return (torch.cat(frames, dim=0), len(frames))
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, image, *args, **kwargs):
|
|
image_path = folder_paths.get_annotated_filepath(image)
|
|
m = hashlib.sha256()
|
|
with open(image_path, "rb") as f:
|
|
m.update(f.read())
|
|
return m.digest().hex()
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(s, video, *args, **kwargs):
|
|
if not folder_paths.exists_annotated_filepath(video):
|
|
return "Invalid video file: {}".format(video)
|
|
|
|
return True
|
|
|
|
|
|
class ImageSizeAndBatchSize:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "Animate Diff/Utils"
|
|
RETURN_TYPES = ("INT", "INT", "INT")
|
|
RETURN_NAMES = ("width", "height", "batch_size")
|
|
FUNCTION = "batch_size"
|
|
|
|
def batch_size(self, image: Tensor):
|
|
(batch_size, height, width) = image.shape[0:3]
|
|
return (width, height, batch_size)
|
|
|
|
|
|
class ImageChunking:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"chunk_size": ("INT", {"default": 16, "min": 1, "max": 1024, "step": 1}),
|
|
"allow_remainder": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "Animate Diff/Utils"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
OUTPUT_IS_LIST = (True,)
|
|
FUNCTION = "chunk"
|
|
|
|
def chunk(self, images: Tensor, chunk_size: int, allow_remainder: bool):
|
|
# Check if tensor is divisible into chunks of chunk_size
|
|
if images.shape[0] % chunk_size != 0 and not allow_remainder:
|
|
raise ValueError("Tensor's first dimension is not divisible by chunk size")
|
|
|
|
# Use torch.chunk to divide the tensor
|
|
chunk_count = images.shape[0] // chunk_size + images.shape[0] % chunk_size
|
|
|
|
print("chunk_count", chunk_count)
|
|
chunks = torch.chunk(images, chunk_count, dim=0)
|
|
|
|
return (list(chunks),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AnimateDiffModuleLoader": AnimateDiffModuleLoader,
|
|
"AnimateDiffLoraLoader": AnimateDiffLoraLoader,
|
|
"AnimateDiffCombine": AnimateDiffCombine,
|
|
"AnimateDiffSampler": AnimateDiffSampler,
|
|
"AnimateDiffSlidingWindowOptions": AnimateDiffSlidingWindowOptions,
|
|
"LoadVideo": LoadVideo,
|
|
"ImageSizeAndBatchSize": ImageSizeAndBatchSize,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AnimateDiffModuleLoader": "Animate Diff Module Loader",
|
|
"AnimateDiffLoraLoader": "Animate Diff Lora Loader",
|
|
"AnimateDiffSampler": "Animate Diff Sampler",
|
|
"AnimateDiffSlidingWindowOptions": "Sliding Window Options",
|
|
"AnimateDiffCombine": "Animate Diff Combine",
|
|
"LoadVideo": "Load Video",
|
|
"ImageSizeAndBatchSize": "Get Image Size + Batch Size",
|
|
}
|