PromptService support for dual Encoders, Tokenizers

saddam213 · saddam213 · commit e4f86a2fca41 · 2023-11-26T14:51:37.000+13:00
diff --git a/OnnxStack.StableDiffusion/Common/IPromptService.cs b/OnnxStack.StableDiffusion/Common/IPromptService.cs
@@ -6,8 +6,6 @@ namespace OnnxStack.StableDiffusion.Common
 {
     public interface IPromptService
     {
-        Task<DenseTensor<float>> CreatePromptAsync(IModelOptions model, PromptOptions promptOptions, bool isGuidanceEnabled);
-        Task<int[]> DecodeTextAsync(IModelOptions model, string inputText);
-        Task<float[]> EncodeTokensAsync(IModelOptions model, int[] tokenizedInput);
+        Task<PromptEmbeddingsResult> CreatePromptAsync(IModelOptions model, PromptOptions promptOptions, bool isGuidanceEnabled);
     }
 }
diff --git a/OnnxStack.StableDiffusion/Common/PromptResult.cs b/OnnxStack.StableDiffusion/Common/PromptResult.cs
@@ -0,0 +1,6 @@
+﻿using Microsoft.ML.OnnxRuntime.Tensors;
+
+namespace OnnxStack.StableDiffusion.Common
+{
+    public record PromptEmbeddingsResult(DenseTensor<float> PromptEmbeds, DenseTensor<float> PooledPromptEmbeds = default);
+}
diff --git a/OnnxStack.StableDiffusion/Diffusers/DiffuserBase.cs b/OnnxStack.StableDiffusion/Diffusers/DiffuserBase.cs
@@ -88,7 +88,7 @@ public DiffuserBase(IOnnxModelService onnxModelService, IPromptService promptSer
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        protected abstract Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, DenseTensor<float> promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default);
+        protected abstract Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, PromptEmbeddingsResult promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default);
 
 
         /// <summary>
