Update app.py

app.py

@@ -2,51 +2,63 @@ import gradio as gr
 import torch
 from diffusers import I2VGenXLPipeline
 from diffusers.utils import export_to_gif, load_image
-import spaces
+import tempfile
 
-# Initialize the pipeline
-pipeline = I2VGenXLPipeline.from_pretrained("ali-vilab/i2vgen-xl", torch_dtype=torch.float16, variant="fp16")
-pipeline.enable_model_cpu_offload()
+# Check if CUDA is available and set the device
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-@spaces.GPU(duration=240)
-def generate_gif(image, prompt, negative_prompt, num_inference_steps, guidance_scale, seed):
-    # Load the image
-    image = load_image(image).convert("RGB")
+# Initialize the pipeline with CUDA support
+pipeline = I2VGenXLPipeline.from_pretrained("ali-vilab/i2vgen-xl", torch_dtype=torch.float16, variant="fp16")
+pipeline.to(device)
 
+def generate_gif(prompt, image, negative_prompt, num_inference_steps, guidance_scale, seed):
     # Set the generator seed
-    generator = torch.manual_seed(seed)
-
-    # Generate the frames
-    frames = pipeline(
-        prompt=prompt,
-        image=image,
-        num_inference_steps=num_inference_steps,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        generator=generator
-    ).frames[0]
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    # Check if an image is provided
+    if image is not None:
+        image = load_image(image).convert("RGB")
+        frames = pipeline(
+            prompt=prompt,
+            image=image,
+            num_inference_steps=num_inference_steps,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            generator=generator
+        ).frames[0]
+    else:
+        frames = pipeline(
+            prompt=prompt,
+            num_inference_steps=num_inference_steps,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            generator=generator
+        ).frames[0]
 
     # Export to GIF
-    gif_path = "i2v.gif"
-    export_to_gif(frames, gif_path)
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".gif") as tmp_gif:
+        gif_path = tmp_gif.name
+        export_to_gif(frames, gif_path)
 
     return gif_path
 
-# Create the Gradio interface
-iface = gr.Interface(
-    fn=generate_gif,
-    inputs=[
-        gr.Image(type="filepath", label="Input Image"),
-        gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt"),
-        gr.Textbox(lines=2, placeholder="Enter your negative prompt here...", label="Negative Prompt"),
-        gr.Slider(1, 50, step=1, value=15, label="Number of Inference Steps"),
-        gr.Slider(1, 10, 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 an image and a prompt using the I2VGen-XL model."
-)
+# Create the Gradio interface with tabs
+with gr.Tabs() as demo:
+    with gr.TabItem("Generate from Text or Image"):
+        interface = gr.Interface(
+            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."
+        )
 
 # Launch the interface
-iface.launch()
+demo.launch()