Remove legacy batch tensor duplication code

saddam213 · saddam213 · commit 7cc75270fc5a · 2023-11-14T08:11:46.000+13:00
diff --git a/OnnxStack.Console/appsettings.json b/OnnxStack.Console/appsettings.json
@@ -26,7 +26,7 @@
 				"InterOpNumThreads": 0,
 				"IntraOpNumThreads": 0,
 				"ExecutionMode": "ORT_SEQUENTIAL",
-				"ExecutionProvider": "Cuda",
+				"ExecutionProvider": "DirectML",
 				"ModelConfigurations": [
 					{
 						"Type": "Tokenizer",
diff --git a/OnnxStack.StableDiffusion/Config/PromptOptions.cs b/OnnxStack.StableDiffusion/Config/PromptOptions.cs
@@ -15,8 +15,6 @@ public class PromptOptions
         [StringLength(512)]
         public string NegativePrompt { get; set; }
 
-        public int BatchCount { get; set; } = 1; // Delete  Me
-
         public InputImage InputImage { get; set; }
 
         public InputImage InputImageMask { get; set; }
diff --git a/OnnxStack.StableDiffusion/Diffusers/DiffuserBase.cs b/OnnxStack.StableDiffusion/Diffusers/DiffuserBase.cs
@@ -119,7 +119,7 @@ public virtual async Task<DenseTensor<float>> DiffuseAsync(IModelOptions modelOp
 
             return schedulerResult;
         }
- 
+
 
         /// <summary>
         /// Runs the stable diffusion batch loop
@@ -210,35 +210,23 @@ protected virtual async Task<DenseTensor<float>> DecodeLatents(IModelOptions mod
             // Scale and decode the image latents with vae.
             latents = latents.MultiplyBy(1.0f / model.ScaleFactor);
 
-            var images = prompt.BatchCount > 1
-                ? latents.Split(prompt.BatchCount)
-                : new[] { latents };
-            var imageTensors = new List<DenseTensor<float>>();
-            foreach (var image in images)
-            {
-                var inputNames = _onnxModelService.GetInputNames(model, OnnxModelType.VaeDecoder);
-                var outputNames = _onnxModelService.GetOutputNames(model, OnnxModelType.VaeDecoder);
+            var inputNames = _onnxModelService.GetInputNames(model, OnnxModelType.VaeDecoder);
+            var outputNames = _onnxModelService.GetOutputNames(model, OnnxModelType.VaeDecoder);
 
-                var outputDim = new[] { 1, 3, options.Height, options.Width };
-                var outputBuffer = new DenseTensor<float>(outputDim);
-                using (var inputTensorValue = image.ToOrtValue())
-                using (var outputTensorValue = outputBuffer.ToOrtValue())
+            var outputDim = new[] { 1, 3, options.Height, options.Width };
+            var outputBuffer = new DenseTensor<float>(outputDim);
+            using (var inputTensorValue = latents.ToOrtValue())
+            using (var outputTensorValue = outputBuffer.ToOrtValue())
+            {
+                var inputs = new Dictionary<string, OrtValue> { { inputNames[0], inputTensorValue } };
+                var outputs = new Dictionary<string, OrtValue> { { outputNames[0], outputTensorValue } };
+                var results = await _onnxModelService.RunInferenceAsync(model, OnnxModelType.VaeDecoder, inputs, outputs);
+                using (var imageResult = results.First())
                 {
-                    var inputs = new Dictionary<string, OrtValue> { { inputNames[0], inputTensorValue } };
-                    var outputs = new Dictionary<string, OrtValue> { { outputNames[0], outputTensorValue } };
-                    var results = await _onnxModelService.RunInferenceAsync(model, OnnxModelType.VaeDecoder, inputs, outputs);
-                    using (var imageResult = results.First())
-                    {
-                        imageTensors.Add(outputBuffer);
-                    }
+                    _logger?.LogEnd("End", timestamp);
+                    return outputBuffer;
                 }
             }
-
-            var result = prompt.BatchCount > 1
-                ? imageTensors.Join()
-                : imageTensors.FirstOrDefault();
-            _logger?.LogEnd("End", timestamp);
-            return result;
         }
 
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/ImageDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/ImageDiffuser.cs
@@ -76,11 +76,7 @@ protected override async Task<DenseTensor<float>> PrepareLatents(IModelOptions m
                        .Add(scheduler.CreateRandomSample(outputBuffer.Dimensions, options.InitialNoiseLevel))
                        .MultiplyBy(model.ScaleFactor);
 
-                    var noisySample = scheduler.AddNoise(scaledSample, scheduler.CreateRandomSample(scaledSample.Dimensions), timesteps);
-                    if (prompt.BatchCount > 1)
-                        return noisySample.Repeat(prompt.BatchCount);
-
-                    return noisySample;
+                    return scheduler.AddNoise(scaledSample, scheduler.CreateRandomSample(scaledSample.Dimensions), timesteps);
                 }
             }
         }
diff --git a/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/TextDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/LatentConsistency/TextDiffuser.cs
@@ -48,7 +48,7 @@ protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, ISc
         /// <returns></returns>
         protected override Task<DenseTensor<float>> PrepareLatents(IModelOptions model, PromptOptions prompt, SchedulerOptions options, IScheduler scheduler, IReadOnlyList<int> timesteps)
         {
-            return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(prompt.BatchCount), scheduler.InitNoiseSigma));
+            return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(), scheduler.InitNoiseSigma));
         }
     }
 }
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/AnimateDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/AnimateDiffuser.cs
@@ -51,7 +51,7 @@ protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, ISc
         /// <returns></returns>
         protected override Task<DenseTensor<float>> PrepareLatents(IModelOptions model, PromptOptions prompt, SchedulerOptions options, IScheduler scheduler, IReadOnlyList<int> timesteps)
         {
-            return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(prompt.BatchCount), scheduler.InitNoiseSigma));
+            return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(), scheduler.InitNoiseSigma));
         }
     }
 }
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/ImageDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/ImageDiffuser.cs
@@ -78,11 +78,7 @@ protected override async Task<DenseTensor<float>> PrepareLatents(IModelOptions m
                        .Add(scheduler.CreateRandomSample(outputBuffer.Dimensions, options.InitialNoiseLevel))
                        .MultiplyBy(model.ScaleFactor);
 
-                    var noisySample = scheduler.AddNoise(scaledSample, scheduler.CreateRandomSample(scaledSample.Dimensions), timesteps);
-                    if (prompt.BatchCount > 1)
-                        return noisySample.Repeat(prompt.BatchCount);
-
-                    return noisySample;
+                    return scheduler.AddNoise(scaledSample, scheduler.CreateRandomSample(scaledSample.Dimensions), timesteps);
                 }
             }
         }
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintDiffuser.cs
@@ -161,9 +161,6 @@ private DenseTensor<float> PrepareMask(IModelOptions modelOptions, PromptOptions
                 });
 
                 imageTensor = imageTensor.MultiplyBy(modelOptions.ScaleFactor);
-                if (promptOptions.BatchCount > 1)
-                    imageTensor = imageTensor.Repeat(promptOptions.BatchCount);
-
                 if (schedulerOptions.GuidanceScale > 1f)
                     imageTensor = imageTensor.Repeat(2);
 
@@ -232,9 +229,6 @@ private DenseTensor<float> PrepareImageMask(IModelOptions modelOptions, PromptOp
                 {
                     var sample = inferResult.FirstElementAs<DenseTensor<float>>();
                     var scaledSample = sample.MultiplyBy(modelOptions.ScaleFactor);
-                    if (promptOptions.BatchCount > 1)
-                        scaledSample = scaledSample.Repeat(promptOptions.BatchCount);
-
                     if (schedulerOptions.GuidanceScale > 1f)
                         scaledSample = scaledSample.Repeat(2);
 
@@ -267,7 +261,7 @@ protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, ISc
         /// <returns></returns>
         protected override Task<DenseTensor<float>> PrepareLatents(IModelOptions model, PromptOptions prompt, SchedulerOptions options, IScheduler scheduler, IReadOnlyList<int> timesteps)
         {
-            return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(prompt.BatchCount), scheduler.InitNoiseSigma));
+            return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(), scheduler.InitNoiseSigma));
         }
 
 
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintLegacyDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintLegacyDiffuser.cs
@@ -173,9 +173,6 @@ protected override async Task<DenseTensor<float>> PrepareLatents(IModelOptions m
                        .Add(scheduler.CreateRandomSample(outputBuffer.Dimensions, options.InitialNoiseLevel))
                        .MultiplyBy(model.ScaleFactor);
 
-                    if (prompt.BatchCount > 1)
-                        return scaledSample.Repeat(prompt.BatchCount);
-
                     return scaledSample;
                 }
             }
@@ -214,9 +211,6 @@ private DenseTensor<float> PrepareMask(IModelOptions modelOptions, PromptOptions
                     }
                 });
 
-                if (promptOptions.BatchCount > 1)
-                    return maskTensor.Repeat(promptOptions.BatchCount);
-
                 return maskTensor;
             }
         }
