New option for DirectML: memory stats provider.

1. Performance Counter.
    Get vram size allocated to & used by python.exe from pdh.dll.
    Generation can be slower than atiadlxx.
    Use memory less greedy then atiadlxx.
    Windows only.
2. atiadlxx.
    Get max vram size and available vram size from AMD GPU driver (atiadlxx.dll).
    Use memory more greedy than Performance Counter.
    Windows & WSL are supported.
3. None.
    Assume available vram size is 8GB.
    Use memory regardless of current vram usage.

modules/dml/__init__.py

@@ -1,11 +1,40 @@
-import os
+from platform import system
 import torch
 from typing import NamedTuple, Callable, Optional
 
 from modules.sd_hijack_utils import CondFunc
 
+memory_providers = ["None", "atiadlxx (AMD only)"]
+default_memory_provider = "None"
+if system() == "Windows":
+    memory_providers.append("Performance Counter")
+    default_memory_provider = "Performance Counter"
 do_nothing = lambda: None
 
+def _set_memory_provider():
+    from modules.shared import opts, cmd_opts, log
+
+    if opts.directml_memory_provider == "Performance Counter":
+        from .backend import pdh_mem_get_info
+        from .memory import MemoryProvider
+        torch.dml.mem_get_info = pdh_mem_get_info
+        if torch.dml.memory_provider is not None:
+            del torch.dml.memory_provider
+        torch.dml.memory_provider = MemoryProvider()
+    elif opts.directml_memory_provider == "atiadlxx (AMD only)":
+        device_name = torch.dml.get_device_name(cmd_opts.device_id)
+        if "AMD" not in device_name and "Radeon" not in device_name:
+            log.warning(f"Memory stats provider is changed to None because the current device is not AMDGPU. Current Device: {device_name}")
+            opts.directml_memory_provider = "None"
+            _set_memory_provider()
+            return
+        from .backend import amd_mem_get_info
+        torch.dml.mem_get_info = amd_mem_get_info
+    else:
+        from .backend import mem_get_info
+        torch.dml.mem_get_info = mem_get_info
+    torch.cuda.mem_get_info = torch.dml.mem_get_info
+
 def directml_init():
     from modules.dml.backend import DirectML # pylint: disable=ungrouped-imports
     # Alternative of torch.cuda for DirectML.
@@ -29,10 +58,6 @@ def directml_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
@@ -42,6 +67,8 @@ def directml_do_hijack():
             lambda orig_func, *args, **kwargs: orig_func(args[0].astype('float32')),
             lambda *args, **kwargs: args[1].dtype == float)
 
+    _set_memory_provider()
+
 class OverrideItem(NamedTuple):
     value: str
     condition: Optional[Callable]
@@ -61,9 +88,9 @@ def directml_override_opts():
 
     count = 0
     for key in opts_override_table:
-        count += 1
         item = opts_override_table[key]
         if getattr(shared.opts, key) != item.value and (item.condition is None or item.condition(shared.opts)):
+            count += 1
             setattr(shared.opts, key, item.value)
             if item.message is not None:
                 shared.log.warning(item.message)
@@ -71,3 +98,5 @@ def directml_override_opts():
     
     if count > 0:
         shared.log.info(f'{count} options are automatically overriden. If you want to keep them from overriding, run with --experimental argument.')
+
+    _set_memory_provider()

modules/dml/backend.py

@@ -1,13 +1,24 @@
 # pylint: disable=no-member,no-self-argument,no-method-argument
-from typing import Optional
+from typing import Optional, Callable
 import torch
 import torch_directml # pylint: disable=import-error
 import modules.dml.amp as amp
 
-from .memctl.unknown import UnknownMemoryControl
 from .utils import rDevice, get_device
 from .device import device
 from .device_properties import DeviceProperties
+from .memory_amd import AMDMemoryProvider
+from .memory import MemoryProvider
+
+def amd_mem_get_info(device: Optional[rDevice]=None) -> tuple[int, int]:
+    return AMDMemoryProvider.mem_get_info(get_device(device).index)
+
+def pdh_mem_get_info(device: Optional[rDevice]=None) -> tuple[int, int]:
+    mem_info = DirectML.memory_provider.get_memory(get_device(device).index)
+    return (mem_info["total_committed"] - mem_info["dedicated_usage"], mem_info["total_committed"])
+
+def mem_get_info(device: Optional[rDevice]=None) -> tuple[int, int]:
+    return (8589934592, 8589934592)
 
 class DirectML:
     amp = amp
@@ -15,29 +26,10 @@ class DirectML:
 
     context_device: Optional[torch.device] = None
 
-    __gpu_memory_bound: Optional[int] = None
-    
     is_autocast_enabled = False
     autocast_gpu_dtype = torch.float16
 
-    def __get_memory_control(device: torch.device):
-        assert device.type == 'privateuseone'
-        try:
-            device_name = torch_directml.device_name(device.index)
-            if 'NVIDIA' in device_name or 'GeForce' in device_name:
-                from .memctl.nvidia import nVidiaMemoryControl as memory_control
-            elif 'AMD' in device_name or 'Radeon' in device_name:
-                from .memctl.amd import AMDMemoryControl as memory_control
-            elif 'Intel' in device_name:
-                from .memctl.intel import IntelMemoryControl as memory_control
-            else:
-                return UnknownMemoryControl
-            return memory_control
-        except Exception:
-            return UnknownMemoryControl
-        
-    def set_gpu_memory_bound(bound: Optional[int]):
-        DirectML.__gpu_memory_bound = bound
+    memory_provider: Optional[MemoryProvider] = None
 
     def is_available() -> bool:
         return torch_directml.is_available()
@@ -73,15 +65,7 @@ class DirectML:
             "num_alloc_retries": mem_stat_fill,
         }
 
-    def mem_get_info(device: Optional[rDevice]=None) -> tuple[int, int]:
-        device = get_device(device)
-        memory_control = DirectML.__get_memory_control(device)
-        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)
+    mem_get_info: Callable = mem_get_info
 
     def memory_allocated(device: Optional[rDevice]=None) -> int:
         return sum(torch_directml.gpu_memory(get_device(device).index)) * (1 << 20)

modules/dml/memctl/amd/__init__.py

@@ -1,8 +0,0 @@
-from modules.dml.memctl.memctl import MemoryControl
-from .driver.atiadlxx import ATIADLxx
-
-class AMDMemoryControl(MemoryControl):
-    driver: ATIADLxx = ATIADLxx()
-    def mem_get_info(index):
-        usage = AMDMemoryControl.driver.get_dedicated_vram_usage(index) * (1 << 20)
-        return (AMDMemoryControl.driver.iHyperMemorySize - usage, AMDMemoryControl.driver.iHyperMemorySize)

modules/dml/memctl/intel/__init__.py

@@ -1,6 +0,0 @@
-from modules.dml.memctl.memctl import MemoryControl
-
-class IntelMemoryControl(MemoryControl):
-    def mem_get_info(index: int):
-        # DML TODO: Implement or find a general (and also lightweight) way.
-        return (1073741824, 1073741824)

modules/dml/memctl/memctl.py

@@ -1,8 +0,0 @@
-from abc import *
-from typing import *
-
-class MemoryControl(metaclass=ABCMeta):
-    driver: Any = None
-    @abstractmethod
-    def mem_get_info(index: int) -> Tuple[int, int]:
-        pass

modules/dml/memctl/nvidia/__init__.py

@@ -1,6 +0,0 @@
-from modules.dml.memctl.memctl import MemoryControl
-
-class nVidiaMemoryControl(MemoryControl):
-    def mem_get_info(index: int):
-        # DML TODO: Implement or find a general (and also lightweight) way.
-        return (1073741824, 1073741824)

modules/dml/memctl/unknown/__init__.py

@@ -1,5 +0,0 @@
-from modules.dml.memctl.memctl import MemoryControl
-
-class UnknownMemoryControl(MemoryControl):
-    def mem_get_info(index: int):
-        return (1073741824, 1073741824)

modules/dml/memory.py

