Support user defined Timesteps

saddam213 · saddam213 · commit 212264189633 · 2024-04-24T10:07:34.000+12:00
diff --git a/OnnxStack.StableDiffusion/Config/SchedulerOptions.cs b/OnnxStack.StableDiffusion/Config/SchedulerOptions.cs
@@ -78,6 +78,7 @@ public record SchedulerOptions
         public PredictionType PredictionType { get; set; } = PredictionType.Epsilon;
         public AlphaTransformType AlphaTransformType { get; set; } = AlphaTransformType.Cosine;
         public float MaximumBeta { get; set; } = 0.999f;
+        public List<int> Timesteps { get; set; }
 
         public int OriginalInferenceSteps { get; set; } = 50;
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/InstaFlow/ControlNetDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/InstaFlow/ControlNetDiffuser.cs
@@ -164,6 +164,9 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             return scheduler.Timesteps;
         }
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/InstaFlow/TextDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/InstaFlow/TextDiffuser.cs
@@ -1,5 +1,6 @@
 ﻿using Microsoft.Extensions.Logging;
 using Microsoft.ML.OnnxRuntime.Tensors;
+using OnnxStack.Core;
 using OnnxStack.Core.Model;
 using OnnxStack.StableDiffusion.Common;
 using OnnxStack.StableDiffusion.Config;
@@ -38,6 +39,9 @@ public TextDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEn
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             return scheduler.Timesteps;
         }
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/ControlNetDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/ControlNetDiffuser.cs
@@ -160,6 +160,9 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             return scheduler.Timesteps;
         }
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/ControlNetImageDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/ControlNetImageDiffuser.cs
@@ -45,6 +45,9 @@ public ControlNetImageDiffuser(ControlNetModel controlNet, UNetConditionModel un
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
             var start = Math.Max(options.InferenceSteps - inittimestep, 0);
             return scheduler.Timesteps.Skip(start).ToList();
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/ImageDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/ImageDiffuser.cs
@@ -44,6 +44,9 @@ public ImageDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoE
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
             var start = Math.Max(options.InferenceSteps - inittimestep, 0);
             return scheduler.Timesteps.Skip(start).ToList();
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/InpaintLegacyDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/InpaintLegacyDiffuser.cs
@@ -47,6 +47,9 @@ public InpaintLegacyDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecode
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
             var start = Math.Max(options.InferenceSteps - inittimestep, 0);
             return scheduler.Timesteps.Skip(start).ToList();
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/TextDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/TextDiffuser.cs
@@ -1,5 +1,6 @@
 ﻿using Microsoft.Extensions.Logging;
 using Microsoft.ML.OnnxRuntime.Tensors;
+using OnnxStack.Core;
 using OnnxStack.Core.Model;
 using OnnxStack.StableDiffusion.Common;
 using OnnxStack.StableDiffusion.Config;
@@ -38,6 +39,9 @@ public TextDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEn
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             return scheduler.Timesteps;
         }
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistencyXL/ControlNetDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistencyXL/ControlNetDiffuser.cs
@@ -166,6 +166,9 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             return scheduler.Timesteps;
         }
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistencyXL/ControlNetImageDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistencyXL/ControlNetImageDiffuser.cs
@@ -45,6 +45,9 @@ public ControlNetImageDiffuser(ControlNetModel controlNet, UNetConditionModel un
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
             var start = Math.Max(options.InferenceSteps - inittimestep, 0);
             return scheduler.Timesteps.Skip(start).ToList();
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistencyXL/ImageDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistencyXL/ImageDiffuser.cs
@@ -45,6 +45,9 @@ public ImageDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoE
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             // Image2Image we narrow step the range by the Strength
             var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
             var start = Math.Max(options.InferenceSteps - inittimestep, 0);
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistencyXL/InpaintLegacyDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistencyXL/InpaintLegacyDiffuser.cs
@@ -141,6 +141,9 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
             var start = Math.Max(options.InferenceSteps - inittimestep, 0);
             return scheduler.Timesteps.Skip(start).ToList();
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistencyXL/TextDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistencyXL/TextDiffuser.cs
@@ -1,5 +1,6 @@
 ﻿using Microsoft.Extensions.Logging;
 using Microsoft.ML.OnnxRuntime.Tensors;
+using OnnxStack.Core;
 using OnnxStack.Core.Model;
 using OnnxStack.StableDiffusion.Common;
 using OnnxStack.StableDiffusion.Config;
@@ -38,6 +39,9 @@ public TextDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEn
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             return scheduler.Timesteps;
         }
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/ControlNetDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/ControlNetDiffuser.cs
@@ -157,6 +157,9 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             return scheduler.Timesteps;
         }
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/ControlNetImageDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/ControlNetImageDiffuser.cs
@@ -45,6 +45,9 @@ public ControlNetImageDiffuser(ControlNetModel controlNet, UNetConditionModel un
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
             var start = Math.Max(options.InferenceSteps - inittimestep, 0);
             return scheduler.Timesteps.Skip(start).ToList();
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/ImageDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/ImageDiffuser.cs
@@ -44,6 +44,9 @@ public ImageDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoE
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             // Image2Image we narrow step the range by the Strength
             var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
             var start = Math.Max(options.InferenceSteps - inittimestep, 0);
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintDiffuser.cs
@@ -252,6 +252,9 @@ private async Task<DenseTensor<float>> PrepareImageMask(PromptOptions promptOpti
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             return scheduler.Timesteps;
         }
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintLegacyDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintLegacyDiffuser.cs
@@ -135,6 +135,9 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
             var start = Math.Max(options.InferenceSteps - inittimestep, 0);
             return scheduler.Timesteps.Skip(start).ToList();
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/TextDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/TextDiffuser.cs
@@ -1,5 +1,6 @@
 ﻿using Microsoft.Extensions.Logging;
 using Microsoft.ML.OnnxRuntime.Tensors;
+using OnnxStack.Core;
 using OnnxStack.Core.Model;
 using OnnxStack.StableDiffusion.Common;
 using OnnxStack.StableDiffusion.Config;
@@ -39,6 +40,9 @@ public TextDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEn
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             return scheduler.Timesteps;
         }
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/ControlNetDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/ControlNetDiffuser.cs
@@ -167,6 +167,9 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             return scheduler.Timesteps;
         }
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/ControlNetImageDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/ControlNetImageDiffuser.cs
@@ -46,6 +46,9 @@ public ControlNetImageDiffuser(ControlNetModel controlNet, UNetConditionModel un
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
             var start = Math.Max(options.InferenceSteps - inittimestep, 0);
             return scheduler.Timesteps.Skip(start).ToList();
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/ImageDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/ImageDiffuser.cs
@@ -45,6 +45,9 @@ public ImageDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoE
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             // Image2Image we narrow step the range by the Strength
             var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
             var start = Math.Max(options.InferenceSteps - inittimestep, 0);
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/InpaintLegacyDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/InpaintLegacyDiffuser.cs
@@ -143,6 +143,9 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
             var start = Math.Max(options.InferenceSteps - inittimestep, 0);
             return scheduler.Timesteps.Skip(start).ToList();
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/TextDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/TextDiffuser.cs
@@ -1,5 +1,6 @@
 ﻿using Microsoft.Extensions.Logging;
 using Microsoft.ML.OnnxRuntime.Tensors;
+using OnnxStack.Core;
 using OnnxStack.Core.Model;
 using OnnxStack.StableDiffusion.Common;
 using OnnxStack.StableDiffusion.Config;
@@ -39,6 +40,9 @@ public TextDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEn
         /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
+            if (!options.Timesteps.IsNullOrEmpty())
+                return options.Timesteps;
+
             return scheduler.Timesteps;
         }
 

Original file line number	Diff line number	Diff line change
`@@ -164,6 +164,9 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt`
`164`	`164`	`/// <returns></returns>`
`165`	`165`	`protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)`
`166`	`166`	`{`
	`167`	`+ if (!options.Timesteps.IsNullOrEmpty())`
	`168`	`+ return options.Timesteps;`
	`169`	`+`
`167`	`170`	`return scheduler.Timesteps;`
`168`	`171`	`}`
`169`	`172`
Original file line number	Diff line number	Diff line change
`@@ -160,6 +160,9 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt`
`160`	`160`	`/// <returns></returns>`
`161`	`161`	`protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)`
`162`	`162`	`{`
	`163`	`+ if (!options.Timesteps.IsNullOrEmpty())`
	`164`	`+ return options.Timesteps;`
	`165`	`+`
`163`	`166`	`return scheduler.Timesteps;`
`164`	`167`	`}`
`165`	`168`
Original file line number	Diff line number	Diff line change
`@@ -166,6 +166,9 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt`
`166`	`166`	`/// <returns></returns>`
`167`	`167`	`protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)`
`168`	`168`	`{`
	`169`	`+ if (!options.Timesteps.IsNullOrEmpty())`
	`170`	`+ return options.Timesteps;`
	`171`	`+`
`169`	`172`	`return scheduler.Timesteps;`
`170`	`173`	`}`
`171`	`174`