Nitro TextEncoder

Author: saddam213

TensorStack.StableDiffusion/Pipelines/Nitro/NitroBase.cs

@@ -1,7 +1,6 @@
 // Copyright (c) TensorStack. All rights reserved.
 // Licensed under the Apache 2.0 License.
 using Microsoft.Extensions.Logging;
-using Microsoft.ML.Tokenizers;
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
@@ -21,7 +20,7 @@ namespace TensorStack.StableDiffusion.Pipelines.Nitro
     public abstract class NitroBase : PipelineBase
     {
         /// <summary>
-        /// Initializes a new instance of the <see cref="vBase"/> class.
+        /// Initializes a new instance of the <see cref="NitroBase"/> class.
         /// </summary>
         /// <param name="transformer">The transformer.</param>
         /// <param name="textEncoder">The text encoder.</param>
@@ -46,6 +45,7 @@ public NitroBase(NitroConfig configuration, ILogger logger = default) : this(
             new TransformerNitroModel(configuration.Transformer),
             new LlamaPipeline(new LlamaConfig
             {
+                OutputLastHiddenStates = true,
                 DecoderConfig = configuration.TextEncoder,
                 Tokenizer = new BPETokenizer(configuration.Tokenizer),
             }),
@@ -128,46 +128,28 @@ protected override void ValidateOptions(GenerateOptions options)
         /// <param name="cancellationToken">The cancellation token.</param>
         protected async Task<PromptResult> CreatePromptAsync(IPipelineOptions options, CancellationToken cancellationToken = default)
         {
-            //// Tokenize2
-            //var promptTokens = await TokenizePromptAsync(options.Prompt, cancellationToken);
-            //var negativePromptTokens = await TokenizePromptAsync(options.NegativePrompt, cancellationToken);
-            //var maxTokenLength = (int)Math.Max(promptTokens.InputIds.Length, negativePromptTokens.InputIds.Length);
-
-            //// Tokenizer2
-            //var prompt2Tokens = await TokenizePrompt2Async(options.Prompt, cancellationToken);
-            //var negativePrompt2Tokens = await TokenizePrompt2Async(options.NegativePrompt, cancellationToken);
-
-            //// TextEncoder
-            //var promptEmbeddings = await EncodePromptAsync(promptTokens, maxTokenLength, cancellationToken);
-            //var negativePromptEmbeddings = await EncodePromptAsync(negativePromptTokens, maxTokenLength, cancellationToken);
-            //if (options.IsLowMemoryEnabled || options.IsLowMemoryTextEncoderEnabled)
-            //    await TextEncoder.UnloadAsync();
-
-
-
-            //// Prompt
-            //var promptEmbeds = prompt2Embeddings.HiddenStates;
-            //var promptPooledEmbeds = promptEmbeddings.TextEmbeds;
-            //promptPooledEmbeds = promptPooledEmbeds.Reshape([promptPooledEmbeds.Dimensions[^2], promptPooledEmbeds.Dimensions[^1]]).FirstBatch();
-
-            //// Negative promt
-            //var negativePromptEmbeds = negativePrompt2Embeddings.HiddenStates;
-            //var negativePromptPooledEmbeds = negativePromptEmbeddings.TextEmbeds;
-            //negativePromptPooledEmbeds = negativePromptPooledEmbeds.Reshape([negativePromptPooledEmbeds.Dimensions[^2], negativePromptPooledEmbeds.Dimensions[^1]]).FirstBatch();
-
-            //return new PromptResult(promptEmbeds, promptPooledEmbeds, negativePromptEmbeds, negativePromptPooledEmbeds);
-
-
-            var result = TextEncoder.RunAsync(new TextGeneration.Common.GenerateOptions
+            // Conditional Prompt
+            var promptEmbeds = await TextEncoder.GetLastHiddenState(new TextGeneration.Common.GenerateOptions
             {
                 Seed = options.Seed,
-                Prompt = options.Prompt
-            });
-
-            return default;
-        }
+                Prompt = options.Prompt,
+                MaxLength = 128
+            }, cancellationToken);
 
+            // Unconditional prompt
+            var negativePromptEmbeds = default(Tensor<float>);
+            if (!string.IsNullOrEmpty(options.NegativePrompt))
+            {
+                negativePromptEmbeds = await TextEncoder.GetLastHiddenState(new TextGeneration.Common.GenerateOptions
+                {
+                    Seed = options.Seed,
+                    Prompt = options.NegativePrompt,
+                    MaxLength = 128
+                }, cancellationToken);
+            }
 
+            return new PromptResult(promptEmbeds, default, negativePromptEmbeds, default);
+        }
 
 
         /// <summary>
@@ -372,7 +354,7 @@ protected override GenerateOptions ConfigureDefaultOptions()
                 Scheduler = SchedulerType.FlowMatchEulerDiscrete
             };
 
-            // SD3-Turbo Models , 4 Steps, No Guidance
+            // Nitro-Distilled Models ,4 Steps, No Guidance
             if (Transformer.ModelType == ModelType.Turbo)
             {
                 return options with

TensorStack.StableDiffusion/Pipelines/Nitro/NitroConfig.cs

@@ -5,7 +5,6 @@
 using TensorStack.StableDiffusion.Config;
 using TensorStack.StableDiffusion.Enums;
 using TensorStack.TextGeneration.Common;
-using TensorStack.TextGeneration.Pipelines.Llama;
 using TensorStack.TextGeneration.Tokenizers;
 
 namespace TensorStack.StableDiffusion.Pipelines.Nitro
@@ -26,8 +25,8 @@ public NitroConfig()
             {
                 NumHeads = 32,
                 NumLayers = 16,
-                HiddenSize = 2048,
                 NumKVHeads = 8,
+                HiddenSize = 2048,
                 VocabSize = 128256
             };
             Transformer = new TransformerModelConfig
@@ -40,8 +39,7 @@ public NitroConfig()
             AutoEncoder = new AutoEncoderModelConfig
             {
                 LatentChannels = 32,
-                ScaleFactor = 0.3611f,
-                ShiftFactor = 0.1159f
+                ScaleFactor = 0.41407f
             };
         }
 
@@ -116,7 +114,7 @@ public static NitroConfig FromFile(string configFile, ExecutionProvider executio
         public static NitroConfig FromFolder(string modelFolder, ModelType modelType, ExecutionProvider executionProvider = default)
         {
             var config = FromDefault(Path.GetFileNameWithoutExtension(modelFolder), modelType, executionProvider);
-            config.Tokenizer.Path = Path.Combine(modelFolder, "tokenizer", "vocab.json");
+            config.Tokenizer.Path = Path.Combine(modelFolder, "tokenizer");
             config.TextEncoder.Path = Path.Combine(modelFolder, "text_encoder", "model.onnx");
             config.Transformer.Path = Path.Combine(modelFolder, "transformer", "model.onnx");
             config.AutoEncoder.DecoderModelPath = Path.Combine(modelFolder, "vae_decoder", "model.onnx");

TensorStack.StableDiffusion/Pipelines/Nitro/NitroPipeline.cs

@@ -44,7 +44,6 @@ public NitroPipeline(NitroConfig configuration, ILogger logger = null)
         public async Task<ImageTensor> RunAsync(GenerateOptions options, IProgress<GenerateProgress> progressCallback = null, CancellationToken cancellationToken = default)
         {
             ValidateOptions(options);
-
             var prompt = await CreatePromptAsync(options, cancellationToken);
             using (var scheduler = CreateScheduler(options))
             {

TensorStack.TextGeneration/Pipelines/Llama/LlamaPipeline.cs

@@ -89,6 +89,21 @@ public async Task<GenerateResult[]> RunAsync(SearchOptions options, IProgress<Ge
         }
 
 
+        /// <summary>
+        /// Gets the LastHiddenState.
+        /// </summary>
+        /// <param name="options">The options.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        public async Task<Tensor<float>> GetLastHiddenState(GenerateOptions options, CancellationToken cancellationToken = default)
+        {
+            await TokenizePromptAsync(options);
+            using (var sequence = await InitializeAsync(options))
+            {
+                return sequence.LastHiddenState;
+            }
+        }
+
+
         /// <summary>
         /// Gets the token processors.
         /// </summary>