From bcd4249827daa64c037f969f1e37e4a45f79c95e Mon Sep 17 00:00:00 2001 From: Simon Lui <502929+simonlui@users.noreply.github.com> Date: Sat, 16 Sep 2023 17:32:45 -0700 Subject: [PATCH] Add polyexponential and vp schedulers and set device to self.device for any k-diffusion scheduler. --- comfy/samplers.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/comfy/samplers.py b/comfy/samplers.py index 57673a029..5cc68d80d 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -544,7 +544,7 @@ def encode_adm(model, conds, batch_size, width, height, device, prompt_type): class KSampler: - SCHEDULERS = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"] + SCHEDULERS = ["normal", "karras", "exponential", "polyexponential", "vp", "sgm_uniform", "simple", "ddim_uniform"] SAMPLERS = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral", "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu", "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "ddim", "uni_pc", "uni_pc_bh2"] @@ -580,9 +580,13 @@ class KSampler: discard_penultimate_sigma = True if self.scheduler == "karras": - sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=self.sigma_min, sigma_max=self.sigma_max) + sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=self.sigma_min, sigma_max=self.sigma_max, device=self.device) elif self.scheduler == "exponential": - sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=self.sigma_min, sigma_max=self.sigma_max) + sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=self.sigma_min, sigma_max=self.sigma_max, device=self.device) + elif self.scheduler == "polyexponential": + sigmas = k_diffusion_sampling.get_sigmas_polyexponential(n=steps, sigma_min=self.sigma_min, sigma_max=self.sigma_max, device=self.device) + elif self.scheduler == "vp": + sigmas = k_diffusion_sampling.get_sigmas_vp(n=steps, device=self.device) elif self.scheduler == "normal": sigmas = self.model_wrap.get_sigmas(steps) elif self.scheduler == "simple":