From 75565346bf1041b52f3da65c3f7cb9a3b9772eae Mon Sep 17 00:00:00 2001 From: Phr00t Date: Sat, 21 Sep 2024 19:59:37 -0400 Subject: [PATCH] adding Heun and DPMSingleStep samplers --- nodes.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/nodes.py b/nodes.py index c61984c..b69b589 100644 --- a/nodes.py +++ b/nodes.py @@ -3,7 +3,7 @@ import torch import folder_paths import comfy.model_management as mm from comfy.utils import ProgressBar, load_torch_file -from diffusers.schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler, DDIMScheduler, PNDMScheduler, DPMSolverMultistepScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, LCMScheduler, UniPCMultistepScheduler, TCDScheduler +from diffusers.schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler, DDIMScheduler, PNDMScheduler, DPMSolverMultistepScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, LCMScheduler, UniPCMultistepScheduler, TCDScheduler, HeunDiscreteScheduler, DPMSolverSinglestepScheduler from diffusers.models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel from .pipeline_cogvideox import CogVideoXPipeline @@ -701,6 +701,8 @@ class CogVideoXFunSampler: "LCM", "UniPC", "TCD", + "Heun", + "DPMSingleStep", "CogVideoXDDIM", "CogVideoXDPMScheduler", ], @@ -762,6 +764,10 @@ class CogVideoXFunSampler: noise_scheduler = UniPCMultistepScheduler.from_config(scheduler_config) elif scheduler == "TCD": noise_scheduler = TCDScheduler.from_config(scheduler_config) + elif scheduler == "Heun": + noise_scheduler = HeunDiscreteScheduler.from_config(scheduler_config) + elif scheduler == "DPMSingleStep": + noise_scheduler = DPMSolverSinglestepScheduler.from_config(scheduler_config) elif scheduler == "CogVideoXDDIM": noise_scheduler = CogVideoXDDIMScheduler.from_config(scheduler_config) elif scheduler == "CogVideoXDPMScheduler": @@ -823,6 +829,8 @@ class CogVideoXFunVid2VidSampler: "LCM", "UniPC", "TCD", + "Heun", + "DPMSingleStep", "CogVideoXDDIM", "CogVideoXDPMScheduler", ], @@ -873,6 +881,10 @@ class CogVideoXFunVid2VidSampler: noise_scheduler = UniPCMultistepScheduler.from_pretrained(base_path, subfolder= 'scheduler') elif scheduler == "TCD": noise_scheduler = TCDScheduler.from_pretrained(base_path, subfolder= 'scheduler') + elif scheduler == "Heun": + noise_scheduler = HeunDiscreteScheduler.from_pretrained(base_path, subfolder= 'scheduler') + elif scheduler == "DPMSingleStep": + noise_scheduler = DPMSolverSinglestepScheduler.from_pretrained(base_path, subfolder= 'scheduler') elif scheduler == "CogVideoXDDIM": noise_scheduler = CogVideoXDDIMScheduler.from_pretrained(base_path, subfolder= 'scheduler') elif scheduler == "CogVideoXDPMScheduler":