From 87fcee08518c3ecbda723e161db55a032213a683 Mon Sep 17 00:00:00 2001 From: bssrdf Date: Tue, 6 Feb 2024 22:17:40 -0500 Subject: [PATCH] making progress; need to implement tokenizer decoder --- clip.hpp | 167 +++++++++++++++++++++++++++++++++++++++++-- model.h | 1 + stable-diffusion.cpp | 54 ++++++++++---- 3 files changed, 203 insertions(+), 19 deletions(-) diff --git a/clip.hpp b/clip.hpp index d655661d..672fb20f 100644 --- a/clip.hpp +++ b/clip.hpp @@ -75,6 +75,7 @@ class CLIPTokenizer { private: SDVersion version = VERSION_1_x; std::map byte_encoder; + std::map byte_decoder; std::map encoder; std::map decoder; std::map, int> bpe_ranks; @@ -123,6 +124,9 @@ public: auto byte_unicode_pairs = bytes_to_unicode(); // printf("byte_unicode_pairs have %lu pairs \n", byte_unicode_pairs.size()); byte_encoder = std::map(byte_unicode_pairs.begin(), byte_unicode_pairs.end()); + for (auto & pair: byte_unicode_pairs) { + byte_decoder[pair.second] = pair.first; + } // for (auto & pair: byte_unicode_pairs) { // std::cout << pair.first << ": " << pair.second << std::endl; // } @@ -292,6 +296,25 @@ public: return tokens; } + std::string decode(const std::vector &tokens){ + std::u32string text = utf8_to_utf32(""); + for(int t : tokens){ + if(t == 49406 || t == 49407) + continue; + std::u32string ts = decoder[t]; + // printf("%d, %s\n", t, ts.c_str()); + text += ts; + } + std::vector bytes; + for (auto c : text){ + bytes.push_back(byte_decoder[c]); + } + + std::string s((char *)bytes.data()); + return s; + } + + std::vector encode(std::string text, on_new_token_cb_t on_new_token_cb) { std::string original_text = text; std::vector bpe_tokens; @@ -339,7 +362,7 @@ public: } ss << "]"; // LOG_DEBUG("split prompt \"%s\" to tokens %s", original_text.c_str(), ss.str().c_str()); - printf("split prompt \"%s\" to tokens %s \n", original_text.c_str(), ss.str().c_str()); + // printf("split prompt \"%s\" to tokens %s \n", original_text.c_str(), ss.str().c_str()); return bpe_tokens; } }; @@ -1215,6 +1238,16 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule { return tokenize(text, text_model.max_position_embeddings, padding); } + std::tuple, std::vector, std::vector> + tokenize_with_trigger_token(std::string text, + int num_input_imgs, + int32_t image_token, + bool padding = false) { + return tokenize_with_trigger_token(text, num_input_imgs, image_token, + text_model.max_position_embeddings, padding); + } + + std::vector convert_token_to_id(std::string text) { auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { size_t word_end = str.find(","); @@ -1243,6 +1276,130 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule { return curr_tokens; } + std::string decode(const std::vector &tokens){ + return tokenizer.decode(tokens); + } + + + std::tuple, std::vector, std::vector> + tokenize_with_trigger_token(std::string text, + int num_input_imgs, + int32_t image_token, + size_t max_length = 0, + bool padding = false) { + auto parsed_attention = parse_prompt_attention(text); + + { + std::stringstream ss; + ss << "["; + for (const auto& item : parsed_attention) { + ss << "['" << item.first << "', " << item.second << "], "; + } + ss << "]"; + LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str()); + } + + auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { + size_t word_end = str.find(","); + std::string embd_name = word_end == std::string::npos ? str : str.substr(0, word_end); + embd_name = trim(embd_name); + std::string embd_path = get_full_path(text_model.embd_dir, embd_name + ".pt"); + if (embd_path.size() == 0) { + embd_path = get_full_path(text_model.embd_dir, embd_name + ".ckpt"); + } + if (embd_path.size() == 0) { + embd_path = get_full_path(text_model.embd_dir, embd_name + ".safetensors"); + } + if (embd_path.size() > 0) { + if (text_model.load_embedding(embd_name, embd_path, bpe_tokens)) { + if (word_end != std::string::npos) { + str = str.substr(word_end); + } else { + str = ""; + } + return true; + } + } + return false; + }; + + std::vector tokens; + std::vector weights; + std::vector class_token_mask; + std::vector clean_input_ids; + std::vector class_token_index; + for (const auto& item : parsed_attention) { + const std::string& curr_text = item.first; + float curr_weight = item.second; + std::vector curr_tokens = tokenizer.encode(curr_text, on_new_token_cb); + int32_t clean_index = 0; + for(uint32_t i = 0; i < curr_tokens.size(); i++){ + int token_id = curr_tokens[i]; + if (token_id == image_token) + class_token_index.push_back(clean_index - 1); + else{ + clean_input_ids.push_back(token_id); + clean_index++; + } + } + GGML_ASSERT(class_token_index.size() == 1); // PhotoMaker currently does not support multiple + // trigger words in a single prompt. + // Expand the class word token and corresponding mask + int class_token = clean_input_ids[class_token_index[0]]; + std::vector clean_input_ids_tmp; + for(uint32_t i = 0; i < class_token_index[0]; i++) + clean_input_ids_tmp.push_back(clean_input_ids[i]); + for(uint32_t i = 0; i < num_input_imgs; i++) + clean_input_ids_tmp.push_back(class_token); + for(uint32_t i = class_token_index[0]+1; i < clean_input_ids.size(); i++) + clean_input_ids_tmp.push_back(clean_input_ids[i]); + clean_input_ids.clear(); + clean_input_ids = clean_input_ids_tmp; + + // class_tokens_mask = [True if class_token_index <= i < class_token_index+num_id_images else False \ + // for i in range(len(clean_input_ids))] + + tokens.insert(tokens.end(), clean_input_ids.begin(), clean_input_ids.end()); + weights.insert(weights.end(), clean_input_ids.size(), curr_weight); + } + tokens.insert(tokens.begin(), BOS_TOKEN_ID); + weights.insert(weights.begin(), 1.0); + + if (max_length > 0) { + if (tokens.size() > max_length - 1) { + tokens.resize(max_length - 1); + weights.resize(max_length - 1); + tokens.push_back(EOS_TOKEN_ID); + weights.push_back(1.0); + } else { + tokens.push_back(EOS_TOKEN_ID); + weights.push_back(1.0); + if (padding) { + int pad_token_id = PAD_TOKEN_ID; + if (version == VERSION_2_x) { + pad_token_id = 0; + } + tokens.insert(tokens.end(), max_length - tokens.size(), pad_token_id); + weights.insert(weights.end(), max_length - weights.size(), 1.0); + } + } + } + + for(uint32_t i = 0; i < tokens.size(); i++){ + if(class_token_index[0] <= i && i < class_token_index[0]+num_input_imgs) + class_token_mask.push_back(true); + else + class_token_mask.push_back(false); + } + + // for (int i = 0; i < tokens.size(); i++) { + // std::cout << tokens[i] << ":" << weights[i] << ", "; + // } + // std::cout << std::endl; + + return std::make_tuple(tokens, weights, class_token_mask); + } + std::pair, std::vector> tokenize(std::string text, size_t max_length = 0, bool padding = false) { @@ -1314,10 +1471,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule { } } - for (int i = 0; i < tokens.size(); i++) { - std::cout << tokens[i] << ":" << weights[i] << ", "; - } - std::cout << std::endl; + // for (int i = 0; i < tokens.size(); i++) { + // std::cout << tokens[i] << ":" << weights[i] << ", "; + // } + // std::cout << std::endl; return {tokens, weights}; } diff --git a/model.h b/model.h index 13665a7e..6a6644ad 100644 --- a/model.h +++ b/model.h @@ -7,6 +7,7 @@ #include #include #include +#include #include "ggml/ggml-backend.h" #include "ggml/ggml.h" diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index c898a0e9..01253a7f 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -474,6 +474,20 @@ public: curr_lora_state = lora_state; } + std::string remove_trigger_from_prompt(ggml_context* work_ctx, + const std::string& prompt){ + + auto image_tokens = cond_stage_model.convert_token_to_id(trigger_word); + GGML_ASSERT(image_tokens.size() == 1); + + auto tokens_and_weights = cond_stage_model.tokenize(prompt, false); + std::vector& tokens = tokens_and_weights.first; + auto it = std::find(tokens.begin(), tokens.end(), image_tokens[0]); + GGML_ASSERT(it != tokens.end()); // prompt must have trigger word + tokens.erase(it); + return cond_stage_model.decode(tokens); + + } std::tuple get_learned_condition_with_trigger(ggml_context* work_ctx, @@ -483,13 +497,21 @@ public: int height, int num_input_imgs, bool force_zero_embeddings = false) { - cond_stage_model.set_clip_skip(clip_skip); - auto image_token_id = cond_stage_model.tokenize(trigger_word, true); - - - auto tokens_and_weights = cond_stage_model.tokenize(text, true); - std::vector& tokens = tokens_and_weights.first; - std::vector& weights = tokens_and_weights.second; + cond_stage_model.set_clip_skip(clip_skip); + auto image_tokens = cond_stage_model.convert_token_to_id(trigger_word); + printf(" length of image tokens: %lu \n", image_tokens.size()); + if(image_tokens.size() == 1){ + printf(" image token id is: %d \n", image_tokens[0]); + } + GGML_ASSERT(image_tokens.size() == 1); + // auto tokens_and_weights = cond_stage_model.tokenize(text, true); + auto tokens_and_weights = cond_stage_model.tokenize_with_trigger_token(text, + num_input_imgs, + image_tokens[0], + true); + std::vector& tokens = std::get<0>(tokens_and_weights); + std::vector& weights = std::get<1>(tokens_and_weights); + std::vector& clsm = std::get<2>(tokens_and_weights); int64_t t0 = ggml_time_ms(); struct ggml_tensor* pooled = NULL; size_t total_hidden_size = cond_stage_model.text_model.hidden_size; @@ -509,7 +531,13 @@ public: // print_ggml_tensor(pooled); // } - ggml_tensor *class_tokens_mask = ggml_dup_tensor(work_ctx, hidden_states); + ggml_tensor *class_tokens_mask = ggml_new_tensor_1d(work_ctx, GGML_TYPE_I32, cond_stage_model.text_model.max_position_embeddings); + for (int i1 = 0; i1 < hidden_states->ne[1]; i1++) { + if(clsm[i1]) + ggml_set_i32_1d(class_tokens_mask, i1, 1); + else + ggml_set_i32_1d(class_tokens_mask, i1, 0); + } int64_t t1 = ggml_time_ms(); LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); @@ -584,12 +612,7 @@ public: int width, int height, bool force_zero_embeddings = false) { - cond_stage_model.set_clip_skip(clip_skip); - auto image_tokens = cond_stage_model.convert_token_to_id(trigger_word); - printf(" length of image tokens: %lu \n", image_tokens.size()); - if(image_tokens.size() == 1){ - printf(" image token id is: %d \n", image_tokens[0]); - } + cond_stage_model.set_clip_skip(clip_skip); auto tokens_and_weights = cond_stage_model.tokenize(text, true); std::vector& tokens = tokens_and_weights.first; @@ -1426,6 +1449,9 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx, __func__, ne[0], ne[1], ne[2], ne[3]); } + std::string prompt_text_only = sd_ctx->sd->remove_trigger_from_prompt(work_ctx, prompt); + printf("%s || %s \n", prompt.c_str(), prompt_text_only.c_str()); + t0 = ggml_time_ms(); auto cond_pair = sd_ctx->sd->get_learned_condition(work_ctx, prompt, clip_skip, width, height); ggml_tensor* c = cond_pair.first;