Update app.py

app.py

@@ -2,15 +2,19 @@ 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
 
 # Load the pipeline
 pipe = StableVideoDiffusionPipeline.from_pretrained(
     "stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16"
 )
+pipe.to("cuda")
+pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
 pipe.enable_model_cpu_offload()
+pipe.unet.enable_forward_chunking()
 
-def generate_video(image, seed=42, fps=7):
+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))
 
@@ -18,11 +22,11 @@ def generate_video(image, seed=42, fps=7):
     generator = torch.manual_seed(seed)
 
     # Generate the frames
-    frames = pipe(image, decode_chunk_size=8, generator=generator).frames[0]
+    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"
-    imageio.mimwrite(output_path, frames, fps=fps)
+    export_to_video(frames, output_path, fps=fps)
 
     return output_path
 
@@ -32,7 +36,9 @@ iface = gr.Interface(
     inputs=[
         gr.Image(type="pil", label="Upload Image"),
         gr.Number(label="Seed", value=42),
-        gr.Number(label="FPS", value=7)
+        gr.Number(label="FPS", value=7),
+        gr.Number(label="Motion Bucket ID", value=180),
+        gr.Number(label="Noise Aug Strength", value=0.1)
     ],
     outputs=gr.Video(label="Generated Video"),
     title="Stable Video Diffusion",