Fix issues found during testing

TensorStack-AI · Apr 25, 2024 · 6626eb5 · 6626eb5
1 parent d54393e
commit 6626eb5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/OnnxStack.StableDiffusion/Pipelines/StableDiffusionXLPipeline.cs b/OnnxStack.StableDiffusion/Pipelines/StableDiffusionXLPipeline.cs
@@ -293,7 +293,7 @@ private async Task<PromptEmbeddingsResult> GenerateEmbedsAsync(TokenizerResult i
             }
 
             var promptTensor = new DenseTensor<float>(promptEmbeddings.ToArray(), new[] { 1, promptEmbeddings.Count / _tokenizer2.TokenizerLength, _tokenizer2.TokenizerLength });
-            var pooledTensor = new DenseTensor<float>(pooledPromptEmbeddings.ToArray(), new[] { 1, pooledPromptEmbeddings.Count });
+            var pooledTensor = new DenseTensor<float>(pooledPromptEmbeddings.Take(_tokenizer2.TokenizerLength).ToArray(), new[] { 1, _tokenizer2.TokenizerLength });
             return new PromptEmbeddingsResult(promptTensor, pooledTensor);
         }