Console ProgressCallback, Fix batch scheduler issue

Author: saddam213

OnnxStack.Console/Examples/StableDebug.cs

@@ -66,7 +66,7 @@ public async Task RunAsync()
                         OutputHelpers.WriteConsole($"Generating {scheduler} Image...", ConsoleColor.Green);
 
                         // Run pipeline
-                        var result = await pipeline.RunAsync(promptOptions, schedulerOptions);
+                        var result = await pipeline.RunAsync(promptOptions, schedulerOptions, progressCallback: OutputHelpers.ProgressCallback);
 
                         // Create Image from Tensor result
                         var image = result.ToImage();

OnnxStack.Console/Examples/StableDiffusionBatch.cs

@@ -56,12 +56,9 @@ public async Task RunAsync()
                 // Preload Models (optional)
                 await pipeline.LoadAsync();
 
-                // Progress Callback (optional)
-                var progressCallback = (DiffusionProgress progress) => OutputHelpers.WriteConsole($"Image: {progress.BatchValue}/{progress.BatchMax} - Step: {progress.StepValue}/{progress.StepMax}", ConsoleColor.Cyan);
-
                 // Run Batch
                 var timestamp = Stopwatch.GetTimestamp();
-                await foreach (var result in pipeline.RunBatchAsync(batchOptions, promptOptions, progressCallback: progressCallback))
+                await foreach (var result in pipeline.RunBatchAsync(batchOptions, promptOptions, progressCallback: OutputHelpers.BatchProgressCallback))
                 {
                     // Create Image from Tensor result
                     var image = result.ImageResult.ToImage();

OnnxStack.Console/Examples/StableDiffusionExample.cs

@@ -67,7 +67,7 @@ public async Task RunAsync()
                         OutputHelpers.WriteConsole($"Generating '{schedulerType}' Image...", ConsoleColor.Green);
 
                         // Run pipeline
-                        var result = await pipeline.RunAsync(promptOptions, schedulerOptions);
+                        var result = await pipeline.RunAsync(promptOptions, schedulerOptions, progressCallback: OutputHelpers.ProgressCallback);
 
                         // Create Image from Tensor result
                         var image = result.ToImage();

OnnxStack.Console/Examples/StableDiffusionGenerator.cs

@@ -55,7 +55,7 @@ public async Task RunAsync()
                     OutputHelpers.WriteConsole($"Generating '{generationPrompt.Key}'", ConsoleColor.Green);
 
                     // Run pipeline
-                    var result = await pipeline.RunAsync(promptOptions);
+                    var result = await pipeline.RunAsync(promptOptions, progressCallback: OutputHelpers.ProgressCallback);
 
                     // Create Image from Tensor result
                     var image = result.ToImage();

OnnxStack.Console/Examples/VideoToVideoExample.cs

@@ -34,10 +34,6 @@ public async Task RunAsync()
             var videoInfo = await _videoService.GetVideoInfoAsync(inputFile);
             var videoInput = await _videoService.CreateFramesAsync(inputFile, videoInfo.FPS);
 
-            // Progress Callback (optional)
-            var progressCallback = (DiffusionProgress progress) => OutputHelpers.WriteConsole($"Frame: {progress.BatchValue}/{progress.BatchMax} - Step: {progress.StepValue}/{progress.StepMax}", ConsoleColor.Cyan);
-
-
             // Loop though the appsettings.json model sets
             foreach (var modelSet in _configuration.ModelSets)
             {
@@ -58,7 +54,7 @@ public async Task RunAsync()
                 };
 
                 // Run pipeline
-                var result = await pipeline.RunAsync(promptOptions, progressCallback: progressCallback);
+                var result = await pipeline.RunAsync(promptOptions, progressCallback: OutputHelpers.FrameProgressCallback);
 
                 // Create Video from Tensor result
                 var videoResult = await _videoService.CreateVideoAsync(result, videoInfo.FPS);

OnnxStack.Console/OutputHelpers.cs

@@ -1,4 +1,6 @@
-namespace OnnxStack.Console
+using OnnxStack.StableDiffusion.Common;
+
+namespace OnnxStack.Console
 {
     internal static class OutputHelpers
     {
@@ -22,5 +24,9 @@ public static void WriteConsole(string value, ConsoleColor color, bool line = tr
                 System.Console.Write(value);
             System.Console.ForegroundColor = previous;
         }
+
+        public static Action<DiffusionProgress> ProgressCallback => (DiffusionProgress progress) => WriteConsole($"Step: {progress.StepValue}/{progress.StepMax}", ConsoleColor.Gray);
+        public static Action<DiffusionProgress> BatchProgressCallback => (DiffusionProgress progress) => WriteConsole($"Batch: {progress.BatchValue}/{progress.BatchMax} - Step: {progress.StepValue}/{progress.StepMax}", ConsoleColor.Gray);
+        public static Action<DiffusionProgress> FrameProgressCallback => (DiffusionProgress progress) => WriteConsole($"Frame: {progress.BatchValue}/{progress.BatchMax} - Step: {progress.StepValue}/{progress.StepMax}", ConsoleColor.Gray);
     }
 }

OnnxStack.StableDiffusion/Extensions.cs

@@ -88,43 +88,5 @@ internal static int[] GetScaledDimension(this SchedulerOptions options, int batc
         {
             return new[] { batch, channels, options.GetScaledHeight(), options.GetScaledWidth() };
         }
-
-
-        /// <summary>
-        /// Gets the pipeline schedulers.
-        /// </summary>
-        /// <param name="pipelineType">Type of the pipeline.</param>
-        /// <returns></returns>
-        public static SchedulerType[] GetSchedulerTypes(this DiffuserPipelineType pipelineType)
-        {
-            switch (pipelineType)
-            {
-                case DiffuserPipelineType.StableDiffusion:
-                case DiffuserPipelineType.StableDiffusionXL:
-                    return new[]
-                    {
-                        SchedulerType.LMS,
-                        SchedulerType.Euler,
-                        SchedulerType.EulerAncestral,
-                        SchedulerType.DDPM,
-                        SchedulerType.DDIM,
-                        SchedulerType.KDPM2
-                    };
-                case DiffuserPipelineType.LatentConsistency:
-                case DiffuserPipelineType.LatentConsistencyXL:
-                    return new[] 
-                    { 
-                        SchedulerType.LCM 
-                    };
-                case DiffuserPipelineType.InstaFlow:
-                    return new[] 
-                    {
-                        SchedulerType.InstaFlow
-                    };
-                default:
-                    return default;
-            }
-        }
-
     }
 }

