Merge branch 'master' of https://github.com/vladmandic/automatic into settings-improvements

.github/workflows/on_pull_request.yaml

@@ -24,15 +24,12 @@ jobs:
           python-version: 3.10.6
           cache: pip
           cache-dependency-path: requirements.txt
-      - name: Install PyLint
-        run: | 
+      - name: Test Startup
+        run: |
+          export COMMANDLINE_ARGS="--debug --test"
+          python launch.py
+      - name: Linting
+        run: |
           python -m pip install --upgrade pip
           pip install pylint
-      # This lets PyLint check to see if it can resolve imports
-      - name: Install dependencies
-        run: |
-          export COMMANDLINE_ARGS="--skip-torch-cuda-test --exit"
-          python launch.py
-      - name: Analysing the code with pylint
-        run: |
           pylint $(git ls-files '*.py')

.pylintrc

@@ -11,6 +11,7 @@ fail-under=10
 ignore=CVS
 ignore-paths=^repositories/.*$,
              ^extensions/.*$,
+             ^extensions-builtin/.*$,
              /usr/lib/.*$,
 ignore-patterns=
 ignored-modules=
@@ -141,6 +142,8 @@ disable=raw-checker-failed,
         consider-using-dict-items,
         dangerous-default-value,
         unnecessary-dunder-call,
+        invalid-name,
+        R0801,
 enable=c-extension-no-member
 
 [METHOD_ARGS]

TODO.md

@@ -4,10 +4,7 @@
 
 Stuff to be fixed...
 
-- Run VAE with hires at 1280
-- Transformers version
 - Move Restart Server from WebUI to Launch and reload modules
-- Follow-up on `p.script_args`
 - Mdularize `cli` scripts
 
 ## Features
@@ -16,9 +13,11 @@ Stuff to be added...
 
 - Update README
 - Add Gradio theme maker
+- Transformers version
 - Create new GitHub hooks/actions for CI/CD  
 - Redo Extensions tab: see <https://vladmandic.github.io/sd-extension-manager/pages/extensions.html>
 - Stream-load models as option for slow storage
+- Auto-test `torch.layer_norm` for FP16
 
 ## Investigate
 

javascript/imageviewer.js

