Try Fix LCM 1 and 2 step inference

Author: saddam213

OnnxStack.StableDiffusion/Schedulers/LatentConsistency/LCMScheduler.cs

@@ -67,19 +67,19 @@ protected override int[] SetTimesteps()
             var timeIncrement = Options.TrainTimesteps / Options.OriginalInferenceSteps;
 
             //# LCM Training Steps Schedule
-            var lcmOriginTimesteps = Enumerable.Range(1, Options.OriginalInferenceSteps)
-                .Select(x => x * timeIncrement - 1f)
+            var lcmOriginTimesteps = Enumerable.Range(0, Options.OriginalInferenceSteps)
+                .Select(x => x * (timeIncrement - 1))
                 .ToArray();
 
             var skippingStep = lcmOriginTimesteps.Length / Options.InferenceSteps;
 
             // LCM Inference Steps Schedule
-            return lcmOriginTimesteps
+            var steps = lcmOriginTimesteps
                 .Where((t, index) => index % skippingStep == 0)
                 .Take(Options.InferenceSteps)
-                .Select(x => (int)x)
                 .OrderByDescending(x => x)
                 .ToArray();
+            return steps;
         }
 
 
@@ -111,19 +111,20 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
 
             // 1. get previous step value
             int prevIndex = Timesteps.IndexOf(currentTimestep) + 1;
-            int previousTimestep = prevIndex < Timesteps.Count 
-                ? Timesteps[prevIndex] 
+            int previousTimestep = prevIndex < Timesteps.Count
+                ? Timesteps[prevIndex]
                 : currentTimestep;
 
             //# 2. compute alphas, betas
             float alphaProdT = _alphasCumProd[currentTimestep];
-            float alphaProdTPrev = previousTimestep >= 0 
-                ? _alphasCumProd[previousTimestep] 
+            float alphaProdTPrev = previousTimestep >= 0
+                ? _alphasCumProd[previousTimestep]
                 : _finalAlphaCumprod;
             float betaProdT = 1f - alphaProdT;
             float betaProdTPrev = 1f - alphaProdTPrev;
-            float alphaSqrt = MathF.Sqrt(alphaProdT);
-            float betaSqrt = MathF.Sqrt(betaProdT);
+
+            float alphaProdTSqrt = MathF.Sqrt(alphaProdT);
+            float betaProdTSqrt = MathF.Sqrt(betaProdT);
             float betaProdTPrevSqrt = MathF.Sqrt(betaProdTPrev);
             float alphaProdTPrevSqrt = MathF.Sqrt(alphaProdTPrev);
 
@@ -137,8 +138,8 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
             if (Options.PredictionType == PredictionType.Epsilon)
             {
                 predOriginalSample = sample
-                    .SubtractTensors(modelOutput.MultiplyTensorByFloat(betaSqrt))
-                    .DivideTensorByFloat(alphaSqrt);
+                    .SubtractTensors(modelOutput.MultiplyTensorByFloat(betaProdTSqrt))
+                    .DivideTensorByFloat(alphaProdTSqrt);
             }
             else if (Options.PredictionType == PredictionType.Sample)
             {
@@ -147,8 +148,8 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
             else if (Options.PredictionType == PredictionType.VariablePrediction)
             {
                 predOriginalSample = sample
-                    .MultiplyTensorByFloat(alphaSqrt)
-                    .SubtractTensors(modelOutput.MultiplyTensorByFloat(betaSqrt));
+                    .MultiplyTensorByFloat(alphaProdTSqrt)
+                    .SubtractTensors(modelOutput.MultiplyTensorByFloat(betaProdTSqrt));
             }
 
 
@@ -163,13 +164,22 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
 
 
             //# 7. Sample and inject noise z ~ N(0, I) for MultiStep Inference
-            var prevSample = Timesteps.Count > 1
-                ? CreateRandomSample(modelOutput.Dimensions)
+            //# Noise is not used on the final timestep of the timestep schedule.
+            //# This also means that noise is not used for one-step sampling.
+            if (Timesteps.IndexOf(currentTimestep) != Options.InferenceSteps - 1)
+            {
+                var noise = CreateRandomSample(modelOutput.Dimensions);
+                predOriginalSample = noise
                     .MultiplyTensorByFloat(betaProdTPrevSqrt)
-                    .AddTensors(denoised.MultiplyTensorByFloat(alphaProdTPrevSqrt))
-                : denoised;
+                    .AddTensors(denoised.MultiplyTensorByFloat(alphaProdTPrevSqrt));
+            }
+            else
+            {
+                predOriginalSample = denoised;
+            }
+
 
-            return new SchedulerStepResult(prevSample, denoised);
+            return new SchedulerStepResult(predOriginalSample, denoised);
         }
 
 
@@ -203,10 +213,12 @@ public override DenseTensor<float> AddNoise(DenseTensor<float> originalSamples,
         {
             //self.sigma_data = 0.5  # Default: 0.5
             var sigmaData = 0.5f;
+            var timestepScaling = 10f;
+            var scaledTimestep = timestepScaling * timestep;
 
-            float c = MathF.Pow(timestep / 0.1f, 2f) + MathF.Pow(sigmaData, 2f);
+            float c = MathF.Pow(scaledTimestep, 2f) + MathF.Pow(sigmaData, 2f);
             float cSkip = MathF.Pow(sigmaData, 2f) / c;
-            float cOut = timestep / 0.1f / MathF.Pow(c, 0.5f);
+            float cOut = scaledTimestep / MathF.Pow(c, 0.5f);
             return (cSkip, cOut);
         }
 

OnnxStack.UI/UserControls/SchedulerControl.xaml.cs

@@ -122,6 +122,7 @@ private void OnModelChanged(ModelOptionsModel model)
             {
                 SchedulerOptions.OriginalInferenceSteps = 50;
                 SchedulerOptions.InferenceSteps = 6;
+                SchedulerOptions.GuidanceScale = 1f;
             }