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

Commit 1f6d747

Browse files
saddam213saddam213
committed
AnnotationPipeline, ControlNet feature extraction example
1 parent cdc80c2 commit 1f6d747

File tree

10 files changed

+394
-276
lines changed

10 files changed

+394
-276
lines changed

‎OnnxStack.Console/Examples/ControlNetExample.cs

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using OnnxStack.Core.Image;
2+
using OnnxStack.StableDiffusion.Common;
23
using OnnxStack.StableDiffusion.Config;
34
using OnnxStack.StableDiffusion.Enums;
45
using OnnxStack.StableDiffusion.Models;
@@ -31,7 +32,7 @@ public ControlNetExample(StableDiffusionConfig configuration)
3132
public async Task RunAsync()
3233
{
3334
// Load Control Image
34-
var controlImage = await File.ReadAllBytesAsync("D:\\Repositories\\OnnxStack\\Assets\\Samples\\OpenPose.png");
35+
var controlImage = await InputImage.FromFileAsync("D:\\Repositories\\OnnxStack\\Assets\\Samples\\OpenPose.png");
3536

3637
// Create ControlNet
3738
var controlNet = ControlNetModel.Create("D:\\Repositories\\controlnet_onnx\\controlnet\\openpose.onnx");
@@ -42,13 +43,15 @@ public async Task RunAsync()
4243
// Prompt
4344
var promptOptions = new PromptOptions
4445
{
45-
Prompt = "Stormtrooper flexing",
46+
Prompt = "Stormtrooper",
4647
DiffuserType = DiffuserType.ControlNet,
47-
InputContolImage = new InputImage(controlImage)
48+
InputContolImage = controlImage
4849
};
4950

51+
52+
5053
// Run pipeline
51-
var result = await pipeline.RunAsync(promptOptions, controlNet: controlNet);
54+
var result = await pipeline.RunAsync(promptOptions, controlNet: controlNet, progressCallback: OutputHelpers.ProgressCallback);
5255

5356
// Create Image from Tensor result
5457
var image = result.ToImage();
@@ -57,6 +60,9 @@ public async Task RunAsync()
5760
var outputFilename = Path.Combine(_outputDirectory, $"Output.png");
5861
await image.SaveAsPngAsync(outputFilename);
5962

63+
//Unload
64+
await controlNet.UnloadAsync();
65+
await pipeline.UnloadAsync();
6066
}
6167
}
6268
}

‎OnnxStack.Console/Examples/ControlNetFeatureExample.cs

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
using OnnxStack.Core.Image;
2+
using OnnxStack.FeatureExtractor.Pipelines;
3+
using OnnxStack.StableDiffusion.Config;
4+
using OnnxStack.StableDiffusion.Enums;
5+
using OnnxStack.StableDiffusion.Models;
6+
using OnnxStack.StableDiffusion.Pipelines;
7+
using SixLabors.ImageSharp;
8+
9+
namespace OnnxStack.Console.Runner
10+
{
11+
public sealed class ControlNetFeatureExample : IExampleRunner
12+
{
13+
private readonly string _outputDirectory;
14+
private readonly StableDiffusionConfig _configuration;
15+
16+
public ControlNetFeatureExample(StableDiffusionConfig configuration)
17+
{
18+
_configuration = configuration;
19+
_outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", nameof(ControlNetFeatureExample));
20+
Directory.CreateDirectory(_outputDirectory);
21+
}
22+
23+
public int Index => 12;
24+
25+
public string Name => "ControlNet + Feature Extraction Example";
26+
27+
public string Description => "ControlNet StableDiffusion with input image Depth feature extraction";
28+
29+
/// <summary>
30+
/// ControlNet Example
31+
/// </summary>
32+
public async Task RunAsync()
33+
{
34+
// Load Control Image
35+
var inputImage = await InputImage.FromFileAsync("D:\\Repositories\\OnnxStack\\Assets\\Samples\\Img2Img_Start.bmp");
36+
37+
// Create Annotation pipeline
38+
var annotationPipeline = AnnotationPipeline.CreatePipeline("D:\\Repositories\\controlnet_onnx\\annotators");
39+
40+
// Create Depth Image
41+
var controlImage = await annotationPipeline.DepthImage(inputImage);
42+
43+
// Save Depth Image (Debug Only)
44+
await controlImage.Image.SaveAsPngAsync(Path.Combine(_outputDirectory, $"Depth.png"));
45+
46+
// Create ControlNet
47+
var controlNet = ControlNetModel.Create("D:\\Repositories\\controlnet_onnx\\controlnet\\depth.onnx");
48+
49+
// Create Pipeline
50+
var pipeline = StableDiffusionPipeline.CreatePipeline("D:\\Repositories\\stable_diffusion_onnx", ModelType.ControlNet);
51+
52+
// Prompt
53+
var promptOptions = new PromptOptions
54+
{
55+
Prompt = "steampunk dog",
56+
DiffuserType = DiffuserType.ControlNet,
57+
InputContolImage = controlImage
58+
};
59+
60+
// Run pipeline
61+
var result = await pipeline.RunAsync(promptOptions, controlNet: controlNet, progressCallback: OutputHelpers.ProgressCallback);
62+
63+
// Create Image from Tensor result
64+
var image = result.ToImage();
65+
66+
// Save Image File
67+
var outputFilename = Path.Combine(_outputDirectory, $"Output.png");
68+
await image.SaveAsPngAsync(outputFilename);
69+
70+
//Unload
71+
await annotationPipeline.UnloadAsync();
72+
await controlNet.UnloadAsync();
73+
await pipeline.UnloadAsync();
74+
}
75+
}
76+
}

