cosmicdream commited on
Commit
466917c
1 Parent(s): 1321dbf

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -26,14 +26,15 @@ import ruclip
26
 
27
  # prepare models:
28
  #device = 'cuda'
29
- device = "cuda" if torch.cuda.is_available() else "cpu"
 
30
  #fp16=True
31
  #if device == 'CPU':
32
  # fp16 = False
33
 
34
  model = get_rudalle_model('Malevich', pretrained=True, fp16=True, device=device)
35
  tokenizer = get_tokenizer()
36
- vae = get_vae(dwt=True).to(device)
37
 
38
  # pipeline utils:
39
  #TEST--------
 
26
 
27
  # prepare models:
28
  #device = 'cuda'
29
+ device = "cuda" if torch.cuda.is_available() else "CPU"
30
+ print('Using device:', device)
31
  #fp16=True
32
  #if device == 'CPU':
33
  # fp16 = False
34
 
35
  model = get_rudalle_model('Malevich', pretrained=True, fp16=True, device=device)
36
  tokenizer = get_tokenizer()
37
+ vae = get_vae(dwt=True).to('cuda')
38
 
39
  # pipeline utils:
40
  #TEST--------