Skip to content
This repository was archived by the owner on Nov 27, 2024. It is now read-only.

Commit fed0a58

Browse files
saddam213saddam213
committed
ImageService for processing of Canny,Hed and Depth image inputs
1 parent 6d266f3 commit fed0a58

File tree

8 files changed

+265
-2
lines changed

8 files changed

+265
-2
lines changed

‎OnnxStack.Core/Config/OnnxModelType.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ public enum OnnxModelType
1010
VaeEncoder = 30,
1111
VaeDecoder = 40,
1212
ControlNet = 50,
13-
Upscaler = 1000,
13+
Annotation = 51,
14+
Upscaler = 1000
1415
}
1516
}

‎OnnxStack.Core/Image/Extensions.cs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,23 @@ public static async Task<byte[]> ToImageBytesAsync(this DenseTensor<float> image
6363
}
6464
}
6565

66+
public static Image<Rgba32> ToImageMask(this DenseTensor<float> imageTensor)
67+
{
68+
var width = imageTensor.Dimensions[3];
69+
var height = imageTensor.Dimensions[2];
70+
using (var result = new Image<L8>(width, height))
71+
{
72+
for (var y = 0; y < height; y++)
73+
{
74+
for (var x = 0; x < width; x++)
75+
{
76+
result[x, y] = new L8((byte)(imageTensor[0, 0, y, x] * 255.0f));
77+
}
78+
}
79+
return result.CloneAs<Rgba32>();
80+
}
81+
}
82+
6683

