From 582756ee757d6cacbd4e574d579b732d7aa45b7c Mon Sep 17 00:00:00 2001 From: Austin Mroz Date: Mon, 18 Sep 2023 15:52:52 -0500 Subject: [PATCH] Add additional video formats and config system Allows for ffmpeg output settings to be set from a series of json config files stored in video_formats. This allows for higher quality output files with newer codecs (like av1), while still ensuring compatibility with older ffmpeg versions In addition to these changes, the logging of ffmpeg is better handled. Error messages from ffmpeg itself are printed and logging from encoders that was improperly filtered have been individually addressed. While h265 has been included, most browsers will be unable to display the resulting video. --- animatediff/model_utils.py | 7 +++++++ animatediff/nodes.py | 38 +++++++++++++++++++++++++------------ video_formats/av1-webm.json | 10 ++++++++++ video_formats/h264-mp4.json | 9 +++++++++ video_formats/h265-mp4.json | 11 +++++++++++ video_formats/webm.json | 9 +++++++++ 6 files changed, 72 insertions(+), 12 deletions(-) create mode 100644 video_formats/av1-webm.json create mode 100644 video_formats/h264-mp4.json create mode 100644 video_formats/h265-mp4.json create mode 100644 video_formats/webm.json 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" +}