@@ -0,0 +1,31 @@
+from os import getpid
+from collections import defaultdict
+
+from modules.dml.pdh import HQuery, HCounter, expand_wildcard_path
+
+class MemoryProvider:
+    hQuery: HQuery
+    hCounters: defaultdict[str, list[HCounter]]
+
+    def __init__(self):
+        self.hQuery = HQuery()
+        self.hCounters = defaultdict(list)
+
+    def get_memory(self, device_id: int) -> dict[str, int]:
+        if len(self.hCounters) == 0:
+            pid = getpid()
+            paths_dedicated = expand_wildcard_path(f"\\GPU Process Memory(pid_{pid}_*_phys_{device_id})\\Dedicated Usage")
+            paths_committed = expand_wildcard_path(f"\\GPU Process Memory(pid_{pid}_*_phys_{device_id})\\Total Committed")
+            for path in paths_dedicated:
+                self.hCounters["dedicated_usage"].append(self.hQuery.add_counter(path))
+            for path in paths_committed:
+                self.hCounters["total_committed"].append(self.hQuery.add_counter(path))
+        self.hQuery.collect_data()
+        result = defaultdict(int)
+        for key in self.hCounters:
+            for hCounter in self.hCounters[key]:
+                result[key] += hCounter.get_formatted_value(int)
+        return dict(result)
+
+    def __del__(self):
+        self.hQuery.close()

modules/dml/memory_amd/__init__.py

@@ -0,0 +1,7 @@
+from .driver.atiadlxx import ATIADLxx
+
+class AMDMemoryProvider:
+    driver: ATIADLxx = ATIADLxx()
+    def mem_get_info(index):
+        usage = AMDMemoryProvider.driver.get_dedicated_vram_usage(index) * (1 << 20)
+        return (AMDMemoryProvider.driver.iHyperMemorySize - usage, AMDMemoryProvider.driver.iHyperMemorySize)

modules/dml/pdh/__init__.py

@@ -0,0 +1,85 @@
+from ctypes import *
+from ctypes.wintypes import *
+from typing import NamedTuple, TypeVar
+
+from .apis import PdhExpandWildCardPathW, PdhOpenQueryW, PdhAddEnglishCounterW, PdhCollectQueryData, PdhGetFormattedCounterValue, PdhGetFormattedCounterArrayW, PdhCloseQuery
+from .structures import PDH_HQUERY, PDH_HCOUNTER, PDH_FMT_COUNTERVALUE, PPDH_FMT_COUNTERVALUE_ITEM_W
+from .defines import *
+from .msvcrt import malloc
+from .errors import PDHError
+
+class __InternalAbstraction(NamedTuple):
+    flag: int
+    attr_name: str
+
+_type_map = {
+    int: __InternalAbstraction(PDH_FMT_LARGE, "largeValue"),
+    float: __InternalAbstraction(PDH_FMT_DOUBLE, "doubleValue"),
+}
+
+def expand_wildcard_path(path: str) -> list[str]:
+    listLength = DWORD(0)
+    if PdhExpandWildCardPathW(None, LPCWSTR(path), None, byref(listLength), PDH_NOEXPANDCOUNTERS) != PDH_MORE_DATA:
+        raise PDHError("Something went wrong.")
+    expanded = (WCHAR * listLength.value)()
+    if PdhExpandWildCardPathW(None, LPCWSTR(path), expanded, byref(listLength), PDH_NOEXPANDCOUNTERS) != PDH_OK:
+        raise PDHError(f"Couldn't expand wildcard path '{path}'")
+    result = list()
+    cur = str()
+    for chr in expanded:
+        if chr == '\0':
+            result.append(cur)
+            cur = str()
+        else:
+            cur += chr
+    result.pop()
+    return result
+
+T = TypeVar("T", *_type_map.keys())
+
+class HCounter(PDH_HCOUNTER):
+    def get_formatted_value(self, type: T) -> T:
+        if type not in _type_map:
+            raise PDHError(f"Invalid value type: {type}")
+        flag, attr_name = _type_map[type]
+        value = PDH_FMT_COUNTERVALUE()
+        if PdhGetFormattedCounterValue(self, DWORD(flag | PDH_FMT_NOSCALE), None, byref(value)) != PDH_OK:
+            raise PDHError("Couldn't get formatted counter value.")
+        return getattr(value.u, attr_name)
+
+    def get_formatted_dict(self, type: T) -> dict[str, T]:
+        if type not in _type_map:
+            raise PDHError(f"Invalid value type: {type}")
+        flag, attr_name = _type_map[type]
+        bufferSize = DWORD(0)
+        itemCount = DWORD(0)
+        if PdhGetFormattedCounterArrayW(self, DWORD(flag | PDH_FMT_NOSCALE), byref(bufferSize), byref(itemCount), None) != PDH_MORE_DATA:
+            raise PDHError("Something went wrong.")
+        itemBuffer = cast(malloc(c_size_t(bufferSize.value)), PPDH_FMT_COUNTERVALUE_ITEM_W)
+        if PdhGetFormattedCounterArrayW(self, DWORD(flag | PDH_FMT_NOSCALE), byref(bufferSize), byref(itemCount), itemBuffer) != PDH_OK:
+            raise PDHError("Couldn't get formatted counter array.")
+        result: dict[str, T] = dict()
+        for i in range(0, itemCount.value):
+            item = itemBuffer[i]
+            result[item.szName] = getattr(item.FmtValue.u, attr_name)
+        return result
+
+class HQuery(PDH_HQUERY):
+    def __init__(self):
+        super(HQuery, self).__init__()
+        if PdhOpenQueryW(None, None, byref(self)) != PDH_OK:
+            raise PDHError("Couldn't open PDH query.")
+
+    def add_counter(self, path: str) -> HCounter:
+        hCounter = HCounter()
+        if PdhAddEnglishCounterW(self, LPCWSTR(path), None, byref(hCounter)) != PDH_OK:
+            raise PDHError("Couldn't add counter query.")
+        return hCounter
+
+    def collect_data(self):
+        if PdhCollectQueryData(self) != PDH_OK:
+            raise PDHError("Couldn't collect query data.")
+
+    def close(self):
+        if PdhCloseQuery(self) != PDH_OK:
+            raise PDHError("Couldn't close PDH query.")

