mrcuddle commited on
Commit
443ba67
·
verified ·
1 Parent(s): dea9c7f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +13 -6
app.py CHANGED
@@ -2,18 +2,23 @@ import torch
2
  import gradio as gr
3
  from diffusers import StableVideoDiffusionPipeline
4
  from diffusers.utils import load_image, export_to_video
 
5
  import spaces
6
 
7
  # Check if GPU is available
8
  device = "cuda" if torch.cuda.is_available() else "cpu"
9
 
10
- # Load the pipeline
11
  pipeline = StableVideoDiffusionPipeline.from_pretrained(
12
  "stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16"
13
  )
14
  pipeline.to(device)
15
 
16
- @spaces.GPU(duration=120)
 
 
 
 
17
  def generate_video(image_path, seed):
18
  # Load and preprocess the image
19
  image = load_image(image_path)
@@ -40,11 +45,13 @@ iface = gr.Interface(
40
  ],
41
  outputs=gr.Video(label="Generated Video"),
42
  title="Stable Video Diffusion",
 
43
  examples=[
44
- ["image.jpeg"],
45
- ["generated.mp4"]
46
  ],
47
- description="Generate a video from an uploaded image using Stable Video Diffusion.",
48
  )
 
49
  # Launch the interface
50
- iface.launch()
 
 
2
  import gradio as gr
3
  from diffusers import StableVideoDiffusionPipeline
4
  from diffusers.utils import load_image, export_to_video
5
+ import os
6
  import spaces
7
 
8
  # Check if GPU is available
9
  device = "cuda" if torch.cuda.is_available() else "cpu"
10
 
11
+ # Load the pipeline with optimizations
12
  pipeline = StableVideoDiffusionPipeline.from_pretrained(
13
  "stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16"
14
  )
15
  pipeline.to(device)
16
 
17
+ # Enable forward chunking to reduce memory usage
18
+ pipeline.unet.enable_forward_chunking(chunk_size=1)
19
+
20
+ @spaces.GPU
21
+ # Define the video generation function
22
  def generate_video(image_path, seed):
23
  # Load and preprocess the image
24
  image = load_image(image_path)
 
45
  ],
46
  outputs=gr.Video(label="Generated Video"),
47
  title="Stable Video Diffusion",
48
+ description="Generate a video from an uploaded image using Stable Video Diffusion.",
49
  examples=[
50
+ ["example.jpeg", 666666]
 
51
  ],
52
+ cache_examples=True # Enable caching of examples
53
  )
54
+
55
  # Launch the interface
56
+ if __name__ == "__main__":
57
+ iface.launch()