make the pipeline a bit more abstract

Author: saddam213

OnnxStack.StableDiffusion/Common/IPipeline.cs

@@ -0,0 +1,29 @@
+using OnnxStack.StableDiffusion.Diffusers;
+using OnnxStack.StableDiffusion.Enums;
+using System.Collections.Concurrent;
+
+namespace OnnxStack.StableDiffusion.Common
+{
+    public interface IPipeline
+    {
+
+        /// <summary>
+        /// Gets the type of the pipeline.
+        /// </summary>
+        DiffuserPipelineType PipelineType { get; }
+
+
+        /// <summary>
+        /// The pipelines diffuser set.
+        /// </summary>
+        ConcurrentDictionary<DiffuserType, IDiffuser> Diffusers { get; }
+
+
+        /// <summary>
+        /// Gets the diffuser.
+        /// </summary>
+        /// <param name="diffuserType">Type of the diffuser.</param>
+        /// <returns></returns>
+        IDiffuser GetDiffuser(DiffuserType diffuserType);
+    }
+}

OnnxStack.StableDiffusion/Enums/SchedulerType.cs

@@ -20,6 +20,9 @@ public enum SchedulerType
         DDIM = 4,
 
         [Display(Name = "KDPM2")]
-        KDPM2 = 5
+        KDPM2 = 5,
+
+         [Display(Name = "LCM")]
+        LCM = 20
     }
 }

OnnxStack.StableDiffusion/Pipelines/LatentConsistency.cs

@@ -0,0 +1,49 @@
+using OnnxStack.Core.Services;
+using OnnxStack.StableDiffusion.Common;
+using OnnxStack.StableDiffusion.Diffusers;
+using OnnxStack.StableDiffusion.Enums;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+
+namespace OnnxStack.StableDiffusion.Pipelines
+{
+    public sealed class LatentConsistencyPipeline : IPipeline
+    {
+        private readonly DiffuserPipelineType _pipelineType;
+        private readonly ConcurrentDictionary<DiffuserType, IDiffuser> _diffusers;
+
+        public LatentConsistencyPipeline(IOnnxModelService onnxModelService, IPromptService promptService)
+        {
+            var diffusers = new Dictionary<DiffuserType, IDiffuser>
+            {
+                //TODO: TextToImage and ImageToImage is supported with LCM
+            };
+            _pipelineType = DiffuserPipelineType.LatentConsistency;
+            _diffusers = new ConcurrentDictionary<DiffuserType, IDiffuser>(diffusers);
+        }
+
+
+        /// <summary>
+        /// Gets the type of the pipeline.
+        /// </summary>
+        public DiffuserPipelineType PipelineType => _pipelineType;
+
+
+        /// <summary>
+        /// Gets the diffusers.
+        /// </summary>
+        public ConcurrentDictionary<DiffuserType, IDiffuser> Diffusers => _diffusers;
+
+
+        /// <summary>
+        /// Gets the diffuser.
+        /// </summary>
+        /// <param name="diffuserType">Type of the diffuser.</param>
+        /// <returns></returns>
+        public IDiffuser GetDiffuser(DiffuserType diffuserType)
+        {
+            _diffusers.TryGetValue(diffuserType, out var diffuser);
+            return diffuser;
+        }
+    }
+}

OnnxStack.StableDiffusion/Pipelines/StableDiffusion.cs

@@ -0,0 +1,53 @@
+using OnnxStack.Core.Services;
+using OnnxStack.StableDiffusion.Common;
+using OnnxStack.StableDiffusion.Diffusers;
+using OnnxStack.StableDiffusion.Enums;
+using OnnxStack.StableDiffusion.Services;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+
+namespace OnnxStack.StableDiffusion.Pipelines
+{
+    public sealed class StableDiffusionPipeline : IPipeline
+    {
+        private readonly DiffuserPipelineType _pipelineType;
+        private readonly ConcurrentDictionary<DiffuserType, IDiffuser> _diffusers;
+
+        public StableDiffusionPipeline(IOnnxModelService onnxModelService, IPromptService promptService)
+        {
+            var diffusers = 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) }
+            };
+            _pipelineType = DiffuserPipelineType.StableDiffusion;
+            _diffusers = new ConcurrentDictionary<DiffuserType, IDiffuser>(diffusers);
+        }
+
+
+        /// <summary>
+        /// Gets the type of the pipeline.
+        /// </summary>
+        public DiffuserPipelineType PipelineType => _pipelineType;
+
+
+        /// <summary>
+        /// Gets the diffusers.
+        /// </summary>
+        public ConcurrentDictionary<DiffuserType, IDiffuser> Diffusers => _diffusers;
+
+
+        /// <summary>
+        /// Gets the diffuser.
+        /// </summary>
+        /// <param name="diffuserType">Type of the diffuser.</param>
+        /// <returns></returns>
+        public IDiffuser GetDiffuser(DiffuserType diffuserType)
+        {
+            _diffusers.TryGetValue(diffuserType, out var diffuser);
+            return diffuser;
+        }
+    }
+}

OnnxStack.StableDiffusion/Services/StableDiffusionService.cs

@@ -2,15 +2,13 @@
 using OnnxStack.Core.Services;
 using OnnxStack.StableDiffusion.Common;
 using OnnxStack.StableDiffusion.Config;
-using OnnxStack.StableDiffusion.Diffusers;
 using OnnxStack.StableDiffusion.Enums;
 using OnnxStack.StableDiffusion.Helpers;
+using OnnxStack.StableDiffusion.Pipelines;
 using SixLabors.ImageSharp;
 using SixLabors.ImageSharp.PixelFormats;
 using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.Collections.Immutable;
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
@@ -25,7 +23,7 @@ public sealed class StableDiffusionService : IStableDiffusionService
     {
         private readonly IOnnxModelService _onnxModelService;
         private readonly StableDiffusionConfig _configuration;
-        private readonly IDictionary<DiffuserPipelineType, IDictionary<DiffuserType, IDiffuser>> _diffusers;
+        private readonly IDictionary<DiffuserPipelineType, IPipeline> _pipelines;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="StableDiffusionService"/> class.
@@ -35,21 +33,8 @@ public StableDiffusionService(StableDiffusionConfig configuration, IOnnxModelSer
         {
             _configuration = configuration;
             _onnxModelService = onnxModelService;
-            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);
+            _pipelines.Add(DiffuserPipelineType.StableDiffusion, new StableDiffusionPipeline(onnxModelService, promptService));
+            _pipelines.Add(DiffuserPipelineType.LatentConsistency, new LatentConsistencyPipeline(onnxModelService, promptService));
         }
 
 
@@ -156,11 +141,11 @@ 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)
         {
-            var pipeline = _diffusers[modelOptions.PipelineType];
+            var pipeline = _pipelines[modelOptions.PipelineType];
             if (pipeline is null)
                 throw new Exception("Pipeline not found or is unsupported");
 
-            var diffuser = pipeline[promptOptions.DiffuserType];
+            var diffuser = pipeline.GetDiffuser(promptOptions.DiffuserType);
             if (diffuser is null)
                 throw new Exception("Diffuser not found or is unsupported");