Update app.py

app.py

@@ -3,15 +3,20 @@ import torch
 from diffusers import I2VGenXLPipeline
 from diffusers.utils import export_to_gif, load_image
 import tempfile
+import spaces
 
+# Function to initialize the pipeline with CUDA support
+@spaces.GPU
 def initialize_pipeline():
+    # Check if CUDA is available and set the device
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Initialize the pipeline with CUDA support
-pipeline = I2VGenXLPipeline.from_pretrained("ali-vilab/i2vgen-xl", torch_dtype=torch.float16, variant="fp16")
-pipeline.to(device)
+    # Initialize the pipeline with CUDA support
+    pipeline = I2VGenXLPipeline.from_pretrained("ali-vilab/i2vgen-xl", torch_dtype=torch.float16, variant="fp16")
+    pipeline.to(device)
+    return pipeline, device
 
-def generate_gif(prompt, image, negative_prompt, num_inference_steps, guidance_scale, seed):
+def generate_gif(pipeline, device, prompt, image, negative_prompt, num_inference_steps, guidance_scale, seed):
     # Set the generator seed
     generator = torch.Generator(device=device).manual_seed(seed)
 
@@ -43,22 +48,28 @@ def generate_gif(prompt, image, negative_prompt, num_inference_steps, guidance_s
     return gif_path
 
 # Create the Gradio interface with tabs
-with gr.Tabs() as demo:
+with gr.Blocks() as demo:
+    pipeline, device = initialize_pipeline()
+
     with gr.TabItem("Generate from Text or Image"):
-        interface = gr.Interface(
+        with gr.Row():
+            with gr.Column():
+                prompt = gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt")
+                image = gr.Image(type="filepath", label="Input Image (optional)")
+                negative_prompt = gr.Textbox(lines=2, placeholder="Enter your negative prompt here...", label="Negative Prompt")
+                num_inference_steps = gr.Slider(1, 100, step=1, value=50, label="Number of Inference Steps")
+                guidance_scale = gr.Slider(1, 20, step=0.1, value=9.0, label="Guidance Scale")
+                seed = gr.Number(label="Seed", value=8888)
+                generate_button = gr.Button("Generate GIF")
+
+            with gr.Column():
+                output_video = gr.Video(label="Generated GIF")
+
+        generate_button.click(
             fn=generate_gif,
-            inputs=[
-                gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt"),
-                gr.Image(type="filepath", label="Input Image (optional)"),
-                gr.Textbox(lines=2, placeholder="Enter your negative prompt here...", label="Negative Prompt"),
-                gr.Slider(1, 100, step=1, value=50, label="Number of Inference Steps"),
-                gr.Slider(1, 20, step=0.1, value=9.0, label="Guidance Scale"),
-                gr.Number(label="Seed", value=8888)
-            ],
-            outputs=gr.Video(label="Generated GIF"),
-            title="I2VGen-XL GIF Generator",
-            description="Generate a GIF from a text prompt and/or an image using the I2VGen-XL model."
+            inputs=[pipeline, device, prompt, image, negative_prompt, num_inference_steps, guidance_scale, seed],
+            outputs=output_video
         )
 
 # Launch the interface
-demo.launch()
+demo.launch()