fixed a bug in FuseBlock forward; also remove diag_mask op in for vision transformer; getting better results

This commit is contained in:
bssrdf
2024-02-15 18:57:19 -05:00
parent 26c5cca043
commit 191c5ec05a
4 changed files with 76 additions and 47 deletions

View File

@@ -482,6 +482,7 @@ struct ResidualAttentionBlock {
int32_t d_model;
int32_t hidden_size; // n_head * d_model
int32_t intermediate_size;
bool apply_diag_mask = true;
// attention
struct ggml_tensor* q_w; // [hidden_size, hidden_size]
@@ -596,8 +597,8 @@ struct ResidualAttentionBlock {
v = ggml_reshape_3d(ctx, v, n_token, d_model, n_head * N); // [N * n_head, d_model, n_token]
struct ggml_tensor* kq = ggml_mul_mat(ctx, k, q); // [N * n_head, n_token, n_token]
kq = ggml_diag_mask_inf_inplace(ctx, kq, 0);
if(apply_diag_mask)
kq = ggml_diag_mask_inf_inplace(ctx, kq, 0);
kq = ggml_soft_max_inplace(ctx, kq);
struct ggml_tensor* kqv = ggml_mul_mat(ctx, v, kq); // [N * n_head, n_token, d_model]
@@ -904,6 +905,7 @@ struct CLIPVisionModel {
resblocks[i].n_head = n_head;
resblocks[i].hidden_size = hidden_size;
resblocks[i].intermediate_size = intermediate_size;
resblocks[i].apply_diag_mask = false;
}
}
@@ -961,7 +963,7 @@ struct CLIPVisionModel {
const int d_head = hidden_size / n_head;
// ggml_set_name(x, "id_pixel");
print_ggml_tensor(x, true, "id_pixel");
// print_ggml_tensor(x, true, "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]);
@@ -971,15 +973,15 @@ struct CLIPVisionModel {
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");
// print_ggml_tensor(inp, true, "inp_cont");
ggml_set_name(inp, "inp_cont");
// ggml_set_name(class_embedding, "class_embedding");
@@ -997,7 +999,7 @@ struct CLIPVisionModel {
ggml_set_name(class_embedding_rep, "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");
@@ -1017,7 +1019,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_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);
@@ -1028,12 +1030,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_encoder");
// print_ggml_tensor(embeddings, true, "embeddings_after_encoder");
// 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_after_cls_token");
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);
@@ -1055,13 +1057,13 @@ struct CLIPVisionModel {
}
class_embedding = ggml_new_tensor_1d(ctx, wtype, hidden_size);
ggml_set_name(class_embedding, "class_embedding");
ggml_set_name(class_embedding, "vision.class_embedding");
patch_embeddings = ggml_new_tensor_4d(ctx, wtype, patch_size, patch_size, 3, hidden_size);
ggml_set_name(patch_embeddings, "patch_embeddings");
ggml_set_name(patch_embeddings, "vision.patch_embeddings");
position_embeddings = ggml_new_tensor_2d(ctx, wtype, hidden_size, 257);
ggml_set_name(position_embeddings, "position_embeddings");
ggml_set_name(position_embeddings, "vision.position_embeddings");
pre_ln_w = ggml_new_tensor_1d(ctx, wtype, hidden_size);

View File

@@ -679,6 +679,18 @@ struct GGMLModule {
ggml_allocr_alloc_graph(compute_allocr, gf);
// struct ggml_init_params params = {};
// params.mem_size += 1024*1024*1024; // 100M
// params.mem_size += 2 * ggml_tensor_overhead();
// params.mem_buffer = NULL;
// params.no_alloc = false;
// struct ggml_context* obs_ctx = ggml_init(params);
// if (!obs_ctx) {
// LOG_ERROR("ggml_init() failed");
// return;
// }
if (ggml_backend_is_cpu(backend)) {
ggml_backend_cpu_set_n_threads(backend, n_threads);
}
@@ -693,26 +705,40 @@ struct GGMLModule {
struct ggml_tensor * imb = NULL;
for (int i = 0; i < gf->n_leafs; i++) {
struct ggml_tensor * t1 = gf->leafs[i];
if(strcmp(ggml_get_name(t1), "id_pixel_values_input") == 0) {
if(strcmp(ggml_get_name(t1), "vision.patch_embeddings") == 0) {
imb = t1;
int64_t stride = imb->ne[0];
int64_t ne3 = imb->ne[3];
float* out_data = new float[ggml_nelements(imb)];
ggml_backend_tensor_get(imb, out_data, 0, ggml_nbytes(imb));
for(int l = 0; l < ne3; ++l){
printf("[");
for(int k = 0; k < 3; ++k){
float mi = 100.f, mx= -100.f;
for(int i = 0; i < stride; i++){
printf("[");
for(int i = 0; i < stride; i++){
for(int j = 0; j < stride; j++){
float val = out_data[k*stride*stride+i*stride+j];
if(mi > val) mi = val;
if(mx < val) mx = val;
float val = out_data[l*3*stride*stride+ k*stride*stride+i*stride+j];
printf("%f, ", val);
}
printf("]\n");
}
}
printf("B. channel, min, max: %d, %f %f \n", k, mi, mx);
// printf("B. channel, min, max: %d, %f %f \n", k, mi, mx);
}
printf("]\n");
}
// for(int k = 0; k < 3; ++k){
// float mi = 100.f, mx= -100.f;
// for(int i = 0; i < stride; i++){
// printf("[");
// for(int j = 0; j < stride; j++){
// float val = out_data[k*stride*stride+i*stride+j];
// if(mi > val) mi = val;
// if(mx < val) mx = val;
// printf("%f, ", val);
// }
// printf("]\n");
// }
// printf("B. channel, min, max: %d, %f %f \n", k, mi, mx);
// }
// printf("[");
// for(int i = 0; i < stride; i++){
@@ -722,7 +748,7 @@ struct GGMLModule {
delete out_data;
}
}
// #if 0
for (int i = 0; i < gf->n_nodes; i++) {
struct ggml_tensor * t1 = gf->nodes[i];
// if(strcmp(ggml_get_name(t1), "embeddings_after_add") == 0) {

View File

@@ -79,13 +79,14 @@ struct FuseBlock {
// x: [N, channels, h, w]
// in_layers
auto h = ggml_nn_layer_norm(ctx, x, ln_w, ln_b);
h = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, h), fc1_b);
h = ggml_gelu_inplace(ctx, h);
h = ggml_add(ctx, ggml_mul_mat(ctx, fc2_w, h), fc2_b);
struct ggml_tensor* r = x;
x = ggml_nn_layer_norm(ctx, x, ln_w, ln_b);
x = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, x), fc1_b);
x = ggml_gelu_inplace(ctx, x);
x = ggml_add(ctx, ggml_mul_mat(ctx, fc2_w, x), fc2_b);
if(use_residue)
x = ggml_add(ctx, x, h);
return h;
x = ggml_add(ctx, x, r);
return x;
}
};
@@ -204,16 +205,16 @@ 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(valid_id_embeds, true, "valid_id_embeds");
// print_ggml_tensor(image_token_embeds, true, "image_token_embeds");
// print_ggml_tensor(valid_id_embeds, true, "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");
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){
@@ -221,19 +222,19 @@ 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(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, true, "prompt_embeds_perm");
// print_ggml_tensor(prompt_embeds, true, "prompt_embeds_perm");
// class_tokens_mask = ggml_repeat(ctx, class_tokens_mask, prompt_embeds_perm);
class_tokens_mask = ggml_cont(ctx, ggml_transpose(ctx, class_tokens_mask));
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");
prompt_embeds = ggml_mul(ctx, prompt_embeds, class_tokens_mask);
@@ -333,7 +334,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");
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]);
@@ -359,11 +360,11 @@ struct PhotoMakerIDEncoder : public GGMLModule {
// 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,

View File

@@ -837,8 +837,8 @@ public:
diffusion_model.compute(out_cond, n_threads, noised_input, NULL, c, control_net.controls, control_strength, t_emb, c_vector);
}
else{
// diffusion_model.compute(out_cond, n_threads, noised_input, NULL, c_id, control_net.controls, control_strength, t_emb, c_vec_id);
diffusion_model.compute(out_cond, n_threads, noised_input, NULL, c_id, control_net.controls, control_strength, t_emb, c_vector);
diffusion_model.compute(out_cond, n_threads, noised_input, NULL, c_id, control_net.controls, control_strength, t_emb, c_vec_id);
// diffusion_model.compute(out_cond, n_threads, noised_input, NULL, c_id, control_net.controls, control_strength, t_emb, c_vector);
}
}else{
diffusion_model.compute(out_cond, n_threads, noised_input, NULL, c, control_net.controls, control_strength, t_emb, c_vector);
@@ -1630,7 +1630,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 = 30.f;
float style_strength_ratio = 20.f;
start_merge_step = int(style_strength_ratio / 100.f * sample_steps);
if(start_merge_step > 30)
start_merge_step = 30;