From d711880aa93c2f331d9224ca10b427534a4adad3 Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Tue, 1 Aug 2023 01:58:04 +0900 Subject: [PATCH] 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 | 41 +++++++-- modules/dml/backend.py | 46 ++++------ modules/dml/memctl/amd/__init__.py | 8 -- modules/dml/memctl/intel/__init__.py | 6 -- modules/dml/memctl/memctl.py | 8 -- modules/dml/memctl/nvidia/__init__.py | 6 -- modules/dml/memctl/unknown/__init__.py | 5 -- modules/dml/memory.py | 31 +++++++ modules/dml/memory_amd/__init__.py | 7 ++ .../amd => memory_amd}/driver/atiadlxx.py | 0 .../driver/atiadlxx_apis.py | 0 .../driver/atiadlxx_defines.py | 0 .../driver/atiadlxx_structures.py | 0 modules/dml/pdh/__init__.py | 85 +++++++++++++++++++ modules/dml/pdh/apis.py | 36 ++++++++ modules/dml/pdh/defines.py | 22 +++++ modules/dml/pdh/errors.py | 3 + modules/dml/pdh/msvcrt.py | 11 +++ modules/dml/pdh/structures.py | 41 +++++++++ modules/shared.py | 3 +- 20 files changed, 288 insertions(+), 71 deletions(-) delete mode 100644 modules/dml/memctl/amd/__init__.py delete mode 100644 modules/dml/memctl/intel/__init__.py delete mode 100644 modules/dml/memctl/memctl.py delete mode 100644 modules/dml/memctl/nvidia/__init__.py delete mode 100644 modules/dml/memctl/unknown/__init__.py create mode 100644 modules/dml/memory.py create mode 100644 modules/dml/memory_amd/__init__.py rename modules/dml/{memctl/amd => memory_amd}/driver/atiadlxx.py (100%) rename modules/dml/{memctl/amd => memory_amd}/driver/atiadlxx_apis.py (100%) rename modules/dml/{memctl/amd => memory_amd}/driver/atiadlxx_defines.py (100%) rename modules/dml/{memctl/amd => memory_amd}/driver/atiadlxx_structures.py (100%) create mode 100644 modules/dml/pdh/__init__.py create mode 100644 modules/dml/pdh/apis.py create mode 100644 modules/dml/pdh/defines.py create mode 100644 modules/dml/pdh/errors.py create mode 100644 modules/dml/pdh/msvcrt.py create mode 100644 modules/dml/pdh/structures.py diff --git a/modules/dml/__init__.py b/modules/dml/__init__.py index 3c8d8f0e1..a36af5da3 100644 --- a/modules/dml/__init__.py +++ b/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() diff --git a/modules/dml/backend.py b/modules/dml/backend.py index feb0a07bd..ef46f288a 100644 --- a/modules/dml/backend.py +++ b/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) diff --git a/modules/dml/memctl/amd/__init__.py b/modules/dml/memctl/amd/__init__.py deleted file mode 100644 index b2de1c1bf..000000000 --- a/modules/dml/memctl/amd/__init__.py +++ /dev/null @@ -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) diff --git a/modules/dml/memctl/intel/__init__.py b/modules/dml/memctl/intel/__init__.py deleted file mode 100644 index 24909c5bf..000000000 --- a/modules/dml/memctl/intel/__init__.py +++ /dev/null @@ -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) diff --git a/modules/dml/memctl/memctl.py b/modules/dml/memctl/memctl.py deleted file mode 100644 index bda31dd20..000000000 --- a/modules/dml/memctl/memctl.py +++ /dev/null @@ -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 diff --git a/modules/dml/memctl/nvidia/__init__.py b/modules/dml/memctl/nvidia/__init__.py deleted file mode 100644 index 3334b3518..000000000 --- a/modules/dml/memctl/nvidia/__init__.py +++ /dev/null @@ -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) diff --git a/modules/dml/memctl/unknown/__init__.py b/modules/dml/memctl/unknown/__init__.py deleted file mode 100644 index 31fda2bcd..000000000 --- a/modules/dml/memctl/unknown/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from modules.dml.memctl.memctl import MemoryControl - -class UnknownMemoryControl(MemoryControl): - def mem_get_info(index: int): - return (1073741824, 1073741824) diff --git a/modules/dml/memory.py b/modules/dml/memory.py new file mode 100644 index 000000000..af2d8060f --- /dev/null +++ b/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() diff --git a/modules/dml/memory_amd/__init__.py b/modules/dml/memory_amd/__init__.py new file mode 100644 index 000000000..9928d5bc5 --- /dev/null +++ b/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) diff --git a/modules/dml/memctl/amd/driver/atiadlxx.py b/modules/dml/memory_amd/driver/atiadlxx.py similarity index 100% rename from modules/dml/memctl/amd/driver/atiadlxx.py rename to modules/dml/memory_amd/driver/atiadlxx.py diff --git a/modules/dml/memctl/amd/driver/atiadlxx_apis.py b/modules/dml/memory_amd/driver/atiadlxx_apis.py similarity index 100% rename from modules/dml/memctl/amd/driver/atiadlxx_apis.py rename to modules/dml/memory_amd/driver/atiadlxx_apis.py diff --git a/modules/dml/memctl/amd/driver/atiadlxx_defines.py b/modules/dml/memory_amd/driver/atiadlxx_defines.py similarity index 100% rename from modules/dml/memctl/amd/driver/atiadlxx_defines.py rename to modules/dml/memory_amd/driver/atiadlxx_defines.py diff --git a/modules/dml/memctl/amd/driver/atiadlxx_structures.py b/modules/dml/memory_amd/driver/atiadlxx_structures.py similarity index 100% rename from modules/dml/memctl/amd/driver/atiadlxx_structures.py rename to modules/dml/memory_amd/driver/atiadlxx_structures.py diff --git a/modules/dml/pdh/__init__.py b/modules/dml/pdh/__init__.py new file mode 100644 index 000000000..0dcd466cb --- /dev/null +++ b/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.") diff --git a/modules/dml/pdh/apis.py b/modules/dml/pdh/apis.py new file mode 100644 index 000000000..87c1d1204 --- /dev/null +++ b/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] diff --git a/modules/dml/pdh/defines.py b/modules/dml/pdh/defines.py new file mode 100644 index 000000000..a5ea1d479 --- /dev/null +++ b/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 diff --git a/modules/dml/pdh/errors.py b/modules/dml/pdh/errors.py new file mode 100644 index 000000000..60d9ab8f7 --- /dev/null +++ b/modules/dml/pdh/errors.py @@ -0,0 +1,3 @@ +class PDHError(Exception): + def __init__(self, message: str): + super(PDHError, self).__init__(message) diff --git a/modules/dml/pdh/msvcrt.py b/modules/dml/pdh/msvcrt.py new file mode 100644 index 000000000..bc5d93031 --- /dev/null +++ b/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] diff --git a/modules/dml/pdh/structures.py b/modules/dml/pdh/structures.py new file mode 100644 index 000000000..8fb09e6cb --- /dev/null +++ b/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) diff --git a/modules/shared.py b/modules/shared.py index 711a77802..e7e585223 100644 --- a/modules/shared.py +++ b/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"), {