diff --git a/model.cpp b/model.cpp index b7422b98..c41dceea 100644 --- a/model.cpp +++ b/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; diff --git a/pmid.hpp b/pmid.hpp index 7eada014..a4980562 100644 --- a/pmid.hpp +++ b/pmid.hpp @@ -548,237 +548,35 @@ struct PhotoMakerLoraModel : public GGMLModule { float multiplier = 1.0f; std::map lora_tensors; std::string file_path; - int in_channels = 4; - // ModelLoader model_loader; + int in_channels = 4; + 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 transformer_depth = {1, 2, 10}; // only for SDXL - std::vector 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 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)) { - // 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; - } + if (!model_loader.init_from_file(file_path, "pmid.")) { + load_failed = true; } - 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 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; - + return PM_LORA_GRAPH_SIZE; } - size_t calculate_mem_size(ggml_type wtype) { - size_t mem_size = 0; - // 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); - + size_t calculate_mem_size() { + return model_loader.cal_mem_size(NULL); } + 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); diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index 72521912..180bb5e5 100644 --- a/stable-diffusion.cpp +++ b/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) - total_params_size += pmid_model.params_buffer_size; + if(stacked_id){ + 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,