From a9fa014e3147b569f9e3832226b83085534a078f Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Thu, 29 Jun 2023 11:10:48 +0900 Subject: [PATCH] Fix image corruption of RESRGAN upscaler in DirectML. --- modules/dml/hijack/kdiffusion.py | 8 ++++---- modules/dml/hijack/realesrgan_model.py | 14 ++++++-------- 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/modules/dml/hijack/kdiffusion.py b/modules/dml/hijack/kdiffusion.py index 78bc9f2b5..19e16b013 100644 --- a/modules/dml/hijack/kdiffusion.py +++ b/modules/dml/hijack/kdiffusion.py @@ -12,8 +12,8 @@ def dpm_solver_adaptive(self, x, t_start, t_end, order=3, rtol=0.05, atol=0.0078 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) + atol = torch.tensor(atol, device=device) + rtol = torch.tensor(rtol, device=device) s = t_start x_prev = x accept = True @@ -67,7 +67,7 @@ def sample_dpm_fast(model, x, sigma_min, sigma_max, n, extra_args=None, callback 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) + return dpm_solver.dpm_solver_fast(x, dpm_solver.t(torch.tensor(sigma_max, device=device)), dpm_solver.t(torch.tensor(sigma_min, device=device)), n, eta, s_noise, noise_sampler) @torch.no_grad() @@ -79,7 +79,7 @@ def sample_dpm_adaptive(model, x, sigma_min, sigma_max, extra_args=None, callbac 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) + x, info = dpm_solver.dpm_solver_adaptive(x, dpm_solver.t(torch.tensor(sigma_max, device=device)), dpm_solver.t(torch.tensor(sigma_min, device=device)), order, rtol, atol, h_init, pcoeff, icoeff, dcoeff, accept_safety, eta, s_noise, noise_sampler) if return_info: return x, info return x diff --git a/modules/dml/hijack/realesrgan_model.py b/modules/dml/hijack/realesrgan_model.py index c82429fda..341e9aded 100644 --- a/modules/dml/hijack/realesrgan_model.py +++ b/modules/dml/hijack/realesrgan_model.py @@ -4,7 +4,7 @@ import torch from realesrgan import RealESRGANer -# DML Solution: Some tensors turn 0 after Extended Slices. Move output to cpu and get it back. +# DML Solution: Some of contents of output tensor turn to 0 after Extended Slices. Move it to cpu. def tile_process(self): batch, channel, height, width = self.img.shape output_height = height * self.scale @@ -12,7 +12,7 @@ def tile_process(self): output_shape = (batch, channel, output_height, output_width) # start with black image - self.output = self.img.new_zeros(output_shape, device='cpu') + self.output = self.img.new_zeros(output_shape) tiles_x = math.ceil(width / self.tile_size) tiles_y = math.ceil(height / self.tile_size) @@ -38,13 +38,12 @@ def tile_process(self): 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] + input_tile = self.img[0:self.img.shape[0], 0:self.img.shape[1], 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}') @@ -61,9 +60,8 @@ def tile_process(self): 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 + self.output = self.output.cpu() # 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) + self.output[0:self.output.shape[0], 0:self.output.shape[1], output_start_y:output_end_y, output_start_x:output_end_x] = output_tile.cpu()[0:output_tile.shape[0], 0:output_tile.shape[1], output_start_y_tile:output_end_y_tile, output_start_x_tile:output_end_x_tile] + self.output = self.output.to(output_tile.device) RealESRGANer.tile_process = tile_process