DirectML update.

DirectML reuses GPU memory instead of returning it.
So prints "practical" GPU memory utilization too.

modules/devices.py

@@ -33,7 +33,7 @@ def get_cuda_device_string():
     elif backend == 'directml' and torch.dml.is_available():
         if shared.cmd_opts.device_id is not None:
             return f"privateuseone:{shared.cmd_opts.device_id}"
-        return torch.dml.get_default_device_string()
+        return torch.dml.get_device_string(torch.dml.default_device().index)
     else:
         if shared.cmd_opts.device_id is not None:
             return f"cuda:{shared.cmd_opts.device_id}"
@@ -70,6 +70,9 @@ def torch_gc(force=False):
     if used > 95:
         shared.log.warning(f'GPU high memory utilization: {used}% {mem}')
         force = True
+        if backend == "directml":
+            practical_used = round(100 * torch.cuda.memory_allocated() / (1 << 30) / gpu.get('total', 1))
+            shared.log.info(f'Practical GPU memory utilization: {practical_used}%')
 
     if shared.opts.disable_gc and not force:
         return

modules/dml/__init__.py

@@ -1,5 +1,10 @@
+import os
 import torch
 
+from modules.sd_hijack_utils import CondFunc
+
+do_nothing = lambda: None
+
 def directml_init():
     from modules.dml.backend import DirectML # pylint: disable=ungrouped-imports
     # Alternative of torch.cuda for DirectML.
@@ -7,18 +12,34 @@ def directml_init():
 
     torch.cuda.is_available = lambda: False
     torch.cuda.device = torch.dml.device
+    torch.cuda.device_count = torch.dml.device_count
     torch.cuda.current_device = torch.dml.current_device
     torch.cuda.get_device_name = torch.dml.get_device_name
     torch.cuda.get_device_properties = torch.dml.get_device_properties
 
+    torch.cuda.empty_cache = do_nothing
+    torch.cuda.ipc_collect = do_nothing
     torch.cuda.memory_stats = torch.dml.memory_stats
     torch.cuda.mem_get_info = torch.dml.mem_get_info
     torch.cuda.memory_allocated = torch.dml.memory_allocated
     torch.cuda.max_memory_allocated = torch.dml.max_memory_allocated
     torch.cuda.reset_peak_memory_stats = torch.dml.reset_peak_memory_stats
+    torch.cuda.utilization = lambda: 0
 
-def directml_hijack_init():
+    torch.Tensor.directml = lambda self: self.to(torch.dml.current_device())
+
+    mem_bound = os.environ.get("DML_GPU_MEMORY_BOUND", None)
+    if mem_bound is not None:
+        torch.dml.set_gpu_memory_bound(int(mem_bound))
+
+def directml_do_hijack():
     import modules.dml.hijack
+    from modules.devices import device
+
+    if not torch.dml.has_float64_support(device):
+        CondFunc('torch.from_numpy',
+            lambda orig_func, *args, **kwargs: orig_func(args[0].astype('float32')),
+            lambda *args, **kwargs: args[1].dtype == float)
 
 def directml_override_opts():
     from modules import shared

modules/dml/backend.py

@@ -15,6 +15,8 @@ class DirectML:
 
     context_device: torch.device | None = None
 
+    __gpu_memory_bound: int | None = None
+    
     is_autocast_enabled = False
     autocast_gpu_dtype = torch.float16
 
@@ -33,18 +35,30 @@ class DirectML:
             return memory_control
         except Exception:
             return UnknownMemoryControl
+        
+    def set_gpu_memory_bound(bound: int | None):
+        DirectML.__gpu_memory_bound = bound
 
     def is_available() -> bool:
         return torch_directml.is_available()
 
+    def is_directml_device(device: torch.device) -> bool:
+        return device.type == "privateuseone"
+
+    def has_float64_support(device: Optional[rDevice]=None) -> bool:
+        return torch_directml.has_float64_support(get_device(device).index)
+
+    def device_count() -> int:
+        return torch_directml.device_count()
+
     def current_device() -> torch.device:
         return DirectML.context_device or DirectML.default_device()
 
     def default_device() -> torch.device:
         return torch_directml.device(torch_directml.default_device())
 
-    def get_default_device_string() -> str:
-        return f"privateuseone:{torch_directml.default_device()}"
+    def get_device_string(device: Optional[rDevice]=None) -> str:
+        return f"privateuseone:{get_device(device).index}"
 
     def get_device_name(device: Optional[rDevice]=None) -> str:
         return torch_directml.device_name(get_device(device))
@@ -59,14 +73,19 @@ class DirectML:
             "num_alloc_retries": mem_stat_fill,
         }
 
-    def mem_get_info(device: Optional[rDevice]=None):
+    def mem_get_info(device: Optional[rDevice]=None) -> tuple[int, int]:
         device = get_device(device)
         memory_control = DirectML.__get_memory_control(device)
-        return memory_control.mem_get_info(device.index)
+        mem_info = memory_control.mem_get_info(device.index)
+        if DirectML.__gpu_memory_bound is None:
+            return mem_info
+        used = mem_info[1] - mem_info[0]
+        available = DirectML.__gpu_memory_bound - used
+        return (0 if available < 0 else available, DirectML.__gpu_memory_bound)
 
-    def memory_allocated(device: Optional[rDevice]=None):
+    def memory_allocated(device: Optional[rDevice]=None) -> int:
         device = get_device(device)
-        return sum(torch_directml.gpu_memory(device.index)) / (1 << 20)
+        return sum(torch_directml.gpu_memory(device.index)) * (1 << 20)
 
     def max_memory_allocated(device: Optional[rDevice]=None):
         return DirectML.memory_allocated(device) # DirectML does not empty GPU memory

modules/dml/device_properties.py

@@ -7,6 +7,13 @@ class DeviceProperties:
     minor: int = 0
     total_memory: int
     multi_processor_count: int = 1
+
     def __init__(self, device: torch.device):
         self.name = torch.dml.get_device_name(device)
         self.total_memory = torch.dml.mem_get_info(device)[0]
+
+    def __str__(self):
+        return f"DeviceProperties(name='{self.name}', total_memory='{self.total_memory}')"
+
+    def __repr__(self):
+        return f"DeviceProperties(name='{self.name}', total_memory='{self.total_memory}')"

modules/dml/hijack/torch.py

@@ -3,3 +3,11 @@ 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')
+
+_new = torch.Tensor.new
+def new(self: torch.Tensor, *args, **kwargs):
+    if torch.dml.is_directml_device(self.device):
+        return _new(self.cpu(), *args, **kwargs).to(self.device)
+    return _new(self, *args, **kwargs)
+
+torch.Tensor.new = new

modules/shared.py

@@ -12,7 +12,7 @@ import requests
 import fasteners
 from modules import errors, ui_components, shared_items, cmd_args
 from modules.paths_internal import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611
-from modules.dml import directml_hijack_init, directml_override_opts
+from modules.dml import directml_do_hijack, directml_override_opts
 import modules.interrogate
 import modules.memmon
 import modules.styles
@@ -803,7 +803,7 @@ parallel_processing_allowed = not cmd_opts.lowvram
 mem_mon = modules.memmon.MemUsageMonitor("MemMon", device, opts)
 mem_mon.start()
 if devices.backend == "directml":
-    directml_hijack_init()
+    directml_do_hijack()
     directml_override_opts()