Implement torch.dml.

VERY UNSTABLE & NOT TESTED.
2026-09-19 17:24:32 +02:00 · 2023-04-26 12:21:44 +09:00
parent 8ea2fe15e7
commit 09ae33cdf7
19 changed files with 347 additions and 97 deletions
@@ -0,0 +1,32 @@
+import torch
+import torch_directml
+
+import modules.dml.kdiffusion
+import modules.dml.stablediffusion
+import modules.dml.torch
+
+from optimizer.unknown import UnknownOptimizer
+
+class DirectML():
+    def get_optimizer(self, 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:
+            return UnknownOptimizer
+
+    def memory_stats(self, device: torch.device):
+        optimizer = self.get_optimizer(device)
+        return optimizer.memory_stats(device.index)
+
+# Alternative of torch.cuda for DirectML.
+torch.dml = DirectML
@@ -2,9 +2,6 @@ import torch
 from tqdm.auto import tqdm

 from modules.shared import device
-from modules.sd_hijack_utils import CondFunc
-
-# k-diffusion
 from k_diffusion import sampling

 def dpm_solver_adaptive(self, x, t_start, t_end, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None):
@@ -89,88 +86,4 @@ def sample_dpm_adaptive(model, x, sigma_min, sigma_max, extra_args=None, callbac

 sampling.DPMSolver.dpm_solver_adaptive = dpm_solver_adaptive
 sampling.sample_dpm_fast = sample_dpm_fast
-sampling.sample_dpm_adaptive = sample_dpm_adaptive
-
-# stablediffusion
-from ldm.models.diffusion.ddim import DDIMSampler
-from ldm.modules.diffusionmodules.util import noise_like
-
-@torch.no_grad()
-def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
-                    temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
-                    unconditional_guidance_scale=1., unconditional_conditioning=None,
-                    dynamic_threshold=None):
-    b, *_, device = *x.shape, x.device
-
-    if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
-        model_output = self.model.apply_model(x, t, c)
-    else:
-        x_in = torch.cat([x] * 2)
-        t_in = torch.cat([t] * 2)
-        if isinstance(c, dict):
-            assert isinstance(unconditional_conditioning, dict)
-            c_in = dict()
-            for k in c:
-                if isinstance(c[k], list):
-                    c_in[k] = [torch.cat([
-                        unconditional_conditioning[k][i],
-                        c[k][i]]) for i in range(len(c[k]))]
-                else:
-                    c_in[k] = torch.cat([
-                            unconditional_conditioning[k],
-                            c[k]])
-        elif isinstance(c, list):
-            c_in = list()
-            assert isinstance(unconditional_conditioning, list)
-            for i in range(len(c)):
-                c_in.append(torch.cat([unconditional_conditioning[i], c[i]]))
-        else:
-            c_in = torch.cat([unconditional_conditioning, c])
-        model_uncond, model_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
-        model_output = model_uncond + unconditional_guidance_scale * (model_t - model_uncond)
-
-    if self.model.parameterization == "v":
-        e_t = self.model.predict_eps_from_z_and_v(x, t, model_output)
-    else:
-        e_t = model_output
-
-    if score_corrector is not None:
-        assert self.model.parameterization == "eps", 'not implemented'
-        e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
-
-    alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
-    alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
-    sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
-    sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
-    # select parameters corresponding to the currently considered timestep
-    print(alphas[index]) # DML Solution: DDIM Sampling does not work without this print.
-    a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
-    a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
-    sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
-    sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
-
-    # current prediction for x_0
-    if self.model.parameterization != "v":
-        pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
-    else:
-        pred_x0 = self.model.predict_start_from_z_and_v(x, t, model_output)
-
-    if quantize_denoised:
-        pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
-
-    if dynamic_threshold is not None:
-        raise NotImplementedError()
-
-    # direction pointing to x_t
-    dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
-    noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
-    if noise_dropout > 0.:
-        noise = torch.nn.functional.dropout(noise, p=noise_dropout)
-    x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
-    return x_prev, pred_x0
-
-DDIMSampler.p_sample_ddim = p_sample_ddim
-
-# torch
-
-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')
+sampling.sample_dpm_adaptive = sample_dpm_adaptive
@@ -0,0 +1,7 @@
+from modules.dml.optimizer.optimizer import Optimizer
+from driver.atiadlxx import ATIADLxx
+
+class AMDOptimizer(Optimizer):
+    driver: ATIADLxx = ATIADLxx()
+    def memory_stats(self, index):
+        return (AMDOptimizer.driver.iHyperMemorySize, AMDOptimizer.driver.get_dedicated_vram_usage(index))
@@ -0,0 +1,43 @@
+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 = []
+        for adapter in AdapterInfoArray:
+            self.devices.append(adapter)
+        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
@@ -0,0 +1,50 @@
+import ctypes as C
+import platform
+from .atiadlxx_structures import *
+
+_platform = platform.system()
+
+try:
+    if _platform == "Windows":
+        atiadlxx = C.WinDLL("atiadlxx.dll")
+
+        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)]
+    else:
+        print("Warning: experimental graphic memory optimization for AMDGPU is disabled. Because this is not Windows platform.")
+except FileNotFoundError:
+    print("Warning: memory optimization for AMDGPU is disabled. Because couldn't find 'atiadlxx.dll'. Please install GPU driver downloaded from AMD.com.")
@@ -0,0 +1 @@
+ADL_OK = 0
@@ -0,0 +1,87 @@
+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
@@ -0,0 +1,7 @@
+from modules.dml.optimizer.optimizer import Optimizer
+
+class IntelOptimizer(Optimizer):
+    def memory_stats():
+        raise NotImplementedError()
+        # DML TODO: Implement
+        return
@@ -0,0 +1,7 @@
+from modules.dml.optimizer.optimizer import Optimizer
+
+class nVidiaOptimizer(Optimizer):
+    def memory_stats():
+        raise NotImplementedError()
+        # DML TODO: Implement
+        return
@@ -0,0 +1,8 @@
+from abc import *
+from typing import *
+
+class Optimizer(metaclass=ABCMeta):
+    driver: Any = None
+    @abstractmethod
+    def memory_stats(self, index: int) -> Tuple[int, int]:
+        pass
@@ -0,0 +1,6 @@
+from modules.dml.optimizer.optimizer import Optimizer
+
+class UnknownOptimizer(Optimizer):
+    def memory_stats():
+        # DML TODO: Implement
+        return (1073741824, 0)
@@ -0,0 +1,80 @@
+import torch
+
+from ldm.models.diffusion.ddim import DDIMSampler
+from ldm.modules.diffusionmodules.util import noise_like
+
+@torch.no_grad()
+def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
+                    temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
+                    unconditional_guidance_scale=1., unconditional_conditioning=None,
+                    dynamic_threshold=None):
+    b, *_, device = *x.shape, x.device
+
+    if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
+        model_output = self.model.apply_model(x, t, c)
+    else:
+        x_in = torch.cat([x] * 2)
+        t_in = torch.cat([t] * 2)
+        if isinstance(c, dict):
+            assert isinstance(unconditional_conditioning, dict)
+            c_in = dict()
+            for k in c:
+                if isinstance(c[k], list):
+                    c_in[k] = [torch.cat([
+                        unconditional_conditioning[k][i],
+                        c[k][i]]) for i in range(len(c[k]))]
+                else:
+                    c_in[k] = torch.cat([
+                            unconditional_conditioning[k],
+                            c[k]])
+        elif isinstance(c, list):
+            c_in = list()
+            assert isinstance(unconditional_conditioning, list)
+            for i in range(len(c)):
+                c_in.append(torch.cat([unconditional_conditioning[i], c[i]]))
+        else:
+            c_in = torch.cat([unconditional_conditioning, c])
+        model_uncond, model_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
+        model_output = model_uncond + unconditional_guidance_scale * (model_t - model_uncond)
+
+    if self.model.parameterization == "v":
+        e_t = self.model.predict_eps_from_z_and_v(x, t, model_output)
+    else:
+        e_t = model_output
+
+    if score_corrector is not None:
+        assert self.model.parameterization == "eps", 'not implemented'
+        e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
+
+    alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+    alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
+    sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
+    sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
+    # select parameters corresponding to the currently considered timestep
+    print(alphas[index]) # DML Solution: DDIM Sampling does not work without this print.
+    a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
+    a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
+    sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
+    sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
+
+    # current prediction for x_0
+    if self.model.parameterization != "v":
+        pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+    else:
+        pred_x0 = self.model.predict_start_from_z_and_v(x, t, model_output)
+
+    if quantize_denoised:
+        pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
+
+    if dynamic_threshold is not None:
+        raise NotImplementedError()
+
+    # direction pointing to x_t
+    dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
+    noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+    if noise_dropout > 0.:
+        noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+    x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+    return x_prev, pred_x0
+
+DDIMSampler.p_sample_ddim = p_sample_ddim
@@ -0,0 +1,5 @@
+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')
@@ -20,6 +20,9 @@ if shared.opts.cross_attention_optimization == "xFormers":
    except Exception:
        pass

