Sigmoid BetaSchedule support, refactor SchedulerBase

Author: saddam213

OnnxStack.StableDiffusion/Enums/BetaScheduleType.cs

@@ -4,6 +4,7 @@ public enum BetaScheduleType
     {
         Linear = 0,
         ScaledLinear = 1,
-        SquaredCosCapV2 = 2
+        SquaredCosCapV2 = 2,
+        Sigmoid = 3
     }
 }

OnnxStack.StableDiffusion/Schedulers/DDPMScheduler.cs

@@ -11,8 +11,7 @@ namespace OnnxStack.StableDiffusion.Schedulers
 {
     internal class DDPMScheduler : SchedulerBase
     {
-        private float[] _betas;
-        private List<float> _alphasCumulativeProducts;
+        private float[] _alphasCumProd;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="DDPMScheduler"/> class.
@@ -33,43 +32,13 @@ public DDPMScheduler(SchedulerOptions options) : base(options) { }
         /// </summary>
         protected override void Initialize()
         {
-            var alphas = new List<float>();
-            if (Options.TrainedBetas != null)
-            {
-                _betas = Options.TrainedBetas.ToArray();
-            }
-            else if (Options.BetaSchedule == BetaScheduleType.Linear)
-            {
-                _betas = np.linspace(Options.BetaStart, Options.BetaEnd, Options.TrainTimesteps).ToArray<float>();
-            }
-            else if (Options.BetaSchedule == BetaScheduleType.ScaledLinear)
-            {
-                // This schedule is very specific to the latent diffusion model.
-                _betas = np.power(np.linspace(MathF.Sqrt(Options.BetaStart), MathF.Sqrt(Options.BetaEnd), Options.TrainTimesteps), 2).ToArray<float>();
-            }
-            else if (Options.BetaSchedule == BetaScheduleType.SquaredCosCapV2)
-            {
-                // Glide cosine schedule
-                _betas = GetBetasForAlphaBar();
-            }
-            //else if (betaSchedule == "sigmoid")
-            //{
-            //    // GeoDiff sigmoid schedule
-            //    var betas = np.linspace(-6, 6, numTrainTimesteps);
-            //    Betas = (np.multiply(np.exp(betas), (betaEnd - betaStart)) + betaStart).ToArray<float>();
-            //}
-
-
-            for (int i = 0; i < Options.TrainTimesteps; i++)
-            {
-                alphas.Add(1.0f - _betas[i]);
-            }
+            _alphasCumProd = null;
 
-            _alphasCumulativeProducts = new List<float> { alphas[0] };
-            for (int i = 1; i < Options.TrainTimesteps; i++)
-            {
-                _alphasCumulativeProducts.Add(_alphasCumulativeProducts[i - 1] * alphas[i]);
-            }
+            var betas = GetBetaSchedule();
+            var alphas = betas.Select(beta => 1.0f - beta);
+            _alphasCumProd = alphas
+                .Select((alpha, i) => alphas.Take(i + 1).Aggregate((a, b) => a * b))
+                .ToArray();
 
             SetInitNoiseSigma(1.0f);
         }
@@ -82,29 +51,8 @@ protected override void Initialize()
         protected override int[] SetTimesteps()
         {
             // Create timesteps based on the specified strategy
-            NDArray timestepsArray = null;
-            if (Options.TimestepSpacing == TimestepSpacingType.Linspace)
-            {
-                timestepsArray = np.linspace(0, Options.TrainTimesteps - 1, Options.InferenceSteps);
-                timestepsArray = np.around(timestepsArray)["::1"];
-            }
-            else if (Options.TimestepSpacing == TimestepSpacingType.Leading)
-            {
-                var stepRatio = Options.TrainTimesteps / Options.InferenceSteps;
-                timestepsArray = np.arange(0, (float)Options.InferenceSteps) * stepRatio;
-                timestepsArray = np.around(timestepsArray)["::1"];
-                timestepsArray += Options.StepsOffset;
-            }
-            else if (Options.TimestepSpacing == TimestepSpacingType.Trailing)
-            {
-                var stepRatio = Options.TrainTimesteps / (Options.InferenceSteps - 1);
-                timestepsArray = np.arange((float)Options.TrainTimesteps, 0, -stepRatio)["::-1"];
-                timestepsArray = np.around(timestepsArray);
-                timestepsArray -= 1;
-            }
-
-            return timestepsArray
-                .ToArray<float>()
+            var timesteps = GetTimesteps();
+            return timesteps
                 .Select(x => (int)x)
                 .OrderByDescending(x => x)
                 .ToArray();
@@ -139,8 +87,8 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time
             int previousTimestep = GetPreviousTimestep(currentTimestep);
 
             //# 1. compute alphas, betas
-            float alphaProdT = _alphasCumulativeProducts[currentTimestep];
-            float alphaProdTPrev = previousTimestep >= 0 ? _alphasCumulativeProducts[previousTimestep] : 1f;
+            float alphaProdT = _alphasCumProd[currentTimestep];
+            float alphaProdTPrev = previousTimestep >= 0 ? _alphasCumProd[previousTimestep] : 1f;
             float betaProdT = 1 - alphaProdT;
             float betaProdTPrev = 1 - alphaProdTPrev;
             float currentAlphaT = alphaProdT / alphaProdTPrev;
@@ -161,27 +109,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time
 
             //# 2. compute predicted original sample from predicted noise also called
             //# "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
-            DenseTensor<float> predOriginalSample = null;
-            if (Options.PredictionType == PredictionType.Epsilon)
-            {
-                //pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
-                var sampleBeta = sample.SubtractTensors(modelOutput.MultipleTensorByFloat((float)Math.Sqrt(betaProdT)));
-                predOriginalSample = sampleBeta.DivideTensorByFloat((float)Math.Sqrt(alphaProdT), sampleBeta.Dimensions);
-            }
-            else if (Options.PredictionType == PredictionType.Sample)
-            {
-                predOriginalSample = modelOutput;
-            }
-            else if (Options.PredictionType == PredictionType.VariablePrediction)
-            {
-                // pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
-                var alphaSqrt = (float)Math.Sqrt(alphaProdT);
-                var betaSqrt = (float)Math.Sqrt(betaProdT);
-                predOriginalSample = sample
-                    .MultipleTensorByFloat(alphaSqrt)
-                    .SubtractTensors(modelOutput.MultipleTensorByFloat(betaSqrt));
-            }
-
+            var predOriginalSample = GetPredictedSample(modelOutput, sample, alphaProdT, betaProdT);
 
             //# 3. Clip or threshold "predicted x_0"
             if (Options.Thresholding)
@@ -234,6 +162,31 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time
         }
 
 
+        private DenseTensor<float> GetPredictedSample(DenseTensor<float> modelOutput, DenseTensor<float> sample, float alphaProdT, float betaProdT)
+        {
+            DenseTensor<float> predOriginalSample = null;
+            if (Options.PredictionType == PredictionType.Epsilon)
+            {
+                //pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
+                var sampleBeta = sample.SubtractTensors(modelOutput.MultipleTensorByFloat((float)Math.Sqrt(betaProdT)));
+                predOriginalSample = sampleBeta.DivideTensorByFloat((float)Math.Sqrt(alphaProdT), sampleBeta.Dimensions);
+            }
+            else if (Options.PredictionType == PredictionType.Sample)
+            {
+                predOriginalSample = modelOutput;
+            }
+            else if (Options.PredictionType == PredictionType.VariablePrediction)
+            {
+                // pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
+                var alphaSqrt = (float)Math.Sqrt(alphaProdT);
+                var betaSqrt = (float)Math.Sqrt(betaProdT);
+                predOriginalSample = sample
+                    .MultipleTensorByFloat(alphaSqrt)
+                    .SubtractTensors(modelOutput.MultipleTensorByFloat(betaSqrt));
+            }
+            return predOriginalSample;
+        }
+
         /// <summary>
         /// Adds noise to the sample.
         /// </summary>
@@ -245,7 +198,7 @@ public override DenseTensor<float> AddNoise(DenseTensor<float> originalSamples,
         {
             // Ref: https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_ddpm.py#L456
             int timestep = timesteps[0];
-            float alphaProd = _alphasCumulativeProducts[timestep];
+            float alphaProd = _alphasCumProd[timestep];
             float sqrtAlpha = (float)Math.Sqrt(alphaProd);
             float sqrtOneMinusAlpha = (float)Math.Sqrt(1.0f - alphaProd);
 
@@ -263,8 +216,8 @@ public override DenseTensor<float> AddNoise(DenseTensor<float> originalSamples,
         private float GetVariance(int timestep, float predictedVariance = 0f)
         {
             int prevTimestep = GetPreviousTimestep(timestep);
-            float alphaProdT = _alphasCumulativeProducts[timestep];
-            float alphaProdTPrev = prevTimestep >= 0 ? _alphasCumulativeProducts[prevTimestep] : 1.0f;
+            float alphaProdT = _alphasCumProd[timestep];
+            float alphaProdTPrev = prevTimestep >= 0 ? _alphasCumProd[prevTimestep] : 1.0f;
             float currentBetaT = 1 - alphaProdT / alphaProdTPrev;
 
             // For t > 0, compute predicted variance βt
@@ -384,8 +337,7 @@ private bool IsVarianceTypeLearned()
 
         protected override void Dispose(bool disposing)
         {
-            _betas = null;
-            _alphasCumulativeProducts = null;
+            _alphasCumProd = null;
             base.Dispose(disposing);
         }
     }

OnnxStack.StableDiffusion/Schedulers/EulerAncestralScheduler.cs

@@ -34,38 +34,15 @@ public EulerAncestralScheduler(SchedulerOptions schedulerOptions) : base(schedul
         protected override void Initialize()
         {
             _sigmas = null;
-            var betas = Enumerable.Empty<float>();
-            if (!Options.TrainedBetas.IsNullOrEmpty())
-            {
-                betas = Options.TrainedBetas;
-            }
-            else if (Options.BetaSchedule == BetaScheduleType.Linear)
-            {
-                betas = np.linspace(Options.BetaStart, Options.BetaEnd, Options.TrainTimesteps).ToArray<float>();
-            }
-            else if (Options.BetaSchedule == BetaScheduleType.ScaledLinear)
-            {
-                var start = (float)Math.Sqrt(Options.BetaStart);
-                var end = (float)Math.Sqrt(Options.BetaEnd);
-                betas = np.linspace(start, end, Options.TrainTimesteps)
-                    .ToArray<float>()
-                    .Select(x => x * x);
-            }
-            else if (Options.BetaSchedule == BetaScheduleType.SquaredCosCapV2)
-            {
-                betas = GetBetasForAlphaBar();
-            }
 
+            var betas = GetBetaSchedule();
             var alphas = betas.Select(beta => 1.0f - beta);
             var alphaCumProd = alphas.Select((alpha, i) => alphas.Take(i + 1).Aggregate((a, b) => a * b));
             _sigmas = alphaCumProd
                  .Select(alpha_prod => (float)Math.Sqrt((1 - alpha_prod) / alpha_prod))
                  .ToArray();
 
-            var maxSigma = _sigmas.Max();
-            var initNoiseSigma = Options.TimestepSpacing == TimestepSpacingType.Linspace || Options.TimestepSpacing == TimestepSpacingType.Trailing
-                ? maxSigma
-                : (float)Math.Sqrt(maxSigma * maxSigma + 1);
+            var initNoiseSigma = GetInitNoiseSigma(_sigmas);
             SetInitNoiseSigma(initNoiseSigma);
         }
 
@@ -76,29 +53,8 @@ protected override void Initialize()
         /// <returns></returns>
         protected override int[] SetTimesteps()
         {
-            NDArray timestepsArray = null;
-            if (Options.TimestepSpacing == TimestepSpacingType.Linspace)
-            {
-                timestepsArray = np.linspace(0, Options.TrainTimesteps - 1, Options.InferenceSteps);
-                timestepsArray = np.around(timestepsArray)["::1"];
-            }
-            else if (Options.TimestepSpacing == TimestepSpacingType.Leading)
-            {
-                var stepRatio = Options.TrainTimesteps / Options.InferenceSteps;
-                timestepsArray = np.arange(0, (float)Options.InferenceSteps) * stepRatio;
-                timestepsArray = np.around(timestepsArray)["::1"];
-                timestepsArray += Options.StepsOffset;
-            }
-            else if (Options.TimestepSpacing == TimestepSpacingType.Trailing)
-            {
-                var stepRatio = Options.TrainTimesteps / (Options.InferenceSteps - 1);
-                timestepsArray = np.arange(Options.TrainTimesteps, 0f, -stepRatio)["::-1"];
-                timestepsArray = np.around(timestepsArray);
-                timestepsArray -= 1;
-            }
-
             var sigmas = _sigmas.ToArray();
-            var timesteps = timestepsArray.ToArray<float>();
+            var timesteps = GetTimesteps();
             var log_sigmas = np.log(sigmas).ToArray<float>();
             var range = np.arange(0, (float)_sigmas.Length).ToArray<float>();
             sigmas = Interpolate(timesteps, range, _sigmas);
@@ -135,9 +91,7 @@ public override DenseTensor<float> ScaleInput(DenseTensor<float> sample, int tim
             sigma = (float)Math.Sqrt(Math.Pow(sigma, 2) + 1);
 
             // Divide sample tensor shape {2,4,(H/8),(W/8)} by sigma
-            sample = sample.DivideTensorByFloat(sigma, sample.Dimensions);
-
-            return sample;
+            return sample.DivideTensorByFloat(sigma, sample.Dimensions);
         }
 
 
@@ -155,23 +109,7 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time
             var sigma = _sigmas[stepIndex];
 
             // 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
-            DenseTensor<float> predOriginalSample = null;
-            if (Options.PredictionType == PredictionType.Epsilon)
-            {
-                predOriginalSample = sample.SubtractTensors(modelOutput.MultipleTensorByFloat(sigma));
-            }
-            else if (Options.PredictionType == PredictionType.VariablePrediction)
-            {
-                var sigmaSqrt = (float)Math.Sqrt(sigma * sigma + 1);
-                predOriginalSample = sample.DivideTensorByFloat(sigmaSqrt)
-                    .AddTensors(modelOutput.MultipleTensorByFloat(-sigma / sigmaSqrt));
-            }
-            else if (Options.PredictionType == PredictionType.Sample)
-            {
-                //prediction_type not implemented yet: sample
-                predOriginalSample = modelOutput.ToDenseTensor();
-            }
-
+            var predOriginalSample = GetPredictedSample(modelOutput, sample, sigma);
 
             var sigmaFrom = _sigmas[stepIndex];
             var sigmaTo = _sigmas[stepIndex + 1];
@@ -195,6 +133,8 @@ public override DenseTensor<float> Step(DenseTensor<float> modelOutput, int time
         }
 
 
+
+
         /// <summary>
         /// Adds noise to the sample.
         /// </summary>