feat: add IP-Adapter support for SD 1.5 and SDXL (#1803)

src/model/adapter/ip_adapter.hpp

@@ -0,0 +1,88 @@
+#ifndef __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
+#define __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
+
+#include "core/ggml_extend.hpp"
+#include "model/common/block.hpp"
+#include "model_loader.h"
+
+namespace IPAdapter {
+
+    struct ImageProjModel : public GGMLBlock {
+        int64_t num_tokens = 4;
+        int64_t ctx_dim    = 768;
+        int64_t clip_dim   = 1024;
+
+        ImageProjModel() {}
+        ImageProjModel(int64_t num_tokens, int64_t ctx_dim, int64_t clip_dim)
+            : num_tokens(num_tokens), ctx_dim(ctx_dim), clip_dim(clip_dim) {
+            blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(clip_dim, num_tokens * ctx_dim, true));
+            blocks["norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(ctx_dim));
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image_embeds) {
+            auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
+            auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
+
+            int64_t n = image_embeds->ne[1];
+            auto x    = proj->forward(ctx, image_embeds);
+            x         = ggml_reshape_3d(ctx->ggml_ctx, x, ctx_dim, num_tokens, n);
+            x         = norm->forward(ctx, x);
+            return x;
+        }
+    };
+
+    struct IPAdapterRunner : public GGMLRunner {
+        ImageProjModel image_proj;
+        int64_t num_tokens = 4;
+        std::string prefix;
+
+        IPAdapterRunner(ggml_backend_t backend,
+                        const String2TensorStorage& tensor_storage_map,
+                        const std::string prefix,
+                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : GGMLRunner(backend, weight_manager), prefix(prefix) {
+            int64_t ctx_dim  = 768;
+            int64_t clip_dim = 1024;
+            int64_t out_dim  = 3072;
+            auto norm_iter   = tensor_storage_map.find(prefix + ".image_proj.norm.weight");
+            if (norm_iter != tensor_storage_map.end()) {
+                ctx_dim = norm_iter->second.ne[0];
+            }
+            auto proj_iter = tensor_storage_map.find(prefix + ".image_proj.proj.weight");
+            if (proj_iter != tensor_storage_map.end()) {
+                clip_dim = proj_iter->second.ne[0];
+                out_dim  = proj_iter->second.ne[1];
+            }
+            num_tokens = out_dim / ctx_dim;
+            image_proj = ImageProjModel(num_tokens, ctx_dim, clip_dim);
+            image_proj.init(params_ctx, tensor_storage_map, prefix + ".image_proj");
+        }
+
+        std::string get_desc() override {
+            return "ip_adapter";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string = "") {
+            image_proj.get_param_tensors(tensors, prefix + ".image_proj");
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& image_embeds_tensor) {
+            ggml_cgraph* gf     = new_graph_custom(1024);
+            ggml_tensor* embeds = make_input(image_embeds_tensor);
+            auto runner_ctx     = get_context();
+            ggml_tensor* out    = image_proj.forward(&runner_ctx, embeds);
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& image_embeds) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(image_embeds);
+            };
+            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
+        }
+    };
+
+}  // namespace IPAdapter
+
+#endif  // __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__

src/model/common/block.hpp

@@ -310,17 +310,33 @@ protected:
     int64_t context_dim;
     int64_t n_head;
     int64_t d_head;
-    bool xtra_dim = false;
+    bool xtra_dim  = false;
+    bool enable_ip = false;
+    bool has_ip    = false;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
+        if (enable_ip &&
+            tensor_storage_map.find(prefix + "to_k_ip.weight") != tensor_storage_map.end()) {
+            has_ip            = true;
+            int64_t inner_dim = d_head * n_head;
+            int64_t ip_dim    = tensor_storage_map.at(prefix + "to_k_ip.weight").ne[0];
+            blocks["to_k_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
+            blocks["to_v_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
+        }
+    }
 
 public:
     CrossAttention(int64_t query_dim,
                    int64_t context_dim,
                    int64_t n_head,
-                   int64_t d_head)
+                   int64_t d_head,
+                   bool enable_ip = false)
         : n_head(n_head),
           d_head(d_head),
           query_dim(query_dim),
