diff --git a/clip.hpp b/clip.hpp index 13398e26..769deb94 100644 --- a/clip.hpp +++ b/clip.hpp @@ -587,7 +587,8 @@ struct ResidualAttentionBlock { struct ggml_tensor* k = ggml_nn_linear(ctx, x, k_w, k_b); k = ggml_reshape_4d(ctx, k, d_model, n_head, n_token, N); // [N, n_token, n_head, d_model] k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [N, n_head, n_token, d_model] - k = ggml_reshape_3d(ctx, k, d_model, n_token, n_head); // [N * n_head, n_token, d_model] + // k = ggml_reshape_3d(ctx, k, d_model, n_token, n_head); // [N * n_head, n_token, d_model] + k = ggml_reshape_3d(ctx, k, d_model, n_token, n_head * N); // [N * n_head, n_token, d_model] struct ggml_tensor* v = ggml_nn_linear(ctx, x, v_w, v_b); v = ggml_reshape_4d(ctx, v, d_model, n_head, n_token, N); // [N, n_token, n_head, d_model] @@ -603,13 +604,14 @@ struct ResidualAttentionBlock { kqv = ggml_reshape_4d(ctx, kqv, d_model, n_token, n_head, N); kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3)); // [N, n_token, n_head, d_model] - x = ggml_reshape_2d(ctx, kqv, d_model * n_head, n_token * N); // // [N * n_token, d_model * n_head] + // x = ggml_reshape_2d(ctx, kqv, d_model * n_head, n_token * N); // // [N * n_token, d_model * n_head] + x = ggml_reshape_3d(ctx, kqv, d_model * n_head, n_token, N); // // [N, n_token, d_model * n_head] } - // attention output + // attention output x = ggml_nn_linear(ctx, x, out_w, out_b); - // residual + // residual x = ggml_add(ctx, x, r); r = x; @@ -868,7 +870,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; + float eps = 1.e-7; // embeddings struct ggml_tensor * class_embedding; // [hidden_size,] @@ -936,13 +938,16 @@ struct CLIPVisionModel { tensors[stripped + "visual_projection.weight"] = visual_projection; } + struct ggml_tensor* visual_project(struct ggml_context* ctx0, struct ggml_tensor* input){ auto h = ggml_mul_mat(ctx0, visual_projection, input); return h; } - struct ggml_tensor* forward(struct ggml_context* ctx0, struct ggml_tensor *x) { + struct ggml_tensor* forward(struct ggml_context* ctx0, struct ggml_tensor *x, + struct ggml_tensor *temp, + struct ggml_tensor *positions) { // input_ids: [N, n_token] // GGML_ASSERT(input_ids->ne[0] <= position_ids->ne[0]); const int image_size = x->ne[0]; @@ -950,167 +955,48 @@ struct CLIPVisionModel { 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); + const int d_head = hidden_size / n_head; - // token_embedding + position_embedding - - 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, 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)); + // inp = ggml_cont(ctx0, ggml_permute(ctx0, inp, 1, 0, 2, 3)); + inp = ggml_cont(ctx0, ggml_permute(ctx0, inp, 2, 0, 3, 1)); // concat class_embeddings and patch_embeddings - struct ggml_tensor * embeddings = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, hidden_size, num_positions, batch_size); + // struct ggml_tensor * embeddings = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, hidden_size, num_positions, batch_size); + + // struct ggml_tensor * temp = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, hidden_size, 1, batch_size); + + ggml_tensor *class_embedding_rep = ggml_repeat(ctx0, class_embedding, temp); - ggml_set_zero(embeddings); - struct ggml_tensor * temp = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, hidden_size, 1, batch_size); + struct ggml_tensor *embeddings = ggml_concat(ctx0, class_embedding_rep, inp); + embeddings = ggml_cont(ctx0, ggml_permute(ctx0, embeddings, 0, 3, 1, 2)); - 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], class_embedding->nb[1]); + // embeddings = ggml_acc(ctx0, embeddings, ggml_repeat(ctx0, class_embedding, temp), embeddings->nb[1], + // embeddings->nb[2], embeddings->nb[3], 0); + + // struct ggml_tensor * positions = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, num_positions); - 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, position_embeddings, positions), embeddings)); - // pre-layernorm - { - embeddings = ggml_norm(ctx0, embeddings, eps); + // pre-layernorm + embeddings = ggml_nn_layer_norm(ctx0, embeddings, pre_ln_w, pre_ln_w); - embeddings = ggml_add(ctx0, ggml_mul(ctx0, ggml_repeat(ctx0, pre_ln_w, embeddings), embeddings), - ggml_repeat(ctx0, pre_ln_b, embeddings)); - } + // transformer + for (int i = 0; i < num_hidden_layers; i++) { + embeddings = resblocks[i].forward(ctx0, embeddings); // [N, n_token, hidden_size] + } - // loop over layers - for (int il = 0; il < n_layer; il++) { - struct ggml_tensor * cur = embeddings; // embeddings = residual, cur = hidden_states + struct ggml_tensor * cur = embeddings; - const size_t nb_q_w = resblocks[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, 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, 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); - 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, 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)); - 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"); + ggml_set_name(cur, "last_hidden_state"); // we return last hidden state instead of image embedding here + // TODO: add other return items return cur; // [batch_size, seq_length, hidden_size] -#endif - return NULL; - } diff --git a/ggml_extend.hpp b/ggml_extend.hpp index 6fbe65f4..3fe5b1cc 100644 --- a/ggml_extend.hpp +++ b/ggml_extend.hpp @@ -77,8 +77,11 @@ __STATIC_INLINE__ ggml_fp16_t ggml_tensor_get_f16(const ggml_tensor* tensor, int return *(ggml_fp16_t*)((char*)(tensor->data) + i * tensor->nb[3] + j * tensor->nb[2] + k * tensor->nb[1] + l * tensor->nb[0]); } -__STATIC_INLINE__ void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only = false) { - printf("shape(%zu, %zu, %zu, %zu)\n", tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); + + + +__STATIC_INLINE__ void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only = false, const char *mark = "") { + printf("%s: shape(%zu, %zu, %zu, %zu)\n", mark, tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); fflush(stdout); if (shape_only) { return; diff --git a/pmid.hpp b/pmid.hpp index 967a9d66..64ac872a 100644 --- a/pmid.hpp +++ b/pmid.hpp @@ -234,19 +234,28 @@ struct PhotoMakerIDEncoder : public GGMLModule { struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* id_pixel_values, struct ggml_tensor* prompt_embeds, - struct ggml_tensor* class_tokens_mask) { + struct ggml_tensor* class_tokens_mask, + struct ggml_tensor* class_embedding_temp, + struct ggml_tensor* positions) { // x: [N, channels, h, w] // in_layers - - struct ggml_tensor *shared_id_embeds = vision_model.forward(ctx, id_pixel_values); // [batch_size, seq_length, hidden_size] + 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, + class_embedding_temp, + positions + ); // [batch_size, seq_length, hidden_size] + print_ggml_tensor(shared_id_embeds, true, "shared_id_embeds"); 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) + print_ggml_tensor(id_embeds, true, "id_embeds"); + print_ggml_tensor(id_embeds_2, true, "id_embeds_2"); 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); @@ -280,15 +289,37 @@ struct PhotoMakerIDEncoder : public GGMLModule { struct ggml_tensor* class_tokens_mask_d = ggml_dup_tensor(ctx0, class_tokens_mask); ggml_allocr_alloc(allocr, class_tokens_mask_d); + const int image_size = id_pixel_values->ne[0]; + int batch_size = id_pixel_values->ne[3]; + const int num_patches = ((image_size / vision_model.patch_size) * (image_size / vision_model.patch_size)); + const int num_positions = num_patches + 1; + + + struct ggml_tensor * class_embedding_temp = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, + vision_model.hidden_size, 1, 1, batch_size); + struct ggml_tensor * positions = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, num_positions); + ggml_allocr_alloc(allocr, class_embedding_temp); + ggml_allocr_alloc(allocr, positions); + + + 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)); + std::vector pos; + for (int i = 0; i < num_positions; i++) { + pos.push_back(i); + } + ggml_backend_tensor_set(positions, pos.data(), 0, ggml_nbytes(positions)); } struct ggml_tensor* updated_prompt_embeds = forward(ctx0, id_pixel_values_d, prompt_embeds_d, - class_tokens_mask_d); + class_tokens_mask_d, + class_embedding_temp, + positions + ); ggml_build_forward_expand(gf, updated_prompt_embeds); ggml_free(ctx0); diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index d2eee1b0..995a468d 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -1487,11 +1487,11 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx, 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); + 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(); - // } + 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());