fixed soem bugs; now photomaker model weights can be loaded into their tensor buffers

clip.hpp

@@ -863,7 +863,7 @@ struct CLIPVisionModel {
         return mem_size;
     }
 
-    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
+    void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix, const std::string stripped) {
         tensors[prefix + "post_layernorm.bias"]                  = post_ln_b;
         tensors[prefix + "post_layernorm.weight"]                = post_ln_w;
         tensors[prefix + "pre_layrnorm.bias"]                    = pre_ln_b;
@@ -876,7 +876,8 @@ struct CLIPVisionModel {
             std::string name = prefix + "encoder.layers." + std::to_string(i) + ".";
             resblocks[i].map_by_name(tensors, prefix + "encoder.layers." + std::to_string(i) + ".");
         }
-        tensors[prefix + "visual_projection"] = visual_projection;
+        // std::string stripped = prefix.replace(prefix.find("vision_model."), sizeof("vision_model.")-1, ""); 
+        tensors[stripped + "visual_projection.weight"] = visual_projection;
     }   
 
 

model.cpp

@@ -1319,7 +1319,6 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend
                 continue;
             }
             // LOG_DEBUG("%s", tensor_storage.name.c_str());
-
             ggml_tensor* dst_tensor = NULL;
 
             success = on_new_tensor_cb(tensor_storage, &dst_tensor);

pmid.hpp

@@ -99,7 +99,7 @@ struct FuseModule{
     FuseModule(int imb_d): 
         embed_dim(imb_d), 
         mlp1(imb_d*2, imb_d, imb_d, false),
-        mlp2(imb_d*2, imb_d, imb_d, true) {
+        mlp2(imb_d, imb_d, imb_d, true) {
 
         // mlp1 = FuseBlock(embed_dim*2, embed_dim, embed_dim, false);
         // mlp2 = FuseBlock(embed_dim*2, embed_dim, embed_dim, true);
@@ -112,6 +112,8 @@ struct FuseModule{
         ln_w = ggml_new_tensor_1d(ctx, wtype, embed_dim);
         ln_b = ggml_new_tensor_1d(ctx, wtype, embed_dim);
         // alloc all tensors linked to this context
+        mlp1.init_params(ctx, wtype, alloc);
+        mlp2.init_params(ctx, wtype, alloc);
         for (struct ggml_tensor* t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
             if (t->data == NULL) {
                 ggml_allocr_alloc(alloc, t);
@@ -123,8 +125,8 @@ struct FuseModule{
 
         tensors[prefix + "layer_norm.weight"] = ln_w;
         tensors[prefix + "layer_norm.bias"]   = ln_b;      
-        mlp1.map_by_name(tensors, prefix + ".mlp1.");
+        mlp1.map_by_name(tensors, prefix + "mlp1.");
-        mlp2.map_by_name(tensors, prefix + ".mlp2.");
+        mlp2.map_by_name(tensors, prefix + "mlp2.");
 
     }
 
@@ -188,14 +190,15 @@ struct PhotoMakerIDEncoder : public GGMLModule {
         
         vision_model.init_params(params_ctx, backend, wtype, alloc);
         fuse_module.init_params(params_ctx, wtype, alloc);
-        visual_projection_2 = ggml_new_tensor_2d(params_ctx, wtype, 1280, 1024); // python [1024, 1280]
+        // visual_projection_2 = ggml_new_tensor_2d(params_ctx, wtype, 1280, 1024); // python [1024, 1280]
+        visual_projection_2 = ggml_new_tensor_2d(params_ctx, wtype, 1024, 1280); // python [1024, 1280]
         ggml_allocr_alloc(alloc, visual_projection_2);
         ggml_allocr_free(alloc); 
     }
 
     void map_by_name(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
-        // vision_model.
+        vision_model.map_by_name(tensors, prefix + "vision_model.", prefix);
-        fuse_module.map_by_name(tensors, prefix + ".fuse_module.");
+        fuse_module.map_by_name(tensors, prefix + "fuse_module.");
         tensors[prefix + "visual_projection_2.weight"] = visual_projection_2;        
     }
 

stable-diffusion.cpp

@@ -202,7 +202,7 @@ public:
 
         if (id_embeddings_path.size() > 0) {
             LOG_INFO("loading stacked ID embedding (PHOTOMAKER) model file from '%s'", id_embeddings_path.c_str());
-            if (!model_loader.init_from_file(id_embeddings_path)) {                
+            if (!model_loader.init_from_file(id_embeddings_path, "pmid.")) {                
                 LOG_WARN("loading stacked ID embedding from '%s' failed", id_embeddings_path.c_str());
             }
             else{
@@ -210,6 +210,14 @@ public:
             }
         }
 
+        if(stacked_id){
+            if (!pmid_model.alloc_params_buffer(backend, model_data_type)){
+                LOG_ERROR(" pmid model params buffer allocation failed");
+                return false;
+            }
+        }
+
+
         ggml_type vae_type = model_data_type;
         if (version == VERSION_XL) {
             vae_type = GGML_TYPE_F32;  // avoid nan, not work...
@@ -235,6 +243,11 @@ public:
                 first_stage_model.init_params();
             }
             first_stage_model.map_by_name(tensors, "first_stage_model.");
+
+            if(stacked_id){
+               pmid_model.init_params();
+               pmid_model.map_by_name(tensors, "pmid.");
+            }
         }
 
         struct ggml_init_params params;
@@ -252,6 +265,7 @@ public:
 
         // load weights
         LOG_DEBUG("loading weights");
+        fprintf(stderr, "%s: loading weights \n", __func__);
         int64_t t0 = ggml_time_ms();
 
         std::map<std::string, struct ggml_tensor*> tensors_need_to_load;
@@ -270,6 +284,14 @@ 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);
         }
         bool success = model_loader.load_tensors(tensors_need_to_load, backend, ignore_tensors);
@@ -278,6 +300,7 @@ public:
             ggml_free(ctx);
             return false;
         }
+        fprintf(stderr, "%s: loading weights sccess\n", __func__);
 
         // LOG_DEBUG("model size = %.2fMB", total_size / 1024.0 / 1024.0);