1adcc4c3a2
This node takes a list of sigmas and a sampler object as input. This lets people easily implement custom schedulers and samplers as nodes. More nodes will be added to it in the future.
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,
|
|
}
|