diff --git a/conditioner.hpp b/conditioner.hpp index e299d367..e4b475ce 100644 --- a/conditioner.hpp +++ b/conditioner.hpp @@ -673,33 +673,80 @@ struct SD3CLIPEmbedder : public Conditioner { bool offload_params_to_cpu, const String2GGMLType& tensor_types = {}) : clip_g_tokenizer(0) { - clip_l = std::make_shared(backend, offload_params_to_cpu, tensor_types, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, false); - clip_g = std::make_shared(backend, offload_params_to_cpu, tensor_types, "text_encoders.clip_g.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false); - t5 = std::make_shared(backend, offload_params_to_cpu, tensor_types, "text_encoders.t5xxl.transformer"); + bool use_clip_l = false; + bool use_clip_g = false; + bool use_t5 = false; + for (auto pair : tensor_types) { + if (pair.first.find("text_encoders.clip_l") != std::string::npos) { + use_clip_l = true; + } else if (pair.first.find("text_encoders.clip_g") != std::string::npos) { + use_clip_g = true; + } else if (pair.first.find("text_encoders.t5xxl") != std::string::npos) { + use_t5 = true; + } + } + if (!use_clip_l && !use_clip_g && !use_t5) { + LOG_WARN("IMPORTANT NOTICE: No text encoders provided, cannot process prompts!"); + return; + } + if (use_clip_l) { + clip_l = std::make_shared(backend, offload_params_to_cpu, tensor_types, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, false); + } + if (use_clip_g) { + clip_g = std::make_shared(backend, offload_params_to_cpu, tensor_types, "text_encoders.clip_g.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false); + } + if (use_t5) { + t5 = std::make_shared(backend, offload_params_to_cpu, tensor_types, "text_encoders.t5xxl.transformer"); + } } void get_param_tensors(std::map& tensors) override { - clip_l->get_param_tensors(tensors, "text_encoders.clip_l.transformer.text_model"); - clip_g->get_param_tensors(tensors, "text_encoders.clip_g.transformer.text_model"); - t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + if (clip_l) { + clip_l->get_param_tensors(tensors, "text_encoders.clip_l.transformer.text_model"); + } + if (clip_g) { + clip_g->get_param_tensors(tensors, "text_encoders.clip_g.transformer.text_model"); + } + if (t5) { + t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } } void alloc_params_buffer() override { - clip_l->alloc_params_buffer(); - clip_g->alloc_params_buffer(); - t5->alloc_params_buffer(); + if (clip_l) { + clip_l->alloc_params_buffer(); + } + if (clip_g) { + clip_g->alloc_params_buffer(); + } + if (t5) { + t5->alloc_params_buffer(); + } } void free_params_buffer() override { - clip_l->free_params_buffer(); - clip_g->free_params_buffer(); - t5->free_params_buffer(); + if (clip_l) { + clip_l->free_params_buffer(); + } + if (clip_g) { + clip_g->free_params_buffer(); + } + if (t5) { + t5->free_params_buffer(); + } } size_t get_params_buffer_size() override { - size_t buffer_size = clip_l->get_params_buffer_size(); - buffer_size += clip_g->get_params_buffer_size(); - buffer_size += t5->get_params_buffer_size(); + size_t buffer_size = 0; + if (clip_l) { + buffer_size += clip_l->get_params_buffer_size(); + } + if (clip_g) { + buffer_size += clip_g->get_params_buffer_size(); + } + if (t5) { + buffer_size += t5->get_params_buffer_size(); + } return buffer_size; } @@ -731,23 +778,32 @@ struct SD3CLIPEmbedder : public Conditioner { for (const auto& item : parsed_attention) { const std::string& curr_text = item.first; float curr_weight = item.second; - - std::vector curr_tokens = clip_l_tokenizer.encode(curr_text, on_new_token_cb); - clip_l_tokens.insert(clip_l_tokens.end(), curr_tokens.begin(), curr_tokens.end()); - clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight); - - curr_tokens = clip_g_tokenizer.encode(curr_text, on_new_token_cb); - clip_g_tokens.insert(clip_g_tokens.end(), curr_tokens.begin(), curr_tokens.end()); - clip_g_weights.insert(clip_g_weights.end(), curr_tokens.size(), curr_weight); - - curr_tokens = t5_tokenizer.Encode(curr_text, true); - t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end()); - t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight); + if (clip_l) { + std::vector curr_tokens = clip_l_tokenizer.encode(curr_text, on_new_token_cb); + clip_l_tokens.insert(clip_l_tokens.end(), curr_tokens.begin(), curr_tokens.end()); + clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight); + } + if (clip_g) { + std::vector curr_tokens = clip_g_tokenizer.encode(curr_text, on_new_token_cb); + clip_g_tokens.insert(clip_g_tokens.end(), curr_tokens.begin(), curr_tokens.end()); + clip_g_weights.insert(clip_g_weights.end(), curr_tokens.size(), curr_weight); + } + if (t5) { + std::vector curr_tokens = t5_tokenizer.Encode(curr_text, true); + t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end()); + t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight); + } } - clip_l_tokenizer.pad_tokens(clip_l_tokens, clip_l_weights, max_length, padding); - clip_g_tokenizer.pad_tokens(clip_g_tokens, clip_g_weights, max_length, padding); - t5_tokenizer.pad_tokens(t5_tokens, t5_weights, nullptr, max_length, padding); + if (clip_l) { + clip_l_tokenizer.pad_tokens(clip_l_tokens, clip_l_weights, max_length, padding); + } + if (clip_g) { + clip_g_tokenizer.pad_tokens(clip_g_tokens, clip_g_weights, max_length, padding); + } + if (t5) { + t5_tokenizer.pad_tokens(t5_tokens, t5_weights, nullptr, max_length, padding); + } // for (int i = 0; i < clip_l_tokens.size(); i++) { // std::cout << clip_l_tokens[i] << ":" << clip_l_weights[i] << ", "; @@ -795,10 +851,10 @@ struct SD3CLIPEmbedder : public Conditioner { std::vector hidden_states_vec; size_t chunk_len = 77; - size_t chunk_count = clip_l_tokens.size() / chunk_len; + size_t chunk_count = std::max(std::max(clip_l_tokens.size(), clip_g_tokens.size()), t5_tokens.size()) / chunk_len; for (int chunk_idx = 0; chunk_idx < chunk_count; chunk_idx++) { // clip_l - { + if (clip_l) { std::vector chunk_tokens(clip_l_tokens.begin() + chunk_idx * chunk_len, clip_l_tokens.begin() + (chunk_idx + 1) * chunk_len); std::vector chunk_weights(clip_l_weights.begin() + chunk_idx * chunk_len, @@ -845,10 +901,17 @@ struct SD3CLIPEmbedder : public Conditioner { &pooled_l, work_ctx); } + } else { + chunk_hidden_states_l = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 768, chunk_len); + ggml_set_f32(chunk_hidden_states_l, 0.f); + if (chunk_idx == 0) { + pooled_l = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, 768); + ggml_set_f32(pooled_l, 0.f); + } } // clip_g - { + if (clip_g) { std::vector chunk_tokens(clip_g_tokens.begin() + chunk_idx * chunk_len, clip_g_tokens.begin() + (chunk_idx + 1) * chunk_len); std::vector chunk_weights(clip_g_weights.begin() + chunk_idx * chunk_len, @@ -896,10 +959,17 @@ struct SD3CLIPEmbedder : public Conditioner { &pooled_g, work_ctx); } + } else { + chunk_hidden_states_g = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 1280, chunk_len); + ggml_set_f32(chunk_hidden_states_g, 0.f); + if (chunk_idx == 0) { + pooled_g = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, 1280); + ggml_set_f32(pooled_g, 0.f); + } } // t5 - { + if (t5) { std::vector chunk_tokens(t5_tokens.begin() + chunk_idx * chunk_len, t5_tokens.begin() + (chunk_idx + 1) * chunk_len); std::vector chunk_weights(t5_weights.begin() + chunk_idx * chunk_len, @@ -927,6 +997,9 @@ struct SD3CLIPEmbedder : public Conditioner { float new_mean = ggml_tensor_mean(tensor); ggml_tensor_scale(tensor, (original_mean / new_mean)); } + } else { + chunk_hidden_states_t5 = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 4096, chunk_len); + ggml_set_f32(chunk_hidden_states_t5, 0.f); } auto chunk_hidden_states_lg_pad = ggml_new_tensor_3d(work_ctx, @@ -969,11 +1042,20 @@ struct SD3CLIPEmbedder : public Conditioner { ((float*)chunk_hidden_states->data) + ggml_nelements(chunk_hidden_states)); } - hidden_states = vector_to_ggml_tensor(work_ctx, hidden_states_vec); - hidden_states = ggml_reshape_2d(work_ctx, - hidden_states, - chunk_hidden_states->ne[0], - ggml_nelements(hidden_states) / chunk_hidden_states->ne[0]); + if (hidden_states_vec.size() > 0) { + hidden_states = vector_to_ggml_tensor(work_ctx, hidden_states_vec); + hidden_states = ggml_reshape_2d(work_ctx, + hidden_states, + chunk_hidden_states->ne[0], + ggml_nelements(hidden_states) / chunk_hidden_states->ne[0]); + } else { + hidden_states = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 4096, 256); + ggml_set_f32(hidden_states, 0.f); + } + if (pooled == nullptr) { + pooled = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, 2048); + ggml_set_f32(pooled, 0.f); + } return {hidden_states, pooled, nullptr}; } @@ -999,28 +1081,68 @@ struct FluxCLIPEmbedder : public Conditioner { FluxCLIPEmbedder(ggml_backend_t backend, bool offload_params_to_cpu, const String2GGMLType& tensor_types = {}) { - clip_l = std::make_shared(backend, offload_params_to_cpu, tensor_types, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, true); - t5 = std::make_shared(backend, offload_params_to_cpu, tensor_types, "text_encoders.t5xxl.transformer"); + bool use_clip_l = false; + bool use_t5 = false; + for (auto pair : tensor_types) { + if (pair.first.find("text_encoders.clip_l") != std::string::npos) { + use_clip_l = true; + } else if (pair.first.find("text_encoders.t5xxl") != std::string::npos) { + use_t5 = true; + } + } + + if (!use_clip_l && !use_t5) { + LOG_WARN("IMPORTANT NOTICE: No text encoders provided, cannot process prompts!"); + return; + } + + if (use_clip_l) { + clip_l = std::make_shared(backend, offload_params_to_cpu, tensor_types, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, true); + } else { + LOG_WARN("clip_l text encoder not found! Prompt adherence might be degraded."); + } + if (use_t5) { + t5 = std::make_shared(backend, offload_params_to_cpu, tensor_types, "text_encoders.t5xxl.transformer"); + } else { + LOG_WARN("t5xxl text encoder not found! Prompt adherence might be degraded."); + } } void get_param_tensors(std::map& tensors) override { - clip_l->get_param_tensors(tensors, "text_encoders.clip_l.transformer.text_model"); - t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + if (clip_l) { + clip_l->get_param_tensors(tensors, "text_encoders.clip_l.transformer.text_model"); + } + if (t5) { + t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } } void alloc_params_buffer() override { - clip_l->alloc_params_buffer(); - t5->alloc_params_buffer(); + if (clip_l) { + clip_l->alloc_params_buffer(); + } + if (t5) { + t5->alloc_params_buffer(); + } } void free_params_buffer() override { - clip_l->free_params_buffer(); - t5->free_params_buffer(); + if (clip_l) { + clip_l->free_params_buffer(); + } + if (t5) { + t5->free_params_buffer(); + } } size_t get_params_buffer_size() override { - size_t buffer_size = clip_l->get_params_buffer_size(); - buffer_size += t5->get_params_buffer_size(); + size_t buffer_size = 0; + if (clip_l) { + buffer_size += clip_l->get_params_buffer_size(); + } + if (t5) { + buffer_size += t5->get_params_buffer_size(); + } return buffer_size; } @@ -1050,18 +1172,24 @@ struct FluxCLIPEmbedder : public Conditioner { for (const auto& item : parsed_attention) { const std::string& curr_text = item.first; float curr_weight = item.second; - - std::vector curr_tokens = clip_l_tokenizer.encode(curr_text, on_new_token_cb); - clip_l_tokens.insert(clip_l_tokens.end(), curr_tokens.begin(), curr_tokens.end()); - clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight); - - curr_tokens = t5_tokenizer.Encode(curr_text, true); - t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end()); - t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight); + if (clip_l) { + std::vector curr_tokens = clip_l_tokenizer.encode(curr_text, on_new_token_cb); + clip_l_tokens.insert(clip_l_tokens.end(), curr_tokens.begin(), curr_tokens.end()); + clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight); + } + if (t5) { + std::vector curr_tokens = t5_tokenizer.Encode(curr_text, true); + t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end()); + t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight); + } } - clip_l_tokenizer.pad_tokens(clip_l_tokens, clip_l_weights, 77, padding); - t5_tokenizer.pad_tokens(t5_tokens, t5_weights, nullptr, max_length, padding); + if (clip_l) { + clip_l_tokenizer.pad_tokens(clip_l_tokens, clip_l_weights, 77, padding); + } + if (t5) { + t5_tokenizer.pad_tokens(t5_tokens, t5_weights, nullptr, max_length, padding); + } // for (int i = 0; i < clip_l_tokens.size(); i++) { // std::cout << clip_l_tokens[i] << ":" << clip_l_weights[i] << ", "; @@ -1096,35 +1224,37 @@ struct FluxCLIPEmbedder : public Conditioner { struct ggml_tensor* pooled = nullptr; // [768,] std::vector hidden_states_vec; - size_t chunk_count = t5_tokens.size() / chunk_len; + size_t chunk_count = std::max(clip_l_tokens.size() > 0 ? chunk_len : 0, t5_tokens.size()) / chunk_len; for (int chunk_idx = 0; chunk_idx < chunk_count; chunk_idx++) { // clip_l if (chunk_idx == 0) { - size_t chunk_len_l = 77; - std::vector chunk_tokens(clip_l_tokens.begin(), - clip_l_tokens.begin() + chunk_len_l); - std::vector chunk_weights(clip_l_weights.begin(), - clip_l_weights.begin() + chunk_len_l); + if (clip_l) { + size_t chunk_len_l = 77; + std::vector chunk_tokens(clip_l_tokens.begin(), + clip_l_tokens.begin() + chunk_len_l); + std::vector chunk_weights(clip_l_weights.begin(), + clip_l_weights.begin() + chunk_len_l); - auto input_ids = vector_to_ggml_tensor_i32(work_ctx, chunk_tokens); - size_t max_token_idx = 0; + auto input_ids = vector_to_ggml_tensor_i32(work_ctx, chunk_tokens); + size_t max_token_idx = 0; - auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_l_tokenizer.EOS_TOKEN_ID); - max_token_idx = std::min(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1); + auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_l_tokenizer.EOS_TOKEN_ID); + max_token_idx = std::min(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1); - clip_l->compute(n_threads, - input_ids, - 0, - nullptr, - max_token_idx, - true, - clip_skip, - &pooled, - work_ctx); + clip_l->compute(n_threads, + input_ids, + 0, + nullptr, + max_token_idx, + true, + clip_skip, + &pooled, + work_ctx); + } } // t5 - { + if (t5) { std::vector chunk_tokens(t5_tokens.begin() + chunk_idx * chunk_len, t5_tokens.begin() + (chunk_idx + 1) * chunk_len); std::vector chunk_weights(t5_weights.begin() + chunk_idx * chunk_len, @@ -1152,6 +1282,9 @@ struct FluxCLIPEmbedder : public Conditioner { float new_mean = ggml_tensor_mean(tensor); ggml_tensor_scale(tensor, (original_mean / new_mean)); } + } else { + chunk_hidden_states = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 4096, chunk_len); + ggml_set_f32(chunk_hidden_states, 0.f); } int64_t t1 = ggml_time_ms(); @@ -1168,11 +1301,20 @@ struct FluxCLIPEmbedder : public Conditioner { ((float*)chunk_hidden_states->data) + ggml_nelements(chunk_hidden_states)); } - hidden_states = vector_to_ggml_tensor(work_ctx, hidden_states_vec); - hidden_states = ggml_reshape_2d(work_ctx, - hidden_states, - chunk_hidden_states->ne[0], - ggml_nelements(hidden_states) / chunk_hidden_states->ne[0]); + if (hidden_states_vec.size() > 0) { + hidden_states = vector_to_ggml_tensor(work_ctx, hidden_states_vec); + hidden_states = ggml_reshape_2d(work_ctx, + hidden_states, + chunk_hidden_states->ne[0], + ggml_nelements(hidden_states) / chunk_hidden_states->ne[0]); + } else { + hidden_states = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 4096, 256); + ggml_set_f32(hidden_states, 0.f); + } + if (pooled == nullptr) { + pooled = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, 768); + ggml_set_f32(pooled, 0.f); + } return {hidden_states, pooled, nullptr}; } @@ -1203,26 +1345,44 @@ struct T5CLIPEmbedder : public Conditioner { int mask_pad = 1, bool is_umt5 = false) : use_mask(use_mask), mask_pad(mask_pad), t5_tokenizer(is_umt5) { - t5 = std::make_shared(backend, offload_params_to_cpu, tensor_types, "text_encoders.t5xxl.transformer", is_umt5); + bool use_t5 = false; + for (auto pair : tensor_types) { + if (pair.first.find("text_encoders.t5xxl") != std::string::npos) { + use_t5 = true; + } + } + + if (!use_t5) { + LOG_WARN("IMPORTANT NOTICE: No text encoders provided, cannot process prompts!"); + return; + } else { + t5 = std::make_shared(backend, offload_params_to_cpu, tensor_types, "text_encoders.t5xxl.transformer", is_umt5); + } } void get_param_tensors(std::map& tensors) override { - t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + if (t5) { + t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } } void alloc_params_buffer() override { - t5->alloc_params_buffer(); + if (t5) { + t5->alloc_params_buffer(); + } } void free_params_buffer() override { - t5->free_params_buffer(); + if (t5) { + t5->free_params_buffer(); + } } size_t get_params_buffer_size() override { size_t buffer_size = 0; - - buffer_size += t5->get_params_buffer_size(); - + if (t5) { + buffer_size += t5->get_params_buffer_size(); + } return buffer_size; } @@ -1248,17 +1408,18 @@ struct T5CLIPEmbedder : public Conditioner { std::vector t5_tokens; std::vector t5_weights; std::vector t5_mask; - for (const auto& item : parsed_attention) { - const std::string& curr_text = item.first; - float curr_weight = item.second; - - std::vector curr_tokens = t5_tokenizer.Encode(curr_text, true); - t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end()); - t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight); - } + if (t5) { + for (const auto& item : parsed_attention) { + const std::string& curr_text = item.first; + float curr_weight = item.second; - t5_tokenizer.pad_tokens(t5_tokens, t5_weights, &t5_mask, max_length, padding); + std::vector curr_tokens = t5_tokenizer.Encode(curr_text, true); + t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end()); + t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight); + } + t5_tokenizer.pad_tokens(t5_tokens, t5_weights, &t5_mask, max_length, padding); + } return {t5_tokens, t5_weights, t5_mask}; } @@ -1282,6 +1443,13 @@ struct T5CLIPEmbedder : public Conditioner { std::tuple, std::vector, std::vector> token_and_weights, int clip_skip, bool zero_out_masked = false) { + if (!t5) { + auto hidden_states = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 4096, 256); + ggml_set_f32(hidden_states, 0.f); + auto t5_attn_mask = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, 256); + ggml_set_f32(t5_attn_mask, -HUGE_VALF); + return {hidden_states, t5_attn_mask, nullptr}; + } auto& t5_tokens = std::get<0>(token_and_weights); auto& t5_weights = std::get<1>(token_and_weights); auto& t5_attn_mask_vec = std::get<2>(token_and_weights);