limit video to 3 seconds

app.py

@@ -14,19 +14,26 @@ pipe_xl.scheduler = DPMSolverMultistepScheduler.from_config(pipe_xl.scheduler.co
 pipe_xl.enable_model_cpu_offload()
 pipe_xl.to("cuda")
 
-def convert_mp4_to_frames(video_path):
+def convert_mp4_to_frames(video_path, duration=3):
     # Read the video file
     video = cv2.VideoCapture(video_path)
 
+    # Get the frames per second (fps) of the video
+    fps = video.get(cv2.CAP_PROP_FPS)
+
+    # Calculate the number of frames to extract
+    num_frames = int(fps * duration)
+
     frames = []
+    frame_count = 0
     
     # Iterate through each frame
     while True:
         # Read a frame
         ret, frame = video.read()
         
-        # If the frame was not successfully read, then we have reached the end of the video
-        if not ret:
+        # If the frame was not successfully read or we have reached the desired duration, break the loop
+        if not ret or frame_count == num_frames:
             break
         
         # Convert BGR to RGB
@@ -35,6 +42,8 @@ def convert_mp4_to_frames(video_path):
         # Append the frame to the list of frames
         frames.append(frame)
 
+        frame_count += 1
+
     # Release the video object
     video.release()
 
@@ -45,7 +54,7 @@ def convert_mp4_to_frames(video_path):
 
 def infer(prompt, video_in):
 
-    video = convert_mp4_to_frames(video_in)
+    video = convert_mp4_to_frames(video_in, duration=3)
     video_resized = [Image.fromarray(frame).resize((1024, 576)) for frame in video]
     video_frames = pipe_xl(prompt, video=video_resized, strength=0.6).frames
     video_path = export_to_video(video_frames, output_video_path="xl_result.mp4")