From cbfa702fada56b0a3cf6cf6b8423dff682f1cd28 Mon Sep 17 00:00:00 2001 From: bssrdf Date: Sun, 4 Feb 2024 16:03:00 -0500 Subject: [PATCH] added PhotoMakerIDEncoder model in SD --- clip.hpp | 162 +++++++++++++++++++++++++++++++++++++++++-- pmid.hpp | 117 ++++++++++++++++++++----------- stable-diffusion.cpp | 4 ++ 3 files changed, 240 insertions(+), 43 deletions(-) diff --git a/clip.hpp b/clip.hpp index 3c475847..95c12cda 100644 --- a/clip.hpp +++ b/clip.hpp @@ -885,14 +885,168 @@ struct CLIPVisionModel { return h; } - struct ggml_tensor* forward(struct ggml_context* ctx0, struct ggml_tensor* input_ids, struct ggml_tensor* tkn_embeddings, uint32_t max_token_idx = 0, bool return_pooled = false) { + 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]); // token_embedding + position_embedding - struct ggml_tensor* x = NULL; +#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); + inp = ggml_cont(ctx0, ggml_permute(ctx0, inp, 1, 0, 2, 3)); + + // concat class_embeddings and patch_embeddings + struct ggml_tensor * embeddings = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, hidden_size, num_positions, batch_size); + + 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->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]); + + struct ggml_tensor * positions = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, num_positions); + for (int i = 0; i < num_positions; i++) { + ggml_set_i32_1d(positions, i, i); + } + + embeddings = + ggml_add(ctx0, embeddings, ggml_repeat(ctx0, ggml_get_rows(ctx0, model.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)); + } + + // 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]; + + ggml_set_scratch(ctx0, {0, scr0_size, scr0}); + + // layernorm1 + { + 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)); + } + + // 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)); + + 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); + Q = ggml_cont(ctx0, ggml_permute(ctx0, Q, 0, 2, 1, 3)); + 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)); + + 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)); + K = ggml_reshape_3d(ctx0, K, d_head, num_positions, n_head * batch_size); + + struct ggml_tensor * V = + ggml_add(ctx0, ggml_repeat(ctx0, model.layers[il].v_b, cur), ggml_mul_mat(ctx0, model.layers[il].v_w, cur)); + + V = ggml_reshape_4d(ctx0, V, d_head, n_head, num_positions, batch_size); + V = ggml_cont(ctx0, ggml_permute(ctx0, V, 1, 2, 0, 3)); + V = ggml_reshape_3d(ctx0, V, num_positions, d_head, n_head * batch_size); + + struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q); + KQ = ggml_soft_max_inplace(ctx0, KQ); + struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ); + KQV = ggml_reshape_4d(ctx0, KQV, d_head, num_positions, n_head, batch_size); + KQV = ggml_cont(ctx0, ggml_permute(ctx0, KQV, 0, 2, 1, 3)); + + cur = ggml_cpy(ctx0, KQV, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, hidden_size, num_positions, batch_size)); + } + + // attention output + cur = ggml_add(ctx0, ggml_repeat(ctx0, model.layers[il].o_b, cur), ggml_mul_mat(ctx0, model.layers[il].o_w, cur)); + + // re-add the layer input, e.g., residual + cur = ggml_add(ctx0, cur, embeddings); + + embeddings = cur; // embeddings = residual, cur = hidden_states + + // layernorm2 + { + cur = ggml_norm(ctx0, cur, eps); + + cur = ggml_add(ctx0, ggml_mul(ctx0, ggml_repeat(ctx0, model.layers[il].ln_2_w, cur), cur), + ggml_repeat(ctx0, model.layers[il].ln_2_b, cur)); + } + + cur = ggml_mul_mat(ctx0, model.layers[il].ff_i_w, cur); + cur = ggml_add(ctx0, ggml_repeat(ctx0, model.layers[il].ff_i_b, cur), cur); + + // if (ctx->use_gelu) { + // cur = ggml_gelu_inplace(ctx0, cur); + // } else { + cur = ggml_gelu_quick_inplace(ctx0, cur); + // } + + cur = ggml_mul_mat(ctx0, model.layers[il].ff_o_w, cur); + cur = ggml_add(ctx0, ggml_repeat(ctx0, model.layers[il].ff_o_b, cur), cur); + + // residual 2 + cur = ggml_add(ctx0, embeddings, cur); + + embeddings = cur; + } + + // get the output of cls token, e.g., 0th index + struct ggml_tensor * cls = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, batch_size); + for (int b = 0; b < batch_size; b++) { + ggml_set_i32_1d(cls, b, b * num_positions); + } + embeddings = ggml_get_rows(ctx0, ggml_reshape_2d(ctx0, embeddings, hidden_size, num_positions * batch_size), cls); + + // post-layernorm + { + embeddings = ggml_norm(ctx0, embeddings, eps); + + embeddings = ggml_add(ctx0, ggml_mul(ctx0, ggml_repeat(ctx0, model.post_ln_w, embeddings), embeddings), + ggml_repeat(ctx0, model.post_ln_b, embeddings)); + } + + ggml_set_scratch(ctx0, {0, 0, nullptr}); + + // final visual projection + embeddings = ggml_mul_mat(ctx0, model.projection, embeddings); + + // normalize output embeddings + struct ggml_tensor * output = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, projection_dim, batch_size); + + for (int b = 0; b < batch_size; b++) { + struct ggml_tensor * embedding = ggml_get_rows(ctx0, embeddings, ggml_new_i32(ctx0, b)); + if (normalize) { + ggml_tensor * length = ggml_sqrt(ctx0, ggml_sum(ctx0, ggml_sqr(ctx0, embedding))); + embedding = ggml_scale_inplace(ctx0, embedding, ggml_div(ctx0, ggml_new_f32(ctx0, 1.0f), length)); + } + output = ggml_acc(ctx0, output, embedding, output->nb[1], output->nb[2], output->nb[3], b * ggml_nbytes(embedding)); + } + ggml_set_name(output, "check"); + + ggml_set_name(cur, "last_hidden"); + // we return last hidden state instead of image embedding here + return cur; // [batch_size, seq_length, hidden_size] +#endif + return NULL; + - return x; // [N, n_token, hidden_size] } void init_params(ggml_context* ctx, ggml_backend_t backend, ggml_type wtype, ggml_allocr* alloc) { @@ -911,7 +1065,7 @@ struct CLIPVisionModel { post_ln_w = ggml_new_tensor_1d(ctx, wtype, hidden_size); post_ln_b = ggml_new_tensor_1d(ctx, wtype, hidden_size); - visual_projection = ggml_new_tensor_2d(ctx, wtype, projection_dim, hidden_size); + visual_projection = ggml_new_tensor_2d(ctx, wtype, hidden_size, projection_dim); // Check!!! // alloc all tensors linked to this context diff --git a/pmid.hpp b/pmid.hpp index ac32d90c..4d75a67b 100644 --- a/pmid.hpp +++ b/pmid.hpp @@ -35,27 +35,30 @@ struct FuseBlock { size_t calculate_mem_size(ggml_type wtype) { size_t mem_size = 0; - mem_size += 2 * ggml_row_size(GGML_TYPE_F32, channels); // in_layer_0_w/b - mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * channels * 3 * 3); // in_layer_2_w - mem_size += 5 * ggml_row_size(GGML_TYPE_F32, out_channels); // in_layer_2_b/emb_layer_1_b/out_layer_0_w/out_layer_0_b/out_layer_3_b - mem_size += ggml_row_size(wtype, out_channels * emb_channels); // emb_layer_1_w - mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * out_channels * 3 * 3); // out_layer_3_w - - if (out_channels != channels) { - mem_size += ggml_row_size(GGML_TYPE_F16, out_channels * channels * 1 * 1); // skip_w - mem_size += ggml_row_size(GGML_TYPE_F32, out_channels); // skip_b - } + mem_size += 2 * ggml_row_size(wtype, in_dim); + mem_size += ggml_row_size(wtype, in_dim*hidden_dim); + mem_size += 5 * ggml_row_size(wtype, in_dim); + mem_size += ggml_row_size(wtype, hidden_dim*out_dim); + mem_size += ggml_row_size(wtype, hidden_dim); + return mem_size; } - void init_params(struct ggml_context* ctx, ggml_type wtype) { - ln_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_dim); - ln_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_dim); + void init_params(struct ggml_context* ctx, ggml_type wtype, ggml_allocr* alloc) { + ln_w = ggml_new_tensor_1d(ctx, wtype, in_dim); + ln_b = ggml_new_tensor_1d(ctx, wtype, in_dim); + + fc1_b = ggml_new_tensor_1d(ctx, wtype, hidden_dim); + fc1_w = ggml_new_tensor_2d(ctx, wtype, in_dim, hidden_dim); + fc2_b = ggml_new_tensor_1d(ctx, wtype, out_dim); + fc2_w = ggml_new_tensor_2d(ctx, wtype, hidden_dim, out_dim); + // 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) { + ggml_allocr_alloc(alloc, t); + } + } - fc1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_dim); - fc1_w = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, in_dim, hidden_dim); - fc2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_dim); - fc2_w = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_dim, out_dim); } void map_by_name(std::map& tensors, const std::string prefix) { @@ -93,17 +96,27 @@ struct FuseModule{ struct ggml_tensor* ln_b; // [in_dim, ] - FuseModule(int imb_d) - : embed_dim(imb_d){ + FuseModule(int imb_d): + embed_dim(imb_d), + mlp1(imb_d*2, imb_d, imb_d, false), + mlp2(imb_d*2, imb_d, imb_d, true) { + + // mlp1 = FuseBlock(embed_dim*2, embed_dim, embed_dim, false); + // mlp2 = FuseBlock(embed_dim*2, embed_dim, embed_dim, true); + } - void init_params(struct ggml_context* ctx, ggml_type wtype) { - mlp1 = FuseBlock(embed_dim*2, embed_dim, embed_dim, false); - mlp2 = FuseBlock(embed_dim*2, embed_dim, embed_dim, true); - ln_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, embed_dim); - ln_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, embed_dim); + void init_params(struct ggml_context* ctx, ggml_type wtype, ggml_allocr* alloc) { + ln_w = ggml_new_tensor_1d(ctx, wtype, embed_dim); + ln_b = ggml_new_tensor_1d(ctx, wtype, embed_dim); + // 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) { + ggml_allocr_alloc(alloc, t); + } + } } void map_by_name(std::map& tensors, const std::string prefix) { @@ -115,6 +128,13 @@ struct FuseModule{ } + size_t calculate_mem_size(ggml_type wtype) { + size_t mem_size = mlp1.calculate_mem_size(wtype); + mem_size += mlp2.calculate_mem_size(wtype); + mem_size += 2 * ggml_row_size(wtype, embed_dim); + return mem_size; + } + struct ggml_tensor* fuse_fn(struct ggml_context* ctx, struct ggml_tensor* prompt_embeds, struct ggml_tensor* id_embeds) { @@ -154,20 +174,23 @@ struct PhotoMakerIDEncoder : public GGMLModule { struct ggml_tensor* visual_projection_2; PhotoMakerIDEncoder(SDVersion version = VERSION_XL) - : version(version){ - - - - - + : version(version), + fuse_module(2048) { + vision_model = CLIPVisionModel(); + // fuse_module = FuseModule(2048); + } - void init_params(struct ggml_context* ctx, ggml_type wtype) { + // void init_params(ggml_context* ctx, ggml_backend_t backend, ggml_type wtype, ggml_allocr* alloc) { + void init_params() { - vision_model = CLIPVisionModel(); - fuse_module = FuseModule(2048); - visual_projection_2 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1280, 1024); // python [1024, 1280] + 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] + ggml_allocr_alloc(alloc, visual_projection_2); + ggml_allocr_free(alloc); } void map_by_name(std::map& tensors, const std::string prefix) { @@ -176,6 +199,22 @@ struct PhotoMakerIDEncoder : public GGMLModule { tensors[prefix + "visual_projection_2.weight"] = visual_projection_2; } + size_t calculate_mem_size() { + size_t mem_size = vision_model.calculate_mem_size(wtype); + + mem_size += fuse_module.calculate_mem_size(wtype); + + mem_size += ggml_row_size(wtype, 1280*1024); + + return mem_size; + } + + size_t get_num_tensors() { + size_t num_tensors = (3 + 2 + 37); + + return num_tensors; + } + struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* id_pixel_values, @@ -185,18 +224,18 @@ struct PhotoMakerIDEncoder : public GGMLModule { // in_layers - struct ggml_tensor *shared_id_embeds = vision_model.forward(ctx, id_pixel_values); // [1] - struct ggml_tensor *id_embeds = vision_model.visual_project(ctx, shared_id_embeds); - struct ggml_tensor *id_embeds_2 = ggml_mul_mat(ctx, visual_projection_2, shared_id_embeds); + struct ggml_tensor *shared_id_embeds = vision_model.forward(ctx, id_pixel_values); // [batch_size, seq_length, hidden_size] + struct ggml_tensor *id_embeds = vision_model.visual_project(ctx, shared_id_embeds); // [batch_size, seq_length, proj_dim(768)] + struct ggml_tensor *id_embeds_2 = ggml_mul_mat(ctx, visual_projection_2, shared_id_embeds); // [batch_size, seq_length, 1280] // id_embeds = id_embeds.view(b, num_inputs, 1, -1) // id_embeds_2 = id_embeds_2.view(b, num_inputs, 1, -1) // id_embeds = torch.cat((id_embeds, id_embeds_2), dim=-1) - id_embeds = ggml_concat(ctx, id_embeds, id_embeds_2); // check whether concat at dim 2 is right + id_embeds = ggml_concat(ctx, id_embeds, id_embeds_2); // [batch_size, seq_length, 1, 2048] check whether concat at dim 2 is right struct ggml_tensor * updated_prompt_embeds = fuse_module.forward(ctx, prompt_embeds, id_embeds, class_tokens_mask); - return updated_prompt_embeds + return updated_prompt_embeds; } diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index b5e61305..37f28000 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -7,6 +7,7 @@ #include "util.h" #include "clip.hpp" +#include "pmid.hpp" #include "control.hpp" #include "denoiser.hpp" #include "esrgan.hpp" @@ -64,6 +65,7 @@ public: FrozenCLIPEmbedderWithCustomWords cond_stage_model; UNetModel diffusion_model; AutoEncoderKL first_stage_model; + PhotoMakerIDEncoder pmid_model; bool use_tiny_autoencoder = false; bool vae_tiling = false; bool stacked_id = false; @@ -167,6 +169,8 @@ public: } cond_stage_model = FrozenCLIPEmbedderWithCustomWords(version); diffusion_model = UNetModel(version); + pmid_model = PhotoMakerIDEncoder(version); + LOG_INFO("Stable Diffusion %s ", model_version_to_str[version]); if (wtype == GGML_TYPE_COUNT) {