@@ -105,8 +105,8 @@ function setupImageForLightbox(e) {
 	var isFirefox = navigator.userAgent.toLowerCase().indexOf('firefox') > -1
 	var event = isFirefox ? 'mousedown' : 'click'
 	e.addEventListener(event, function (evt) {
-		if(!opts.js_modal_lightbox || evt.button != 0) return;
-		modalZoomSet(gradioApp().getElementById('modalImage'), opts.js_modal_lightbox_initially_zoomed)
+		if (evt.button != 0) return;
+		modalZoomSet(gradioApp().getElementById('modalImage'), true)
 		evt.preventDefault()
 		showModal(evt)
 	}, true);

launch.py

@@ -97,5 +97,9 @@ if __name__ == "__main__":
     setup.log.info(f"Server arguments: {sys.argv[1:]}")
     setup.log.debug('Starting WebUI')
     logging.disable(logging.INFO)
+    if args.test:
+        setup.log.info("Test only")
+        import webui
+        exit(0)
     import webui
     webui.webui()

modules/cmd_args.py

@@ -74,6 +74,7 @@ def compatibility_args(opts, args):
     parser.add_argument("--sub-quad-q-chunk-size", help=argparse.SUPPRESS, default=opts.sub_quad_q_chunk_size)
     parser.add_argument("--sub-quad-kv-chunk-size", help=argparse.SUPPRESS, default=opts.sub_quad_kv_chunk_size)
     parser.add_argument("--sub-quad-chunk-threshold", help=argparse.SUPPRESS, default=opts.sub_quad_chunk_threshold)
+    parser.add_argument('--debug', default = False, action='store_true', help = "Run installer with debug logging, default: %(default)s")
 
     opts.use_old_emphasis_implementation = False
     opts.use_old_karras_scheduler_sigmas = False

modules/devices.py

@@ -27,12 +27,26 @@ def get_cuda_device_string():
     return "cuda"
 
 
+def get_dml_device_string():
+    from modules import shared
+    if shared.cmd_opts.device_id is not None:
+        return f"privateuseone:{shared.cmd_opts.device_id}"
+    return "privateuseone:0"
+
+
 def get_optimal_device_name():
     if torch.cuda.is_available():
         return get_cuda_device_string()
     if has_mps():
         return "mps"
-    return "cpu"
+    try:
+        import torch_directml
+        if torch_directml.is_available():
+            return get_dml_device_string()
+        else:
+            return "cpu"
+    except:
+        return "cpu"
 
 
 def get_optimal_device():

modules/dml/__init__.py

@@ -0,0 +1,30 @@
+import torch
+import torch_directml
+
+import modules.dml.hijack
+
+from .optimizer.unknown import UnknownOptimizer
+
+class DirectML():
+    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:
+            return UnknownOptimizer
+
+    def memory_stats(device: torch.device):
+        optimizer = DirectML.get_optimizer(device)
+        return optimizer.memory_stats(device.index)
+
+# Alternative of torch.cuda for DirectML.
+torch.dml = DirectML

modules/dml/hijack/__init__.py

@@ -0,0 +1,5 @@
+import modules.dml.hijack.kdiffusion
+import modules.dml.hijack.stablediffusion
+import modules.dml.hijack.torch
+import modules.dml.hijack.realesrgan_model
+import modules.dml.hijack.plms

modules/dml/hijack/kdiffusion.py

@@ -0,0 +1,89 @@
+import torch
+from tqdm.auto import tqdm
+
+from modules.shared import device
+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):
+    noise_sampler = sampling.default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    if order not in {2, 3}:
+        raise ValueError('order should be 2 or 3')
+    forward = t_end > t_start
+    if not forward and eta:
+        raise ValueError('eta must be 0 for reverse sampling')
+    h_init = abs(h_init) * (1 if forward else -1)
+    atol = torch.tensor(atol).to(device)
+    rtol = torch.tensor(rtol).to(device)
+    s = t_start
+    x_prev = x
+    accept = True
+    pid = sampling.PIDStepSizeController(h_init, pcoeff, icoeff, dcoeff, 1.5 if eta else order, accept_safety)
+    info = {'steps': 0, 'nfe': 0, 'n_accept': 0, 'n_reject': 0}
+
+    while s < t_end - 1e-5 if forward else s > t_end + 1e-5:
+        eps_cache = {}
+        t = torch.minimum(t_end, s + pid.h) if forward else torch.maximum(t_end, s + pid.h)
+        if eta:
+            sd, su = sampling.get_ancestral_step(self.sigma(s), self.sigma(t), eta)
+            t_ = torch.minimum(t_end, self.t(sd))
+            su = (self.sigma(t) ** 2 - self.sigma(t_) ** 2) ** 0.5
+        else:
+            t_, su = t, 0.
+
+        eps, eps_cache = self.eps(eps_cache, 'eps', x, s)
+        denoised = x - self.sigma(s) * eps
+
+        if order == 2:
+            x_low, eps_cache = self.dpm_solver_1_step(x, s, t_, eps_cache=eps_cache)
+            x_high, eps_cache = self.dpm_solver_2_step(x, s, t_, eps_cache=eps_cache)
+        else:
+            x_low, eps_cache = self.dpm_solver_2_step(x, s, t_, r1=1 / 3, eps_cache=eps_cache)
+            x_high, eps_cache = self.dpm_solver_3_step(x, s, t_, eps_cache=eps_cache)
+        delta = torch.maximum(atol, rtol * torch.maximum(x_low.abs(), x_prev.abs()))
+        error = torch.linalg.norm((x_low - x_high) / delta) / x.numel() ** 0.5
+        accept = pid.propose_step(error)
+        if accept:
+            x_prev = x_low
+            x = x_high + su * s_noise * noise_sampler(self.sigma(s), self.sigma(t))
+            s = t
+            info['n_accept'] += 1
+        else:
+            info['n_reject'] += 1
+        info['nfe'] += order
+        info['steps'] += 1
+
+        if self.info_callback is not None:
+            self.info_callback({'x': x, 'i': info['steps'] - 1, 't': s, 't_up': s, 'denoised': denoised, 'error': error, 'h': pid.h, **info})
+
+    return x, info
+
+
+@torch.no_grad()
+def sample_dpm_fast(model, x, sigma_min, sigma_max, n, extra_args=None, callback=None, disable=None, eta=0., s_noise=1., noise_sampler=None):
+    """DPM-Solver-Fast (fixed step size). See https://arxiv.org/abs/2206.00927."""
+    if sigma_min <= 0 or sigma_max <= 0:
+        raise ValueError('sigma_min and sigma_max must not be 0')
+    with tqdm(total=n, disable=disable) as pbar:
+        dpm_solver = sampling.DPMSolver(model, extra_args, eps_callback=pbar.update)
+        if callback is not None:
+            dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
+        return dpm_solver.dpm_solver_fast(x, dpm_solver.t(torch.tensor(sigma_max).to(device)), dpm_solver.t(torch.tensor(sigma_min).to(device)), n, eta, s_noise, noise_sampler)
+
+
+@torch.no_grad()
+def sample_dpm_adaptive(model, x, sigma_min, sigma_max, extra_args=None, callback=None, disable=None, 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, return_info=False):
+    """DPM-Solver-12 and 23 (adaptive step size). See https://arxiv.org/abs/2206.00927."""
+    if sigma_min <= 0 or sigma_max <= 0:
+        raise ValueError('sigma_min and sigma_max must not be 0')
+    with tqdm(disable=disable) as pbar:
+        dpm_solver = sampling.DPMSolver(model, extra_args, eps_callback=pbar.update)
+        if callback is not None:
+            dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
+        x, info = dpm_solver.dpm_solver_adaptive(x, dpm_solver.t(torch.tensor(sigma_max).to(device)), dpm_solver.t(torch.tensor(sigma_min).to(device)), order, rtol, atol, h_init, pcoeff, icoeff, dcoeff, accept_safety, eta, s_noise, noise_sampler)
+    if return_info:
+        return x, info
+    return x
+
+sampling.DPMSolver.dpm_solver_adaptive = dpm_solver_adaptive
+sampling.sample_dpm_fast = sample_dpm_fast
+sampling.sample_dpm_adaptive = sample_dpm_adaptive

