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

Commit b1bd7ac

Browse files
saddam213saddam213
authored
Merge pull request #15 from jdluzen/master
Vectorize most of TensorHelper
2 parents 429b910 + b78508d commit b1bd7ac

File tree

9 files changed

+158
-105
lines changed

9 files changed

+158
-105
lines changed

‎OnnxStack.StableDiffusion/Helpers/TensorHelper.cs‎

Lines changed: 111 additions & 58 deletions
Large diffs are not rendered by default.

‎OnnxStack.StableDiffusion/Schedulers/LatentConsistency/LCMScheduler.cs‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
137137
if (Options.PredictionType == PredictionType.Epsilon)
138138
{
139139
predOriginalSample = sample
140-
.SubtractTensors(modelOutput.MultipleTensorByFloat(betaSqrt))
140+
.SubtractTensors(modelOutput.MultiplyTensorByFloat(betaSqrt))
141141
.DivideTensorByFloat(alphaSqrt);
142142
}
143143
else if (Options.PredictionType == PredictionType.Sample)
@@ -147,8 +147,8 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
147147
else if (Options.PredictionType == PredictionType.VariablePrediction)
148148
{
149149
predOriginalSample = sample
150-
.MultipleTensorByFloat(alphaSqrt)
151-
.SubtractTensors(modelOutput.MultipleTensorByFloat(betaSqrt));
150+
.MultiplyTensorByFloat(alphaSqrt)
151+
.SubtractTensors(modelOutput.MultiplyTensorByFloat(betaSqrt));
152152
}
153153

154154

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

159159
//# 6. Denoise model output using boundary conditions
160160
var denoised = sample
161-
.MultipleTensorByFloat(cSkip)
162-
.AddTensors(predOriginalSample.MultipleTensorByFloat(cOut));
161+
.MultiplyTensorByFloat(cSkip)
162+
.AddTensors(predOriginalSample.MultiplyTensorByFloat(cOut));
163163

164164

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

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

191191
return noise
192-
.MultipleTensorByFloat(sqrtOneMinusAlpha)
193-
.AddTensors(originalSamples.MultipleTensorByFloat(sqrtAlpha));
192+
.MultiplyTensorByFloat(sqrtOneMinusAlpha)
193+
.AddTensors(originalSamples.MultiplyTensorByFloat(sqrtAlpha));
194194
}
195195

196196

‎OnnxStack.StableDiffusion/Schedulers/SchedulerBase.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -205,13 +205,13 @@ protected virtual DenseTensor<float> GetPredictedSample(DenseTensor<float> model
205205
DenseTensor<float> predOriginalSample = null;
206206
if (Options.PredictionType == PredictionType.Epsilon)
207207
{
208-
predOriginalSample = sample.SubtractTensors(modelOutput.MultipleTensorByFloat(sigma));
208+
predOriginalSample = sample.SubtractTensors(modelOutput.MultiplyTensorByFloat(sigma));
209209
}
210210
else if (Options.PredictionType == PredictionType.VariablePrediction)
211211
{
212212
var sigmaSqrt = (float)Math.Sqrt(sigma * sigma + 1);
213213
predOriginalSample = sample.DivideTensorByFloat(sigmaSqrt)
214-
.AddTensors(modelOutput.MultipleTensorByFloat(-sigma / sigmaSqrt));
214+
.AddTensors(modelOutput.MultiplyTensorByFloat(-sigma / sigmaSqrt));
215215
}
216216
else if (Options.PredictionType == PredictionType.Sample)
217217
{

‎OnnxStack.StableDiffusion/Schedulers/StableDiffusion/DDIMScheduler.cs‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -114,27 +114,27 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
114114
DenseTensor<float> predOriginalSample = null;
115115
if (Options.PredictionType == PredictionType.Epsilon)
116116
{
117-
var sampleBeta = sample.SubtractTensors(modelOutput.MultipleTensorByFloat((float)Math.Sqrt(betaProdT)));
117+
var sampleBeta = sample.SubtractTensors(modelOutput.MultiplyTensorByFloat((float)Math.Sqrt(betaProdT)));
118118
predOriginalSample = sampleBeta.DivideTensorByFloat((float)Math.Sqrt(alphaProdT));
119119
predEpsilon = modelOutput;
120120
}
121121
else if (Options.PredictionType == PredictionType.Sample)
122122
{
123123
predOriginalSample = modelOutput;
124124
predEpsilon = sample.SubtractTensors(predOriginalSample
125-
.MultipleTensorByFloat((float)Math.Sqrt(alphaProdT)))
125+
.MultiplyTensorByFloat((float)Math.Sqrt(alphaProdT)))
126126
.DivideTensorByFloat((float)Math.Sqrt(betaProdT));
127127
}
128128
else if (Options.PredictionType == PredictionType.VariablePrediction)
129129
{
130130
var alphaSqrt = (float)Math.Sqrt(alphaProdT);
131131
var betaSqrt = (float)Math.Sqrt(betaProdT);
132132
predOriginalSample = sample
133-
.MultipleTensorByFloat(alphaSqrt)
134-
.SubtractTensors(modelOutput.MultipleTensorByFloat(betaSqrt));
133+
.MultiplyTensorByFloat(alphaSqrt)
134+
.SubtractTensors(modelOutput.MultiplyTensorByFloat(betaSqrt));
135135
predEpsilon = modelOutput
136-
.MultipleTensorByFloat(alphaSqrt)
137-
.AddTensors(sample.MultipleTensorByFloat(betaSqrt));
136+
.MultiplyTensorByFloat(alphaSqrt)
137+
.AddTensors(sample.MultiplyTensorByFloat(betaSqrt));
138138
}
139139

140140

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

168168

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

172172

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

176176
if (eta > 0)
177-
prevSample = prevSample.AddTensors(CreateRandomSample(modelOutput.Dimensions).MultipleTensorByFloat(stdDevT));
177+
prevSample = prevSample.AddTensors(CreateRandomSample(modelOutput.Dimensions).MultiplyTensorByFloat(stdDevT));
178178

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

198198
return noise
199-
.MultipleTensorByFloat(sqrtOneMinusAlpha)
200-
.AddTensors(originalSamples.MultipleTensorByFloat(sqrtAlpha));
199+
.MultiplyTensorByFloat(sqrtOneMinusAlpha)
200+
.AddTensors(originalSamples.MultiplyTensorByFloat(sqrtAlpha));
201201
}
202202

203203

‎OnnxStack.StableDiffusion/Schedulers/StableDiffusion/DDPMScheduler.cs‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -132,8 +132,8 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
132132
//# See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
133133
//pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample
134134
var predPrevSample = sample
135-
.MultipleTensorByFloat(currentSampleCoeff)
136-
.AddTensors(predOriginalSample.MultipleTensorByFloat(predOriginalSampleCoeff));
135+
.MultiplyTensorByFloat(currentSampleCoeff)
136+
.AddTensors(predOriginalSample.MultiplyTensorByFloat(predOriginalSampleCoeff));
137137

138138

139139
//# 6. Add noise
@@ -144,17 +144,17 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
144144
if (Options.VarianceType == VarianceType.FixedSmallLog)
145145
{
146146
var v = GetVariance(currentTimestep, predictedVariance);
147-
variance = varianceNoise.MultipleTensorByFloat(v);
147+
variance = varianceNoise.MultiplyTensorByFloat(v);
148148
}
149149
else if (Options.VarianceType == VarianceType.LearnedRange)
150150
{
151151
var v = (float)Math.Exp(0.5 * GetVariance(currentTimestep, predictedVariance));
152-
variance = varianceNoise.MultipleTensorByFloat(v);
152+
variance = varianceNoise.MultiplyTensorByFloat(v);
153153
}
154154
else
155155
{
156156
var v = (float)Math.Sqrt(GetVariance(currentTimestep, predictedVariance));
157-
variance = varianceNoise.MultipleTensorByFloat(v);
157+
variance = varianceNoise.MultiplyTensorByFloat(v);
158158
}
159159
predPrevSample = predPrevSample.AddTensors(variance);
160160
}
@@ -169,7 +169,7 @@ private DenseTensor<float> GetPredictedSample(DenseTensor<float> modelOutput, De
169169
if (Options.PredictionType == PredictionType.Epsilon)
170170
{
171171
//pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
172-
var sampleBeta = sample.SubtractTensors(modelOutput.MultipleTensorByFloat((float)Math.Sqrt(betaProdT)));
172+
var sampleBeta = sample.SubtractTensors(modelOutput.MultiplyTensorByFloat((float)Math.Sqrt(betaProdT)));
173173
predOriginalSample = sampleBeta.DivideTensorByFloat((float)Math.Sqrt(alphaProdT));
174174
}
175175
else if (Options.PredictionType == PredictionType.Sample)
@@ -182,8 +182,8 @@ private DenseTensor<float> GetPredictedSample(DenseTensor<float> modelOutput, De
182182
var alphaSqrt = (float)Math.Sqrt(alphaProdT);
183183
var betaSqrt = (float)Math.Sqrt(betaProdT);
184184
predOriginalSample = sample
185-
.MultipleTensorByFloat(alphaSqrt)
186-
.SubtractTensors(modelOutput.MultipleTensorByFloat(betaSqrt));
185+
.MultiplyTensorByFloat(alphaSqrt)
186+
.SubtractTensors(modelOutput.MultiplyTensorByFloat(betaSqrt));
187187
}
188188
return predOriginalSample;
189189
}
@@ -204,8 +204,8 @@ public override DenseTensor<float> AddNoise(DenseTensor<float> originalSamples,
204204
float sqrtOneMinusAlpha = (float)Math.Sqrt(1.0f - alphaProd);
205205

206206
return noise
207-
.MultipleTensorByFloat(sqrtOneMinusAlpha)
208-
.AddTensors(originalSamples.MultipleTensorByFloat(sqrtAlpha));
207+
.MultiplyTensorByFloat(sqrtOneMinusAlpha)
208+
.AddTensors(originalSamples.MultiplyTensorByFloat(sqrtAlpha));
209209
}
210210