diff --git a/OnnxStack.StableDiffusion/Diffusers/InstaFlow/InstaFlowDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/InstaFlow/InstaFlowDiffuser.cs
@@ -45,7 +45,7 @@ public InstaFlowDiffuser(IOnnxModelService onnxModelService, IPromptService prom
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, DenseTensor<float> promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
+        protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, PromptEmbeddingsResult promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
         {
             // Get Scheduler
             using (var scheduler = GetScheduler(schedulerOptions))
@@ -81,7 +81,7 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio
                     {
                         inferenceParameters.AddInputTensor(inputTensor);
                         inferenceParameters.AddInputTensor(timestepTensor);
-                        inferenceParameters.AddInputTensor(promptEmbeddings);
+                        inferenceParameters.AddInputTensor(promptEmbeddings.PromptEmbeds);
                         inferenceParameters.AddOutputBuffer(outputDimension);
 
                         var results = await _onnxModelService.RunInferenceAsync(modelOptions, OnnxModelType.Unet, inferenceParameters);
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/InpaintLegacyDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/InpaintLegacyDiffuser.cs
@@ -65,7 +65,7 @@ protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, ISc
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, DenseTensor<float> promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
+        protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, PromptEmbeddingsResult promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
         {
             using (var scheduler = GetScheduler(schedulerOptions))
             {
@@ -111,7 +111,7 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio
                     {
                         inferenceParameters.AddInputTensor(inputTensor);
                         inferenceParameters.AddInputTensor(timestepTensor);
-                        inferenceParameters.AddInputTensor(promptEmbeddings);
+                        inferenceParameters.AddInputTensor(promptEmbeddings.PromptEmbeds);
                         inferenceParameters.AddInputTensor(guidanceEmbeddings);
                         inferenceParameters.AddOutputBuffer(outputDimension);
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/LatentConsistencyDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/LatentConsistencyDiffuser.cs
@@ -88,7 +88,7 @@ protected override bool ShouldPerformGuidance(SchedulerOptions schedulerOptions)
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, DenseTensor<float> promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
+        protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, PromptEmbeddingsResult promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
         {
             // Get Scheduler
             using (var scheduler = GetScheduler(schedulerOptions))
@@ -126,7 +126,7 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio
                     {
                         inferenceParameters.AddInputTensor(inputTensor);
                         inferenceParameters.AddInputTensor(timestepTensor);
-                        inferenceParameters.AddInputTensor(promptEmbeddings);
+                        inferenceParameters.AddInputTensor(promptEmbeddings.PromptEmbeds);
                         inferenceParameters.AddInputTensor(guidanceEmbeddings);
                         inferenceParameters.AddOutputBuffer(outputDimension);
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintDiffuser.cs
@@ -53,7 +53,7 @@ public InpaintDiffuser(IOnnxModelService onnxModelService, IPromptService prompt
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, DenseTensor<float> promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
+        protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, PromptEmbeddingsResult promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
         {
             // Get Scheduler
             using (var scheduler = GetScheduler(schedulerOptions))
@@ -93,7 +93,7 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio
                     {
                         inferenceParameters.AddInputTensor(inputTensor);
                         inferenceParameters.AddInputTensor(timestepTensor);
-                        inferenceParameters.AddInputTensor(promptEmbeddings);
+                        inferenceParameters.AddInputTensor(promptEmbeddings.PromptEmbeds);
                         inferenceParameters.AddOutputBuffer(outputDimension);
 
                         var results = await _onnxModelService.RunInferenceAsync(modelOptions, OnnxModelType.Unet, inferenceParameters);
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintLegacyDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintLegacyDiffuser.cs
@@ -50,7 +50,7 @@ public InpaintLegacyDiffuser(IOnnxModelService onnxModelService, IPromptService
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, DenseTensor<float> promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
+        protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, PromptEmbeddingsResult promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
         {
             using (var scheduler = GetScheduler(schedulerOptions))
             {
@@ -91,7 +91,7 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio
                     {
                         inferenceParameters.AddInputTensor(inputTensor);
                         inferenceParameters.AddInputTensor(timestepTensor);
-                        inferenceParameters.AddInputTensor(promptEmbeddings);
+                        inferenceParameters.AddInputTensor(promptEmbeddings.PromptEmbeds);
                         inferenceParameters.AddOutputBuffer(outputDimension);
 
                         var results = await _onnxModelService.RunInferenceAsync(modelOptions, OnnxModelType.Unet, inferenceParameters);
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/StableDiffusionDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/StableDiffusionDiffuser.cs
@@ -48,7 +48,7 @@ public StableDiffusionDiffuser(IOnnxModelService onnxModelService, IPromptServic
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, DenseTensor<float> promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
+        protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, PromptEmbeddingsResult promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
         {
             // Get Scheduler
             using (var scheduler = GetScheduler(schedulerOptions))
@@ -81,7 +81,7 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio
                     {
                         inferenceParameters.AddInputTensor(inputTensor);
                         inferenceParameters.AddInputTensor(timestepTensor);
-                        inferenceParameters.AddInputTensor(promptEmbeddings);
+                        inferenceParameters.AddInputTensor(promptEmbeddings.PromptEmbeds);
                         inferenceParameters.AddOutputBuffer(outputDimension);
 
                         var results = await _onnxModelService.RunInferenceAsync(modelOptions, OnnxModelType.Unet, inferenceParameters);
diff --git a/OnnxStack.StableDiffusion/Services/PromptService.cs b/OnnxStack.StableDiffusion/Services/PromptService.cs

Original file line number	Diff line number	Diff line change
`@@ -6,8 +6,6 @@ namespace OnnxStack.StableDiffusion.Common`
`6`	`6`	`{`
`7`	`7`	`public interface IPromptService`
`8`	`8`	`{`
`9`		`- Task<DenseTensor<float>> CreatePromptAsync(IModelOptions model, PromptOptions promptOptions, bool isGuidanceEnabled);`
`10`		`- Task<int[]> DecodeTextAsync(IModelOptions model, string inputText);`
`11`		`- Task<float[]> EncodeTokensAsync(IModelOptions model, int[] tokenizedInput);`
	`9`	`+ Task<PromptEmbeddingsResult> CreatePromptAsync(IModelOptions model, PromptOptions promptOptions, bool isGuidanceEnabled);`
`12`	`10`	`}`
`13`	`11`	`}`
Original file line number	Diff line number	Diff line change
`@@ -65,7 +65,7 @@ protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, ISc`
`65`	`65`	`/// <param name="progressCallback">The progress callback.</param>`
`66`	`66`	`/// <param name="cancellationToken">The cancellation token.</param>`
`67`	`67`	`/// <returns></returns>`
`68`		`- protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, DenseTensor<float> promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)`
	`68`	`+ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, PromptEmbeddingsResult promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)`
`69`	`69`	`{`
`70`	`70`	`using (var scheduler = GetScheduler(schedulerOptions))`
`71`	`71`	`{`
`@@ -111,7 +111,7 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio`
`111`	`111`	`{`
`112`	`112`	`inferenceParameters.AddInputTensor(inputTensor);`
`113`	`113`	`inferenceParameters.AddInputTensor(timestepTensor);`
`114`		`- inferenceParameters.AddInputTensor(promptEmbeddings);`
	`114`	`+ inferenceParameters.AddInputTensor(promptEmbeddings.PromptEmbeds);`
`115`	`115`	`inferenceParameters.AddInputTensor(guidanceEmbeddings);`
`116`	`116`	`inferenceParameters.AddOutputBuffer(outputDimension);`
`117`	`117`
Original file line number	Diff line number	Diff line change
`@@ -50,7 +50,7 @@ public InpaintLegacyDiffuser(IOnnxModelService onnxModelService, IPromptService`
`50`	`50`	`/// <param name="progressCallback">The progress callback.</param>`
`51`	`51`	`/// <param name="cancellationToken">The cancellation token.</param>`
`52`	`52`	`/// <returns></returns>`
`53`		`- protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, DenseTensor<float> promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)`
	`53`	`+ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, PromptEmbeddingsResult promptEmbeddings, bool performGuidance, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)`
`54`	`54`	`{`
`55`	`55`	`using (var scheduler = GetScheduler(schedulerOptions))`
`56`	`56`	`{`
`@@ -91,7 +91,7 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(IModelOptio`
`91`	`91`	`{`
`92`	`92`	`inferenceParameters.AddInputTensor(inputTensor);`
`93`	`93`	`inferenceParameters.AddInputTensor(timestepTensor);`
`94`		`- inferenceParameters.AddInputTensor(promptEmbeddings);`
	`94`	`+ inferenceParameters.AddInputTensor(promptEmbeddings.PromptEmbeds);`
`95`	`95`	`inferenceParameters.AddOutputBuffer(outputDimension);`
`96`	`96`
`97`	`97`	`var results = await _onnxModelService.RunInferenceAsync(modelOptions, OnnxModelType.Unet, inferenceParameters);`