modules/dml/pdh/apis.py

@@ -0,0 +1,36 @@
+from ctypes import *
+from ctypes.wintypes import *
+from typing import Callable
+
+from .structures import *
+from .defines import *
+
+pdh = CDLL("pdh.dll")
+
+PdhExpandWildCardPathW: Callable = pdh.PdhExpandWildCardPathW
+PdhExpandWildCardPathW.restype = PDH_FUNCTION
+PdhExpandWildCardPathW.argtypes = [LPCWSTR, LPCWSTR, PZZWSTR, LPDWORD, DWORD]
+
+PdhOpenQueryW: Callable = pdh.PdhOpenQueryW
+PdhOpenQueryW.restype = PDH_FUNCTION
+PdhOpenQueryW.argtypes = [LPCWSTR, DWORD_PTR, POINTER(PDH_HQUERY)]
+
+PdhAddEnglishCounterW: Callable = pdh.PdhAddEnglishCounterW
+PdhAddEnglishCounterW.restype = PDH_FUNCTION
+PdhAddEnglishCounterW.argtypes = [PDH_HQUERY, LPCWSTR, DWORD_PTR, POINTER(PDH_HCOUNTER)]
+
+PdhCollectQueryData: Callable = pdh.PdhCollectQueryData
+PdhCollectQueryData.restype = PDH_FUNCTION
+PdhCollectQueryData.argtypes = [PDH_HQUERY]
+
+PdhGetFormattedCounterValue: Callable = pdh.PdhGetFormattedCounterValue
+PdhGetFormattedCounterValue.restype = PDH_FUNCTION
+PdhGetFormattedCounterValue.argtypes = [PDH_HCOUNTER, DWORD, LPDWORD, PPDH_FMT_COUNTERVALUE]
+
+PdhGetFormattedCounterArrayW: Callable = pdh.PdhGetFormattedCounterArrayW
+PdhGetFormattedCounterArrayW.restype = PDH_FUNCTION
+PdhGetFormattedCounterArrayW.argtypes = [PDH_HCOUNTER, DWORD, LPDWORD, LPDWORD, PPDH_FMT_COUNTERVALUE_ITEM_W]
+
+PdhCloseQuery: Callable = pdh.PdhCloseQuery
+PdhCloseQuery.restype = PDH_FUNCTION
+PdhCloseQuery.argtypes = [PDH_HQUERY]

