pineconeT94 commited on
Commit
35c77e1
·
1 Parent(s): 1e88599

auto detect cuda cpu.

Browse files
Files changed (1) hide show
  1. app.py +9 -4
app.py CHANGED
@@ -13,9 +13,14 @@ exec python tools/webui.py ${DEVICE}
13
 
14
  import os
15
 
16
- CUDA_ENABLED = os.environ.get("CUDA_ENABLED", "true")
17
- DEVICE = ""
18
- if CUDA_ENABLED != "true":
19
- DEVICE = "--device cpu"
 
 
 
 
 
20
 
21
  os.system(f"python tools/webui.py {DEVICE}")
 
13
 
14
  import os
15
 
16
+ import torch
17
+
18
+ if torch.cuda.is_available():
19
+ DEVICE = "cuda"
20
+ else:
21
+ DEVICE = "cpu"
22
+
23
+ DEVICE = f"--device {DEVICE}"
24
+
25
 
26
  os.system(f"python tools/webui.py {DEVICE}")