diff --git a/TODO_DML.md b/TODO_DML.md new file mode 100644 index 000000000..af6559ac1 --- /dev/null +++ b/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. diff --git a/modules/devices.py b/modules/devices.py index 7c869d60a..5f5e18768 100644 --- a/modules/devices.py +++ b/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(): diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py index 86007f906..ac1102e38 100644 --- a/modules/realesrgan_model.py +++ b/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] diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index f817b7afd..cd859a6c3 100644 --- a/modules/sd_hijack.py +++ b/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) diff --git a/modules/sd_hijack_inpainting.py b/modules/sd_hijack_inpainting.py index 4b23c132d..0c37fef89 100644 --- a/modules/sd_hijack_inpainting.py +++ b/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) diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index 776d91d6c..0f1f2a757 100644 --- a/modules/sd_hijack_optimizations.py +++ b/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) diff --git a/modules/shared.py b/modules/shared.py index a2e202735..f147317cd 100644 --- a/modules/shared.py +++ b/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") })) diff --git a/setup.py b/setup.py index 73269b0df..aa01a3de5 100644 --- a/setup.py +++ b/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")