modules/dml/hijack/plms.py

@@ -0,0 +1,91 @@
+import torch
+
+from ldm.models.diffusion.ddim import noise_like
+
+import modules.sd_hijack_inpainting as plms_hijack
+
+
+@torch.no_grad()
+def p_sample_plms(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, old_eps=None, t_next=None, dynamic_threshold=None):
+    b, *_, device = *x.shape, x.device
+
+    def get_model_output(x, t):
+        if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
+            e_t = 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]])
+            else:
+                c_in = torch.cat([unconditional_conditioning, c])
+
+            e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
+            e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
+
+        if score_corrector is not None:
+            assert self.model.parameterization == "eps"
+            e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
+
+        return e_t
+
+    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
+
+    def get_x_prev_and_pred_x0(e_t, index):
+        # select parameters corresponding to the currently considered timestep
+        print(alphas[index]) # DML Solution: PLMS 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
+        pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+        if quantize_denoised:
+            pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
+        if dynamic_threshold is not None:
+            from ldm.models.diffusion.sampling_util import norm_thresholding
+            pred_x0 = norm_thresholding(pred_x0, dynamic_threshold)
+        # 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
+
+    e_t = get_model_output(x, t)
+    if len(old_eps) == 0:
+        # Pseudo Improved Euler (2nd order)
+        x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
+        e_t_next = get_model_output(x_prev, t_next)
+        e_t_prime = (e_t + e_t_next) / 2
+    elif len(old_eps) == 1:
+        # 2nd order Pseudo Linear Multistep (Adams-Bashforth)
+        e_t_prime = (3 * e_t - old_eps[-1]) / 2
+    elif len(old_eps) == 2:
+        # 3nd order Pseudo Linear Multistep (Adams-Bashforth)
+        e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
+    elif len(old_eps) >= 3:
+        # 4nd order Pseudo Linear Multistep (Adams-Bashforth)
+        e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24
+
+    x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
+
+    return x_prev, pred_x0, e_t
+plms_hijack.p_sample_plms = p_sample_plms

