Add ProgressCallback and Cancellation support

Author: saddam213

OnnxStack.StableDiffusion/Common/ISchedulerService.cs

@@ -1,5 +1,7 @@
 using Microsoft.ML.OnnxRuntime.Tensors;
 using OnnxStack.StableDiffusion.Config;
+using System;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace OnnxStack.StableDiffusion.Common
@@ -13,6 +15,6 @@ public interface ISchedulerService
         /// <param name="prompt">The prompt.</param>
         /// <param name="options">The options.</param>
         /// <returns></returns>
-        Task<DenseTensor<float>> RunAsync(PromptOptions prompt, SchedulerOptions options);
+        Task<DenseTensor<float>> RunAsync(PromptOptions prompt, SchedulerOptions options, Action<int, int> progress = null, CancellationToken cancellationToken = default);
     }
 }

OnnxStack.StableDiffusion/Common/IStableDiffusionService.cs

@@ -1,6 +1,8 @@
 using OnnxStack.StableDiffusion.Config;
 using OnnxStack.StableDiffusion.Results;
+using System;
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace OnnxStack.StableDiffusion.Common
@@ -9,8 +11,10 @@ public interface IStableDiffusionService
     {
         Task<ImageResult> TextToImage(PromptOptions prompt);
         Task<ImageResult> TextToImage(PromptOptions prompt, SchedulerOptions options);
+        Task<ImageResult> TextToImage(PromptOptions prompt, SchedulerOptions options, Action<int, int> progress = null, CancellationToken cancellationToken = default);
 
         Task<ImageResult> TextToImageFile(PromptOptions prompt, string outputFile);
         Task<ImageResult> TextToImageFile(PromptOptions prompt, SchedulerOptions options, string outputFile);
+        Task<ImageResult> TextToImageFile(PromptOptions prompt, SchedulerOptions options, string outputFile, Action<int, int> progress = null, CancellationToken cancellationToken = default);
     }
 }

OnnxStack.StableDiffusion/Services/SchedulerService.cs

@@ -11,6 +11,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 
 
@@ -41,7 +42,7 @@ public SchedulerService(OnnxStackConfig configuration, IOnnxModelService onnxMod
         /// <param name="promptOptions">The options.</param>
         /// <param name="schedulerOptions">The scheduler configuration.</param>
         /// <returns></returns>
-        public async Task<DenseTensor<float>> RunAsync(PromptOptions promptOptions, SchedulerOptions schedulerOptions)
+        public async Task<DenseTensor<float>> RunAsync(PromptOptions promptOptions, SchedulerOptions schedulerOptions, Action<int, int> progress = null, CancellationToken cancellationToken = default)
         {
             // Create random seed if none was set
             schedulerOptions.Seed = schedulerOptions.Seed > 0 ? schedulerOptions.Seed : Random.Shared.Next();
@@ -63,6 +64,8 @@ public async Task<DenseTensor<float>> RunAsync(PromptOptions promptOptions, Sche
                 var step = 0;
                 foreach (var timestep in timesteps)
                 {
+                    cancellationToken.ThrowIfCancellationRequested();
+
                     // Create input tensor.
                     var inputTensor = scheduler.ScaleInput(latentSample.Duplicate(schedulerOptions.GetScaledDimension(2)), timestep);
 
@@ -90,6 +93,7 @@ public async Task<DenseTensor<float>> RunAsync(PromptOptions promptOptions, Sche
                     }
 
                     Console.WriteLine($"Step: {++step}/{timesteps.Count}");
+                    progress?.Invoke(step, timesteps.Count);
                 }
 
                 // Decode Latents

OnnxStack.StableDiffusion/Services/StableDiffusionService.cs

@@ -2,6 +2,8 @@
 using OnnxStack.StableDiffusion.Config;
 using OnnxStack.StableDiffusion.Helpers;
 using OnnxStack.StableDiffusion.Results;
+using System;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace OnnxStack.StableDiffusion.Services
@@ -25,6 +27,14 @@ public Task<ImageResult> TextToImage(PromptOptions prompt, SchedulerOptions opti
             return TextToImageInternal(prompt, options);
         }
 
+        public Task<ImageResult> TextToImage(PromptOptions prompt, SchedulerOptions options, Action<int, int> progress = null, CancellationToken cancellationToken = default)
+        {
+            return TextToImageInternal(prompt, options, progress, cancellationToken);
+        }
+
+
+
+
         public Task<ImageResult> TextToImageFile(PromptOptions prompt, string outputFile)
         {
             return TextToImageFileInternal(prompt, new SchedulerOptions(), outputFile);
@@ -35,16 +45,21 @@ public Task<ImageResult> TextToImageFile(PromptOptions prompt, SchedulerOptions
             return TextToImageFileInternal(prompt, options, outputFile);
         }
 
+        public Task<ImageResult> TextToImageFile(PromptOptions prompt, SchedulerOptions options, string outputFile, Action<int, int> progress = null, CancellationToken cancellationToken = default)
+        {
+            return TextToImageFileInternal(prompt, options, outputFile, progress, cancellationToken);
+        }
+
 
-        private async Task<ImageResult> TextToImageInternal(PromptOptions prompt, SchedulerOptions options)
+        private async Task<ImageResult> TextToImageInternal(PromptOptions prompt, SchedulerOptions options, Action<int, int> progress = null, CancellationToken cancellationToken = default)
         {
-            var imageTensorData = await _schedulerService.RunAsync(prompt, options).ConfigureAwait(false);
+            var imageTensorData = await _schedulerService.RunAsync(prompt, options, progress, cancellationToken).ConfigureAwait(false);
             return ImageHelpers.TensorToImage(options, imageTensorData);
         }
 
-        private async Task<ImageResult> TextToImageFileInternal(PromptOptions prompt, SchedulerOptions options, string outputFile)
+        private async Task<ImageResult> TextToImageFileInternal(PromptOptions prompt, SchedulerOptions options, string outputFile, Action<int, int> progress = null, CancellationToken cancellationToken = default)
         {
-            var result = await TextToImageInternal(prompt, options);
+            var result = await TextToImageInternal(prompt, options, progress, cancellationToken);
             if (result is null)
                 return null;