Update app.py

app.py

@@ -2,104 +2,51 @@ import gradio as gr
 import torch
 from diffusers import I2VGenXLPipeline
 from diffusers.utils import export_to_gif, load_image
-import tempfile
 import spaces
 
-@spaces.GPU
-def initialize_pipeline():
-    # Initialize the pipeline without CUDA support
-    pipeline = I2VGenXLPipeline.from_pretrained("ali-vilab/i2vgen-xl", torch_dtype=torch.float16, variant="fp16")
-    return pipeline
-
-import gradio as gr
-import torch
-from diffusers import I2VGenXLPipeline
-from diffusers.utils import export_to_gif, load_image
-import tempfile
-import spaces
+# Initialize the pipeline
+pipeline = I2VGenXLPipeline.from_pretrained("ali-vilab/i2vgen-xl", torch_dtype=torch.float16, variant="fp16")
+pipeline.enable_model_cpu_offload()
 
 @spaces.GPU
-def initialize_pipeline():
-    # Initialize the pipeline without CUDA support
-    pipeline = I2VGenXLPipeline.from_pretrained("ali-vilab/i2vgen-xl", torch_dtype=torch.float16, variant="fp16")
-    return pipeline
-
-def generate_gif(prompt, image, negative_prompt, num_inference_steps, guidance_scale, seed):
-    # Initialize the pipeline within the function
-    pipeline = initialize_pipeline()
+def generate_gif(image, prompt, negative_prompt, num_inference_steps, guidance_scale, seed):
+    # Load the image
+    image = load_image(image).convert("RGB")
 
     # Set the generator seed
-    generator = torch.Generator().manual_seed(seed)
-
-    # Check if an image is provided
-    if image is not None and image != "":  # Ensure image is not empty
-        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]
+    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]
 
     # Export to GIF
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".gif") as tmp_gif:
-        gif_path = tmp_gif.name
-        export_to_gif(frames, gif_path)
+    gif_path = "i2v.gif"
+    export_to_gif(frames, gif_path)
 
     return gif_path
 
-# Create the Gradio interface with tabs
-with gr.Blocks() as demo:
-    with gr.Tabs():
-        with gr.TabItem("Generate from Text"):
-            with gr.Row():
-                with gr.Column():
-                    text_prompt = gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt")
-                    text_negative_prompt = gr.Textbox(lines=2, placeholder="Enter your negative prompt here...", label="Negative Prompt")
-                    text_num_inference_steps = gr.Slider(1, 100, step=1, value=50, label="Number of Inference Steps")
-                    text_guidance_scale = gr.Slider(1, 20, step=0.1, value=9.0, label="Guidance Scale")
-                    text_seed = gr.Number(label="Seed", value=8888)
-                    text_generate_button = gr.Button("Generate GIF")
-
-                with gr.Column():
-                    text_output_video = gr.Video(label="Generated GIF")
-
-            # When generating from text, pass an empty string as the image input
-            text_generate_button.click(
-                fn=generate_gif,
-                inputs=[text_prompt, "", text_negative_prompt, text_num_inference_steps, text_guidance_scale, text_seed],
-                outputs=text_output_video
-            )
-
-        with gr.TabItem("Generate from Image"):
-            with gr.Row():
-                with gr.Column():
-                    image_prompt = gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt")
-                    image_input = gr.Image(type="filepath", label="Input Image")
-                    image_negative_prompt = gr.Textbox(lines=2, placeholder="Enter your negative prompt here...", label="Negative Prompt")
-                    image_num_inference_steps = gr.Slider(1, 100, step=1, value=50, label="Number of Inference Steps")
-                    image_guidance_scale = gr.Slider(1, 20, step=0.1, value=9.0, label="Guidance Scale")
-                    image_seed = gr.Number(label="Seed", value=8888)
-                    image_generate_button = gr.Button("Generate GIF")
-
-                with gr.Column():
-                    image_output_video = gr.Video(label="Generated GIF")
-
-            image_generate_button.click(
-                fn=generate_gif,
-                inputs=[image_prompt, image_input, image_negative_prompt, image_num_inference_steps, image_guidance_scale, image_seed],
-                outputs=image_output_video
-            )
+# 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, 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.File(label="Generated GIF"),
+    title="I2VGen-XL GIF Generator",
+    description="Generate a GIF from an image and a prompt using the I2VGen-XL model."
+)
 
 # Launch the interface
-demo.launch()
+iface.launch()