‎OnnxStack.Console/OnnxStack.Console.csproj

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
<ItemGroup>
1919
<PackageReference Include="OnnxStack.StableDiffusion" Version="0.15.0" Condition=" '$(Configuration)' == 'Release' OR '$(Configuration)' == 'Release-Cuda' OR '$(Configuration)' == 'Release-TensorRT'" />
2020
<PackageReference Include="OnnxStack.ImageUpscaler" Version="0.15.0" Condition=" '$(Configuration)' == 'Release' OR '$(Configuration)' == 'Release-Cuda' OR '$(Configuration)' == 'Release-TensorRT'" />
21+
<ProjectReference Include="..\OnnxStack.FeatureExtractor\OnnxStack.FeatureExtractor.csproj" />
2122
<ProjectReference Include="..\OnnxStack.StableDiffusion\OnnxStack.StableDiffusion.csproj" Condition=" '$(Configuration)' == 'Debug' OR '$(Configuration)' == 'Debug-Cuda' OR '$(Configuration)' == 'Debug-TensorRT'" />
2223
<ProjectReference Include="..\OnnxStack.ImageUpscaler\OnnxStack.ImageUpscaler.csproj" Condition=" '$(Configuration)' == 'Debug' OR '$(Configuration)' == 'Debug-Cuda' OR '$(Configuration)' == 'Debug-TensorRT'" />
2324
</ItemGroup>

‎OnnxStack.Core/Image/InputImage.cs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
using SixLabors.ImageSharp.PixelFormats;
44
using System.IO;
55
using System.Text.Json.Serialization;
6+
using System.Threading.Tasks;
67

78
namespace OnnxStack.Core.Image
89
{
@@ -89,5 +90,16 @@ public InputImage() { }
8990
|| ImageBytes != null
9091
|| ImageStream != null
9192
|| ImageTensor != null;
93+
94+
95+
/// <summary>
96+
/// Create an image from file
97+
/// </summary>
98+
/// <param name="filePath">The file path.</param>
99+
/// <returns></returns>
100+
public static async Task<InputImage> FromFileAsync(string filePath)
101+
{
102+
return new InputImage(await File.ReadAllBytesAsync(filePath));
103+
}
92104
}
93105
}

‎OnnxStack.FeatureExtractor/Common/AnnotationModelSet.cs

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
using Microsoft.ML.OnnxRuntime;
2+
using OnnxStack.Core.Config;
3+
using System.Text.Json.Serialization;
4+
5+
namespace OnnxStack.FeatureExtractor.Common
6+
{
7+
public record AnnotationModelSet : IOnnxModelSetConfig
8+
{
9+
public string Name { get; set; }
10+
public bool IsEnabled { get; set; }
11+
public int DeviceId { get; set; }
12+
public int InterOpNumThreads { get; set; }
13+
public int IntraOpNumThreads { get; set; }
14+
public ExecutionMode ExecutionMode { get; set; }
15+
public ExecutionProvider ExecutionProvider { get; set; }
16+
17+
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
18+
public FeatureExtractorModelConfig CannyImageConfig { get; set; }
19+
20+
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
21+
public FeatureExtractorModelConfig HedImageConfig { get; set; }
22+
23+
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
24+
public FeatureExtractorModelConfig DepthImageConfig { get; set; }
25+
}
26+
27+
28+
}

‎OnnxStack.FeatureExtractor/Common/FeatureExtractorModel.cs

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
using Microsoft.ML.OnnxRuntime;
2+
using OnnxStack.Core.Config;
3+
using OnnxStack.Core.Model;
4+
5+
namespace OnnxStack.FeatureExtractor.Common
6+
{
7+
public class FeatureExtractorModel : OnnxModelSession
8+
{
9+
private readonly int _sampleSize;
10+
11+
public FeatureExtractorModel(FeatureExtractorModelConfig configuration)
12+
: base(configuration)
13+
{
14+
_sampleSize = configuration.SampleSize;
15+
}
16+
17+
public int SampleSize => _sampleSize;
18+
19+
public static FeatureExtractorModel Create(FeatureExtractorModelConfig configuration)
20+
{
21+
return new FeatureExtractorModel(configuration);
22+
}
23+
24+
public static FeatureExtractorModel Create(string modelFile, int deviceId = 0, ExecutionProvider executionProvider = ExecutionProvider.DirectML)
25+
{
26+
var configuration = new FeatureExtractorModelConfig
27+
{
28+
SampleSize = 512,
29+
DeviceId = deviceId,
30+
ExecutionProvider = executionProvider,
31+
ExecutionMode = ExecutionMode.ORT_SEQUENTIAL,
32+
InterOpNumThreads = 0,
33+
IntraOpNumThreads = 0,
34+
OnnxModelPath = modelFile
35+
};
36+
return new FeatureExtractorModel(configuration);
37+
}
38+
}
39+
40+
public record FeatureExtractorModelConfig : OnnxModelConfig
41+
{
42+
public int SampleSize { get; set; }
43+
}
44+
}

‎OnnxStack.FeatureExtractor/Common/FeatureExtractorModelSet.cs

Lines changed: 0 additions & 18 deletions
This file was deleted.

0 commit comments

Comments
 (0)