Merge pull request #88 from saddam213/ControlNet

Add ControlNet Support

Author: saddam213

OnnxStack.Console/Examples/StableDebug.cs

@@ -77,7 +77,7 @@ private async Task<bool> GenerateImage(StableDiffusionModelSet model, PromptOpti
         {
             var timestamp = Stopwatch.GetTimestamp();
             var outputFilename = Path.Combine(_outputDirectory, $"{model.Name}_{options.Seed}_{options.SchedulerType}.png");
-            var result = await _stableDiffusionService.GenerateAsImageAsync(model, prompt, options);
+            var result = await _stableDiffusionService.GenerateAsImageAsync(new ModelOptions(model), prompt, options);
             if (result is not null)
             {
                 await result.SaveAsPngAsync(outputFilename);

OnnxStack.Console/Examples/StableDiffusionBatch.cs

@@ -66,7 +66,7 @@ public async Task RunAsync()
                         OutputHelpers.WriteConsole($"Image: {progress.BatchValue}/{progress.BatchMax} - Step: {progress.StepValue}/{progress.StepMax}", ConsoleColor.Cyan);
                     };
 
-                    await foreach (var result in _stableDiffusionService.GenerateBatchAsync(model, promptOptions, schedulerOptions, batchOptions, default))
+                    await foreach (var result in _stableDiffusionService.GenerateBatchAsync(new ModelOptions(model), promptOptions, schedulerOptions, batchOptions, default))
                     {
                         var outputFilename = Path.Combine(_outputDirectory, $"{batchIndex}_{result.SchedulerOptions.Seed}.png");
                         var image = result.ImageResult.ToImage();

OnnxStack.Console/Examples/StableDiffusionExample.cs

@@ -70,7 +70,7 @@ public async Task RunAsync()
         private async Task<bool> GenerateImage(StableDiffusionModelSet model, PromptOptions prompt, SchedulerOptions options)
         {
             var outputFilename = Path.Combine(_outputDirectory, $"{options.Seed}_{options.SchedulerType}.png");
-            var result = await _stableDiffusionService.GenerateAsImageAsync(model, prompt, options);
+            var result = await _stableDiffusionService.GenerateAsImageAsync(new ModelOptions(model), prompt, options);
             if (result == null)
                 return false;
 

OnnxStack.Console/Examples/StableDiffusionGenerator.cs

@@ -67,7 +67,7 @@ public async Task RunAsync()
         private async Task<bool> GenerateImage(StableDiffusionModelSet model, PromptOptions prompt, SchedulerOptions options, string key)
         {
             var outputFilename = Path.Combine(_outputDirectory, $"{options.Seed}_{options.SchedulerType}_{key}.png");
-            var result = await _stableDiffusionService.GenerateAsImageAsync(model, prompt, options);
+            var result = await _stableDiffusionService.GenerateAsImageAsync(new ModelOptions(model), prompt, options);
             if (result == null)
                 return false;
 

OnnxStack.Console/Examples/StableDiffusionGif.cs

@@ -102,7 +102,7 @@ public async Task RunAsync()
 
                         // Set prompt Image, Run Diffusion
                         promptOptions.InputImage = new InputImage(mergedFrame.CloneAs<Rgba32>());
-                        var result = await _stableDiffusionService.GenerateAsImageAsync(model, promptOptions, schedulerOptions);
+                        var result = await _stableDiffusionService.GenerateAsImageAsync(new ModelOptions(model), promptOptions, schedulerOptions);
 
                         // Save Debug Output
                         await result.SaveAsPngAsync(Path.Combine(_outputDirectory, $"Debug-Output.png"));

OnnxStack.Core/Config/OnnxModelType.cs

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

OnnxStack.Core/Extensions/OrtValueExtensions.cs

@@ -59,6 +59,11 @@ public static OrtValue ToOrtValue(this DenseTensor<long> tensor, OnnxNamedMetada
             return OrtValue.CreateTensorValueFromMemory(OrtMemoryInfo.DefaultInstance, tensor.Buffer, tensor.Dimensions.ToLong());
         }
 
+        public static OrtValue ToOrtValue(this DenseTensor<double> tensor, OnnxNamedMetadata metadata)
+        {
+            return OrtValue.CreateTensorValueFromMemory(OrtMemoryInfo.DefaultInstance, tensor.Buffer, tensor.Dimensions.ToLong());
+        }
+
 
         /// <summary>
         /// Creates and allocates the output tensors buffer.

OnnxStack.Core/Image/Extensions.cs

@@ -63,6 +63,23 @@ public static async Task<byte[]> ToImageBytesAsync(this DenseTensor<float> image
             }
         }
 
+        public static Image<Rgba32> ToImageMask(this DenseTensor<float> imageTensor)
+        {
+            var width = imageTensor.Dimensions[3];
+            var height = imageTensor.Dimensions[2];
+            using (var result = new Image<L8>(width, height))
+            {
+                for (var y = 0; y < height; y++)
+                {
+                    for (var x = 0; x < width; x++)
+                    {
+                        result[x, y] = new L8((byte)(imageTensor[0, 0, y, x] * 255.0f));
+                    }
+                }
+                return result.CloneAs<Rgba32>();
+            }
+        }
+
 
         private static byte CalculateByte(Tensor<float> imageTensor, int index, int y, int x)
         {

OnnxStack.Core/Model/OnnxInferenceParameters.cs

@@ -45,6 +45,12 @@ public void AddInputTensor(DenseTensor<float> value)
             _inputs.Add(metaData, value.ToOrtValue(metaData));
         }
 
+        public void AddInputTensor(DenseTensor<double> value)
+        {
+            var metaData = GetNextInputMetadata();
+            _inputs.Add(metaData, value.ToOrtValue(metaData));
+        }
+
 
         /// <summary>
         /// Adds the input tensor.

OnnxStack.StableDiffusion/Common/IControlNetImageService.cs

@@ -0,0 +1,20 @@
+using OnnxStack.Core.Image;
+using OnnxStack.StableDiffusion.Config;
+using System.Threading.Tasks;
+
+namespace OnnxStack.StableDiffusion.Common
+{
+    public interface IControlNetImageService
+    {
+
+        /// <summary>
+        /// Prepares the ContolNet input image, If the ControlNetModelSet has a configure Annotation model this will be used to process the image
+        /// </summary>
+        /// <param name="controlNetModel">The control net model.</param>
+        /// <param name="inputImage">The input image.</param>
+        /// <param name="height">The height.</param>
+        /// <param name="width">The width.</param>
+        /// <returns></returns>
+        Task<InputImage> PrepareInputImage(ControlNetModelSet controlNetModel, InputImage inputImage, int height, int width);
+    }
+}

OnnxStack.StableDiffusion/Common/IStableDiffusionService.cs

@@ -1,4 +1,5 @@
 using Microsoft.ML.OnnxRuntime.Tensors;
+using OnnxStack.Core.Config;
 using OnnxStack.StableDiffusion.Config;
 using OnnxStack.StableDiffusion.Models;
 using SixLabors.ImageSharp;
@@ -18,15 +19,15 @@ public interface IStableDiffusionService
         /// </summary>
         /// <param name="modelOptions">The model options.</param>
         /// <returns></returns>
-        Task<bool> LoadModelAsync(StableDiffusionModelSet model);
+        Task<bool> LoadModelAsync(IOnnxModelSetConfig model);
 
 
         /// <summary>
         /// Unloads the model.
         /// </summary>
         /// <param name="modelOptions">The model options.</param>
         /// <returns></returns>
-        Task<bool> UnloadModelAsync(StableDiffusionModelSet model);
+        Task<bool> UnloadModelAsync(IOnnxModel model);
 
         /// <summary>
         /// Determines whether the specified model is loaded
@@ -35,7 +36,7 @@ public interface IStableDiffusionService
         /// <returns>
         ///   <c>true</c> if the specified model is loaded; otherwise, <c>false</c>.
         /// </returns>
-        bool IsModelLoaded(StableDiffusionModelSet model);
+        bool IsModelLoaded(IOnnxModel model);
 
         /// <summary>
         /// Generates the StableDiffusion image using the prompt and options provided.
@@ -45,7 +46,7 @@ public interface IStableDiffusionService
         /// <param name="progressCallback">The callback used to provide progess of the current InferenceSteps.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns>The diffusion result as <see cref="DenseTensor<float>"/></returns>
-        Task<DenseTensor<float>> GenerateAsync(StableDiffusionModelSet model, PromptOptions prompt, SchedulerOptions options, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+        Task<DenseTensor<float>> GenerateAsync(ModelOptions model, PromptOptions prompt, SchedulerOptions options, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
 
         /// <summary>
         /// Generates the StableDiffusion image using the prompt and options provided.
@@ -55,7 +56,7 @@ public interface IStableDiffusionService
         /// <param name="progressCallback">The callback used to provide progess of the current InferenceSteps.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns>The diffusion result as <see cref="SixLabors.ImageSharp.Image<Rgba32>"/></returns>
-        Task<Image<Rgba32>> GenerateAsImageAsync(StableDiffusionModelSet model, PromptOptions prompt, SchedulerOptions options, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+        Task<Image<Rgba32>> GenerateAsImageAsync(ModelOptions model, PromptOptions prompt, SchedulerOptions options, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
 
         /// <summary>
         /// Generates the StableDiffusion image using the prompt and options provided.
@@ -65,7 +66,7 @@ public interface IStableDiffusionService
         /// <param name="progressCallback">The callback used to provide progess of the current InferenceSteps.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns>The diffusion result as <see cref="byte[]"/></returns>
-        Task<byte[]> GenerateAsBytesAsync(StableDiffusionModelSet model, PromptOptions prompt, SchedulerOptions options, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+        Task<byte[]> GenerateAsBytesAsync(ModelOptions model, PromptOptions prompt, SchedulerOptions options, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
 
         /// <summary>
         /// Generates the StableDiffusion image using the prompt and options provided.
@@ -75,7 +76,7 @@ public interface IStableDiffusionService
         /// <param name="progressCallback">The callback used to provide progess of the current InferenceSteps.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns>The diffusion result as <see cref="System.IO.Stream"/></returns>
-        Task<Stream> GenerateAsStreamAsync(StableDiffusionModelSet model, PromptOptions prompt, SchedulerOptions options, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+        Task<Stream> GenerateAsStreamAsync(ModelOptions model, PromptOptions prompt, SchedulerOptions options, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
 
         /// <summary>
         /// Generates a batch of StableDiffusion image using the prompt and options provided.
@@ -87,7 +88,7 @@ public interface IStableDiffusionService
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        IAsyncEnumerable<BatchResult> GenerateBatchAsync(StableDiffusionModelSet model, PromptOptions promptOptions, SchedulerOptions schedulerOptions, BatchOptions batchOptions, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+        IAsyncEnumerable<BatchResult> GenerateBatchAsync(ModelOptions model, PromptOptions promptOptions, SchedulerOptions schedulerOptions, BatchOptions batchOptions, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
 
         /// <summary>
         /// Generates a batch of StableDiffusion image using the prompt and options provided.
@@ -99,7 +100,7 @@ public interface IStableDiffusionService
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        IAsyncEnumerable<Image<Rgba32>> GenerateBatchAsImageAsync(StableDiffusionModelSet model, PromptOptions promptOptions, SchedulerOptions schedulerOptions, BatchOptions batchOptions, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+        IAsyncEnumerable<Image<Rgba32>> GenerateBatchAsImageAsync(ModelOptions model, PromptOptions promptOptions, SchedulerOptions schedulerOptions, BatchOptions batchOptions, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
 
         /// <summary>
         /// Generates a batch of StableDiffusion image using the prompt and options provided.
@@ -111,7 +112,7 @@ public interface IStableDiffusionService
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        IAsyncEnumerable<byte[]> GenerateBatchAsBytesAsync(StableDiffusionModelSet model, PromptOptions promptOptions, SchedulerOptions schedulerOptions, BatchOptions batchOptions, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+        IAsyncEnumerable<byte[]> GenerateBatchAsBytesAsync(ModelOptions model, PromptOptions promptOptions, SchedulerOptions schedulerOptions, BatchOptions batchOptions, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
 
         /// <summary>
         /// Generates a batch of StableDiffusion image using the prompt and options provided.
@@ -123,6 +124,6 @@ public interface IStableDiffusionService
         /// <param name="progressCallback">The progress callback.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        IAsyncEnumerable<Stream> GenerateBatchAsStreamAsync(StableDiffusionModelSet model, PromptOptions promptOptions, SchedulerOptions schedulerOptions, BatchOptions batchOptions, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+        IAsyncEnumerable<Stream> GenerateBatchAsStreamAsync(ModelOptions model, PromptOptions promptOptions, SchedulerOptions schedulerOptions, BatchOptions batchOptions, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
     }
 }

OnnxStack.StableDiffusion/Config/ControlNetModelSet.cs

@@ -0,0 +1,20 @@
+using Microsoft.ML.OnnxRuntime;
+using OnnxStack.Core.Config;
+using OnnxStack.StableDiffusion.Enums;
+using System.Collections.Generic;
+
+namespace OnnxStack.StableDiffusion.Config
+{
+    public record ControlNetModelSet : IOnnxModelSetConfig
+    {
+        public ControlNetType Type { get; set; }
+        public string Name { get; set; }
+        public bool IsEnabled { get; set; }
+        public int DeviceId { get; set; }
+        public int InterOpNumThreads { get; set; }
+        public int IntraOpNumThreads { get; set; }
+        public ExecutionMode ExecutionMode { get; set; }
+        public ExecutionProvider ExecutionProvider { get; set; }
+        public List<OnnxModelConfig> ModelConfigurations { get; set; }
+    }
+}

OnnxStack.StableDiffusion/Config/ModelOptions.cs

@@ -0,0 +1,12 @@
+using OnnxStack.StableDiffusion.Enums;
+
+namespace OnnxStack.StableDiffusion.Config
+{
+    public record ModelOptions(StableDiffusionModelSet BaseModel, ControlNetModelSet ControlNetModel = default)
+    {
+        public string Name => BaseModel.Name;
+        public DiffuserPipelineType PipelineType => BaseModel.PipelineType;
+        public float ScaleFactor => BaseModel.ScaleFactor;
+        public ModelType ModelType => BaseModel.ModelType;
+    }
+}

OnnxStack.StableDiffusion/Config/PromptOptions.cs

@@ -20,6 +20,8 @@ public class PromptOptions
 
         public InputImage InputImageMask { get; set; }
 
+        public InputImage InputContolImage { get; set; }
+
         public VideoInput InputVideo { get; set; }
 
         public float VideoInputFPS { get; set; }

OnnxStack.StableDiffusion/Config/SchedulerOptions.cs

@@ -84,6 +84,9 @@ public record SchedulerOptions
         public float AestheticScore { get; set; } = 6f;
         public float AestheticNegativeScore { get; set; } = 2.5f;
 
+        public float ConditioningScale { get; set; } = 0.7f;
+        public bool IsControlImageProcessingEnabled { get; set; }
+
         public bool IsKarrasScheduler
         {
             get