First DirectML implementation.

Unstable and not tested.

TODO_DML.md

@@ -0,0 +1,22 @@
+# TODO
+
+## Issues
+
+Stuff to be fixed...
+
+- `mat1 and mat2 must have the same dtype` error (half mode)
+- Some samplers won't work (test later)
+
+## Something needs discussion
+
+- About memory optimization.
+
+Basically, we cannot get detailed vram information from `torch-directml`.
+
+It has `gpu_memory` method which returns an array contains used memory size, but it is almostly useless without any other information.
+
+What should we do?
+
+1. Use any fixed value as the available memory capacity.
+2. Use `atiadlxx`(AMD/ATI GPU driver library) to infer vram information as similar as possible to the actual value. (works for AMDGPUs)
+3. or another better way.

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"
+
+
 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/realesrgan_model.py

@@ -6,10 +6,76 @@ 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
 
 
+# DML ISSUE: Some tensors turn 0 after Extended Slices.
+def realesrgan_tile_process_dml_fix(self):
+    import math
+    import torch
+    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' if self.device.type == 'privateuseone' else self.device)
+    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)
+
+
 class UpscalerRealESRGAN(Upscaler):
     def __init__(self, path):
         self.name = "RealESRGAN"
@@ -37,6 +103,8 @@ class UpscalerRealESRGAN(Upscaler):
 
         try:
             from realesrgan import RealESRGANer
+            if device.type == 'privateuseone':
+                RealESRGANer.tile_process = realesrgan_tile_process_dml_fix
         except:
             print("Error importing Real-ESRGAN:", file=sys.stderr)
             return img
@@ -53,6 +121,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.py

@@ -257,6 +257,9 @@ class EmbeddingsWithFixes(torch.nn.Module):
             for offset, embedding in fixes:
                 emb = devices.cond_cast_unet(embedding.vec)
                 emb_len = min(tensor.shape[0] - offset - 1, emb.shape[0])
+                # DML ISSUE: type mismatch on half mode
+                if tensor.dtype == torch.float16 and emb.dtype == torch.float32 and not shared.cmd_opts.no_half:
+                    emb = emb.half()
                 tensor = torch.cat([tensor[0:offset + 1], emb[0:emb_len], tensor[offset + 1 + emb_len:]])
 
             vecs.append(tensor)

modules/sd_hijack_inpainting.py

@@ -53,6 +53,7 @@ def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=F
 
     def get_x_prev_and_pred_x0(e_t, index):
         # select parameters corresponding to the currently considered timestep
+        print(alphas[index]) # DML ISSUE: 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)

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':
+        # DML ISSUE: There's no way to get any memory info.
+        return 1048576
     else:
         return psutil.virtual_memory().available
 
@@ -195,6 +198,10 @@ 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):
+    # DML ISSUE: There's no way to get any memory info.
+    return einsum_op_tensor_mem(q, k, v, 1024)
+
 def einsum_op(q, k, v):
     if q.device.type == 'cuda':
         return einsum_op_cuda(q, k, v)
@@ -204,6 +211,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

@@ -422,7 +422,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"]}),
+    "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.
     "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")
 }))

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 = {}
@@ -188,17 +189,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==1.13.1 torchvision==0.14.1 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:
@@ -207,6 +212,16 @@ 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)}')
+            except:
+                log.warning("Torch repoorts CUDA not available")
     except Exception as e:
         log.error(f'Could not load torch: {e}')
         exit(1)
@@ -239,14 +254,14 @@ def install_repositories():
         return os.path.join(os.path.dirname(__file__), 'repositories', name)
     log.info('Installing repositories')
     os.makedirs(os.path.join(os.path.dirname(__file__), 'repositories'), exist_ok=True)
-    stable_diffusion_repo = os.environ.get('STABLE_DIFFUSION_REPO', "https://github.com/Stability-AI/stablediffusion.git")
-    stable_diffusion_commit = os.environ.get('STABLE_DIFFUSION_COMMIT_HASH', "cf1d67a6fd5ea1aa600c4df58e5b47da45f6bdbf")
+    stable_diffusion_repo = os.environ.get('STABLE_DIFFUSION_REPO', "https://github.com/Stability-AI/stablediffusion.git") # DML TODO: check samplers work well
+    stable_diffusion_commit = os.environ.get('STABLE_DIFFUSION_COMMIT_HASH', "d4c168b2ad29d82e5fdfea4d598075f40a3b0341")
     clone(stable_diffusion_repo, d('stable-diffusion-stability-ai'), stable_diffusion_commit)
     taming_transformers_repo = os.environ.get('TAMING_TRANSFORMERS_REPO', "https://github.com/CompVis/taming-transformers.git")
     taming_transformers_commit = os.environ.get('TAMING_TRANSFORMERS_COMMIT_HASH', "3ba01b241669f5ade541ce990f7650a3b8f65318")
     clone(taming_transformers_repo, d('taming-transformers'), taming_transformers_commit)
-    k_diffusion_repo = os.environ.get('K_DIFFUSION_REPO', 'https://github.com/crowsonkb/k-diffusion.git')
-    k_diffusion_commit = os.environ.get('K_DIFFUSION_COMMIT_HASH', "b43db16749d51055f813255eea2fdf1def801919")
+    k_diffusion_repo = os.environ.get('K_DIFFUSION_REPO', 'https://github.com/crowsonkb/k-diffusion.git') # DML TODO: check samplers work well
+    k_diffusion_commit = os.environ.get('K_DIFFUSION_COMMIT_HASH', "47b6ef08bca986ff5e72815e74a419ef6616bdbb")
     clone(k_diffusion_repo, d('k-diffusion'), k_diffusion_commit)
     codeformer_repo = os.environ.get('CODEFORMER_REPO', 'https://github.com/sczhou/CodeFormer.git')
     codeformer_commit = os.environ.get('CODEFORMER_COMMIT_HASH', "c5b4593074ba6214284d6acd5f1719b6c5d739af")
@@ -481,6 +496,7 @@ def parse_args():
     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")