diff --git a/requirements.txt b/requirements.txt index 0db2596c7..b66f58234 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,7 +9,6 @@ flax>=0.10.2 optax>=0.2.3 torch>=2.6.0 torchvision>=0.20.1 -mediapy ftfy tensorboard>=2.17.0 tensorboardx>=2.6.2.2 diff --git a/requirements_with_jax_ai_image.txt b/requirements_with_jax_ai_image.txt index 0e52ab7ed..ea42e1cd4 100644 --- a/requirements_with_jax_ai_image.txt +++ b/requirements_with_jax_ai_image.txt @@ -3,7 +3,6 @@ --extra-index-url https://download.pytorch.org/whl/cpu jax>=0.6.2 jaxlib>=0.4.30 -mediapy grain google-cloud-storage>=2.17.0 absl-py diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py index 888b228a2..a9bcf366c 100644 --- a/src/maxdiffusion/generate_wan.py +++ b/src/maxdiffusion/generate_wan.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os from typing import Sequence import jax import time @@ -50,11 +49,9 @@ def run(config, pipeline=None, filename_prefix=""): print("compile time: ", (time.perf_counter() - s0)) saved_video_path = [] - for i, video in enumerate(videos): + for i in range(len(videos)): video_path = f"{filename_prefix}wan_output_{config.seed}_{i}.mp4" - if os.path.exists(f"{config.base_output_dir}"): - video_path = f"{config.base_output_dir}/{video_path}" - export_to_video(video, video_path, fps=config.fps) + export_to_video(videos[i], video_path, fps=config.fps) saved_video_path.append(video_path) s0 = time.perf_counter() diff --git a/src/maxdiffusion/utils/export_utils.py b/src/maxdiffusion/utils/export_utils.py index 353fcf485..5dfa3562f 100644 --- a/src/maxdiffusion/utils/export_utils.py +++ b/src/maxdiffusion/utils/export_utils.py @@ -5,7 +5,6 @@ from contextlib import contextmanager from typing import List, Optional, Union -import mediapy import numpy as np import PIL.Image @@ -145,8 +144,65 @@ def export_to_video( quality: float = 5.0, bitrate: Optional[int] = None, macro_block_size: Optional[int] = 16, -): - mediapy.write_video(path = output_video_path, - images=video_frames, - fps=fps, - ) \ No newline at end of file +) -> str: + """ + quality: + Video output quality. Default is 5. Uses variable bit rate. Highest quality is 10, lowest is 0. Set to None to + prevent variable bitrate flags to FFMPEG so you can manually specify them using output_params instead. + Specifying a fixed bitrate using `bitrate` disables this parameter. + + bitrate: + Set a constant bitrate for the video encoding. Default is None causing `quality` parameter to be used instead. + Better quality videos with smaller file sizes will result from using the `quality` variable bitrate parameter + rather than specifiying a fixed bitrate with this parameter. + + macro_block_size: + Size constraint for video. Width and height, must be divisible by this number. If not divisible by this number + imageio will tell ffmpeg to scale the image up to the next closest size divisible by this number. Most codecs + are compatible with a macroblock size of 16 (default), some can go smaller (4, 8). To disable this automatic + feature set it to None or 1, however be warned many players can't decode videos that are odd in size and some + codecs will produce poor results or fail. See https://en.wikipedia.org/wiki/Macroblock. + """ + # TODO: Dhruv. Remove by Diffusers release 0.33.0 + # Added to prevent breaking existing code + if not is_imageio_available(): + logger.warning( + ( + "It is recommended to use `export_to_video` with `imageio` and `imageio-ffmpeg` as a backend. \n" + "These libraries are not present in your environment. Attempting to use legacy OpenCV backend to export video. \n" + "Support for the OpenCV backend will be deprecated in a future Diffusers version" + ) + ) + return _legacy_export_to_video(video_frames, output_video_path, fps) + + if is_imageio_available(): + import imageio + else: + raise ImportError(BACKENDS_MAPPING["imageio"][1].format("export_to_video")) + + try: + imageio.plugins.ffmpeg.get_exe() + except AttributeError: + raise AttributeError( + ( + "Found an existing imageio backend in your environment. Attempting to export video with imageio. \n" + "Unable to find a compatible ffmpeg installation in your environment to use with imageio. Please install via `pip install imageio-ffmpeg" + ) + ) + + if output_video_path is None: + output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name + + if isinstance(video_frames[0], np.ndarray): + video_frames = [(frame * 255).astype(np.uint8) for frame in video_frames] + + elif isinstance(video_frames[0], PIL.Image.Image): + video_frames = [np.array(frame) for frame in video_frames] + + with imageio.get_writer( + output_video_path, fps=fps, quality=quality, bitrate=bitrate, macro_block_size=macro_block_size + ) as writer: + for frame in video_frames: + writer.append_data(frame) + + return output_video_path