From 270e351bde6205b0762759a222ab86d2a751c64b Mon Sep 17 00:00:00 2001
From: wzds2015 <wenliang.zhao.nyu@gmail.com>
Date: Fri, 22 Mar 2024 11:21:37 +0000
Subject: [PATCH] fix video writing issue

---
 scripts/sampling/simple_video_sample.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/scripts/sampling/simple_video_sample.py b/scripts/sampling/simple_video_sample.py
index 29a8b8581..ace2c2137 100644
--- a/scripts/sampling/simple_video_sample.py
+++ b/scripts/sampling/simple_video_sample.py
@@ -15,10 +15,11 @@
 from omegaconf import OmegaConf
 from PIL import Image
 from rembg import remove
+from torchvision.transforms import ToTensor
+
 from scripts.util.detection.nsfw_and_watermark_dectection import DeepFloydDataFiltering
 from sgm.inference.helpers import embed_watermark
 from sgm.util import default, instantiate_from_config
-from torchvision.transforms import ToTensor
 
 
 def sample(
@@ -271,7 +272,16 @@ def denoiser(input, sigma, c):
                     .astype(np.uint8)
                 )
                 video_path = os.path.join(output_folder, f"{base_count:06d}.mp4")
-                imageio.mimwrite(video_path, vid)
+                frame0 = vid[0, :, :, :].squeeze()
+                out = cv2.VideoWriter(
+                    video_path,
+                    cv2.VideoWriter_fourcc(*"MP4V"),
+                    20.0,
+                    (frame0.shape[1], frame0.shape[0]),
+                )
+                for frame in vid:
+                    out.write(frame[:, :, ::-1])
+                out.release()
 
 
 def get_unique_embedder_keys_from_conditioner(conditioner):