JiayiGuo821 commited on
Commit
8964f12
β€’
1 Parent(s): cdbeb67
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -26,6 +26,7 @@ from diffusers import DDIMScheduler
26
  from app_utils import auto_dropdown
27
 
28
  from huggingface_hub import hf_hub_download
 
29
 
30
  version = "Smooth Diffusion Demo v1.0"
31
  refresh_symbol = "\U0001f504" # πŸ”„
@@ -248,7 +249,7 @@ def t2i_core(model, xt, emb, nemb, step=30, cfg_scale=7.5, return_list=False):
248
  ########
249
  # main #
250
  ########
251
-
252
  class wrapper(object):
253
  def __init__(self,
254
  fp16=False,
@@ -256,8 +257,7 @@ class wrapper(object):
256
  tag_lora=None,
257
  tag_scheduler=None,):
258
 
259
- # self.device = "cuda"
260
- self.device = "cpu"
261
  if fp16:
262
  self.torch_dtype = torch.float16
263
  else:
 
26
  from app_utils import auto_dropdown
27
 
28
  from huggingface_hub import hf_hub_download
29
+ import spaces
30
 
31
  version = "Smooth Diffusion Demo v1.0"
32
  refresh_symbol = "\U0001f504" # πŸ”„
 
249
  ########
250
  # main #
251
  ########
252
+ @spaces.GPU()
253
  class wrapper(object):
254
  def __init__(self,
255
  fp16=False,
 
257
  tag_lora=None,
258
  tag_scheduler=None,):
259
 
260
+ self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
261
  if fp16:
262
  self.torch_dtype = torch.float16
263
  else: