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

Vectorize most of TensorHelper #15

Merged
merged 2 commits into from
Nov 16, 2023
Merged

Vectorize most of TensorHelper #15

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
31562c1
Vectorize most of TensorHelper
jdluzen Nov 10, 2023
b78508d
Merge branch 'saddam213:master' into master
jdluzen Nov 10, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
169 changes: 111 additions & 58 deletions OnnxStack.StableDiffusion/Helpers/TensorHelper.cs
View file
Open in desktop

Large diffs are not rendered by default.

18 changes: 9 additions & 9 deletions OnnxStack.StableDiffusion/Schedulers/LatentConsistency/LCMScheduler.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
if (Options.PredictionType == PredictionType.Epsilon)
{
predOriginalSample = sample
.SubtractTensors(modelOutput.MultipleTensorByFloat(betaSqrt))
.SubtractTensors(modelOutput.MultiplyTensorByFloat(betaSqrt))
.DivideTensorByFloat(alphaSqrt);
}
else if (Options.PredictionType == PredictionType.Sample)
Expand All @@ -147,8 +147,8 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
else if (Options.PredictionType == PredictionType.VariablePrediction)
{
predOriginalSample = sample
.MultipleTensorByFloat(alphaSqrt)
.SubtractTensors(modelOutput.MultipleTensorByFloat(betaSqrt));
.MultiplyTensorByFloat(alphaSqrt)
.SubtractTensors(modelOutput.MultiplyTensorByFloat(betaSqrt));
}


Expand All @@ -158,15 +158,15 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim

//# 6. Denoise model output using boundary conditions
var denoised = sample
.MultipleTensorByFloat(cSkip)
.AddTensors(predOriginalSample.MultipleTensorByFloat(cOut));
.MultiplyTensorByFloat(cSkip)
.AddTensors(predOriginalSample.MultiplyTensorByFloat(cOut));


//# 7. Sample and inject noise z ~ N(0, I) for MultiStep Inference
var prevSample = Timesteps.Count > 1
? CreateRandomSample(modelOutput.Dimensions)
.MultipleTensorByFloat(betaProdTPrevSqrt)
.AddTensors(denoised.MultipleTensorByFloat(alphaProdTPrevSqrt))
.MultiplyTensorByFloat(betaProdTPrevSqrt)
.AddTensors(denoised.MultiplyTensorByFloat(alphaProdTPrevSqrt))
: denoised;

return new SchedulerStepResult(prevSample, denoised);
Expand All @@ -189,8 +189,8 @@ public override DenseTensor<float> AddNoise(DenseTensor<float> originalSamples,
float sqrtOneMinusAlpha = MathF.Sqrt(1.0f - alphaProd);

return noise
.MultipleTensorByFloat(sqrtOneMinusAlpha)
.AddTensors(originalSamples.MultipleTensorByFloat(sqrtAlpha));
.MultiplyTensorByFloat(sqrtOneMinusAlpha)
.AddTensors(originalSamples.MultiplyTensorByFloat(sqrtAlpha));
}


