stable-diffusion-webui/modules/devices.py

Ignoring revisions in .git-blame-ignore-revs. Click here to bypass and see the normal blame view.

296 lines
8.6 KiB
Python
Raw Normal View History

import sys
import contextlib
from functools import lru_cache
2022-09-10 23:11:27 -06:00
import torch
2024-01-27 02:21:32 -07:00
from modules import errors, shared, npu_specific
if sys.platform == "darwin":
from modules import mac_specific
if shared.cmd_opts.use_ipex:
from modules import xpu_specific
def has_xpu() -> bool:
return shared.cmd_opts.use_ipex and xpu_specific.has_xpu
def has_mps() -> bool:
if sys.platform != "darwin":
return False
else:
return mac_specific.has_mps
2022-09-11 09:48:36 -06:00
2023-10-28 01:24:26 -06:00
def cuda_no_autocast(device_id=None) -> bool:
if device_id is None:
device_id = get_cuda_device_id()
return (
2023-11-19 00:50:06 -07:00
torch.cuda.get_device_capability(device_id) == (7, 5)
2023-10-28 01:24:26 -06:00
and torch.cuda.get_device_name(device_id).startswith("NVIDIA GeForce GTX 16")
)
def get_cuda_device_id():
return (
2023-11-19 00:50:06 -07:00
int(shared.cmd_opts.device_id)
if shared.cmd_opts.device_id is not None and shared.cmd_opts.device_id.isdigit()
2023-10-28 01:24:26 -06:00
else 0
) or torch.cuda.current_device()
2022-11-27 03:08:54 -07:00
def get_cuda_device_string():
if shared.cmd_opts.device_id is not None:
return f"cuda:{shared.cmd_opts.device_id}"
2022-11-27 03:08:54 -07:00
return "cuda"
2022-11-27 03:08:54 -07:00
def get_optimal_device_name():
2022-11-27 03:08:54 -07:00
if torch.cuda.is_available():
return get_cuda_device_string()
2022-09-11 09:48:36 -06:00
if has_mps():
return "mps"
if has_xpu():
2023-11-09 20:06:26 -07:00
return xpu_specific.get_xpu_device_string()
2024-01-27 02:21:32 -07:00
if npu_specific.has_npu:
return npu_specific.get_npu_device_string()
return "cpu"
2022-09-11 09:48:36 -06:00
def get_optimal_device():
return torch.device(get_optimal_device_name())
def get_device_for(task):
if task in shared.cmd_opts.use_cpu or "all" in shared.cmd_opts.use_cpu:
return cpu
return get_optimal_device()
def torch_gc():
if torch.cuda.is_available():
2022-11-27 03:08:54 -07:00
with torch.cuda.device(get_cuda_device_string()):
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
if has_mps():
mac_specific.torch_mps_gc()
if has_xpu():
xpu_specific.torch_xpu_gc()
2024-01-27 02:21:32 -07:00
if npu_specific.has_npu:
2024-01-30 19:46:53 -07:00
torch_npu_set_device()
2024-01-27 02:21:32 -07:00
npu_specific.torch_npu_gc()
2024-01-30 19:46:53 -07:00
def torch_npu_set_device():
# Work around due to bug in torch_npu, revert me after fixed, @see https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue
if npu_specific.has_npu:
torch.npu.set_device(0)
def enable_tf32():
if torch.cuda.is_available():
# enabling benchmark option seems to enable a range of cards to do fp16 when they otherwise can't
# see https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/4407
2023-10-28 01:24:26 -06:00
if cuda_no_autocast():
2022-11-07 19:06:48 -07:00
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
errors.run(enable_tf32, "Enabling TF32")
2022-09-12 11:09:32 -06:00
cpu: torch.device = torch.device("cpu")
fp8: bool = False
2024-05-16 17:50:06 -06:00
# Force fp16 for all models in inference. No casting during inference.
# This flag is controlled by "--precision half" command line arg.
force_fp16: bool = False
device: torch.device = None
device_interrogate: torch.device = None
device_gfpgan: torch.device = None
device_esrgan: torch.device = None
device_codeformer: torch.device = None
dtype: torch.dtype = torch.float16
dtype_vae: torch.dtype = torch.float16
dtype_unet: torch.dtype = torch.float16
dtype_inference: torch.dtype = torch.float16
unet_needs_upcast = False
2022-09-12 11:09:32 -06:00
def cond_cast_unet(input):
2024-05-16 17:50:06 -06:00
if force_fp16:
return input.to(torch.float16)
return input.to(dtype_unet) if unet_needs_upcast else input
def cond_cast_float(input):
return input.float() if unet_needs_upcast else input
nv_rng = None
2023-10-28 01:24:26 -06:00
patch_module_list = [
torch.nn.Linear,
torch.nn.Conv2d,
torch.nn.MultiheadAttention,
torch.nn.GroupNorm,
torch.nn.LayerNorm,
]
2023-11-19 00:50:06 -07:00
def manual_cast_forward(target_dtype):
def forward_wrapper(self, *args, **kwargs):
2024-01-09 07:39:39 -07:00
if any(
isinstance(arg, torch.Tensor) and arg.dtype != target_dtype
for arg in args
):
args = [arg.to(target_dtype) if isinstance(arg, torch.Tensor) else arg for arg in args]
kwargs = {k: v.to(target_dtype) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
2024-01-09 07:53:58 -07:00
org_dtype = target_dtype
for param in self.parameters():
if param.dtype != target_dtype:
org_dtype = param.dtype
break
2024-01-29 07:27:53 -07:00
2024-01-09 07:39:39 -07:00
if org_dtype != target_dtype:
self.to(target_dtype)
result = self.org_forward(*args, **kwargs)
2024-01-09 07:39:39 -07:00
if org_dtype != target_dtype:
self.to(org_dtype)
2024-01-09 07:53:58 -07:00
if target_dtype != dtype_inference:
if isinstance(result, tuple):
result = tuple(
i.to(dtype_inference)
if isinstance(i, torch.Tensor)
else i
for i in result
)
elif isinstance(result, torch.Tensor):
result = result.to(dtype_inference)
return result
return forward_wrapper
2023-11-19 00:50:06 -07:00
2023-10-28 01:24:26 -06:00
@contextlib.contextmanager
def manual_cast(target_dtype):
2024-01-20 01:31:12 -07:00
applied = False
2023-10-28 01:24:26 -06:00
for module_type in patch_module_list:
2024-01-18 09:14:03 -07:00
if hasattr(module_type, "org_forward"):
continue
2024-01-20 01:31:12 -07:00
applied = True
2023-10-28 01:24:26 -06:00
org_forward = module_type.forward
2024-01-29 07:27:53 -07:00
if module_type == torch.nn.MultiheadAttention:
module_type.forward = manual_cast_forward(torch.float32)
else:
module_type.forward = manual_cast_forward(target_dtype)
2023-10-28 01:24:26 -06:00
module_type.org_forward = org_forward
try:
yield None
finally:
2024-01-20 01:33:59 -07:00
if applied:
for module_type in patch_module_list:
if hasattr(module_type, "org_forward"):
module_type.forward = module_type.org_forward
delattr(module_type, "org_forward")
2022-10-10 07:11:14 -06:00
def autocast(disable=False):
if disable:
return contextlib.nullcontext()
2024-05-16 17:50:06 -06:00
if force_fp16:
# No casting during inference if force_fp16 is enabled.
# All tensor dtype conversion happens before inference.
return contextlib.nullcontext()
2023-10-28 01:24:26 -06:00
if fp8 and device==cpu:
return torch.autocast("cpu", dtype=torch.bfloat16, enabled=True)
if fp8 and dtype_inference == torch.float32:
return manual_cast(dtype)
if dtype == torch.float32 or dtype_inference == torch.float32:
return contextlib.nullcontext()
2024-01-09 08:30:55 -07:00
if has_xpu() or has_mps() or cuda_no_autocast():
return manual_cast(dtype)
return torch.autocast("cuda")
def without_autocast(disable=False):
return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext()
class NansException(Exception):
pass
def test_for_nans(x, where):
2023-01-17 01:04:56 -07:00
if shared.cmd_opts.disable_nan_check:
return
2024-06-08 03:41:28 -06:00
if not torch.isnan(x[(0, ) * len(x.shape)]):
return
if where == "unet":
2024-06-08 03:41:28 -06:00
message = "A tensor with NaNs was produced in Unet."
if not shared.cmd_opts.no_half:
message += " This could be either because there's not enough precision to represent the picture, or because your video card does not support half type. Try setting the \"Upcast cross attention layer to float32\" option in Settings > Stable Diffusion or using the --no-half commandline argument to fix this."
elif where == "vae":
2024-06-08 03:41:28 -06:00
message = "A tensor with NaNs was produced in VAE."
if not shared.cmd_opts.no_half and not shared.cmd_opts.no_half_vae:
message += " This could be because there's not enough precision to represent the picture. Try adding --no-half-vae commandline argument to fix this."
else:
2024-06-08 03:41:28 -06:00
message = "A tensor with NaNs was produced."
message += " Use --disable-nan-check commandline argument to disable this check."
raise NansException(message)
@lru_cache
def first_time_calculation():
"""
just do any calculation with pytorch layers - the first time this is done it allocaltes about 700MB of memory and
2024-03-03 23:37:23 -07:00
spends about 2.7 seconds doing that, at least with NVidia.
"""
x = torch.zeros((1, 1)).to(device, dtype)
linear = torch.nn.Linear(1, 1).to(device, dtype)
linear(x)
x = torch.zeros((1, 1, 3, 3)).to(device, dtype)
conv2d = torch.nn.Conv2d(1, 1, (3, 3)).to(device, dtype)
conv2d(x)
2024-05-16 17:50:06 -06:00
def force_model_fp16():
"""
ldm and sgm has modules.diffusionmodules.util.GroupNorm32.forward, which
force conversion of input to float32. If force_fp16 is enabled, we need to
prevent this casting.
"""
assert force_fp16
import sgm.modules.diffusionmodules.util as sgm_util
import ldm.modules.diffusionmodules.util as ldm_util
sgm_util.GroupNorm32 = torch.nn.GroupNorm
ldm_util.GroupNorm32 = torch.nn.GroupNorm
print("ldm/sgm GroupNorm32 replaced with normal torch.nn.GroupNorm due to `--precision half`.")