Wrap batch SchedulerOptions with result

Author: saddam213

OnnxStack.StableDiffusion/Common/BatchResult.cs

@@ -1,7 +1,11 @@
 using Microsoft.ML.OnnxRuntime.Tensors;
+using OnnxStack.Core.Image;
+using OnnxStack.Core.Video;
 using OnnxStack.StableDiffusion.Config;
 
 namespace OnnxStack.StableDiffusion.Common
 {
     public record BatchResult(SchedulerOptions SchedulerOptions, DenseTensor<float> Result);
+    public record BatchImageResult(SchedulerOptions SchedulerOptions, OnnxImage Result);
+    public record BatchVideoResult(SchedulerOptions SchedulerOptions, OnnxVideo Result);
 }

OnnxStack.StableDiffusion/Pipelines/Base/IPipeline.cs

@@ -101,7 +101,7 @@ public interface IPipeline
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        IAsyncEnumerable<OnnxImage> GenerateImageBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+        IAsyncEnumerable<BatchImageResult> GenerateImageBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
 
 
         /// <summary>
@@ -126,6 +126,6 @@ public interface IPipeline
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        IAsyncEnumerable<OnnxVideo> GenerateVideoBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+        IAsyncEnumerable<BatchVideoResult> GenerateVideoBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
     }
 }

OnnxStack.StableDiffusion/Pipelines/Base/PipelineBase.cs

@@ -133,7 +133,7 @@ protected PipelineBase(PipelineOptions pipelineOptions, ILogger logger)
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        public abstract IAsyncEnumerable<OnnxImage> GenerateImageBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+        public abstract IAsyncEnumerable<BatchImageResult> GenerateImageBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
 
 
         /// <summary>
@@ -158,7 +158,7 @@ protected PipelineBase(PipelineOptions pipelineOptions, ILogger logger)
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        public abstract IAsyncEnumerable<OnnxVideo> GenerateVideoBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+        public abstract IAsyncEnumerable<BatchVideoResult> GenerateVideoBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
 
 
         /// <summary>

OnnxStack.StableDiffusion/Pipelines/StableDiffusionPipeline.cs

@@ -291,7 +291,7 @@ public override async Task<OnnxImage> GenerateImageAsync(PromptOptions promptOpt
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        public override async IAsyncEnumerable<OnnxImage> GenerateImageBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        public override async IAsyncEnumerable<BatchImageResult> GenerateImageBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
         {
             var diffuseBatchTime = _logger?.LogBegin("Batch Diffuser starting...");
             var options = GetSchedulerOptionsOrDefault(schedulerOptions);
@@ -316,7 +316,7 @@ public override async IAsyncEnumerable<OnnxImage> GenerateImageBatchAsync(BatchO
             foreach (var batchSchedulerOption in batchSchedulerOptions)
             {
                 var tensorResult = await DiffuseImageAsync(diffuser, promptOptions, batchSchedulerOption, promptEmbeddings, performGuidance, progressCallback, cancellationToken);
-                yield return new OnnxImage(tensorResult);
+                yield return new BatchImageResult(batchSchedulerOption, new OnnxImage(tensorResult));
                 batchIndex++;
             }
 
@@ -367,7 +367,7 @@ public override async Task<OnnxVideo> GenerateVideoAsync(PromptOptions promptOpt
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        public override async IAsyncEnumerable<OnnxVideo> GenerateVideoBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        public override async IAsyncEnumerable<BatchVideoResult> GenerateVideoBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
         {
             var diffuseBatchTime = _logger?.LogBegin("Batch Diffuser starting...");
             var options = GetSchedulerOptionsOrDefault(schedulerOptions);
@@ -392,11 +392,11 @@ public override async IAsyncEnumerable<OnnxVideo> GenerateVideoBatchAsync(BatchO
             foreach (var batchSchedulerOption in batchSchedulerOptions)
             {
                 var frames = new List<OnnxImage>();
-                await foreach (var frameTensor in DiffuseVideoAsync(diffuser, promptOptions, options, promptEmbeddings, performGuidance, progressCallback, cancellationToken))
+                await foreach (var frameTensor in DiffuseVideoAsync(diffuser, promptOptions, batchSchedulerOption, promptEmbeddings, performGuidance, progressCallback, cancellationToken))
                 {
                     frames.Add(new OnnxImage(frameTensor));
                 }
-                yield return new OnnxVideo(promptOptions.InputVideo.Info, frames);
+                yield return new BatchVideoResult(batchSchedulerOption, new OnnxVideo(promptOptions.InputVideo.Info, frames));
                 batchIndex++;
             }