211211
/// <summary>

‎OnnxStack.StableDiffusion/Schedulers/StableDiffusion/EulerAncestralScheduler.cs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -127,9 +127,9 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
127127
.DivideTensorByFloat(sigma);
128128

129129
var delta = sigmaDown - sigma;
130-
var prevSample = sample.AddTensors(derivative.MultipleTensorByFloat(delta));
130+
var prevSample = sample.AddTensors(derivative.MultiplyTensorByFloat(delta));
131131
var noise = CreateRandomSample(prevSample.Dimensions);
132-
prevSample = prevSample.AddTensors(noise.MultipleTensorByFloat(sigmaUp));
132+
prevSample = prevSample.AddTensors(noise.MultiplyTensorByFloat(sigmaUp));
133133
return new SchedulerStepResult(prevSample);
134134
}
135135

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

154154
return noise
155-
.MultipleTensorByFloat(sigma)
155+
.MultiplyTensorByFloat(sigma)
156156
.AddTensors(originalSamples);
157157
}
158158

‎OnnxStack.StableDiffusion/Schedulers/StableDiffusion/EulerScheduler.cs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -121,17 +121,17 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
121121

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

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

130130

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

136136

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

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

146146

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

162162
return noise
163-
.MultipleTensorByFloat(sigma)
163+
.MultiplyTensorByFloat(sigma)
164164
.AddTensors(originalSamples);
165165
}
166166

‎OnnxStack.StableDiffusion/Schedulers/StableDiffusion/KDPM2Scheduler.cs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -142,13 +142,13 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
142142
DenseTensor<float> predOriginalSample;
143143
if (Options.PredictionType == PredictionType.Epsilon)
144144
{
145-
predOriginalSample = sample.SubtractTensors(modelOutput.MultipleTensorByFloat(sigmaInput));
145+
predOriginalSample = sample.SubtractTensors(modelOutput.MultiplyTensorByFloat(sigmaInput));
146146
}
147147
else if (Options.PredictionType == PredictionType.VariablePrediction)
148148
{
149149
var sigmaSqrt = (float)Math.Sqrt(sigmaInput * sigmaInput + 1f);
150150
predOriginalSample = sample.DivideTensorByFloat(sigmaSqrt)
151-
.AddTensors(modelOutput.MultipleTensorByFloat(-sigmaInput / sigmaSqrt));
151+
.AddTensors(modelOutput.MultiplyTensorByFloat(-sigmaInput / sigmaSqrt));
152152
}
153153
else
154154
{
@@ -177,7 +177,7 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
177177
}
178178

179179
_stepIndex += 1;
180-
return new SchedulerStepResult(sample.AddTensors(derivative.MultipleTensorByFloat(dt)));
180+
return new SchedulerStepResult(sample.AddTensors(derivative.MultiplyTensorByFloat(dt)));
181181
}
182182

183183

@@ -192,7 +192,7 @@ public override DenseTensor<float> AddNoise(DenseTensor<float> originalSamples,
192192
{
193193
var sigma = _sigmas[_stepIndex];
194194
return noise
195-
.MultipleTensorByFloat(sigma)
195+
.MultiplyTensorByFloat(sigma)
196196
.AddTensors(originalSamples);
197197
}
198198

‎OnnxStack.StableDiffusion/Schedulers/StableDiffusion/LMSScheduler.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,7 @@ public override SchedulerStepResult Step(DenseTensor<float> modelOutput, int tim
144144
{
145145
// Multiply to coeff by each derivatives to create the new tensors
146146
var (lmsCoeff, derivative) = lmsCoeffsAndDerivatives[i];
147-
lmsDerProduct[i] = derivative.MultipleTensorByFloat((float)lmsCoeff);
147+
lmsDerProduct[i] = derivative.MultiplyTensorByFloat((float)lmsCoeff);
148148
}
149149

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

170170
return noise
171-
.MultipleTensorByFloat(sigma)
171+
.MultiplyTensorByFloat(sigma)
172172
.AddTensors(originalSamples);
173173
}
174174

0 commit comments

Comments
 (0)