+if shared.device.type == 'privateuseone':
+    import dml
+

 def get_available_vram():
    if shared.device.type == 'cuda':
@@ -31,8 +34,8 @@ def get_available_vram():
        mem_free_total = mem_free_cuda + mem_free_torch
        return mem_free_total
    elif shared.device.type == 'privateuseone':
-        # DML ISSUE: There's no way to get any memory info.
-        return 1073741824
+        mem_total, mem_active = torch.dml.memory_stats(shared.device)
+        return mem_total - mem_active * (1 << 20)
    else:
        return psutil.virtual_memory().available

@@ -199,8 +202,9 @@ def einsum_op_cuda(q, k, v):
    return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20))

 def einsum_op_dml(q, k, v):
-    # DML ISSUE: There's no way to get any memory info.
-    return einsum_op_tensor_mem(q, k, v, 1073741824)
+    mem_total, mem_active = devices.adl.memory_stats()
+    mem_reserved = mem_total / (1 << 20) * 0.7
+    return einsum_op_tensor_mem(q, k, v, (mem_reserved - mem_active) if mem_reserved > mem_active else 1)

 def einsum_op(q, k, v):
    if q.device.type == 'cuda':
@@ -64,7 +64,7 @@ clip_model = None


 if device.type == 'privateuseone':
