feat: allow models to run without all text encoder(s) (#645)

conditioner.hpp

@@ -673,33 +673,80 @@ struct SD3CLIPEmbedder : public Conditioner {
                     bool offload_params_to_cpu,
                     const String2GGMLType& tensor_types = {})
         : clip_g_tokenizer(0) {
+        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<CLIPTextModelRunner>(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<CLIPTextModelRunner>(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<T5Runner>(backend, offload_params_to_cpu, tensor_types, "text_encoders.t5xxl.transformer");
         }
+    }
 
     void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+        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 {
+        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 {
+        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();
+        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;
-
+            if (clip_l) {
                 std::vector<int> 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);
+            if (clip_g) {
+                std::vector<int> 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);
+            if (t5) {
+                std::vector<int> 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) {
             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<float> 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<int> chunk_tokens(clip_l_tokens.begin() + chunk_idx * chunk_len,
                                               clip_l_tokens.begin() + (chunk_idx + 1) * chunk_len);
                 std::vector<float> 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<int> chunk_tokens(clip_g_tokens.begin() + chunk_idx * chunk_len,
                                               clip_g_tokens.begin() + (chunk_idx + 1) * chunk_len);
                 std::vector<float> 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<int> chunk_tokens(t5_tokens.begin() + chunk_idx * chunk_len,
                                               t5_tokens.begin() + (chunk_idx + 1) * chunk_len);
                 std::vector<float> 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));
         }
 
+        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 = {}) {
+        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<CLIPTextModelRunner>(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<T5Runner>(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<std::string, struct ggml_tensor*>& tensors) override {
+        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 {
+        if (clip_l) {
             clip_l->alloc_params_buffer();
+        }
+        if (t5) {
             t5->alloc_params_buffer();
         }
+    }
 
     void free_params_buffer() override {
+        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();
+        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;
-
+            if (clip_l) {
                 std::vector<int> 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);
+            if (t5) {
+                std::vector<int> 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) {
             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,10 +1224,11 @@ struct FluxCLIPEmbedder : public Conditioner {
         struct ggml_tensor* pooled              = nullptr;  // [768,]
         std::vector<float> 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) {
+                if (clip_l) {
                     size_t chunk_len_l = 77;
                     std::vector<int> chunk_tokens(clip_l_tokens.begin(),
                                                   clip_l_tokens.begin() + chunk_len_l);
@@ -1122,9 +1251,10 @@ struct FluxCLIPEmbedder : public Conditioner {
                                     &pooled,
                                     work_ctx);
                 }
+            }
 
             // t5
-            {
+            if (t5) {
                 std::vector<int> chunk_tokens(t5_tokens.begin() + chunk_idx * chunk_len,
                                               t5_tokens.begin() + (chunk_idx + 1) * chunk_len);
                 std::vector<float> 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));
         }
 
+        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) {
+        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<T5Runner>(backend, offload_params_to_cpu, tensor_types, "text_encoders.t5xxl.transformer", is_umt5);
         }
+    }
 
     void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
+        if (t5) {
             t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
         }
+    }
 
     void alloc_params_buffer() override {
+        if (t5) {
             t5->alloc_params_buffer();
         }
+    }
 
     void free_params_buffer() override {
+        if (t5) {
             t5->free_params_buffer();
         }
+    }
 
     size_t get_params_buffer_size() override {
         size_t buffer_size = 0;
-
+        if (t5) {
             buffer_size += t5->get_params_buffer_size();
-
+        }
         return buffer_size;
     }
 
@@ -1248,6 +1408,7 @@ struct T5CLIPEmbedder : public Conditioner {
         std::vector<int> t5_tokens;
         std::vector<float> t5_weights;
         std::vector<float> t5_mask;
+        if (t5) {
             for (const auto& item : parsed_attention) {
                 const std::string& curr_text = item.first;
                 float curr_weight            = item.second;
@@ -1258,7 +1419,7 @@ struct T5CLIPEmbedder : public Conditioner {
             }
 
             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<int>, std::vector<float>, std::vector<float>> 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);