Get/Set prompt cache

Author: saddam213

TensorStack.StableDiffusion/Pipelines/Flux/FluxBase.cs

@@ -148,6 +148,10 @@ protected override void ValidateOptions(GenerateOptions options)
         /// <param name="cancellationToken">The cancellation token.</param>
         protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, CancellationToken cancellationToken = default)
         {
+            var cachedPrompt = GetPromptCache(options);
+            if (cachedPrompt is not null)
+                return cachedPrompt;
+
             // Tokenize2
             var promptTokens = await TokenizePromptAsync(options.Prompt, cancellationToken);
             var negativePromptTokens = await TokenizePromptAsync(options.NegativePrompt, cancellationToken);
@@ -179,7 +183,7 @@ protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, C
             var negativePromptPooledEmbeds = negativePromptEmbeddings.TextEmbeds;
             negativePromptPooledEmbeds = negativePromptPooledEmbeds.Reshape([negativePromptPooledEmbeds.Dimensions[^2], negativePromptPooledEmbeds.Dimensions[^1]]).FirstBatch();
 
-            return new PromptResult(promptEmbeds, promptPooledEmbeds, negativePromptEmbeds, negativePromptPooledEmbeds);
+            return SetPromptCache(options, new PromptResult(promptEmbeds, promptPooledEmbeds, negativePromptEmbeds, negativePromptPooledEmbeds));
         }
 
 

TensorStack.StableDiffusion/Pipelines/Nitro/NitroBase.cs

@@ -137,6 +137,10 @@ protected override void ValidateOptions(GenerateOptions options)
         /// <param name="cancellationToken">The cancellation token.</param>
         protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, CancellationToken cancellationToken = default)
         {
+            var cachedPrompt = GetPromptCache(options);
+            if (cachedPrompt is not null)
+                return cachedPrompt;
+
             // Conditional Prompt
             var promptEmbeds = await TextEncoder.GetLastHiddenState(new TextGeneration.Common.GenerateOptions
             {
@@ -159,7 +163,7 @@ protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, C
                 }, cancellationToken);
             }
 
-            return new PromptResult(promptEmbeds, default, negativePromptEmbeds, default);
+            return SetPromptCache(options, new PromptResult(promptEmbeds, default, negativePromptEmbeds, default));
         }
 
 

TensorStack.StableDiffusion/Pipelines/StableCascade/StableCascadeBase.cs

@@ -147,6 +147,10 @@ protected override void ValidateOptions(GenerateOptions options)
         /// <param name="cancellationToken">The cancellation token.</param>
         protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, CancellationToken cancellationToken = default)
         {
+            var cachedPrompt = GetPromptCache(options);
+            if (cachedPrompt is not null)
+                return cachedPrompt;
+
             // Tokenizer
             var promptTokens = await TokenizePromptAsync(options.Prompt, cancellationToken);
             var negativePromptTokens = await TokenizePromptAsync(options.NegativePrompt, cancellationToken);
@@ -167,7 +171,7 @@ protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, C
                ? negativePromptEmbeddings.TextEmbeds
                : negativePromptEmbeddings.TextEmbeds.Reshape([1, .. negativePromptEmbeddings.TextEmbeds.Dimensions]);
 
-            return new PromptResult(promptEmbeddings.HiddenStates, textEmbeds, negativePromptEmbeddings.HiddenStates, negativeTextEmbeds);
+            return SetPromptCache(options, new PromptResult(promptEmbeddings.HiddenStates, textEmbeds, negativePromptEmbeddings.HiddenStates, negativeTextEmbeds));
         }
 
 

TensorStack.StableDiffusion/Pipelines/StableDiffusion/StableDiffusionBase.cs

@@ -146,6 +146,10 @@ protected override void ValidateOptions(GenerateOptions options)
         /// <param name="cancellationToken">The cancellation token.</param>
         protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, CancellationToken cancellationToken = default)
         {
+            var cachedPrompt = GetPromptCache(options);
+            if (cachedPrompt is not null)
+                return cachedPrompt;
+
             // Tokenizer
             var promptTokens = await TokenizePromptAsync(options.Prompt, cancellationToken);
             var negativePromptTokens = await TokenizePromptAsync(options.NegativePrompt, cancellationToken);
@@ -157,7 +161,7 @@ protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, C
             if (options.IsLowMemoryEnabled || options.IsLowMemoryTextEncoderEnabled)
                 await TextEncoder.UnloadAsync();
 
-            return new PromptResult(promptEmbeddings.HiddenStates, promptEmbeddings.TextEmbeds, negativePromptEmbeddings.HiddenStates, negativePromptEmbeddings.TextEmbeds);
+            return SetPromptCache(options, new PromptResult(promptEmbeddings.HiddenStates, promptEmbeddings.TextEmbeds, negativePromptEmbeddings.HiddenStates, negativePromptEmbeddings.TextEmbeds));
         }
 
 

TensorStack.StableDiffusion/Pipelines/StableDiffusion3/StableDiffusion3Base.cs

@@ -165,6 +165,10 @@ protected override void ValidateOptions(GenerateOptions options)
         /// <param name="cancellationToken">The cancellation token.</param>
         protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, CancellationToken cancellationToken = default)
         {
+            var cachedPrompt = GetPromptCache(options);
+            if (cachedPrompt is not null)
+                return cachedPrompt;
+
             // Tokenizer
             var promptTokens = await TokenizePromptAsync(options.Prompt, cancellationToken);
             var negativePromptTokens = await TokenizePromptAsync(options.NegativePrompt, cancellationToken);
@@ -217,7 +221,7 @@ protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, C
             negativePromptPooledEmbeds = negativePromptPooledEmbeds.Reshape([negativePromptPooledEmbeds.Dimensions[^2], negativePromptPooledEmbeds.Dimensions[^1]]).FirstBatch();
             negativePromptPooledEmbeds = negativePromptPooledEmbeds.Concatenate(negativePromptPooledEmbeds2, 1);
 
-            return new PromptResult(promptEmbeds, promptPooledEmbeds, negativePromptEmbeds, negativePromptPooledEmbeds);
+            return SetPromptCache(options, new PromptResult(promptEmbeds, promptPooledEmbeds, negativePromptEmbeds, negativePromptPooledEmbeds));
         }
 
 

TensorStack.StableDiffusion/Pipelines/StableDiffusionXL/StableDiffusionXLBase.cs

@@ -146,6 +146,10 @@ protected override void ValidateOptions(GenerateOptions options)
         /// <param name="cancellationToken">The cancellation token.</param>
         protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, CancellationToken cancellationToken = default)
         {
+            var cachedPrompt = GetPromptCache(options);
+            if (cachedPrompt is not null)
+                return cachedPrompt;
+
             // Tokenizer
             var promptTokens = await TokenizePromptAsync(options.Prompt, cancellationToken);
             var negativePromptTokens = await TokenizePromptAsync(options.NegativePrompt, cancellationToken);
@@ -176,7 +180,7 @@ protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, C
             var pooledNegativePromptEmbeds = negativePrompt2Embeddings.TextEmbeds;
             var negativePromptEmbeddings = negativePrompt1Embeddings.HiddenStates.Concatenate(negativePrompt2Embeddings.HiddenStates, 2);
 
-            return new PromptResult(promptEmbeddings, pooledPromptEmbeds, negativePromptEmbeddings, pooledNegativePromptEmbeds);
+            return SetPromptCache(options, new PromptResult(promptEmbeddings, pooledPromptEmbeds, negativePromptEmbeddings, pooledNegativePromptEmbeds));
         }