Support for LCM and other pipelines

Author: saddam213

OnnxStack.Console/appsettings.json

@@ -17,6 +17,7 @@
 				"TokenizerLimit": 77,
 				"EmbeddingsLength": 768,
 				"ScaleFactor": 0.18215,
+				"PipelineType": "StableDiffusion",
 				"Diffusers": [
 					"TextToImage",
 					"ImageToImage",

OnnxStack.Core/Extensions.cs

@@ -4,7 +4,6 @@
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
-using System.Xml.Linq;
 
 namespace OnnxStack.Core
 {

OnnxStack.StableDiffusion/Common/IModelOptions.cs

@@ -14,6 +14,7 @@ public interface IModelOptions : IOnnxModel
         int TokenizerLimit { get; set; }
         int InputTokenLimit { get; set; }
         int EmbeddingsLength { get; set; }
+        DiffuserPipelineType PipelineType { get; set; }
         List<DiffuserType> Diffusers { get; set; }
         ImmutableArray<int> BlankTokenValueArray { get; set; }
     }

OnnxStack.StableDiffusion/Common/IScheduler.cs

@@ -1,11 +1,17 @@
 using Microsoft.ML.OnnxRuntime.Tensors;
+using OnnxStack.StableDiffusion.Enums;
 using System;
 using System.Collections.Generic;
 
 namespace OnnxStack.StableDiffusion.Common
 {
     public interface IScheduler : IDisposable
     {
+        /// <summary>
+        /// Gets the compatible pipeline
+        /// </summary>
+        DiffuserPipelineType PipelineType { get; }
+
         /// <summary>
         /// Gets the initial noise sigma.
         /// </summary>

OnnxStack.StableDiffusion/Config/ModelOptions.cs

@@ -19,6 +19,7 @@ public class ModelOptions : IModelOptions, IOnnxModelSetConfig
         public int TokenizerLimit { get; set; }
         public int EmbeddingsLength { get; set; }
         public float ScaleFactor { get; set; }
+        public DiffuserPipelineType PipelineType { get; set; }
         public List<DiffuserType> Diffusers { get; set; } = new List<DiffuserType>();
 
         public int DeviceId { get; set; }

OnnxStack.StableDiffusion/Enums/DiffuserPipelineType.cs

@@ -0,0 +1,8 @@
+namespace OnnxStack.StableDiffusion.Enums
+{
+    public enum DiffuserPipelineType
+    {
+        StableDiffusion = 0,
+        LatentConsistency = 10
+    }
+}

OnnxStack.StableDiffusion/Schedulers/SchedulerBase.cs

@@ -50,6 +50,11 @@ public SchedulerBase(SchedulerOptions schedulerOptions)
         /// </summary>
         public IReadOnlyList<int> Timesteps => _timesteps;
 
+        /// <summary>
+        /// Gets the compatible pipeline.
+        /// </summary>
+        public virtual DiffuserPipelineType PipelineType => DiffuserPipelineType.StableDiffusion;
+
         /// <summary>
         /// Scales the input.
         /// </summary>

OnnxStack.StableDiffusion/Services/StableDiffusionService.cs

@@ -25,7 +25,7 @@ public sealed class StableDiffusionService : IStableDiffusionService
     {
         private readonly IOnnxModelService _onnxModelService;
         private readonly StableDiffusionConfig _configuration;
-        private readonly IDictionary<DiffuserType, IDiffuser> _diffusers;
+        private readonly IDictionary<DiffuserPipelineType, IDictionary<DiffuserType, IDiffuser>> _diffusers;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="StableDiffusionService"/> class.
@@ -35,11 +35,21 @@ public StableDiffusionService(StableDiffusionConfig configuration, IOnnxModelSer
         {
             _configuration = configuration;
             _onnxModelService = onnxModelService;
-            _diffusers = new ConcurrentDictionary<DiffuserType, IDiffuser>();
-            _diffusers.Add(DiffuserType.TextToImage, new TextDiffuser(onnxModelService, promptService));
-            _diffusers.Add(DiffuserType.ImageToImage, new ImageDiffuser(onnxModelService, promptService));
-            _diffusers.Add(DiffuserType.ImageInpaint, new InpaintDiffuser(onnxModelService, promptService));
-            _diffusers.Add(DiffuserType.ImageInpaintLegacy, new InpaintLegacyDiffuser(onnxModelService, promptService));
+            var stableDiffusionPipeline = new Dictionary<DiffuserType, IDiffuser>
+            {
+                { DiffuserType.TextToImage, new TextDiffuser(onnxModelService, promptService) },
+                { DiffuserType.ImageToImage, new ImageDiffuser(onnxModelService, promptService) },
+                { DiffuserType.ImageInpaint, new InpaintDiffuser(onnxModelService, promptService) },
+                { DiffuserType.ImageInpaintLegacy, new InpaintLegacyDiffuser(onnxModelService, promptService) }
+            };
+
+            var latentConsistancyPipeline = new Dictionary<DiffuserType, IDiffuser>
+            {
+                //TODO: TextToImage and ImageToImage is supported with LCM
+            };
+
+            _diffusers.Add(DiffuserPipelineType.StableDiffusion, stableDiffusionPipeline);
+            _diffusers.Add(DiffuserPipelineType.LatentConsistency, latentConsistancyPipeline);
         }
 
 
@@ -146,8 +156,15 @@ public async Task<Stream> GenerateAsStreamAsync(IModelOptions model, PromptOptio
 
         private async Task<DenseTensor<float>> DiffuseAsync(IModelOptions modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, Action<int, int> progress = null, CancellationToken cancellationToken = default)
         {
-            return await _diffusers[promptOptions.DiffuserType]
-                .DiffuseAsync(modelOptions, promptOptions, schedulerOptions, progress, cancellationToken);
+            var pipeline = _diffusers[modelOptions.PipelineType];
+            if (pipeline is null)
+                throw new Exception("Pipeline not found or is unsupported");
+
+            var diffuser = pipeline[promptOptions.DiffuserType];
+            if (diffuser is null)
+                throw new Exception("Diffuser not found or is unsupported");
+
+            return await diffuser.DiffuseAsync(modelOptions, promptOptions, schedulerOptions, progress, cancellationToken);
         }
     }
 }

OnnxStack.UI/appsettings.json

@@ -17,6 +17,7 @@
         "TokenizerLimit": 77,
         "EmbeddingsLength": 768,
         "ScaleFactor": 0.18215,
+        "PipelineType": "StableDiffusion",
         "Diffusers": [
           "TextToImage",
           "ImageToImage",

OnnxStack.WebUI/appsettings.json

@@ -17,6 +17,7 @@
 				"TokenizerLimit": 77,
 				"EmbeddingsLength": 768,
 				"ScaleFactor": 0.18215,
+				"PipelineType": "StableDiffusion",
 				"Diffusers": [
 					"TextToImage",
 					"ImageToImage",