refector: remove some useless code

* apply pmid lora only once for multiple txt2img calls

* add better support for SDXL LoRA

* fix for some sdxl lora, like lcm-lora-xl

---------

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

clip.hpp

@@ -939,22 +939,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule {
         return "clip";
     }
 
-    size_t get_params_mem_size() {
-        size_t params_mem_size = text_model.get_params_mem_size();
-        if (version == VERSION_XL) {
-            params_mem_size += text_model2.get_params_mem_size();
-        }
-        return params_mem_size;
-    }
-
-    size_t get_params_num() {
-        size_t params_num = text_model.get_params_num();
-        if (version == VERSION_XL) {
-            params_num += text_model2.get_params_num();
-        }
-        return params_num;
-    }
-
     void set_clip_skip(int clip_skip) {
         text_model.set_clip_skip(clip_skip);
         if (version == VERSION_XL) {
@@ -1385,14 +1369,6 @@ struct FrozenCLIPVisionEmbedder : public GGMLModule {
         return "clip_vision";
     }
 
-    size_t get_params_mem_size() {
-        return vision_model.get_params_mem_size();
-    }
-
-    size_t get_params_num() {
-        return vision_model.get_params_num();
-    }
-
     void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
         vision_model.get_param_tensors(tensors, prefix + "transformer");
     }

control.hpp

@@ -369,14 +369,6 @@ struct ControlNet : public GGMLModule {
         return "control_net";
     }
 
-    size_t get_params_mem_size() {
-        return control_net.get_params_mem_size();
-    }
-
-    size_t get_params_num() {
-        return control_net.get_params_num();
-    }
-
     void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
         control_net.get_param_tensors(tensors, prefix);
     }

esrgan.hpp

@@ -152,14 +152,6 @@ struct ESRGAN : public GGMLModule {
         return "esrgan";
     }
 
-    size_t get_params_mem_size() {
-        return rrdb_net.get_params_mem_size();
-    }
-
-    size_t get_params_num() {
-        return rrdb_net.get_params_num();
-    }
-
     bool load_from_file(const std::string& file_path) {
         LOG_INFO("loading esrgan from '%s'", file_path.c_str());
 

examples/cli/main.cpp

@@ -686,27 +686,26 @@ int main(int argc, const char* argv[]) {
         // Resize input image ...
         if (params.height % 64 != 0 || params.width % 64 != 0) {
             int resized_height = params.height + (64 - params.height % 64);
-            int resized_width = params.width + (64 - params.width % 64);
+            int resized_width  = params.width + (64 - params.width % 64);
 
-            uint8_t *resized_image_buffer = (uint8_t *)malloc(resized_height * resized_width * 3);
+            uint8_t* resized_image_buffer = (uint8_t*)malloc(resized_height * resized_width * 3);
             if (resized_image_buffer == NULL) {
                 fprintf(stderr, "error: allocate memory for resize input image\n");
                 free(input_image_buffer);
                 return 1;
             }
-            stbir_resize(input_image_buffer, params.width, params.height, 0, 
-                resized_image_buffer, resized_width, resized_height, 0, STBIR_TYPE_UINT8, 
-                3 /*RGB channel*/, STBIR_ALPHA_CHANNEL_NONE, 0,
-                STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP,
-                STBIR_FILTER_BOX, STBIR_FILTER_BOX,
-                STBIR_COLORSPACE_SRGB, nullptr
-            );
+            stbir_resize(input_image_buffer, params.width, params.height, 0,
+                         resized_image_buffer, resized_width, resized_height, 0, STBIR_TYPE_UINT8,
+                         3 /*RGB channel*/, STBIR_ALPHA_CHANNEL_NONE, 0,
+                         STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP,
+                         STBIR_FILTER_BOX, STBIR_FILTER_BOX,
+                         STBIR_COLORSPACE_SRGB, nullptr);
 
             // Save resized result
             free(input_image_buffer);
             input_image_buffer = resized_image_buffer;
-            params.height = resized_height;
-            params.width = resized_width;
+            params.height      = resized_height;
+            params.width       = resized_width;
         }
     }
 

ggml_extend.hpp

@@ -484,6 +484,7 @@ __STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const
     if (tile_count < num_tiles) {
         pretty_progress(num_tiles, num_tiles, last_time);
     }
+    ggml_free(tiles_ctx);
 }
 
 __STATIC_INLINE__ struct ggml_tensor* ggml_group_norm_32(struct ggml_context* ctx,
@@ -751,18 +752,15 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_timestep_embedding(
     return ggml_timestep_embedding(ctx, timesteps, dim, max_period);
 }
 
-// struct GGMLComputeGraph {
-//     virtual void init(struct ggml_context* ctx, ggml_type wtype) = 0;
-//     virtual std::string get_desc() = 0;
-//     virtual size_t get_params_mem_size()   = 0;
-//     virtual size_t get_params_num() = 0;
-//     virtual struct ggml_cgraph* get_ggml_cgraph() = 0;
-// };
 
-/*
-#define MAX_PARAMS_TENSOR_NUM 10240
-#define MAX_GRAPH_SIZE 10240
-*/
+__STATIC_INLINE__ size_t ggml_tensor_num(ggml_context * ctx) {
+    size_t num = 0;
+    for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+        num++;
+    }
+    return num;
+}
+
 /* SDXL with LoRA requires more space */
 #define MAX_PARAMS_TENSOR_NUM 15360
 #define MAX_GRAPH_SIZE 15360
