diff --git a/pmid.hpp b/pmid.hpp index 7dec72dc..94f3e8bf 100644 --- a/pmid.hpp +++ b/pmid.hpp @@ -61,6 +61,11 @@ struct FuseBlock { } + + size_t get_num_tensors() { + return 6; + } + void map_by_name(std::map& tensors, const std::string prefix) { tensors[prefix + "fc1.weight"] = fc1_w; tensors[prefix + "fc1.bias"] = fc1_b; @@ -130,6 +135,13 @@ struct FuseModule{ } + size_t get_num_tensors() { + size_t n = mlp1.get_num_tensors(); + n += mlp2.get_num_tensors(); + n += 2; + return n; + } + size_t calculate_mem_size(ggml_type wtype) { size_t mem_size = mlp1.calculate_mem_size(wtype); mem_size += mlp2.calculate_mem_size(wtype); @@ -213,8 +225,8 @@ struct PhotoMakerIDEncoder : public GGMLModule { } size_t get_num_tensors() { - size_t num_tensors = (3 + 2 + 37); - + size_t num_tensors = (3 + 2 + 37 * vision_model.num_hidden_layers); + num_tensors += fuse_module.get_num_tensors() + 1; return num_tensors; } diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index 61d153bd..21a6d4c0 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -86,6 +86,8 @@ public: ControlNet control_net; + std::string trigger_word = "img"; // should be user settable + StableDiffusionGGML() = default; StableDiffusionGGML(int n_threads,