Make ControlNet diffuser based not pipeline based

Author: saddam213

OnnxStack.Core/Config/OnnxModelType.cs

@@ -9,8 +9,7 @@ public enum OnnxModelType
         TextEncoder2 = 21,
         VaeEncoder = 30,
         VaeDecoder = 40,
-        Control = 50,
-        Annotation = 100,
+        ControlNet = 50,
         Upscaler = 1000,
     }
 }

OnnxStack.StableDiffusion/Diffusers/ControlNet/ControlNetDiffuser.cs

@@ -9,32 +9,38 @@
 using OnnxStack.StableDiffusion.Enums;
 using OnnxStack.StableDiffusion.Helpers;
 using OnnxStack.StableDiffusion.Models;
+using OnnxStack.StableDiffusion.Schedulers.StableDiffusion;
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 
-namespace OnnxStack.StableDiffusion.Diffusers.ControlNet
+namespace OnnxStack.StableDiffusion.Diffusers.StableDiffusion
 {
-    public sealed class TextDiffuser : ControlNetDiffuser
+    public class ControlNetDiffuser : DiffuserBase
     {
         /// <summary>
-        /// Initializes a new instance of the <see cref="TextDiffuser"/> class.
+        /// Initializes a new instance of the <see cref="ControlNetDiffuser"/> class.
         /// </summary>
         /// <param name="configuration">The configuration.</param>
         /// <param name="onnxModelService">The onnx model service.</param>
-        public TextDiffuser(IOnnxModelService onnxModelService, IPromptService promptService, ILogger<TextDiffuser> logger)
-            : base(onnxModelService, promptService, logger)
-        {
-        }
+        public ControlNetDiffuser(IOnnxModelService onnxModelService, IPromptService promptService, ILogger<ControlNetDiffuser> logger)
+            : base(onnxModelService, promptService, logger) { }
+
+
+        /// <summary>
+        /// Gets the type of the pipeline.
+        /// </summary>
+        public override DiffuserPipelineType PipelineType => DiffuserPipelineType.StableDiffusion;
 
 
         /// <summary>
         /// Gets the type of the diffuser.
         /// </summary>
-        public override DiffuserType DiffuserType => DiffuserType.ImageToImage;
+        public override DiffuserType DiffuserType => DiffuserType.ControlNet;
+
 
         /// <summary>
         /// Called on each Scheduler step.
@@ -47,7 +53,7 @@ public TextDiffuser(IOnnxModelService onnxModelService, IPromptService promptSer
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        /// <exception cref="System.NotImplementedException"></exception>
+        /// <exception cref="NotImplementedException"></exception>
         protected override async Task<DenseTensor<float>> SchedulerStepAsync(StableDiffusionModelSet modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, PromptEmbeddingsResult promptEmbeddings, bool performGuidance, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default)
         {
             // Get Scheduler
@@ -63,10 +69,10 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(StableDiffu
                 var metadata = _onnxModelService.GetModelMetadata(modelOptions, OnnxModelType.Unet);
 
                 // Get Model metadata
-                var controlNetMetadata = _onnxModelService.GetModelMetadata(modelOptions, OnnxModelType.Control);
- 
+                var controlNetMetadata = _onnxModelService.GetModelMetadata(modelOptions, OnnxModelType.ControlNet);
+
                 // Control Image
-                var controlImage = promptOptions.InputImage.ToDenseTensor(new[] { 1, 3, schedulerOptions.Height, schedulerOptions.Width }, false);
+                var controlImage = PrepareControlImage(promptOptions, schedulerOptions);
 
                 // Loop though the timesteps
                 var step = 0;
@@ -98,14 +104,15 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(StableDiffu
                             controlNetParameters.AddInputTensor(timestepTensor);
                             controlNetParameters.AddInputTensor(promptEmbeddings.PromptEmbeds);
                             controlNetParameters.AddInputTensor(controlImage);
-                            controlNetParameters.AddInputTensor(conditioningScale);
+                            if (controlNetMetadata.Inputs.Count == 5)
+                                controlNetParameters.AddInputTensor(conditioningScale);
 
                             // Optimization: Pre-allocate device buffers for inputs
                             foreach (var item in controlNetMetadata.Outputs)
                                 controlNetParameters.AddOutputBuffer();
 
                             // ControlNet inference
-                            var controlNetResults = _onnxModelService.RunInference(modelOptions, OnnxModelType.Control, controlNetParameters);
+                            var controlNetResults = _onnxModelService.RunInference(modelOptions, OnnxModelType.ControlNet, controlNetParameters);
 
                             // Add ControlNet outputs to Unet input
                             foreach (var item in controlNetResults)
@@ -139,14 +146,75 @@ protected override async Task<DenseTensor<float>> SchedulerStepAsync(StableDiffu
             }
         }
 
+
+        /// <summary>
+        /// Gets the timesteps.
+        /// </summary>
+        /// <param name="options">The options.</param>
+        /// <param name="scheduler">The scheduler.</param>
+        /// <returns></returns>
         protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
         {
             return scheduler.Timesteps;
         }
 
+
+        /// <summary>
+        /// Prepares the input latents.
+        /// </summary>
+        /// <param name="model">The model.</param>
+        /// <param name="prompt">The prompt.</param>
+        /// <param name="options">The options.</param>
+        /// <param name="scheduler">The scheduler.</param>
+        /// <param name="timesteps">The timesteps.</param>
+        /// <returns></returns>
         protected override Task<DenseTensor<float>> PrepareLatentsAsync(StableDiffusionModelSet model, PromptOptions prompt, SchedulerOptions options, IScheduler scheduler, IReadOnlyList<int> timesteps)
         {
             return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(), scheduler.InitNoiseSigma));
         }
+
+
+        /// <summary>
+        /// Creates the Conditioning Scale tensor.
+        /// </summary>
+        /// <param name="conditioningScale">The conditioningScale.</param>
+        /// <returns></returns>
+        protected static DenseTensor<double> CreateConditioningScaleTensor(float conditioningScale)
+        {
+            return TensorHelper.CreateTensor(new double[] { conditioningScale }, new int[] { 1 });
+        }
+
+
+        /// <summary>
+        /// Prepares the control image.
+        /// </summary>
+        /// <param name="promptOptions">The prompt options.</param>
+        /// <param name="schedulerOptions">The scheduler options.</param>
+        /// <returns></returns>
+        protected DenseTensor<float> PrepareControlImage(PromptOptions promptOptions, SchedulerOptions schedulerOptions)
+        {
+            return promptOptions.InputImage.ToDenseTensor(new[] { 1, 3, schedulerOptions.Height, schedulerOptions.Width }, false);
+        }
+
+
+        /// <summary>
+        /// Gets the scheduler.
+        /// </summary>
+        /// <param name="options">The options.</param>
+        /// <param name="schedulerConfig">The scheduler configuration.</param>
+        /// <returns></returns>
+        protected override IScheduler GetScheduler(SchedulerOptions options)
+        {
+            return options.SchedulerType switch
+            {
+                SchedulerType.LMS => new LMSScheduler(options),
+                SchedulerType.Euler => new EulerScheduler(options),
+                SchedulerType.EulerAncestral => new EulerAncestralScheduler(options),
+                SchedulerType.DDPM => new DDPMScheduler(options),
+                SchedulerType.DDIM => new DDIMScheduler(options),
+                SchedulerType.KDPM2 => new KDPM2Scheduler(options),
+                _ => default
+            };
+        }
     }
 }

OnnxStack.StableDiffusion/Diffusers/StableDiffusion/AnimateDiffuser.cs

@@ -0,0 +1,83 @@
+using Microsoft.Extensions.Logging;
+using Microsoft.ML.OnnxRuntime.Tensors;
+using OnnxStack.Core;
+using OnnxStack.Core.Config;
+using OnnxStack.Core.Model;
+using OnnxStack.Core.Services;
+using OnnxStack.StableDiffusion.Common;
+using OnnxStack.StableDiffusion.Config;
+using OnnxStack.StableDiffusion.Enums;
+using OnnxStack.StableDiffusion.Helpers;
+using SixLabors.ImageSharp;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+
+namespace OnnxStack.StableDiffusion.Diffusers.StableDiffusion
+{
+    public sealed class ControlNetImageDiffuser : ControlNetDiffuser
+    {
+        /// <summary>
+        /// Initializes a new instance of the <see cref="ControlNetImageDiffuser"/> class.
+        /// </summary>
+        /// <param name="configuration">The configuration.</param>
+        /// <param name="onnxModelService">The onnx model service.</param>
+        public ControlNetImageDiffuser(IOnnxModelService onnxModelService, IPromptService promptService, ILogger<ControlNetDiffuser> logger)
+            : base(onnxModelService, promptService, logger)
+        {
+        }
+
+
+        /// <summary>
+        /// Gets the type of the diffuser.
+        /// </summary>
+        public override DiffuserType DiffuserType => DiffuserType.ControlNetImage;
+
+
+        /// <summary>
+        /// Gets the timesteps.
+        /// </summary>
+        /// <param name="prompt">The prompt.</param>
+        /// <param name="options">The options.</param>
+        /// <param name="scheduler">The scheduler.</param>
+        /// <returns></returns>
+        protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, IScheduler scheduler)
+        {
+            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();
+        }
+
+
+        /// <summary>
+        /// Prepares the latents for inference.
+        /// </summary>
+        /// <param name="prompt">The prompt.</param>
+        /// <param name="options">The options.</param>
+        /// <param name="scheduler">The scheduler.</param>
+        /// <returns></returns>
+        protected override async Task<DenseTensor<float>> PrepareLatentsAsync(StableDiffusionModelSet model, PromptOptions prompt, SchedulerOptions options, IScheduler scheduler, IReadOnlyList<int> timesteps)
+        {
+            var imageTensor = prompt.InputImage.ToDenseTensor(new[] { 1, 3, options.Height, options.Width });
+
+            //TODO: Model Config, Channels
+            var outputDimension = options.GetScaledDimension();
+            var metadata = _onnxModelService.GetModelMetadata(model, OnnxModelType.VaeEncoder);
+            using (var inferenceParameters = new OnnxInferenceParameters(metadata))
+            {
+                inferenceParameters.AddInputTensor(imageTensor);
+                inferenceParameters.AddOutputBuffer(outputDimension);
+
+                var results = await _onnxModelService.RunInferenceAsync(model, OnnxModelType.VaeEncoder, inferenceParameters);
+                using (var result = results.First())
+                {
+                    var outputResult = result.ToDenseTensor();
+                    var scaledSample = outputResult.MultiplyBy(model.ScaleFactor);
+                    return scheduler.AddNoise(scaledSample, scheduler.CreateRandomSample(scaledSample.Dimensions), timesteps);
+                }
+            }
+        }
+
+    }
+}

OnnxStack.StableDiffusion/Diffusers/ControlNet/TextDiffuser.cs renamed to OnnxStack.StableDiffusion/Diffusers/StableDiffusion/ControlNetDiffuser.cs

@@ -6,7 +6,6 @@ public enum DiffuserPipelineType
         StableDiffusionXL = 1,
         LatentConsistency = 10,
         LatentConsistencyXL = 11,
-        ControlNet = 20,
         InstaFlow = 30,
     }
 }