Expand Down
4 changes: 2 additions & 2 deletions OnnxStack.StableDiffusion/Schedulers/SchedulerBase.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,13 +205,13 @@ protected virtual DenseTensor<float> GetPredictedSample(DenseTensor<float> model
DenseTensor<float> predOriginalSample = null;
if (Options.PredictionType == PredictionType.Epsilon)
{
predOriginalSample = sample.SubtractTensors(modelOutput.MultipleTensorByFloat(sigma));
predOriginalSample = sample.SubtractTensors(modelOutput.MultiplyTensorByFloat(sigma));
}
else if (Options.PredictionType == PredictionType.VariablePrediction)
{
var sigmaSqrt = (float)Math.Sqrt(sigma * sigma + 1);
predOriginalSample = sample.DivideTensorByFloat(sigmaSqrt)
.AddTensors(modelOutput.MultipleTensorByFloat(-sigma / sigmaSqrt));
.AddTensors(modelOutput.MultiplyTensorByFloat(-sigma / sigmaSqrt));
}
else if (Options.PredictionType == PredictionType.Sample)
{
Expand Down
24 changes: 12 additions & 12 deletions OnnxStack.StableDiffusion/Schedulers/StableDiffusion/DDIMScheduler.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,27 +114,27 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
DenseTensor<float> predOriginalSample = null;
if (Options.PredictionType == PredictionType.Epsilon)
{
var sampleBeta = sample.SubtractTensors(modelOutput.MultipleTensorByFloat((float)Math.Sqrt(betaProdT)));
var sampleBeta = sample.SubtractTensors(modelOutput.MultiplyTensorByFloat((float)Math.Sqrt(betaProdT)));
predOriginalSample = sampleBeta.DivideTensorByFloat((float)Math.Sqrt(alphaProdT));
predEpsilon = modelOutput;
}
else if (Options.PredictionType == PredictionType.Sample)
{
predOriginalSample = modelOutput;
predEpsilon = sample.SubtractTensors(predOriginalSample
.MultipleTensorByFloat((float)Math.Sqrt(alphaProdT)))
.MultiplyTensorByFloat((float)Math.Sqrt(alphaProdT)))
.DivideTensorByFloat((float)Math.Sqrt(betaProdT));
}
else if (Options.PredictionType == PredictionType.VariablePrediction)
{
var alphaSqrt = (float)Math.Sqrt(alphaProdT);
var betaSqrt = (float)Math.Sqrt(betaProdT);
predOriginalSample = sample
.MultipleTensorByFloat(alphaSqrt)
.SubtractTensors(modelOutput.MultipleTensorByFloat(betaSqrt));
.MultiplyTensorByFloat(alphaSqrt)
.SubtractTensors(modelOutput.MultiplyTensorByFloat(betaSqrt));
predEpsilon = modelOutput
.MultipleTensorByFloat(alphaSqrt)
.AddTensors(sample.MultipleTensorByFloat(betaSqrt));
.MultiplyTensorByFloat(alphaSqrt)
.AddTensors(sample.MultiplyTensorByFloat(betaSqrt));
}


Expand All @@ -161,20 +161,20 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
{
//# the pred_epsilon is always re-derived from the clipped x_0 in Glide
predEpsilon = sample
.SubtractTensors(predOriginalSample.MultipleTensorByFloat((float)Math.Sqrt(alphaProdT)))
.SubtractTensors(predOriginalSample.MultiplyTensorByFloat((float)Math.Sqrt(alphaProdT)))
.DivideTensorByFloat((float)Math.Sqrt(betaProdT));
}


//# 5. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
var predSampleDirection = predEpsilon.MultipleTensorByFloat((float)Math.Sqrt(1f - alphaProdTPrev - Math.Pow(stdDevT, 2f)));
var predSampleDirection = predEpsilon.MultiplyTensorByFloat((float)Math.Sqrt(1f - alphaProdTPrev - Math.Pow(stdDevT, 2f)));


//# 6. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
var prevSample = predSampleDirection.AddTensors(predOriginalSample.MultipleTensorByFloat((float)Math.Sqrt(alphaProdTPrev)));
var prevSample = predSampleDirection.AddTensors(predOriginalSample.MultiplyTensorByFloat((float)Math.Sqrt(alphaProdTPrev)));

if (eta > 0)
prevSample = prevSample.AddTensors(CreateRandomSample(modelOutput.Dimensions).MultipleTensorByFloat(stdDevT));
prevSample = prevSample.AddTensors(CreateRandomSample(modelOutput.Dimensions).MultiplyTensorByFloat(stdDevT));

return new SchedulerStepResult(prevSample);
}
Expand All @@ -196,8 +196,8 @@ public override DenseTensor<float> AddNoise(DenseTensor<float> originalSamples,
float sqrtOneMinusAlpha = (float)Math.Sqrt(1.0f - alphaProd);

return noise
.MultipleTensorByFloat(sqrtOneMinusAlpha)
.AddTensors(originalSamples.MultipleTensorByFloat(sqrtAlpha));
.MultiplyTensorByFloat(sqrtOneMinusAlpha)
.AddTensors(originalSamples.MultiplyTensorByFloat(sqrtAlpha));
}


Expand Down
20 changes: 10 additions & 10 deletions OnnxStack.StableDiffusion/Schedulers/StableDiffusion/DDPMScheduler.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,8 +132,8 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
//# See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
//pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample
var predPrevSample = sample
.MultipleTensorByFloat(currentSampleCoeff)
.AddTensors(predOriginalSample.MultipleTensorByFloat(predOriginalSampleCoeff));
.MultiplyTensorByFloat(currentSampleCoeff)
.AddTensors(predOriginalSample.MultiplyTensorByFloat(predOriginalSampleCoeff));


