.

options/Video_model/Model.py

@@ -1,4 +1,4 @@
-import torch,os,cv2
+import torch,os,imageio
 from diffusers import StableVideoDiffusionPipeline
 from diffusers.utils import load_image, export_to_video
 from PIL import Image
@@ -7,6 +7,12 @@ login(token=os.getenv("TOKEN"))
 # Check if CUDA (GPU) is available, otherwise use CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+def save_video(frames, save_path, fps, quality=9):
+    writer = imageio.get_writer(save_path, fps=fps, quality=quality)
+    for frame in frames:
+        frame = np.array(frame)
+        writer.append_data(frame)
+    writer.close()
 
 # Function to generate the video
 def Video(image):
@@ -36,6 +42,6 @@ def Video(image):
     frames = pipeline(image, decode_chunk_size=8, generator=generator).frames[0]
     path="generated.mp4"
     # Export the frames to a video file
-    export_to_video(frames, path, fps=7)
-    vid=cv2.VideoCapture(path)
-    return vid
+    save_video(frames, path, fps=7)
+    # vid=cv2.VideoCapture(path)
+    return path