upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-07-29 06:10:52 -05:00
Code Issues Packages Projects Releases Wiki Activity

added PhotoMakerIDEncoder model in SD

Browse Source
This commit is contained in:
bssrdf
2024-02-04 16:03:00 -05:00
parent 5dbbea0ca4
commit cbfa702fad
3 changed files with 240 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

162
clip.hpp
View File

@@ -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

117
pmid.hpp
View File

@@ -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<std::string, struct ggml_tensor*>& 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<std::string, struct ggml_tensor*>& 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<std::string, struct ggml_tensor*>& 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;
}

4
stable-diffusion.cpp
View File

@@ -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) {