eggman-poff commited on
Commit
77ceff1
·
verified ·
1 Parent(s): da9f3bc

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -7,6 +7,8 @@ if not hasattr(hf_hub, "model_info"):
7
 
8
  import gradio as gr
9
  import torch
 
 
10
  import tempfile
11
  from diffusers import StableVideoDiffusionPipeline
12
  from diffusers.utils import export_to_video
@@ -14,10 +16,10 @@ from diffusers.utils import export_to_video
14
  # Use the official SVD-XT img2vid-xt model
15
  MODEL = "stabilityai/stable-video-diffusion-img2vid-xt"
16
 
17
- # Load pipeline in half-precision on GPU
18
  pipe = StableVideoDiffusionPipeline.from_pretrained(
19
  MODEL, torch_dtype=torch.float16
20
- ).to("cuda")
21
 
22
  def infer(first_image, last_image, prompt, guidance=7.5, frames=25):
23
  # Generate the in-between frames
 
7
 
8
  import gradio as gr
9
  import torch
10
+ # Determine device: use GPU if available, otherwise CPU
11
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
  import tempfile
13
  from diffusers import StableVideoDiffusionPipeline
14
  from diffusers.utils import export_to_video
 
16
  # Use the official SVD-XT img2vid-xt model
17
  MODEL = "stabilityai/stable-video-diffusion-img2vid-xt"
18
 
19
+ # Load pipeline in half-precision on GPU or CPU
20
  pipe = StableVideoDiffusionPipeline.from_pretrained(
21
  MODEL, torch_dtype=torch.float16
22
+ ).to(device)
23
 
24
  def infer(first_image, last_image, prompt, guidance=7.5, frames=25):
25
  # Generate the in-between frames