|
|
|
|
import torch
|
|
|
|
|
|
|
|
|
|
def common_upscale(samples, width, height, upscale_method, crop):
|
|
|
|
|
if crop == "center":
|
|
|
|
|
old_width = samples.shape[3]
|
|
|
|
|
old_height = samples.shape[2]
|
|
|
|
|
old_aspect = old_width / old_height
|
|
|
|
|
new_aspect = width / height
|
|
|
|
|
x = 0
|
|
|
|
|
y = 0
|
|
|
|
|
if old_aspect > new_aspect:
|
|
|
|
|
x = round((old_width - old_width * (new_aspect / old_aspect)) / 2)
|
|
|
|
|
elif old_aspect < new_aspect:
|
|
|
|
|
y = round((old_height - old_height * (old_aspect / new_aspect)) / 2)
|
|
|
|
|
s = samples[:,:,y:old_height-y,x:old_width-x]
|
|
|
|
|
else:
|
|
|
|
|
s = samples
|
|
|
|
|
return torch.nn.functional.interpolate(s, size=(height, width), mode=upscale_method)
|
|
|
|
|
|
|
|
|
|
@torch.inference_mode()
|
|
|
|
|
def tiled_scale(samples, function, tile_x=64, tile_y=64, overlap = 8, upscale_amount = 4, out_channels = 3):
|
|
|
|
|
output = torch.empty((samples.shape[0], out_channels, round(samples.shape[2] * upscale_amount), round(samples.shape[3] * upscale_amount)), device="cpu")
|
|
|
|
|
for b in range(samples.shape[0]):
|
|
|
|
|
s = samples[b:b+1]
|
|
|
|
|
out = torch.zeros((s.shape[0], out_channels, round(s.shape[2] * upscale_amount), round(s.shape[3] * upscale_amount)), device="cpu")
|
|
|
|
|
out_div = torch.zeros((s.shape[0], out_channels, round(s.shape[2] * upscale_amount), round(s.shape[3] * upscale_amount)), device="cpu")
|
|
|
|
|
for y in range(0, s.shape[2], tile_y - overlap):
|
|
|
|
|
for x in range(0, s.shape[3], tile_x - overlap):
|
|
|
|
|
s_in = s[:,:,y:y+tile_y,x:x+tile_x]
|
|
|
|
|
|
|
|
|
|
ps = function(s_in).cpu()
|
|
|
|
|
mask = torch.ones_like(ps)
|
|
|
|
|
feather = round(overlap * upscale_amount)
|
|
|
|
|
for t in range(feather):
|
|
|
|
|
mask[:,:,t:1+t,:] *= ((1.0/feather) * (t + 1))
|
|
|
|
|
mask[:,:,mask.shape[2] -1 -t: mask.shape[2]-t,:] *= ((1.0/feather) * (t + 1))
|
|
|
|
|
mask[:,:,:,t:1+t] *= ((1.0/feather) * (t + 1))
|
|
|
|
|
mask[:,:,:,mask.shape[3]- 1 - t: mask.shape[3]- t] *= ((1.0/feather) * (t + 1))
|
|
|
|
|
out[:,:,round(y*upscale_amount):round((y+tile_y)*upscale_amount),round(x*upscale_amount):round((x+tile_x)*upscale_amount)] += ps * mask
|
|
|
|
|
out_div[:,:,round(y*upscale_amount):round((y+tile_y)*upscale_amount),round(x*upscale_amount):round((x+tile_x)*upscale_amount)] += mask
|
|
|
|
|
|
|
|
|
|
output[b:b+1] = out/out_div
|
|
|
|
|
return output
|
