219 lines
8.3 KiB
Python
219 lines
8.3 KiB
Python
import torch
|
|
|
|
import torch.nn.functional as F
|
|
from comfy.sampler_helpers import prepare_mask
|
|
from comfy.k_diffusion import sampling
|
|
|
|
class LogSigmas:
|
|
"""For testing, simply prints the input sigmas"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "sigmas": ("SIGMAS",),
|
|
}}
|
|
|
|
FUNCTION = "log_sigmas"
|
|
OUTPUT_NODE = True
|
|
RETURN_TYPES = ()
|
|
CATEGORY = "_for_testing"
|
|
|
|
def log_sigmas(self, sigmas):
|
|
print(sigmas)
|
|
return ()
|
|
|
|
class RerangeSigmas:
|
|
"""Given a set of input sigmas, produce a new set of sigmas that cover the same range"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"sigmas": ("SIGMAS",),
|
|
"steps": ("INT", {"default": 10, "min": 1})}}
|
|
FUNCTION = "rerange_sigmas"
|
|
RETURN_TYPES = ("SIGMAS",)
|
|
CATEGORY = "sampling/custom_sampling/sigmas"
|
|
def rerange_sigmas(self, sigmas, steps):
|
|
assert(len(sigmas)>1)
|
|
(s_max, s_min) = (sigmas[0], sigmas[-1])
|
|
full_denoise = False
|
|
if s_min == 0:
|
|
assert(len(sigmas)>2)
|
|
full_denoise = True
|
|
s_min = sigmas[-2]
|
|
else:
|
|
steps+=1
|
|
|
|
#TODO: Implement scheduling method with a more uniform distribution
|
|
sigmas = sampling.get_sigmas_exponential(n=steps, sigma_min=s_min, sigma_max=s_max)
|
|
if not full_denoise:
|
|
sigmas = sigmas[:-1]
|
|
return (sigmas,)
|
|
|
|
|
|
#Blur functions shameless borrowed from comfy_extras/nodes_post_processing
|
|
#with slight modifications for latent dimensions
|
|
def gaussian_kernel(sigma, size=5):
|
|
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)
|
|
return mask / mask.sum()
|
|
|
|
|
|
|
|
|
|
def gaussian_blur(latents, kernel, radius=5):
|
|
padded_latents = F.pad(latents, [radius]*4, 'reflect')
|
|
blurred = F.conv2d(padded_latents, kernel, padding=(radius*2+1) // 2, groups=4)
|
|
return blurred[:, :, radius:-radius, radius:-radius]
|
|
|
|
class SpliceLatents:
|
|
"""Performs a fast approximate splice of 2 latents by bluring."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
#These numbers are likely flawed
|
|
return {"required": {"sigmas": ("SIGMAS",),
|
|
"radius": ("INT", {"default": 4, "min": 1, "step": 1}),
|
|
"wetness": ("FLOAT", {"default": 1.0, "max": 1,
|
|
"min": 0, "precision": 3,
|
|
"step": 0.1, "round": .01}),
|
|
"texture_override": (["None", "Upper", "Lower"],)},
|
|
"optional": {"lower": ("LATENT",),
|
|
"upper": ("LATENT",)}}
|
|
|
|
FUNCTION = "splice_latents"
|
|
RETURN_TYPES = ("LATENT",)
|
|
CATEGORY = "latent/advanced"
|
|
|
|
def splice_latents(s, sigmas, radius, texture_override, wetness=1.0, lower=None, upper=None):
|
|
#TODO: Find solution to prevent errors when nodes are muted in workflow
|
|
if lower is None and upper is None:
|
|
raise "lower and upper can't both be none"
|
|
if lower is None:
|
|
lower = torch.zeros_like(upper['samples'])
|
|
else:
|
|
lower = lower['samples']
|
|
if upper is None:
|
|
upper = torch.zeros_like(lower)
|
|
else:
|
|
upper = upper['samples']
|
|
length = radius * 2 + 1
|
|
#for 1.5, channel 3 is texture. Its feasible to
|
|
#create a kernel that blurs channels [0,1,3] and zeros 2
|
|
#this conceptually delegates sub-pixel texture to upper always,
|
|
#but is less viable for mixed configuration
|
|
#further experimentation is needed
|
|
mask = gaussian_kernel(sigmas[-1], length)
|
|
kernel = torch.stack((mask, mask, mask, mask)).unsqueeze(1)
|
|
|
|
lower_b = gaussian_blur(lower, kernel, radius)
|
|
upper_b = gaussian_blur(upper, kernel, radius)
|
|
upper_e = upper - upper_b
|
|
lower_out = lower_b * wetness + lower * (1 - wetness)
|
|
upper_out = upper_e * wetness + upper * (1 - wetness)
|
|
out = lower_out + upper_out
|
|
if texture_override == "Upper":
|
|
out[:,2] = upper[:,2]
|
|
elif texture_override == "Lower":
|
|
out[:,2] = lower[:,2]
|
|
|
|
return ({"samples": out},)
|
|
|
|
class TemporalSplice:
|
|
"""Areas of low movement are passed from lower"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
#These numbers are likely flawed
|
|
return {"required": {"sigma": ("FLOAT", {"default": 1.0, "step": .01, "min": 0}),
|
|
"wetness": ("FLOAT", {"default": 1.0, "max": 1,
|
|
"min": 0, "precision": 3,
|
|
"step": 0.1, "round": .01}),},
|
|
"optional": {"lower": ("LATENT",),
|
|
"upper": ("LATENT",)}}
|
|
FUNCTION = "temporal_splice"
|
|
RETURN_TYPES = ("LATENT",)
|
|
CATEGORY = "latent/advanced"
|
|
def temporal_splice(s, sigma, wetness, lower=None, upper=None):
|
|
if lower is None and upper is None:
|
|
raise "lower and upper can't both be none"
|
|
if lower is None:
|
|
lower = torch.zeros_like(upper['samples'])
|
|
else:
|
|
lower = lower['samples']
|
|
if upper is None:
|
|
upper = torch.zeros_like(lower)
|
|
else:
|
|
upper = upper['samples']
|
|
#Ensure odd with no overlap
|
|
length = max(lower.shape[0], upper.shape[0])
|
|
radius = (length + 1) // 2
|
|
t = torch.linspace(-radius, radius, 2 * radius + 1)
|
|
d = torch.exp(-t*t/(2*sigma*sigma))
|
|
mask = d / d.sum()
|
|
kernel = torch.stack([mask] * 4).unsqueeze(1)
|
|
def temporal_blur(tensor, kernel, radius):
|
|
#latent is B C H W, but HW C B is desired
|
|
tensor = tensor.permute((3,2,1,0))
|
|
shape = tensor.shape
|
|
tensor = tensor.reshape((shape[0]*shape[1],shape[2],shape[3]))
|
|
tensor = F.pad(tensor, [radius]*2, 'circular')
|
|
tensor = F.conv1d(tensor, kernel, padding=(radius*2+1) // 2, groups=4)
|
|
tensor = tensor[:,:,radius:-radius]
|
|
|
|
#Test code to force full blur
|
|
#tensor = tensor.mean(dim=3).unsqueeze(3).repeat(1,1,1,shape[3])
|
|
|
|
tensor = tensor.reshape(shape)
|
|
tensor = tensor.permute((3,2,1,0))
|
|
return tensor
|
|
if len(lower) == 1:
|
|
lower_b = lower
|
|
else:
|
|
if len(lower) < length //2:
|
|
lower = lower.repeat(2,1,1,1)
|
|
lower_b = temporal_blur(lower, kernel, radius)
|
|
if len(upper) == 1:
|
|
upper_b = upper
|
|
else:
|
|
if len(upper) < radius:
|
|
upper = upper.repeat(2,1,1,1)
|
|
upper_b = temporal_blur(upper, kernel, radius)
|
|
upper_e = upper - upper_b
|
|
lower_out = lower_b * wetness + lower * (1 - wetness)
|
|
upper_out = upper_e * wetness + upper * (1 - wetness)
|
|
out = lower_out + upper_out
|
|
#TODO: copy other items (mask,batch) from inputs? (also splice_latents)
|
|
return ({"samples": out},)
|
|
|
|
class SpliceDenoised:
|
|
"""A convenience node to splice latents when both noised and denoised outputs exist"""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"noised_latent" : ("LATENT",),
|
|
"denoised_latent" : ("LATENT",),
|
|
"donor_latent" : ("LATENT",),
|
|
}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "splice_denoised"
|
|
CATEGORY = "latent/advanced"
|
|
|
|
def splice_denoised(self, noised_latent, denoised_latent, donor_latent):
|
|
#Partial mask support for donor latent
|
|
donor_samples = donor_latent['samples']
|
|
if 'noise_mask' in donor_latent:
|
|
mask = prepare_mask(donor_latent['noise_mask'], donor_samples.shape, 'cpu')
|
|
donor_samples = donor_samples * (1 - mask)
|
|
donor_samples = donor_samples + denoised_latent['samples'] * mask
|
|
samples = noised_latent['samples'] - denoised_latent['samples'] + donor_samples
|
|
out = noised_latent.copy()
|
|
out['samples'] = samples
|
|
return (out,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RerangeSigmas": RerangeSigmas,
|
|
"LogSigmas": LogSigmas,
|
|
"SpliceLatents": SpliceLatents,
|
|
"SpliceDenoised": SpliceDenoised,
|
|
"TemporalSplice": TemporalSplice
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|