cosmicdream commited on
Commit
06ee334
1 Parent(s): 0424af2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -24,8 +24,8 @@ from rudalle.utils import seed_everything
24
  #from ruclip import load as get_ruclip
25
 
26
  # prepare models:
27
- device = 'cuda'
28
- #device = torch.device("cuda:0" if torch.cuda.is_available() else "CPU")
29
  model = get_rudalle_model('Malevich', pretrained=True, fp16=True, device=device)
30
  tokenizer = get_tokenizer()
31
  vae = get_vae(dwt=True).to(device)
 
24
  #from ruclip import load as get_ruclip
25
 
26
  # prepare models:
27
+ #device = 'cuda'
28
+ device = "cuda" if torch.cuda.is_available() else "cpu"
29
  model = get_rudalle_model('Malevich', pretrained=True, fp16=True, device=device)
30
  tokenizer = get_tokenizer()
31
  vae = get_vae(dwt=True).to(device)