Update app.py

app.py

@@ -1,51 +1,34 @@
-import torch
-from PIL import Image
-import imageio
-from diffusers import StableVideoDiffusionPipeline
-from diffusers.utils import load_image, export_to_video
 import gradio as gr
-import spaces
+from diffusers import StableVideoDiffusionPipeline, EulerDiscreteScheduler
+import torch
 
-# Load the pipeline
-pipe = StableVideoDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16"
-)
+# Load the Stable Video Diffusion model
+model_id = "stabilityai/stable-video-diffusion-img2vid-xt"
+pipe = StableVideoDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="main")
+pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
 pipe.to("cuda")
-pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-pipe.enable_model_cpu_offload()
-pipe.unet.enable_forward_chunking()
-
-@spaces.GPU(duration=300)
-def generate_video(image, seed=42, fps=7, motion_bucket_id=180, noise_aug_strength=0.1):
-    # Resize the image
-    image = image.resize((1024, 576))
-
-    # Set the generator seed
-    generator = torch.manual_seed(seed)
-
-    # Generate the frames
-    frames = pipe(image, decode_chunk_size=2, generator=generator, num_frames=25, motion_bucket_id=motion_bucket_id, noise_aug_strength=noise_aug_strength).frames[0]
-
-    # Export the frames to a video
-    output_path = "generated.mp4"
-    export_to_video(frames, output_path, fps=fps)
 
-    return output_path
+def generate_video(image, prompt, num_frames=25, resolution=(576, 1024)):
+    # Generate the video
+    video_frames = pipe(prompt, image=image, num_frames=num_frames, height=resolution[0], width=resolution[1]).frames
+    return video_frames
 
 # Create the Gradio interface
 iface = gr.Interface(
     fn=generate_video,
     inputs=[
-        gr.Image(type="pil", label="Upload Image"),
-        gr.Number(label="Seed", value=42),
-        gr.Number(label="FPS", value=7),
-        gr.Number(label="Motion Bucket ID", value=180),
-        gr.Number(label="Noise Aug Strength", value=0.1)
+        gr.inputs.Image(type="pil", label="Upload Image"),
+        gr.inputs.Textbox(lines=2, placeholder="Enter prompt...", label="Prompt"),
+        gr.inputs.Slider(1, 50, step=1, default=25, label="Number of Frames"),
+        gr.inputs.Number(label="Resolution Height", default=576),
+        gr.inputs.Number(label="Resolution Width", default=1024)
     ],
-    outputs=gr.Video(label="Generated Video"),
-    title="Stable Video Diffusion",
-    description="Generate a video from an uploaded image using Stable Video Diffusion."
+    outputs=gr.outputs.Video(label="Generated Video"),
+    title="Image to Video with Stable Diffusion XT",
+    description="Upload an image and enter a prompt to generate a video."
 )
 
 # Launch the interface
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()
+