feat: implement the complete bpe function (#119)

* implement the complete bpe function
---------

Co-authored-by: leejet <leejet714@gmail.com>

.clang-format

@@ -7,7 +7,7 @@ IndentCaseLabels: false
 ColumnLimit: 0
 AccessModifierOffset: -4
 NamespaceIndentation: All
-FixNamespaceComments: false 
+FixNamespaceComments: false
 AlignAfterOpenBracket: true
 AlignConsecutiveAssignments: true
 IndentCaseLabels: true

README.md

@@ -11,6 +11,7 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
 - Plain C/C++ implementation based on [ggml](https://github.com/ggerganov/ggml), working in the same way as [llama.cpp](https://github.com/ggerganov/llama.cpp)
 - Super lightweight and without external dependencies
 - SD1.x and SD2.x support
+- [SD-Turbo](https://huggingface.co/stabilityai/sd-turbo) support
 - 16-bit, 32-bit float support
 - 4-bit, 5-bit and 8-bit integer quantization support
 - Accelerated memory-efficient CPU inference
@@ -50,7 +51,6 @@ Inference of [Stable Diffusion](https://github.com/CompVis/stable-diffusion) in
     - The current implementation of ggml_conv_2d is slow and has high memory usage
     - Implement Winograd Convolution 2D for 3x3 kernel filtering
 - [ ] Continuing to reduce memory usage (quantizing the weights of ggml_conv_2d)
-- [ ] Implement BPE Tokenizer
 - [ ] Implement [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN/tree/master) upscaler
 - [ ] k-quants support
 

model.cpp

@@ -54,7 +54,6 @@ std::string self_attn_names[] = {
     "self_attn.q_proj.weight",
     "self_attn.k_proj.weight",
     "self_attn.v_proj.weight",
-
     "self_attn.q_proj.bias",
     "self_attn.k_proj.bias",
     "self_attn.v_proj.bias",
@@ -75,6 +74,8 @@ const char* unused_tensors[] = {
     "cond_stage_model.transformer.text_model.embeddings.position_ids",
     "cond_stage_model.model.logit_scale",
     "cond_stage_model.model.text_projection",
+    "conditioner.embedders.0.model.logit_scale",
+    "conditioner.embedders.0.model.text_projection",
     "model.diffusion_model.time_embedding.cond_proj.weight",
     "unet.time_embedding.cond_proj.weight",
     "model_ema.decay",
@@ -82,6 +83,7 @@ const char* unused_tensors[] = {
     "model_ema.diffusion_model",
     "control_model",
     "embedding_manager",
+    "denoiser.sigmas",
 };
 
 bool is_unused_tensor(std::string name) {
@@ -126,16 +128,19 @@ std::unordered_map<std::string, std::string> vae_decoder_name_map = {
 };
 
 std::string convert_open_clip_to_hf_clip(const std::string& name) {
-    std::string new_name                  = name;
+    std::string new_name = name;
+    if (starts_with(new_name, "conditioner.embedders.0.")) {
+        new_name = "cond_stage_model." + new_name.substr(strlen("conditioner.embedders.0."));
+    }
     std::string open_clip_resblock_prefix = "cond_stage_model.model.transformer.resblocks.";
     std::string hf_clip_resblock_prefix   = "cond_stage_model.transformer.text_model.encoder.layers.";
 
-    if (open_clip_to_hf_clip_model.find(name) != open_clip_to_hf_clip_model.end()) {
-        new_name = open_clip_to_hf_clip_model[name];
+    if (open_clip_to_hf_clip_model.find(new_name) != open_clip_to_hf_clip_model.end()) {
+        new_name = open_clip_to_hf_clip_model[new_name];
     }
 
-    if (name.find(open_clip_resblock_prefix) == 0) {
-        std::string remain = name.substr(open_clip_resblock_prefix.length());
+    if (new_name.find(open_clip_resblock_prefix) == 0) {
+        std::string remain = new_name.substr(open_clip_resblock_prefix.length());
         std::string idx    = remain.substr(0, remain.find("."));
         std::string suffix = remain.substr(idx.length() + 1);
 
@@ -349,7 +354,7 @@ std::string convert_diffusers_name_to_compvis(const std::string& key, char seq)
 
 std::string convert_tensor_name(const std::string& name) {
     std::string new_name;
-    if (starts_with(name, "cond_stage_model.model")) {
+    if (starts_with(name, "cond_stage_model.model") || starts_with(name, "conditioner.embedders.0.model")) {
         new_name = convert_open_clip_to_hf_clip(name);
     } else if (starts_with(name, "first_stage_model.decoder")) {
         new_name = convert_vae_decoder_name(name);
@@ -1159,7 +1164,8 @@ SDVersion ModelLoader::get_sd_version() {
         if (tensor_storage.name == "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight" ||
             tensor_storage.name == "cond_stage_model.model.token_embedding.weight" ||
             tensor_storage.name == "text_model.embeddings.token_embedding.weight" ||
-            tensor_storage.name == "te.text_model.embeddings.token_embedding.weight") {
+            tensor_storage.name == "te.text_model.embeddings.token_embedding.weight" ||
+            tensor_storage.name == "conditioner.embedders.0.model.token_embedding.weight") {
             token_embedding_weight = tensor_storage;
             break;
         }
@@ -1186,20 +1192,9 @@ ggml_type ModelLoader::get_sd_wtype() {
     return GGML_TYPE_COUNT;
 }
 
-bool ModelLoader::load_vocab(on_new_token_cb_t on_new_token_cb) {
-    char* vocab_buffer          = reinterpret_cast<char*>(vocab_json);
-    nlohmann::json vocab        = nlohmann::json::parse(vocab_buffer);
-    std::map<char, int> decoder = unicode_to_byte();
-    for (auto& it : vocab.items()) {
-        int token_id          = it.value();
-        std::string token_str = it.key();
-        std::string token     = "";
-        for (char c : token_str) {
-            token += decoder[c];
-        }
-        on_new_token_cb(token, token_id);
-    }
-    return true;
+std::string ModelLoader::load_merges() {
+    std::string merges_utf8_str(reinterpret_cast<const char*>(merges_utf8_c_str), sizeof(merges_utf8_c_str));
+    return merges_utf8_str;
 }
 
 bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb) {

model.h

@@ -115,7 +115,7 @@ public:
     bool init_from_file(const std::string& file_path, const std::string& prefix = "");
     SDVersion get_sd_version();
     ggml_type get_sd_wtype();
-    bool load_vocab(on_new_token_cb_t on_new_token_cb);
+    std::string load_merges();
     bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb);
     int64_t cal_mem_size(ggml_backend_t backend);
     ~ModelLoader() = default;

stable-diffusion.cpp

@@ -493,12 +493,39 @@ const int BOS_TOKEN_ID = 49406;
 const int EOS_TOKEN_ID = 49407;
 const int PAD_TOKEN_ID = 49407;
 
+std::vector<std::pair<int, std::u32string>> bytes_to_unicode() {
+    std::vector<std::pair<int, std::u32string>> byte_unicode_pairs;
+    std::set<int> byte_set;
+    for (int b = static_cast<int>('!'); b <= static_cast<int>('~'); ++b) {
+        byte_set.insert(b);
+        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
+    }
+    for (int b = 161; b <= 172; ++b) {
+        byte_set.insert(b);
+        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
+    }
+    for (int b = 174; b <= 255; ++b) {
+        byte_set.insert(b);
+        byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
+    }
+    int n = 0;
+    for (int b = 0; b < 256; ++b) {
+        if (byte_set.find(b) == byte_set.end()) {
+            byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(n + 256)));
+            ++n;
+        }
+    }
+    // LOG_DEBUG("byte_unicode_pairs %d", byte_unicode_pairs.size());
+    return byte_unicode_pairs;
+}
+
 // Ref: https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py
-// TODO: implement bpe
 class CLIPTokenizer {
 private:
     SDVersion version = VERSION_1_x;
-    std::map<std::string, int32_t> encoder;
+    std::map<int, std::u32string> byte_encoder;
+    std::map<std::u32string, int> encoder;
+    std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;
     std::regex pat;
 
     static std::string strip(const std::string& str) {
@@ -519,21 +546,145 @@ private:
         return text;
     }
 
-public:
-    CLIPTokenizer(SDVersion version = VERSION_1_x)
-        : version(version){};
-    std::string bpe(std::string token) {
-        std::string word = token + "</w>";
-        if (encoder.find(word) != encoder.end()) {
-            return word;
-        } else if (encoder.find(token) != encoder.end()) {
-            return token;
+    static std::set<std::pair<std::u32string, std::u32string>> get_pairs(const std::vector<std::u32string>& subwords) {
+        std::set<std::pair<std::u32string, std::u32string>> pairs;
+        if (subwords.size() == 0) {
+            return pairs;
         }
-        return UNK_TOKEN;
+        std::u32string prev_subword = subwords[0];
+        for (int i = 1; i < subwords.size(); i++) {
+            std::u32string subword = subwords[i];
+            std::pair<std::u32string, std::u32string> pair(prev_subword, subword);
+            pairs.insert(pair);
+            prev_subword = subword;
+        }
+        return pairs;
     }
 
-    void add_token(std::string token, int32_t token_id) {
-        encoder[token] = token_id;
+public:
+    CLIPTokenizer(SDVersion version = VERSION_1_x)
+        : version(version) {}
+
+    void load_from_merges(const std::string& merges_utf8_str) {
+        auto byte_unicode_pairs = bytes_to_unicode();
+        byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
+        // for (auto & pair: byte_unicode_pairs) {
+        //     std::cout << pair.first << ": " << pair.second << std::endl;
+        // }
+        std::vector<std::u32string> merges;
+        size_t start = 0;
+        size_t pos;
+        std::u32string merges_utf32_str = utf8_to_utf32(merges_utf8_str);
+        while ((pos = merges_utf32_str.find('\n', start)) != std::string::npos) {
+            merges.push_back(merges_utf32_str.substr(start, pos - start));
+            start = pos + 1;
+        }
+        // LOG_DEBUG("merges size %llu", merges.size());
+        GGML_ASSERT(merges.size() == 48895);
+        merges = std::vector<std::u32string>(merges.begin() + 1, merges.end());
+        std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
+        for (const auto& merge : merges) {
+            size_t space_pos = merge.find(' ');
+            merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
+            // LOG_DEBUG("%s", utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
+        }
+        std::vector<std::u32string> vocab;
+        for (const auto& pair : byte_unicode_pairs) {
+            vocab.push_back(pair.second);
+        }
+        for (const auto& pair : byte_unicode_pairs) {
+            vocab.push_back(pair.second + utf8_to_utf32("</w>"));
+        }
+        for (const auto& merge : merge_pairs) {
+            vocab.push_back(merge.first + merge.second);
+        }
+        vocab.push_back(utf8_to_utf32("<|startoftext|>"));
+        vocab.push_back(utf8_to_utf32("<|endoftext|>"));
+        LOG_DEBUG("vocab size: %llu", vocab.size());
+        int i = 0;
+        for (const auto& token : vocab) {
+            encoder[token] = i++;
+        }
+
+        int rank = 0;
+        for (const auto& merge : merge_pairs) {
+            bpe_ranks[merge] = rank++;
+        }
+    };
+
+    std::u32string bpe(const std::u32string& token) {
+        std::vector<std::u32string> word;
+
+        for (int i = 0; i < token.size() - 1; i++) {
+            word.emplace_back(1, token[i]);
+        }
+        word.push_back(token.substr(token.size() - 1) + utf8_to_utf32("</w>"));
+
+        std::set<std::pair<std::u32string, std::u32string>> pairs = get_pairs(word);
+
+        if (pairs.empty()) {
+            return token + utf8_to_utf32("</w>");
+        }
+
+        while (true) {
+            auto min_pair_iter = std::min_element(pairs.begin(),
+                                                  pairs.end(),
+                                                  [&](const std::pair<std::u32string, std::u32string>& a,
+                                                      const std::pair<std::u32string, std::u32string>& b) {
+                                                      if (bpe_ranks.find(a) == bpe_ranks.end()) {
+                                                          return false;
+                                                      } else if (bpe_ranks.find(b) == bpe_ranks.end()) {
+                                                          return true;
+                                                      }
+                                                      return bpe_ranks.at(a) < bpe_ranks.at(b);
+                                                  });
+
+            const std::pair<std::u32string, std::u32string>& bigram = *min_pair_iter;
+
+            if (bpe_ranks.find(bigram) == bpe_ranks.end()) {
+                break;
+            }
+
+            std::u32string first  = bigram.first;
+            std::u32string second = bigram.second;
+            std::vector<std::u32string> new_word;
+            int32_t i = 0;
+
+            while (i < word.size()) {
+                auto it = std::find(word.begin() + i, word.end(), first);
+                if (it == word.end()) {
+                    new_word.insert(new_word.end(), word.begin() + i, word.end());
+                    break;
+                }
+                new_word.insert(new_word.end(), word.begin() + i, it);
+                i = static_cast<int32_t>(std::distance(word.begin(), it));
+
+                if (word[i] == first && i < static_cast<int32_t>(word.size()) - 1 && word[i + 1] == second) {
+                    new_word.push_back(first + second);
+                    i += 2;
+                } else {
+                    new_word.push_back(word[i]);
+                    i += 1;
+                }
+            }
+
+            word = new_word;
+
+            if (word.size() == 1) {
+                break;
+            }
+            pairs = get_pairs(word);
+        }
+
+        std::u32string result;
+        for (int i = 0; i < word.size(); i++) {
+            result += word[i];
+            if (i != word.size() - 1) {
+                result += utf8_to_utf32(" ");
+            }
+        }
+
+        return result;
     }
 
     std::vector<int> tokenize(std::string text, size_t max_length = 0, bool padding = false) {
@@ -571,13 +722,25 @@ public:
         std::vector<std::string> token_strs;
         while (std::regex_search(str, matches, pat)) {
             for (auto& token : matches) {
-                std::istringstream iss(bpe(token));
-                std::vector<std::string> tokens{std::istream_iterator<std::string>{iss},
-                                                std::istream_iterator<std::string>{}};
-                for (const auto& bpe_token : tokens) {
-                    bpe_tokens.push_back(encoder[bpe_token]);
-                    token_strs.push_back(bpe_token);
+                std::string token_str = token.str();
+                std::u32string utf32_token;
+                for (int i = 0; i < token_str.length(); i++) {
+                    char b = token_str[i];
+                    utf32_token += byte_encoder[b];
                 }
+                auto bpe_strs = bpe(utf32_token);
+                size_t start  = 0;
+                size_t pos;
+                while ((pos = bpe_strs.find(' ', start)) != std::u32string::npos) {
+                    auto bpe_str = bpe_strs.substr(start, pos - start);
+                    bpe_tokens.push_back(encoder[bpe_str]);
+                    token_strs.push_back(utf32_to_utf8(bpe_str));
+
+                    start = pos + 1;
+                }
+                auto bpe_str = bpe_strs.substr(start, bpe_strs.size() - start);
+                bpe_tokens.push_back(encoder[bpe_str]);
+                token_strs.push_back(utf32_to_utf8(bpe_str));
             }
             str = matches.suffix();
         }
@@ -4323,15 +4486,14 @@ public:
         LOG_INFO("Stable Diffusion weight type: %s", ggml_type_name(model_data_type));
 
         LOG_DEBUG("loading vocab");
-        auto add_token = [&](const std::string& token, int32_t token_id) {
-            cond_stage_model.tokenizer.add_token(token, token_id);
-        };
-        bool success = model_loader.load_vocab(add_token);
-        if (!success) {
-            LOG_ERROR("get vocab from file failed: '%s'", model_path.c_str());
+        std::string merges_utf8_str = model_loader.load_merges();
+        if (merges_utf8_str.size() == 0) {
+            LOG_ERROR("get merges failed: '%s'", model_path.c_str());
             return false;
         }
 
+        cond_stage_model.tokenizer.load_from_merges(merges_utf8_str);
+
         // create the ggml context for network params
         LOG_DEBUG("ggml tensor size = %d bytes", (int)sizeof(ggml_tensor));
 
@@ -4431,7 +4593,7 @@ public:
 
         // print_ggml_tensor(alphas_cumprod_tensor);
 
-        success = model_loader.load_tensors(on_new_tensor_cb);
+        bool success = model_loader.load_tensors(on_new_tensor_cb);
         if (!success) {
             LOG_ERROR("load tensors from file failed");
             ggml_free(ctx);

stable-diffusion.h

@@ -5,6 +5,8 @@
 #include <string>
 #include <vector>
 
+#include "ggml/ggml.h"
+
 enum RNGType {
     STD_DEFAULT_RNG,
     CUDA_RNG
@@ -42,10 +44,12 @@ public:
                     bool free_params_immediately = false,
                     std::string lora_model_dir   = "",
                     RNGType rng_type             = STD_DEFAULT_RNG);
+
     bool load_from_file(const std::string& model_path,
                         const std::string& vae_path,
                         ggml_type wtype,
                         Schedule d = DEFAULT);
+
     std::vector<uint8_t*> txt2img(
         std::string prompt,
         std::string negative_prompt,

util.cpp

@@ -1,7 +1,9 @@
 #include "util.h"
 
 #include <stdarg.h>
+#include <codecvt>
 #include <fstream>
+#include <locale>
 #include <thread>
 #include <unordered_set>
 #include <vector>
@@ -119,6 +121,21 @@ int32_t get_num_physical_cores() {
     return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4;
 }
 
+std::u32string utf8_to_utf32(const std::string& utf8_str) {
+    std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
+    return converter.from_bytes(utf8_str);
+}
+
+std::string utf32_to_utf8(const std::u32string& utf32_str) {
+    std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
+    return converter.to_bytes(utf32_str);
+}
+
+std::u32string unicode_value_to_utf32(int unicode_value) {
+    std::u32string utf32_string = {static_cast<char32_t>(unicode_value)};
+    return utf32_string;
+}
+
 std::string basename(const std::string& path) {
     size_t pos = path.find_last_of('/');
     if (pos != std::string::npos) {

util.h

@@ -2,6 +2,7 @@
 #define __UTIL_H__
 
 #include <string>
+#include <cstdint>
 
 bool ends_with(const std::string& str, const std::string& ending);
 bool starts_with(const std::string& str, const std::string& start);
@@ -13,6 +14,10 @@ void replace_all_chars(std::string& str, char target, char replacement);
 bool file_exists(const std::string& filename);
 bool is_directory(const std::string& path);
 
+std::u32string utf8_to_utf32(const std::string& utf8_str);
+std::string utf32_to_utf8(const std::u32string& utf32_str);
+std::u32string unicode_value_to_utf32(int unicode_value);
+
 std::string basename(const std::string& path);
 
 std::string path_join(const std::string& p1, const std::string& p2);
@@ -34,4 +39,4 @@ void log_printf(SDLogLevel level, const char* file, int line, const char* format
 #define LOG_INFO(format, ...) log_printf(SDLogLevel::INFO, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_WARN(format, ...) log_printf(SDLogLevel::WARN, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_ERROR(format, ...) log_printf(SDLogLevel::ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__)
-#endif  // __UTIL_H__
+#endif  // __UTIL_H__