diff options
| author | AUTOMATIC1111 <16777216c@gmail.com> | 2022-10-06 17:30:29 +0000 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2022-10-06 17:30:29 +0000 |
| commit | ab4ddbf333eef170804ef8de67001f77c8fdd64c (patch) | |
| tree | 21cb1109f8eae463aa4066eec0926cd71ab81740 /modules/devices.py | |
| parent | 2a7f48cdb8dcf9acb02610cccae0d1ee5d260bc2 (diff) | |
| parent | cf7c784fcc0c84a8a4edd8d3aca4dda4c7025c43 (diff) | |
| download | stable-diffusion-webui-gfx803-ab4ddbf333eef170804ef8de67001f77c8fdd64c.tar.gz stable-diffusion-webui-gfx803-ab4ddbf333eef170804ef8de67001f77c8fdd64c.tar.bz2 stable-diffusion-webui-gfx803-ab4ddbf333eef170804ef8de67001f77c8fdd64c.zip | |
Merge branch 'master' into gallery-styling
Diffstat (limited to 'modules/devices.py')
| -rw-r--r-- | modules/devices.py | 18 |
1 files changed, 13 insertions, 5 deletions
diff --git a/modules/devices.py b/modules/devices.py index 07bb2339..0158b11f 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -1,8 +1,10 @@ +import contextlib + import torch -# 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") @@ -32,10 +34,8 @@ 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): # Pytorch currently doesn't handle setting randomness correctly when the metal backend is used. @@ -58,3 +58,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") |