From 7cf10c64d630aa6e2cf9b5cb3d34f6e61894c5d0 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 18 Sep 2024 17:57:07 +0300
Subject: [PATCH] Update nodes.py

---
 nodes.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index 71ce2c4..a4dec80 100644
--- a/nodes.py
+++ b/nodes.py
@@ -455,6 +455,8 @@ class CogVideoXFunSampler:
                         "DPM++",
                         "PNDM",
                         "DDIM",
+                        "CogVideoXDDIM",
+                        "CogVideoXDPMScheduler",
                     ],
                     {
                         "default": 'DDIM'
@@ -502,7 +504,11 @@ class CogVideoXFunSampler:
         elif scheduler == "PNDM":
             noise_scheduler = PNDMScheduler.from_pretrained(base_path, subfolder= 'scheduler')
         elif scheduler == "DDIM":
-            noise_scheduler = DDIMScheduler.from_pretrained(base_path, subfolder= 'scheduler')
+            noise_scheduler = CogVideoXDDIMScheduler.from_pretrained(base_path, subfolder= 'scheduler')
+        elif scheduler == "CogVideoXDDIM":
+            noise_scheduler = CogVideoXDDIMScheduler.from_pretrained(base_path, subfolder= 'scheduler')
+        elif scheduler == "CogVideoXDPMScheduler":
+            noise_scheduler = CogVideoXDPMScheduler.from_pretrained(base_path, subfolder= 'scheduler')
         pipe.scheduler = noise_scheduler
 
         #if not pipeline["cpu_offloading"]: