Update Video Loading/Export to use imageio (#9094)

* update

* update

* update

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

Author: 2 people

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 BACKENDS_MAPPING, is_imageio_available, is_opencv_available
 from .logging import get_logger
 
 
@@ -112,9 +112,9 @@ def export_to_obj(mesh, output_obj_path: str = None):
         f.writelines("\n".join(combined_data))
 
 
-def export_to_video(
+def _legacy_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
     else:
@@ -134,4 +134,51 @@ def export_to_video(
     for i in range(len(video_frames)):
         img = cv2.cvtColor(video_frames[i], cv2.COLOR_RGB2BGR)
         video_writer.write(img)
+
+    return output_video_path
+
+
+def export_to_video(
+    video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], output_video_path: str = None, fps: int = 10
+) -> str:
+    # 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) 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
@@ -575,6 +588,11 @@ def is_sentencepiece_available():
 {0} requires the bitsandbytes library but it was not found in your environment. You can install it with pip: `pip install bitsandbytes`
 """
 
+# docstyle-ignore
+IMAGEIO_IMPORT_ERROR = """
+{0} requires the imageio library and ffmpeg but it was not found in your environment. You can install it with pip: `pip install imageio imageio-ffmpeg`
+"""
+
 BACKENDS_MAPPING = OrderedDict(
     [
         ("bs4", (is_bs4_available, BS4_IMPORT_ERROR)),
@@ -599,6 +617,7 @@ def is_sentencepiece_available():
         ("safetensors", (is_safetensors_available, SAFETENSORS_IMPORT_ERROR)),
         ("bitsandbytes", (is_bitsandbytes_available, BITSANDBYTES_IMPORT_ERROR)),
         ("sentencepiece", (is_sentencepiece_available, SENTENCEPIECE_IMPORT_ERROR)),
+        ("imageio", (is_imageio_available, IMAGEIO_IMPORT_ERROR)),
     ]
 )
 

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 BACKENDS_MAPPING, 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,22 @@ 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"))
+            raise ImportError(BACKENDS_MAPPING["imageio"][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()
+        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)