OnnxStack.StableDiffusion/Helpers/BatchGenerator.cs

@@ -1,5 +1,6 @@
 using OnnxStack.StableDiffusion.Config;
 using OnnxStack.StableDiffusion.Enums;
+using OnnxStack.StableDiffusion.Pipelines;
 using System;
 using System.Collections.Generic;
 using System.Linq;
@@ -14,7 +15,7 @@ public static class BatchGenerator
         /// <param name="batchOptions">The batch options.</param>
         /// <param name="schedulerOptions">The scheduler options.</param>
         /// <returns></returns>
-        public static List<SchedulerOptions> GenerateBatch(DiffuserPipelineType pipelineType, BatchOptions batchOptions, SchedulerOptions schedulerOptions)
+        public static List<SchedulerOptions> GenerateBatch(IPipeline pipeline, BatchOptions batchOptions, SchedulerOptions schedulerOptions)
         {
             if (batchOptions.BatchType == BatchOptionType.Seed)
             {
@@ -46,7 +47,7 @@ public static List<SchedulerOptions> GenerateBatch(DiffuserPipelineType pipeline
             }
             else if (batchOptions.BatchType == BatchOptionType.Scheduler)
             {
-                return pipelineType.GetSchedulerTypes()
+                return pipeline.SupportedSchedulers
                   .Select(x => schedulerOptions with { SchedulerType = x })
                   .ToList();
             }

OnnxStack.StableDiffusion/Pipelines/StableDiffusionPipeline.cs

@@ -209,7 +209,7 @@ public override async IAsyncEnumerable<BatchResult> RunBatchAsync(BatchOptions b
             var promptEmbeddings = await CreatePromptEmbedsAsync(promptOptions, performGuidance);
 
             // Generate batch options
-            var batchSchedulerOptions = BatchGenerator.GenerateBatch(DiffuserPipelineType.StableDiffusion, batchOptions, schedulerOptions);
+            var batchSchedulerOptions = BatchGenerator.GenerateBatch(this, batchOptions, schedulerOptions);
 
             // Create Diffuser
             var diffuser = CreateDiffuser(promptOptions.DiffuserType, controlNet);