first efforts at implementing photomaker; lots more to do

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)

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<ResidualAttentionBlock> resblocks;
+    
+    struct ggml_tensor * pre_ln_w; // [hidden_size,]
+    struct ggml_tensor * pre_ln_b; // [hidden_size,]
+    // std::vector<clip_layer> 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<std::string, struct ggml_tensor*>& 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<int> 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 {

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,

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<int64_t>();
         }
-
         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());

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<std::string, struct ggml_tensor*>& 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<std::string, struct ggml_tensor*>& 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<std::string, struct ggml_tensor*>& 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__
+

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<std::string, struct ggml_tensor*> 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,

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,