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);
+    }
+}