gitea-mirror/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-12-21 23:00:59 -08:00
Code Issues Packages Projects Releases Wiki Activity

Modular device management

Browse Source
This commit is contained in:
Abdullah Barhoum
2022-09-11 07:11:27 +02:00
committed by AUTOMATIC1111
parent 065e310a3f
commit b5d1af11b7
4 changed files with 19 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
modules/devices.py Normal file
View File

@@ -0,0 +1,12 @@
import torch
# has_mps is only available in nightly pytorch (for now), `getattr` for compatibility
has_mps = getattr(torch, 'has_mps', False)
def get_optimal_device():
if torch.cuda.is_available():
return torch.device("cuda")
if has_mps:
return torch.device("mps")
return torch.device("cpu")