Improve Florence2 segmentation, add missing merges.txt

Author: saddam213

TensorStack.TextGeneration/Pipelines/DecoderPipeline.cs

@@ -1,7 +1,6 @@
 // Copyright (c) TensorStack. All rights reserved.
 // Licensed under the Apache 2.0 License.
 
-using Google.Protobuf.WellKnownTypes;
 using System;
 using System.Collections.Generic;
 using System.Linq;
@@ -245,7 +244,7 @@ protected virtual async Task<Sequence[]> BeamSearchAsync(GenerateOptions options
         /// </summary>
         /// <param name="candidates">The sequences.</param>
         /// <param name="options">The options.</param>
-        protected IEnumerable<Sequence> GetSequenceCandidates(SequenceCollection candidates, GenerateOptions options)
+        protected virtual IEnumerable<Sequence> GetSequenceCandidates(SequenceCollection candidates, GenerateOptions options)
         {
             // TODO: Diversity Penalty
             _sequenceComparer.SetLength(options.DiversityLength);

TensorStack.TextGeneration/Pipelines/Florence/FlorenceOptions.cs

@@ -13,8 +13,14 @@ public record FlorenceOptions : GenerateOptions
         public CoordinateBox<float> Region { get; set; }
     }
 
+    public record FlorenceSearchOptions : FlorenceOptions;
+
+
+
+
     public enum TaskType
     {
+        NONE,
         OCR,
         OCR_WITH_REGION,
         CAPTION,

TensorStack.TextGeneration/Pipelines/Florence/FlorencePipeline.cs

@@ -2,8 +2,10 @@
 // Licensed under the Apache 2.0 License.
 using Microsoft.ML.OnnxRuntime;
 using System;
+using System.Collections.Generic;
 using System.IO;
 using System.Linq;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using TensorStack.Common;
@@ -15,7 +17,9 @@
 
 namespace TensorStack.TextGeneration.Pipelines.Florence
 {
-    public class FlorencePipeline : EncoderDecoderPipeline
+    public class FlorencePipeline : EncoderDecoderPipeline,
+         IPipeline<GenerateResult, FlorenceOptions>,
+         IPipelineStream<GenerateResult, FlorenceSearchOptions>
     {
         private readonly FlorenceConfig _configuration;
         private readonly PreProcessor _preProcessor;
@@ -103,6 +107,35 @@ public virtual async Task<GenerateResult> RunAsync(FlorenceOptions options, IPro
         }
 
 
+        public virtual async IAsyncEnumerable<GenerateResult> RunAsync(FlorenceSearchOptions options, IProgress<RunProgress> progressCallback = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            var textPrompt = _preProcessor.ProcessPrompt(options);
+            var imagePrompt = _preProcessor.ProcessImage(options);
+
+            TokenizerOutput = await Tokenizer.EncodeAsync(textPrompt);
+            var embedsOutput = await RunTextEmbedAsync(TokenizerOutput.InputIds);
+            _visionOutput = await RunVisionEncoderAsync(embedsOutput, imagePrompt);
+            EncoderOutput = await RunEncoderAsync();
+
+            var sequences = await BeamSearchAsync(options, cancellationToken);
+            foreach (var sequence in sequences)
+            {
+                using (sequence)
+                {
+                    var processedBeamOutput = _postProcessor.Process(options, sequence.Tokens);
+                    yield return new GenerateResult
+                    {
+                        Beam = sequence.Id,
+                        Score = sequence.Score,
+                        PenaltyScore = sequence.PenaltyScore,
+                        Result = processedBeamOutput.Result,
+                        CoordinateResults = processedBeamOutput.CoordinateResults
+                    };
+                }
+            }
+        }
+
+
         /// <summary>
         /// Run encoder model.
         /// </summary>
@@ -233,14 +266,20 @@ protected override void Dispose(bool disposing)
         /// <param name="embedModel">The embed model.</param>
         /// <param name="visionModel">The vision model.</param>
         /// <returns>FlorencePipeline.</returns>
-        public static FlorencePipeline Create(ExecutionProvider provider, string modelPath, string encoderModel = "encoder_model.onnx", string decoderModel = "decoder_model_merged.onnx", string embedModel = "embed_tokens.onnx", string visionModel = "vision_encoder.onnx")
+        public static FlorencePipeline Create(ExecutionProvider provider, string modelPath, FlorenceType modelType, string encoderModel = "encoder_model.onnx", string decoderModel = "decoder_model_merged.onnx", string embedModel = "embed_tokens.onnx", string visionModel = "vision_encoder.onnx")
         {
-            // Florence Large
-            //NumLayers = 12,
-            //NumHeads = 16,
-            //NumKVHeads = 16,
-            //HiddenSize = 768,
-            //VocabSize = 51289
+            var numLayers = 6;
+            var numHeads = 12;
+            var numKVHeads = 12;
+            var hiddenSize = 768;
+            var vocabSize = 51289;
+            if (modelType == FlorenceType.Large)
+            {
+                numLayers = 12;
+                numHeads = 16;
+                numKVHeads = 16;
+                hiddenSize = 1024;
+            }
 
             var config = new FlorenceConfig
             {
@@ -251,20 +290,20 @@ public static FlorencePipeline Create(ExecutionProvider provider, string modelPa
                 }),
                 EncoderConfig = new EncoderConfig
                 {
-                    NumLayers = 6,
-                    NumHeads = 12,
-                    NumKVHeads = 12,
-                    HiddenSize = 768,
-                    VocabSize = 51289,
+                    NumLayers = numLayers,
+                    NumHeads = numHeads,
+                    NumKVHeads = numKVHeads,
+                    HiddenSize = hiddenSize,
+                    VocabSize = vocabSize,
                     Path = Path.Combine(modelPath, encoderModel),
                 },
                 DecoderConfig = new DecoderConfig
                 {
-                    NumLayers = 6,
-                    NumHeads = 12,
-                    NumKVHeads = 12,
-                    HiddenSize = 768,
-                    VocabSize = 51289,
+                    NumLayers = numLayers,
+                    NumHeads = numHeads,
+                    NumKVHeads = numKVHeads,
+                    HiddenSize = hiddenSize,
+                    VocabSize = vocabSize,
                     Path = Path.Combine(modelPath, decoderModel),
                 },
                 EmbedsConfig = new ModelConfig
@@ -278,6 +317,7 @@ public static FlorencePipeline Create(ExecutionProvider provider, string modelPa
             };
 
             config.EncoderConfig.SetProvider(provider);
+            //config.DecoderConfig.SetProvider(provider);
             config.DecoderConfig.SetProvider(ProviderCPU()); // TODO
             config.EmbedsConfig.SetProvider(provider);
             config.VisionConfig.SetProvider(provider);
@@ -302,4 +342,10 @@ private static ExecutionProvider ProviderCPU()
             });
         }
     }
+
+    public enum FlorenceType
+    {
+        Base = 0,
+        Large = 1
+    }
 }

