Revert "Use mediapy"

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

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

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()

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,
-                      )
+) -> 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