From 910bc0e4d1020b58ce2ed3d4d01d2162e02ab4b9 Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Mon, 29 Apr 2024 20:16:03 +0900
Subject: [PATCH] feat: ImpactSchedulerAdapter

https://github.com/ltdrdata/ComfyUI-Impact-Pack/issues/576
https://github.com/ltdrdata/ComfyUI-Impact-Pack/issues/568
---
 __init__.py                   |  8 ++++++--
 modules/impact/config.py      |  2 +-
 modules/impact/impact_pack.py | 23 +++++++++++++++++++++++
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/__init__.py b/__init__.py
index 847fa8a..02be7f3 100644
--- a/__init__.py
+++ b/__init__.py
@@ -306,7 +306,9 @@ NODE_CLASS_MAPPINGS = {
     "ImpactRemoteInt": ImpactRemoteInt,
 
     "ImpactHFTransformersClassifierProvider": HF_TransformersClassifierProvider,
-    "ImpactSEGSClassify": SEGS_Classify
+    "ImpactSEGSClassify": SEGS_Classify,
+
+    "ImpactSchedulerAdapter": ImpactSchedulerAdapter
 }
 
 
@@ -433,7 +435,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "LatentSwitch": "Switch (latent/legacy)",
     "SEGSSwitch": "Switch (SEGS/legacy)",
 
-    "SEGSPreviewCNet": "SEGSPreview (CNET Image)"
+    "SEGSPreviewCNet": "SEGSPreview (CNET Image)",
+
+    "ImpactSchedulerAdapter": "Impact Scheduler Adapter",
 }
 
 if not impact.config.get_config()['mmdet_skip']:
diff --git a/modules/impact/config.py b/modules/impact/config.py
index 6bc8b2f..5ac4103 100644
--- a/modules/impact/config.py
+++ b/modules/impact/config.py
@@ -2,7 +2,7 @@ import configparser
 import os
 
 
-version_code = [5, 0, 3]
+version_code = [5, 1, 3]
 version = f"V{version_code[0]}.{version_code[1]}" + (f'.{version_code[2]}' if len(version_code) > 2 else '')
 
 dependency_version = 20
diff --git a/modules/impact/impact_pack.py b/modules/impact/impact_pack.py
index 8c7bd34..e193bb0 100644
--- a/modules/impact/impact_pack.py
+++ b/modules/impact/impact_pack.py
@@ -2150,3 +2150,26 @@ class ImpactWildcardEncode:
         populated = kwargs['populated_text']
         model, clip, conditioning = impact.wildcards.process_with_loras(populated, kwargs['model'], kwargs['clip'])
         return (model, clip, conditioning, populated)
+
+
+class ImpactSchedulerAdapter:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"defaultInput": True,}),
+            "ays_scheduler": (['None', 'AYS SDXL', 'AYS SD1', 'AYS SVD'],),
+        }}
+
+    CATEGORY = "ImpactPack/Util"
+
+    RETURN_TYPES = (core.SCHEDULERS,)
+    RETURN_NAMES = ("scheduler",)
+
+    FUNCTION = "doit"
+
+    def doit(self, scheduler, ays_scheduler):
+        if ays_scheduler != 'None':
+            return (ays_scheduler,)
+
+        return (scheduler,)
+