diff --git a/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/TextDiffuser.cs b/OnnxStack.StableDiffusion/Diffusers/StableDiffusion/TextDiffuser.cs
@@ -50,7 +50,7 @@ protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, ISc
         /// <returns></returns>
         protected override Task<DenseTensor<float>> PrepareLatents(IModelOptions model, PromptOptions prompt, SchedulerOptions options, IScheduler scheduler, IReadOnlyList<int> timesteps)
         {
-            return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(prompt.BatchCount), scheduler.InitNoiseSigma));
+            return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(), scheduler.InitNoiseSigma));
         }
     }
 }
diff --git a/OnnxStack.StableDiffusion/Services/PromptService.cs b/OnnxStack.StableDiffusion/Services/PromptService.cs
@@ -46,13 +46,6 @@ public async Task<DenseTensor<float>> CreatePromptAsync(IModelOptions model, Pro
             var promptEmbeddings = await GenerateEmbedsAsync(model, promptTokens, maxPromptTokenCount);
             var negativePromptEmbeddings = await GenerateEmbedsAsync(model, negativePromptTokens, maxPromptTokenCount);
 
-            // If we have a batch, repeat the prompt embeddings
-            if (promptOptions.BatchCount > 1)
-            {
-                promptEmbeddings = promptEmbeddings.Repeat(promptOptions.BatchCount);
-                negativePromptEmbeddings = negativePromptEmbeddings.Repeat(promptOptions.BatchCount);
-            }
-
             // If we are doing guided diffusion, concatenate the negative prompt embeddings
             // If not we ingore the negative prompt embeddings
             if (isGuidanceEnabled)
@@ -166,6 +159,5 @@ private static IReadOnlyCollection<NamedOnnxValue> CreateInputParameters(params
         {
             return parameters.ToList().AsReadOnly();
         }
-
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -26,7 +26,7 @@`
`26`	`26`	`"InterOpNumThreads": 0,`
`27`	`27`	`"IntraOpNumThreads": 0,`
`28`	`28`	`"ExecutionMode": "ORT_SEQUENTIAL",`
`29`		`- "ExecutionProvider": "Cuda",`
	`29`	`+ "ExecutionProvider": "DirectML",`
`30`	`30`	`"ModelConfigurations": [`
`31`	`31`	`{`
`32`	`32`	`"Type": "Tokenizer",`
Original file line number	Diff line number	Diff line change
`@@ -76,11 +76,7 @@ protected override async Task<DenseTensor<float>> PrepareLatents(IModelOptions m`
`76`	`76`	`.Add(scheduler.CreateRandomSample(outputBuffer.Dimensions, options.InitialNoiseLevel))`
`77`	`77`	`.MultiplyBy(model.ScaleFactor);`
`78`	`78`
`79`		`- var noisySample = scheduler.AddNoise(scaledSample, scheduler.CreateRandomSample(scaledSample.Dimensions), timesteps);`
`80`		`- if (prompt.BatchCount > 1)`
`81`		`- return noisySample.Repeat(prompt.BatchCount);`
`82`		`-`
`83`		`- return noisySample;`
	`79`	`+ return scheduler.AddNoise(scaledSample, scheduler.CreateRandomSample(scaledSample.Dimensions), timesteps);`
`84`	`80`	`}`
`85`	`81`	`}`
`86`	`82`	`}`
Original file line number	Diff line number	Diff line change
`@@ -48,7 +48,7 @@ protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, ISc`
`48`	`48`	`/// <returns></returns>`
`49`	`49`	`protected override Task<DenseTensor<float>> PrepareLatents(IModelOptions model, PromptOptions prompt, SchedulerOptions options, IScheduler scheduler, IReadOnlyList<int> timesteps)`
`50`	`50`	`{`
`51`		`- return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(prompt.BatchCount), scheduler.InitNoiseSigma));`
	`51`	`+ return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(), scheduler.InitNoiseSigma));`
`52`	`52`	`}`
`53`	`53`	`}`
`54`	`54`	`}`
Original file line number	Diff line number	Diff line change
`@@ -51,7 +51,7 @@ protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, ISc`
`51`	`51`	`/// <returns></returns>`
`52`	`52`	`protected override Task<DenseTensor<float>> PrepareLatents(IModelOptions model, PromptOptions prompt, SchedulerOptions options, IScheduler scheduler, IReadOnlyList<int> timesteps)`
`53`	`53`	`{`
`54`		`- return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(prompt.BatchCount), scheduler.InitNoiseSigma));`
	`54`	`+ return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(), scheduler.InitNoiseSigma));`
`55`	`55`	`}`
`56`	`56`	`}`
`57`	`57`	`}`
Original file line number	Diff line number	Diff line change
`@@ -78,11 +78,7 @@ protected override async Task<DenseTensor<float>> PrepareLatents(IModelOptions m`
`78`	`78`	`.Add(scheduler.CreateRandomSample(outputBuffer.Dimensions, options.InitialNoiseLevel))`
`79`	`79`	`.MultiplyBy(model.ScaleFactor);`
`80`	`80`
`81`		`- var noisySample = scheduler.AddNoise(scaledSample, scheduler.CreateRandomSample(scaledSample.Dimensions), timesteps);`
`82`		`- if (prompt.BatchCount > 1)`
`83`		`- return noisySample.Repeat(prompt.BatchCount);`
`84`		`-`
`85`		`- return noisySample;`
	`81`	`+ return scheduler.AddNoise(scaledSample, scheduler.CreateRandomSample(scaledSample.Dimensions), timesteps);`
`86`	`82`	`}`
`87`	`83`	`}`
`88`	`84`	`}`
Original file line number	Diff line number	Diff line change
`@@ -173,9 +173,6 @@ protected override async Task<DenseTensor<float>> PrepareLatents(IModelOptions m`
`173`	`173`	`.Add(scheduler.CreateRandomSample(outputBuffer.Dimensions, options.InitialNoiseLevel))`
`174`	`174`	`.MultiplyBy(model.ScaleFactor);`
`175`	`175`
`176`		`- if (prompt.BatchCount > 1)`
`177`		`- return scaledSample.Repeat(prompt.BatchCount);`
`178`		`-`
`179`	`176`	`return scaledSample;`
`180`	`177`	`}`
`181`	`178`	`}`
`@@ -214,9 +211,6 @@ private DenseTensor<float> PrepareMask(IModelOptions modelOptions, PromptOptions`
`214`	`211`	`}`
`215`	`212`	`});`
`216`	`213`
`217`		`- if (promptOptions.BatchCount > 1)`
`218`		`- return maskTensor.Repeat(promptOptions.BatchCount);`
`219`		`-`
`220`	`214`	`return maskTensor;`
`221`	`215`	`}`
`222`	`216`	`}`
Original file line number	Diff line number	Diff line change
`@@ -50,7 +50,7 @@ protected override IReadOnlyList<int> GetTimesteps(SchedulerOptions options, ISc`
`50`	`50`	`/// <returns></returns>`
`51`	`51`	`protected override Task<DenseTensor<float>> PrepareLatents(IModelOptions model, PromptOptions prompt, SchedulerOptions options, IScheduler scheduler, IReadOnlyList<int> timesteps)`
`52`	`52`	`{`
`53`		`- return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(prompt.BatchCount), scheduler.InitNoiseSigma));`
	`53`	`+ return Task.FromResult(scheduler.CreateRandomSample(options.GetScaledDimension(), scheduler.InitNoiseSigma));`
`54`	`54`	`}`
`55`	`55`	`}`
`56`	`56`	`}`
Original file line number	Diff line number	Diff line change
`@@ -46,13 +46,6 @@ public async Task<DenseTensor<float>> CreatePromptAsync(IModelOptions model, Pro`
`46`	`46`	`var promptEmbeddings = await GenerateEmbedsAsync(model, promptTokens, maxPromptTokenCount);`
`47`	`47`	`var negativePromptEmbeddings = await GenerateEmbedsAsync(model, negativePromptTokens, maxPromptTokenCount);`
`48`	`48`
`49`		`- // If we have a batch, repeat the prompt embeddings`
`50`		`- if (promptOptions.BatchCount > 1)`
`51`		`- {`
`52`		`- promptEmbeddings = promptEmbeddings.Repeat(promptOptions.BatchCount);`
`53`		`- negativePromptEmbeddings = negativePromptEmbeddings.Repeat(promptOptions.BatchCount);`
`54`		`- }`
`55`		`-`
`56`	`49`	`// If we are doing guided diffusion, concatenate the negative prompt embeddings`
`57`	`50`	`// If not we ingore the negative prompt embeddings`
`58`	`51`	`if (isGuidanceEnabled)`
`@@ -166,6 +159,5 @@ private static IReadOnlyCollection<NamedOnnxValue> CreateInputParameters(params`
`166`	`159`	`{`
`167`	`160`	`return parameters.ToList().AsReadOnly();`
`168`	`161`	`}`
`169`		`-`
`170`	`162`	`}`
`171`	`163`	`}`