From de6dea5425f3a8834d1467981a4e5eb560691f95 Mon Sep 17 00:00:00 2001 From: Austin Mroz Date: Thu, 30 May 2024 15:12:33 -0500 Subject: [PATCH] Oversample in second pass In an attempt to improve visual quality, the oversampling is instead performed as a second pass at the end. The hope was that this would combine the promising selection code with an existent and known working solution for improving the visual fidelity of the subregion, but performance is still worse than expected. --- .gitignore | 1 + nodes.py | 78 ++++++++++++++++++++++++++++-------------------------- 2 files changed, 41 insertions(+), 38 deletions(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c18dd8d --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__/ diff --git a/nodes.py b/nodes.py index 5c84abe..34e0b6a 100644 --- a/nodes.py +++ b/nodes.py @@ -8,6 +8,9 @@ from comfy.k_diffusion import sampling from comfy.samplers import KSAMPLER from tqdm.auto import trange +from comfy import latent_formats +latent_factors = torch.tensor(latent_formats.SD15().latent_rgb_factors) + class LatentFFTAsImage: """Takes a latent as input , """ @classmethod @@ -33,6 +36,7 @@ def dump_image(tensor): if len(tensor.shape) == 3: tensor = tensor.unsqueeze(1).expand(-1,3,-1,-1) tensor = tensor.permute(0,2,3,1)[0] + tensor @= latent_factors h,w,c = tensor.shape b = np.clip(tensor.numpy() * 255, 0, 255).astype(np.uint8) Image.fromarray(b).save("test.png") @@ -46,6 +50,8 @@ def sample_dynamic(model, x, sigmas, extra_args=None, callback=None, disable=Non hist = torch.zeros((n,h,w), device=x.device, dtype=x.dtype) kx, ky = torch.meshgrid(torch.linspace(-h/2,h/2,h),torch.linspace(-w/2,w/2,w), indexing="ij") base_mask = (kx*kx + ky*ky).sqrt().to(x.device) + kx, ky = torch.meshgrid(torch.linspace(-h/4,h/4,h//2),torch.linspace(-w/4,w/4,w//2), indexing="ij") + base_submask = (kx*kx + ky*ky).sqrt().to(x.device) cw = w-w//4+1 ch = h-h//4+1 kernel = torch.ones((1,1,h//4,w//4), dtype=x.dtype, device=x.device) @@ -65,44 +71,8 @@ def sample_dynamic(model, x, sigmas, extra_args=None, callback=None, disable=Non filt = torch.fft.ifft2(mask*torch.fft.fft2(denoised)).real if prev_denoised is not None: ext = prev_denoised+filt-denoised - dist = (ext*ext).sum(1)/-dt + dist = (ext*ext).sum(1)/math.sqrt(-dt) hist+= dist - summed = torch.nn.functional.conv2d(dist, kernel, groups=1) - focal = torch.unravel_index(summed.view((n,ch*cw)).argmax(1), (ch,cw)) - focal = torch.stack(focal).transpose(0,1) - sublocs = [] - subinds = [] - focalm = [] - for j,ind in enumerate(focal): - focalm.append(summed[j, *ind]) - if summed[j, *ind] > 1000: - subx = x[j,:,ind[0]:ind[0]+h//4,ind[1]:ind[1]+w//4] - #scale+noise subx - sublocs.append(subx) - subinds.append(j) - if len(sublocs) > 0: - early_terminate = True - subx = torch.stack(sublocs) - #subx = subx.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3) - subx = interpolate(subx, size=(h//2,w//2), mode='bicubic') - subx += torch.randn_like(subx)*(1-math.sqrt(2))*sigma_hat*s_in - print("doing subrender") - subdn = model(subx, sigma_hat*s_in*math.sqrt(2), **extra_args) - #downscale - subdn = interpolate(subdn, size=(h//4,w//4), mode='bicubic') - #TODO: Does added noise need to be filtered out? - for j,dn in zip(subinds,subdn): - subx = denoised[j,:,focal[j][0]:focal[j][0]+h//4,focal[j][1]:focal[j][1]+w//4] - sub_weight = -1 - if sub_weight == -1: - subx[:] = dn - else: - subx[:] += sub_weight*dn - subx /= sub_weight+1 - focal[:,0] += h//8 - focal[:,1] += w//8 - print(focal*8) - print(focalm) prev_denoised = denoised d = sampling.to_d(x, sigma_hat, denoised) @@ -114,7 +84,39 @@ def sample_dynamic(model, x, sigmas, extra_args=None, callback=None, disable=Non if early_terminate: break summed = torch.nn.functional.conv2d(hist, kernel, groups=1) - dump_image(hist) + focal = torch.unravel_index(summed.view((n,ch*cw)).argmax(1), (ch,cw)) + focal = torch.stack(focal).transpose(0,1) + sublocs = [] + subinds = [] + focalm = [] + for j,ind in enumerate(focal): + focalm.append(summed[j, *ind]) + if summed[j, *ind] > 1000: + subx = x[j,:,ind[0]:ind[0]+h//4,ind[1]:ind[1]+w//4] + sublocs.append(subx) + subinds.append(j) + if len(sublocs) > 0: + print("doing subrender") + subx = torch.stack(sublocs) + subx = interpolate(subx, size=(h//2,w//2), mode='bicubic') + subnoise = model.noise[:,:,:h//2,:w//2] + subnoise = model.inner_model.inner_model.model_sampling.noise_scaling(sigmas[len(sigmas)//2], subnoise, subx, 0) + subres = sampling.sample_euler(model, subx, sigmas[len(sigmas)//2:], extra_args, callback, disable, s_churn, s_tmin, s_tmax, s_noise) + #downscale + subres = interpolate(subres, size=(h//4,w//4)) + #TODO: Does added noise need to be filtered out? + for j,dn in zip(subinds,subres): + subx = x[j,:,focal[j][0]:focal[j][0]+h//4,focal[j][1]:focal[j][1]+w//4] + sub_weight = -1 + if sub_weight == -1: + subx[:] = dn + else: + subx[:] += sub_weight*dn + subx /= sub_weight+1 + focal[:,0] += h//8 + focal[:,1] += w//8 + print(focal*8) + print(focalm) return x class DynamicSampler: