Support multi-tensor scheduler results

saddam213 · saddam213 · commit 76e34b431e28 · 2023-11-02T20:07:07.000+13:00
diff --git a/OnnxStack.StableDiffusion/Common/IScheduler.cs b/OnnxStack.StableDiffusion/Common/IScheduler.cs
@@ -1,5 +1,6 @@
 ﻿using Microsoft.ML.OnnxRuntime.Tensors;
 using OnnxStack.StableDiffusion.Enums;
+using OnnxStack.StableDiffusion.Schedulers;
 using System;
 using System.Collections.Generic;
 
@@ -38,7 +39,7 @@ public interface IScheduler : IDisposable
         /// <param name="sample">The sample.</param>
         /// <param name="order">The order.</param>
         /// <returns></returns>
-        DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4);
+        SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4);
 
         /// <summary>
         /// Adds noise to the sample.
diff --git a/OnnxStack.StableDiffusion/Diffusers/DiffuserBase.cs b/OnnxStack.StableDiffusion/Diffusers/DiffuserBase.cs
@@ -106,7 +106,7 @@ public virtual async Task<DenseTensor<float>> DiffuseAsync(IModelOptions modelOp
                             noisePred = PerformGuidance(noisePred, schedulerOptions.GuidanceScale);
 
                         // Scheduler Step
-                        latents = scheduler.Step(noisePred, timestep, latents);
+                        latents = scheduler.Step(noisePred, timestep, latents).PreviousSample;
                     }
 
                     progressCallback?.Invoke(step, timesteps.Count);
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintDiffuser.cs
@@ -87,7 +87,7 @@ public override async Task<DenseTensor<float>> DiffuseAsync(IModelOptions modelO
                             noisePred = PerformGuidance(noisePred, schedulerOptions.GuidanceScale);
 
                         // Scheduler Step
-                        latents = scheduler.Step(noisePred, timestep, latents);
+                        latents = scheduler.Step(noisePred, timestep, latents).PreviousSample;
                     }
 
                     progress?.Invoke(++step, timesteps.Count);
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintLegacyDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintLegacyDiffuser.cs
@@ -82,7 +82,7 @@ public override async Task<DenseTensor<float>> DiffuseAsync(IModelOptions modelO
                             noisePred = PerformGuidance(noisePred, schedulerOptions.GuidanceScale);
 
                         // Scheduler Step
-                        var steplatents = scheduler.Step(noisePred, timestep, latents);
+                        var steplatents = scheduler.Step(noisePred, timestep, latents).PreviousSample;
 
                         // Add noise to original latent
                         var initLatentsProper = scheduler.AddNoise(latentsOriginal, noise, new[] { timestep });
diff --git a/OnnxStack.StableDiffusion/Schedulers/DDIMScheduler.cs b/OnnxStack.StableDiffusion/Schedulers/DDIMScheduler.cs
@@ -86,7 +86,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim
         /// <returns></returns>
         /// <exception cref="System.ArgumentException">Invalid prediction_type: {SchedulerOptions.PredictionType}</exception>
         /// <exception cref="System.NotImplementedException">DDIMScheduler Thresholding currently not implemented</exception>
-        public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
+        public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
         {
             //# See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
             //# Ideally, read DDIM paper in-detail understanding
@@ -176,7 +176,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time
             if (eta > 0)
                 prevSample = prevSample.AddTensors(CreateRandomSample(modelOutput.Dimensions).MultipleTensorByFloat(stdDevT));
 
-            return prevSample;
+            return new SchedulerStepResult(prevSample);
         }
 
 
diff --git a/OnnxStack.StableDiffusion/Schedulers/DDPMScheduler.cs b/OnnxStack.StableDiffusion/Schedulers/DDPMScheduler.cs
@@ -81,7 +81,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim
         /// <returns></returns>
         /// <exception cref="System.ArgumentException">Invalid prediction_type: {SchedulerOptions.PredictionType}</exception>
         /// <exception cref="System.NotImplementedException">DDPMScheduler Thresholding currently not implemented</exception>
-        public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
+        public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
         {
             int currentTimestep = timestep;
             int previousTimestep = GetPreviousTimestep(currentTimestep);
@@ -159,7 +159,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time
                 predPrevSample = predPrevSample.AddTensors(variance);
             }
 
-            return predPrevSample;
+            return new SchedulerStepResult(predPrevSample);
         }
 
 