modules/dml/hijack/realesrgan_model.py

@@ -0,0 +1,69 @@
+import math
+import torch
+
+from realesrgan import RealESRGANer
+
+
+# DML Solution: Some tensors turn 0 after Extended Slices. Move output to cpu and get it back.
+def tile_process(self):
+    batch, channel, height, width = self.img.shape
+    output_height = height * self.scale
+    output_width = width * self.scale
+    output_shape = (batch, channel, output_height, output_width)
+
+    # start with black image
+    self.output = self.img.new_zeros(output_shape, device='cpu')
+    tiles_x = math.ceil(width / self.tile_size)
+    tiles_y = math.ceil(height / self.tile_size)
+
+    # loop over all tiles
+    for y in range(tiles_y):
+        for x in range(tiles_x):
+            # extract tile from input image
+            ofs_x = x * self.tile_size
+            ofs_y = y * self.tile_size
+            # input tile area on total image
+            input_start_x = ofs_x
+            input_end_x = min(ofs_x + self.tile_size, width)
+            input_start_y = ofs_y
+            input_end_y = min(ofs_y + self.tile_size, height)
+
+            # input tile area on total image with padding
+            input_start_x_pad = max(input_start_x - self.tile_pad, 0)
+            input_end_x_pad = min(input_end_x + self.tile_pad, width)
+            input_start_y_pad = max(input_start_y - self.tile_pad, 0)
+            input_end_y_pad = min(input_end_y + self.tile_pad, height)
+
+            # input tile dimensions
+            input_tile_width = input_end_x - input_start_x
+            input_tile_height = input_end_y - input_start_y
+            tile_idx = y * tiles_x + x + 1
+            input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad]
+
+            # upscale tile
+            try:
+                with torch.no_grad():
+                    output_tile = self.model(input_tile)
+                    output_tile = output_tile.cpu()
+            except RuntimeError as error:
+                print('Error', error)
+            print(f'\tTile {tile_idx}/{tiles_x * tiles_y}')
+
+            # output tile area on total image
+            output_start_x = input_start_x * self.scale
+            output_end_x = input_end_x * self.scale
+            output_start_y = input_start_y * self.scale
+            output_end_y = input_end_y * self.scale
+
+            # output tile area without padding
+            output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale
+            output_end_x_tile = output_start_x_tile + input_tile_width * self.scale
+            output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale
+            output_end_y_tile = output_start_y_tile + input_tile_height * self.scale
+
+            # put tile into output image
+            self.output[:, :, output_start_y:output_end_y,
+                        output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile,
+                                                                    output_start_x_tile:output_end_x_tile]
+    self.output = self.output.to(self.device)
+RealESRGANer.tile_process = tile_process

modules/dml/hijack/stablediffusion.py

@@ -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

modules/dml/hijack/torch.py

@@ -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')

modules/dml/optimizer/amd/__init__.py

@@ -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(index):
+        return (AMDOptimizer.driver.iHyperMemorySize, AMDOptimizer.driver.get_dedicated_vram_usage(index))

modules/dml/optimizer/amd/driver/atiadlxx.py

@@ -0,0 +1,46 @@
+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

