From 836b26aff3944a06f592d7446b2f68159ab48bb8 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Fri, 10 Jan 2025 00:36:58 +0200
Subject: [PATCH] scheduler fixes

---
 nodes.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/nodes.py b/nodes.py
index d0290be..4f805b9 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1177,6 +1177,10 @@ class HyVideoSampler:
         scheduler_config["flow_shift"] = flow_shift
         if scheduler == "SDE-DPMSolverMultistepScheduler":
             scheduler_config["algorithm_type"] = "sde-dpmsolver++"
+        elif scheduler == "SASolverScheduler":
+            scheduler_config["algorithm_type"] = "data_prediction"
+        else:
+            scheduler_config.pop("algorithm_type", None)
         #model["scheduler_config"]["use_beta_flow_sigmas"] = True
         
         noise_scheduler = scheduler_mapping[scheduler].from_config(scheduler_config)