Fix getting steps for SamplerCustom node.

shiimizu · Nov 25, 2023 · b567095 · b567095
1 parent e9ecab1
commit b567095
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/smZNodes.py b/smZNodes.py
@@ -772,7 +772,9 @@ def find_nearest_ksampler(clip_id):
         def get_steps(graph, node_id):
             node = graph.get(str(node_id), {})
             steps_input_value = node.get("inputs", {}).get("steps", None)
-
+            if steps_input_value is None:
+                steps_input_value = node.get("inputs", {}).get("sigmas", None)
+
             while(True):
                 # Base case: it's a direct value
                 if isinstance(steps_input_value, (int, float, str)):
@@ -782,9 +784,11 @@ def get_steps(graph, node_id):
                 elif isinstance(steps_input_value, list):
                     ref_node_id, ref_input_index = steps_input_value
                     ref_node = graph.get(str(ref_node_id), {})
-                    keys = list(ref_node.get("inputs", {}).keys())
-                    ref_input_key = keys[ref_input_index % len(keys)]
-                    steps_input_value = ref_node.get("inputs", {}).get(ref_input_key)
+                    steps_input_value = ref_node.get("inputs", {}).get("steps", None)
+                    if steps_input_value is None:
+                        keys = list(ref_node.get("inputs", {}).keys())
+                        ref_input_key = keys[ref_input_index % len(keys)]
+                        steps_input_value = ref_node.get("inputs", {}).get(ref_input_key)
                 else:
                     return None