Update app.py

app.py

@@ -2,23 +2,18 @@ import torch
 import gradio as gr
 from diffusers import StableVideoDiffusionPipeline
 from diffusers.utils import load_image, export_to_video
-import os
 import spaces
 
 # Check if GPU is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the pipeline with optimizations
+# Load the pipeline
 pipeline = StableVideoDiffusionPipeline.from_pretrained(
     "stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16"
 )
 pipeline.to(device)
 
-# Enable forward chunking to reduce memory usage
-pipeline.unet.enable_forward_chunking(chunk_size=1)
-
-@spaces.GPU
-# Define the video generation function
+@spaces.GPU(duration=120)
 def generate_video(image_path, seed):
     # Load and preprocess the image
     image = load_image(image_path)
@@ -37,21 +32,28 @@ def generate_video(image_path, seed):
     return output_video_path
 
 # Create the Gradio interface
-iface = gr.Interface(
-    fn=generate_video,
-    inputs=[
-        gr.Image(type="filepath", label="Upload Image"),
-        gr.Number(label="Seed", value=666666)
-    ],
-    outputs=gr.Video(label="Generated Video"),
-    title="Stable Video Diffusion",
-    description="Generate a video from an uploaded image using Stable Video Diffusion.",
-    examples=[
-        ["example.jpeg", 666666]
-    ],
-    cache_examples=True  # Enable caching of examples
-)
+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)
+            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")
+
+    generate_button.click(
+        fn=generate_video,
+        inputs=[image_input, seed_input],
+        outputs=video_output
+    )
 
 # Launch the interface
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()