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

Commit 2623d32

Browse files
committed
Use base64 for web uploads, ImageInpaint UI started
1 parent 83190b7 commit 2623d32

22 files changed

+1262
-673
lines changed

OnnxStack.StableDiffusion/Helpers/ImageHelpers.cs

Lines changed: 49 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -67,11 +67,18 @@ public static Stream ToImageStream(this DenseTensor<float> imageTensor)
6767
}
6868

6969

70+
/// <summary>
71+
/// Converts to <see cref="DenseTensor<float>"/>
72+
/// </summary>
73+
/// <param name="imageData">The image data.</param>
74+
/// <param name="width">The width.</param>
75+
/// <param name="height">The height.</param>
76+
/// <returns></returns>
7077
public static DenseTensor<float> ToDenseTensor(this InputImage imageData, int width, int height)
7178
{
72-
if (!string.IsNullOrEmpty(imageData.ImagePath))
73-
return TensorFromFile(imageData.ImagePath, width, height);
74-
if(imageData.ImageBytes != null)
79+
if (!string.IsNullOrEmpty(imageData.ImageBase64))
80+
return TensorFromBase64(imageData.ImageBase64, width, height);
81+
if (imageData.ImageBytes != null)
7582
return TensorFromBytes(imageData.ImageBytes, width, height);
7683
if (imageData.ImageStream != null)
7784
return TensorFromStream(imageData.ImageStream, width, height);
@@ -145,6 +152,30 @@ public static DenseTensor<float> TensorFromFile(string filename, int width, int
145152
}
146153

147154

155+
/// <summary>
156+
/// Create an image Tensor from base64 string.
157+
/// </summary>
158+
/// <param name="base64Image">The base64 image.</param>
159+
/// <param name="width">The width.</param>
160+
/// <param name="height">The height.</param>
161+
public static DenseTensor<float> TensorFromBase64(string base64Image, int width, int height)
162+
{
163+
using (Image<Rgb24> image = Image.Load<Rgb24>(Convert.FromBase64String(base64Image.Split(',')[1])))
164+
{
165+
Resize(image, width, height);
166+
return ProcessPixels(width, height, image);
167+
}
168+
}
169+
170+
171+
/// <summary>
172+
/// Create an image Tensor from bytes.
173+
/// </summary>
174+
/// <param name="imageBytes">The image bytes.</param>
175+
/// <param name="width">The width.</param>
176+
/// <param name="height">The height.</param>
177+
/// <returns>
178+
/// </returns>
148179
public static DenseTensor<float> TensorFromBytes(byte[] imageBytes, int width, int height)
149180
{
150181
using (var image = Image.Load<Rgb24>(imageBytes))
@@ -154,6 +185,14 @@ public static DenseTensor<float> TensorFromBytes(byte[] imageBytes, int width, i
154185
}
155186
}
156187

188+
189+
/// <summary>
190+
/// Create an image Tensor from stream.
191+
/// </summary>
192+
/// <param name="imageStream">The image stream.</param>
193+
/// <param name="width">The width.</param>
194+
/// <param name="height">The height.</param>
195+
/// <returns></returns>
157196
public static DenseTensor<float> TensorFromStream(Stream imageStream, int width, int height)
158197
{
159198
using (var image = Image.Load<Rgb24>(imageStream))
@@ -164,8 +203,13 @@ public static DenseTensor<float> TensorFromStream(Stream imageStream, int width,
164203
}
165204

166205

167-
168-
206+
/// <summary>
207+
/// Processes the pixels.
208+
/// </summary>
209+
/// <param name="width">The width.</param>
210+
/// <param name="height">The height.</param>
211+
/// <param name="image">The image.</param>
212+
/// <returns></returns>
169213
private static DenseTensor<float> ProcessPixels(int width, int height, Image<Rgb24> image)
170214
{
171215
var imageArray = new DenseTensor<float>(new[] { 1, 3, width, height });

OnnxStack.StableDiffusion/Models/InputImage.cs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@ public InputImage() { }
2222
/// <summary>
2323
/// Initializes a new instance of the <see cref="InputImage"/> class.
2424
/// </summary>
25-
/// <param name="imagePath">The image path.</param>
26-
public InputImage(string imagePath) => ImagePath = imagePath;
25+
/// <param name="imageBase64">The image in base64 format.</param>
26+
public InputImage(string imageBase64) => ImageBase64 = imageBase64;
2727

2828
/// <summary>
2929
/// Initializes a new instance of the <see cref="InputImage"/> class.
@@ -51,9 +51,9 @@ public InputImage() { }
5151

5252

5353
/// <summary>
54-
/// Gets the image path.
54+
/// Gets the image base64 string.
5555
/// </summary>
56-
public string ImagePath { get; set; }
56+
public string ImageBase64 { get; set; }
5757

5858

5959
/// <summary>
@@ -85,7 +85,7 @@ public InputImage() { }
8585
/// </value>
8686
[JsonIgnore]
8787
public bool HasImage => Image != null
88-
|| !string.IsNullOrEmpty(ImagePath)
88+
|| !string.IsNullOrEmpty(ImageBase64)
8989
|| ImageBytes != null
9090
|| ImageStream != null
9191
|| ImageTensor != null;

OnnxStack.WebUI/Hubs/StableDiffusionHub.cs

Lines changed: 93 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
using Microsoft.AspNetCore.SignalR;
22
using OnnxStack.StableDiffusion.Common;
33
using OnnxStack.StableDiffusion.Config;
4-
using OnnxStack.StableDiffusion.Models;
54
using OnnxStack.WebUI.Models;
65
using Services;
76
using SixLabors.ImageSharp;
@@ -96,6 +95,62 @@ public async IAsyncEnumerable<StableDiffusionResult> OnExecuteImageToImage(Promp
9695
}
9796

9897

98+
/// <summary>
99+
/// Execute Image-Inpaint Stable Diffusion
100+
/// </summary>
101+
/// <param name="options">The options.</param>
102+
/// <param name="cancellationToken">The cancellation token.</param>
103+
/// <returns></returns>
104+
[HubMethodName("ExecuteImageInpaint")]
105+
public async IAsyncEnumerable<StableDiffusionResult> OnExecuteImageInpaint(PromptOptions promptOptions, SchedulerOptions schedulerOptions, [EnumeratorCancellation] CancellationToken cancellationToken)
106+
{
107+
_logger.Log(LogLevel.Information, "[ExecuteImageInpaint] - New request received, Connection: {0}", Context.ConnectionId);
108+
var cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(Context.ConnectionAborted, cancellationToken);
109+
110+
// TODO: Add support for multiple results
111+
var result = await GenerateImageInpaintResult(promptOptions, schedulerOptions, cancellationTokenSource.Token);
112+
if (!result.IsError)
113+
yield return result;
114+
115+
await Clients.Caller.OnError(result.Error);
116+
yield break;
117+
}
118+
119+
120+
/// <summary>
121+
/// Generates the text to image result.
122+
/// </summary>
123+
/// <param name="options">The options.</param>
124+
/// <param name="cancellationToken">The cancellation token.</param>
125+
/// <returns></returns>
126+
private async Task<StableDiffusionResult> GenerateTextToImageResult(PromptOptions promptOptions, SchedulerOptions schedulerOptions, CancellationToken cancellationToken)
127+
{
128+
var timestamp = Stopwatch.GetTimestamp();
129+
schedulerOptions.Seed = GenerateSeed(schedulerOptions.Seed);
130+
131+
//1. Create filenames
132+
var random = await _fileService.CreateRandomName();
133+
var outputImage = $"{random}-output.png";
134+
var outputBlueprint = $"{random}-output.json";
135+
var outputImageUrl = await _fileService.CreateOutputUrl(outputImage);
136+
var outputImageFile = await _fileService.UrlToPhysicalPath(outputImageUrl);
137+
var outputImageLink = await _fileService.CreateOutputUrl(outputImage, false);
138+
139+
//2. Generate blueprint
140+
var blueprint = new ImageBlueprint(promptOptions, schedulerOptions, outputImageLink);
141+
var bluprintFile = await _fileService.SaveBlueprintFile(blueprint, outputBlueprint);
142+
if (bluprintFile is null)
143+
return new StableDiffusionResult("Failed to save blueprint");
144+
145+
//3. Run stable diffusion
146+
if (!await RunStableDiffusion(promptOptions, schedulerOptions, outputImageFile, cancellationToken))
147+
return new StableDiffusionResult("Failed to run stable diffusion");
148+
149+
//4. Return result
150+
return new StableDiffusionResult(outputImage, outputImageUrl, blueprint, bluprintFile.Filename, bluprintFile.FileUrl, GetElapsed(timestamp));
151+
}
152+
153+
99154
/// <summary>
100155
/// Generates the image to image result.
101156
/// </summary>
@@ -110,72 +165,79 @@ private async Task<StableDiffusionResult> GenerateImageToImageResult(PromptOptio
110165

111166
//1. Create filenames
112167
var random = await _fileService.CreateRandomName();
113-
var output = $"Output-{random}";
114-
var outputImage = $"{output}.png";
115-
var outputBlueprint = $"{output}.json";
116-
var inputImage = $"Input-{random}.png";
117-
var uploadImage = Path.GetFileName(promptOptions.InputImage.ImagePath);
168+
var inputImage = $"{random}-input.png";
169+
var outputImage = $"{random}-output.png";
170+
var outputBlueprint = $"{random}-output.json";
118171
var outputImageUrl = await _fileService.CreateOutputUrl(outputImage);
119172
var outputImageFile = await _fileService.UrlToPhysicalPath(outputImageUrl);
120-
var inputOriginaUrl = await _fileService.CreateOutputUrl(uploadImage, false);
173+
var outputImageLink = await _fileService.CreateOutputUrl(outputImage, false);
174+
var inputImageLink = await _fileService.CreateOutputUrl(inputImage, false);
121175

122-
//2. Copy input image to new file
123-
var inputImageFile = await _fileService.CopyInputImageFile(promptOptions.InputImage.ImagePath, inputImage);
176+
// 2. Save Input Image to disk
177+
var inputImageFile = await _fileService.SaveImageFile(promptOptions.InputImage.ImageBase64, inputImage);
124178
if (inputImageFile is null)
125-
return new StableDiffusionResult("Failed to copy input image");
179+
return new StableDiffusionResult("Failed to save input image");
126180

127-
//3. Generate blueprint
128-
var inputImageLink = await _fileService.CreateOutputUrl(inputImage, false);
129-
var outputImageLink = await _fileService.CreateOutputUrl(outputImage, false);
130-
promptOptions.InputImage = new InputImage(inputOriginaUrl);
181+
//3. Generate and save blueprint file
131182
var blueprint = new ImageBlueprint(promptOptions, schedulerOptions, outputImageLink, inputImageLink);
132183
var bluprintFile = await _fileService.SaveBlueprintFile(blueprint, outputBlueprint);
133184
if (bluprintFile is null)
134185
return new StableDiffusionResult("Failed to save blueprint");
135186

136-
//4. Set full path of input image
137-
promptOptions.InputImage = new InputImage(inputImageFile.FilePath);
138-
139-
//5. Run stable diffusion
187+
//4. Run stable diffusion
140188
if (!await RunStableDiffusion(promptOptions, schedulerOptions, outputImageFile, cancellationToken))
141189
return new StableDiffusionResult("Failed to run stable diffusion");
142190

143-
//6. Return result
191+
//5. Return result
144192
return new StableDiffusionResult(outputImage, outputImageUrl, blueprint, bluprintFile.Filename, bluprintFile.FileUrl, GetElapsed(timestamp));
145193
}
146194

147195

148196
/// <summary>
149-
/// Generates the text to image result.
197+
/// Generates the image inpaint result.
150198
/// </summary>
151-
/// <param name="options">The options.</param>
199+
/// <param name="promptOptions">The prompt options.</param>
200+
/// <param name="schedulerOptions">The scheduler options.</param>
152201
/// <param name="cancellationToken">The cancellation token.</param>
153202
/// <returns></returns>
154-
private async Task<StableDiffusionResult> GenerateTextToImageResult(PromptOptions promptOptions, SchedulerOptions schedulerOptions, CancellationToken cancellationToken)
203+
private async Task<StableDiffusionResult> GenerateImageInpaintResult(PromptOptions promptOptions, SchedulerOptions schedulerOptions, CancellationToken cancellationToken)
155204
{
156205
var timestamp = Stopwatch.GetTimestamp();
157206
schedulerOptions.Seed = GenerateSeed(schedulerOptions.Seed);
158207

159-
//1. Create filenames
208+
// 1. Create filenames
160209
var random = await _fileService.CreateRandomName();
161-
var output = $"Output-{random}";
162-
var outputImage = $"{output}.png";
163-
var outputBlueprint = $"{output}.json";
210+
var maskImage = $"{random}-mask.png";
211+
var inputImage = $"{random}-input.png";
212+
var outputImage = $"{random}-output.png";
213+
var outputBlueprint = $"{random}-output.json";
164214
var outputImageUrl = await _fileService.CreateOutputUrl(outputImage);
165215
var outputImageFile = await _fileService.UrlToPhysicalPath(outputImageUrl);
166-
167-
//2. Generate blueprint
168216
var outputImageLink = await _fileService.CreateOutputUrl(outputImage, false);
169-
var blueprint = new ImageBlueprint(promptOptions, schedulerOptions, outputImageLink);
217+
var inputImageLink = await _fileService.CreateOutputUrl(inputImage, false);
218+
var maskImageLink = await _fileService.CreateOutputUrl(maskImage, false);
219+
220+
// 2. Save Input Image to disk
221+
var inputImageFile = await _fileService.SaveImageFile(promptOptions.InputImage.ImageBase64, inputImage);
222+
if (inputImageFile is null)
223+
return new StableDiffusionResult("Failed to save input image");
224+
225+
// 3. Save Mask Image to disk
226+
var maskImageFile = await _fileService.SaveImageFile(promptOptions.InputImageMask.ImageBase64, maskImage);
227+
if (maskImageFile is null)
228+
return new StableDiffusionResult("Failed to save mask image");
229+
230+
// 4. Generate and save blueprint file
231+
var blueprint = new ImageBlueprint(promptOptions, schedulerOptions, outputImageLink, inputImageLink, maskImageLink);
170232
var bluprintFile = await _fileService.SaveBlueprintFile(blueprint, outputBlueprint);
171233
if (bluprintFile is null)
172234
return new StableDiffusionResult("Failed to save blueprint");
173235

174-
//3. Run stable diffusion
236+
// 5. Run stable diffusion
175237
if (!await RunStableDiffusion(promptOptions, schedulerOptions, outputImageFile, cancellationToken))
176238
return new StableDiffusionResult("Failed to run stable diffusion");
177239

178-
//4. Return result
240+
// 6. Return result
179241
return new StableDiffusionResult(outputImage, outputImageUrl, blueprint, bluprintFile.Filename, bluprintFile.FileUrl, GetElapsed(timestamp));
180242
}
181243

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,29 @@
11
using OnnxStack.StableDiffusion.Config;
2+
using OnnxStack.StableDiffusion.Enums;
23

34
namespace OnnxStack.WebUI.Models
45
{
5-
public record ImageBlueprint(PromptOptions Prompt, SchedulerOptions Options, string OutputImageUrl, string InputImageUrl = null)
6+
public record ImageBlueprint
67
{
7-
public DateTime Timestamp { get; } = DateTime.UtcNow;
8+
public ImageBlueprint(PromptOptions prompt, SchedulerOptions options, string outputImageUrl, string inputImageUrl = null, string maskImageUrl = null)
9+
{
10+
Timestamp = DateTime.Now;
11+
SchedulerOptions = options;
12+
MaskImageUrl = maskImageUrl;
13+
InputImageUrl = inputImageUrl;
14+
OutputImageUrl = outputImageUrl;
15+
Prompt = prompt.Prompt;
16+
NegativePrompt = prompt.NegativePrompt;
17+
SchedulerType = prompt.SchedulerType;
18+
}
19+
20+
public DateTime Timestamp { get; }
21+
public string Prompt { get; set; }
22+
public string NegativePrompt { get; set; }
23+
public SchedulerType SchedulerType { get; set; }
24+
public string MaskImageUrl { get; set; }
25+
public string InputImageUrl { get; set; }
26+
public string OutputImageUrl { get; set; }
27+
public SchedulerOptions SchedulerOptions { get; set; }
828
}
929
}
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
namespace Models
22
{
3-
public record InitialImageModel(string Name, string Url, int Width, int Height);
3+
public record InitialImageModel(string Url, int Width, int Height);
44
}

OnnxStack.WebUI/Pages/Error.cshtml.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ namespace OnnxStack.WebUI.Pages
88
[IgnoreAntiforgeryToken]
99
public class ErrorModel : PageModel
1010
{
11-
public string? RequestId { get; set; }
11+
public string RequestId { get; set; }
1212

1313
public bool ShowRequestId => !string.IsNullOrEmpty(RequestId);
1414

OnnxStack.WebUI/Pages/Index.cshtml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,11 +43,11 @@
4343
<div class="card border-1" style="width: 18rem;border-left:1px">
4444
<img src="~/images/placeholder.jpg" class="card-img-top" alt="...">
4545
<div class="card-body">
46-
<h5 class="card-title">In-Painting</h5>
46+
<h5 class="card-title">Image Inpainting</h5>
4747
<p class="card-text">Inpainting is a computer vision technique that reconstructs missing or damaged parts of an image, seamlessly filling in the gaps to restore visual completeness</p>
4848
</div>
4949
<div class="card-footer">
50-
<button href="#" class="d-block btn btn-outline-primary w-100" disabled>Demo</button>
50+
<a href="/StableDiffusion/ImageInpaint" class=" d-block btn btn-outline-primary">Demo</a>
5151
</div>
5252
</div>
5353
</div>

0 commit comments

Comments
 (0)