Update app.py

app.py

@@ -14,7 +14,7 @@ pipeline = StableVideoDiffusionPipeline.from_pretrained(
 pipeline.to(device)
 
 @spaces.GPU(duration=120)
-def generate_video(image_path, seed):
+def generate_video(image_path, seed, fps, duration, use_duration):
     # Load and preprocess the image
     image = load_image(image_path)
     image = image.resize((1024, 576))
@@ -25,9 +25,13 @@ def generate_video(image_path, seed):
     # Generate the video frames
     frames = pipeline(image, decode_chunk_size=8, generator=generator).frames[0]
 
+    # Calculate fps if duration is specified
+    if use_duration:
+        fps = len(frames) / duration
+
     # Export the frames to a video file
     output_video_path = "generated.mp4"
-    export_to_video(frames, output_video_path, fps=25)
+    export_to_video(frames, output_video_path, fps=fps)
 
     return output_video_path
 
@@ -35,25 +39,40 @@ def generate_video(image_path, seed):
 with gr.Blocks() as demo:
     gr.Markdown("# Stable Video Diffusion")
     gr.Markdown("Generate a video from an uploaded image using Stable Video Diffusion.")
-    
+
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="filepath", label="Upload Image")
             seed_input = gr.Number(label="Seed", value=666666)
+            use_duration_toggle = gr.Checkbox(label="Use Duration Instead of FPS")
+            fps_input = gr.Number(label="FPS", value=25, minimum=1, maximum=60, visible=True)
+            duration_input = gr.Number(label="Duration (seconds)", value=1, minimum=1, maximum=60, visible=False)
             generate_button = gr.Button("Generate Video")
         with gr.Column():
             video_output = gr.Video(label="Generated Video")
-    
+
     with gr.Row():
-        example_image = gr.Image("example.jpeg", label="Example Image")
-        example_video = gr.Video("generated.mp4", label="Example Video")
+        example_image = gr.Image(value="example.jpeg", label="Example Image")
+        example_video = gr.Video(value="generated.mp4", label="Example Video")
+
+    def toggle_visibility(use_duration):
+        return {
+            fps_input: gr.update(visible=not use_duration),
+            duration_input: gr.update(visible=use_duration)
+        }
+
+    use_duration_toggle.change(
+        fn=toggle_visibility,
+        inputs=use_duration_toggle,
+        outputs=[fps_input, duration_input]
+    )
 
     generate_button.click(
         fn=generate_video,
-        inputs=[image_input, seed_input],
+        inputs=[image_input, seed_input, fps_input, duration_input, use_duration_toggle],
         outputs=video_output
     )
 
 # Launch the interface
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()