-          context_dim(context_dim) {
+          context_dim(context_dim),
+          enable_ip(enable_ip) {
         int64_t inner_dim = d_head * n_head;
         if (context_dim == 320 && d_head == 320) {
             // LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09");
@@ -363,6 +379,15 @@ public:
         }
         x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, inner_dim]
 
+        if (has_ip && ctx->ip_context != nullptr && ctx->ip_scale != 0.0f) {
+            auto to_k_ip = std::dynamic_pointer_cast<Linear>(blocks["to_k_ip"]);
+            auto to_v_ip = std::dynamic_pointer_cast<Linear>(blocks["to_v_ip"]);
+            auto k_ip    = to_k_ip->forward(ctx, ctx->ip_context);
+            auto v_ip    = to_v_ip->forward(ctx, ctx->ip_context);
+            auto x_ip    = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_ip, v_ip, n_head, nullptr, false, ctx->flash_attn_enabled);
+            x            = ggml_add(ctx->ggml_ctx, x, ggml_scale(ctx->ggml_ctx, x_ip, ctx->ip_scale));
+        }
+
         x = to_out_0->forward(ctx, x);  // [N, n_token, query_dim]
         return x;
     }
@@ -387,7 +412,7 @@ public:
         // inner_dim is always None or equal to dim
         // gated_ff is always True
         blocks["attn1"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, dim, n_head, d_head));
-        blocks["attn2"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, context_dim, n_head, d_head));
+        blocks["attn2"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, context_dim, n_head, d_head, true));
         blocks["ff"]    = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim));
         blocks["norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
         blocks["norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));

src/model/diffusion/model.hpp

@@ -46,6 +46,8 @@ struct UNetDiffusionExtra {
     int num_video_frames                           = -1;
     const std::vector<sd::Tensor<float>>* controls = nullptr;
     float control_strength                         = 0.f;
+    const sd::Tensor<float>* ip_context            = nullptr;
+    float ip_scale                                 = 1.f;
 };
 
 struct SkipLayerDiffusionExtra {

src/model/diffusion/unet.hpp

@@ -775,14 +775,17 @@ struct UNetModelRunner : public DiffusionModelRunner {
                              const sd::Tensor<float>& y_tensor                     = {},
                              int num_video_frames                                  = -1,
                              const std::vector<sd::Tensor<float>>& controls_tensor = {},
-                             float control_strength                                = 0.f) {
+                             float control_strength                                = 0.f,
+                             const sd::Tensor<float>& ip_context_tensor            = {},
+                             float ip_scale                                        = 1.f) {
         ggml_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE);
 
-        ggml_tensor* x         = make_input(x_tensor);
-        ggml_tensor* timesteps = make_input(timesteps_tensor);
-        ggml_tensor* context   = make_optional_input(context_tensor);
-        ggml_tensor* c_concat  = make_optional_input(c_concat_tensor);
-        ggml_tensor* y         = make_optional_input(y_tensor);
+        ggml_tensor* x          = make_input(x_tensor);
+        ggml_tensor* timesteps  = make_input(timesteps_tensor);
+        ggml_tensor* context    = make_optional_input(context_tensor);
+        ggml_tensor* c_concat   = make_optional_input(c_concat_tensor);
+        ggml_tensor* y          = make_optional_input(y_tensor);
+        ggml_tensor* ip_context = make_optional_input(ip_context_tensor);
         std::vector<ggml_tensor*> controls;
         controls.reserve(controls_tensor.size());
         for (const auto& control_tensor : controls_tensor) {
@@ -793,7 +796,9 @@ struct UNetModelRunner : public DiffusionModelRunner {
             num_video_frames = static_cast<int>(x->ne[3]);
         }
 
-        auto runner_ctx = get_context();
+        auto runner_ctx       = get_context();
+        runner_ctx.ip_context = ip_context;
+        runner_ctx.ip_scale   = ip_scale;
 
         ggml_tensor* out = unet.forward(&runner_ctx,
                                         x,
@@ -818,14 +823,16 @@ struct UNetModelRunner : public DiffusionModelRunner {
                               const sd::Tensor<float>& y                     = {},
                               int num_video_frames                           = -1,
                               const std::vector<sd::Tensor<float>>& controls = {},
-                              float control_strength                         = 0.f) {
+                              float control_strength                         = 0.f,
+                              const sd::Tensor<float>& ip_context            = {},
+                              float ip_scale                                 = 1.f) {
         // x: [N, in_channels, h, w]
         // timesteps: [N, ]
         // context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
         // c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
         // y: [N, adm_in_channels] or [1, adm_in_channels]
         auto get_graph = [&]() -> ggml_cgraph* {
-            return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength);
+            return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, ip_context, ip_scale);
         };
 
         return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -845,7 +852,9 @@ struct UNetModelRunner : public DiffusionModelRunner {
                        tensor_or_empty(diffusion_params.y),
                        extra->num_video_frames,
                        extra->controls ? *extra->controls : empty_controls,
-                       extra->control_strength);
+                       extra->control_strength,
+                       extra->ip_context ? *extra->ip_context : sd::Tensor<float>{},
+                       extra->ip_scale);
     }
 
     void test() {