TensorStack.TextGeneration/Pipelines/Florence/FlorenceTokenizer.cs

@@ -21,6 +21,7 @@ public class FlorenceTokenizer : ITokenizer
         private readonly MapCollection<byte, char> _unicodeMap;
         private readonly MapCollection<long, int> _coordinateMap;
         private readonly MapCollection<long, string> _vocabularyMap;
+        private readonly Dictionary<MergeToken, int> _mergesMap;
         private readonly MapCollection<long, string> _specialTokensMap;
 
         /// <summary>
@@ -34,6 +35,7 @@ public FlorenceTokenizer(TokenizerConfig configuration)
             _specialTokensMap = CreateSpecialTokenMapping();
             _vocabularyMap = CreateVocabMapping();
             _coordinateMap = CreateCoordinateMapping();
+            _mergesMap = CreateMergesMapping();
             _preTokenizeRegex = new Regex(@"'s|'t|'re|'ve|'m|'ll|'d|<loc_[\p{L}\p{N}_]+>| ?[\p{L}_][\p{L}\p{N}_]*|[^ \s\p{L}\p{N}]+|\s+(?!\S)|\s+", RegexOptions.Compiled);
         }
 
@@ -214,6 +216,7 @@ private string TokensToString(IEnumerable<string> tokens)
         private long[] StringToTokens(ReadOnlySpan<char> input)
         {
             var tokens = PreTokenize(input)
+                .SelectMany(ApplyMerges)
                 .Select(TokenToId)
                 .Prepend(_configuration.BOS)
                 .Append(_configuration.EOS)
@@ -241,6 +244,42 @@ private string[] PreTokenize(ReadOnlySpan<char> input)
         }
 
 
+        /// <summary>
+        /// Applies the merges.
+        /// </summary>
+        /// <param name="token">The token.</param>
+        private List<string> ApplyMerges(string token)
+        {
+            if (_specialTokensMap.ContainsKey(token))
+                return [token];
+
+            var symbols = token.Select(c => c.ToString()).ToList();
+            while (symbols.Count > 1)
+            {
+                int bestIndex = -1;
+                int bestRank = int.MaxValue;
+                for (int i = 0; i < symbols.Count - 1; i++)
+                {
+                    var pair = new MergeToken(symbols[i], symbols[i + 1]);
+                    if (_mergesMap.TryGetValue(pair, out int rank) && rank < bestRank)
+                    {
+                        bestRank = rank;
+                        bestIndex = i;
+                    }
+                }
+
+                if (bestIndex == -1)
+                    break;
+
+                // Merge the pair
+                string merged = symbols[bestIndex] + symbols[bestIndex + 1];
+                symbols[bestIndex] = merged;
+                symbols.RemoveAt(bestIndex + 1);
+            }
+            return symbols;
+        }
+
+
         /// <summary>
         /// Parses the coordinate.
         /// </summary>
