fix directml backend

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):

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