diff --git a/scripts/animatediff.py b/scripts/animatediff.py index 5267421..439a011 100644 --- a/scripts/animatediff.py +++ b/scripts/animatediff.py @@ -127,6 +127,45 @@ def on_ui_settings(): section=section ) ) + shared.opts.add_option( + key="animatediff_mp4_crf", + info=shared.OptionInfo( + default=23, + label="MP4 Quality (CRF)", + component=gr.Slider, + component_args={ + "minimum": 0, + "maximum": 51, + "step": 1}, + section=section + ) + .link("docs", "https://trac.ffmpeg.org/wiki/Encode/H.264#crf") + .info("17 for best quality, up to 28 for smaller size") + ) + shared.opts.add_option( + key="animatediff_mp4_preset", + info=shared.OptionInfo( + default="", + label="MP4 Encoding Preset", + component=gr.Dropdown, + component_args={"choices": ["", 'veryslow', 'slower', 'slow', 'medium', 'fast', 'faster', 'veryfast', 'superfast', 'ultrafast']}, + section=section, + ) + .link("docs", "https://trac.ffmpeg.org/wiki/Encode/H.264#Preset") + .info("encoding speed, use the slowest you can tolerate") + ) + shared.opts.add_option( + key="animatediff_mp4_tune", + info=shared.OptionInfo( + default="", + label="MP4 Tune encoding for content type", + component=gr.Dropdown, + component_args={"choices": ["", "film", "animation", "grain"]}, + section=section + ) + .link("docs", "https://trac.ffmpeg.org/wiki/Encode/H.264#Tune") + .info("optimize for specific content types") + ) shared.opts.add_option( "animatediff_webp_quality", shared.OptionInfo( diff --git a/scripts/animatediff_output.py b/scripts/animatediff_output.py index 061fc31..f759717 100644 --- a/scripts/animatediff_output.py +++ b/scripts/animatediff_output.py @@ -233,14 +233,31 @@ class AnimateDiffOutput: except ImportError: from launch import run_pip run_pip( - "install imageio[pyav]", - "sd-webui-animatediff MP4 save requirement: imageio[pyav]", + "install pyav", + "sd-webui-animatediff MP4 save requirement: PyAV", ) - with imageio.imopen(video_path_mp4, 'w', plugin='pyav') as file: - if use_infotext: - file.container_metadata["Comment"] = infotext - logger.info(f"Saving {video_path_mp4}") - file.write(video_array, codec='h264', fps=params.fps) + import av + options = { + "crf": str(shared.opts.data.get("animatediff_mp4_crf", 23)) + } + preset = shared.opts.data.get("animatediff_mp4_preset", "") + if preset != "": options["preset"] = preset + tune = shared.opts.data.get("animatediff_mp4_tune", "") + if tune != "": options["tune"] = tune + output = av.open(video_path_mp4, "w") + logger.info(f"Saving {video_path_mp4}") + if use_infotext: + output.metadata["Comment"] = infotext + stream = output.add_stream('libx264', params.fps, options=options) + stream.width = frame_list[0].width + stream.height = frame_list[0].height + for img in video_array: + frame = av.VideoFrame.from_ndarray(img) + packet = stream.encode(frame) + output.mux(packet) + packet = stream.encode(None) + output.mux(packet) + output.close() if "TXT" in params.format and res.images[index].info is not None: video_path_txt = str(video_path_prefix) + ".txt"