@@ -337,6 +376,22 @@ private MapCollection<long, string> CreateVocabMapping()
         }
 
 
+        /// <summary>
+        /// Creates the merges mapping.
+        /// </summary>
+        private Dictionary<MergeToken, int> CreateMergesMapping()
+        {
+            var mergesFile = Path.Combine(_configuration.Path, "merges.txt");
+            return File.ReadLines(mergesFile)
+                .Where(line => !string.IsNullOrWhiteSpace(line) && !line.StartsWith("#"))
+                .Select((line, index) =>
+                {
+                    var parts = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);
+                    return (new MergeToken(parts[0], parts[1]), index);
+                }).ToDictionary(x => x.Item1, x => x.index);
+        }
+
+
         /// <summary>
         /// Creates the coordinate mapping.
         /// </summary>
@@ -374,5 +429,8 @@ private record AddedTokenJson
             [JsonPropertyName("content")]
             public string Content { get; set; }
         }
+
+        private record MergeToken(string PartA, string PartB);
+
     }
 }

TensorStack.TextGeneration/Pipelines/Florence/PostProcessor.cs

@@ -133,9 +133,11 @@ public List<CoordinateResult> ExtractCoordinatePolygons(IReadOnlyCollection<long
             var currentLabel = new List<long>();
             var currentCoordinates = new List<int>();
             var results = new List<CoordinateResult>();
+            var index = 0;
             foreach (var tokenId in tokens)
             {
-                if (!_tokenizer.TryGetCoordinate(tokenId, out var coordinate))
+                index++;
+                if (!_tokenizer.TryGetCoordinate(tokenId, out var coordinate) || index == tokens.Count)
                 {
                     if (currentCoordinates.Count > 0)
                     {
@@ -149,6 +151,9 @@ public List<CoordinateResult> ExtractCoordinatePolygons(IReadOnlyCollection<long
                             coordinates.Add(new Coordinate<int>(position[0], position[1]));
                         }
 
+                        if (coordinates.Count == 0)
+                            return new List<CoordinateResult>();
+
                         var scaledCoordinates = _coordinateScaler.ScaleUp(coordinates.ToArray(), sourceImage);
                         var coordinateBox = new CoordinateBox<int>
                         (

TensorStack.TextGeneration/Pipelines/Florence/PreProcessor.cs

@@ -43,7 +43,7 @@ public string ProcessPrompt(FlorenceOptions options)
                 TaskType.REGION_PROPOSAL => "Locate the region proposals in the image.",
                 TaskType.CAPTION_TO_PHRASE_GROUNDING => $"Locate the phrases in the caption: {options.Prompt}",
                 TaskType.REFERRING_EXPRESSION_SEGMENTATION => $"Locate {options.Prompt} in the image with mask",
-                TaskType.REGION_TO_SEGMENTATION => $"What is the mask of region {regionTokens}",
+                TaskType.REGION_TO_SEGMENTATION => $"What is the polygon mask of region {regionTokens}",
                 TaskType.OPEN_VOCABULARY_DETECTION => $"Locate {options.Prompt} in the image.",
                 TaskType.REGION_TO_CATEGORY => $"What is the region {regionTokens}?",
                 TaskType.REGION_TO_DESCRIPTION => $"What does the region {regionTokens} describe?",

TensorStack.TextGeneration/Processing/KVCacheEncoderDecoder.cs

@@ -84,6 +84,9 @@ public void Update(OrtValue[] currentValues, bool useBranchCache)
             {
                 if (i % 4 == 0)
                 {
+
+                    // TODO: Allocate entire Maxlength and update the buffer
+
                     _values[i].Dispose();
                     _values[i + 1].Dispose();
 

TensorStack.TextGeneration/Processing/Logit/BOSLogitsProcessor.cs

@@ -27,6 +27,7 @@ public void Process(List<long> inputs, Tensor<float> logits)
         {
             if (inputs.Count == 0)
             {
+                inputs.Add(_bosTokenId);
                 logits.Fill(float.NegativeInfinity);
                 logits[0, 0] = float.NegativeZero;
             }