From a6a676b646628ab6cac9c0e012526621ede738e7 Mon Sep 17 00:00:00 2001 From: bssrdf Date: Mon, 19 Feb 2024 12:49:41 -0500 Subject: [PATCH] add style_ratio as a cli param; reworked encode with trigger for attention weights --- clip.hpp | 64 ++++++++++++++++++++++++------------------- examples/cli/main.cpp | 10 +++++++ ggml | 2 +- lora.hpp | 5 +++- pmid.hpp | 19 ++++++------- stable-diffusion.cpp | 5 ++-- stable-diffusion.h | 1 + 7 files changed, 63 insertions(+), 43 deletions(-) diff --git a/clip.hpp b/clip.hpp index 6d9c9890..0da75169 100644 --- a/clip.hpp +++ b/clip.hpp @@ -976,21 +976,21 @@ struct CLIPVisionModel { // patch_embeddings_f16 = ggml_reshape_4d(ctx0, patch_embeddings_f16, ne[0], ne[1], ne[2], ne[3]); struct ggml_tensor *patch_embeddings_f16 = ggml_cast(ctx0, patch_embeddings, GGML_TYPE_F16); ggml_set_name(patch_embeddings_f16, "patch_embeddings_f16"); - print_ggml_tensor(patch_embeddings_f16, true, "patch_embeddings_f16"); + // print_ggml_tensor(patch_embeddings_f16, true, "patch_embeddings_f16"); struct ggml_tensor * inp = ggml_conv_2d(ctx0, patch_embeddings_f16, x, patch_size, patch_size, 0, 0, 1, 1); ggml_set_name(inp, "inp_conv_2d"); - print_ggml_tensor(inp, true, "inp_conv_2d"); + // print_ggml_tensor(inp, true, "inp_conv_2d"); inp = ggml_reshape_3d(ctx0, inp, num_patches, hidden_size, batch_size); ggml_set_name(inp, "inp_reshape_3d"); - print_ggml_tensor(inp, true, "inp_reshape_3d"); + // print_ggml_tensor(inp, true, "inp_reshape_3d"); // inp = ggml_cont(ctx0, ggml_permute(ctx0, inp, 1, 0, 2, 3)); // print_ggml_tensor(ggml_permute(ctx0, inp, 2, 0, 1, 3), true, "inp_permute"); inp = ggml_cont(ctx0, ggml_permute(ctx0, inp, 2, 0, 1, 3)); - print_ggml_tensor(inp, true, "inp_cont_perm+cont"); + // print_ggml_tensor(inp, true, "inp_cont_perm+cont"); ggml_set_name(inp, "inp_cont"); // ggml_set_name(class_embedding, "class_embedding"); @@ -1000,18 +1000,18 @@ struct CLIPVisionModel { // struct ggml_tensor * temp = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, hidden_size, 1, batch_size); // ggml_tensor *class_embedding_rep = ggml_add(ctx0, cast_f32_class, class_embedding); // ggml_set_name(class_embedding_rep, "add_class_embedding_to_zero"); - print_ggml_tensor(class_embedding, true, "model.class_embedding"); + // print_ggml_tensor(class_embedding, true, "model.class_embedding"); // print_ggml_tensor(class_embedding_rep, true, "class_embedding_rep_bef_repeat"); - print_ggml_tensor(temp, true, "temp"); + // print_ggml_tensor(temp, true, "temp"); ggml_tensor *class_embedding_rep = ggml_repeat(ctx0, class_embedding, temp); ggml_set_name(class_embedding_rep, "class_embedding_rep"); - print_ggml_tensor(class_embedding_rep, true, "class_embedding_rep"); + // print_ggml_tensor(class_embedding_rep, true, "class_embedding_rep"); // class_embedding_rep = ggml_cast(ctx0, class_embedding_rep, inp->type); // print_ggml_tensor(class_embedding_rep, true, "class_embedding_rep_aft_casting"); struct ggml_tensor *embeddings = ggml_concat(ctx0, class_embedding_rep, inp); ggml_set_name(embeddings, "embeddings_after_concat"); - print_ggml_tensor(embeddings, true, "embeddings_after_concat"); + // print_ggml_tensor(embeddings, true, "embeddings_after_concat"); // print_ggml_tensor(ggml_permute(ctx0, embeddings, 0, 3, 1, 2), true, "embeddings_after_concat_permute"); embeddings = ggml_cont(ctx0, ggml_permute(ctx0, embeddings, 0, 2, 1, 3)); ggml_set_name(embeddings, "embeddings_after_permute"); @@ -1028,7 +1028,7 @@ struct CLIPVisionModel { // ggml_add(ctx0, embeddings, ggml_repeat(ctx0, ggml_get_rows(ctx0, position_embeddings, positions), embeddings)); ggml_add(ctx0, embeddings, ggml_get_rows(ctx0, position_embeddings, positions)); ggml_set_name(embeddings, "embeddings_to_transformer"); - print_ggml_tensor(embeddings, true, "embeddings_to_transformer"); + // print_ggml_tensor(embeddings, true, "embeddings_to_transformer"); // pre-layernorm embeddings = ggml_nn_layer_norm(ctx0, embeddings, pre_ln_w, pre_ln_w); @@ -1039,12 +1039,12 @@ struct CLIPVisionModel { embeddings = resblocks[i].forward(ctx0, embeddings); // [N, n_token, hidden_size] } ggml_set_name(embeddings, "embeddings_after_transformer"); - print_ggml_tensor(embeddings, true, "embeddings_after_transformer"); + // print_ggml_tensor(embeddings, true, "embeddings_after_transformer"); // get the output of cls token, e.g., 0th index embeddings = ggml_get_rows(ctx0, ggml_reshape_2d(ctx0, embeddings, hidden_size, num_positions * batch_size), cls); ggml_set_name(embeddings, "embeddings_pooled_after_encoder"); - print_ggml_tensor(embeddings, true, "embeddings_pooled_after_encoder"); + // print_ggml_tensor(embeddings, true, "embeddings_pooled_after_encoder"); // post-layernorm embeddings = ggml_nn_layer_norm(ctx0, embeddings, post_ln_w, post_ln_b); @@ -1291,11 +1291,13 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule { std::vector tokens; std::vector weights; std::vector class_token_mask; - std::vector clean_input_ids; - std::vector class_token_index; + int32_t class_idx = -1, tokens_acc = 0; for (const auto& item : parsed_attention) { + std::vector class_token_index; + std::vector clean_input_ids; const std::string& curr_text = item.first; float curr_weight = item.second; + // printf(" %s: %f \n", curr_text.c_str(), curr_weight); std::vector curr_tokens = tokenizer.encode(curr_text, on_new_token_cb); int32_t clean_index = 0; for(uint32_t i = 0; i < curr_tokens.size(); i++){ @@ -1307,23 +1309,29 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule { clean_index++; } } - GGML_ASSERT(class_token_index.size() == 1); // PhotoMaker currently does not support multiple + //GGML_ASSERT(class_token_index.size() == 1); // PhotoMaker currently does not support multiple // trigger words in a single prompt. - // Expand the class word token and corresponding mask - int class_token = clean_input_ids[class_token_index[0]]; - std::vector clean_input_ids_tmp; - for(uint32_t i = 0; i < class_token_index[0]; i++) - clean_input_ids_tmp.push_back(clean_input_ids[i]); - for(uint32_t i = 0; i < num_input_imgs; i++) - clean_input_ids_tmp.push_back(class_token); - for(uint32_t i = class_token_index[0]+1; i < clean_input_ids.size(); i++) - clean_input_ids_tmp.push_back(clean_input_ids[i]); - clean_input_ids.clear(); - clean_input_ids = clean_input_ids_tmp; - + if(class_token_index.size() == 1){ + // Expand the class word token and corresponding mask + int class_token = clean_input_ids[class_token_index[0]]; + class_idx = tokens_acc + class_token_index[0]; + std::vector clean_input_ids_tmp; + for(uint32_t i = 0; i < class_token_index[0]; i++) + clean_input_ids_tmp.push_back(clean_input_ids[i]); + for(uint32_t i = 0; i < num_input_imgs; i++) + clean_input_ids_tmp.push_back(class_token); + for(uint32_t i = class_token_index[0]+1; i < clean_input_ids.size(); i++) + clean_input_ids_tmp.push_back(clean_input_ids[i]); + clean_input_ids.clear(); + clean_input_ids = clean_input_ids_tmp; + } + tokens_acc += clean_index; // class_tokens_mask = [True if class_token_index <= i < class_token_index+num_id_images else False \ // for i in range(len(clean_input_ids))] - + // printf("["); + // for(int i = 0; i < clean_input_ids.size(); ++i) + // printf("%d, ", clean_input_ids[i]); + // printf("]\n"); tokens.insert(tokens.end(), clean_input_ids.begin(), clean_input_ids.end()); weights.insert(weights.end(), clean_input_ids.size(), curr_weight); } @@ -1351,7 +1359,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule { } for(uint32_t i = 0; i < tokens.size(); i++){ - if(class_token_index[0]+1 <= i && i < class_token_index[0]+1+num_input_imgs) + if(class_idx+1 <= i && i < class_idx+1+num_input_imgs) class_token_mask.push_back(true); else class_token_mask.push_back(false); diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index b1bec31a..cc95ce1f 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -74,6 +74,7 @@ struct SDParams { std::string prompt; std::string negative_prompt; float cfg_scale = 7.0f; + float style_ratio = 20.f; int clip_skip = -1; // <= 0 represents unspecified int width = 512; int height = 512; @@ -105,6 +106,7 @@ void print_params(SDParams params) { printf(" embeddings_path: %s\n", params.embeddings_path.c_str()); printf(" stacked_id_embeddings_path: %s\n", params.stacked_id_embeddings_path.c_str()); printf(" input_id_images_path: %s\n", params.input_id_images_path.c_str()); + printf(" style ratio: %.2f\n", params.style_ratio); printf(" output_path: %s\n", params.output_path.c_str()); printf(" init_img: %s\n", params.input_path.c_str()); printf(" control_image: %s\n", params.control_image_path.c_str()); @@ -152,6 +154,7 @@ void print_usage(int argc, const char* argv[]) { printf(" -n, --negative-prompt PROMPT the negative prompt (default: \"\")\n"); printf(" --cfg-scale SCALE unconditional guidance scale: (default: 7.0)\n"); printf(" --strength STRENGTH strength for noising/unnoising (default: 0.75)\n"); + printf(" --style-ratio STYLE-RATIO strength for keeping input identity (default: 20%%)\n"); printf(" --control-strength STRENGTH strength to apply Control Net (default: 0.9)\n"); printf(" 1.0 corresponds to full destruction of information in init image\n"); printf(" -H, --height H image height, in pixel space (default: 512)\n"); @@ -322,6 +325,12 @@ void parse_args(int argc, const char** argv, SDParams& params) { break; } params.strength = std::stof(argv[i]); + }else if (arg == "--style-ratio") { + if (++i >= argc) { + invalid_arg = true; + break; + } + params.style_ratio = std::stof(argv[i]); } else if (arg == "--control-strength") { if (++i >= argc) { invalid_arg = true; @@ -664,6 +673,7 @@ int main(int argc, const char* argv[]) { params.batch_count, control_image, params.control_strength, + params.style_ratio, input_id_images); } else { sd_image_t input_image = {(uint32_t)params.width, diff --git a/ggml b/ggml index d2545e80..93f0fde5 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit d2545e80469ab8e66341d4d389c76f67c7a69610 +Subproject commit 93f0fde5ec4d00998b81d3aba8617ae4104d84db diff --git a/lora.hpp b/lora.hpp index 5587b3af..e8b15ac1 100644 --- a/lora.hpp +++ b/lora.hpp @@ -46,7 +46,6 @@ struct LoraModel : public GGMLModule { struct ggml_tensor* real = ggml_new_tensor(params_ctx, tensor_storage.type, tensor_storage.n_dims, tensor_storage.ne); ggml_allocr_alloc(alloc, real); - *dst_tensor = real; lora_tensors[name] = real; @@ -86,6 +85,10 @@ struct LoraModel : public GGMLModule { continue; } k_tensor = k_tensor.substr(0, k_pos); + // some lora models use 'unet' instead of "model.diffusion_model" + // TODO: make both work + if(starts_with(k_tensor, "model.diffusion_model")) + k_tensor.replace(0, strlen("model.diffusion_model"), "unet"); 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"; diff --git a/pmid.hpp b/pmid.hpp index 31e049de..e59d3fcc 100644 --- a/pmid.hpp +++ b/pmid.hpp @@ -205,20 +205,17 @@ struct FuseModule{ struct ggml_tensor * valid_id_embeds = id_embeds; // # slice out the image token embeddings struct ggml_tensor * image_token_embeds = ggml_get_rows(ctx, prompt_embeds, class_tokens_mask_pos); - print_ggml_tensor(image_token_embeds, true, "image_token_embeds"); + // print_ggml_tensor(image_token_embeds, true, "image_token_embeds"); ggml_set_name(image_token_embeds, "image_token_embeds"); - print_ggml_tensor(valid_id_embeds, true, "valid_id_embeds"); + // print_ggml_tensor(valid_id_embeds, true, "valid_id_embeds"); ggml_set_name(valid_id_embeds, "valid_id_embeds"); struct ggml_tensor *stacked_id_embeds = fuse_fn(ctx, image_token_embeds, valid_id_embeds); - print_ggml_tensor(stacked_id_embeds, true, "stacked_id_embeds_before_concat"); + // print_ggml_tensor(stacked_id_embeds, true, "stacked_id_embeds_before_concat"); ggml_set_name(stacked_id_embeds, "stacked_id_embeds_fuse_fn"); stacked_id_embeds = ggml_cont(ctx, ggml_permute(ctx, stacked_id_embeds, 0, 2, 1, 3)); - print_ggml_tensor(stacked_id_embeds, true, "stacked_id_embeds_after_permute"); if(left && right){ - print_ggml_tensor(left, true, "left"); - print_ggml_tensor(right, true, "right"); stacked_id_embeds = ggml_concat(ctx, left, stacked_id_embeds); stacked_id_embeds = ggml_concat(ctx, stacked_id_embeds, right); }else if(left){ @@ -261,10 +258,12 @@ struct PhotoMakerIDEncoder : public GGMLModule { CLIPVisionModel vision_model; FuseModule fuse_module; struct ggml_tensor* visual_projection_2; + float style_strength; - PhotoMakerIDEncoder(SDVersion version = VERSION_XL) + PhotoMakerIDEncoder(SDVersion version = VERSION_XL, float sty = 20.f) : version(version), - fuse_module(2048) { + fuse_module(2048), + style_strength(sty){ vision_model = CLIPVisionModel(); // fuse_module = FuseModule(2048); // wtype = GGML_TYPE_F32; @@ -341,7 +340,7 @@ struct PhotoMakerIDEncoder : public GGMLModule { class_embedding_temp, positions ); // [batch_size, seq_length, hidden_size] - print_ggml_tensor(shared_id_embeds, true, "shared_id_embeds_from_vision"); + // print_ggml_tensor(shared_id_embeds, true, "shared_id_embeds_from_vision"); ggml_set_name(shared_id_embeds, "shared_id_embeds_from_vision"); // if(class_tokens_mask->backend == GGML_BACKEND_GPU){ @@ -359,8 +358,6 @@ struct PhotoMakerIDEncoder : public GGMLModule { // id_embeds_2 = id_embeds_2.view(b, num_inputs, 1, -1) // id_embeds = torch.cat((id_embeds, id_embeds_2), dim=-1) - print_ggml_tensor(id_embeds, true, "id_embeds_proj1"); - print_ggml_tensor(id_embeds_2, true, "id_embeds_proj2"); ggml_set_name(id_embeds, "id_embeds_proj1"); ggml_set_name(id_embeds_2, "id_embeds_proj2"); diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index c38d17bf..9451b0eb 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -1445,6 +1445,7 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx, int batch_count, const sd_image_t* control_cond, float control_strength, + float style_ratio, std::vector &input_id_images) { LOG_DEBUG("txt2img %dx%d", width, height); if (sd_ctx == NULL) { @@ -1507,6 +1508,7 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx, // ggml_tensor* class_tokens_mask = NULL; std::vector class_tokens_mask; if(sd_ctx->sd->stacked_id){ + sd_ctx->sd->pmid_model.style_strength = style_ratio; int32_t w = input_id_images[0]->width; int32_t h = input_id_images[0]->height; int32_t channels = input_id_images[0]->channel; @@ -1630,8 +1632,7 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx, int start_merge_step = -1; if(sd_ctx->sd->stacked_id){ - float style_strength_ratio = 20.f; - start_merge_step = int(style_strength_ratio / 100.f * sample_steps); + start_merge_step = int(sd_ctx->sd->pmid_model.style_strength / 100.f * sample_steps); if(start_merge_step > 30) start_merge_step = 30; LOG_INFO("PHOTOMAKER: start_merge_step: %d", start_merge_step); diff --git a/stable-diffusion.h b/stable-diffusion.h index 92e32681..4a7c518a 100644 --- a/stable-diffusion.h +++ b/stable-diffusion.h @@ -133,6 +133,7 @@ SD_API sd_image_t* txt2img(sd_ctx_t* sd_ctx, int batch_count, const sd_image_t* control_cond, float control_strength, + float style_strength, std::vector &input_id_images); SD_API sd_image_t* img2img(sd_ctx_t* sd_ctx,