Li Yan commited on
Commit
0e59027
·
1 Parent(s): 2c5ac74

correct device

Browse files
Files changed (1) hide show
  1. app.py +4 -1
app.py CHANGED
@@ -1,11 +1,14 @@
1
  import gradio as gr
 
2
  from diffusers import DiffusionPipeline, DDIMScheduler
3
 
 
 
4
  scheduler = DDIMScheduler.from_pretrained('li-yan/diffusion-aurora-256')
5
  scheduler.set_timesteps(num_inference_steps=40)
6
 
7
  pipeline = DiffusionPipeline.from_pretrained(
8
- 'li-yan/diffusion-aurora-256', scheduler=scheduler).to("cuda")
9
 
10
  def image_gen(name):
11
  images = pipeline(num_inference_steps=40).images
 
1
  import gradio as gr
2
+ import torch
3
  from diffusers import DiffusionPipeline, DDIMScheduler
4
 
5
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
6
+
7
  scheduler = DDIMScheduler.from_pretrained('li-yan/diffusion-aurora-256')
8
  scheduler.set_timesteps(num_inference_steps=40)
9
 
10
  pipeline = DiffusionPipeline.from_pretrained(
11
+ 'li-yan/diffusion-aurora-256', scheduler=scheduler).to(device)
12
 
13
  def image_gen(name):
14
  images = pipeline(num_inference_steps=40).images