94 lines
2.8 KiB
Python
94 lines
2.8 KiB
Python
import torch
|
|
from tqdm import trange
|
|
|
|
from comfy.samplers import KSAMPLER
|
|
|
|
|
|
class SamplerInversedEulerNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {}}
|
|
|
|
RETURN_TYPES = ("SAMPLER",)
|
|
CATEGORY = "sampling/custom_sampling/samplers"
|
|
|
|
FUNCTION = "get_sampler"
|
|
|
|
def get_sampler(self):
|
|
@torch.no_grad()
|
|
def sample_inversed_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
|
|
"""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]])
|
|
for i in trange(1, len(sigmas), disable=disable):
|
|
sigma_in = sigmas[i-1]
|
|
|
|
if i == 1:
|
|
sigma_t = sigmas[i]
|
|
else:
|
|
sigma_t = sigma_in
|
|
|
|
denoised = model(x, sigma_t * s_in, **extra_args)
|
|
|
|
if i == 1:
|
|
d = (x - denoised) / (2 * sigmas[i])
|
|
else:
|
|
d = (x - denoised) / sigmas[i-1]
|
|
|
|
dt = sigmas[i] - sigmas[i-1]
|
|
x = x + d * dt
|
|
if callback is not None:
|
|
callback(
|
|
{'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
|
|
return x / sigmas[-1]
|
|
|
|
ksampler = KSAMPLER(sample_inversed_euler)
|
|
return (ksampler, )
|
|
|
|
|
|
class MixNoiseNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"latent": ("LATENT", ),
|
|
"randomness": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01, "round": False}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
CATEGORY = "noise"
|
|
|
|
FUNCTION = "mix"
|
|
|
|
def mix(self, latent, randomness):
|
|
latent_image = latent['samples']
|
|
rand_noise = torch.randn_like(latent_image)
|
|
combined_noise = ((1 - randomness) * latent_image + randomness *
|
|
rand_noise) / ((randomness**2 + (1-randomness)**2) ** 0.5)
|
|
latent = latent.copy()
|
|
latent['samples'] = combined_noise
|
|
return (latent, )
|
|
|
|
|
|
class CombineNoiseLatentNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"latent": ("LATENT", ),
|
|
"noise": ("LATENT",),
|
|
"sigmas": ("SIGMAS", ),
|
|
}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("noise",)
|
|
CATEGORY = "noise"
|
|
|
|
FUNCTION = "combine"
|
|
|
|
def combine(self, latent, noise, sigmas):
|
|
latent_image = latent['samples']
|
|
noise_image = noise['samples']
|
|
noise_dt = noise_image - (latent_image / sigmas[0])
|
|
noise = noise.copy()
|
|
noise['samples'] = noise_dt
|
|
return (noise, )
|