Use mediapy (#219)

Signed-off-by: Kunjan patel <kunjanp@google.com>

Author: coolkp

requirements.txt

@@ -9,6 +9,7 @@ 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

requirements_with_jax_ai_image.txt

@@ -3,6 +3,7 @@
 --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

src/maxdiffusion/generate_wan.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from typing import Sequence
 import jax
 import time
@@ -49,9 +50,11 @@ def run(config, pipeline=None, filename_prefix=""):
 
   print("compile time: ", (time.perf_counter() - s0))
   saved_video_path = []
-  for i in range(len(videos)):
+  for i, video in enumerate(videos):
     video_path = f"{filename_prefix}wan_output_{config.seed}_{i}.mp4"
-    export_to_video(videos[i], video_path, fps=config.fps)
+    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)
     saved_video_path.append(video_path)
 
   s0 = time.perf_counter()

src/maxdiffusion/utils/export_utils.py

@@ -5,6 +5,7 @@
 from contextlib import contextmanager
 from typing import List, Optional, Union
 
+import mediapy
 import numpy as np
 
 import PIL.Image
@@ -144,65 +145,8 @@ def export_to_video(
     quality: float = 5.0,
     bitrate: Optional[int] = None,
     macro_block_size: Optional[int] = 16,
-) -> 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
+):
+  mediapy.write_video(path = output_video_path,
+                      images=video_frames,
+                      fps=fps,
+                      )