diff --git a/modules/dml/hijack/torch.py b/modules/dml/hijack/torch.py index e216de46f..1786a1499 100644 --- a/modules/dml/hijack/torch.py +++ b/modules/dml/hijack/torch.py @@ -1,12 +1,17 @@ import torch - from modules.sd_hijack_utils import CondFunc + CondFunc('torchsde._brownian.brownian_interval._randn', lambda _, size, dtype, device, seed: torch.randn(size, dtype=dtype, device=torch.device("cpu"), generator=torch.Generator(torch.device("cpu")).manual_seed(int(seed))).to(device), lambda _, size, dtype, device, seed: device.type == 'privateuseone') # https://github.com/microsoft/DirectML/issues/400 CondFunc('torch.Tensor.new', lambda orig, self, *args, **kwargs: orig(self.cpu(), *args, **kwargs).to(self.device), lambda orig, self, *args, **kwargs: torch.dml.is_directml_device(self.device)) +def cuda(self: torch.Tensor): + return self.to(torch.dml.current_device()) +torch.Tensor.cuda = cuda + + # https://github.com/lshqqytiger/stable-diffusion-webui-directml/issues/436 _pow_ = torch.Tensor.pow_ def pow_(self: torch.Tensor, *args, **kwargs): diff --git a/modules/sd_offload.py b/modules/sd_offload.py index b200bed1a..7e603f4ec 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -162,6 +162,12 @@ class OffloadHook(accelerate.hooks.ModelHook): if device_map is None or max_memory != getattr(module, "balanced_offload_max_memory", None): device_map = accelerate.infer_auto_device_map(module, max_memory=max_memory) offload_dir = getattr(module, "offload_dir", os.path.join(shared.opts.accelerate_offload_path, module.__class__.__name__)) + keys = device_map.keys() + for v in keys: + if isinstance(device_map[v], int): + # int implies CUDA device, but it will break DirectML backend. + # Therefore, the type of device should be added. + device_map[v] = f"{devices.device.type}:{device_map[v]}" module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir) module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access module.balanced_offload_device_map = device_map