diff --git a/OnnxStack.StableDiffusion/Schedulers/EulerAncestralScheduler.cs b/OnnxStack.StableDiffusion/Schedulers/EulerAncestralScheduler.cs
@@ -103,7 +103,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim
         /// <param name="sample">The sample.</param>
         /// <param name="order">The order.</param>
         /// <returns></returns>
-        public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
+        public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
         {
             var stepIndex = Timesteps.IndexOf(timestep);
             var sigma = _sigmas[stepIndex];
@@ -129,7 +129,8 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time
             var delta = sigmaDown - sigma;
             var prevSample = sample.AddTensors(derivative.MultipleTensorByFloat(delta));
             var noise = CreateRandomSample(prevSample.Dimensions);
-            return prevSample.AddTensors(noise.MultipleTensorByFloat(sigmaUp));
+            prevSample = prevSample.AddTensors(noise.MultipleTensorByFloat(sigmaUp));
+            return new SchedulerStepResult(prevSample);
         }
 
 
diff --git a/OnnxStack.StableDiffusion/Schedulers/EulerScheduler.cs b/OnnxStack.StableDiffusion/Schedulers/EulerScheduler.cs
@@ -108,7 +108,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim
         /// <param name="sample">The sample.</param>
         /// <param name="order">The order.</param>
         /// <returns></returns>
-        public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
+        public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
         {
             // TODO: Implement "extended settings for scheduler types"
             float s_churn = 0f;
@@ -140,7 +140,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time
                 .DivideTensorByFloat(sigmaHat);
 
             var delta = _sigmas[stepIndex + 1] - sigmaHat;
-            return sample.AddTensors(derivative.MultipleTensorByFloat(delta));
+            return new SchedulerStepResult(sample.AddTensors(derivative.MultipleTensorByFloat(delta)));
         }
 
 
diff --git a/OnnxStack.StableDiffusion/Schedulers/KDPM2Scheduler.cs b/OnnxStack.StableDiffusion/Schedulers/KDPM2Scheduler.cs
@@ -116,7 +116,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim
         /// <returns></returns>
         /// <exception cref="System.ArgumentException">Invalid prediction_type: {SchedulerOptions.PredictionType}</exception>
         /// <exception cref="System.NotImplementedException">KDPM2Scheduler Thresholding currently not implemented</exception>
-        public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
+        public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
         {
             float sigma;
             float sigmaInterpol;
@@ -177,7 +177,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time
             }
 
             _stepIndex += 1;
-            return sample.AddTensors(derivative.MultipleTensorByFloat(dt));
+            return new SchedulerStepResult(sample.AddTensors(derivative.MultipleTensorByFloat(dt)));
         }
 
 
diff --git a/OnnxStack.StableDiffusion/Schedulers/LCMScheduler.cs b/OnnxStack.StableDiffusion/Schedulers/LCMScheduler.cs
@@ -107,7 +107,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim
         /// <param name="sample">The sample.</param>
         /// <param name="order">The order.</param>
         /// <returns></returns>
-        public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
+        public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
         {
             //# Latent Consistency Models paper https://arxiv.org/abs/2310.04378
 
@@ -164,7 +164,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time
                     .AddTensors(denoised.MultipleTensorByFloat(MathF.Sqrt(alphaProdTPrev)))
                 : denoised;
 
-            return prevSample;
+            return new SchedulerStepResult(prevSample, denoised);
         }
 
 
diff --git a/OnnxStack.StableDiffusion/Schedulers/LMSScheduler.cs b/OnnxStack.StableDiffusion/Schedulers/LMSScheduler.cs
@@ -107,7 +107,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim
         /// <param name="sample">The sample.</param>
         /// <param name="order">The order.</param>
         /// <returns></returns>
