Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix image corruption of RESRGAN upscaler in DirectML.

Browse Source
This commit is contained in:
Seunghoon Lee
2023-06-29 11:10:48 +09:00
parent 20163c9a32
commit a9fa014e31
2 changed files with 10 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/dml/hijack/kdiffusion.py
+4 -4
View File
@@ -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
modules/dml/hijack/realesrgan_model.py
+6 -8
View File
@@ -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