modules/dml/optimizer/amd/driver/atiadlxx_apis.py

@@ -0,0 +1,45 @@
+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)]

modules/dml/optimizer/amd/driver/atiadlxx_defines.py

@@ -0,0 +1 @@
+ADL_OK = 0

modules/dml/optimizer/amd/driver/atiadlxx_structures.py

@@ -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

modules/dml/optimizer/intel/__init__.py

@@ -0,0 +1,7 @@
+from modules.dml.optimizer.optimizer import Optimizer
+
+class IntelOptimizer(Optimizer):
+    def memory_stats(index):
+        raise NotImplementedError()
+        # DML TODO: Implement
+        return

modules/dml/optimizer/nvidia/__init__.py

@@ -0,0 +1,7 @@
+from modules.dml.optimizer.optimizer import Optimizer
+
+class nVidiaOptimizer(Optimizer):
+    def memory_stats(index):
+        raise NotImplementedError()
+        # DML TODO: Implement
+        return

modules/dml/optimizer/optimizer.py

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

modules/dml/optimizer/unknown/__init__.py

@@ -0,0 +1,6 @@
+from modules.dml.optimizer.optimizer import Optimizer
+
+class UnknownOptimizer(Optimizer):
+    def memory_stats(index):
+        # DML TODO: Implement
+        return (1073741824, 0)

modules/realesrgan_model.py

@@ -6,7 +6,7 @@ from PIL import Image
 from basicsr.utils.download_util import load_file_from_url
 
 from modules.upscaler import Upscaler, UpscalerData
-from modules.shared import cmd_opts, opts
+from modules.shared import cmd_opts, opts, device
 import modules.errors as errors
 
 
@@ -53,6 +53,7 @@ class UpscalerRealESRGAN(Upscaler):
             half=not cmd_opts.no_half and not opts.upcast_sampling,
             tile=opts.ESRGAN_tile,
             tile_pad=opts.ESRGAN_tile_overlap,
+            device=device,
         )
 
         upsampled = upsampler.enhance(np.array(img), outscale=info.scale)[0]

modules/sd_hijack_optimizations.py

@@ -30,6 +30,9 @@ def get_available_vram():
         mem_free_torch = mem_reserved - mem_active
         mem_free_total = mem_free_cuda + mem_free_torch
         return mem_free_total
+    elif shared.device.type == 'privateuseone':
+        mem_total, mem_active = torch.dml.memory_stats(shared.device)
+        return mem_total - mem_active * (1 << 20)
     else:
         return psutil.virtual_memory().available
 
@@ -195,6 +198,11 @@ def einsum_op_cuda(q, k, v):
     # Divide factor of safety as there's copying and fragmentation
     return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20))
 
+def einsum_op_dml(q, k, v):
+    mem_total, mem_active = torch.dml.memory_stats(q.device)
+    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':
         return einsum_op_cuda(q, k, v)
@@ -204,6 +212,9 @@ def einsum_op(q, k, v):
             return einsum_op_mps_v1(q, k, v)
         return einsum_op_mps_v2(q, k, v)
 
+    if q.device.type == 'privateuseone':
+        return einsum_op_dml(q, k, v)
+
     # Smaller slices are faster due to L2/L3/SLC caches.
     # Tested on i7 with 8MB L3 cache.
     return einsum_op_tensor_mem(q, k, v, 32)

modules/shared.py