-        public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
+        public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)
         {
             int stepIndex = Timesteps.IndexOf(timestep);
             var sigma = _sigmas[stepIndex];
@@ -148,7 +148,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time
             }
 
             // Add the sumed tensor to the sample
-            return sample.AddTensors(lmsDerProduct.SumTensors(modelOutput.Dimensions));
+            return new SchedulerStepResult(sample.AddTensors(lmsDerProduct.SumTensors(modelOutput.Dimensions)));
         }
 
 
diff --git a/OnnxStack.StableDiffusion/Schedulers/SchedulerBase.cs b/OnnxStack.StableDiffusion/Schedulers/SchedulerBase.cs
@@ -71,7 +71,7 @@ public SchedulerBase(SchedulerOptions schedulerOptions)
         /// <param name="sample">The sample.</param>
         /// <param name="order">The order.</param>
         /// <returns></returns>
-        public abstract DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4);
+        public abstract SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4);
 
 
         /// <summary>
@@ -443,4 +443,22 @@ protected virtual void Dispose(bool disposing)
 
         #endregion
     }
+
+
+    public class SchedulerStepResult
+    {
+        public SchedulerStepResult(DenseTensor<float> previousSample)
+        {
+            PreviousSample = previousSample;
+        }
+
+        public SchedulerStepResult(DenseTensor<float> previousSample, DenseTensor<float> extraSample)
+        {
+            ExtraSample = extraSample;
+            PreviousSample = previousSample;
+        }
+
+        public DenseTensor<float> PreviousSample { get; set; }
+        public DenseTensor<float> ExtraSample { get; set; }
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -106,7 +106,7 @@ public virtual async Task<DenseTensor<float>> DiffuseAsync(IModelOptions modelOp`
`106`	`106`	`noisePred = PerformGuidance(noisePred, schedulerOptions.GuidanceScale);`
`107`	`107`
`108`	`108`	`// Scheduler Step`
`109`		`- latents = scheduler.Step(noisePred, timestep, latents);`
	`109`	`+ latents = scheduler.Step(noisePred, timestep, latents).PreviousSample;`
`110`	`110`	`}`
`111`	`111`
`112`	`112`	`progressCallback?.Invoke(step, timesteps.Count);`
Original file line number	Diff line number	Diff line change
`@@ -87,7 +87,7 @@ public override async Task<DenseTensor<float>> DiffuseAsync(IModelOptions modelO`
`87`	`87`	`noisePred = PerformGuidance(noisePred, schedulerOptions.GuidanceScale);`
`88`	`88`
`89`	`89`	`// Scheduler Step`
`90`		`- latents = scheduler.Step(noisePred, timestep, latents);`
	`90`	`+ latents = scheduler.Step(noisePred, timestep, latents).PreviousSample;`
`91`	`91`	`}`
`92`	`92`
`93`	`93`	`progress?.Invoke(++step, timesteps.Count);`
Original file line number	Diff line number	Diff line change
`@@ -86,7 +86,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim`
`86`	`86`	`/// <returns></returns>`
`87`	`87`	`/// <exception cref="System.ArgumentException">Invalid prediction_type: {SchedulerOptions.PredictionType}</exception>`
`88`	`88`	`/// <exception cref="System.NotImplementedException">DDIMScheduler Thresholding currently not implemented</exception>`
`89`		`- public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)`
	`89`	`+ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)`
`90`	`90`	`{`
`91`	`91`	`//# See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf`
`92`	`92`	`//# Ideally, read DDIM paper in-detail understanding`
`@@ -176,7 +176,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time`
`176`	`176`	`if (eta > 0)`
`177`	`177`	`prevSample = prevSample.AddTensors(CreateRandomSample(modelOutput.Dimensions).MultipleTensorByFloat(stdDevT));`
`178`	`178`
`179`		`- return prevSample;`
	`179`	`+ return new SchedulerStepResult(prevSample);`
`180`	`180`	`}`
`181`	`181`
`182`	`182`
Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim`
`81`	`81`	`/// <returns></returns>`
`82`	`82`	`/// <exception cref="System.ArgumentException">Invalid prediction_type: {SchedulerOptions.PredictionType}</exception>`
`83`	`83`	`/// <exception cref="System.NotImplementedException">DDPMScheduler Thresholding currently not implemented</exception>`
`84`		`- public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)`
	`84`	`+ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)`
`85`	`85`	`{`
`86`	`86`	`int currentTimestep = timestep;`
`87`	`87`	`int previousTimestep = GetPreviousTimestep(currentTimestep);`
`@@ -159,7 +159,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time`
`159`	`159`	`predPrevSample = predPrevSample.AddTensors(variance);`
`160`	`160`	`}`
`161`	`161`
`162`		`- return predPrevSample;`
	`162`	`+ return new SchedulerStepResult(predPrevSample);`
`163`	`163`	`}`
`164`	`164`
`165`	`165`
Original file line number	Diff line number	Diff line change
`@@ -108,7 +108,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim`
`108`	`108`	`/// <param name="sample">The sample.</param>`
`109`	`109`	`/// <param name="order">The order.</param>`
`110`	`110`	`/// <returns></returns>`
`111`		`- public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)`
	`111`	`+ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)`
`112`	`112`	`{`
`113`	`113`	`// TODO: Implement "extended settings for scheduler types"`
`114`	`114`	`float s_churn = 0f;`
`@@ -140,7 +140,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time`
`140`	`140`	`.DivideTensorByFloat(sigmaHat);`
`141`	`141`
`142`	`142`	`var delta = _sigmas[stepIndex + 1] - sigmaHat;`
`143`		`- return sample.AddTensors(derivative.MultipleTensorByFloat(delta));`
	`143`	`+ return new SchedulerStepResult(sample.AddTensors(derivative.MultipleTensorByFloat(delta)));`
`144`	`144`	`}`
`145`	`145`
`146`	`146`
Original file line number	Diff line number	Diff line change
`@@ -116,7 +116,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim`
`116`	`116`	`/// <returns></returns>`
`117`	`117`	`/// <exception cref="System.ArgumentException">Invalid prediction_type: {SchedulerOptions.PredictionType}</exception>`
`118`	`118`	`/// <exception cref="System.NotImplementedException">KDPM2Scheduler Thresholding currently not implemented</exception>`
`119`		`- public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)`
	`119`	`+ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)`
`120`	`120`	`{`
`121`	`121`	`float sigma;`
`122`	`122`	`float sigmaInterpol;`
`@@ -177,7 +177,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time`
`177`	`177`	`}`
`178`	`178`
`179`	`179`	`_stepIndex += 1;`
`180`		`- return sample.AddTensors(derivative.MultipleTensorByFloat(dt));`
	`180`	`+ return new SchedulerStepResult(sample.AddTensors(derivative.MultipleTensorByFloat(dt)));`
`181`	`181`	`}`
`182`	`182`
`183`	`183`
Original file line number	Diff line number	Diff line change
`@@ -107,7 +107,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim`
`107`	`107`	`/// <param name="sample">The sample.</param>`
`108`	`108`	`/// <param name="order">The order.</param>`
`109`	`109`	`/// <returns></returns>`
`110`		`- public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)`
	`110`	`+ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int timestep, DenseTensor<float> sample, int order = 4)`
`111`	`111`	`{`
`112`	`112`	`//# Latent Consistency Models paper https://arxiv.org/abs/2310.04378`
`113`	`113`
`@@ -164,7 +164,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time`
`164`	`164`	`.AddTensors(denoised.MultipleTensorByFloat(MathF.Sqrt(alphaProdTPrev)))`
`165`	`165`	`: denoised;`
`166`	`166`
`167`		`- return prevSample;`
	`167`	`+ return new SchedulerStepResult(prevSample, denoised);`
`168`	`168`	`}`
`169`	`169`
`170`	`170`