From 78e5747a92b818c6b995ba48040d710ad9f3e25b Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Thu, 9 May 2024 01:46:39 +0900 Subject: [PATCH] DirectML fix some samplers (unipc, euler, etc) --- modules/dml/backend.py | 3 + modules/dml/hijack/diffusers.py | 203 ++++++++++++++------------ modules/dml/hijack/plms.py | 2 +- modules/dml/hijack/stablediffusion.py | 4 +- modules/dml/hijack/torch.py | 25 +++- modules/processing_diffusers.py | 2 + 6 files changed, 142 insertions(+), 97 deletions(-) diff --git a/modules/dml/backend.py b/modules/dml/backend.py index 7947dc81b..5f34b4542 100644 --- a/modules/dml/backend.py +++ b/modules/dml/backend.py @@ -78,3 +78,6 @@ class DirectML: def reset_peak_memory_stats(device: Optional[rDevice]=None): return + + def synchronize_tensor(tensor: torch.Tensor) -> None: + tensor.__str__() diff --git a/modules/dml/hijack/diffusers.py b/modules/dml/hijack/diffusers.py index 881ac44fd..4896ccd3e 100644 --- a/modules/dml/hijack/diffusers.py +++ b/modules/dml/hijack/diffusers.py @@ -1,7 +1,7 @@ +from typing import Optional, Union, Tuple import torch import diffusers import diffusers.utils.torch_utils -from typing import Optional, Union, Tuple def PNDMScheduler__get_prev_sample(self, sample: torch.FloatTensor, timestep, prev_timestep, model_output): @@ -17,7 +17,7 @@ def PNDMScheduler__get_prev_sample(self, sample: torch.FloatTensor, timestep, pr # sample -> x_t # model_output -> e_θ(x_t, t) # prev_sample -> x_(t−δ) - sample.__str__() # PNDM Sampling does not work without 'stringify'. (because it depends on PLMS) + torch.dml.synchronize_tensor(sample) # DML synchronize alpha_prod_t = self.alphas_cumprod[timestep] alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod beta_prod_t = 1 - alpha_prod_t @@ -53,30 +53,43 @@ diffusers.PNDMScheduler._get_prev_sample = PNDMScheduler__get_prev_sample # pyli def UniPCMultistepScheduler_multistep_uni_p_bh_update( - self: diffusers.UniPCMultistepScheduler, + self, model_output: torch.FloatTensor, - prev_timestep: int, - sample: torch.FloatTensor, - order: int, + *args, + sample: torch.FloatTensor = None, + order: int = None, + **kwargs, ) -> torch.FloatTensor: """ One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified. Args: model_output (`torch.FloatTensor`): - direct outputs from learned diffusion model at the current timestep. - prev_timestep (`int`): previous discrete timestep in the diffusion chain. + The direct output from the learned diffusion model at the current timestep. + prev_timestep (`int`): + The previous discrete timestep in the diffusion chain. sample (`torch.FloatTensor`): - current instance of sample being created by diffusion process. - order (`int`): the order of UniP at this step, also the p in UniPC-p. + A current instance of a sample created by the diffusion process. + order (`int`): + The order of UniP at this timestep (corresponds to the *p* in UniPC-p). Returns: - `torch.FloatTensor`: the sample tensor at the previous timestep. + `torch.FloatTensor`: + The sample tensor at the previous timestep. """ - timestep_list = self.timestep_list + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError(" missing `sample` as a required keyward argument") + if order is None: + if len(args) > 2: + order = args[2] + else: + raise ValueError(" missing `order` as a required keyward argument") model_output_list = self.model_outputs - s0, t = self.timestep_list[-1], prev_timestep + s0 = self.timestep_list[-1] m0 = model_output_list[-1] x = sample @@ -84,10 +97,13 @@ def UniPCMultistepScheduler_multistep_uni_p_bh_update( x_t = self.solver_p.step(model_output, s0, x).prev_sample return x_t - sample.__str__() # UniPC Sampling does not work without 'stringify'. - lambda_t, lambda_s0 = self.lambda_t[t], self.lambda_t[s0] - alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0] - sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0] + torch.dml.synchronize_tensor(sample) # DML synchronize + sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) h = lambda_t - lambda_s0 device = sample.device @@ -95,9 +111,10 @@ def UniPCMultistepScheduler_multistep_uni_p_bh_update( rks = [] D1s = [] for i in range(1, order): - si = timestep_list[-(i + 1)] + si = self.step_index - i mi = model_output_list[-(i + 1)] - lambda_si = self.lambda_t[si] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) rk = (lambda_si - lambda_s0) / h rks.append(rk) D1s.append((mi - m0) / rk) @@ -143,14 +160,14 @@ def UniPCMultistepScheduler_multistep_uni_p_bh_update( if self.predict_x0: x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 if D1s is not None: - pred_res = torch.einsum("k,bkchw->bchw", rhos_p, D1s) + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) else: pred_res = 0 x_t = x_t_ - alpha_t * B_h * pred_res else: x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 if D1s is not None: - pred_res = torch.einsum("k,bkchw->bchw", rhos_p, D1s) + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) else: pred_res = 0 x_t = x_t_ - sigma_t * B_h * pred_res @@ -170,91 +187,91 @@ def LCMScheduler_step( generator: Optional[torch.Generator] = None, return_dict: bool = True, ) -> Union[diffusers.schedulers.scheduling_lcm.LCMSchedulerOutput, Tuple]: - """ - Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion - process from the learned model outputs (most often the predicted noise). + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). - Args: - model_output (`torch.FloatTensor`): - The direct output from learned diffusion model. - timestep (`float`): - The current discrete timestep in the diffusion chain. - sample (`torch.FloatTensor`): - A current instance of a sample created by the diffusion process. - generator (`torch.Generator`, *optional*): - A random number generator. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`. - Returns: - [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a - tuple is returned where the first element is the sample tensor. - """ - if self.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + generator (`torch.Generator`, *optional*): + A random number generator. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`. + Returns: + [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a + tuple is returned where the first element is the sample tensor. + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) - if self.step_index is None: - self._init_step_index(timestep) + if self.step_index is None: + self._init_step_index(timestep) - # 1. get previous step value - prev_step_index = self.step_index + 1 - if prev_step_index < len(self.timesteps): - prev_timestep = self.timesteps[prev_step_index] - else: - prev_timestep = timestep + # 1. get previous step value + prev_step_index = self.step_index + 1 + if prev_step_index < len(self.timesteps): + prev_timestep = self.timesteps[prev_step_index] + else: + prev_timestep = timestep - # 2. compute alphas, betas - sample.__str__() - alpha_prod_t = self.alphas_cumprod[timestep] - alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod + # 2. compute alphas, betas + torch.dml.synchronize_tensor(sample) # DML synchronize + alpha_prod_t = self.alphas_cumprod[timestep] + alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev + beta_prod_t = 1 - alpha_prod_t + beta_prod_t_prev = 1 - alpha_prod_t_prev - # 3. Get scalings for boundary conditions - c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep) + # 3. Get scalings for boundary conditions + c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep) - # 4. Compute the predicted original sample x_0 based on the model parameterization - if self.config.prediction_type == "epsilon": # noise-prediction - predicted_original_sample = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt() - elif self.config.prediction_type == "sample": # x-prediction - predicted_original_sample = model_output - elif self.config.prediction_type == "v_prediction": # v-prediction - predicted_original_sample = alpha_prod_t.sqrt() * sample - beta_prod_t.sqrt() * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or" - " `v_prediction` for `LCMScheduler`." - ) + # 4. Compute the predicted original sample x_0 based on the model parameterization + if self.config.prediction_type == "epsilon": # noise-prediction + predicted_original_sample = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt() + elif self.config.prediction_type == "sample": # x-prediction + predicted_original_sample = model_output + elif self.config.prediction_type == "v_prediction": # v-prediction + predicted_original_sample = alpha_prod_t.sqrt() * sample - beta_prod_t.sqrt() * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or" + " `v_prediction` for `LCMScheduler`." + ) - # 5. Clip or threshold "predicted x_0" - if self.config.thresholding: - predicted_original_sample = self._threshold_sample(predicted_original_sample) - elif self.config.clip_sample: - predicted_original_sample = predicted_original_sample.clamp( - -self.config.clip_sample_range, self.config.clip_sample_range - ) + # 5. Clip or threshold "predicted x_0" + if self.config.thresholding: + predicted_original_sample = self._threshold_sample(predicted_original_sample) + elif self.config.clip_sample: + predicted_original_sample = predicted_original_sample.clamp( + -self.config.clip_sample_range, self.config.clip_sample_range + ) - # 6. Denoise model output using boundary conditions - denoised = c_out * predicted_original_sample + c_skip * sample + # 6. Denoise model output using boundary conditions + denoised = c_out * predicted_original_sample + c_skip * sample - # 7. Sample and inject noise z ~ N(0, I) for MultiStep Inference - # Noise is not used for one-step sampling. - if len(self.timesteps) > 1: - noise = diffusers.utils.torch_utils.randn_tensor(model_output.shape, generator=generator, device=model_output.device) - prev_sample = alpha_prod_t_prev.sqrt() * denoised + beta_prod_t_prev.sqrt() * noise - else: - prev_sample = denoised + # 7. Sample and inject noise z ~ N(0, I) for MultiStep Inference + # Noise is not used for one-step sampling. + if len(self.timesteps) > 1: + noise = diffusers.utils.torch_utils.randn_tensor(model_output.shape, generator=generator, device=model_output.device) + prev_sample = alpha_prod_t_prev.sqrt() * denoised + beta_prod_t_prev.sqrt() * noise + else: + prev_sample = denoised - # upon completion increase step index by one - self._step_index += 1 + # upon completion increase step index by one + self._step_index += 1 - if not return_dict: - return (prev_sample, denoised) + if not return_dict: + return (prev_sample, denoised) - return diffusers.schedulers.scheduling_lcm.LCMSchedulerOutput(prev_sample=prev_sample, denoised=denoised) + return diffusers.schedulers.scheduling_lcm.LCMSchedulerOutput(prev_sample=prev_sample, denoised=denoised) diffusers.LCMScheduler.step = LCMScheduler_step diff --git a/modules/dml/hijack/plms.py b/modules/dml/hijack/plms.py index a8afcbd05..d19c46629 100644 --- a/modules/dml/hijack/plms.py +++ b/modules/dml/hijack/plms.py @@ -47,7 +47,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 - alphas[index].__str__() # DML Solution: PLMS Sampling does not work without this 'stringify'. + torch.dml.synchronize_tensor(alphas[index]) # DML synchronize 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/dml/hijack/stablediffusion.py b/modules/dml/hijack/stablediffusion.py index 02c394e9e..818250181 100644 --- a/modules/dml/hijack/stablediffusion.py +++ b/modules/dml/hijack/stablediffusion.py @@ -52,7 +52,7 @@ def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=F 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 - alphas[index].__str__() # DML Solution: DDIM Sampling does not work without this 'stringify'. + torch.dml.synchronize_tensor(alphas[index]) # DML synchronize 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) @@ -68,7 +68,7 @@ def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=F pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0) if dynamic_threshold is not None: - raise NotImplementedError + raise NotImplementedError() # direction pointing to x_t dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t diff --git a/modules/dml/hijack/torch.py b/modules/dml/hijack/torch.py index 084671425..e4913a357 100644 --- a/modules/dml/hijack/torch.py +++ b/modules/dml/hijack/torch.py @@ -3,10 +3,10 @@ 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') - # https://github.com/microsoft/DirectML/issues/400 CondFunc('torch.Tensor.new', lambda orig, self, *args, **kwargs: orig(self.cpu(), *args, **kwargs).to(self.device), lambda orig, self, *args, **kwargs: torch.dml.is_directml_device(self.device)) + _lerp = torch.lerp def lerp(*args, **kwargs) -> torch.Tensor: rep = None @@ -40,3 +40,26 @@ def lerp(*args, **kwargs) -> torch.Tensor: return _lerp(*args, **kwargs).to(rep.device).type(rep.dtype) return _lerp(*args, **kwargs) torch.lerp = lerp + + +# https://github.com/lshqqytiger/stable-diffusion-webui-directml/issues/436 +_pow_ = torch.Tensor.pow_ +def pow_(self: torch.Tensor, *args, **kwargs): + if self.dtype == torch.float64: + return _pow_(self.cpu(), *args, **kwargs).to(self.device) + return _pow_(self, *args, **kwargs) +torch.Tensor.pow_ = pow_ + + +_load = torch.load +def load(f, map_location = "cpu", *args, **kwargs): + if type(map_location) in (str, torch.device,): + device = torch.device(map_location) + if device.type == "privateuseone": + data = _load(f, *args, map_location="cpu", **kwargs) + for k in data: + for weight in data[k]: + data[k][weight] = data[k][weight].to(device) + return data + return _load(f, *args, map_location=map_location, **kwargs) +torch.load = load diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index b5f8ffce8..3324abfa1 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -62,6 +62,8 @@ def process_diffusers(p: processing.StableDiffusionProcessing): def diffusers_callback(pipe, step: int, timestep: int, kwargs: dict): latents = kwargs.get('latents', None) + if torch.is_tensor(latents) and latents.device.type == "privateuseone": + torch.dml.synchronize_tensor(latents) # DML synchronize debug_callback(f'Callback: step={step} timestep={timestep} latents={latents.shape if latents is not None else None} kwargs={list(kwargs)}') shared.state.sampling_step = step if shared.state.interrupted or shared.state.skipped: