diff --git a/.github/workflows/on_pull_request.yaml b/.github/workflows/on_pull_request.yaml index 011fe9b2d..6cb1005b0 100644 --- a/.github/workflows/on_pull_request.yaml +++ b/.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') diff --git a/.pylintrc b/.pylintrc index 72e1224e8..515dce01d 100644 --- a/.pylintrc +++ b/.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] diff --git a/TODO.md b/TODO.md index 167a823ec..0a5b0e99d 100644 --- a/TODO.md +++ b/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 - Stream-load models as option for slow storage +- Auto-test `torch.layer_norm` for FP16 ## Investigate diff --git a/javascript/imageviewer.js b/javascript/imageviewer.js index c928d9c5f..149177430 100644 --- a/javascript/imageviewer.js +++ b/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); diff --git a/launch.py b/launch.py index bf34ee507..8df3f787c 100644 --- a/launch.py +++ b/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() diff --git a/modules/cmd_args.py b/modules/cmd_args.py index 3bca9b357..9b8438800 100644 --- a/modules/cmd_args.py +++ b/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 diff --git a/modules/devices.py b/modules/devices.py index 7c869d60a..e317d91f4 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: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(): diff --git a/modules/dml/__init__.py b/modules/dml/__init__.py new file mode 100644 index 000000000..b77db18d0 --- /dev/null +++ b/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 diff --git a/modules/dml/hijack/__init__.py b/modules/dml/hijack/__init__.py new file mode 100644 index 000000000..d8cc0913a --- /dev/null +++ b/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 diff --git a/modules/dml/hijack/kdiffusion.py b/modules/dml/hijack/kdiffusion.py new file mode 100644 index 000000000..2eced885f --- /dev/null +++ b/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 \ No newline at end of file diff --git a/modules/dml/hijack/plms.py b/modules/dml/hijack/plms.py new file mode 100644 index 000000000..49979db28 --- /dev/null +++ b/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 diff --git a/modules/dml/hijack/realesrgan_model.py b/modules/dml/hijack/realesrgan_model.py new file mode 100644 index 000000000..c82429fda --- /dev/null +++ b/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 diff --git a/modules/dml/hijack/stablediffusion.py b/modules/dml/hijack/stablediffusion.py new file mode 100644 index 000000000..cbc4b85fb --- /dev/null +++ b/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 diff --git a/modules/dml/hijack/torch.py b/modules/dml/hijack/torch.py new file mode 100644 index 000000000..04e777258 --- /dev/null +++ b/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') diff --git a/modules/dml/optimizer/amd/__init__.py b/modules/dml/optimizer/amd/__init__.py new file mode 100644 index 000000000..23c6e57e3 --- /dev/null +++ b/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)) diff --git a/modules/dml/optimizer/amd/driver/atiadlxx.py b/modules/dml/optimizer/amd/driver/atiadlxx.py new file mode 100644 index 000000000..81a578a5c --- /dev/null +++ b/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 diff --git a/modules/dml/optimizer/amd/driver/atiadlxx_apis.py b/modules/dml/optimizer/amd/driver/atiadlxx_apis.py new file mode 100644 index 000000000..25e6390ef --- /dev/null +++ b/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)] diff --git a/modules/dml/optimizer/amd/driver/atiadlxx_defines.py b/modules/dml/optimizer/amd/driver/atiadlxx_defines.py new file mode 100644 index 000000000..c242b9819 --- /dev/null +++ b/modules/dml/optimizer/amd/driver/atiadlxx_defines.py @@ -0,0 +1 @@ +ADL_OK = 0 \ No newline at end of file diff --git a/modules/dml/optimizer/amd/driver/atiadlxx_structures.py b/modules/dml/optimizer/amd/driver/atiadlxx_structures.py new file mode 100644 index 000000000..a68392ec4 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/modules/dml/optimizer/intel/__init__.py b/modules/dml/optimizer/intel/__init__.py new file mode 100644 index 000000000..bffda69f1 --- /dev/null +++ b/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 diff --git a/modules/dml/optimizer/nvidia/__init__.py b/modules/dml/optimizer/nvidia/__init__.py new file mode 100644 index 000000000..e5fda97b8 --- /dev/null +++ b/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 diff --git a/modules/dml/optimizer/optimizer.py b/modules/dml/optimizer/optimizer.py new file mode 100644 index 000000000..db88c6579 --- /dev/null +++ b/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 diff --git a/modules/dml/optimizer/unknown/__init__.py b/modules/dml/optimizer/unknown/__init__.py new file mode 100644 index 000000000..40ff476cd --- /dev/null +++ b/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) diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py index 11b76ca3d..c9b6cddf0 100644 --- a/modules/realesrgan_model.py +++ b/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] diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index 776d91d6c..12ee9f956 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': + 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) diff --git a/modules/shared.py b/modules/shared.py index 0dfefa8c4..f3ebbce55 100644 --- a/modules/shared.py +++ b/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)"), diff --git a/setup.py b/setup.py index 804ee35b0..cfaa1aaa0 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 = {} @@ -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():