diff --git a/animatediff/model_utils.py b/animatediff/model_utils.py index 74b487b..21de6a9 100644 --- a/animatediff/model_utils.py +++ b/animatediff/model_utils.py @@ -46,6 +46,13 @@ folder_names_and_paths[Folders.MODELS] = ([MODEL_DIR], folder_paths.supported_pt filename_list_cache = {} +#Register video_formats folder +folder_paths.folder_names_and_paths["video_formats"] = ( + [ + os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "video_formats"), + ], + [".json"] +) def get_filename_list_(folder_name): global folder_names_and_paths diff --git a/animatediff/nodes.py b/animatediff/nodes.py index 7cee39f..7b46f2f 100644 --- a/animatediff/nodes.py +++ b/animatediff/nodes.py @@ -422,6 +422,13 @@ class AnimateDiffUnload: class AnimateDiffCombine: @classmethod def INPUT_TYPES(s): + ffmpeg_path = shutil.which("ffmpeg") + #Hide ffmpeg formats if ffmpeg isn't available + if ffmpeg_path is not None: + ffmpeg_formats = ["video/"+x[:-5] for x in folder_paths.get_filename_list("video_formats")] + else: + ffmpeg_formats = [] + logger.warning("ffmpeg could not be found. Outputs that require it have been disabled") return { "required": { "images": ("IMAGE",), @@ -431,7 +438,7 @@ class AnimateDiffCombine: ), "loop_count": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), "filename_prefix": ("STRING", {"default": "AnimateDiff"}), - "format": (["image/gif", "image/webp", "video/webm"],), + "format": (["image/gif", "image/webp"] + ffmpeg_formats,), "pingpong": ("BOOLEAN", {"default": False}), "save_image": ("BOOLEAN", {"default": True}), }, @@ -465,7 +472,7 @@ class AnimateDiffCombine: img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) frames.append(img) - # save image + # get output information output_dir = ( folder_paths.get_output_directory() if save_image @@ -497,11 +504,11 @@ class AnimateDiffCombine: if pingpong: frames = frames + frames[-2:0:-1] - # save gif format_type, format_ext = format.split("/") file = f"{filename}_{counter:05}_.{format_ext}" file_path = os.path.join(full_output_folder, file) if format_type == "image": + # Use pillow directly to save an animated image frames[0].save( file_path, format=format_ext.upper(), @@ -512,18 +519,26 @@ class AnimateDiffCombine: compress_level=4, ) else: - # save webm - import shutil - import subprocess - + # Use ffmpeg to save a video ffmpeg_path = shutil.which("ffmpeg") if ffmpeg_path is None: + #Should never be reachable raise ProcessLookupError("Could not find ffmpeg") - dimensions = f"{frames[0].width}x{frames[0].height}" - args = [ffmpeg_path, "-v", "panic", "-n", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", - dimensions, "-r", str(frame_rate), "-i", "-", "-pix_fmt", "yuv420p", file_path] - with subprocess.Popen(args, stdin=subprocess.PIPE) as proc: + video_format_path = folder_paths.get_full_path("video_formats", 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.copy() + 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()) @@ -535,7 +550,6 @@ class AnimateDiffCombine: "format": format, } ] - print(previews) return {"ui": {"gifs": previews}} class CheckpointLoaderSimpleWithNoiseSelect: diff --git a/video_formats/av1-webm.json b/video_formats/av1-webm.json new file mode 100644 index 0000000..137ad18 --- /dev/null +++ b/video_formats/av1-webm.json @@ -0,0 +1,10 @@ +{ + "main_pass": + [ + "-n", "-c:v", "libsvtav1", + "-pix_fmt", "yuv420p10le", + "-crf", "23" + ], + "extension": "webm", + "environment": {"SVT_LOG": "1"} +} diff --git a/video_formats/h264-mp4.json b/video_formats/h264-mp4.json new file mode 100644 index 0000000..6b50b12 --- /dev/null +++ b/video_formats/h264-mp4.json @@ -0,0 +1,9 @@ +{ + "main_pass": + [ + "-n", "-c:v", "libx264", + "-pix_fmt", "yuv420p", + "-crf", "19" + ], + "extension": "mp4" +} diff --git a/video_formats/h265-mp4.json b/video_formats/h265-mp4.json new file mode 100644 index 0000000..a8b677b --- /dev/null +++ b/video_formats/h265-mp4.json @@ -0,0 +1,11 @@ +{ + "main_pass": + [ + "-n", "-c:v", "libx265", + "-pix_fmt", "yuv420p10le", + "-preset", "medium", + "-crf", "22", + "-x265-params", "log-level=quiet" + ], + "extension": "mp4" +} diff --git a/video_formats/webm.json b/video_formats/webm.json new file mode 100644 index 0000000..1551e2c --- /dev/null +++ b/video_formats/webm.json @@ -0,0 +1,9 @@ +{ + "main_pass": + [ + "-n", + "-pix_fmt", "yuv420p", + "-crf", "23" + ], + "extension": "webm" +}