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

Commit 71eb138

Browse files
committed
Add LCM InpaintLegacyDiffuser base
1 parent 1e89247 commit 71eb138

File tree

2 files changed

+73
-4
lines changed

2 files changed

+73
-4
lines changed
Lines changed: 72 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,80 @@
1-
using System;
1+
using Microsoft.Extensions.Logging;
2+
using Microsoft.ML.OnnxRuntime;
3+
using Microsoft.ML.OnnxRuntime.Tensors;
4+
using OnnxStack.Core.Config;
5+
using OnnxStack.Core.Services;
6+
using OnnxStack.StableDiffusion.Common;
7+
using OnnxStack.StableDiffusion.Config;
8+
using OnnxStack.StableDiffusion.Enums;
9+
using OnnxStack.StableDiffusion.Helpers;
10+
using SixLabors.ImageSharp;
11+
using System;
212
using System.Collections.Generic;
313
using System.Linq;
4-
using System.Text;
5-
using System.Threading.Tasks;
614

715
namespace OnnxStack.StableDiffusion.Diffusers.LatentConsistency
816
{
9-
internal class InpaintLegacyDiffuser
17+
public sealed class InpaintLegacyDiffuser : LatentConsistencyDiffuser
1018
{
19+
/// <summary>
20+
/// Initializes a new instance of the <see cref="InpaintLegacyDiffuser"/> class.
21+
/// </summary>
22+
/// <param name="configuration">The configuration.</param>
23+
/// <param name="onnxModelService">The onnx model service.</param>
24+
public InpaintLegacyDiffuser(IOnnxModelService onnxModelService, IPromptService promptService, ILogger<LatentConsistencyDiffuser> logger)
25+
: base(onnxModelService, promptService, logger)
26+
{
27+
}
28+
29+
30+
/// <summary>
31+
/// Gets the type of the diffuser.
32+
/// </summary>
33+
public override DiffuserType DiffuserType => DiffuserType.ImageInpaintLegacy;
34+
35+
36+
/// <summary>
37+
/// Gets the timesteps.
38+
/// </summary>
39+
/// <param name="prompt">The prompt.</param>
40+
/// <param name="options">The options.</param>
41+
/// <param name="scheduler">The scheduler.</param>
42+
/// <returns></returns>
43+
protected override IReadOnlyList<int> GetTimesteps(PromptOptions prompt, SchedulerOptions options, IScheduler scheduler)
44+
{
45+
// Image2Image we narrow step the range by the Strength
46+
var inittimestep = Math.Min((int)(options.InferenceSteps * options.Strength), options.InferenceSteps);
47+
var start = Math.Max(options.InferenceSteps - inittimestep, 0);
48+
return scheduler.Timesteps.Skip(start).ToList();
49+
}
50+
51+
52+
/// <summary>
53+
/// Prepares the latents for inference.
54+
/// </summary>
55+
/// <param name="prompt">The prompt.</param>
56+
/// <param name="options">The options.</param>
57+
/// <param name="scheduler">The scheduler.</param>
58+
/// <returns></returns>
59+
protected override DenseTensor<float> PrepareLatents(IModelOptions model, PromptOptions prompt, SchedulerOptions options, IScheduler scheduler, IReadOnlyList<int> timesteps)
60+
{
61+
// Image input, decode, add noise, return as latent 0
62+
var imageTensor = prompt.InputImage.ToDenseTensor(new[] { 1, 3, options.Height, options.Width });
63+
var inputNames = _onnxModelService.GetInputNames(model, OnnxModelType.VaeEncoder);
64+
var inputParameters = CreateInputParameters(NamedOnnxValue.CreateFromTensor(inputNames[0], imageTensor));
65+
using (var inferResult = _onnxModelService.RunInference(model, OnnxModelType.VaeEncoder, inputParameters))
66+
{
67+
var sample = inferResult.FirstElementAs<DenseTensor<float>>();
68+
var scaledSample = sample
69+
.Add(scheduler.CreateRandomSample(sample.Dimensions, options.InitialNoiseLevel))
70+
.MultiplyBy(model.ScaleFactor);
71+
72+
var noisySample = scheduler.AddNoise(scaledSample, scheduler.CreateRandomSample(scaledSample.Dimensions), timesteps);
73+
if (prompt.BatchCount > 1)
74+
return noisySample.Repeat(prompt.BatchCount);
75+
76+
return noisySample;
77+
}
78+
}
1179
}
1280
}

OnnxStack.StableDiffusion/Registration.cs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ public static void AddOnnxStackStableDiffusion(this IServiceCollection serviceCo
4242
//LatentConsistency
4343
serviceCollection.AddSingleton<IDiffuser, StableDiffusion.Diffusers.LatentConsistency.TextDiffuser>();
4444
serviceCollection.AddSingleton<IDiffuser, StableDiffusion.Diffusers.LatentConsistency.ImageDiffuser>();
45+
serviceCollection.AddSingleton<IDiffuser, StableDiffusion.Diffusers.LatentConsistency.InpaintLegacyDiffuser>();
4546
}
4647

4748

0 commit comments

Comments
 (0)