220 lines
9.9 KiB
Python
220 lines
9.9 KiB
Python
import torch
|
|
from PIL import Image
|
|
import numpy as np
|
|
from torch.nn.functional import interpolate
|
|
import math
|
|
|
|
from comfy.k_diffusion import sampling
|
|
from comfy.samplers import KSAMPLER
|
|
from tqdm.auto import trange
|
|
|
|
def dump_image(tensor):
|
|
if (tensor == tensor.flatten()[0]).all():
|
|
print("trivial dump")
|
|
return
|
|
tensor = tensor.to('cpu', copy=True)
|
|
tensor -= tensor.min()
|
|
tensor /= tensor.max()
|
|
if len(tensor.shape) == 3:
|
|
tensor = tensor.unsqueeze(1).expand(-1,3,-1,-1)
|
|
tensor = tensor.permute(0,2,3,1)[0]
|
|
h,w,c = tensor.shape
|
|
b = np.clip(tensor.numpy() * 255, 0, 255).astype(np.uint8)
|
|
Image.fromarray(b).save("test.png")
|
|
|
|
@torch.no_grad()
|
|
def sample_dynamic(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., sub_weight=3, early_terminate=False, resolution_mult=1.25):
|
|
"""Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
|
|
extra_args = {} if extra_args is None else extra_args
|
|
s_in = x.new_ones([x.shape[0]])
|
|
n,c,h,w = x.shape
|
|
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)
|
|
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)
|
|
prev_denoised = None
|
|
oversample_performed = False
|
|
for i in trange(len(sigmas) - 1, disable=disable):
|
|
gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
|
|
sigma_hat = sigmas[i] * (gamma + 1)
|
|
if gamma > 0:
|
|
eps = torch.randn_like(x) * s_noise
|
|
x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
|
|
denoised = model(x, sigma_hat * s_in, **extra_args)
|
|
maskl = (i/len(sigmas)*base_mask.max()-base_mask).clip(max=1,min=0).unsqueeze(0).unsqueeze(0)
|
|
maskh = ((i+1)/len(sigmas)*base_mask.max()-base_mask).clip(max=1,min=0).unsqueeze(0).unsqueeze(0)
|
|
#A filter of where changes are expected
|
|
mask = torch.fft.ifftshift(maskh*(1-maskl))
|
|
#Just the expected changes
|
|
filt = torch.fft.ifft2(mask*torch.fft.fft2(denoised)).real
|
|
if prev_denoised is not None:
|
|
#changes made outside the current sigma schedule
|
|
ext = prev_denoised+filt-denoised
|
|
#The distance of changes relative to the current step
|
|
dist = (ext*ext).sum(1)/math.sqrt(-dt)
|
|
hist+= dist
|
|
#The sum of distances that would be modified by a subrender for each center point
|
|
summed = torch.nn.functional.conv2d(dist, kernel, groups=1)
|
|
#The most valuable of possible center points
|
|
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])
|
|
#TODO: allow for simultaneous selection of multiple zones
|
|
#An arbitrary cutoff
|
|
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:
|
|
oversample_performed = True
|
|
subx = torch.stack(sublocs)
|
|
sh,sw = h//4,w//4
|
|
subx = interpolate(subx, size=(int(sh*resolution_mult),
|
|
int(sw*resolution_mult)), mode='bicubic')
|
|
#The upscaled/interpolated pixels lack higher frequency noise,
|
|
#so additional noise must be added
|
|
subx += torch.randn_like(subx)*(1-math.sqrt(2))*sigma_hat*s_in
|
|
subdn = model(subx, sigma_hat*s_in*math.sqrt(2), **extra_args)
|
|
#downscale
|
|
subdn = interpolate(subdn, size=(sh,sw), mode='bicubic')
|
|
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]
|
|
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, focalm)
|
|
prev_denoised = denoised
|
|
d = sampling.to_d(x, sigma_hat, denoised)
|
|
|
|
if callback is not None:
|
|
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
|
|
dt = sigmas[i + 1] - sigma_hat
|
|
# Euler method
|
|
x = x + d * dt
|
|
if early_terminate and oversample_performed:
|
|
break
|
|
#Debugging code to dump the measured distances
|
|
#summed = torch.nn.functional.conv2d(hist, kernel, groups=1)
|
|
#dump_image(hist)
|
|
return x
|
|
|
|
class DynamicSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"sub_weight": ("FLOAT", {"default": 3, "min": -1, "step": 0.01}),
|
|
"resolution_mult": ("FLOAT", {"default": 1.25, "min": 1, "step": 0.01}),
|
|
"early_terminate": ("BOOLEAN", {"default": False}),}}
|
|
RETURN_TYPES = ("SAMPLER",)
|
|
CATEGORY = "sampling/custom_sampling/samplers"
|
|
|
|
FUNCTION = "get_sampler"
|
|
def get_sampler(self, sub_weight, early_terminate, resolution_mult):
|
|
return (KSAMPLER(lambda *args, **kwargs: sample_dynamic(*args, sub_weight=sub_weight, early_terminate=early_terminate, resolution_mult=resolution_mult, **kwargs)),)
|
|
|
|
class MeasuredSampler(KSAMPLER):
|
|
def __init__(self, sampler):
|
|
self.sampler = sampler.sampler_function
|
|
self.prev_denoised = None
|
|
self.prev_sigma = 0
|
|
super().__init__(self.wrapped_sample)
|
|
def wrapped_sample(self, *args, **kwargs):
|
|
original_callback = kwargs.get("callback", None)
|
|
def callback(args):
|
|
self.callback(args)
|
|
if original_callback is not None:
|
|
original_callback(args)
|
|
self.sigmas = args[2]
|
|
self.steps = len(args[2])
|
|
kwargs["callback"] = callback
|
|
x = args[1]
|
|
n,c,h,w = x.shape
|
|
self.hist = torch.zeros((n,h,w), device=x.device, dtype=torch.float32)
|
|
kx, ky = torch.meshgrid(torch.linspace(-h/2,h/2,h),torch.linspace(-w/2,w/2,w), indexing="ij")
|
|
self.base_mask = (kx*kx + ky*ky).sqrt().to(x.device)
|
|
self.prev_denoised = None
|
|
x = self.sampler(*args, **kwargs)
|
|
return x
|
|
|
|
def callback(self, args):
|
|
denoised = args["denoised"]
|
|
i = args["i"]
|
|
dt = self.sigmas[i+1] - args['sigma_hat']
|
|
if self.prev_denoised is not None:
|
|
maskl = (i/self.steps*self.base_mask.max()-self.base_mask).clip(max=1,min=0).unsqueeze(0).unsqueeze(0)
|
|
maskh = ((i+1)/self.steps*self.base_mask.max()-self.base_mask).clip(max=1,min=0).unsqueeze(0).unsqueeze(0)
|
|
#A filter of where changes are expected
|
|
mask = torch.fft.ifftshift(maskh*(1-maskl))
|
|
#Just the expected changes
|
|
filt = torch.fft.ifft2(mask*torch.fft.fft2(denoised)).real
|
|
#changes made outside the current sigma schedule
|
|
ext = self.prev_denoised+filt-denoised
|
|
#The distance of changes relative to the current step
|
|
dist = (ext*ext).sum(1)/math.sqrt(-dt)
|
|
self.hist+= dist
|
|
self.prev_denoised = denoised
|
|
|
|
#Blur functions borrowed from comfy_extras/nodes_post_processing
|
|
#with slight modifications
|
|
def gaussian_blur(latents, sigma=1):
|
|
radius = min(latents.shape[2:])//2
|
|
size = radius*2+1
|
|
maxl = size // 2
|
|
x, y = torch.meshgrid(torch.linspace(-maxl, maxl, size), torch.linspace(-maxl, maxl, size), indexing="ij")
|
|
d = (x * x + y * y) / (2 * sigma * sigma)
|
|
mask = torch.exp(-d) / (2 * torch.pi * sigma * sigma)
|
|
kernel = (mask / mask.sum())[None, None].to(latents.device)
|
|
padded_latents = torch.nn.functional.pad(latents, [radius]*4, 'reflect')
|
|
blurred = torch.nn.functional.conv2d(padded_latents, kernel, padding=(radius*2+1) // 2, groups=1)
|
|
return blurred[:, :, radius:-radius, radius:-radius]
|
|
|
|
class MeasuredSamplerNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"sampler": ("SAMPLER",),}}
|
|
RETURN_TYPES = ("SAMPLER", "MASK_PROMISE")
|
|
CATEGORY = "sampling/custom_sampling/samplers"
|
|
|
|
FUNCTION = "get_sampler"
|
|
def get_sampler(self, sampler):
|
|
s = MeasuredSampler(sampler)
|
|
return (s, s)
|
|
|
|
class ResolveMaskPromise:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"latent": ("LATENT",), "mask_promise": ("MASK_PROMISE",),
|
|
"upper_threshold": ("FLOAT", {"default": .8, "step": .01, "min": 0, "max": 1}),
|
|
"lower_threshold": ("FLOAT", {"default": .2, "step": .01, "min": 0, "max": 1}),
|
|
"blur_sigma": ("FLOAT", {"default": 0, "min": 0, "step": .01}),}}
|
|
RETURN_TYPES = ("MASK",)
|
|
|
|
FUNCTION = "get_mask"
|
|
def get_mask(self, latent, mask_promise, lower_threshold, upper_threshold, blur_sigma):
|
|
#NOTE: latent is only used to ensure this executes after sampling
|
|
hist = mask_promise.hist
|
|
if blur_sigma > 0:
|
|
hist = gaussian_blur(hist.unsqueeze(1), blur_sigma).squeeze(1)
|
|
sorted_hist = hist.flatten(start_dim=1).sort().values
|
|
lower = sorted_hist[:,int((sorted_hist.size(1)-1)*lower_threshold)][:,None,None]
|
|
upper = sorted_hist[:,int((sorted_hist.size(1)-1)*upper_threshold)][:,None,None]
|
|
mask = ((hist-lower)/(upper-lower)).clip(max=1, min=0)
|
|
return (mask.cpu(),)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DynamicSampler": DynamicSampler,
|
|
"MeasuredSampler": MeasuredSamplerNode,
|
|
"ResolveMaskPromise": ResolveMaskPromise,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|