diff options
| author | Jairo Correa <jn.j41r0@gmail.com> | 2022-10-04 22:53:52 +0000 |
|---|---|---|
| committer | Jairo Correa <jn.j41r0@gmail.com> | 2022-10-04 22:53:52 +0000 |
| commit | 1f50971fb8c83c255c2819dd0b3f29a46b74f7d9 (patch) | |
| tree | fd57f40a1ffa2b28105ec0bb3f7f3ab4a742681a /modules/devices.py | |
| parent | ad0cc85d1f0bd52877963f296eb1257a0c2b012b (diff) | |
| parent | ef40e4cd4d383a3405e03f1da3f5b5a1820a8f53 (diff) | |
| download | stable-diffusion-webui-gfx803-1f50971fb8c83c255c2819dd0b3f29a46b74f7d9.tar.gz stable-diffusion-webui-gfx803-1f50971fb8c83c255c2819dd0b3f29a46b74f7d9.tar.bz2 stable-diffusion-webui-gfx803-1f50971fb8c83c255c2819dd0b3f29a46b74f7d9.zip | |
Merge branch 'master' into fix-vram
Diffstat (limited to 'modules/devices.py')
| -rw-r--r-- | modules/devices.py | 15 |
1 files changed, 12 insertions, 3 deletions
diff --git a/modules/devices.py b/modules/devices.py index ebf40082..6db4e57c 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -1,9 +1,11 @@ +import contextlib + import torch import gc -# has_mps is only available in nightly pytorch (for now), `getattr` for compatibility from modules import errors +# has_mps is only available in nightly pytorch (for now), `getattr` for compatibility has_mps = getattr(torch, 'has_mps', False) cpu = torch.device("cpu") @@ -33,8 +35,7 @@ def enable_tf32(): errors.run(enable_tf32, "Enabling TF32") -device = get_optimal_device() -device_codeformer = cpu if has_mps else device +device = device_gfpgan = device_bsrgan = device_esrgan = device_scunet = device_codeformer = get_optimal_device() dtype = torch.float16 def randn(seed, shape): @@ -58,3 +59,11 @@ def randn_without_seed(shape): return torch.randn(shape, device=device) + +def autocast(): + from modules import shared + + if dtype == torch.float32 or shared.cmd_opts.precision == "full": + return contextlib.nullcontext() + + return torch.autocast("cuda") |