132 lines
6.1 KiB
Python
132 lines
6.1 KiB
Python
from nodes import VAEDecode, EmptyLatentImage, CLIPTextEncode
|
|
from comfy_extras.nodes_custom_sampler import SamplerCustom, BasicScheduler, PolyexponentialScheduler, VPScheduler
|
|
from comfy_extras.nodes_align_your_steps import AlignYourStepsScheduler
|
|
from server import PromptServer
|
|
from ..sn0w import Logger, MessageHolder
|
|
from .custom_schedulers.get_sigmas_sigmoid import get_sigmas_sigmoid
|
|
import comfy.samplers
|
|
|
|
class SimpleSamplerCustom:
|
|
logger = Logger()
|
|
scheduler_list = comfy.samplers.KSampler.SCHEDULERS + ["align_your_steps", "polyexponential", "vp", "sigmoid"]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"model_type": (["SD1", "SDXL", "SVD"], ),
|
|
"clip": ("CLIP", ),
|
|
"vae": ("VAE", ),
|
|
"add_noise": ("BOOLEAN", {"default": True}),
|
|
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"steps": ("INT", {"default": 10, "min": 10, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
|
|
"scheduler_name": (cls.scheduler_list, ),
|
|
"width": ("INT", {"default": 0, "min": 0, "step":64}),
|
|
"height": ("INT", {"default": 0, "min": 0, "step":64}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step":0.01, "round": 0.01}),
|
|
},
|
|
"optional": {
|
|
"positive": ("*"),
|
|
"negative": ("*"),
|
|
"sigmas (optional)": ("SIGMAS", ),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "LATENT", "CONDITIONING", "CONDITIONING")
|
|
RETURN_NAMES = ("IMAGE", "LATENT", "POSITIVE", "NEGATIVE")
|
|
|
|
FUNCTION = "sample"
|
|
|
|
CATEGORY = "sampling/custom_sampling"
|
|
|
|
def sample(self, model, model_type, clip, vae, add_noise, noise_seed, steps, cfg, sampler_name, scheduler_name, width, height, denoise, **kwargs):
|
|
custom_sampler = SamplerCustom()
|
|
vae_decode = VAEDecode()
|
|
text_encode = CLIPTextEncode()
|
|
|
|
# Encode inputs
|
|
positive_prompt = self.get_prompt("positive", text_encode, clip, kwargs)
|
|
negative_prompt = self.get_prompt("negative", text_encode, clip, kwargs)
|
|
|
|
# Create latent
|
|
latent_image = EmptyLatentImage().generate(width, height)[0]
|
|
|
|
# Get sampler and sigmas
|
|
sampler = self.get_sampler(sampler_name)[0]
|
|
if 'sigmas (optional)' in kwargs and kwargs['sigmas (optional)'] is not None:
|
|
sigmas = kwargs['sigmas (optional)']
|
|
else:
|
|
sigmas = self.get_custom_sigmas(model, model_type, scheduler_name, steps, denoise, kwargs["unique_id"])
|
|
|
|
# Generate and decode image
|
|
samples = custom_sampler.sample(model, add_noise, noise_seed, cfg, positive_prompt, negative_prompt, sampler, sigmas, latent_image)
|
|
image = vae_decode.decode(vae, samples[1])
|
|
|
|
return (image[0], samples[1], positive_prompt, negative_prompt)
|
|
|
|
def get_sampler(self, sampler_name):
|
|
sampler = comfy.samplers.sampler_object(sampler_name)
|
|
return (sampler, )
|
|
|
|
def get_scheduler_values(self, unique_id, widgets_needed):
|
|
PromptServer.instance.send_sync("get_scheduler_values", {
|
|
"id": unique_id,
|
|
"widgets_needed": widgets_needed,
|
|
})
|
|
outputs = MessageHolder.waitForMessage(unique_id)
|
|
return outputs
|
|
|
|
def get_custom_sigmas(self, model, model_type, scheduler_name, steps, denoise, unique_id):
|
|
if scheduler_name == "align_your_steps":
|
|
sigmas = AlignYourStepsScheduler.get_sigmas(self, model_type, steps, denoise)[0]
|
|
elif scheduler_name == "polyexponential":
|
|
values = self.get_scheduler_values(unique_id, 3)
|
|
|
|
sigma_max = values["sigma_max_poly"]["value"]
|
|
sigma_min = values["sigma_min_poly"]["value"]
|
|
rho = values["rho"]["value"]
|
|
self.logger.log(f"Sigma Max: {sigma_max}, Sigma Min: {sigma_min}, Rho: {rho}", "DEBUG")
|
|
|
|
sigmas = PolyexponentialScheduler.get_sigmas(self, steps, sigma_max, sigma_min, rho)[0]
|
|
elif scheduler_name == "vp":
|
|
values = self.get_scheduler_values(unique_id, 3)
|
|
|
|
beta_d = values["beta_d"]["value"]
|
|
beta_min = values["beta_min"]["value"]
|
|
eps_s = values["eps_s"]["value"]
|
|
self.logger.log(f"Beta D: {beta_d}, Beta Min: {beta_min}, Eps S: {eps_s}", "DEBUG")
|
|
|
|
sigmas = VPScheduler.get_sigmas(self, steps, beta_d, beta_min, eps_s)[0]
|
|
elif scheduler_name == "sigmoid":
|
|
values = self.get_scheduler_values(unique_id, 4)
|
|
|
|
midpoint_ratio = values["midpoint_ratio"]["value"]
|
|
sigma_max = values["sigma_max_sig"]["value"]
|
|
sigma_min = values["sigma_min_sig"]["value"]
|
|
steepness = values["steepness"]["value"]
|
|
self.logger.log(f"Sigma Max: {sigma_max}, Sigma Min: {sigma_min}, Steepness: {steepness}, Midpoint Ratio: {midpoint_ratio}", "DEBUG")
|
|
|
|
sigmas = get_sigmas_sigmoid(self, steps, sigma_max, sigma_min, steepness, midpoint_ratio)[0]
|
|
else:
|
|
sigmas = BasicScheduler.get_sigmas(self, model, scheduler_name, steps, denoise)[0]
|
|
|
|
self.logger.print_sigmas_differences(scheduler_name, sigmas)
|
|
return sigmas
|
|
|
|
def get_prompt(self, name, text_encode, clip, kwargs):
|
|
if name in kwargs and kwargs[name] is not None:
|
|
if isinstance(kwargs[name], str):
|
|
return text_encode.encode(clip, kwargs[name])[0]
|
|
elif isinstance(kwargs[name], list):
|
|
return kwargs[name]
|
|
else:
|
|
log = f"{name} prompt cannot be {type(kwargs[name]).__name__}, it has to be either a string or conditioning."
|
|
self.logger.log(log, "ERROR")
|
|
raise TypeError(log)
|