conditionner: make t5 optional for chroma

Author: stduhpf

conditioner.hpp

@@ -666,9 +666,6 @@ struct SD3CLIPEmbedder : public Conditioner {
                     bool offload_params_to_cpu,
                     const String2GGMLType& tensor_types = {})
         : clip_g_tokenizer(0) {
-        if (clip_skip <= 0) {
-            clip_skip = 2;
-        }
 
         for (auto pair : tensor_types) {
             if (pair.first.find("text_encoders.clip_l") != std::string::npos) {
@@ -684,12 +681,12 @@ struct SD3CLIPEmbedder : public Conditioner {
             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, clip_skip, false);
+            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);
         } else {
             LOG_WARN("clip_l text encoder not found! Prompt adherence might be degraded.");
         }
         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, clip_skip, false);
+            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);
         } else {
             LOG_WARN("clip_g text encoder not found! Prompt adherence might be degraded.");
         }
@@ -698,19 +695,6 @@ struct SD3CLIPEmbedder : public Conditioner {
         } else {
             LOG_WARN("t5xxl text encoder not found! Prompt adherence might be degraded.");
         }
-        set_clip_skip(clip_skip);
-    }
-
-    void set_clip_skip(int clip_skip) {
-        if (clip_skip <= 0) {
-            clip_skip = 2;
-        }
-        if (use_clip_l) {
-            clip_l->set_clip_skip(clip_skip);
-        }
-        if (use_clip_g) {
-            clip_g->set_clip_skip(clip_skip);
-        }
     }
 
     void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) {
@@ -1113,10 +1097,30 @@ struct FluxCLIPEmbedder : public Conditioner {
     FluxCLIPEmbedder(ggml_backend_t backend,
                      bool offload_params_to_cpu,
                      const String2GGMLType& tensor_types = {}) {
-        clip_l = std::make_shared<CLIPTextModelRunner>(backend, tensor_types, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, true);
-        t5     = std::make_shared<T5Runner>(backend, tensor_types, "text_encoders.t5xxl.transformer");
-    }
+        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) {
         if (use_clip_l) {
@@ -1296,7 +1300,6 @@ struct FluxCLIPEmbedder : public Conditioner {
                 ggml_set_f32(chunk_hidden_states, 0.f);
             }
 
-            
             int64_t t1 = ggml_time_ms();
             LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
             if (zero_out_masked) {
@@ -1305,12 +1308,12 @@ struct FluxCLIPEmbedder : public Conditioner {
                     vec[i] = 0;
                 }
             }
-            
+
             hidden_states_vec.insert(hidden_states_vec.end(),
-            (float*)chunk_hidden_states->data,
-            ((float*)chunk_hidden_states->data) + ggml_nelements(chunk_hidden_states));
+                                     (float*)chunk_hidden_states->data,
+                                     ((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,
@@ -1364,7 +1367,8 @@ struct T5CLIPEmbedder : public Conditioner {
     size_t chunk_len = 512;
     bool use_mask    = false;
     int mask_pad     = 1;
-    bool is_umt5     = false;
+    bool use_t5      = false;
+
 
     T5CLIPEmbedder(ggml_backend_t backend,
                    bool offload_params_to_cpu,
@@ -1373,26 +1377,43 @@ 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<T5Runner>(backend, offload_params_to_cpu, tensor_types, "text_encoders.t5xxl.transformer", is_umt5);
+        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) {
-        t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        if (use_t5) {
+            t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
     }
 
     void alloc_params_buffer() {
-        t5->alloc_params_buffer();
+        if (use_t5) {
+            t5->alloc_params_buffer();
+        }
     }
 
     void free_params_buffer() {
-        t5->free_params_buffer();
+        if (use_t5) {
+            t5->free_params_buffer();
+        }
     }
 
     size_t get_params_buffer_size() {
         size_t buffer_size = 0;
-
-        buffer_size += t5->get_params_buffer_size();
-
+        if (use_t5) {
+            buffer_size += t5->get_params_buffer_size();
+        }
         return buffer_size;
     }
 
@@ -1418,17 +1439,18 @@ struct T5CLIPEmbedder : public Conditioner {
         std::vector<int> t5_tokens;
         std::vector<float> t5_weights;
         std::vector<float> t5_mask;
-        for (const auto& item : parsed_attention) {
-            const std::string& curr_text = item.first;
-            float curr_weight            = item.second;
-
-            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 (use_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<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);
+            }
 
+            t5_tokenizer.pad_tokens(t5_tokens, t5_weights, &t5_mask, max_length, padding);
+        }
         return {t5_tokens, t5_weights, t5_mask};
     }
 
@@ -1465,66 +1487,78 @@ struct T5CLIPEmbedder : public Conditioner {
         std::vector<float> hidden_states_vec;
 
         size_t chunk_count = t5_tokens.size() / chunk_len;
-
         for (int chunk_idx = 0; chunk_idx < chunk_count; chunk_idx++) {
             // 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,
-                                             t5_weights.begin() + (chunk_idx + 1) * chunk_len);
-            std::vector<float> chunk_mask(t5_attn_mask_vec.begin() + chunk_idx * chunk_len,
-                                          t5_attn_mask_vec.begin() + (chunk_idx + 1) * chunk_len);
-
-            auto input_ids          = vector_to_ggml_tensor_i32(work_ctx, chunk_tokens);
-            auto t5_attn_mask_chunk = use_mask ? vector_to_ggml_tensor(work_ctx, chunk_mask) : NULL;
-
-            t5->compute(n_threads,
-                        input_ids,
-                        t5_attn_mask_chunk,
-                        &chunk_hidden_states,
-                        work_ctx);
-            {
-                auto tensor         = chunk_hidden_states;
-                float original_mean = ggml_tensor_mean(tensor);
-                for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
-                    for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
-                        for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
-                            float value = ggml_tensor_get_f32(tensor, i0, i1, i2);
-                            value *= chunk_weights[i1];
-                            ggml_tensor_set_f32(tensor, value, i0, i1, i2);
+
+            if (use_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,
+                                                 t5_weights.begin() + (chunk_idx + 1) * chunk_len);
+                std::vector<float> chunk_mask(t5_attn_mask_vec.begin() + chunk_idx * chunk_len,
+                                              t5_attn_mask_vec.begin() + (chunk_idx + 1) * chunk_len);
+
+                auto input_ids          = vector_to_ggml_tensor_i32(work_ctx, chunk_tokens);
+                auto t5_attn_mask_chunk = use_mask ? vector_to_ggml_tensor(work_ctx, chunk_mask) : NULL;
+                t5->compute(n_threads,
+                            input_ids,
+                            t5_attn_mask_chunk,
+                            &chunk_hidden_states,
+                            work_ctx);
+                {
+                    auto tensor         = chunk_hidden_states;
+                    float original_mean = ggml_tensor_mean(tensor);
+                    for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
+                        for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
+                            for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
+                                float value = ggml_tensor_get_f32(tensor, i0, i1, i2);
+                                value *= chunk_weights[i1];
+                                ggml_tensor_set_f32(tensor, value, i0, i1, i2);
+                            }
                         }
                     }
+                    float new_mean = ggml_tensor_mean(tensor);
+                    ggml_tensor_scale(tensor, (original_mean / new_mean));
                 }
-                float new_mean = ggml_tensor_mean(tensor);
-                ggml_tensor_scale(tensor, (original_mean / new_mean));
-            }
-
-            int64_t t1 = ggml_time_ms();
-            LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
-            if (zero_out_masked) {
-                auto tensor = chunk_hidden_states;
-                for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
-                    for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
-                        for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
-                            if (chunk_mask[i1] < 0.f) {
-                                ggml_tensor_set_f32(tensor, 0.f, i0, i1, i2);
+                if (zero_out_masked) {
+                    auto tensor = chunk_hidden_states;
+                    for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
+                        for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
+                            for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
+                                if (chunk_mask[i1] < 0.f) {
+                                    ggml_tensor_set_f32(tensor, 0.f, i0, i1, i2);
+                                }
                             }
                         }
                     }
                 }
+            } else {
+                chunk_hidden_states = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 4096, chunk_len);
+                ggml_set_f32(chunk_hidden_states, 0.f);
+                t5_attn_mask = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, chunk_len);
+                ggml_set_f32(t5_attn_mask, -HUGE_VALF);
             }
 
+            int64_t t1 = ggml_time_ms();
+            LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
+
             hidden_states_vec.insert(hidden_states_vec.end(),
                                      (float*)chunk_hidden_states->data,
                                      ((float*)chunk_hidden_states->data) + ggml_nelements(chunk_hidden_states));
         }
 
-        GGML_ASSERT(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]);
+        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, chunk_len);
+            ggml_set_f32(hidden_states, 0.f);
+            t5_attn_mask = ggml_new_tensor_1d(work_ctx, GGML_TYPE_F32, chunk_len);
+            ggml_set_f32(t5_attn_mask, -HUGE_VALF);
+        }
 
         modify_mask_to_attend_padding(t5_attn_mask, ggml_nelements(t5_attn_mask), mask_pad);