diff --git a/CMakeLists.txt b/CMakeLists.txt index 8545ef6e..a5807921 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -34,7 +34,7 @@ option(BUILD_SHARED_LIBS "sd: build shared libs" OFF) if(SD_CUBLAS) message("Use CUBLAS as backend stable-diffusion") set(GGML_CUBLAS ON) - add_definitions(-DSD_USE_CUBLAS) + add_definitions(-DSD_USE_CUBLAS -DGGML_CUDA_FORCE_MMQ) endif() if(SD_METAL) diff --git a/clip.hpp b/clip.hpp index 546e944b..3c475847 100644 --- a/clip.hpp +++ b/clip.hpp @@ -803,6 +803,138 @@ struct CLIPTextModel { } }; + +struct CLIPVisionModel { + CLIPVersion version = OPENAI_CLIP_VIT_L_14; + // network hparams + int32_t hidden_size = 1024; + int32_t intermediate_size = 4096; // for OPEN_CLIP_VIT_H_14 + int32_t n_head = 16; // num_attention_heads, 16 for OPEN_CLIP_VIT_H_14 + 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 + + // embeddings + struct ggml_tensor * class_embedding; // [hidden_size,] + struct ggml_tensor * patch_embeddings; // [patch_size,patch_size,3,hidden_size] + struct ggml_tensor * position_embeddings; // [hidden_size, 257] + + // transformer + std::vector resblocks; + + struct ggml_tensor * pre_ln_w; // [hidden_size,] + struct ggml_tensor * pre_ln_b; // [hidden_size,] + // std::vector layers; + + struct ggml_tensor * post_ln_w; // [hidden_size,] + struct ggml_tensor * post_ln_b; // [hidden_size,] + + struct ggml_tensor* visual_projection; + + + CLIPVisionModel(CLIPVersion version = OPENAI_CLIP_VIT_L_14) + : version(version){ + + resblocks.resize(num_hidden_layers); + set_resblocks_hp_params(); + } + + void set_resblocks_hp_params() { + int d_model = hidden_size / n_head; // 64 / SDXL is 40 for CLIPTextModelWithProjection + for (int i = 0; i < num_hidden_layers; i++) { + resblocks[i].d_model = d_model; + resblocks[i].n_head = n_head; + resblocks[i].hidden_size = hidden_size; + resblocks[i].intermediate_size = intermediate_size; + } + } + + size_t calculate_mem_size(ggml_type wtype) { + size_t mem_size = 0; + for (int i = 0; i < num_hidden_layers; i++) { + mem_size += resblocks[i].calculate_mem_size(wtype); + } + mem_size += 4 * ggml_row_size(wtype, hidden_size); // + mem_size += 1 * ggml_row_size(wtype, hidden_size); // + mem_size += ggml_row_size(wtype, hidden_size*patch_size*patch_size*3); // + mem_size += ggml_row_size(wtype, hidden_size*257); // + mem_size += ggml_row_size(wtype, hidden_size * projection_dim); // visual_projection + + return mem_size; + } + + void map_by_name(std::map& tensors, const std::string prefix) { + tensors[prefix + "post_layernorm.bias"] = post_ln_b; + tensors[prefix + "post_layernorm.weight"] = post_ln_w; + tensors[prefix + "pre_layrnorm.bias"] = pre_ln_b; + tensors[prefix + "pre_layrnorm.weight"] = pre_ln_w; + + tensors[prefix + "embeddings.patch_embedding.weight"] = patch_embeddings; + tensors[prefix + "embeddings.position_embedding.weight"] = position_embeddings; + tensors[prefix + "embeddings.class_embedding"] = class_embedding; + for (int i = 0; i < num_hidden_layers; i++) { + std::string name = prefix + "encoder.layers." + std::to_string(i) + "."; + resblocks[i].map_by_name(tensors, prefix + "encoder.layers." + std::to_string(i) + "."); + } + tensors[prefix + "visual_projection"] = 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* input_ids, struct ggml_tensor* tkn_embeddings, uint32_t max_token_idx = 0, bool return_pooled = false) { + // input_ids: [N, n_token] + // GGML_ASSERT(input_ids->ne[0] <= position_ids->ne[0]); + + // token_embedding + position_embedding + struct ggml_tensor* x = NULL; + + return x; // [N, n_token, hidden_size] + } + + void init_params(ggml_context* ctx, ggml_backend_t backend, ggml_type wtype, ggml_allocr* alloc) { + class_embedding = ggml_new_tensor_1d(ctx, wtype, hidden_size); + + patch_embeddings = ggml_new_tensor_4d(ctx, wtype, patch_size, patch_size, 3, hidden_size); + + position_embeddings = ggml_new_tensor_2d(ctx, wtype, hidden_size, 257); + + for (int i = 0; i < num_hidden_layers; i++) { + resblocks[i].init_params(ctx, alloc, wtype); + } + + pre_ln_w = ggml_new_tensor_1d(ctx, wtype, hidden_size); + pre_ln_b = ggml_new_tensor_1d(ctx, wtype, hidden_size); + 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); + + + // 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); + } + } + + // if (ggml_backend_is_cpu(backend)) { + // for (int i = 0; i < max_position_embeddings; i++) { + // ggml_set_i32_1d(position_ids, i, i); + // } + // } else { + // std::vector pos_temp; + // for (int i = 0; i < max_position_embeddings; i++) { + // pos_temp.push_back(i); + // } + // ggml_backend_tensor_set(position_ids, pos_temp.data(), 0, ggml_nbytes(position_ids)); + // } + } +}; + // ldm.modules.encoders.modules.FrozenCLIPEmbedder // Ref: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/cad87bf4e3e0b0a759afa94e933527c3123d59bc/modules/sd_hijack_clip.py#L283 struct FrozenCLIPEmbedderWithCustomWords : public GGMLModule { diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 7acc4449..3198907c 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -63,6 +63,7 @@ struct SDParams { std::string esrgan_path; std::string controlnet_path; std::string embeddings_path; + std::string stacked_id_embeddings_path; sd_type_t wtype = SD_TYPE_COUNT; std::string lora_model_dir; std::string output_path = "output.png"; @@ -101,6 +102,7 @@ void print_params(SDParams params) { printf(" esrgan_path: %s\n", params.esrgan_path.c_str()); printf(" controlnet_path: %s\n", params.controlnet_path.c_str()); printf(" embeddings_path: %s\n", params.embeddings_path.c_str()); + printf(" stacked_id_embeddings_path: %s\n", params.stacked_id_embeddings_path.c_str()); printf(" output_path: %s\n", params.output_path.c_str()); printf(" init_img: %s\n", params.input_path.c_str()); printf(" control_image: %s\n", params.control_image_path.c_str()); @@ -135,6 +137,7 @@ void print_usage(int argc, const char* argv[]) { printf(" --taesd [TAESD_PATH] path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)\n"); printf(" --control-net [CONTROL_PATH] path to control net model\n"); printf(" --embd-dir [EMBEDDING_PATH] path to embeddings.\n"); + printf(" --stacked-id-embd-dir [ID_EMBEDDING_PATH] path to photomakerstacked id embeddings.\n"); printf(" --upscale-model [ESRGAN_PATH] path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now.\n"); printf(" --type [TYPE] weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)\n"); printf(" If not specified, the default is the type of the weight file.\n"); @@ -231,6 +234,12 @@ void parse_args(int argc, const char** argv, SDParams& params) { break; } params.embeddings_path = argv[i]; + } else if (arg == "--stacked-id-embd-dir") { + if (++i >= argc) { + invalid_arg = true; + break; + } + params.stacked_id_embeddings_path = argv[i]; } else if (arg == "--type") { if (++i >= argc) { invalid_arg = true; @@ -573,6 +582,7 @@ int main(int argc, const char* argv[]) { params.controlnet_path.c_str(), params.lora_model_dir.c_str(), params.embeddings_path.c_str(), + params.stacked_id_embeddings_path.c_str(), vae_decode_only, params.vae_tiling, true, diff --git a/model.cpp b/model.cpp index b89edf27..a889f764 100644 --- a/model.cpp +++ b/model.cpp @@ -788,7 +788,6 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const for (int i = 0; i < n_dims; i++) { ne[i] = shape[i].get(); } - TensorStorage tensor_storage(prefix + name, type, ne, n_dims, file_index, ST_HEADER_SIZE_LEN + header_size_ + begin); tensor_storage.reverse_ne(); @@ -1127,6 +1126,8 @@ bool ModelLoader::parse_data_pkl(uint8_t* buffer, if (reader.phase == PickleTensorReader::READ_DIMENS) { reader.tensor_storage.reverse_ne(); reader.tensor_storage.file_index = file_index; + // if(strcmp(prefix.c_str(), "pm") == 0) + printf(" got tensor %s \n ", reader.tensor_storage.name.c_str()); reader.tensor_storage.name = prefix + reader.tensor_storage.name; tensor_storages.push_back(reader.tensor_storage); // LOG_DEBUG("%s", reader.tensor_storage.name.c_str()); diff --git a/pmid.hpp b/pmid.hpp new file mode 100644 index 00000000..ac32d90c --- /dev/null +++ b/pmid.hpp @@ -0,0 +1,208 @@ +#ifndef __PMI_HPP__ +#define __PMI_HPP__ + +#include "ggml_extend.hpp" + +#include "clip.hpp" + + +struct FuseBlock { + // network hparams + int in_dim; + int out_dim; + int hidden_dim; + bool use_residue; + + // network params + // in_layers + + // layer norm + struct ggml_tensor* ln_w; // [in_dim, ] + struct ggml_tensor* ln_b; // [in_dim, ] + + struct ggml_tensor* fc1_w; // [in_dim, hidden_dim] + struct ggml_tensor* fc1_b; // [in_dim, ] + struct ggml_tensor* fc2_w; // [hidden_dim, out_dim ] + struct ggml_tensor* fc2_b; // [hidden_dim, ] + + + FuseBlock(int i_d, int o_d, int h_d, bool use_residue = true) + : in_dim(i_d), out_dim(o_d), hidden_dim(h_d), + use_residue(use_residue){ + + } + + + 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 + } + 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); + + 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) { + tensors[prefix + "fc1.weight"] = fc1_w; + tensors[prefix + "fc1.bias"] = fc1_b; + tensors[prefix + "fc2.weight"] = fc2_w; + tensors[prefix + "fc2.bias"] = fc2_b; + tensors[prefix + "layernorm.weight"] = ln_w; + tensors[prefix + "layernorm.bias"] = ln_b; + } + + struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) { + // x: [N, channels, h, w] + + // in_layers + auto h = ggml_nn_group_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); + if(use_residue) + x = ggml_add(ctx, x, h); + return h; + } + +}; + +struct FuseModule{ + // network hparams + int embed_dim; + + struct FuseBlock mlp1; + struct FuseBlock mlp2; + // layer norm + struct ggml_tensor* ln_w; // [in_dim, ] + struct ggml_tensor* ln_b; // [in_dim, ] + + + FuseModule(int imb_d) + : embed_dim(imb_d){ + + } + + + 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 map_by_name(std::map& tensors, const std::string prefix) { + + tensors[prefix + "layer_norm.weight"] = ln_w; + tensors[prefix + "layer_norm.bias"] = ln_b; + mlp1.map_by_name(tensors, prefix + ".mlp1."); + mlp2.map_by_name(tensors, prefix + ".mlp2."); + + } + + struct ggml_tensor* fuse_fn(struct ggml_context* ctx, + struct ggml_tensor* prompt_embeds, + struct ggml_tensor* id_embeds) { + // x: [N, channels, h, w] + + // in_layers + auto stacked_id_embeds = ggml_concat(ctx, prompt_embeds, id_embeds); // check whether concat at dim 2 is right + + stacked_id_embeds = mlp1.forward(ctx, stacked_id_embeds); + stacked_id_embeds = ggml_add(ctx, stacked_id_embeds, prompt_embeds); + stacked_id_embeds = mlp2.forward(ctx, stacked_id_embeds); + stacked_id_embeds = ggml_nn_group_norm(ctx, stacked_id_embeds, ln_w, ln_b); + return stacked_id_embeds; + + } + + + struct ggml_tensor* forward(struct ggml_context* ctx, + struct ggml_tensor* prompt_embeds, + struct ggml_tensor* id_embeds, + struct ggml_tensor* class_tokens_mask) { + // x: [N, channels, h, w] + + // in_layers + struct ggml_tensor* h = NULL; + return h; + } + + +}; + + +struct PhotoMakerIDEncoder : public GGMLModule { + SDVersion version = VERSION_XL; + CLIPVisionModel vision_model; + FuseModule fuse_module; + struct ggml_tensor* visual_projection_2; + + PhotoMakerIDEncoder(SDVersion version = VERSION_XL) + : version(version){ + + + + + + } + + void init_params(struct ggml_context* ctx, ggml_type wtype) { + + vision_model = CLIPVisionModel(); + fuse_module = FuseModule(2048); + visual_projection_2 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1280, 1024); // python [1024, 1280] + + } + + void map_by_name(std::map& tensors, const std::string prefix) { + // vision_model. + fuse_module.map_by_name(tensors, prefix + ".fuse_module."); + tensors[prefix + "visual_projection_2.weight"] = visual_projection_2; + } + + + 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) { + // x: [N, channels, h, w] + + // 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); + + // 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 + struct ggml_tensor * updated_prompt_embeds = fuse_module.forward(ctx, prompt_embeds, id_embeds, class_tokens_mask); + + return updated_prompt_embeds + + } + +}; + + + +#endif // __PMI_HPP__ + diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index 8dd5f16e..b5e61305 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -66,6 +66,7 @@ public: AutoEncoderKL first_stage_model; bool use_tiny_autoencoder = false; bool vae_tiling = false; + bool stacked_id = false; std::map tensors; @@ -111,6 +112,7 @@ public: const std::string& vae_path, const std::string control_net_path, const std::string embeddings_path, + const std::string id_embeddings_path, const std::string& taesd_path, bool vae_tiling_, ggml_type wtype, @@ -194,6 +196,16 @@ public: cond_stage_model.text_model.embd_dir = embeddings_path; + if (id_embeddings_path.size() > 0) { + LOG_INFO("loading stacked ID embedding (PHOTOMAKER) model file from '%s'", id_embeddings_path.c_str()); + if (!model_loader.init_from_file(id_embeddings_path)) { + LOG_WARN("loading stacked ID embedding from '%s' failed", id_embeddings_path.c_str()); + } + else{ + stacked_id = true; + } + } + ggml_type vae_type = model_data_type; if (version == VERSION_XL) { vae_type = GGML_TYPE_F32; // avoid nan, not work... @@ -1142,6 +1154,7 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str, const char* control_net_path_c_str, const char* lora_model_dir_c_str, const char* embed_dir_c_str, + const char* id_embed_dir_c_str, bool vae_decode_only, bool vae_tiling, bool free_params_immediately, @@ -1159,6 +1172,7 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str, std::string taesd_path(taesd_path_c_str); std::string control_net_path(control_net_path_c_str); std::string embd_path(embed_dir_c_str); + std::string id_embd_path(id_embed_dir_c_str); std::string lora_model_dir(lora_model_dir_c_str); sd_ctx->sd = new StableDiffusionGGML(n_threads, @@ -1174,6 +1188,7 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str, vae_path, control_net_path, embd_path, + id_embd_path, taesd_path, vae_tiling, (ggml_type)wtype, diff --git a/stable-diffusion.h b/stable-diffusion.h index a8c9f532..a395417b 100644 --- a/stable-diffusion.h +++ b/stable-diffusion.h @@ -108,6 +108,7 @@ SD_API sd_ctx_t* new_sd_ctx(const char* model_path, const char* control_net_path_c_str, const char* lora_model_dir, const char* embed_dir_c_str, + const char* stacked_id_embed_dir_c_str, bool vae_decode_only, bool vae_tiling, bool free_params_immediately,