modules/dml/pdh/defines.py

@@ -0,0 +1,22 @@
+from ctypes import *
+from ctypes.wintypes import *
+
+PDH_FUNCTION = c_int
+PDH_OK = 0x00000000
+PDH_MORE_DATA = -2147481646#0x800007D2
+
+DWORD_PTR = POINTER(DWORD)
+PWSTR = POINTER(WCHAR)
+PZZWSTR = POINTER(WCHAR)
+
+PDH_NOEXPANDCOUNTERS = 1
+PDH_NOEXPANDINSTANCES = 2
+PDH_REFRESHCOUNTERS = 4
+
+PDH_FMT_LONG = 0x00000100
+PDH_FMT_DOUBLE = 0x00000200
+PDH_FMT_LARGE = 0x00000400
+
+PDH_FMT_NOSCALE = 0x00001000
+PDH_FMT_1000 = 0x00002000
+PDH_FMT_NOCAP100 = 0x00008000

modules/dml/pdh/errors.py

@@ -0,0 +1,3 @@
+class PDHError(Exception):
+    def __init__(self, message: str):
+        super(PDHError, self).__init__(message)

modules/dml/pdh/msvcrt.py

@@ -0,0 +1,11 @@
+from ctypes import *
+
+msvcrt = CDLL("msvcrt")
+
+malloc = msvcrt.malloc
+malloc.restype = c_void_p
+malloc.argtypes = [c_size_t]
+
+free = msvcrt.free
+free.restype = None
+free.argtypes = [c_void_p]

modules/dml/pdh/structures.py

@@ -0,0 +1,41 @@
+from ctypes import *
+from ctypes.wintypes import *
+
+PDH_HQUERY = HANDLE
+PDH_HCOUNTER = HANDLE
+
+class PDH_FMT_COUNTERVALUE_U(Union):
+    _fields_ = [
+        ("longValue", LONG),
+        ("doubleValue", c_double),
+        ("largeValue", c_longlong),
+        ("AnsiStringValue", LPCSTR),
+        ("WideStringValue", LPCWSTR),
+    ]
+
+    longValue: int
+    doubleValue: float
+    largeValue: int
+    AnsiStringValue: LPCSTR
+    WideStringValue: LPCWSTR
+
+class PDH_FMT_COUNTERVALUE(Structure):
+    _anonymous_ = ("u",)
+    _fields_ = [
+        ("CStatus", DWORD),
+        ("u", PDH_FMT_COUNTERVALUE_U),
+    ]
+
+    CStatus: DWORD
+    u: PDH_FMT_COUNTERVALUE_U
+PPDH_FMT_COUNTERVALUE = POINTER(PDH_FMT_COUNTERVALUE)
+
+class PDH_FMT_COUNTERVALUE_ITEM_W(Structure):
+    _fields_ = [
+        ("szName", LPWSTR),
+        ("FmtValue", PDH_FMT_COUNTERVALUE),
+    ]
+
+    szName: str
+    FmtValue: PDH_FMT_COUNTERVALUE
+PPDH_FMT_COUNTERVALUE_ITEM_W = POINTER(PDH_FMT_COUNTERVALUE_ITEM_W)

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_do_hijack
+from modules.dml import memory_providers, default_memory_provider, directml_do_hijack
 import modules.interrogate
 import modules.memmon
 import modules.styles
@@ -390,6 +390,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "cuda_compile_verbose": OptionInfo(False, "Model compile verbose mode"),
     "cuda_compile_errors": OptionInfo(True, "Model compile suppress errors"),
     "disable_gc": OptionInfo(True, "Disable Torch memory garbage collection"),
+    "directml_memory_provider": OptionInfo(default_memory_provider, '[DirectML] Memory stats provider', gr.Dropdown, lambda: {"choices": memory_providers}),
 }))
 
 options_templates.update(options_section(('diffusers', "Diffusers Settings"), {