Fix step automation increment

TensorStack-AI · Nov 16, 2023 · eb09a0a · eb09a0a
1 parent 9417b17
commit eb09a0a
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/OnnxStack.StableDiffusion/Helpers/BatchGenerator.cs b/OnnxStack.StableDiffusion/Helpers/BatchGenerator.cs
@@ -26,8 +26,9 @@ public static List<SchedulerOptions> GenerateBatch(IModelOptions modelOptions, B
             }
             else if (batchOptions.BatchType == BatchOptionType.Step)
             {
-                return Enumerable.Range(Math.Max(0, (int)batchOptions.ValueFrom), Math.Max(1, (int)batchOptions.ValueTo))
-                  .Select(x => schedulerOptions with { InferenceSteps = x })
+                var totalIncrements = (int)Math.Max(1, (batchOptions.ValueTo - batchOptions.ValueFrom) / batchOptions.Increment);
+                return Enumerable.Range(0, totalIncrements)
+                  .Select(x => schedulerOptions with { InferenceSteps = (int)(batchOptions.ValueFrom + (batchOptions.Increment * x)) })
                   .ToList();
             }
             else if (batchOptions.BatchType == BatchOptionType.Guidance)