opencv instead of imageio

app.py

@@ -28,7 +28,7 @@ import spaces
 from huggingface_hub import hf_hub_download
 import imageio
 
-# import cv2
+import cv2
 
 
 @spaces.GPU
@@ -137,17 +137,17 @@ def do_sample(
                 .numpy()
                 .astype(np.uint8)
             )
-            # write_video(video_path, frames, fps=6)
-            # writer = cv2.VideoWriter(
-            #     video_path,
-            #     cv2.VideoWriter_fourcc(*"MP4V"),
-            #     6,
-            #     (frames.shape[-1], frames.shape[-2]),
-            # )
-            # for fr in frames:
-            #     writer.write(cv2.cvtColor(fr, cv2.COLOR_RGB2BGR))
-            # writer.release()
-            imageio.mimwrite(video_path, frames)
+            write_video(video_path, frames, fps=6)
+            writer = cv2.VideoWriter(
+                video_path,
+                cv2.VideoWriter_fourcc(*"MP4V"),
+                6,
+                (frames.shape[-1], frames.shape[-2]),
+            )
+            for fr in frames:
+                writer.write(cv2.cvtColor(fr, cv2.COLOR_RGB2BGR))
+            writer.release()
+            # imageio.mimwrite(video_path, frames)
 
     return video_path