Inpainting prototype

Author: saddam213

OnnxStack.StableDiffusion/Common/IStableDiffusionService.cs

@@ -28,8 +28,8 @@ public interface IStableDiffusionService
         /// <param name="options">The Scheduler options.</param>
         /// <param name="progressCallback">The callback used to provide progess of the current InferenceSteps.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
-        /// <returns>The diffusion result as <see cref="SixLabors.ImageSharp.Image<Rgb24>"/></returns>
-        Task<Image<Rgb24>> GenerateAsImageAsync(PromptOptions prompt, SchedulerOptions options, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default);
+        /// <returns>The diffusion result as <see cref="SixLabors.ImageSharp.Image<Rgba32>"/></returns>
+        Task<Image<Rgba32>> GenerateAsImageAsync(PromptOptions prompt, SchedulerOptions options, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default);
 
         /// <summary>
         /// Generates the StableDiffusion image using the prompt and options provided.

OnnxStack.StableDiffusion/Diffusers/DiffuserBase.cs

@@ -66,8 +66,7 @@ public virtual async Task<DenseTensor<float>> DiffuseAsync(PromptOptions promptO
         {
             // Create random seed if none was set
             schedulerOptions.Seed = schedulerOptions.Seed > 0 ? schedulerOptions.Seed : Random.Shared.Next();
-            Console.WriteLine($"Scheduler: {promptOptions.SchedulerType}, Size: {schedulerOptions.Width}x{schedulerOptions.Height}, Seed: {schedulerOptions.Seed}, Steps: {schedulerOptions.InferenceSteps}, Guidance: {schedulerOptions.GuidanceScale}");
-
+     
             // Get Scheduler
             using (var scheduler = GetScheduler(promptOptions, schedulerOptions))
             {
@@ -78,7 +77,7 @@ public virtual async Task<DenseTensor<float>> DiffuseAsync(PromptOptions promptO
                 var timesteps = GetTimesteps(promptOptions, schedulerOptions, scheduler);
 
                 // Create latent sample
-                var latentSample = PrepareLatents(promptOptions, schedulerOptions, scheduler, timesteps);
+                var latents = PrepareLatents(promptOptions, schedulerOptions, scheduler, timesteps);
 
                 // Loop though the timesteps
                 var step = 0;
@@ -87,8 +86,9 @@ public virtual async Task<DenseTensor<float>> DiffuseAsync(PromptOptions promptO
                     cancellationToken.ThrowIfCancellationRequested();
 
                     // Create input tensor.
-                    var inputTensor = scheduler.ScaleInput(latentSample.Duplicate(schedulerOptions.GetScaledDimension(2)), timestep);
+                    var inputTensor = scheduler.ScaleInput(latents.Duplicate(schedulerOptions.GetScaledDimension(2)), timestep);
 
+                    // Create Input Parameters
                     var inputNames = _onnxModelService.GetInputNames(OnnxModelType.Unet);
                     var inputParameters = CreateInputParameters(
                          NamedOnnxValue.CreateFromTensor(inputNames[0], inputTensor),
@@ -98,27 +98,24 @@ public virtual async Task<DenseTensor<float>> DiffuseAsync(PromptOptions promptO
                     // Run Inference
                     using (var inferResult = await _onnxModelService.RunInferenceAsync(OnnxModelType.Unet, inputParameters))
                     {
-                        var resultTensor = inferResult.FirstElementAs<DenseTensor<float>>();
-
-                        // Split tensors from 2,4,(H/8),(W/8) to 1,4,(H/8),(W/8)
-                        var splitTensors = resultTensor.SplitTensor(schedulerOptions.GetScaledDimension(), schedulerOptions.GetScaledHeight(), schedulerOptions.GetScaledWidth());
-                        var noisePred = splitTensors.Item1;
-                        var noisePredText = splitTensors.Item2;
+                        var noisePred = inferResult.FirstElementAs<DenseTensor<float>>();
 
                         // Perform guidance
-                        noisePred = noisePred.PerformGuidance(noisePredText, schedulerOptions.GuidanceScale);
+                        if (schedulerOptions.GuidanceScale > 1.0f)
+                        {
+                            var (noisePredUncond, noisePredText) = noisePred.SplitTensor(schedulerOptions.GetScaledDimension());
+                            noisePred = noisePredUncond.PerformGuidance(noisePredText, schedulerOptions.GuidanceScale);
+                        }
 
-                        // LMS Scheduler Step
-                        latentSample = scheduler.Step(noisePred, timestep, latentSample);
-                        // ImageHelpers.TensorToImageDebug(latentSample, 64, $@"Examples\StableDebug\Latent_{step}.png");
+                        // Scheduler Step
+                        latents = scheduler.Step(noisePred, timestep, latents);
                     }
 
-                    Console.WriteLine($"Step: {++step}/{timesteps.Count}");
-                    progress?.Invoke(step, timesteps.Count);
+                    progress?.Invoke(++step, timesteps.Count);
                 }
 
                 // Decode Latents
-                return await DecodeLatents(schedulerOptions, latentSample);
+                return await DecodeLatents(schedulerOptions, latents);
             }
         }
 
@@ -192,7 +189,7 @@ protected static DenseTensor<float> ClipImageFeatureExtractor(SchedulerOptions o
             using (var image = imageTensor.ToImage())
             {
                 // Resize image
-                ImageHelpers.Resize(image, 224, 224);
+                ImageHelpers.Resize(image, new[] { 1, 3, 224, 224 });
 
                 // Preprocess image
                 var input = new DenseTensor<float>(new[] { 1, 3, 224, 224 });

OnnxStack.StableDiffusion/Diffusers/ImageDiffuser.cs

@@ -53,7 +53,7 @@ protected override IReadOnlyList<int> GetTimesteps(PromptOptions prompt, Schedul
         protected override DenseTensor<float> PrepareLatents(PromptOptions prompt, SchedulerOptions options, IScheduler scheduler, IReadOnlyList<int> timesteps)
         {
             // Image input, decode, add noise, return as latent 0
-            var imageTensor = prompt.InputImage.ToDenseTensor(options.Width, options.Height);
+            var imageTensor = prompt.InputImage.ToDenseTensor(new[] { 1, 3, options.Width, options.Height });
             var inputNames = _onnxModelService.GetInputNames(OnnxModelType.VaeEncoder);
             var inputParameters = CreateInputParameters(NamedOnnxValue.CreateFromTensor(inputNames[0], imageTensor));
             using (var inferResult = _onnxModelService.RunInference(OnnxModelType.VaeEncoder, inputParameters))