DirectML rework.

2026-09-19 09:14:35 +02:00 · 2023-07-15 16:19:38 +09:00
parent 558b71f088
commit 578f206bf7
11 changed files with 9 additions and 249 deletions
@@ -1,39 +1,9 @@
-# pylint: disable=no-member,no-self-argument,no-method-argument
 import torch
-import torch_directml # pylint: disable=import-error
-import modules.dml.hijack
-import modules.dml.amp as amp
-from modules.dml.opts import override_opts

-from .optimizer.unknown import UnknownOptimizer
+def directml_init():
+    from modules.dml.backend import DirectML # pylint: disable=ungrouped-imports
+    from modules.dml.opts import override_opts # pylint: disable=ungrouped-imports
+    # Alternative of torch.cuda for DirectML.
+    torch.dml = DirectML

-class DirectML():
-    is_autocast_enabled = False
-    autocast_gpu_dtype = torch.float16
-
-    def get_optimizer(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 .optimizer.nvidia import nVidiaOptimizer as optimizer
-            elif 'AMD' in device_name or 'Radeon' in device_name:
-                from .optimizer.amd import AMDOptimizer as optimizer
-            elif 'Intel' in device_name:
-                from .optimizer.intel import IntelOptimizer as optimizer
-            else:
-                return UnknownOptimizer
-            return optimizer
-        except Exception:
-            return UnknownOptimizer
-
-    def memory_stats(device: torch.device):
-        optimizer = DirectML.get_optimizer(device)
-        return optimizer.memory_stats(device.index)
-
-
-DirectML.amp = amp
-# Alternative of torch.cuda for DirectML.
-torch.dml = DirectML
-
-override_opts()
+    override_opts()
@@ -1,7 +0,0 @@
-from modules.dml.optimizer.optimizer import Optimizer
-from .driver.atiadlxx import ATIADLxx
-
-class AMDOptimizer(Optimizer):
-    driver: ATIADLxx = ATIADLxx()
-    def memory_stats(index):
-        return (AMDOptimizer.driver.iHyperMemorySize, AMDOptimizer.driver.get_dedicated_vram_usage(index))
@@ -1,46 +0,0 @@
-import ctypes as C
-from .atiadlxx_apis import *
-from .atiadlxx_structures import *
-from .atiadlxx_defines import *
-
-class ATIADLxx(object):
-    iHyperMemorySize = 0
-
-    def __init__(self):
-        self.context = ADL_CONTEXT_HANDLE()
-        ADL2_Main_Control_Create(ADL_Main_Memory_Alloc, 1, C.byref(self.context))
-        num_adapters = C.c_int(-1)
-        ADL2_Adapter_NumberOfAdapters_Get(self.context, C.byref(num_adapters))
-        AdapterInfoArray = (AdapterInfo * num_adapters.value)()
-        ADL2_Adapter_AdapterInfo_Get(self.context, C.cast(AdapterInfoArray, LPAdapterInfo), C.sizeof(AdapterInfoArray))
-        self.devices = []
-        busNumbers = []
-        for adapter in AdapterInfoArray:
-            if adapter.iBusNumber not in busNumbers: # filter duplicate device
-                self.devices.append(adapter)
-                busNumbers.append(adapter.iBusNumber)
-        self.iHyperMemorySize = self.get_memory_info2(0).iHyperMemorySize
-
-    def get_memory_info2(self, adapterIndex: int) -> ADLMemoryInfo2:
-        info = ADLMemoryInfo2()
-
-        if ADL2_Adapter_MemoryInfo2_Get(self.context, adapterIndex, C.byref(info)) != ADL_OK:
-            raise RuntimeError("ADL2: Failed to get MemoryInfo2")
-        
-        return info
-
-    def get_dedicated_vram_usage(self, index: int) -> int:
-        usage = C.c_int(-1)
-
-        if ADL2_Adapter_DedicatedVRAMUsage_Get(self.context, self.devices[index].iAdapterIndex, C.byref(usage)) != ADL_OK:
-            raise RuntimeError("ADL2: Failed to get DedicatedVRAMUsage")
-
-        return usage.value
-
-    def get_vram_usage(self, index: int) -> int:
-        usage = C.c_int(-1)
-
-        if ADL2_Adapter_VRAMUsage_Get(self.context, self.devices[index].iAdapterIndex, C.byref(usage)) != ADL_OK:
-            raise RuntimeError("ADL2: Failed to get VRAMUsage")
-
-        return usage.value
@@ -1,45 +0,0 @@
-import ctypes as C
-from platform import platform
-from .atiadlxx_structures import *
-
-if 'Windows' in platform():
-    atiadlxx = C.WinDLL("atiadlxx.dll")
-else:
-    atiadlxx = C.CDLL("libatiadlxx.so") # Not tested on Linux system. But will be supported.
-
-ADL_MAIN_MALLOC_CALLBACK = C.CFUNCTYPE(C.c_void_p, C.c_int)
-ADL_MAIN_FREE_CALLBACK = C.CFUNCTYPE(None, C.POINTER(C.c_void_p))
-
-@ADL_MAIN_MALLOC_CALLBACK
-def ADL_Main_Memory_Alloc(iSize):
-    return C._malloc(iSize)
-
-@ADL_MAIN_FREE_CALLBACK
-def ADL_Main_Memory_Free(lpBuffer):
-    if lpBuffer[0] is not None:
-        C._free(lpBuffer[0])
-        lpBuffer[0] = None
-
-ADL2_Main_Control_Create = atiadlxx.ADL2_Main_Control_Create
-ADL2_Main_Control_Create.restype = C.c_int
-ADL2_Main_Control_Create.argtypes = [ADL_MAIN_MALLOC_CALLBACK, C.c_int, ADL_CONTEXT_HANDLE]
-
-ADL2_Adapter_NumberOfAdapters_Get = atiadlxx.ADL2_Adapter_NumberOfAdapters_Get
-ADL2_Adapter_NumberOfAdapters_Get.restype = C.c_int
-ADL2_Adapter_NumberOfAdapters_Get.argtypes = [ADL_CONTEXT_HANDLE, C.POINTER(C.c_int)]
-
-ADL2_Adapter_AdapterInfo_Get = atiadlxx.ADL2_Adapter_AdapterInfo_Get
-ADL2_Adapter_AdapterInfo_Get.restype = C.c_int
-ADL2_Adapter_AdapterInfo_Get.argtypes = [ADL_CONTEXT_HANDLE, LPAdapterInfo, C.c_int]
-
-ADL2_Adapter_MemoryInfo2_Get = atiadlxx.ADL2_Adapter_MemoryInfo2_Get
-ADL2_Adapter_MemoryInfo2_Get.restype = C.c_int
-ADL2_Adapter_MemoryInfo2_Get.argtypes = [ADL_CONTEXT_HANDLE, C.c_int, C.POINTER(ADLMemoryInfo2)]
-
-ADL2_Adapter_DedicatedVRAMUsage_Get = atiadlxx.ADL2_Adapter_DedicatedVRAMUsage_Get
-ADL2_Adapter_DedicatedVRAMUsage_Get.restype = C.c_int
-ADL2_Adapter_DedicatedVRAMUsage_Get.argtypes = [ADL_CONTEXT_HANDLE, C.c_int, C.POINTER(C.c_int)]
-
-ADL2_Adapter_VRAMUsage_Get = atiadlxx.ADL2_Adapter_VRAMUsage_Get
-ADL2_Adapter_VRAMUsage_Get.restype = C.c_int
-ADL2_Adapter_VRAMUsage_Get.argtypes = [ADL_CONTEXT_HANDLE, C.c_int, C.POINTER(C.c_int)]
@@ -1 +0,0 @@
-ADL_OK = 0
@@ -1,87 +0,0 @@
-import ctypes as C
-
-class _ADLPMActivity(C.Structure):
-    __slot__ = [
-        'iActivityPercent',
-        'iCurrentBusLanes',
-        'iCurrentBusSpeed',
-        'iCurrentPerformanceLevel',
-        'iEngineClock',
-        'iMaximumBusLanes',
-        'iMemoryClock',
-        'iReserved',
-        'iSize',
-        'iVddc',
-    ]
-_ADLPMActivity._fields_ = [
-    ('iActivityPercent', C.c_int),
-    ('iCurrentBusLanes', C.c_int),
-    ('iCurrentBusSpeed', C.c_int),
-    ('iCurrentPerformanceLevel', C.c_int),
-    ('iEngineClock', C.c_int),
-    ('iMaximumBusLanes', C.c_int),
-    ('iMemoryClock', C.c_int),
-    ('iReserved', C.c_int),
-    ('iSize', C.c_int),
-    ('iVddc', C.c_int),
-]
-ADLPMActivity = _ADLPMActivity
-
-class _ADLMemoryInfo2(C.Structure):
-    __slot__ = [
-        'iHyperMemorySize',
-        'iInvisibleMemorySize',
-        'iMemoryBandwidth',
-        'iMemorySize',
-        'iVisibleMemorySize',
-        'strMemoryType'
-    ]
-_ADLMemoryInfo2._fields_ = [
-    ('iHyperMemorySize', C.c_longlong),
-    ('iInvisibleMemorySize', C.c_longlong),
-    ('iMemoryBandwidth', C.c_longlong),
-    ('iMemorySize', C.c_longlong),
-    ('iVisibleMemorySize', C.c_longlong),
-    ('strMemoryType', C.c_char * 256)
-]
-ADLMemoryInfo2 = _ADLMemoryInfo2
-
-class _AdapterInfo(C.Structure):
-    __slot__ = [
-        'iSize',
-        'iAdapterIndex',
-        'strUDID',
-        'iBusNumber',
-        'iDeviceNumber',
-        'iFunctionNumber',
-        'iVendorID',
-        'strAdapterName',
-        'strDisplayName',
-        'iPresent',
-        'iExist',
-        'strDriverPath',
-        'strDriverPathExt',
-        'strPNPString',
-        'iOSDisplayIndex',
-    ]
-_AdapterInfo._fields_ = [
-    ('iSize', C.c_int),
-    ('iAdapterIndex', C.c_int),
-    ('strUDID', C.c_char * 256),
-    ('iBusNumber', C.c_int),
-    ('iDeviceNumber', C.c_int),
-    ('iFunctionNumber', C.c_int),
-    ('iVendorID', C.c_int),
-    ('strAdapterName', C.c_char * 256),
-    ('strDisplayName', C.c_char * 256),
-    ('iPresent', C.c_int),
-    ('iExist', C.c_int),
-    ('strDriverPath', C.c_char * 256),
-    ('strDriverPathExt', C.c_char * 256),
-    ('strPNPString', C.c_char * 256),
-    ('iOSDisplayIndex', C.c_int)
-]
-AdapterInfo = _AdapterInfo
-LPAdapterInfo = C.POINTER(_AdapterInfo)
-
-ADL_CONTEXT_HANDLE = C.c_void_p
@@ -1,6 +0,0 @@
-from modules.dml.optimizer.optimizer import Optimizer
-
-class IntelOptimizer(Optimizer):
-    def memory_stats(index: int):
-        # DML TODO: Implement or find a general (and also lightweight) way.
-        return (1073741824, 0)
@@ -1,6 +0,0 @@
-from modules.dml.optimizer.optimizer import Optimizer
-
-class nVidiaOptimizer(Optimizer):
-    def memory_stats(index: int):
-        # DML TODO: Implement or find a general (and also lightweight) way.
-        return (1073741824, 0)
@@ -1,8 +0,0 @@
-from abc import *
-from typing import *
-
-class Optimizer(metaclass=ABCMeta):
-    driver: Any = None
-    @abstractmethod
-    def memory_stats(index: int) -> Tuple[int, int]:
-        pass
@@ -1,5 +0,0 @@
-from modules.dml.optimizer.optimizer import Optimizer
-
-class UnknownOptimizer(Optimizer):
-    def memory_stats(index: int):
-        return (1073741824, 0)
@@ -11,6 +11,7 @@ import tqdm
 import requests
 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_init
 import modules.interrogate
 import modules.memmon
 import modules.styles
@@ -766,8 +767,8 @@ batch_cond_uncond = opts.always_batch_cond_uncond or not (cmd_opts.lowvram or cm
 parallel_processing_allowed = not cmd_opts.lowvram and not cmd_opts.medvram
 mem_mon = modules.memmon.MemUsageMonitor("MemMon", device, opts)
 mem_mon.start()
-if device.type == 'privateuseone':
-    import modules.dml # pylint: disable=ungrouped-imports
+if devices.backend == "directml":
+    directml_init()


 def reload_gradio_theme(theme_name=None):