reworked pmid lora

model.cpp

@@ -790,7 +790,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
         }
         TensorStorage tensor_storage(prefix + name, type, ne, n_dims, file_index, ST_HEADER_SIZE_LEN + header_size_ + begin);
         // if(strcmp(prefix.c_str(), "pmid") == 0)
-        // printf(" %s %s \n", (prefix + name).c_str(), ggml_type_name(type));
+        // printf(" %s:%s [%zu, %zu, %zu, %zu]\n", (prefix + name).c_str(), ggml_type_name(type), ne[0], ne[1],ne[2], ne[3]);
         tensor_storage.reverse_ne();
 
         size_t tensor_data_size = end - begin;

pmid.hpp

@@ -548,237 +548,35 @@ struct PhotoMakerLoraModel : public GGMLModule {
     float multiplier = 1.0f;
     std::map<std::string, struct ggml_tensor*> lora_tensors;
     std::string file_path;
-     int in_channels                        = 4;
+    int in_channels                        = 4;
-    // ModelLoader model_loader;
+    ModelLoader model_loader;
     bool load_failed = false;
     int model_channels                     = 320;  // only for SDXL
     int num_heads                          = -1;   // only for SDXL
     int num_head_channels                  = 64;   // only for SDXL
     int context_dim                        = 2048; // only for SDXL
 
-    std::vector<int>  transformer_depth     = {1, 2, 10}; // only for SDXL
-    std::vector<int>  channel_mult          = {1, 2, 4}; // only for SDXL
 
-    struct SpatialTransformerLora{
+    
-        int in_channels;        // mult * model_channels
-        int n_head;             // num_heads
-        int d_head;             // in_channels // n_heads
-        int depth       = 1;    // 1
-        int context_dim = 768;  // hidden_size, 1024 for VERSION_2_x
-          
-
-        struct TransformerLora {
-        
-            // attn1
-            struct ggml_tensor* attn1_q_w_up;  // [in_channels, in_channels]
-            struct ggml_tensor* attn1_q_w_dn;  // [in_channels, in_channels]
-            struct ggml_tensor* attn1_k_w_up;  // [in_channels, in_channels]
-            struct ggml_tensor* attn1_k_w_dn;  // [in_channels, in_channels]
-            struct ggml_tensor* attn1_v_w_up;  // [in_channels, in_channels]
-            struct ggml_tensor* attn1_v_w_dn;  // [in_channels, in_channels]
-
-            struct ggml_tensor* attn1_out_w_up;  // [in_channels, in_channels]
-            struct ggml_tensor* attn1_out_w_dn;  // [in_channels, in_channels]
-
-            // attn2
-            struct ggml_tensor* attn2_q_w_up;  // [in_channels, in_channels]
-            struct ggml_tensor* attn2_q_w_dn;  // [in_channels, in_channels]
-            struct ggml_tensor* attn2_k_w_up;  // [in_channels, context_dim]
-            struct ggml_tensor* attn2_k_w_dn;  // [in_channels, context_dim]
-            struct ggml_tensor* attn2_v_w_up;  // [in_channels, context_dim]
-            struct ggml_tensor* attn2_v_w_dn;  // [in_channels, context_dim]
-
-            struct ggml_tensor* attn2_out_w_up;  // [in_channels, in_channels]
-            struct ggml_tensor* attn2_out_w_dn;  // [in_channels, in_channels]
-
-
-
-        };
-
-        std::vector<TransformerLora> transformers;
-
-        SpatialTransformerLora(int depth = -1)
-        : depth(depth) {
-            if(depth > 0)
-                transformers.resize(depth);
-        }
-
-        int get_num_tensors() {
-            return depth * 16;
-        }
-
-        size_t calculate_mem_size(ggml_type wtype) {
-            size_t mem_size = 0;            
-            // transformer
-            for (auto& transformer : transformers) {
-                mem_size += 12 * ggml_row_size(wtype, in_channels * in_channels);      // attn1_q/k/v/out_w attn2_q/out_w
-                mem_size +=  4 * ggml_row_size(wtype, in_channels * context_dim);      // attn2_k/v_w
-            }
-            return mem_size;
-        }
-
-       void init_params(struct ggml_context* ctx, ggml_type wtype) {
-            // transformer
-            for (auto& transformer : transformers) {
-
-                transformer.attn1_q_w_up = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-                transformer.attn1_q_w_dn = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-                transformer.attn1_k_w_up = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-                transformer.attn1_k_w_dn = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-                transformer.attn1_v_w_up = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-                transformer.attn1_v_w_dn = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-
-                transformer.attn1_out_w_up = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-                transformer.attn1_out_w_dn = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-
-            
-
-                transformer.attn2_q_w_up = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-                transformer.attn2_q_w_dn = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-                transformer.attn2_k_w_up = ggml_new_tensor_2d(ctx, wtype, context_dim, in_channels);
-                transformer.attn2_k_w_dn = ggml_new_tensor_2d(ctx, wtype, context_dim, in_channels);
-                transformer.attn2_v_w_up = ggml_new_tensor_2d(ctx, wtype, context_dim, in_channels);
-                transformer.attn2_v_w_dn = ggml_new_tensor_2d(ctx, wtype, context_dim, in_channels);
-
-                transformer.attn2_out_w_up = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-                transformer.attn2_out_w_dn = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels);
-            
-            }
-        }
-
-
-    };
-
-
-    SpatialTransformerLora down_blocks[3][2];
-    SpatialTransformerLora mid_block;
-    SpatialTransformerLora up_blocks[3][3];
 
 
 
     PhotoMakerLoraModel(const std::string file_path = "")
         : file_path(file_path) {
         name = "photomaker lora";
-        // if (!model_loader.init_from_file(file_path)) {
+        if (!model_loader.init_from_file(file_path, "pmid.")) {
-        //     load_failed = true;
+            load_failed = true;
-        // }
-        int ch = model_channels;
-        for(int i = 0; i < 3; i++){
-            int mult = channel_mult[i];
-            if(i == 0) continue;
-            for(int j = 0; j < 2; j++){
-                ch = mult * model_channels;
-                int n_head = num_heads;
-                int d_head = ch / num_heads;
-                if (num_head_channels != -1) {
-                    d_head = num_head_channels;
-                    n_head = ch / d_head;
-                }
-                down_blocks[i][j] = SpatialTransformerLora(transformer_depth[i]);
-                down_blocks[i][j].in_channels = ch;
-                down_blocks[i][j].n_head      = n_head;
-                down_blocks[i][j].d_head      = d_head;
-                down_blocks[i][j].context_dim = context_dim;
-            }
         }
-        int n_head = num_heads;
-        int d_head = ch / num_heads;
-        if (num_head_channels != -1) {
-            d_head = num_head_channels;
-            n_head = ch / d_head;
-        }
-        mid_block = SpatialTransformerLora(transformer_depth[transformer_depth.size()-1]);
-        mid_block.in_channels = ch;
-        mid_block.n_head      = n_head;
-        mid_block.d_head      = d_head;
-        mid_block.context_dim = context_dim;
-
-
-        for(int i = 2; i >= 0 ; i--){
-            int mult = channel_mult[i];
-            if(i == 0) continue;
-            for(int j = 0; j < 2; j++){
-                ch = mult * model_channels;
-                int n_head = num_heads;
-                int d_head = ch / num_heads;
-                if (num_head_channels != -1) {
-                    d_head = num_head_channels;
-                    n_head = ch / d_head;
-                }
-                up_blocks[i][j] = SpatialTransformerLora(transformer_depth[i]);
-                up_blocks[i][j].in_channels = ch;
-                up_blocks[i][j].n_head      = n_head;
-                up_blocks[i][j].d_head      = d_head;
-                up_blocks[i][j].context_dim = context_dim;
-            }
-        }
-      
-        
     }
 
     size_t get_num_tensors() {
-        size_t num_tensors = 0;         
+        return PM_LORA_GRAPH_SIZE;
-        // return model_loader.cal_mem_size(NULL);
-        for(int i = 0; i < 3; i++){
-            if(i == 0) continue;
-            for(int j = 0; j < 2; j++)               
-                num_tensors += down_blocks[i][j].get_num_tensors();
-        }
-        for(int i = 2; i >= 0 ; i--){
-            if(i == 0) continue;
-            for(int j = 0; j < 2; j++)
-                num_tensors += up_blocks[i][j].get_num_tensors();
-        }
-        num_tensors += mid_block.get_num_tensors();
-        return num_tensors;     
-               
     }
 
-    size_t calculate_mem_size(ggml_type wtype) {
+    size_t calculate_mem_size() {
-        size_t mem_size = 0;         
+        return model_loader.cal_mem_size(NULL);
-        // return model_loader.cal_mem_size(NULL);
-        for(int i = 0; i < 3; i++){
-            if(i == 0) continue;
-            for(int j = 0; j < 2; j++)               
-                mem_size += down_blocks[i][j].calculate_mem_size(wtype);
-        }
-        for(int i = 2; i >= 0 ; i--){
-            if(i == 0) continue;
-            for(int j = 0; j < 2; j++)
-                mem_size += up_blocks[i][j].calculate_mem_size(wtype);
-               
-        }
-        mem_size += mid_block.calculate_mem_size(wtype);
-        return mem_size;       
-    }
-
-
-    void init_params(ggml_type wtype){
-        ggml_allocr* alloc = ggml_allocr_new_from_buffer(params_buffer);
-
-        for(int i = 0; i < 3; i++){
-            if(i == 0) continue;
-            for(int j = 0; j < 2; j++)               
-                down_blocks[i][j].init_params(params_ctx, wtype);
-        }
-        for(int i = 2; i >= 0 ; i--){
-            if(i == 0) continue;
-            for(int j = 0; j < 2; j++)
-                up_blocks[i][j].init_params(params_ctx, wtype);
-        }
-
-        mid_block.init_params(params_ctx, wtype);
-
-        // alloc all tensors linked to this context
-        for (struct ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != NULL; t = ggml_get_next_tensor(params_ctx, t)) {
-            if (t->data == NULL) {
-                ggml_allocr_alloc(alloc, t);
-            }
-        }
-        
-        ggml_allocr_free(alloc); 
-
     }
+   
 
     bool load_from_file(ggml_backend_t backend) {
         if (!alloc_params_buffer(backend)) {
@@ -795,7 +593,14 @@ struct PhotoMakerLoraModel : public GGMLModule {
 
         auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
             const std::string& name = tensor_storage.name;
+            // LOG_INFO("loading LoRA tesnor '%s'", name.c_str());
+            if (!starts_with(name, "pmid.unet")){
+                // LOG_INFO("skipping LoRA tesnor '%s'", name.c_str());
+                return true;
+            }
 
+               
+            // LOG_INFO("loading LoRA tesnor '%s'", name.c_str());
             struct ggml_tensor* real = ggml_new_tensor(params_ctx, tensor_storage.type, tensor_storage.n_dims, tensor_storage.ne);
             ggml_allocr_alloc(alloc, real);
 
@@ -805,7 +610,7 @@ struct PhotoMakerLoraModel : public GGMLModule {
             return true;
         };
 
-        // model_loader.load_tensors(on_new_tensor_cb, backend);
+        model_loader.load_tensors(on_new_tensor_cb, backend);
 
         LOG_DEBUG("finished loaded lora");
         ggml_allocr_free(alloc);

stable-diffusion.cpp

@@ -70,6 +70,7 @@ public:
     UNetModel diffusion_model;
     AutoEncoderKL first_stage_model;
     PhotoMakerIDEncoder pmid_model;
+    PhotoMakerLoraModel pmid_lora;
     bool use_tiny_autoencoder = false;
     bool vae_tiling           = false;
     bool stacked_id           = false;
@@ -176,7 +177,13 @@ public:
         cond_stage_model = FrozenCLIPEmbedderWithCustomWords(version);
         diffusion_model  = UNetModel(version);
         pmid_model = PhotoMakerIDEncoder(version);
-
+        if (id_embeddings_path.size() > 0) {
+            pmid_lora = PhotoMakerLoraModel(id_embeddings_path);
+            if (!pmid_lora.load_from_file(backend)) {
+                LOG_WARN("load photomaker lora tensors from %s failed", id_embeddings_path.c_str());
+                return false;
+            }
+        }
 
         LOG_INFO("Stable Diffusion %s ", model_version_to_str[version]);
         if (wtype == GGML_TYPE_COUNT) {
@@ -218,7 +225,7 @@ public:
         }
 
         if(stacked_id){
-            if (!pmid_model.alloc_params_buffer(backend, model_data_type)){
+            if (!pmid_model.alloc_params_buffer(backend, GGML_TYPE_F32)){
                 LOG_ERROR(" pmid model params buffer allocation failed");
                 return false;
             }
@@ -226,7 +233,6 @@ public:
                  pmid_model.params_buffer_size / 1024.0 / 1024.0);
         }
 
-
         ggml_type vae_type = model_data_type;
         if (version == VERSION_XL) {
             vae_type = GGML_TYPE_F32;  // avoid nan, not work...
@@ -254,12 +260,9 @@ public:
             first_stage_model.map_by_name(tensors, "first_stage_model.");
             
             if(stacked_id){
-            //    printf("preparing memory for pmid \n");  
                pmid_model.init_params(GGML_TYPE_F32);
-            //    printf("done preparing memory for pmid \n");  
+               pmid_model.map_by_name(tensors, "pmid.");               
-               pmid_model.map_by_name(tensors, "pmid.");
             }
-
         }
 
         struct ggml_init_params params;
@@ -296,16 +299,16 @@ public:
                 continue;
             }
 
-            // temporaly ignore lora weights from photomaker
-            // add them later
-            if (stacked_id && (starts_with(name, "pmid.unet"))) {
-                ignore_tensors.insert(name);
-                continue;
-            }
-
-
             tensors_need_to_load.insert(pair);
         }
+        if (stacked_id) {
+            for (auto& pair : pmid_lora.lora_tensors){
+                const std::string& name = pair.first;
+                ignore_tensors.insert(name);  
+            }
+        }
+
+
         bool success = model_loader.load_tensors(tensors_need_to_load, backend, ignore_tensors);
         if (!success) {
             LOG_ERROR("load tensors from model loader failed");
@@ -320,8 +323,9 @@ public:
             cond_stage_model.params_buffer_size +
             diffusion_model.params_buffer_size +
             first_stage_model.params_buffer_size;
-        if(stacked_id)
+        if(stacked_id){
-            total_params_size += pmid_model.params_buffer_size;
+            total_params_size += pmid_model.params_buffer_size;            
+        }
         LOG_INFO("total memory buffer size = %.2fMB (clip %.2fMB, unet %.2fMB, vae %.2fMB)",
                  total_params_size / 1024.0 / 1024.0,
                  cond_stage_model.params_buffer_size / 1024.0 / 1024.0,