Add a basic node to generate sigmas from scheduler.

This commit is contained in:
comfyanonymous 2023-09-28 00:30:45 -04:00
parent d234ca558a
commit 2bf051fda8
1 changed files with 21 additions and 0 deletions

View File

@ -4,6 +4,26 @@ from comfy.k_diffusion import sampling as k_diffusion_sampling
import latent_preview
import torch
class BasicScheduler:
@classmethod
def INPUT_TYPES(s):
return {"required":
{"model": ("MODEL",),
"scheduler": (comfy.samplers.SCHEDULER_NAMES, ),
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
}
}
RETURN_TYPES = ("SIGMAS",)
CATEGORY = "_for_testing/custom_sampling"
FUNCTION = "get_sigmas"
def get_sigmas(self, model, scheduler, steps):
sigmas = comfy.samplers.calculate_sigmas_scheduler(model.model, scheduler, steps).cpu()
return (sigmas, )
class KarrasScheduler:
@classmethod
def INPUT_TYPES(s):
@ -95,4 +115,5 @@ NODE_CLASS_MAPPINGS = {
"SamplerCustom": SamplerCustom,
"KarrasScheduler": KarrasScheduler,
"KSamplerSelect": KSamplerSelect,
"BasicScheduler": BasicScheduler,
}