-    import modules.dml_specific
+    import modules.dml
    is_device_dml = True


@@ -429,7 +429,7 @@ options_templates.update(options_section(('ui', "Live previews"), {
    "live_previews_enable": OptionInfo(True, "Show live previews of the created image"),
    "show_progress_grid": OptionInfo(True, "Show previews of all images generated in a batch as a grid"),
    "show_progress_every_n_steps": OptionInfo(1, "Show new live preview image every N sampling steps. Set to -1 to show after completion of batch.", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),
-    "show_progress_type": OptionInfo("Approx NN", "Image creation progress preview mode", gr.Radio, {"choices": ["Full", "Approx NN", "Approx cheap"]}), # DML ISSUE: Approx NN does not work well on DirectML device.
+    "show_progress_type": OptionInfo("Approx cheap" if is_device_dml else "Approx NN", "Image creation progress preview mode", gr.Radio, {"choices": ["Full", "Approx NN", "Approx cheap"]}), # DML Solution: Use Approx cheap instead of Approx NN as a default progress type.
    "live_preview_content": OptionInfo("Combined", "Live preview subject", gr.Radio, {"choices": ["Combined", "Prompt", "Negative prompt"]}),
    "live_preview_refresh_period": OptionInfo(250, "Progressbar/preview update period, in milliseconds")
 }))