6784
private static byte CalculateByte(Tensor<float> imageTensor, int index, int y, int x)
6885
{

‎OnnxStack.StableDiffusion/Common/IImageService.cs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
using OnnxStack.Core.Image;
2+
using OnnxStack.StableDiffusion.Config;
3+
using System.Threading.Tasks;
4+
5+
namespace OnnxStack.StableDiffusion.Common
6+
{
7+
public interface IImageService
8+
{
9+
10+
/// <summary>
11+
/// Prepares the ContolNet input image, If the ControlNetModelSet has a configure Annotation model this will be used to process the image
12+
/// </summary>
13+
/// <param name="controlNetModel">The control net model.</param>
14+
/// <param name="inputImage">The input image.</param>
15+
/// <param name="height">The height.</param>
16+
/// <param name="width">The width.</param>
17+
/// <returns></returns>
18+
Task<InputImage> PrepareInputImage(ControlNetModelSet controlNetModel, InputImage inputImage, int height, int width);
19+
}
20+
}

‎OnnxStack.StableDiffusion/Config/ControlNetModelSet.cs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
using Microsoft.ML.OnnxRuntime;
22
using OnnxStack.Core.Config;
3+
using OnnxStack.StableDiffusion.Enums;
34
using System.Collections.Generic;
45

56
namespace OnnxStack.StableDiffusion.Config
67
{
78
public record ControlNetModelSet : IOnnxModelSetConfig
89
{
10+
public ControlNetType Type { get; set; }
911
public string Name { get; set; }
1012
public bool IsEnabled { get; set; }
1113
public int DeviceId { get; set; }

‎OnnxStack.StableDiffusion/Enums/ControlNetType.cs‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
namespace OnnxStack.StableDiffusion.Enums
2+
{
3+
public enum ControlNetType
4+
{
5+
Canny = 0,
6+
Depth = 1,
7+
HED = 2,
8+
Inpaint = 3,
9+
MLSD = 4,
10+
Normal = 5,
11+
OpenPose = 6,
12+
Scribble = 7,
13+
Seg = 8
14+
}
15+
}

‎OnnxStack.StableDiffusion/Registration.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ private static void RegisterServices(this IServiceCollection serviceCollection)
4545
ConfigureLibraries();
4646

4747
// Services
48+
serviceCollection.AddSingleton<IImageService, ImageService>();
4849
serviceCollection.AddSingleton<IVideoService, VideoService>();
4950
serviceCollection.AddSingleton<IPromptService, PromptService>();
5051
serviceCollection.AddSingleton<IStableDiffusionService, StableDiffusionService>();

‎OnnxStack.StableDiffusion/Services/ImageService.cs‎

Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
using Microsoft.Extensions.Logging;
2+
using Microsoft.ML.OnnxRuntime.Tensors;
3+
using OnnxStack.Core;
4+
using OnnxStack.Core.Config;
5+
using OnnxStack.Core.Image;
6+
using OnnxStack.Core.Model;
7+
using OnnxStack.Core.Services;
8+
using OnnxStack.StableDiffusion.Common;
9+
using OnnxStack.StableDiffusion.Config;
10+
using OnnxStack.StableDiffusion.Enums;
11+
using OnnxStack.StableDiffusion.Helpers;
12+
using System.Linq;
13+
using System.Threading.Tasks;
14+
15+
namespace OnnxStack.StableDiffusion.Services
16+
{
17+
/// <summary>
18+
/// Service for handing images for input and output of the diffusion process
19+
/// </summary>
20+
/// <seealso cref="OnnxStack.StableDiffusion.Common.IImageService" />
21+
public class ImageService : IImageService
22+
{
23+
private readonly ILogger<ImageService> _logger;
24+
private readonly IOnnxModelService _onnxModelService;
25+
26+
27+
/// <summary>
28+
/// Initializes a new instance of the <see cref="ImageService"/> class.
29+
/// </summary>
30+
/// <param name="onnxModelService">The onnx model service.</param>
31+
public ImageService(IOnnxModelService onnxModelService, ILogger<ImageService> logger)
32+
{
33+
_logger = logger;
34+
_onnxModelService = onnxModelService;
35+
}
36+
37+
38+
/// <summary>
39+
/// Prepares the ControlNet input image.
40+
/// </summary>
41+
/// <param name="controlNetModel">The control net model.</param>
42+
/// <param name="inputImage">The input image.</param>
43+
/// <param name="height">The height.</param>
44+
/// <param name="width">The width.</param>
45+
/// <returns></returns>
46+
public async Task<InputImage> PrepareInputImage(ControlNetModelSet controlNetModel, InputImage inputImage, int height, int width)
47+
{
48+
var annotationModel = controlNetModel.ModelConfigurations.FirstOrDefault(x => x.Type == OnnxModelType.Annotation);
49+
if (annotationModel is not null)
50+
{
51+
_logger.LogInformation($"[PrepareInputImage] - ControlNet {controlNetModel.Type} annotation model found.");
52+
return controlNetModel.Type switch
53+
{
54+
ControlNetType.Canny => await GenerateCannyImage(controlNetModel, inputImage, height, width),
55+
ControlNetType.HED => await GenerateHardEdgeImage(controlNetModel, inputImage, height, width),
56+
ControlNetType.Depth => await GenerateDepthImage(controlNetModel, inputImage, height, width),
57+
_ => PrepareInputImage(inputImage, height, width)
58+
};
59+
}
60+
61+
return PrepareInputImage(inputImage, height, width);
62+
}
63+
64+
65+
/// <summary>
66+
/// Generates the canny image mask.
67+
/// </summary>
68+
/// <param name="controlNetModel">The control net model.</param>
69+
/// <param name="inputImage">The input image.</param>
70+
/// <param name="height">The height.</param>
71+
/// <param name="width">The width.</param>
72+
/// <returns></returns>
73+
private async Task<InputImage> GenerateCannyImage(ControlNetModelSet controlNetModel, InputImage inputImage, int height, int width)
74+
{
75+
_logger.LogInformation($"[GenerateCannyImage] - Generating Canny image...");
76+
var controlImage = inputImage.ToDenseTensor(new[] { 1, 3, height, width }, false);
77+
var metadata = _onnxModelService.GetModelMetadata(controlNetModel, OnnxModelType.Annotation);
78+
using (var inferenceParameters = new OnnxInferenceParameters(metadata))
79+
{
80+
inferenceParameters.AddInputTensor(controlImage);
81+
inferenceParameters.AddOutputBuffer(new[] { 1, 1, height, width });
82+
83+
var results = await _onnxModelService.RunInferenceAsync(controlNetModel, OnnxModelType.Annotation, inferenceParameters);
84+
using (var result = results.First())
85+
{
86+
var testImage = result.ToDenseTensor().Repeat(3);
87+
var imageTensor = new DenseTensor<float>(controlImage.Dimensions);
88+
for (int i = 0; i < testImage.Length; i++)
89+
imageTensor.SetValue(i, testImage.GetValue(i));
90+
91+
var maskImage = imageTensor.ToImageMask();
92+
//await maskImage.SaveAsPngAsync("D:\\Canny.png");
93+
_logger.LogInformation($"[GenerateCannyImage] - Canny image generation complete.");
94+
return new InputImage(maskImage);
95+
}
96+
}
97+
}
98+
99+
100+
/// <summary>
101+
/// Generates the hard edge image mask.
102+
/// </summary>
103+
/// <param name="controlNetModel">The control net model.</param>
104+
/// <param name="inputImage">The input image.</param>
105+
/// <param name="height">The height.</param>
106+
/// <param name="width">The width.</param>
107+
/// <returns></returns>
108+
private async Task<InputImage> GenerateHardEdgeImage(ControlNetModelSet controlNetModel, InputImage inputImage, int height, int width)
109+
{
110+
_logger.LogInformation($"[GenerateHardEdgeImage] - Generating HardEdge image...");
111+
var controlImage = inputImage.ToDenseTensor(new[] { 1, 3, height, width }, false);
112+
var metadata = _onnxModelService.GetModelMetadata(controlNetModel, OnnxModelType.Annotation);
113+
using (var inferenceParameters = new OnnxInferenceParameters(metadata))
114+
{
115+
inferenceParameters.AddInputTensor(controlImage);
116+
inferenceParameters.AddOutputBuffer(new[] { 1, 1, height, width });
117+
118+
var results = await _onnxModelService.RunInferenceAsync(controlNetModel, OnnxModelType.Annotation, inferenceParameters);
119+
using (var result = results.First())
120+
{
121+
var testImage = result.ToDenseTensor().Repeat(3);
122+
var imageTensor = new DenseTensor<float>(controlImage.Dimensions);
123+
for (int i = 0; i < testImage.Length; i++)
124+
imageTensor.SetValue(i, testImage.GetValue(i));
125+
126+
var maskImage = imageTensor.ToImageMask();
127+
//await maskImage.SaveAsPngAsync("D:\\Hed.png");
128+
_logger.LogInformation($"[GenerateHardEdgeImage] - HardEdge image generation complete.");
129+
return new InputImage(maskImage);
130+
}
131+
}
132+
}
133+
134+
135+
/// <summary>
136+
/// Generates the depth image mask.
137+
/// </summary>
138+
/// <param name="controlNetModel">The control net model.</param>
139+
/// <param name="inputImage">The input image.</param>
140+
/// <param name="height">The height.</param>
141+
/// <param name="width">The width.</param>
142+
/// <returns></returns>
143+
private async Task<InputImage> GenerateDepthImage(ControlNetModelSet controlNetModel, InputImage inputImage, int height, int width)
144+
{
145+
_logger.LogInformation($"[GenerateDepthImage] - Generating Depth image...");
146+
var controlImage = inputImage.ToDenseTensor(new[] { 1, 3, height, width }, false);
147+
var metadata = _onnxModelService.GetModelMetadata(controlNetModel, OnnxModelType.Annotation);
148+
using (var inferenceParameters = new OnnxInferenceParameters(metadata))
149+
{
150+
inferenceParameters.AddInputTensor(controlImage);
151+
inferenceParameters.AddOutputBuffer(new[] { 1, 1, height, width });
152+
153+
var results = await _onnxModelService.RunInferenceAsync(controlNetModel, OnnxModelType.Annotation, inferenceParameters);
154+
using (var result = results.First())
155+
{
156+
var imageResult = result.ToDenseTensor();
157+
var imageTensor = new DenseTensor<float>(controlImage.Dimensions);
158+
for (int i = 0; i < imageResult.Length; i++)
159+
{
160+
imageTensor.SetValue(i, imageResult.GetValue(i));
161+
}
162+
163+
NormalizeDepthTensor(imageTensor);
164+
var maskImage = imageTensor.ToImageMask();
165+
//await maskImage.SaveAsPngAsync("D:\\Depth.png");
166+
_logger.LogInformation($"[GenerateDepthImage] - Depth image generation complete.");
167+
return new InputImage(maskImage);
168+
}
169+
}
170+
}
171+
172+
173+
/// <summary>
174+
/// Prepares the input image.
175+
/// </summary>
176+
/// <param name="inputImage">The input image.</param>
177+
/// <param name="height">The height.</param>
178+
/// <param name="width">The width.</param>
179+
/// <returns></returns>
180+
private static InputImage PrepareInputImage(InputImage inputImage, int height, int width)
181+
{
182+
return new InputImage(inputImage.ToDenseTensor(new[] { 1, 3, height, width }, false));
183+
}
184+
185+
186+
/// <summary>
187+
/// Normalizes the depth tensor.
188+
/// </summary>
189+
/// <param name="value">The value.</param>
190+
public static void NormalizeDepthTensor(DenseTensor<float> value)
191+
{
192+
var values = value.Buffer.Span;
193+
float min = float.PositiveInfinity, max = float.NegativeInfinity;
194+
foreach (var val in values)
195+
{
196+
if (min > val) min = val;
197+
if (max < val) max = val;
198+
}
199+
200+
var range = max - min;
201+
for (var i = 0; i < values.Length; i++)
202+
{
203+
values[i] = (values[i] - min) / range;
204+
}
205+
}
206+
}
207+
}

‎OnnxStack.UI/UserControls/SchedulerControl.xaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@
134134
<Label>Conditioning Scale</Label>
135135
<TextBlock Text="{Binding ElementName=SliderInitialConditioningScale, Path=Value, StringFormat={}{0:F2}}" VerticalAlignment="Bottom" HorizontalAlignment="Right" FontSize="10" Margin="0,0,6,0" FontWeight="Medium" />
136136
</DockPanel>
137-
<Slider Name="SliderInitialConditioningScale" Value="{Binding SchedulerOptions.ConditioningScale}" Minimum="0" Maximum="2" TickFrequency="0.01" IsSnapToTickEnabled="true" SmallChange="0.01" LargeChange="0.01">
137+
<Slider Name="SliderInitialConditioningScale" Value="{Binding SchedulerOptions.ConditioningScale}" Minimum="0" Maximum="1" TickFrequency="0.01" IsSnapToTickEnabled="true" SmallChange="0.01" LargeChange="0.01">
138138
<i:Interaction.Behaviors>
139139
<behaviors:SliderMouseWheelBehavior />
140140
</i:Interaction.Behaviors>

0 commit comments

Comments
 (0)