@@ -853,8 +851,6 @@ protected:
     }
 
 public:
-    virtual size_t get_params_mem_size() = 0;
-    virtual size_t get_params_num()      = 0;
     virtual std::string get_desc()       = 0;
 
     GGMLModule(ggml_backend_t backend, ggml_type wtype = GGML_TYPE_F32)
@@ -875,7 +871,7 @@ public:
     }
 
     bool alloc_params_buffer() {
-        size_t num_tensors = get_params_num();
+        size_t num_tensors = ggml_tensor_num(params_ctx);
         params_buffer      = ggml_backend_alloc_ctx_tensors(params_ctx, backend);
         if (params_buffer == NULL) {
             LOG_ERROR("%s alloc params backend buffer failed", get_desc().c_str());
@@ -897,6 +893,13 @@ public:
         }
     }
 
+    size_t get_params_buffer_size() {
+        if (params_buffer != NULL) {
+            return ggml_backend_buffer_get_size(params_buffer);
+        }
+        return 0;
+    }
+
     void free_compute_buffer() {
         if (compute_allocr != NULL) {
             ggml_gallocr_free(compute_allocr);
@@ -967,19 +970,6 @@ public:
 };
 
 class GGMLBlock {
-private:
-    static char temp_buffer[1024 * 1024 * 10];
-    ggml_context* get_temp_ctx() {
-        struct ggml_init_params params;
-        params.mem_size   = sizeof(temp_buffer);
-        params.mem_buffer = temp_buffer;
-        params.no_alloc   = true;
-
-        ggml_context* temp_ctx = ggml_init(params);
-        GGML_ASSERT(temp_ctx != NULL);
-        return temp_ctx;
-    }
-
 protected:
     typedef std::unordered_map<std::string, struct ggml_tensor*> ParameterMap;
     typedef std::unordered_map<std::string, std::shared_ptr<GGMLBlock>> GGMLBlockMap;
@@ -1002,14 +992,6 @@ public:
         init_params(ctx, wtype);
     }
 
-    std::tuple<size_t, size_t> get_params_info(ggml_type wtype) {
-        ggml_context* temp_ctx = get_temp_ctx();
-        init(temp_ctx, wtype);
-        size_t num_tensors = get_params_num();
-        size_t mem_size    = get_params_mem_size();
-        return {num_tensors, mem_size};
-    }
-
     size_t get_params_num() {
         size_t num_tensors = params.size();
         for (auto& pair : blocks) {

lora.hpp

@@ -11,7 +11,7 @@ struct LoraModel : public GGMLModule {
     std::string file_path;
     ModelLoader model_loader;
     bool load_failed = false;
-    bool applied = false;
+    bool applied     = false;
 
     LoraModel(ggml_backend_t backend,
               ggml_type wtype,
@@ -27,14 +27,6 @@ struct LoraModel : public GGMLModule {
         return "lora";
     }
 
-    size_t get_params_num() {
-        return LORA_GRAPH_SIZE;
-    }
-
-    size_t get_params_mem_size() {
-        return model_loader.get_params_mem_size(NULL);
-    }
-
     bool load_from_file(bool filter_tensor = false) {
         LOG_INFO("loading LoRA from '%s'", file_path.c_str());
 
@@ -91,10 +83,15 @@ struct LoraModel : public GGMLModule {
             k_tensor = k_tensor.substr(0, k_pos);
             replace_all_chars(k_tensor, '.', '_');
             // LOG_DEBUG("k_tensor %s", k_tensor.c_str());
-            if (k_tensor == "model_diffusion_model_output_blocks_2_2_conv") {  // fix for SDXL
-                k_tensor = "model_diffusion_model_output_blocks_2_1_conv";
+            std::string lora_up_name = "lora." + k_tensor + ".lora_up.weight";
+            if (lora_tensors.find(lora_up_name) == lora_tensors.end()) {
+                if (k_tensor == "model_diffusion_model_output_blocks_2_2_conv") {
+                    // fix for some sdxl lora, like lcm-lora-xl
+                    k_tensor     = "model_diffusion_model_output_blocks_2_1_conv";
+                    lora_up_name = "lora." + k_tensor + ".lora_up.weight";
+                }
             }
-            std::string lora_up_name   = "lora." + k_tensor + ".lora_up.weight";
+
             std::string lora_down_name = "lora." + k_tensor + ".lora_down.weight";
             std::string alpha_name     = "lora." + k_tensor + ".alpha";
             std::string scale_name     = "lora." + k_tensor + ".scale";

model.cpp

@@ -211,6 +211,8 @@ std::string convert_sdxl_lora_name(std::string tensor_name) {
         {"unet", "model_diffusion_model"},
         {"te2", "cond_stage_model_1_transformer"},
         {"te1", "cond_stage_model_transformer"},
+        {"text_encoder_2", "cond_stage_model_1_transformer"},
+        {"text_encoder", "cond_stage_model_transformer"},
     };
     for (auto& pair_i : sdxl_lora_name_lookup) {
         if (tensor_name.compare(0, pair_i.first.length(), pair_i.first) == 0) {
@@ -446,18 +448,25 @@ std::string convert_tensor_name(const std::string& name) {
         } else {
             new_name = name;
         }
-    } else if (contains(name, "lora_up") || contains(name, "lora_down") || contains(name, "lora.up") || contains(name, "lora.down")) {
+    } else if (contains(name, "lora_up") || contains(name, "lora_down") ||
+               contains(name, "lora.up") || contains(name, "lora.down") ||
+               contains(name, "lora_linear")) {
         size_t pos = new_name.find(".processor");
         if (pos != std::string::npos) {
             new_name.replace(pos, strlen(".processor"), "");
         }
-        pos = new_name.find_last_of('_');
+        pos = new_name.rfind("lora");
         if (pos != std::string::npos) {
-            std::string name_without_network_parts = new_name.substr(0, pos);
-            std::string network_part               = new_name.substr(pos + 1);
+            std::string name_without_network_parts = new_name.substr(0, pos - 1);
+            std::string network_part               = new_name.substr(pos);
             // LOG_DEBUG("%s %s", name_without_network_parts.c_str(), network_part.c_str());
             std::string new_key = convert_diffusers_name_to_compvis(name_without_network_parts, '.');
+            new_key             = convert_sdxl_lora_name(new_key);
             replace_all_chars(new_key, '.', '_');
+            size_t npos = network_part.rfind("_linear_layer");
+            if (npos != std::string::npos) {
+                network_part.replace(npos, strlen("_linear_layer"), "");
+            }
             if (starts_with(network_part, "lora.")) {
                 network_part = "lora_" + network_part.substr(5);
             }

pmid.hpp

@@ -186,16 +186,6 @@ public:
         return "pmid";
     }
 
-    size_t get_params_mem_size() {
-        size_t params_mem_size = id_encoder.get_params_mem_size();
-        return params_mem_size;
-    }
-
-    size_t get_params_num() {
-        size_t params_num = id_encoder.get_params_num();
-        return params_num;
-    }
-
     void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
         id_encoder.get_param_tensors(tensors, prefix);
     }

stable-diffusion.cpp

@@ -42,8 +42,6 @@ const char* sampling_methods_str[] = {
     "LCM",
 };
 
-char GGMLBlock::temp_buffer[1024 * 1024 * 10];
-
 /*================================================== Helper Functions ================================================*/
 
 void calculate_alphas_cumprod(float* alphas_cumprod,
@@ -353,27 +351,27 @@ public:
             // first_stage_model->test();
             // return false;
         } else {
-            size_t clip_params_mem_size = cond_stage_model->get_params_mem_size();
-            size_t unet_params_mem_size = diffusion_model->get_params_mem_size();
+            size_t clip_params_mem_size = cond_stage_model->get_params_buffer_size();
+            size_t unet_params_mem_size = diffusion_model->get_params_buffer_size();
             size_t vae_params_mem_size  = 0;
             if (!use_tiny_autoencoder) {
-                vae_params_mem_size = first_stage_model->get_params_mem_size();
+                vae_params_mem_size = first_stage_model->get_params_buffer_size();
             } else {
                 if (!tae_first_stage->load_from_file(taesd_path)) {
                     return false;
                 }
-                vae_params_mem_size = tae_first_stage->get_params_mem_size();
+                vae_params_mem_size = tae_first_stage->get_params_buffer_size();
             }
             size_t control_net_params_mem_size = 0;
             if (control_net) {
                 if (!control_net->load_from_file(control_net_path)) {
                     return false;
                 }
-                control_net_params_mem_size = control_net->get_params_mem_size();
+                control_net_params_mem_size = control_net->get_params_buffer_size();
             }
             size_t pmid_params_mem_size = 0;
             if (stacked_id) {
-                pmid_params_mem_size = pmid_model->get_params_mem_size();
+                pmid_params_mem_size = pmid_model->get_params_buffer_size();
             }
 
             size_t total_params_ram_size  = 0;
@@ -1610,7 +1608,7 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
     if (sd_ctx->sd->stacked_id && !sd_ctx->sd->pmid_lora->applied) {
         t0 = ggml_time_ms();
         sd_ctx->sd->pmid_lora->apply(sd_ctx->sd->tensors, sd_ctx->sd->n_threads);
-        t1 = ggml_time_ms();
+        t1                             = ggml_time_ms();
         sd_ctx->sd->pmid_lora->applied = true;
         LOG_INFO("pmid_lora apply completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
         if (sd_ctx->sd->free_params_immediately) {

tae.hpp

@@ -200,14 +200,6 @@ struct TinyAutoEncoder : public GGMLModule {
         return "taesd";
     }
 
-    size_t get_params_mem_size() {
-        return taesd.get_params_mem_size();
-    }
-
-    size_t get_params_num() {
-        return taesd.get_params_num();
-    }
-
     bool load_from_file(const std::string& file_path) {
         LOG_INFO("loading taesd from '%s'", file_path.c_str());
         alloc_params_buffer();

unet.hpp

@@ -543,14 +543,6 @@ struct UNetModel : public GGMLModule {
         return "unet";
     }
 
-    size_t get_params_mem_size() {
-        return unet.get_params_mem_size();
-    }
-
-    size_t get_params_num() {
-        return unet.get_params_num();
-    }
-
     void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
         unet.get_param_tensors(tensors, prefix);
     }

vae.hpp

@@ -526,14 +526,6 @@ struct AutoEncoderKL : public GGMLModule {
         return "vae";
     }
 
-    size_t get_params_mem_size() {
-        return ae.get_params_mem_size();
-    }
-
-    size_t get_params_num() {
-        return ae.get_params_num();
-    }
-
     void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
         ae.get_param_tensors(tensors, prefix);
     }