@@ -52,11 +52,17 @@ ui_reorder_categories = [
 cmd_opts.disable_extension_access = (cmd_opts.share or cmd_opts.listen or cmd_opts.server_name) and not cmd_opts.enable_insecure
 devices.device, devices.device_interrogate, devices.device_gfpgan, devices.device_esrgan, devices.device_codeformer = (devices.cpu if any(y in cmd_opts.use_cpu for y in [x, 'all']) else devices.get_optimal_device() for x in ['sd', 'interrogate', 'gfpgan', 'esrgan', 'codeformer'])
 device = devices.device
+is_device_dml = False
 sd_upscalers = []
 sd_model = None
 clip_model = None
 
 
+if device.type == 'privateuseone':
+    import modules.dml
+    is_device_dml = True
+
+
 def reload_hypernetworks():
     from modules.hypernetworks import hypernetwork
     global hypernetworks # pylint: disable=W0603
@@ -232,7 +238,7 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "comma_padding_backtrack": OptionInfo(20, "Increase coherency by padding from the last comma within n tokens when using more than 75 tokens", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1 }),
     "CLIP_stop_at_last_layers": OptionInfo(1, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1, "visible": False}),
     "upcast_attn": OptionInfo(False, "Upcast cross attention layer to float32"),
-    "cross_attention_optimization": OptionInfo("Scaled-Dot-Product", "Cross-attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention() }),
+    "cross_attention_optimization": OptionInfo("Sub-quadratic" if is_device_dml else "Scaled-Dot-Product", "Cross-attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention() }),
     "cross_attention_options": OptionInfo([], "Cross-attention advanced options", gr.CheckboxGroup, lambda: {"choices": ['xFormers enable flash Attention', 'SDP disable memory attention']}),
     "sub_quad_q_chunk_size": OptionInfo(512, "Sub-quadratic cross-attention query chunk size for the  layer optimization to use", gr.Slider, {"minimum": 16, "maximum": 8192, "step": 8}),
     "sub_quad_kv_chunk_size": OptionInfo(512, "Sub-quadratic cross-attentionkv chunk size for the sub-quadratic cross-attention layer optimization to use", gr.Slider, {"minimum": 0, "maximum": 8192, "step": 8}),
@@ -310,8 +316,8 @@ options_templates.update(options_section(('cuda', "CUDA Settings"), {
     "memmon_poll_rate": OptionInfo(2, "VRAM usage polls per second during generation. Set to 0 to disable.", gr.Slider, {"minimum": 0, "maximum": 40, "step": 1}),
     "precision": OptionInfo("Autocast", "Precision type", gr.Radio, lambda: {"choices": ["Autocast", "Full"]}),
     "cuda_dtype": OptionInfo("FP32" if sys.platform == "darwin" else "FP16", "Device precision type", gr.Radio, lambda: {"choices": ["FP32", "FP16", "BF16"]}),
-    "no_half": OptionInfo(False, "Use full precision for model (--no-half)"),
-    "no_half_vae": OptionInfo(False, "Use full precision for VAE (--no-half-vae)"),
+    "no_half": OptionInfo(True if is_device_dml else False, "Use full precision for model (--no-half)", None, None, lambda: print("Warning: Most of DirectML devices do not fully support half mode. Recommend to use full precision to model.") if is_device_dml else None),
+    "no_half_vae": OptionInfo(True if is_device_dml else False, "Use full precision for VAE (--no-half-vae)"),
     "upcast_sampling": OptionInfo(True if sys.platform == "darwin" else False, "Enable upcast sampling. Usually produces similar results to --no-half with better performance while using less memory"),
     "disable_nan_check": OptionInfo(True, "Do not check if produced images/latent spaces have NaN values"),
     "rollback_vae": OptionInfo(False, "Attempt to roll back VAE when produced NaN values, requires NaN check (experimental)"),

setup.py

@@ -4,6 +4,7 @@ import json
 import time
 import shutil
 import logging
+import platform
 import subprocess
 
 try:
@@ -20,7 +21,7 @@ class Dot(dict): # dot notation access to dictionary attributes
 
 
 log = logging.getLogger("sd")
-args = Dot({ 'debug': False, 'upgrade': False, 'noupdate': False, 'skip-extensions': False, 'skip-requirements': False, 'reset': False })
+args = Dot({ 'debug': False, 'upgrade': False, 'noupdate': False, 'nodirectml': False, 'skip-extensions': False, 'skip-requirements': False, 'reset': False })
 quick_allowed = True
 errors = 0
 opts = {}
@@ -197,17 +198,21 @@ def check_torch():
         torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.4.2')
         xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
     else:
-        log.info('Using CPU-only Torch')
-        torch_command = os.environ.get('TORCH_COMMAND', 'torch torchaudio torchvision')
-        xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
+        machine = platform.machine()
+        if 'arm' not in machine and 'aarch' not in machine and not args.nodirectml: # torch-directml is available on AMD64
+            log.info('Using DirectML Backend')
+            torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.0 torchvision torch-directml')
+            xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
+        else:
+            log.info('Using CPU-only Torch')
+            torch_command = os.environ.get('TORCH_COMMAND', 'torch torchaudio torchvision')
+            xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
     if 'torch' in torch_command:
         install(torch_command, 'torch torchvision torchaudio')
     try:
         import torch
         log.info(f'Torch {torch.__version__}')
-        if not torch.cuda.is_available():
-            log.warning("Torch repoorts CUDA not available")
-        else:
+        if torch.cuda.is_available():
             if torch.version.cuda:
                 log.info(f'Torch backend: nVidia CUDA {torch.version.cuda} cuDNN {torch.backends.cudnn.version() if torch.backends.cudnn.is_available() else "N/A"}')
             elif torch.version.hip:
@@ -216,6 +221,17 @@ def check_torch():
                 log.warning('Unknown Torch backend')
             for device in [torch.cuda.device(i) for i in range(torch.cuda.device_count())]:
                 log.info(f'Torch detected GPU: {torch.cuda.get_device_name(device)} VRAM {round(torch.cuda.get_device_properties(device).total_memory / 1024 / 1024)} Arch {torch.cuda.get_device_capability(device)} Cores {torch.cuda.get_device_properties(device).multi_processor_count}')
+        else:
+            try:
+                import torch_directml
+                import pkg_resources
+                version = pkg_resources.get_distribution("torch-directml")
+                log.info(f'Torch backend: DirectML ({version})')
+                for i in range(0, torch_directml.device_count()):
+                    log.info(f'Torch detected GPU: {torch_directml.device_name(i)}')
+                log.info(f'DirectML default device: {torch_directml.device_name(torch_directml.default_device())}')
+            except:
+                log.warning("Torch repoorts CUDA not available")
     except Exception as e:
         log.error(f'Could not load torch: {e}')
         exit(1)
@@ -487,15 +503,16 @@ def check_timestamp():
 
 
 def add_args():
-    if vars(parser)['_option_string_actions'].get('--debug', None) is None:
-        parser.add_argument('--debug', default = False, action='store_true', help = "Run installer with debug logging, default: %(default)s")
+    parser.add_argument('--debug', default = False, action='store_true', help = "Run installer with debug logging, default: %(default)s")
     parser.add_argument('--reset', default = False, action='store_true', help = "Reset main repository to latest version, default: %(default)s")
     parser.add_argument('--upgrade', default = False, action='store_true', help = "Upgrade main repository to latest version, default: %(default)s")
     parser.add_argument('--noupdate', default = False, action='store_true', help = "Skip update of extensions and submodules, default: %(default)s")
+    parser.add_argument('--nodirectml', default = False, action='store_true', help = "Although nVidia and AMD toolkit aren't detected, use CPU not DirectML, default: %(default)s")
     parser.add_argument('--skip-requirements', default = False, action='store_true', help = "Skips checking and installing requirements, default: %(default)s")
     parser.add_argument('--skip-extensions', default = False, action='store_true', help = "Skips running individual extension installers, default: %(default)s")
     parser.add_argument('--skip-git', default = False, action='store_true', help = "Skips running all GIT operations, default: %(default)s")
     parser.add_argument('--experimental', default = False, action='store_true', help = "Allow unsupported versions of libraries, default: %(default)s")
+    parser.add_argument('--test', default = False, action='store_true', help = "Run test only, default: %(default)s")
 
 
 def parse_args():