diff --git a/pmid.hpp b/pmid.hpp index 3e4fca29..fc6585df 100644 --- a/pmid.hpp +++ b/pmid.hpp @@ -547,6 +547,7 @@ struct PhotoMakerIDEncoder : public GGMLModule { struct PhotoMakerLoraModel : public GGMLModule { float multiplier = 1.0f; std::map lora_tensors; + std::vector lora_tensors_to_be_ignored; std::string file_path; ModelLoader model_loader; bool load_failed = false; @@ -582,7 +583,7 @@ struct PhotoMakerLoraModel : public GGMLModule { ggml_allocr* alloc = ggml_allocr_new_from_buffer(params_buffer); auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool { - const std::string& name = tensor_storage.name; + 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()); @@ -595,7 +596,12 @@ struct PhotoMakerLoraModel : public GGMLModule { ggml_allocr_alloc(alloc, real); *dst_tensor = real; - + lora_tensors_to_be_ignored.push_back(name); + size_t k_pos = name.find(".processor"); + if(k_pos != std::string::npos) + name.replace(k_pos, strlen(".processor"), ""); + // if(starts_with(name, "pmid.unet.down_blocks.2.attentions.1.transformer_blocks.9.attn2")) + // print_ggml_tensor(real, true, name.c_str()); lora_tensors[name] = real; return true; }; @@ -607,6 +613,28 @@ struct PhotoMakerLoraModel : public GGMLModule { return true; } + std::pair find_ij0(int n){ + int i, j; + for(i = 0; i < 3; i++){ + for(j = 0; j < 2; j++){ + if((i*3+j+1) == n) + return {i,j}; + } + } + return {-1, -1}; + } + + std::pair find_ij(int n){ + int i, j; + for(i = 0; i < 2; i++){ + for(j = 0; j < 3; j++){ + if((i*3+j) == n) + return {i,j}; + } + } + return {-1, -1}; + } + struct ggml_cgraph* build_graph(std::map model_tensors) { // make a graph to compute all lora, expected lora and models tensors are in the same backend // since we are using ggml-alloc, this buffer only needs enough space to hold the ggml_tensor and ggml_cgraph structs, but not the tensor data @@ -627,15 +655,81 @@ struct PhotoMakerLoraModel : public GGMLModule { for (auto it : model_tensors) { std::string k_tensor = it.first; struct ggml_tensor* weight = model_tensors[it.first]; - - size_t k_pos = k_tensor.find(".weight"); + std::string full_name = k_tensor; + // size_t k_pos = k_tensor.find(".weight"); + size_t k_pos = k_tensor.find(".attn1"); if (k_pos == std::string::npos) { - continue; + k_pos = k_tensor.find(".attn2"); + if (k_pos == std::string::npos) { + continue; + } } - k_tensor = k_tensor.substr(0, k_pos); - replace_all_chars(k_tensor, '.', '_'); - std::string lora_up_name = "lora." + k_tensor + ".lora_up.weight"; - std::string lora_down_name = "lora." + k_tensor + ".lora_down.weight"; + if(ends_with(k_tensor, "bias")) + continue; + int block_kind = -1; + int block_id = -1; + if ((k_pos = k_tensor.find("input_blocks")) != std::string::npos) { + block_id = atoi(k_tensor.substr(k_pos+strlen("input_blocks")+1).c_str()); + block_kind = 0; // input -> down block + + }else if ((k_pos = k_tensor.find("output_blocks")) != std::string::npos) { + block_id = atoi(k_tensor.substr(k_pos+strlen("output_blocks")+1).c_str()); + block_kind = 1; // output -> up block + }else{ + k_pos = k_tensor.find("transformer_blocks"); + block_id = atoi(k_tensor.substr(k_pos-4,1).c_str()); + block_kind = 2; // middle block + } + + + std::string lora_up_name; + std::string lora_down_name; + std::string prefix = "pmid.unet"; + if (block_kind == 0){ + prefix = prefix + ".down_blocks"; + k_pos = k_tensor.find(".weight"); + k_tensor = k_tensor.substr(0, k_pos); + k_pos = k_tensor.find("transformer_blocks"); + k_tensor = k_tensor.substr(k_pos); + if(ends_with(k_tensor, "0")){ + k_tensor = k_tensor.substr(0, k_tensor.length()-2); + } + std::pair ij = find_ij0(block_id); + if(ij.first == -1) + continue; + prefix = prefix + "."+std::to_string(ij.first)+".attentions."+std::to_string(ij.second)+"."; + lora_up_name = prefix + k_tensor + "_lora.up.weight"; + lora_down_name = prefix + k_tensor + "_lora.down.weight"; + }else if (block_kind == 1){ + prefix = prefix + ".up_blocks"; + k_pos = k_tensor.find(".weight"); + k_tensor = k_tensor.substr(0, k_pos); + k_pos = k_tensor.find("transformer_blocks"); + k_tensor = k_tensor.substr(k_pos); + if(ends_with(k_tensor, "0")){ + k_tensor = k_tensor.substr(0, k_tensor.length()-2); + } + std::pair ij = find_ij(block_id); + if(ij.first == -1) + continue; + prefix = prefix + "."+std::to_string(ij.first)+".attentions."+std::to_string(ij.second)+"."; + lora_up_name = prefix + k_tensor + "_lora.up.weight"; + lora_down_name = prefix + k_tensor + "_lora.down.weight"; + }else{ + prefix = prefix + ".mid_block" + ".attentions.0."; + k_pos = k_tensor.find(".weight"); + k_tensor = k_tensor.substr(0, k_pos); + k_pos = k_tensor.find("transformer_blocks"); + k_tensor = k_tensor.substr(k_pos); + if(ends_with(k_tensor, "0")){ + k_tensor = k_tensor.substr(0, k_tensor.length()-2); + } + lora_up_name = prefix + k_tensor + "_lora.up.weight"; + lora_down_name = prefix + k_tensor + "_lora.down.weight"; + } + // LOG_INFO("unet transformer tensor: %s ", full_name.c_str()); + // LOG_INFO("corresponding up tensor: %s ", lora_up_name.c_str()); + // LOG_INFO("corresponding dn tensor: %s ", lora_down_name.c_str()); ggml_tensor* lora_up = NULL; ggml_tensor* lora_down = NULL; @@ -649,14 +743,17 @@ struct PhotoMakerLoraModel : public GGMLModule { } if (lora_up == NULL || lora_down == NULL) { + LOG_WARN("can not find: %s and %s k,id = (%d, %d)", lora_down_name.c_str(), lora_up_name.c_str(), block_kind, block_id); continue; } + ggml_tensor* lora_up_orig = lora_up; + applied_lora_tensors.insert(lora_up_name); applied_lora_tensors.insert(lora_down_name); // ggml_mul_mat requires tensor b transposed - lora_up = ggml_cont(ctx0, ggml_transpose(ctx0, lora_up)); + lora_down = ggml_cont(ctx0, ggml_transpose(ctx0, lora_down)); struct ggml_tensor* updown = ggml_mul_mat(ctx0, lora_down, lora_up); updown = ggml_cont(ctx0, updown); GGML_ASSERT(ggml_nelements(updown) == ggml_nelements(weight)); diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index 180bb5e5..56219dd9 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -302,8 +302,7 @@ public: tensors_need_to_load.insert(pair); } if (stacked_id) { - for (auto& pair : pmid_lora.lora_tensors){ - const std::string& name = pair.first; + for (auto& name : pmid_lora.lora_tensors_to_be_ignored){ ignore_tensors.insert(name); } } @@ -1458,6 +1457,15 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx, sd_ctx->sd->apply_loras(lora_f2m); int64_t t1 = ggml_time_ms(); LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); + + if(sd_ctx->sd->stacked_id){ + sd_ctx->sd->pmid_lora.apply(sd_ctx->sd->tensors, sd_ctx->sd->n_threads); + sd_ctx->sd->pmid_lora.free_compute_buffer(); + if (sd_ctx->sd->free_params_immediately) { + sd_ctx->sd->pmid_lora.free_params_buffer(); + } + } + struct ggml_init_params params; params.mem_size = static_cast(10 * 1024 * 1024); // 10 MB params.mem_size += width * height * 3 * sizeof(float);