From 7f5f580532bb29ab8d7c2c409123cda2f4e40272 Mon Sep 17 00:00:00 2001 From: bssrdf Date: Wed, 7 Feb 2024 17:25:06 -0500 Subject: [PATCH] making more progress; finishing vision model forward --- clip.hpp | 50 ++++++++++++++++++------------ pmid.hpp | 67 ++++++++++++++++++++++++++++++++++++++++ stable-diffusion.cpp | 73 +++++++++++++++++++++++++++++++++++++++++--- 3 files changed, 166 insertions(+), 24 deletions(-) diff --git a/clip.hpp b/clip.hpp index 672fb20f..13398e26 100644 --- a/clip.hpp +++ b/clip.hpp @@ -146,7 +146,7 @@ public: size_t space_pos = merge.find(' '); merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1)); // LOG_DEBUG("%s", utf32_to_utf8(merge.substr(space_pos + 1)).c_str()); - // printf("%s :: %s, %s \n", utf32_to_utf8(merge).c_str(), utf32_to_utf8(merge.substr(0, space_pos)).c_str(), + // printf("%s :: %s | %s \n", utf32_to_utf8(merge).c_str(), utf32_to_utf8(merge.substr(0, space_pos)).c_str(), // utf32_to_utf8(merge.substr(space_pos + 1)).c_str()); } std::vector vocab; @@ -297,21 +297,22 @@ public: } std::string decode(const std::vector &tokens){ - std::u32string text = utf8_to_utf32(""); + std::string text = ""; for(int t : tokens){ if(t == 49406 || t == 49407) continue; std::u32string ts = decoder[t]; - // printf("%d, %s\n", t, ts.c_str()); - text += ts; + // printf("%d, %s \n", t, utf32_to_utf8(ts).c_str()); + text += " "+utf32_to_utf8(ts).replace(ts.length()-4, ts.length()-1, "") ; } std::vector bytes; - for (auto c : text){ - bytes.push_back(byte_decoder[c]); - } + // for (auto c : text){ + // bytes.push_back(byte_decoder[c]); + // } - std::string s((char *)bytes.data()); - return s; + // std::string s((char *)bytes.data()); + // std::string s = ""; + return trim(text); } @@ -867,6 +868,7 @@ struct CLIPVisionModel { int32_t num_hidden_layers = 24; // 24 for OPEN_CLIP_VIT_H_14 int32_t patch_size = 14; int32_t projection_dim = 768; // only for OPEN_CLIP_VIT_BIGG_14 + float eps = 1.e-8; // embeddings struct ggml_tensor * class_embedding; // [hidden_size,] @@ -943,9 +945,17 @@ struct CLIPVisionModel { struct ggml_tensor* forward(struct ggml_context* ctx0, struct ggml_tensor *x) { // input_ids: [N, n_token] // GGML_ASSERT(input_ids->ne[0] <= position_ids->ne[0]); + const int image_size = x->ne[0]; + int batch_size = x->ne[3]; + const int num_patches = ((image_size / patch_size) * (image_size / patch_size)); + const int num_positions = num_patches + 1; + int32_t n_layer = num_hidden_layers; +#if 0 + static size_t scr0_size = get_scr_buf_req_by_size((struct clip_ctx *)ctx); + static void * scr0 = malloc(scr0_size); // token_embedding + position_embedding -#if 0 + struct ggml_tensor * inp = ggml_conv_2d(ctx0, patch_embeddings, x, patch_size, patch_size, 0, 0, 1, 1); inp = ggml_reshape_3d(ctx0, inp, num_patches, hidden_size, batch_size); @@ -957,10 +967,10 @@ struct CLIPVisionModel { ggml_set_zero(embeddings); struct ggml_tensor * temp = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, hidden_size, 1, batch_size); - embeddings = ggml_acc(ctx0, embeddings, ggml_repeat(ctx0, model.class_embedding, temp), embeddings->nb[1], + embeddings = ggml_acc(ctx0, embeddings, ggml_repeat(ctx0, class_embedding, temp), embeddings->nb[1], embeddings->nb[2], embeddings->nb[3], 0); embeddings = - ggml_acc(ctx0, embeddings, inp, embeddings->nb[1], embeddings->nb[2], embeddings->nb[3], model.class_embedding->nb[1]); + ggml_acc(ctx0, embeddings, inp, embeddings->nb[1], embeddings->nb[2], embeddings->nb[3], class_embedding->nb[1]); struct ggml_tensor * positions = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, num_positions); for (int i = 0; i < num_positions; i++) { @@ -968,21 +978,21 @@ struct CLIPVisionModel { } embeddings = - ggml_add(ctx0, embeddings, ggml_repeat(ctx0, ggml_get_rows(ctx0, model.position_embeddings, positions), embeddings)); + ggml_add(ctx0, embeddings, ggml_repeat(ctx0, ggml_get_rows(ctx0, position_embeddings, positions), embeddings)); // pre-layernorm { embeddings = ggml_norm(ctx0, embeddings, eps); - embeddings = ggml_add(ctx0, ggml_mul(ctx0, ggml_repeat(ctx0, model.pre_ln_w, embeddings), embeddings), - ggml_repeat(ctx0, model.pre_ln_b, embeddings)); + embeddings = ggml_add(ctx0, ggml_mul(ctx0, ggml_repeat(ctx0, pre_ln_w, embeddings), embeddings), + ggml_repeat(ctx0, pre_ln_b, embeddings)); } // loop over layers for (int il = 0; il < n_layer; il++) { struct ggml_tensor * cur = embeddings; // embeddings = residual, cur = hidden_states - const size_t nb_q_w = model.layers[il].q_w->nb[0]; + const size_t nb_q_w = resblocks[il].q_w->nb[0]; ggml_set_scratch(ctx0, {0, scr0_size, scr0}); @@ -990,15 +1000,15 @@ struct CLIPVisionModel { { cur = ggml_norm(ctx0, cur, eps); - cur = ggml_add(ctx0, ggml_mul(ctx0, ggml_repeat(ctx0, model.layers[il].ln_1_w, cur), cur), - ggml_repeat(ctx0, model.layers[il].ln_1_b, cur)); + cur = ggml_add(ctx0, ggml_mul(ctx0, ggml_repeat(ctx0, resblocks[il].ln_1_w, cur), cur), + ggml_repeat(ctx0, resblocks[il].ln_1_b, cur)); } // self-attention { struct ggml_tensor * Q = - ggml_add(ctx0, ggml_repeat(ctx0, model.layers[il].q_b, cur), ggml_mul_mat(ctx0, model.layers[il].q_w, cur)); + ggml_add(ctx0, ggml_repeat(ctx0, resblocks[il].q_b, cur), ggml_mul_mat(ctx0, resblocks[il].q_w, cur)); Q = ggml_scale_inplace(ctx0, Q, ggml_new_f32(ctx0, 1.0f / sqrt((float)d_head))); Q = ggml_reshape_4d(ctx0, Q, d_head, n_head, num_positions, batch_size); @@ -1006,7 +1016,7 @@ struct CLIPVisionModel { Q = ggml_reshape_3d(ctx0, Q, d_head, num_positions, n_head * batch_size); struct ggml_tensor * K = - ggml_add(ctx0, ggml_repeat(ctx0, model.layers[il].k_b, cur), ggml_mul_mat(ctx0, model.layers[il].k_w, cur)); + ggml_add(ctx0, ggml_repeat(ctx0, resblocks[il].k_b, cur), ggml_mul_mat(ctx0, resblocks[il].k_w, cur)); K = ggml_reshape_4d(ctx0, K, d_head, n_head, num_positions, batch_size); K = ggml_cont(ctx0, ggml_permute(ctx0, K, 0, 2, 1, 3)); diff --git a/pmid.hpp b/pmid.hpp index 94f3e8bf..967a9d66 100644 --- a/pmid.hpp +++ b/pmid.hpp @@ -254,6 +254,73 @@ struct PhotoMakerIDEncoder : public GGMLModule { } + struct ggml_cgraph* build_graph(struct ggml_allocr* allocr, + struct ggml_tensor* id_pixel_values, + struct ggml_tensor* prompt_embeds, + struct ggml_tensor* class_tokens_mask + ) { + // 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 + static size_t buf_size = ggml_tensor_overhead() * GGML_DEFAULT_GRAPH_SIZE + ggml_graph_overhead(); + static std::vector buf(buf_size); + + struct ggml_init_params params = { + /*.mem_size =*/buf_size, + /*.mem_buffer =*/buf.data(), + /*.no_alloc =*/true, // the tensors will be allocated later by ggml_allocr_alloc_graph() + }; + + struct ggml_context* ctx0 = ggml_init(params); + + struct ggml_cgraph* gf = ggml_new_graph(ctx0); + + struct ggml_tensor* id_pixel_values_d = ggml_dup_tensor(ctx0, id_pixel_values); + ggml_allocr_alloc(allocr, id_pixel_values_d); + struct ggml_tensor* prompt_embeds_d = ggml_dup_tensor(ctx0, prompt_embeds); + ggml_allocr_alloc(allocr, prompt_embeds_d); + struct ggml_tensor* class_tokens_mask_d = ggml_dup_tensor(ctx0, class_tokens_mask); + ggml_allocr_alloc(allocr, class_tokens_mask_d); + + if (!ggml_allocr_is_measure(allocr)) { + ggml_backend_tensor_set(id_pixel_values_d, id_pixel_values->data, 0, ggml_nbytes(id_pixel_values)); + ggml_backend_tensor_set(prompt_embeds_d, prompt_embeds->data, 0, ggml_nbytes(prompt_embeds)); + ggml_backend_tensor_set(class_tokens_mask_d, class_tokens_mask->data, 0, ggml_nbytes(class_tokens_mask_d)); + } + struct ggml_tensor* updated_prompt_embeds = forward(ctx0, + id_pixel_values_d, + prompt_embeds_d, + class_tokens_mask_d); + + ggml_build_forward_expand(gf, updated_prompt_embeds); + ggml_free(ctx0); + + return gf; + } + + void alloc_compute_buffer(ggml_context* work_ctx, + struct ggml_tensor* id_pixel_values, + struct ggml_tensor* prompt_embeds, + struct ggml_tensor* class_tokens_mask) { + auto get_graph = [&]() -> struct ggml_cgraph* { + + return build_graph(compute_allocr, id_pixel_values, prompt_embeds, class_tokens_mask); + }; + GGMLModule::alloc_compute_buffer(get_graph); + } + + void compute(const int n_threads, + struct ggml_tensor* id_pixel_values, + struct ggml_tensor* prompt_embeds, + struct ggml_tensor* class_tokens_mask, + struct ggml_tensor* updated_prompt_embeds) { + + auto get_graph = [&]() -> struct ggml_cgraph* { + return build_graph(compute_allocr, id_pixel_values, prompt_embeds, class_tokens_mask); + }; + + GGMLModule::compute(get_graph, n_threads, updated_prompt_embeds); + + } + }; diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index 01253a7f..d2eee1b0 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -606,6 +606,25 @@ public: return std::make_tuple(result, vec, class_tokens_mask); } + ggml_tensor* id_encoder(ggml_context* work_ctx, + ggml_tensor* init_img, + ggml_tensor* prompts_embeds, + ggml_tensor* class_tokens_mask){ + + size_t total_hidden_size = cond_stage_model.text_model.hidden_size; + if (version == VERSION_XL) { + total_hidden_size += cond_stage_model.text_model2.hidden_size; + } + ggml_tensor *res = ggml_new_tensor_2d(work_ctx, + GGML_TYPE_F32, + total_hidden_size, + cond_stage_model.text_model.max_position_embeddings); + pmid_model.alloc_compute_buffer(work_ctx, init_img, prompts_embeds, class_tokens_mask); + pmid_model.compute(n_threads, init_img, prompts_embeds, class_tokens_mask, res); + pmid_model.free_compute_buffer(); + return res; + } + std::pair get_learned_condition(ggml_context* work_ctx, const std::string& text, int clip_skip, @@ -1431,13 +1450,19 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx, srand((int)time(NULL)); seed = rand(); } + + std::string prompt_text_only; ggml_tensor* init_img = NULL; + ggml_tensor* prompts_embeds = NULL; + ggml_tensor* pooled_prompts_embeds = NULL; + ggml_tensor* class_tokens_mask = NULL; if(sd_ctx->sd->stacked_id){ int32_t width = input_id_images[0]->width; int32_t height = input_id_images[0]->height; int32_t channels = input_id_images[0]->channel; int32_t num_input_images = input_id_images.size(); init_img = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, channels, num_input_images); + // TODO: move these to somewhere else and be user settable float mean[] = {0.48145466, 0.4578275, 0.40821073}; float std[] = {0.26862954, 0.26130258, 0.27577711}; for(int i = 0; i < num_input_images; i++) { @@ -1447,15 +1472,45 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx, 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]); + 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); + + // if (sd_ctx->sd->free_params_immediately) { + // sd_ctx->sd->pmid_model.free_params_buffer(); + // } + // Encode input prompt without the trigger word for delayed conditioning + prompt_text_only = sd_ctx->sd->remove_trigger_from_prompt(work_ctx, prompt); + // printf("%s || %s \n", prompt.c_str(), prompt_text_only.c_str()); + prompt = prompt_text_only; // } - std::string prompt_text_only = sd_ctx->sd->remove_trigger_from_prompt(work_ctx, prompt); - printf("%s || %s \n", prompt.c_str(), prompt_text_only.c_str()); - t0 = ggml_time_ms(); + + t0 = ggml_time_ms(); 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) { @@ -1496,7 +1551,17 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx, std::vector sigmas = sd_ctx->sd->denoiser->schedule->get_sigmas(sample_steps); - struct ggml_tensor* x_0 = sd_ctx->sd->sample(work_ctx, x_t, NULL, c, c_vector, uc, uc_vector, image_hint, cfg_scale, sample_method, sigmas, control_strength); + struct ggml_tensor* x_0 = sd_ctx->sd->sample(work_ctx, + x_t, + NULL, + c, + c_vector, + uc, uc_vector, + image_hint, + cfg_scale, + sample_method, + sigmas, + control_strength); // struct ggml_tensor* x_0 = load_tensor_from_file(ctx, "samples_ddim.bin"); // print_ggml_tensor(x_0); int64_t sampling_end = ggml_time_ms();