From 5a13b48b6e1dfba7bbd01414200b4b85f48d24f3 Mon Sep 17 00:00:00 2001 From: bssrdf Date: Sun, 11 Feb 2024 12:54:18 -0500 Subject: [PATCH] fixed wtype issue; input images can only be 1 because issue with transformer when batch size > 1 (to be investigated) --- clip.hpp | 57 ++++++++++++++++++++++-------------- ggml_extend.hpp | 9 +++--- pmid.hpp | 70 +++++++++++++++++++++----------------------- stable-diffusion.cpp | 38 +++++++----------------- 4 files changed, 83 insertions(+), 91 deletions(-) diff --git a/clip.hpp b/clip.hpp index 39e869c8..38fadd8e 100644 --- a/clip.hpp +++ b/clip.hpp @@ -908,7 +908,9 @@ struct CLIPVisionModel { } size_t calculate_mem_size(ggml_type wtype) { - size_t mem_size = 0; + + size_t mem_size = 0; + for (int i = 0; i < num_hidden_layers; i++) { mem_size += resblocks[i].calculate_mem_size(wtype); } @@ -917,7 +919,6 @@ struct CLIPVisionModel { mem_size += ggml_row_size(wtype, hidden_size*patch_size*patch_size*3); // mem_size += ggml_row_size(wtype, hidden_size*257); // mem_size += ggml_row_size(wtype, hidden_size * projection_dim); // visual_projection - return mem_size; } @@ -961,19 +962,23 @@ struct CLIPVisionModel { ggml_set_name(x, "id_pixel"); ggml_set_name(temp, "temp_input"); + int64_t* ne = patch_embeddings->ne; + struct ggml_tensor *patch_embeddings_f16 = ggml_reshape_3d(ctx0, patch_embeddings, ne[0], ne[1], ne[2]*ne[3]); + patch_embeddings_f16 = ggml_cast(ctx0, patch_embeddings_f16, GGML_TYPE_F16); + patch_embeddings_f16 = ggml_reshape_4d(ctx0, patch_embeddings_f16, ne[0], ne[1], ne[2], ne[3]); - struct ggml_tensor * inp = ggml_conv_2d(ctx0, patch_embeddings, x, patch_size, patch_size, 0, 0, 1, 1); + 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"); + // 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"); + // print_ggml_tensor(inp, true, "inp_cont"); ggml_set_name(inp, "inp_cont"); ggml_set_name(class_embedding, "class_embedding"); @@ -983,18 +988,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"); + // 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"); @@ -1004,13 +1009,13 @@ struct CLIPVisionModel { // struct ggml_tensor * positions = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, num_positions); // print_ggml_tensor(embeddings, true, "embeddings_bef_add"); - print_ggml_tensor(ggml_get_rows(ctx0, position_embeddings, positions), true, "embeddings_add_src1"); + // print_ggml_tensor(ggml_get_rows(ctx0, position_embeddings, positions), true, "embeddings_add_src1"); // print_ggml_tensor(ggml_repeat(ctx0, ggml_get_rows(ctx0, position_embeddings, positions), embeddings), true, "embeddings_add_src1"); // embeddings = ggml_cont(ctx0, ggml_permute(ctx0, embeddings, 0, 1, 3, 2)); embeddings = ggml_add(ctx0, embeddings, ggml_repeat(ctx0, ggml_get_rows(ctx0, position_embeddings, positions), embeddings)); ggml_set_name(embeddings, "embeddings_after_add"); - print_ggml_tensor(embeddings, true, "embeddings_after_add"); + // print_ggml_tensor(embeddings, true, "embeddings_after_add"); // pre-layernorm embeddings = ggml_nn_layer_norm(ctx0, embeddings, pre_ln_w, pre_ln_w); @@ -1040,25 +1045,33 @@ struct CLIPVisionModel { } void init_params(ggml_context* ctx, ggml_backend_t backend, ggml_type wtype, ggml_allocr* alloc) { - wtype = GGML_TYPE_F32; - class_embedding = ggml_new_tensor_1d(ctx, wtype, hidden_size); - - patch_embeddings = ggml_new_tensor_4d(ctx, wtype, patch_size, patch_size, 3, hidden_size); - - position_embeddings = ggml_new_tensor_2d(ctx, wtype, hidden_size, 257); - + for (int i = 0; i < num_hidden_layers; i++) { resblocks[i].init_params(ctx, alloc, wtype); } + class_embedding = ggml_new_tensor_1d(ctx, wtype, hidden_size); + ggml_set_name(class_embedding, "class_embedding"); + + patch_embeddings = ggml_new_tensor_4d(ctx, wtype, patch_size, patch_size, 3, hidden_size); + ggml_set_name(patch_embeddings, "patch_embeddings"); + + position_embeddings = ggml_new_tensor_2d(ctx, wtype, hidden_size, 257); + ggml_set_name(position_embeddings, "position_embeddings"); + + pre_ln_w = ggml_new_tensor_1d(ctx, wtype, hidden_size); + ggml_set_name(pre_ln_w, "vision.pre_ln_w"); pre_ln_b = ggml_new_tensor_1d(ctx, wtype, hidden_size); + ggml_set_name(pre_ln_b, "vision.pre_ln_b"); post_ln_w = ggml_new_tensor_1d(ctx, wtype, hidden_size); + ggml_set_name(post_ln_w, "vision.post_ln_w"); post_ln_b = ggml_new_tensor_1d(ctx, wtype, hidden_size); + ggml_set_name(post_ln_b, "vision.post_ln_b"); visual_projection = ggml_new_tensor_2d(ctx, wtype, hidden_size, projection_dim); // Check!!! + ggml_set_name(visual_projection, "vision.visual_projection"); - // alloc all tensors linked to this context for (struct ggml_tensor* t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { if (t->data == NULL) { diff --git a/ggml_extend.hpp b/ggml_extend.hpp index a9e5c19b..ce184a83 100644 --- a/ggml_extend.hpp +++ b/ggml_extend.hpp @@ -651,19 +651,18 @@ struct GGMLModule { ggml_backend_metal_set_n_cb(backend, n_threads); } #endif - printf("about to do ggml_backend_graph_compute \n"); ggml_backend_graph_compute(backend, gf); #ifdef GGML_PERF ggml_graph_print(gf); #endif - if (output != NULL) - print_ggml_tensor(output, true, "output_before_get"); + // if (output != NULL) + // print_ggml_tensor(output, true, "output_before_get"); if (output != NULL) { ggml_backend_tensor_get_and_sync(backend, gf->nodes[gf->n_nodes - 1], output->data, 0, ggml_nbytes(output)); } - if (output != NULL) - print_ggml_tensor(output, true, "output_get"); + // if (output != NULL) + // print_ggml_tensor(output, true, "output_get"); } void free_compute_buffer() { diff --git a/pmid.hpp b/pmid.hpp index 6d4bed85..85c2da91 100644 --- a/pmid.hpp +++ b/pmid.hpp @@ -204,15 +204,15 @@ 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"); 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"); 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"); + // 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"); + // 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){ @@ -220,24 +220,24 @@ struct FuseModule{ }else if(right){ stacked_id_embeds = ggml_concat(ctx, stacked_id_embeds, right); } - print_ggml_tensor(stacked_id_embeds, true, "stacked_id_embeds_after_concat"); + // print_ggml_tensor(stacked_id_embeds, true, "stacked_id_embeds_after_concat"); 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_2"); - print_ggml_tensor(class_tokens_mask, true, "class_tokens_mask"); - print_ggml_tensor(prompt_embeds, true, "prompt_embeds"); + // print_ggml_tensor(stacked_id_embeds, true, "stacked_id_embeds_after_permute_2"); + // print_ggml_tensor(class_tokens_mask, true, "class_tokens_mask"); + // print_ggml_tensor(prompt_embeds, true, "prompt_embeds"); // assert class_tokens_mask.sum() == stacked_id_embeds.shape[0], f"{class_tokens_mask.sum()} != {stacked_id_embeds.shape[0]}" // prompt_embeds.masked_scatter_(class_tokens_mask[:, None], stacked_id_embeds.to(prompt_embeds.dtype)) // struct ggml_tensor *prompt_embeds_perm = ggml_cont(ctx, ggml_permute(ctx, prompt_embeds, 1, 0, 2, 3)); struct ggml_tensor *prompt_embeds_perm = ggml_cont(ctx, ggml_transpose(ctx, prompt_embeds)); - print_ggml_tensor(prompt_embeds_perm, true, "prompt_embeds_perm"); + // print_ggml_tensor(prompt_embeds_perm, true, "prompt_embeds_perm"); class_tokens_mask = ggml_repeat(ctx, class_tokens_mask, prompt_embeds_perm); - print_ggml_tensor(class_tokens_mask, true, "class_tokens_mask_repeat"); + // print_ggml_tensor(class_tokens_mask, true, "class_tokens_mask_repeat"); class_tokens_mask = ggml_cont(ctx, ggml_transpose(ctx, class_tokens_mask)); - print_ggml_tensor(class_tokens_mask, true, "class_tokens_mask_transpose"); + // print_ggml_tensor(class_tokens_mask, true, "class_tokens_mask_transpose"); prompt_embeds = ggml_mul(ctx, prompt_embeds, class_tokens_mask); - print_ggml_tensor(prompt_embeds, true, "prompt_embeds_after_mul"); + // print_ggml_tensor(prompt_embeds, true, "prompt_embeds_after_mul"); struct ggml_tensor * updated_prompt_embeds = ggml_add(ctx, prompt_embeds, stacked_id_embeds); - print_ggml_tensor(updated_prompt_embeds, true, "updated_prompt_embeds"); + // print_ggml_tensor(updated_prompt_embeds, true, "updated_prompt_embeds"); // updated_prompt_embeds = prompt_embeds.view(batch_size, seq_length, -1) return updated_prompt_embeds; } @@ -257,18 +257,18 @@ struct PhotoMakerIDEncoder : public GGMLModule { fuse_module(2048) { vision_model = CLIPVisionModel(); // fuse_module = FuseModule(2048); + // wtype = GGML_TYPE_F32; } // void init_params(ggml_context* ctx, ggml_backend_t backend, ggml_type wtype, ggml_allocr* alloc) { - void init_params() { - + void init_params(ggml_type wtype) { + // LOG_INFO(" PMID wtype: %s", ggml_type_name(wtype)); ggml_allocr* alloc = ggml_allocr_new_from_buffer(params_buffer); - 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, 1024, 1280); // 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); } @@ -279,17 +279,15 @@ struct PhotoMakerIDEncoder : public GGMLModule { tensors[prefix + "visual_projection_2.weight"] = visual_projection_2; } - size_t calculate_mem_size() { + size_t calculate_mem_size() { + + wtype = GGML_TYPE_F32; + // LOG_INFO(" PMID wtype: %s", ggml_type_name(wtype)); + size_t mem_size = vision_model.calculate_mem_size(wtype); - LOG_INFO("pmid vision memory buffer size = %.2fMB ", - vision_model.calculate_mem_size(wtype) / 1024.0 / 1024.0); mem_size += fuse_module.calculate_mem_size(wtype); - LOG_INFO("pmid fuse memory buffer size = %.2fMB ", - fuse_module.calculate_mem_size(wtype) / 1024.0 / 1024.0); mem_size += ggml_row_size(wtype, 1280*1024); - LOG_INFO("pmid porject buffer size = %.2fMB ", - ggml_row_size(wtype, 1280*1024) / 1024.0 / 1024.0); return mem_size; } @@ -325,15 +323,15 @@ struct PhotoMakerIDEncoder : public GGMLModule { ggml_set_name(right, "right_input"); - print_ggml_tensor(prompt_embeds, true, "prompt_embeds"); - print_ggml_tensor(class_embedding_temp, true, "class_embedding_temp"); + // print_ggml_tensor(prompt_embeds, true, "prompt_embeds"); + // print_ggml_tensor(class_embedding_temp, true, "class_embedding_temp"); struct ggml_tensor *shared_id_embeds = vision_model.forward(ctx, id_pixel_values, cls, class_embedding_temp, positions ); // [batch_size, seq_length, hidden_size] - print_ggml_tensor(shared_id_embeds, true, "shared_id_embeds"); + // print_ggml_tensor(shared_id_embeds, true, "shared_id_embeds"); // if(class_tokens_mask->backend == GGML_BACKEND_GPU){ // int *ctm = (int *)malloc(class_tokens_mask->ne[0]); @@ -350,20 +348,20 @@ 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"); - print_ggml_tensor(id_embeds_2, true, "id_embeds_2"); + // print_ggml_tensor(id_embeds, true, "id_embeds"); + // print_ggml_tensor(id_embeds_2, true, "id_embeds_2"); id_embeds = ggml_cont(ctx, ggml_permute(ctx, id_embeds, 2, 0, 1, 3)); id_embeds_2 = ggml_cont(ctx, ggml_permute(ctx, id_embeds_2, 2, 0, 1, 3)); - print_ggml_tensor(id_embeds, true, "id_embeds_after_perm"); - print_ggml_tensor(id_embeds_2, true, "id_embeds_2_after_perm"); + // print_ggml_tensor(id_embeds, true, "id_embeds_after_perm"); + // print_ggml_tensor(id_embeds_2, true, "id_embeds_2_after_perm"); id_embeds = ggml_concat(ctx, id_embeds, id_embeds_2); // [batch_size, seq_length, 1, 2048] check whether concat at dim 2 is right - print_ggml_tensor(id_embeds, true, "id_embeds_after_cat"); + // print_ggml_tensor(id_embeds, true, "id_embeds_after_cat"); id_embeds = ggml_cont(ctx, ggml_permute(ctx, id_embeds, 1, 2, 0, 3)); - print_ggml_tensor(id_embeds, true, "id_embeds_after_cat+perm"); + // print_ggml_tensor(id_embeds, true, "id_embeds_after_cat+perm"); struct ggml_tensor * updated_prompt_embeds = fuse_module.forward(ctx, @@ -371,7 +369,7 @@ struct PhotoMakerIDEncoder : public GGMLModule { class_tokens_mask, class_tokens_mask_pos, left, right); - print_ggml_tensor(updated_prompt_embeds, true, "updated_prompt_embeds_returned"); + // print_ggml_tensor(updated_prompt_embeds, true, "updated_prompt_embeds_returned"); return updated_prompt_embeds; @@ -508,7 +506,7 @@ struct PhotoMakerIDEncoder : public GGMLModule { positions, left, right ); - print_ggml_tensor(updated_prompt_embeds, true, "updated_prompt_embeds_returned_forward"); + // print_ggml_tensor(updated_prompt_embeds, true, "updated_prompt_embeds_returned_forward"); ggml_build_forward_expand(gf, updated_prompt_embeds); // ggml_graph_dump_dot(gf, NULL, "id_encoder.dot"); ggml_free(ctx0); diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index bbce9318..23a0823e 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -250,10 +250,10 @@ public: first_stage_model.map_by_name(tensors, "first_stage_model."); if(stacked_id){ - printf("preparing memory for pmid \n"); - pmid_model.init_params(); + // 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."); - printf("done preparing memory for pmid \n"); } } @@ -506,10 +506,10 @@ public: bool force_zero_embeddings = false) { cond_stage_model.set_clip_skip(clip_skip); auto image_tokens = cond_stage_model.convert_token_to_id(trigger_word); - printf(" length of image tokens: %lu \n", image_tokens.size()); - if(image_tokens.size() == 1){ - printf(" image token id is: %d \n", image_tokens[0]); - } + // printf(" length of image tokens: %lu \n", image_tokens.size()); + // if(image_tokens.size() == 1){ + // printf(" image token id is: %d \n", image_tokens[0]); + // } GGML_ASSERT(image_tokens.size() == 1); // auto tokens_and_weights = cond_stage_model.tokenize(text, true); auto tokens_and_weights = cond_stage_model.tokenize_with_trigger_token(text, @@ -1475,25 +1475,13 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx, float std[] = {0.26862954, 0.26130258, 0.27577711}; for(int i = 0; i < num_input_images; i++) { sd_image_t* init_image = input_id_images[i]; - sd_mul_images_to_tensor(init_image->data, init_img, i, mean, std); - } - int64_t *ne = init_img->ne; - fprintf(stderr, "%s: input id image tensor ne [%ld, %ld, %ld, %ld] \n", - __func__, ne[0], ne[1], ne[2], ne[3]); + sd_mul_images_to_tensor(init_image->data, init_img, i, mean, std); + } auto cond_tup = sd_ctx->sd->get_learned_condition_with_trigger(work_ctx, prompt, clip_skip, width, height, num_input_images ); prompts_embeds = std::get<0>(cond_tup); pooled_prompts_embeds = std::get<1>(cond_tup); // [adm_in_channels, ] class_tokens_mask = std::get<2>(cond_tup); // - ne = prompts_embeds->ne; - fprintf(stderr, "%s: text embedding tensor ne [%ld, %ld, %ld, %ld] \n", - __func__, ne[0], ne[1], ne[2], ne[3]); - ne = pooled_prompts_embeds->ne; - fprintf(stderr, "%s: SDXL pooled text embedding tensor ne [%ld, %ld, %ld, %ld] \n", - __func__, ne[0], ne[1], ne[2], ne[3]); - // ne = class_tokens_mask->ne; - // fprintf(stderr, "%s: class token mask tensor ne [%ld, %ld, %ld, %ld] \n", - // __func__, ne[0], ne[1], ne[2], ne[3]); prompts_embeds = sd_ctx->sd->id_encoder(work_ctx, init_img, prompts_embeds, class_tokens_mask); @@ -1512,13 +1500,7 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx, auto cond_pair = sd_ctx->sd->get_learned_condition(work_ctx, prompt, clip_skip, width, height); ggml_tensor* c = cond_pair.first; ggml_tensor* c_vector = cond_pair.second; // [adm_in_channels, ] - - int64_t *ne = c->ne; - fprintf(stderr, "%s: text embedding tensor ne [%ld, %ld, %ld, %ld] \n", - __func__, ne[0], ne[1], ne[2], ne[3]); - ne = c_vector->ne; - fprintf(stderr, "%s: SDXL pooled text embedding tensor ne [%ld, %ld, %ld, %ld] \n", - __func__, ne[0], ne[1], ne[2], ne[3]); + struct ggml_tensor* uc = NULL; struct ggml_tensor* uc_vector = NULL; if (cfg_scale != 1.0) {