Update Video Loading/Export to use imageio

src/diffusers/utils/export_utils.py

@@ -9,7 +9,7 @@
 import PIL.Image
 import PIL.ImageOps
 
-from .import_utils import BACKENDS_MAPPING, is_opencv_available
+from .import_utils import is_imageio_available
 from .logging import get_logger
 
 
@@ -115,10 +115,22 @@ def export_to_obj(mesh, output_obj_path: str = None):
 def export_to_video(
     video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], output_video_path: str = None, fps: int = 10
 ) -> str:
-    if is_opencv_available():
-        import cv2
+    if is_imageio_available():
+        import imageio
     else:
-        raise ImportError(BACKENDS_MAPPING["opencv"][1].format("export_to_video"))
+        raise ImportError(
+            (
+                "`export_to_video` requires imageio and ffmpeg to be installed on your machine. "
+                "Please install via `pip install imageio imageio-ffmpeg`"
+            )
+        )
+    try:
+        imageio.plugins.ffmpeg.get_exe()
+    except AttributeError:
+        raise AttributeError(
+            "`Unable to find an ffmpeg installation on your machine. Please install via `pip install imageio-ffmpeg"
+        )
+
     if output_video_path is None:
         output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name
 
@@ -128,10 +140,8 @@ def export_to_video(
     elif isinstance(video_frames[0], PIL.Image.Image):
         video_frames = [np.array(frame) for frame in video_frames]
 
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    h, w, c = video_frames[0].shape
-    video_writer = cv2.VideoWriter(output_video_path, fourcc, fps=fps, frameSize=(w, h))
-    for i in range(len(video_frames)):
-        img = cv2.cvtColor(video_frames[i], cv2.COLOR_RGB2BGR)
-        video_writer.write(img)
+    with imageio.get_writer(output_video_path, fps=fps) as writer:
+        for frame in video_frames:
+            writer.append_data(frame)
+
     return output_video_path

src/diffusers/utils/import_utils.py

@@ -330,6 +330,15 @@ def is_timm_available():
 
 _is_google_colab = "google.colab" in sys.modules or any(k.startswith("COLAB_") for k in os.environ)
 
+_imageio_available = importlib.util.find_spec("imageio") is not None
+if _imageio_available:
+    try:
+        _imageio_version = importlib_metadata.version("imageio")
+        logger.debug(f"Successfully imported imageio version {_imageio_version}")
+
+    except importlib_metadata.PackageNotFoundError:
+        _imageio_available = False
+
 
 def is_torch_available():
     return _torch_available
@@ -447,6 +456,10 @@ def is_sentencepiece_available():
     return _sentencepiece_available
 
 
+def is_imageio_available():
+    return _imageio_available
+
+
 # docstyle-ignore
 FLAX_IMPORT_ERROR = """
 {0} requires the FLAX library but it was not found in your environment. Checkout the instructions on the

src/diffusers/utils/loading_utils.py

@@ -6,7 +6,7 @@
 import PIL.ImageOps
 import requests
 
-from .import_utils import BACKENDS_MAPPING, is_opencv_available
+from .import_utils import is_imageio_available
 
 
 def load_image(
@@ -81,7 +81,8 @@ def load_video(
 
     if is_url:
         video_data = requests.get(video, stream=True).raw
-        video_path = tempfile.NamedTemporaryFile(suffix=os.path.splitext(video)[1], delete=False).name
+        suffix = os.path.splitext(video)[1] or ".mp4"
+        video_path = tempfile.NamedTemporaryFile(suffix=suffix, delete=False).name
         was_tempfile_created = True
         with open(video_path, "wb") as f:
             f.write(video_data.read())
@@ -99,19 +100,26 @@ def load_video(
             pass
 
     else:
-        if is_opencv_available():
-            import cv2
+        if is_imageio_available():
+            import imageio
         else:
-            raise ImportError(BACKENDS_MAPPING["opencv"][1].format("load_video"))
-
-        video_capture = cv2.VideoCapture(video)
-        success, frame = video_capture.read()
-        while success:
-            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            pil_images.append(PIL.Image.fromarray(frame))
-            success, frame = video_capture.read()
+            raise ImportError(
+                (
+                    "`load_video` requires imageio and ffmpeg to be installed on your machine. "
+                    "Please install via `pip install imageio imageio-ffmpeg`"
+                )
+            )
+        try:
+            imageio.plugins.ffmpeg.get_exe()
+        except AttributeError:
+            raise AttributeError(
+                "`Unable to find an ffmpeg installation on your machine. Please install via `pip install imageio-ffmpeg"
+            )
 
-        video_capture.release()
+        with imageio.get_reader(video) as reader:
+            # Read all frames
+            for frame in reader:
+                pil_images.append(PIL.Image.fromarray(frame))
 
     if was_tempfile_created:
         os.remove(video_path)