//# 6. Add noise
Expand All @@ -144,17 +144,17 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
if (Options.VarianceType == VarianceType.FixedSmallLog)
{
var v = GetVariance(currentTimestep, predictedVariance);
variance = varianceNoise.MultipleTensorByFloat(v);
variance = varianceNoise.MultiplyTensorByFloat(v);
}
else if (Options.VarianceType == VarianceType.LearnedRange)
{
var v = (float)Math.Exp(0.5 * GetVariance(currentTimestep, predictedVariance));
variance = varianceNoise.MultipleTensorByFloat(v);
variance = varianceNoise.MultiplyTensorByFloat(v);
}
else
{
var v = (float)Math.Sqrt(GetVariance(currentTimestep, predictedVariance));
variance = varianceNoise.MultipleTensorByFloat(v);
variance = varianceNoise.MultiplyTensorByFloat(v);
}
predPrevSample = predPrevSample.AddTensors(variance);
}
Expand All @@ -169,7 +169,7 @@ private DenseTensor<float> GetPredictedSample(DenseTensor<float> modelOutput, De
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)));
var sampleBeta = sample.SubtractTensors(modelOutput.MultiplyTensorByFloat((float)Math.Sqrt(betaProdT)));
predOriginalSample = sampleBeta.DivideTensorByFloat((float)Math.Sqrt(alphaProdT));
}
else if (Options.PredictionType == PredictionType.Sample)
Expand All @@ -182,8 +182,8 @@ private DenseTensor<float> GetPredictedSample(DenseTensor<float> modelOutput, De
var alphaSqrt = (float)Math.Sqrt(alphaProdT);
var betaSqrt = (float)Math.Sqrt(betaProdT);
predOriginalSample = sample
.MultipleTensorByFloat(alphaSqrt)
.SubtractTensors(modelOutput.MultipleTensorByFloat(betaSqrt));
.MultiplyTensorByFloat(alphaSqrt)
.SubtractTensors(modelOutput.MultiplyTensorByFloat(betaSqrt));
}
return predOriginalSample;
}
Expand All @@ -204,8 +204,8 @@ public override DenseTensor<float> AddNoise(DenseTensor<float> originalSamples,
float sqrtOneMinusAlpha = (float)Math.Sqrt(1.0f - alphaProd);

return noise
.MultipleTensorByFloat(sqrtOneMinusAlpha)
.AddTensors(originalSamples.MultipleTensorByFloat(sqrtAlpha));
.MultiplyTensorByFloat(sqrtOneMinusAlpha)
.AddTensors(originalSamples.MultiplyTensorByFloat(sqrtAlpha));
}

/// <summary>
Expand Down
6 changes: 3 additions & 3 deletions OnnxStack.StableDiffusion/Schedulers/StableDiffusion/EulerAncestralScheduler.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,9 +127,9 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
.DivideTensorByFloat(sigma);

var delta = sigmaDown - sigma;
var prevSample = sample.AddTensors(derivative.MultipleTensorByFloat(delta));
var prevSample = sample.AddTensors(derivative.MultiplyTensorByFloat(delta));
var noise = CreateRandomSample(prevSample.Dimensions);
prevSample = prevSample.AddTensors(noise.MultipleTensorByFloat(sigmaUp));
prevSample = prevSample.AddTensors(noise.MultiplyTensorByFloat(sigmaUp));
return new SchedulerStepResult(prevSample);
}

Expand All @@ -152,7 +152,7 @@ public override DenseTensor<float> AddNoise(DenseTensor<float> originalSamples,
.Max();

return noise
.MultipleTensorByFloat(sigma)
.MultiplyTensorByFloat(sigma)
.AddTensors(originalSamples);
}

Expand Down
10 changes: 5 additions & 5 deletions OnnxStack.StableDiffusion/Schedulers/StableDiffusion/EulerScheduler.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,17 +121,17 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim

float gamma = s_tmin <= sigma && sigma <= s_tmax ? (float)Math.Min(s_churn / (_sigmas.Length - 1f), Math.Sqrt(2.0f) - 1.0f) : 0f;
var noise = CreateRandomSample(modelOutput.Dimensions);
var epsilon = noise.MultipleTensorByFloat(s_noise);
var epsilon = noise.MultiplyTensorByFloat(s_noise);
float sigmaHat = sigma * (1.0f + gamma);

if (gamma > 0)
sample = sample.AddTensors(epsilon.MultipleTensorByFloat((float)Math.Sqrt(Math.Pow(sigmaHat, 2f) - Math.Pow(sigma, 2f))));
sample = sample.AddTensors(epsilon.MultiplyTensorByFloat((float)Math.Sqrt(Math.Pow(sigmaHat, 2f) - Math.Pow(sigma, 2f))));


// 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
var predOriginalSample = Options.PredictionType != PredictionType.Epsilon
? GetPredictedSample(modelOutput, sample, sigma)
: sample.SubtractTensors(modelOutput.MultipleTensorByFloat(sigmaHat));
: sample.SubtractTensors(modelOutput.MultiplyTensorByFloat(sigmaHat));


// 2. Convert to an ODE derivative
Expand All @@ -140,7 +140,7 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
.DivideTensorByFloat(sigmaHat);

var delta = _sigmas[stepIndex + 1] - sigmaHat;
return new SchedulerStepResult(sample.AddTensors(derivative.MultipleTensorByFloat(delta)));
return new SchedulerStepResult(sample.AddTensors(derivative.MultiplyTensorByFloat(delta)));
}


Expand All @@ -160,7 +160,7 @@ public override DenseTensor<float> AddNoise(DenseTensor<float> originalSamples,
.Max();

return noise
.MultipleTensorByFloat(sigma)
.MultiplyTensorByFloat(sigma)
.AddTensors(originalSamples);
}

Expand Down
8 changes: 4 additions & 4 deletions OnnxStack.StableDiffusion/Schedulers/StableDiffusion/KDPM2Scheduler.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,13 +142,13 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
DenseTensor<float> predOriginalSample;
if (Options.PredictionType == PredictionType.Epsilon)
{
predOriginalSample = sample.SubtractTensors(modelOutput.MultipleTensorByFloat(sigmaInput));
predOriginalSample = sample.SubtractTensors(modelOutput.MultiplyTensorByFloat(sigmaInput));
}
else if (Options.PredictionType == PredictionType.VariablePrediction)
{
var sigmaSqrt = (float)Math.Sqrt(sigmaInput * sigmaInput + 1f);
predOriginalSample = sample.DivideTensorByFloat(sigmaSqrt)
.AddTensors(modelOutput.MultipleTensorByFloat(-sigmaInput / sigmaSqrt));
.AddTensors(modelOutput.MultiplyTensorByFloat(-sigmaInput / sigmaSqrt));
}
else
{
Expand Down Expand Up @@ -177,7 +177,7 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
}

_stepIndex += 1;
return new SchedulerStepResult(sample.AddTensors(derivative.MultipleTensorByFloat(dt)));
return new SchedulerStepResult(sample.AddTensors(derivative.MultiplyTensorByFloat(dt)));
}


Expand All @@ -192,7 +192,7 @@ public override DenseTensor<float> AddNoise(DenseTensor<float> originalSamples,
{
var sigma = _sigmas[_stepIndex];
return noise
.MultipleTensorByFloat(sigma)
.MultiplyTensorByFloat(sigma)
.AddTensors(originalSamples);
}

Expand Down
4 changes: 2 additions & 2 deletions OnnxStack.StableDiffusion/Schedulers/StableDiffusion/LMSScheduler.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
{
// Multiply to coeff by each derivatives to create the new tensors
var (lmsCoeff, derivative) = lmsCoeffsAndDerivatives[i];
lmsDerProduct[i] = derivative.MultipleTensorByFloat((float)lmsCoeff);
lmsDerProduct[i] = derivative.MultiplyTensorByFloat((float)lmsCoeff);
}

// Add the sumed tensor to the sample
Expand All @@ -168,7 +168,7 @@ public override DenseTensor<float> AddNoise(DenseTensor<float> originalSamples,
.Max();

return noise
.MultipleTensorByFloat(sigma)
.MultiplyTensorByFloat(sigma)
.AddTensors(originalSamples);
}

Expand Down