99 lines
3.6 KiB
Python
99 lines
3.6 KiB
Python
|
import comfy.samplers
|
||
|
|
import comfy.sample
|
||
|
|
from comfy.k_diffusion import sampling as k_diffusion_sampling
|
||
|
|
import latent_preview
|
||
|
|
|
||
|
|
|
||
|
|
class KarrasScheduler:
|
||
|
|
@classmethod
|
||
|
|
def INPUT_TYPES(s):
|
||
|
|
return {"required":
|
||
|
|
{"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
||
|
|
"sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
|
||
|
|
"sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
|
||
|
|
"rho": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
RETURN_TYPES = ("SIGMAS",)
|
||
|
|
CATEGORY = "_for_testing/custom_sampling"
|
||
|
|
|
||
|
|
FUNCTION = "get_sigmas"
|
||
|
|
|
||
|
|
def get_sigmas(self, steps, sigma_max, sigma_min, rho):
|
||
|
|
sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho)
|
||
|
|
return (sigmas, )
|
||
|
|
|
||
|
|
|
||
|
|
class KSamplerSelect:
|
||
|
|
@classmethod
|
||
|
|
def INPUT_TYPES(s):
|
||
|
|
return {"required":
|
||
|
|
{"sampler_name": (comfy.samplers.KSAMPLER_NAMES, ),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
RETURN_TYPES = ("SAMPLER",)
|
||
|
|
CATEGORY = "_for_testing/custom_sampling"
|
||
|
|
|
||
|
|
FUNCTION = "get_sampler"
|
||
|
|
|
||
|
|
def get_sampler(self, sampler_name):
|
||
|
|
sampler = comfy.samplers.ksampler(sampler_name)()
|
||
|
|
return (sampler, )
|
||
|
|
|
||
|
|
class SamplerCustom:
|
||
|
|
@classmethod
|
||
|
|
def INPUT_TYPES(s):
|
||
|
|
return {"required":
|
||
|
|
{"model": ("MODEL",),
|
||
|
|
"add_noise": (["enable", "disable"], ),
|
||
|
|
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.5, "round": 0.01}),
|
||
|
|
"positive": ("CONDITIONING", ),
|
||
|
|
"negative": ("CONDITIONING", ),
|
||
|
|
"sampler": ("SAMPLER", ),
|
||
|
|
"sigmas": ("SIGMAS", ),
|
||
|
|
"latent_image": ("LATENT", ),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
RETURN_TYPES = ("LATENT","LATENT")
|
||
|
|
RETURN_NAMES = ("output", "denoised_output")
|
||
|
|
|
||
|
|
FUNCTION = "sample"
|
||
|
|
|
||
|
|
CATEGORY = "_for_testing/custom_sampling"
|
||
|
|
|
||
|
|
def sample(self, model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image):
|
||
|
|
latent = latent_image
|
||
|
|
latent_image = latent["samples"]
|
||
|
|
if add_noise == "disable":
|
||
|
|
noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
|
||
|
|
else:
|
||
|
|
batch_inds = latent["batch_index"] if "batch_index" in latent else None
|
||
|
|
noise = comfy.sample.prepare_noise(latent_image, noise_seed, batch_inds)
|
||
|
|
|
||
|
|
noise_mask = None
|
||
|
|
if "noise_mask" in latent:
|
||
|
|
noise_mask = latent["noise_mask"]
|
||
|
|
|
||
|
|
x0_output = {}
|
||
|
|
callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
|
||
|
|
|
||
|
|
disable_pbar = False
|
||
|
|
samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
|
||
|
|
|
||
|
|
out = latent.copy()
|
||
|
|
out["samples"] = samples
|
||
|
|
if "x0" in x0_output:
|
||
|
|
out_denoised = latent.copy()
|
||
|
|
out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu())
|
||
|
|
else:
|
||
|
|
out_denoised = out
|
||
|
|
return (out, out_denoised)
|
||
|
|
|
||
|
|
NODE_CLASS_MAPPINGS = {
|
||
|
|
"SamplerCustom": SamplerCustom,
|
||
|
|
"KarrasScheduler": KarrasScheduler,
|
||
|
|
"KSamplerSelect": KSamplerSelect,
|
||
|
|
}
|