feat: add linear multi-step sampling method (#1843)

Co-authored-by: leejet <leejet714@gmail.com>

README.md

@@ -70,7 +70,7 @@ API and command-line option may change frequently.***
     - [HunyuanVideo 1.5](./docs/hunyuan_video.md)
     - [LingBot-Video](./docs/lingbot_video.md)
   - [PhotoMaker](./docs/photo_maker.md) support.
-  - [IP-Adapter](./docs/ip_adapter.md) support (SD 1.5 and SDXL)
+  - [IP-Adapter](./docs/ip_adapter.md) support (SD 1.5 and SDXL, including Plus)
   - Control Net support with SD 1.5
   - [ADetailer](./docs/adetailer.md)
   - LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora)

docs/ip_adapter.md

@@ -11,6 +11,10 @@ through a decoupled cross-attention added to every attn2 layer of the
 UNet. It composes with Control Net, so a reference image (appearance) and
 an OpenPose hint (pose) can be combined in a single generation.
 
+Both the classic adapters and the higher-fidelity **Plus** adapters are
+supported; see [Plus variants](#plus-variants) below. The variant is
+detected from the weight file, so the same options work for both.
+
 ## Required weights
 
 1. A base SD 1.5 or SDXL model.
@@ -21,6 +25,11 @@ an OpenPose hint (pose) can be combined in a single generation.
    [h94/IP-Adapter](https://huggingface.co/h94/IP-Adapter):
    - SD 1.5: `models/ip-adapter_sd15.safetensors`
    - SDXL: `sdxl_models/ip-adapter_sdxl_vit-h.safetensors`
+   - SD 1.5 Plus: `models/ip-adapter-plus_sd15.safetensors`
+   - SDXL Plus: `sdxl_models/ip-adapter-plus_sdxl_vit-h.safetensors`
+
+   The Plus files (`ip-adapter-plus_*`) are used exactly like the classic
+   ones; see [Plus variants](#plus-variants).
 
 ## Options
 
@@ -45,6 +54,29 @@ sd-cli -m ..\models\sdxl.safetensors --clip_vision ..\models\clip_vision_h.safet
 The SDXL VAE decode at 1024x1024 is memory heavy; add `--vae-tiling` (and
 `--offload-to-cpu`) on GPUs with limited VRAM.
 
+## Plus variants
+
+The Plus adapters (`ip-adapter-plus_sd15`, `ip-adapter-plus_sdxl_vit-h`)
+replace the small linear image projection with a Resampler (a
+Perceiver-style module with learned latent queries). Instead of pooling the
+CLIP-Vision output into one vector, the Resampler attends over the full grid
+of penultimate CLIP-Vision hidden states and emits more image tokens (16
+instead of 4). The result transfers finer detail and layout from the
+reference, at a small extra cost in the image-projection step.
+
+No extra flags are needed. The variant is detected from the weight file (the
+Resampler's `image_proj.latents` tensor), and every Resampler dimension is
+read from the tensor shapes, so the same `--ip-adapter`,
+`--ip-adapter-image`, and `--ip-adapter-strength` options apply. Plus
+composes with Control Net in the same way as the classic adapters.
+
+```
+sd-cli -m ..\models\sd_v1.5.safetensors --clip_vision ..\models\clip_vision_h.safetensors --ip-adapter ..\models\ip-adapter-plus_sd15.safetensors --ip-adapter-image ..\assets\reference.png --ip-adapter-strength 0.8 -p "a woman, best quality" -n "lowres, bad anatomy" --cfg-scale 7 --steps 30 --sampling-method dpm++2m --scheduler karras -W 512 -H 512
+```
+
+The startup log line `IP-Adapter: 16 image tokens` (versus `4` for the
+classic adapters) confirms a Plus file was loaded.
+
 ## Combining with Control Net
 
 Add the usual Control Net options to keep the reference appearance while

examples/common/common.cpp

@@ -1008,7 +1008,7 @@ ArgOptions SDGenerationParams::get_options() {
          &hires_upscaler},
         {"",
          "--extra-sample-args",
-         "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware",
+         "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_divisions",
          (int)',',
          &extra_sample_args},
         {"",
@@ -1538,12 +1538,12 @@ ArgOptions SDGenerationParams::get_options() {
          on_seed_arg},
         {"",
          "--sampling-method",
-         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
+         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp, lms]"
          "(default: euler for Flux/SD3/Wan, euler_a otherwise)",
          on_sample_method_arg},
         {"",
          "--high-noise-sampling-method",
-         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
+         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp, lms]"
          " default: euler for Flux/SD3/Wan, euler_a otherwise",
          on_high_noise_sample_method_arg},
         {"",
@@ -1901,6 +1901,7 @@ bool SDGenerationParams::from_json_str(
 
     load_if_exists("strength", strength);
     load_if_exists("control_strength", control_strength);
+    load_if_exists("ip_adapter_strength", ip_adapter_strength);
     load_if_exists("moe_boundary", moe_boundary);
     load_if_exists("vace_strength", vace_strength);
 
@@ -2072,6 +2073,10 @@ bool SDGenerationParams::from_json_str(
         LOG_ERROR("invalid control_image");
         return false;
     }
+    if (!parse_image_json_field(j, "ip_adapter_image", 3, width, height, ip_adapter_image)) {
+        LOG_ERROR("invalid ip_adapter_image");
+        return false;
+    }
 
     return true;
 }
@@ -2807,6 +2812,7 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
     root["clip_skip"]             = gen_params.clip_skip;
     root["strength"]              = gen_params.strength;
     root["control_strength"]      = gen_params.control_strength;
+    root["ip_adapter_strength"]   = gen_params.ip_adapter_strength;
     root["auto_resize_ref_image"] = gen_params.auto_resize_ref_image;
     root["increase_ref_index"]    = gen_params.increase_ref_index;
     if (mode == VID_GEN) {

examples/server/api.md

@@ -528,6 +528,7 @@ Shared default fields used by both `img_gen` and `vid_gen`:
 | `auto_resize_ref_image` | `boolean` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
+| `ip_adapter_strength` | `number` |
 | `hires` | `object` |
 | `hires.enabled` | `boolean` |
 | `hires.upscaler` | `string` |
@@ -567,6 +568,7 @@ Fields returned in `features_by_mode.img_gen`:
 - `init_image`
 - `mask_image`
 - `control_image`
+- `ip_adapter_image`
 - `ref_images`
 - `lora`
 - `vae_tiling`
@@ -653,12 +655,14 @@ Example:
   "auto_resize_ref_image": true,
   "increase_ref_index": false,
   "control_strength": 0.9,
+  "ip_adapter_strength": 1.0,
   "embed_image_metadata": true,
 
   "init_image": null,
   "ref_images": [],
   "mask_image": null,
   "control_image": null,
+  "ip_adapter_image": null,
 
   "sample_params": {
     "scheduler": "discrete",
@@ -733,6 +737,7 @@ Channel expectations:
 - `init_image`: 3 channels
 - `ref_images[]`: 3 channels
 - `control_image`: 3 channels
+- `ip_adapter_image`: 3 channels
 - `mask_image`: 1 channel
 
 If omitted or null:
@@ -757,6 +762,7 @@ Top-level scalar fields:
 | `auto_resize_ref_image` | `boolean` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
+| `ip_adapter_strength` | `number` |
 | `embed_image_metadata` | `boolean` |
 
 Image fields:
@@ -767,6 +773,7 @@ Image fields:
 | `ref_images` | `array<string>` |
 | `mask_image` | `string \| null` |
 | `control_image` | `string \| null` |
+| `ip_adapter_image` | `string \| null` |
 
 LoRA fields:
 
@@ -958,7 +965,7 @@ Response fields:
 Compared with `img_gen`, the `vid_gen` request body:
 
 - `vid_gen` is a single video sequence job, so `batch_count` is not part of the request schema
-- `ref_images`, `mask_image`, `control_image`, `control_strength`, and `embed_image_metadata` are not part of the request schema
+- `ref_images`, `mask_image`, `control_image`, `control_strength`, `ip_adapter_image`, `ip_adapter_strength`, and `embed_image_metadata` are not part of the request schema
 - `vid_gen` adds `end_image`, `control_frames`, `high_noise_sample_params`, `video_frames`, `fps`, `moe_boundary`, and `vace_strength`
 
 Example:

examples/server/routes_sdcpp.cpp

@@ -130,6 +130,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
         {"auto_resize_ref_image", defaults.auto_resize_ref_image},
         {"increase_ref_index", defaults.increase_ref_index},
         {"control_strength", defaults.control_strength},
+        {"ip_adapter_strength", defaults.ip_adapter_strength},
         {"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers)},
         {"hires", make_hires_json(defaults)},
         {"vae_tiling_params", make_vae_tiling_json(defaults.vae_tiling_params)},
@@ -173,6 +174,7 @@ static json make_img_gen_features_json() {
         {"init_image", true},
         {"mask_image", true},
         {"control_image", true},
+        {"ip_adapter_image", true},
         {"ref_images", true},
         {"lora", true},
         {"vae_tiling", true},

include/stable-diffusion.h

@@ -56,6 +56,7 @@ enum sample_method_t {
     EULER_GE_SAMPLE_METHOD,
     DPMPP2M_SDE_SAMPLE_METHOD,
     DPMPP2M_SDE_BT_SAMPLE_METHOD,
+    LMS_SAMPLE_METHOD,
     SAMPLE_METHOD_COUNT
 };
 

src/conditioning/conditioner.hpp

@@ -117,6 +117,7 @@ public:
     virtual SDCondition get_learned_condition(int n_threads,
                                               const ConditionerParams& conditioner_params) = 0;
     virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors)           = 0;
+    virtual void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {}
     virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
     virtual void set_stream_layers_enabled(bool enabled) {}
     virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
@@ -1664,6 +1665,10 @@ struct AnimaConditioner : public Conditioner {
         llm->get_param_tensors(tensors, "text_encoders.llm");
     }
 
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
+        llm->get_param_tensor_ops(tensor_ops);
+    }
+
     void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
         llm->set_max_graph_vram_bytes(max_vram_bytes);
     }
@@ -1847,6 +1852,10 @@ struct LLMEmbedder : public Conditioner {
         }
     }
 
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
+        llm->get_param_tensor_ops(tensor_ops);
+    }
+
     void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
         llm->set_max_graph_vram_bytes(max_vram_bytes);
         if (byt5) {
@@ -2828,6 +2837,10 @@ struct LTXAVEmbedder : public Conditioner {
         projector->get_param_tensors(tensors, "text_embedding_projection");
     }
 
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
+        llm->get_param_tensor_ops(tensor_ops);
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
         projector->set_flash_attention_enabled(enabled);

src/core/ggml_extend.hpp

@@ -1753,7 +1753,7 @@ protected:
     std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
 
     std::vector<ggml_backend_t> extra_runtime_backends;  // borrowed (SDBackendManager-owned)
-    ggml_backend_sched_t sched             = nullptr;    // owned, multi-device only
+    ggml_backend_sched_t sched             = nullptr;    // owned
     ggml_backend_t cpu_fallback_backend    = nullptr;    // owned, sched requires a trailing CPU backend
     bool multi_device_eval_callback_warned = false;
 
@@ -2147,8 +2147,22 @@ protected:
         return !extra_runtime_backends.empty();
     }
 
+    bool graph_requires_backend_fallback(ggml_cgraph* gf) const {
+        if (gf == nullptr || sd_backend_is_cpu(runtime_backend)) {
+            return false;
+        }
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        for (int i = 0; i < n_nodes; ++i) {
+            ggml_tensor* node = ggml_graph_node(gf, i);
+            if (node != nullptr && !ggml_backend_supports_op(runtime_backend, node)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     bool alloc_compute_buffer(ggml_cgraph* gf) {
-        if (is_multi_device()) {
+        if (sched != nullptr || is_multi_device() || graph_requires_backend_fallback(gf)) {
             // The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve
             // the graph here: reserve runs split_graph, which rewires the
             // graph's src pointers to sched-internal copy tensors, and the
@@ -2156,6 +2170,10 @@ protected:
             // rewired graph, silently corrupting every cross-backend input. A
             // graph must be split at most once; the alloc in execute_graph
             // performs the real allocation.
+            if (compute_allocr != nullptr) {
+                ggml_gallocr_free(compute_allocr);
+                compute_allocr = nullptr;
+            }
             return ensure_sched(gf);
         }
         if (compute_allocr != nullptr) {
@@ -2753,7 +2771,7 @@ protected:
         };
         ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer);
 
-        if (is_multi_device()) {
+        if (sched != nullptr) {
             ggml_backend_sched_reset(sched);
             pin_multi_device_nodes(gf);  // reset clears the pins; re-apply before alloc
             if (!ggml_backend_sched_alloc_graph(sched, gf)) {
@@ -2774,9 +2792,9 @@ protected:
         }
 
         ggml_status status;
-        if (is_multi_device()) {
+        if (sched != nullptr) {
             if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
-                LOG_WARN("%s: eval callback is not supported with multiple runtime backends; ignoring",
+                LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring",
                          get_desc().c_str());
                 multi_device_eval_callback_warned = true;
             }
@@ -3018,12 +3036,9 @@ public:
 
     // do copy after alloc graph
     void set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
-        if (is_multi_device()) {
-            // The sched only assigns a backend (and thus a buffer) to tensors
-            // that participate in the graph; flag standalone data tensors as
-            // inputs so they get one.
-            ggml_set_input(tensor);
-        }
+        // The scheduler only allocates standalone data tensors when they are
+        // marked as graph inputs. The flag is harmless for single-backend graphs.
+        ggml_set_input(tensor);
         backend_tensor_data_map[tensor] = data;
     }
 
@@ -3240,6 +3255,11 @@ protected:
 
     virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {}
 
+    virtual enum ggml_op param_usage_op(const std::string& name) const {
+        (void)name;
+        return GGML_OP_NONE;
+    }
+
 public:
     void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") {
         if (prefix.size() > 0) {
@@ -3290,6 +3310,18 @@ public:
         }
     }
 
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
+        for (auto& pair : blocks) {
+            pair.second->get_param_tensor_ops(tensor_ops);
+        }
+        for (auto& pair : params) {
+            enum ggml_op op = param_usage_op(pair.first);
+            if (op != GGML_OP_NONE) {
+                tensor_ops[pair.second] = op;
+            }
+        }
+    }
+
     virtual std::string get_desc() {
         return "GGMLBlock";
     }
@@ -3417,6 +3449,10 @@ protected:
         params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
     }
 
+    enum ggml_op param_usage_op(const std::string& name) const override {
+        return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
+    }
+
 public:
     Embedding(int64_t num_embeddings, int64_t embedding_dim)
         : embedding_dim(embedding_dim),

src/extensions/photomaker_extension.cpp

@@ -56,7 +56,7 @@ tokenize_photomaker_trigger(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
                                           true);
     std::vector<bool> class_token_mask;
     for (int i = 0; i < tokens.size(); i++) {
-        class_token_mask.push_back(class_idx + 1 <= i && i < class_idx + 1 + trigger_token_count);
+        class_token_mask.push_back(class_idx >= 0 && class_idx + 1 <= i && i < class_idx + 1 + trigger_token_count);
     }
 
     return std::make_tuple(tokens, weights, class_token_mask);

src/model/adapter/ip_adapter.hpp

@@ -31,8 +31,92 @@ namespace IPAdapter {
         }
     };
 
+    struct Resampler : public GGMLBlock {
+        int64_t dim         = 1280;
+        int64_t depth       = 4;
+        int64_t num_queries = 16;
+        int64_t embed_dim   = 1280;
+        int64_t output_dim  = 2048;
+        int64_t ff_inner    = 5120;
+        int64_t dim_head    = 64;
+        int64_t heads       = 20;
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            params["latents"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, dim, num_queries, 1);
+        }
+
+        Resampler() {}
+        Resampler(int64_t dim, int64_t depth, int64_t num_queries, int64_t embed_dim, int64_t output_dim, int64_t ff_inner)
+            : dim(dim), depth(depth), num_queries(num_queries), embed_dim(embed_dim), output_dim(output_dim), ff_inner(ff_inner) {
+            heads              = dim / dim_head;
+            blocks["proj_in"]  = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, dim, true));
+            blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(dim, output_dim, true));
+            blocks["norm_out"] = std::shared_ptr<GGMLBlock>(new LayerNorm(output_dim));
+            for (int64_t i = 0; i < depth; i++) {
+                std::string p           = "layers." + std::to_string(i);
+                blocks[p + ".0.norm1"]  = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
+                blocks[p + ".0.norm2"]  = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
+                blocks[p + ".0.to_q"]   = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, false));
+                blocks[p + ".0.to_kv"]  = std::shared_ptr<GGMLBlock>(new Linear(dim, dim * 2, false));
+                blocks[p + ".0.to_out"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, false));
+                blocks[p + ".1.0"]      = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
+                blocks[p + ".1.1"]      = std::shared_ptr<GGMLBlock>(new Linear(dim, ff_inner, false));
+                blocks[p + ".1.3"]      = std::shared_ptr<GGMLBlock>(new Linear(ff_inner, dim, false));
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image_embeds) {
+            int64_t N     = image_embeds->ne[2];
+            auto proj_in  = std::dynamic_pointer_cast<Linear>(blocks["proj_in"]);
+            auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
+            auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
+
+            ggml_tensor* x       = proj_in->forward(ctx, image_embeds);
+            ggml_tensor* latents = params["latents"];
+            if (N > 1) {
+                latents = ggml_repeat(ctx->ggml_ctx, latents, ggml_new_tensor_3d(ctx->ggml_ctx, GGML_TYPE_F32, dim, num_queries, N));
+            }
+
+            for (int64_t i = 0; i < depth; i++) {
+                std::string p = "layers." + std::to_string(i);
+                auto norm1    = std::dynamic_pointer_cast<LayerNorm>(blocks[p + ".0.norm1"]);
+                auto norm2    = std::dynamic_pointer_cast<LayerNorm>(blocks[p + ".0.norm2"]);
+                auto to_q     = std::dynamic_pointer_cast<Linear>(blocks[p + ".0.to_q"]);
+                auto to_kv    = std::dynamic_pointer_cast<Linear>(blocks[p + ".0.to_kv"]);
+                auto to_out   = std::dynamic_pointer_cast<Linear>(blocks[p + ".0.to_out"]);
+
+                ggml_tensor* xn    = norm1->forward(ctx, x);
+                ggml_tensor* ln    = norm2->forward(ctx, latents);
+                ggml_tensor* q     = to_q->forward(ctx, ln);
+                ggml_tensor* kv_in = ggml_concat(ctx->ggml_ctx, xn, ln, 1);
+                ggml_tensor* kv    = to_kv->forward(ctx, kv_in);
+                int64_t L          = kv->ne[1];
+                ggml_tensor* k     = ggml_cont(ctx->ggml_ctx, ggml_view_3d(ctx->ggml_ctx, kv, dim, L, N, kv->nb[1], kv->nb[2], 0));
+                ggml_tensor* v     = ggml_cont(ctx->ggml_ctx, ggml_view_3d(ctx->ggml_ctx, kv, dim, L, N, kv->nb[1], kv->nb[2], dim * kv->nb[0]));
+                ggml_tensor* attn  = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, heads, nullptr, false, false);
+                attn               = to_out->forward(ctx, attn);
+                latents            = ggml_add(ctx->ggml_ctx, latents, attn);
+
+                auto ff_norm   = std::dynamic_pointer_cast<LayerNorm>(blocks[p + ".1.0"]);
+                auto ff_fc1    = std::dynamic_pointer_cast<Linear>(blocks[p + ".1.1"]);
+                auto ff_fc2    = std::dynamic_pointer_cast<Linear>(blocks[p + ".1.3"]);
+                ggml_tensor* h = ff_norm->forward(ctx, latents);
+                h              = ff_fc1->forward(ctx, h);
+                h              = ggml_gelu_erf(ctx->ggml_ctx, h);
+                h              = ff_fc2->forward(ctx, h);
+                latents        = ggml_add(ctx->ggml_ctx, latents, h);
+            }
+
+            latents = proj_out->forward(ctx, latents);
+            latents = norm_out->forward(ctx, latents);
+            return latents;
+        }
+    };
+
     struct IPAdapterRunner : public GGMLRunner {
         ImageProjModel image_proj;
+        Resampler resampler;
+        bool is_plus       = false;
         int64_t num_tokens = 4;
         std::string prefix;
 
@@ -41,21 +125,54 @@ namespace IPAdapter {
                         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];
+            is_plus = tensor_storage_map.find(prefix + ".image_proj.latents") != tensor_storage_map.end();
+            if (is_plus) {
+                int64_t dim         = 1280;
+                int64_t num_queries = 16;
+                int64_t embed_dim   = 1280;
+                int64_t output_dim  = 2048;
+                int64_t ff_inner    = 5120;
+                auto latents_iter   = tensor_storage_map.find(prefix + ".image_proj.latents");
+                if (latents_iter != tensor_storage_map.end()) {
+                    dim         = latents_iter->second.ne[0];
+                    num_queries = latents_iter->second.ne[1];
+                }
+                auto proj_in_iter = tensor_storage_map.find(prefix + ".image_proj.proj_in.weight");
+                if (proj_in_iter != tensor_storage_map.end()) {
+                    embed_dim = proj_in_iter->second.ne[0];
+                }
+                auto proj_out_iter = tensor_storage_map.find(prefix + ".image_proj.proj_out.weight");
+                if (proj_out_iter != tensor_storage_map.end()) {
+                    output_dim = proj_out_iter->second.ne[1];
+                }
+                auto ff_iter = tensor_storage_map.find(prefix + ".image_proj.layers.0.1.1.weight");
+                if (ff_iter != tensor_storage_map.end()) {
+                    ff_inner = ff_iter->second.ne[1];
+                }
+                int64_t depth = 0;
+                while (tensor_storage_map.find(prefix + ".image_proj.layers." + std::to_string(depth) + ".0.to_q.weight") != tensor_storage_map.end()) {
+                    depth++;
+                }
+                num_tokens = num_queries;
+                resampler  = Resampler(dim, depth, num_queries, embed_dim, output_dim, ff_inner);
+                resampler.init(params_ctx, tensor_storage_map, prefix + ".image_proj");
+            } else {
+                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");
             }
-            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 {
@@ -63,14 +180,18 @@ namespace IPAdapter {
         }
 
         void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string = "") {
-            image_proj.get_param_tensors(tensors, prefix + ".image_proj");
+            if (is_plus) {
+                resampler.get_param_tensors(tensors, prefix + ".image_proj");
+            } else {
+                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_tensor* out    = is_plus ? resampler.forward(&runner_ctx, embeds) : image_proj.forward(&runner_ctx, embeds);
             ggml_build_forward_expand(gf, out);
             return gf;
         }

src/model/adapter/lora.hpp

@@ -14,6 +14,8 @@ struct LoraModel : public GGMLRunner {
     std::unordered_map<std::string, ggml_tensor*> lora_tensors;
     std::map<ggml_tensor*, ggml_tensor*> original_tensor_to_final_tensor;
     std::set<std::string> applied_lora_tensors;
+    std::set<std::string> skipped_incompatible_lora_tensors;
+    std::set<std::string> warned_incompatible_model_tensors;
     std::string file_path;
     std::shared_ptr<ModelManager> model_manager;
     ggml_backend_t params_backend = nullptr;
@@ -133,6 +135,8 @@ struct LoraModel : public GGMLRunner {
         lora_tensors.clear();
         original_tensor_to_final_tensor.clear();
         applied_lora_tensors.clear();
+        skipped_incompatible_lora_tensors.clear();
+        warned_incompatible_model_tensors.clear();
         applied             = false;
         tensor_preprocessed = false;
     }
@@ -338,7 +342,9 @@ struct LoraModel : public GGMLRunner {
             iter = lora_tensors.find(hada_1_mid_name);
             if (iter != lora_tensors.end()) {
                 hada_1_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
-                hada_1_up  = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up));
+                if (hada_1_up != nullptr) {
+                    hada_1_up = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up));
+                }
             }
 
             iter = lora_tensors.find(hada_2_down_name);
@@ -354,7 +360,9 @@ struct LoraModel : public GGMLRunner {
             iter = lora_tensors.find(hada_2_mid_name);
             if (iter != lora_tensors.end()) {
                 hada_2_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
-                hada_2_up  = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
+                if (hada_2_up != nullptr) {
+                    hada_2_up = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
+                }
             }
 
             if (hada_1_up == nullptr || hada_1_down == nullptr || hada_2_up == nullptr || hada_2_down == nullptr) {
@@ -546,7 +554,27 @@ struct LoraModel : public GGMLRunner {
                 }
             }
 
-            GGML_ASSERT(ggml_nelements(diff) == ggml_nelements(model_tensor));
+            if (ggml_nelements(diff) != ggml_nelements(model_tensor)) {
+                const std::string lora_tensor_prefix = "lora." + model_tensor_name + ".";
+                for (const auto& tensor_name : applied_lora_tensors) {
+                    if (starts_with(tensor_name, lora_tensor_prefix)) {
+                        skipped_incompatible_lora_tensors.insert(tensor_name);
+                    }
+                }
+                if (warned_incompatible_model_tensors.insert(model_tensor_name).second) {
+                    LOG_WARN("skip incompatible LoRA tensor |%s|: model shape = [%lld, %lld, %lld, %lld], LoRA shape = [%lld, %lld, %lld, %lld]",
+                             model_tensor_name.c_str(),
+                             static_cast<long long>(model_tensor->ne[0]),
+                             static_cast<long long>(model_tensor->ne[1]),
+                             static_cast<long long>(model_tensor->ne[2]),
+                             static_cast<long long>(model_tensor->ne[3]),
+                             static_cast<long long>(diff->ne[0]),
+                             static_cast<long long>(diff->ne[1]),
+                             static_cast<long long>(diff->ne[2]),
+                             static_cast<long long>(diff->ne[3]));
+                }
+                return nullptr;
+            }
             diff = ggml_reshape(ctx, diff, model_tensor);
         }
         return diff;
@@ -555,10 +583,15 @@ struct LoraModel : public GGMLRunner {
     ggml_tensor* get_out_diff(ggml_context* ctx,
                               ggml_backend_t backend,
                               ggml_tensor* x,
+                              ggml_tensor* model_weight,
                               WeightAdapter::ForwardParams forward_params,
                               const std::string& model_tensor_name) {
         ggml_tensor* out_diff = nullptr;
         int index             = 0;
+
+        std::vector<std::string> used_tensors;
+        bool is_conv2d = forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+
         while (true) {
             std::string key;
             if (index == 0) {
@@ -566,7 +599,6 @@ struct LoraModel : public GGMLRunner {
             } else {
                 key = model_tensor_name + "." + std::to_string(index);
             }
-            bool is_conv2d = forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
 
             std::string lokr_w1_name   = "lora." + key + ".lokr_w1";
             std::string lokr_w1_a_name = "lora." + key + ".lokr_w1_a";
@@ -634,7 +666,6 @@ struct LoraModel : public GGMLRunner {
                 if (iter != lora_tensors.end()) {
                     float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
                     scale_value = alpha / rank;
-                    applied_lora_tensors.insert(alpha_name);
                 }
 
                 if (rank == 1) {
@@ -649,19 +680,27 @@ struct LoraModel : public GGMLRunner {
                     out_diff = ggml_concat(ctx, out_diff, curr_out_diff, 0);
                 }
 
-                if (lokr_w1)
-                    applied_lora_tensors.insert(lokr_w1_name);
-                if (lokr_w1_a)
-                    applied_lora_tensors.insert(lokr_w1_a_name);
-                if (lokr_w1_b)
-                    applied_lora_tensors.insert(lokr_w1_b_name);
-                if (lokr_w2)
-                    applied_lora_tensors.insert(lokr_w2_name);
-                if (lokr_w2_a)
-                    applied_lora_tensors.insert(lokr_w2_a_name);
-                if (lokr_w2_b)
-                    applied_lora_tensors.insert(lokr_w2_b_name);
-                applied_lora_tensors.insert(alpha_name);
+                if (lokr_w1) {
+                    used_tensors.push_back(lokr_w1_name);
+                }
+                if (lokr_w1_a) {
+                    used_tensors.push_back(lokr_w1_a_name);
+                }
+                if (lokr_w1_b) {
+                    used_tensors.push_back(lokr_w1_b_name);
+                }
+                if (lokr_w2) {
+                    used_tensors.push_back(lokr_w2_name);
+                }
+                if (lokr_w2_a) {
+                    used_tensors.push_back(lokr_w2_a_name);
+                }
+                if (lokr_w2_b) {
+                    used_tensors.push_back(lokr_w2_b_name);
+                }
+                if (iter != lora_tensors.end()) {
+                    used_tensors.push_back(alpha_name);
+                }
 
                 index++;
                 continue;
@@ -707,27 +746,60 @@ struct LoraModel : public GGMLRunner {
                 break;
             }
 
-            applied_lora_tensors.insert(lora_up_name);
-            applied_lora_tensors.insert(lora_down_name);
+            if (!is_conv2d) {
+                const int64_t down_in  = lora_down->ne[0];
+                const int64_t down_out = lora_down->ne[1];
+                const int64_t up_in    = lora_up->ne[0];
 
-            if (lora_mid) {
-                applied_lora_tensors.insert(lora_mid_name);
+                bool compatible = down_in == model_weight->ne[0];
+                if (lora_mid != nullptr) {
+                    compatible = compatible &&
+                                 lora_mid->ne[0] == down_out &&
+                                 up_in == lora_mid->ne[1];
+                } else {
+                    compatible = compatible && up_in == down_out;
+                }
+
+                if (!compatible) {
+                    skipped_incompatible_lora_tensors.insert(lora_down_name);
+                    skipped_incompatible_lora_tensors.insert(lora_up_name);
+                    if (lora_mid != nullptr) {
+                        skipped_incompatible_lora_tensors.insert(lora_mid_name);
+                    }
+                    if (lora_tensors.find(scale_name) != lora_tensors.end()) {
+                        skipped_incompatible_lora_tensors.insert(scale_name);
+                    } else if (lora_tensors.find(alpha_name) != lora_tensors.end()) {
+                        skipped_incompatible_lora_tensors.insert(alpha_name);
+                    }
+                    if (warned_incompatible_model_tensors.insert(model_tensor_name).second) {
+                        LOG_WARN("skip incompatible LoRA tensor |%s|: model input dim = %lld, down shape = [%lld, %lld], up shape = [%lld, %lld]",
+                                 model_tensor_name.c_str(),
+                                 static_cast<long long>(model_weight->ne[0]),
+                                 static_cast<long long>(down_in),
+                                 static_cast<long long>(down_out),
+                                 static_cast<long long>(up_in),
+                                 static_cast<long long>(lora_up->ne[1]));
+                    }
+                    index++;
+                    continue;
+                }
             }
 
             float scale_value = 1.0f;
+            std::string scale_tensor_name;
 
             int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
             iter         = lora_tensors.find(scale_name);
             if (iter != lora_tensors.end()) {
-                scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
-                applied_lora_tensors.insert(scale_name);
+                scale_value       = ggml_ext_backend_tensor_get_f32(iter->second);
+                scale_tensor_name = scale_name;
             } else {
                 iter = lora_tensors.find(alpha_name);
                 if (iter != lora_tensors.end()) {
-                    float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
-                    scale_value = alpha / rank;
+                    float alpha       = ggml_ext_backend_tensor_get_f32(iter->second);
+                    scale_value       = alpha / rank;
+                    scale_tensor_name = alpha_name;
                     // LOG_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
-                    applied_lora_tensors.insert(alpha_name);
                 }
             }
             scale_value *= multiplier;
@@ -787,15 +859,45 @@ struct LoraModel : public GGMLRunner {
             }
 
             auto curr_out_diff = ggml_ext_scale(ctx, lx, scale_value, true);
-
             if (out_diff == nullptr) {
                 out_diff = curr_out_diff;
             } else {
-                out_diff = ggml_concat(ctx, out_diff, curr_out_diff, 0);
+                out_diff = ggml_concat(ctx, out_diff, curr_out_diff, is_conv2d ? 2 : 0);
+            }
+
+            used_tensors.push_back(lora_up_name);
+            used_tensors.push_back(lora_down_name);
+            if (lora_mid) {
+                used_tensors.push_back(lora_mid_name);
+            }
+            if (!scale_tensor_name.empty()) {
+                used_tensors.push_back(scale_tensor_name);
             }
 
             index++;
         }
+
+        if (out_diff == nullptr)
+            return nullptr;
+
+        int64_t expected_out_dim = is_conv2d ? model_weight->ne[3] : model_weight->ne[1];
+        int64_t actual_out_dim   = out_diff->ne[is_conv2d ? 2 : 0];
+
+        if (actual_out_dim != expected_out_dim) {
+            for (const auto& name : used_tensors) {
+                skipped_incompatible_lora_tensors.insert(name);
+            }
+            if (warned_incompatible_model_tensors.insert(model_tensor_name).second) {
+                LOG_WARN("skip incompatible LoRA tensors for |%s|: output dim %lld != model dim %lld",
+                         model_tensor_name.c_str(), actual_out_dim, expected_out_dim);
+            }
+            return nullptr;
+        }
+
+        for (const auto& name : used_tensors) {
+            applied_lora_tensors.insert(name);
+        }
+
         return out_diff;
     }
 
@@ -869,10 +971,13 @@ struct LoraModel : public GGMLRunner {
     void stat(bool at_runntime = false) {
         size_t total_lora_tensors_count   = 0;
         size_t applied_lora_tensors_count = 0;
+        size_t skipped_lora_tensors_count = 0;
 
         for (auto& kv : lora_tensors) {
             total_lora_tensors_count++;
-            if (applied_lora_tensors.find(kv.first) == applied_lora_tensors.end()) {
+            if (skipped_incompatible_lora_tensors.find(kv.first) != skipped_incompatible_lora_tensors.end()) {
+                skipped_lora_tensors_count++;
+            } else if (applied_lora_tensors.find(kv.first) == applied_lora_tensors.end()) {
                 if (!at_runntime) {
                     LOG_WARN("unused lora tensor |%s|", kv.first.c_str());
                     print_ggml_tensor(kv.second, true);
@@ -884,12 +989,17 @@ struct LoraModel : public GGMLRunner {
         /* Don't worry if this message shows up twice in the logs per LoRA,
          * this function is called once to calculate the required buffer size
          * and then again to actually generate a graph to be used */
-        if (!at_runntime && applied_lora_tensors_count != total_lora_tensors_count) {
+        size_t compatible_lora_tensors_count = total_lora_tensors_count - skipped_lora_tensors_count;
+        if (!at_runntime && applied_lora_tensors_count != compatible_lora_tensors_count) {
             LOG_WARN("Only (%lu / %lu) LoRA tensors have been applied, lora_file_path = %s",
-                     applied_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
+                     applied_lora_tensors_count, compatible_lora_tensors_count, file_path.c_str());
         } else {
             LOG_INFO("(%lu / %lu) LoRA tensors have been applied, lora_file_path = %s",
-                     applied_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
+                     applied_lora_tensors_count, compatible_lora_tensors_count, file_path.c_str());
+        }
+        if (skipped_lora_tensors_count > 0) {
+            LOG_WARN("(%lu / %lu) incompatible LoRA tensors have been skipped, lora_file_path = %s",
+                     skipped_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
         }
     }
 };
@@ -953,7 +1063,7 @@ public:
                                    forward_params.conv2d.scale);
         }
         for (auto& lora_model : lora_models) {
-            ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, forward_params, prefix + "weight");
+            ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight");
             if (out_diff == nullptr) {
                 continue;
             }

src/model/diffusion/krea2.hpp

@@ -180,9 +180,12 @@ namespace Krea2 {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             ggml_tensor* scale = params["scale"];
-            scale              = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1));
-            x                  = ggml_rms_norm(ctx->ggml_ctx, x, eps);
-            x                  = ggml_mul_inplace(ctx->ggml_ctx, x, scale);
+            if (ctx->weight_adapter) {
+                scale = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, scale, prefix + "scale.weight");
+            }
+            scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1));
+            x     = ggml_rms_norm(ctx->ggml_ctx, x, eps);
+            x     = ggml_mul_inplace(ctx->ggml_ctx, x, scale);
             return x;
         }
     };
@@ -295,10 +298,11 @@ namespace Krea2 {
     class KreaDoubleSharedModulation : public GGMLBlock {
     protected:
         int64_t dim;
+        std::string prefix;
 
         void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
             GGML_UNUSED(tensor_storage_map);
-            GGML_UNUSED(prefix);
+            this->prefix  = prefix;
             params["lin"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim * 6);
         }
 
@@ -307,7 +311,11 @@ namespace Krea2 {
             : dim(dim) {}
 
         std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
-            auto lin = ggml_repeat(ctx->ggml_ctx, params["lin"], vec);
+            auto lin = params["lin"];
+            if (ctx->weight_adapter) {
+                lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight");
+            }
+            lin      = ggml_repeat(ctx->ggml_ctx, lin, vec);
             auto out = ggml_add(ctx->ggml_ctx, vec, lin);
             return ggml_ext_chunk(ctx->ggml_ctx, out, 6, 0);
         }
@@ -316,10 +324,11 @@ namespace Krea2 {
     class KreaFinalModulation : public GGMLBlock {
     protected:
         int64_t dim;
+        std::string prefix;
 
         void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
             GGML_UNUSED(tensor_storage_map);
-            GGML_UNUSED(prefix);
+            this->prefix  = prefix;
             params["lin"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, 2);
         }
 
@@ -328,7 +337,11 @@ namespace Krea2 {
             : dim(dim) {}
 
         std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
-            auto out = ggml_add(ctx->ggml_ctx, params["lin"], vec);
+            auto lin = params["lin"];
+            if (ctx->weight_adapter) {
+                lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight");
+            }
+            auto out = ggml_add(ctx->ggml_ctx, lin, vec);
             return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1);
         }
     };

src/model/te/llm.hpp

@@ -1657,6 +1657,10 @@ namespace LLM {
             model.get_param_tensors(tensors, prefix);
         }
 
+        void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
+            model.get_param_tensor_ops(tensor_ops);
+        }
+
         ggml_tensor* forward(GGMLRunnerContext* ctx,
                              ggml_tensor* input_ids,
                              ggml_tensor* input_pos,

src/model_io/pickle_io.cpp

@@ -2,6 +2,7 @@
 
 #include <cstdlib>
 #include <cstring>
+#include <limits>
 #include <string>
 #include <unordered_map>
 #include <utility>
@@ -512,8 +513,51 @@ static bool parse_storage_type(const std::string& global_name, PickleStorageInfo
     return false;
 }
 
-static bool tensor_is_contiguous(const PickleTensorInfo& tensor) {
-    if (tensor.tensor_storage.nelements() == 0) {
+static bool checked_pickle_byte_count(int64_t element_count,
+                                      uint64_t element_nbytes,
+                                      uint64_t* byte_count) {
+    if (element_count < 0 || element_nbytes == 0) {
+        return false;
+    }
+
+    uint64_t count = static_cast<uint64_t>(element_count);
+    if (count > std::numeric_limits<uint64_t>::max() / element_nbytes) {
+        return false;
+    }
+
+    *byte_count = count * element_nbytes;
+    return true;
+}
+
+static bool tensor_layout_is_valid(const PickleTensorInfo& tensor, uint64_t raw_element_nbytes) {
+    if (raw_element_nbytes == 0) {
+        return false;
+    }
+
+    bool has_zero_dimension = false;
+    uint64_t element_count  = 1;
+    for (int i = 0; i < tensor.tensor_storage.n_dims; ++i) {
+        int64_t dimension = tensor.tensor_storage.ne[i];
+        if (dimension < 0) {
+            return false;
+        }
+        if (dimension == 0) {
+            has_zero_dimension = true;
+            continue;
+        }
+
+        uint64_t size = static_cast<uint64_t>(dimension);
+        if (element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / size) {
+            return false;
+        }
+        element_count *= size;
+    }
+
+    if (!has_zero_dimension &&
+        element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / raw_element_nbytes) {
+        return false;
+    }
+    if (has_zero_dimension) {
         return true;
     }
     if (tensor.stride_n_dims != tensor.tensor_storage.n_dims) {
@@ -932,7 +976,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
                 if (storage.key.empty() || !parse_storage_type(pid.items[1].str_value, &storage)) {
                     return false;
                 }
-                storage.nbytes              = (uint64_t)pid.items[4].int_value * storage.raw_element_nbytes;
+                if (!checked_pickle_byte_count(pid.items[4].int_value,
+                                               storage.raw_element_nbytes,
+                                               &storage.nbytes)) {
+                    set_error(error, "invalid storage size in torch pickle");
+                    return false;
+                }
                 storage_nbytes[storage.key] = storage.nbytes;
                 stack.push_back(make_storage_value(storage));
             } break;
@@ -963,7 +1012,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
                     tensor.tensor_storage.is_f64      = args.items[0].storage.is_f64;
                     tensor.tensor_storage.is_i64      = args.items[0].storage.is_i64;
                     tensor.tensor_storage.storage_key = args.items[0].storage.key;
-                    tensor.tensor_storage.offset      = (uint64_t)args.items[1].int_value * args.items[0].storage.raw_element_nbytes;
+                    if (!checked_pickle_byte_count(args.items[1].int_value,
+                                                   args.items[0].storage.raw_element_nbytes,
+                                                   &tensor.tensor_storage.offset)) {
+                        set_error(error, "invalid tensor storage offset in torch pickle");
+                        return false;
+                    }
 
                     for (const auto& item : args.items[2].items) {
                         if (item.kind != PickleValue::INT || tensor.tensor_storage.n_dims >= SD_MAX_DIMS) {
@@ -979,7 +1033,8 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
                         tensor.stride[tensor.stride_n_dims++] = item.int_value;
                     }
 
-                    if (!tensor_is_contiguous(tensor)) {
+                    if (!tensor_layout_is_valid(tensor, args.items[0].storage.raw_element_nbytes)) {
+                        set_error(error, "invalid tensor shape or stride in torch pickle");
                         return false;
                     }
                     stack.push_back(make_tensor_value(tensor));

src/model_io/torch_legacy_io.cpp

@@ -139,11 +139,16 @@ bool read_torch_legacy_file(const std::string& file_path,
             if (it == legacy_storage_map.end()) {
                 return false;
             }
-            if (current_offset + LEGACY_STORAGE_HEADER_SIZE + it->second > file_size) {
+            if (current_offset > file_size ||
+                LEGACY_STORAGE_HEADER_SIZE > file_size - current_offset) {
                 return false;
             }
-            storage_offsets[storage_key] = current_offset + LEGACY_STORAGE_HEADER_SIZE;
-            current_offset += LEGACY_STORAGE_HEADER_SIZE + it->second;
+            uint64_t storage_offset = current_offset + LEGACY_STORAGE_HEADER_SIZE;
+            if (it->second > file_size - storage_offset) {
+                return false;
+            }
+            storage_offsets[storage_key] = storage_offset;
+            current_offset               = storage_offset + it->second;
         }
 
         for (auto& tensor_storage : tensor_storages) {
@@ -159,8 +164,10 @@ bool read_torch_legacy_file(const std::string& file_path,
 
             uint64_t base_offset    = it_offset->second;
             uint64_t storage_nbytes = it_size->second;
-            uint64_t tensor_nbytes  = tensor_storage.nbytes_to_read();
-            if (tensor_storage.offset + tensor_nbytes > storage_nbytes) {
+            int64_t tensor_nbytes   = tensor_storage.nbytes_to_read();
+            if (tensor_nbytes < 0 ||
+                tensor_storage.offset > storage_nbytes ||
+                static_cast<uint64_t>(tensor_nbytes) > storage_nbytes - tensor_storage.offset) {
                 return false;
             }
 

src/model_io/torch_zip_io.cpp

@@ -76,8 +76,10 @@ static bool parse_zip_data_pkl(const uint8_t* buffer,
             return false;
         }
 
-        uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
-        if (tensor_storage.offset + tensor_nbytes > entry_size) {
+        int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
+        if (tensor_nbytes < 0 ||
+            tensor_storage.offset > entry_size ||
+            static_cast<uint64_t>(tensor_nbytes) > entry_size - tensor_storage.offset) {
             set_error(error, "tensor '" + tensor_storage.name + "' exceeds storage entry '" + entry_name + "'");
             return false;
         }

src/model_loader.cpp

@@ -1053,7 +1053,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
         std::atomic<size_t> tensor_idx(0);
         std::atomic<bool> failed(false);
         std::vector<std::thread> workers;
-        std::mutex rpc_backend_mutex;
+        std::mutex backend_tensor_set_mutex;
 
         for (int i = 0; i < n_threads; ++i) {
             workers.emplace_back([&, file_path, is_zip]() {
@@ -1077,6 +1077,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
 
                 std::vector<uint8_t> read_buffer;
                 std::vector<uint8_t> convert_buffer;
+                std::vector<uint8_t> zip_entry_buffer;
 
                 while (true) {
                     int64_t t0, t1;
@@ -1115,34 +1116,60 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
 
                     size_t nbytes_to_read = tensor_storage.nbytes_to_read();
 
-                    auto read_data = [&](char* buf, size_t n) {
+                    auto read_data = [&](char* buf, size_t n) -> bool {
                         if (zip != nullptr) {
-                            zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
+                            if (zip_entry_openbyindex(zip, tensor_storage.index_in_zip) != 0) {
+                                LOG_ERROR("failed to open zip entry for tensor '%s'", tensor_storage.name.c_str());
+                                return false;
+                            }
                             size_t entry_size = zip_entry_size(zip);
+                            if (tensor_storage.offset > entry_size) {
+                                LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
+                                zip_entry_close(zip);
+                                return false;
+                            }
+                            size_t tensor_offset = static_cast<size_t>(tensor_storage.offset);
+                            if (n > entry_size - tensor_offset) {
+                                LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
+                                zip_entry_close(zip);
+                                return false;
+                            }
+
                             if (entry_size != n) {
                                 int64_t t_memcpy_start;
-                                read_buffer.resize(entry_size);
-                                zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
+                                zip_entry_buffer.resize(entry_size);
+                                auto bytes_read = zip_entry_noallocread(zip, (void*)zip_entry_buffer.data(), entry_size);
+                                if (bytes_read < 0 || static_cast<size_t>(bytes_read) != entry_size) {
+                                    LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
+                                    zip_entry_close(zip);
+                                    return false;
+                                }
                                 t_memcpy_start = ggml_time_ms();
-                                memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
+                                memcpy((void*)buf, (void*)(zip_entry_buffer.data() + tensor_offset), n);
                                 memcpy_time_ms.fetch_add(ggml_time_ms() - t_memcpy_start);
                             } else {
-                                zip_entry_noallocread(zip, (void*)buf, n);
+                                auto bytes_read = zip_entry_noallocread(zip, (void*)buf, n);
+                                if (bytes_read < 0 || static_cast<size_t>(bytes_read) != n) {
+                                    LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
+                                    zip_entry_close(zip);
+                                    return false;
+                                }
                             }
                             zip_entry_close(zip);
                         } else if (mmapped) {
                             if (!mmapped->copy_data(buf, n, tensor_storage.offset)) {
                                 LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
-                                failed = true;
+                                return false;
                             }
                         } else {
                             file.seekg(tensor_storage.offset);
                             file.read(buf, n);
                             if (!file) {
                                 LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
-                                failed = true;
+                                return false;
                             }
                         }
+                        return true;
                     };
 
                     char* read_buf    = nullptr;
@@ -1176,7 +1203,10 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                     }
 
                     t0 = ggml_time_ms();
-                    read_data(read_buf, nbytes_to_read);
+                    if (!read_data(read_buf, nbytes_to_read)) {
+                        failed = true;
+                        break;
+                    }
                     t1 = ggml_time_ms();
                     read_time_ms.fetch_add(t1 - t0);
 
@@ -1214,17 +1244,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                     if (dst_tensor->buffer != nullptr && !ggml_backend_buffer_is_host(dst_tensor->buffer)) {
                         t0 = ggml_time_ms();
 
-                        // RPC backends require serialized access to prevent concurrency issues
-                        const char* buffer_type_name = ggml_backend_buft_name(ggml_backend_buffer_get_type(dst_tensor->buffer));
-                        bool is_rpc_buffer           = buffer_type_name != nullptr &&
-                                             std::string(buffer_type_name).find("RPC") != std::string::npos;
-
-                        if (is_rpc_buffer) {
-                            std::lock_guard<std::mutex> lock(rpc_backend_mutex);
-                            ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
-                        } else {
-                            ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
-                        }
+                        std::lock_guard<std::mutex> lock(backend_tensor_set_mutex);
+                        ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
 
                         t1 = ggml_time_ms();
                         copy_to_backend_time_ms.fetch_add(t1 - t0);

src/model_manager.cpp

@@ -53,6 +53,48 @@ static bool backend_supports_host_buffer(ggml_backend_t backend) {
     return props.caps.buffer_from_host_ptr;
 }
 
+static bool device_supports_param_op(ggml_backend_dev_t device,
+                                     ggml_tensor* weight,
+                                     enum ggml_op op,
+                                     ggml_backend_buffer_type_t buft) {
+    if (op == GGML_OP_NONE) {
+        return true;
+    }
+    if (device == nullptr || weight == nullptr || buft == nullptr || weight->buffer != nullptr) {
+        return false;
+    }
+
+    ggml_init_params params;
+    params.mem_size   = ggml_tensor_overhead() * 2;
+    params.mem_buffer = nullptr;
+    params.no_alloc   = true;
+    ggml_context* ctx = ggml_init(params);
+    if (ctx == nullptr) {
+        return false;
+    }
+
+    ggml_tensor* op_tensor = nullptr;
+    if (op == GGML_OP_GET_ROWS) {
+        ggml_tensor* indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
+        op_tensor            = ggml_get_rows(ctx, weight, indices);
+    }
+    if (op_tensor == nullptr) {
+        ggml_free(ctx);
+        return false;
+    }
+
+    weight->buffer = ggml_backend_buft_alloc_buffer(buft, 0);
+    if (weight->buffer == nullptr) {
+        ggml_free(ctx);
+        return false;
+    }
+    bool supported = ggml_backend_dev_supports_op(device, op_tensor);
+    ggml_backend_buffer_free(weight->buffer);
+    weight->buffer = nullptr;
+    ggml_free(ctx);
+    return supported;
+}
+
 ModelManager::~ModelManager() {
     release_all();
 }
@@ -135,7 +177,8 @@ bool ModelManager::register_param_tensors(const std::string& desc,
                                           ggml_backend_t params_backend,
                                           size_t* registered_tensor_size,
                                           bool allow_split_buffer,
-                                          bool params_follow_compute_backend) {
+                                          bool params_follow_compute_backend,
+                                          const std::map<ggml_tensor*, enum ggml_op>* tensor_ops) {
     if (desc.empty()) {
         LOG_ERROR("model manager tensor desc is empty");
         return false;
@@ -168,6 +211,12 @@ bool ModelManager::register_param_tensors(const std::string& desc,
         state->params_backend                = params_backend;
         state->allow_split_buffer            = allow_split_buffer;
         state->params_follow_compute_backend = params_follow_compute_backend;
+        if (tensor_ops != nullptr) {
+            auto op_it = tensor_ops->find(tensor);
+            if (op_it != tensor_ops->end()) {
+                state->usage_op = op_it->second;
+            }
+        }
         new_states.push_back(std::move(state));
     }
 
@@ -844,6 +893,22 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat
     if (params_buft == nullptr) {
         params_buft = ggml_backend_get_default_buffer_type(state.params_backend);
     }
+    if (state.usage_op != GGML_OP_NONE &&
+        state.compute_backend != nullptr) {
+        ggml_backend_dev_t compute_dev = ggml_backend_get_device(state.compute_backend);
+        if (device_supports_param_op(compute_dev, state.tensor, state.usage_op, params_buft)) {
+            return params_buft;
+        }
+
+        ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+        params_buft                = cpu_dev != nullptr ? ggml_backend_dev_buffer_type(cpu_dev) : nullptr;
+        if (!device_supports_param_op(cpu_dev, state.tensor, state.usage_op, params_buft)) {
+            LOG_ERROR("model manager has no compatible buffer for tensor '%s' used by %s",
+                      state.name.c_str(),
+                      ggml_op_name(state.usage_op));
+            return nullptr;
+        }
+    }
     return params_buft;
 }
 

src/model_manager.h

@@ -39,6 +39,7 @@ private:
         bool allow_split_buffer            = false;
         bool params_follow_compute_backend = false;
         bool metadata_validated            = false;
+        enum ggml_op usage_op              = GGML_OP_NONE;
 
         int active_prepare_count = 0;
 
@@ -130,9 +131,10 @@ public:
                                 ResidencyMode residency_mode,
                                 ggml_backend_t compute_backend,
                                 ggml_backend_t params_backend,
-                                size_t* registered_tensor_size     = nullptr,
-                                bool allow_split_buffer            = false,
-                                bool params_follow_compute_backend = false);
+                                size_t* registered_tensor_size                         = nullptr,
+                                bool allow_split_buffer                                = false,
+                                bool params_follow_compute_backend                     = false,
+                                const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr);
 
     bool unregister_param_tensors(const std::string& desc,
                                   size_t* registered_tensor_size = nullptr);

src/name_conversion.cpp

@@ -1384,6 +1384,8 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
             {".lora_B.weight", ".weight.lora_up"},
             {".lora_A.default.weight", ".weight.lora_down"},
             {".lora_B.default.weight", ".weight.lora_up"},
+            {".lora_A", ".weight.lora_down"},
+            {".lora_B", ".weight.lora_up"},
             {".lora_linear", ".weight.alpha"},
             {".alpha", ".weight.alpha"},
             {".scale", ".weight.scale"},
@@ -1449,12 +1451,21 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
         {"te2.", "cond_stage_model.1.transformer."},
         {"te1.", "cond_stage_model.transformer."},
         {"te3.", "text_encoders.t5xxl.transformer."},
+        {"clip_vision.", "cond_stage_model.transformer."},
     };
 
     if (sd_version_is_flux(version)) {
         prefix_map["te1."] = "text_encoders.clip_l.transformer.";
     }
 
+    if (sd_version_is_unet(version)) {
+        prefix_map["clip_l."] = "cond_stage_model.transformer.";
+        prefix_map["clip_g."] = "cond_stage_model.1.transformer.";
+    } else {
+        prefix_map["clip_l."] = "text_encoders.clip_l.transformer.";
+        prefix_map["clip_g."] = "text_encoders.clip_g.transformer.";
+    }
+
     replace_with_prefix_map(name, prefix_map);
 
     if (sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version)) {

src/runtime/denoiser.hpp

@@ -306,8 +306,33 @@ struct KarrasScheduler : SigmaScheduler {
 };
 
 struct BetaScheduler : SigmaScheduler {
-    static constexpr double alpha = 0.6;
-    static constexpr double beta  = 0.6;
+    double alpha = 0.6;
+    double beta  = 0.6;
+
+    explicit BetaScheduler(const char* extra_sample_args = nullptr) {
+        parse_extra_sample_args(extra_sample_args);
+        LOG_DEBUG("Beta scheduler: alpha=%.4f, beta=%.4f", alpha, beta);
+    }
+
+    void parse_extra_sample_args(const char* extra_sample_args) {
+        for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "beta scheduler arg")) {
+            if (key == "alpha") {
+                float parsed;
+                if (!parse_strict_float(value, parsed) || parsed <= 0.0) {
+                    LOG_WARN("ignoring invalid beta scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                } else {
+                    alpha = static_cast<double>(parsed);
+                }
+            } else if (key == "beta") {
+                float parsed;
+                if (!parse_strict_float(value, parsed) || parsed <= 0.0) {
+                    LOG_WARN("ignoring invalid beta scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                } else {
+                    beta = static_cast<double>(parsed);
+                }
+            }
+        }
+    }
 
     static double log_beta(double a, double b) {
         return std::lgamma(a) + std::lgamma(b) - std::lgamma(a + b);
@@ -1032,7 +1057,7 @@ struct Denoiser {
                 break;
             case BETA_SCHEDULER:
                 LOG_INFO("get_sigmas with Beta scheduler");
-                scheduler = std::make_shared<BetaScheduler>();
+                scheduler = std::make_shared<BetaScheduler>(extra_sample_args);
                 break;
             case EXPONENTIAL_SCHEDULER:
                 LOG_INFO("get_sigmas exponential scheduler");
@@ -2553,6 +2578,88 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
     return x;
 }
 
+static sd::Tensor<float> sample_lms(denoise_cb_t model,
+                                    sd::Tensor<float> x,
+                                    const std::vector<float>& sigmas,
+                                    const SamplerExtraArgs& extra_sample_args) {
+    // Linear Multi-Step from https://github.com/crowsonkb/k-diffusion
+
+    int divisions = 1000;
+    for (const auto& [key, value] : extra_sample_args) {
+        int parsed = 0;
+        if (key == "lms_divisions") {
+            if (!parse_strict_int(value, parsed)) {
+                LOG_WARN("ignoring invalid lms extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                continue;
+            }
+            divisions = parsed;  // std::max(1, parsed);
+            // values above 35M produce noise, can be fixed by double precision
+            // values < 1 always produce noise
+        }
+    }
+    LOG_DEBUG("linear multi-step sampler: integrating using %i division%s", divisions, (divisions == 1) ? "" : "s");
+
+    auto linear_multistep_coeff = [=](const int order, const int m, const int j) -> float {
+        if (!divisions)
+            return sigmas[m + 1] - sigmas[m];  // delta / 0 * 0
+#define LMS_PRECISION float                    // double
+        const LMS_PRECISION a = sigmas[m], dx = (sigmas[m + 1] - a) / divisions, s = sigmas[m - j];
+        const LMS_PRECISION b0 = a + 0.5f * dx;  // using Riemann middle integral
+        LMS_PRECISION sum      = 0.0f;
+        for (int h = 0; h < divisions; h++) {
+            const LMS_PRECISION b = h * dx + b0;
+            LMS_PRECISION prod    = 1.0f;
+            for (int k = 0; k < j; k++) {
+                const LMS_PRECISION t = sigmas[m - k];
+                prod *= (b - t) / (s - t);
+            }
+            for (int k = j + 1; k < order; k++) {
+                const LMS_PRECISION t = sigmas[m - k];
+                prod *= (b - t) / (s - t);
+            }
+            sum += prod;
+        }
+        return sum * dx;
+    };
+
+    const int max_order = 4;
+    float lms_coeff[max_order];
+    std::vector<sd::Tensor<float>> hist = {};
+
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        const float sigma = sigmas[i];
+
+        auto denoised_opt = model(x, sigma, i + 1);
+        if (denoised_opt.pred.empty()) {
+            return {};
+        }
+        sd::Tensor<float> denoised = std::move(denoised_opt.pred);
+
+        const int order = std::min(max_order, i + 1);
+        for (int c = 0; c < order; c++)  // computing coefficients
+            lms_coeff[c] = linear_multistep_coeff(order, i, c);
+
+        sd::Tensor<float> d_cur = (x - denoised) / sigma;
+        switch (order) {
+            case 4:  // derivative + 3 history points
+                x += hist[hist.size() - 2] * lms_coeff[3];
+            case 3:
+                x += hist[hist.size() - 1] * lms_coeff[2];
+            case 2:
+                x += hist.back() * lms_coeff[1];
+            case 1:
+                x += d_cur * lms_coeff[0];
+        }
+
+        if (hist.size() == static_cast<size_t>(max_order - 1)) {
+            hist.erase(hist.begin());
+        }
+        hist.push_back(std::move(d_cur));
+    }
+    return x;
+}
+
 static sd::Tensor<float> sample_euler_cfg_pp(denoise_cb_t model,
                                              sd::Tensor<float> x,
                                              const std::vector<float>& sigmas) {
@@ -2714,6 +2821,8 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
             return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
         case TCD_SAMPLE_METHOD:
             return sample_tcd(model, std::move(x), sigmas, rng, eta);
+        case LMS_SAMPLE_METHOD:
+            return sample_lms(model, std::move(x), sigmas, extra_args);
         case EULER_CFG_PP_SAMPLE_METHOD:
             return sample_euler_cfg_pp(model, std::move(x), sigmas);
         case EULER_A_CFG_PP_SAMPLE_METHOD:

src/stable-diffusion.cpp

@@ -141,6 +141,9 @@ const char* sampling_methods_str[] = {
     "Euler CFG++",
     "Euler A CFG++",
     "Euler GE",
+    "DPM++ (2M) SDE",
+    "DPM++ (2M) SDE BT",
+    "LMS",
 };
 
 /*================================================== Helper Functions ================================================*/
@@ -224,6 +227,7 @@ public:
     std::shared_ptr<ControlNet> control_net;
     std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
     sd::Tensor<float> ip_adapter_tokens;
+    sd::Tensor<float> ip_adapter_uncond_tokens;
     float ip_adapter_strength = 1.0f;
     std::vector<std::shared_ptr<GenerationExtension>> generation_extensions;
     std::vector<std::shared_ptr<LoraModel>> runtime_lora_models;
@@ -317,7 +321,11 @@ public:
             return true;
         }
         std::map<std::string, ggml_tensor*> group_tensors;
+        std::map<ggml_tensor*, enum ggml_op> tensor_ops;
         model->get_param_tensors(group_tensors);
+        if constexpr (std::is_base_of_v<Conditioner, T>) {
+            model->get_param_tensor_ops(tensor_ops);
+        }
         if (model_manager == nullptr) {
             return true;
         }
@@ -334,6 +342,7 @@ public:
                                                                 module,
                                                                 module_backends,
                                                                 std::move(group_tensors),
+                                                                tensor_ops,
                                                                 residency_mode,
                                                                 params_mem_size);
                     }
@@ -342,6 +351,7 @@ public:
                                                               module,
                                                               module_backends,
                                                               std::move(group_tensors),
+                                                              tensor_ops,
                                                               residency_mode,
                                                               params_mem_size);
                 }
@@ -355,7 +365,10 @@ public:
                                                      residency_mode,
                                                      backend_for(module),
                                                      params_backend_for(module),
-                                                     params_mem_size);
+                                                     params_mem_size,
+                                                     false,
+                                                     false,
+                                                     &tensor_ops);
     }
 
     template <typename T>
@@ -364,6 +377,7 @@ public:
                                           SDBackendModule module,
                                           const std::vector<ggml_backend_t>& module_backends,
                                           std::map<std::string, ggml_tensor*> group_tensors,
+                                          const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
                                           ModelManager::ResidencyMode residency_mode,
                                           size_t* params_mem_size) {
         ggml_backend_t main_backend = module_backends[0];
@@ -375,6 +389,7 @@ public:
                                                       module,
                                                       module_backends,
                                                       std::move(group_tensors),
+                                                      tensor_ops,
                                                       residency_mode,
                                                       params_mem_size);
         };
@@ -449,7 +464,9 @@ public:
                                                    main_backend,
                                                    params_backend_for(module),
                                                    params_mem_size,
-                                                   /*allow_split_buffer=*/true)) {
+                                                   /*allow_split_buffer=*/true,
+                                                   false,
+                                                   &tensor_ops)) {
             return false;
         }
         return model_manager->register_param_tensors(desc,
@@ -457,7 +474,10 @@ public:
                                                      residency_mode,
                                                      main_backend,
                                                      params_backend_for(module),
-                                                     params_mem_size);
+                                                     params_mem_size,
+                                                     false,
+                                                     false,
+                                                     &tensor_ops);
     }
 
     // Register graph-cut layer-split tensors on the primary backend first.
@@ -469,6 +489,7 @@ public:
                                             SDBackendModule module,
                                             const std::vector<ggml_backend_t>& module_backends,
                                             std::map<std::string, ggml_tensor*> group_tensors,
+                                            const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
                                             ModelManager::ResidencyMode residency_mode,
                                             size_t* params_mem_size) {
         bool has_cpu_device = false;
@@ -490,7 +511,10 @@ public:
                                                          residency_mode,
                                                          module_backends[0],
                                                          params_backend_for(module),
-                                                         params_mem_size);
+                                                         params_mem_size,
+                                                         false,
+                                                         false,
+                                                         &tensor_ops);
         }
 
         model->set_runtime_backends(module_backends);
@@ -515,7 +539,8 @@ public:
                                                      initial_params_backend,
                                                      params_mem_size,
                                                      false,
-                                                     params_follow_runtime);
+                                                     params_follow_runtime,
+                                                     &tensor_ops);
     }
 
     bool unload_control_net() {
@@ -738,28 +763,23 @@ public:
             }
         }
 
-        bool is_unet = sd_version_is_unet(model_loader.get_sd_version());
-
         if (strlen(SAFE_STR(sd_ctx_params->clip_l_path)) > 0) {
             LOG_INFO("loading clip_l from '%s'", sd_ctx_params->clip_l_path);
-            std::string prefix = is_unet ? "cond_stage_model.transformer." : "text_encoders.clip_l.transformer.";
-            if (!model_loader.init_from_file(sd_ctx_params->clip_l_path, prefix)) {
+            if (!model_loader.init_from_file(sd_ctx_params->clip_l_path, "clip_l.")) {
                 LOG_WARN("loading clip_l from '%s' failed", sd_ctx_params->clip_l_path);
             }
         }
 
         if (strlen(SAFE_STR(sd_ctx_params->clip_g_path)) > 0) {
             LOG_INFO("loading clip_g from '%s'", sd_ctx_params->clip_g_path);
-            std::string prefix = is_unet ? "cond_stage_model.1.transformer." : "text_encoders.clip_g.transformer.";
-            if (!model_loader.init_from_file(sd_ctx_params->clip_g_path, prefix)) {
+            if (!model_loader.init_from_file(sd_ctx_params->clip_g_path, "clip_g.")) {
                 LOG_WARN("loading clip_g from '%s' failed", sd_ctx_params->clip_g_path);
             }
         }
 
         if (strlen(SAFE_STR(sd_ctx_params->clip_vision_path)) > 0) {
             LOG_INFO("loading clip_vision from '%s'", sd_ctx_params->clip_vision_path);
-            std::string prefix = "cond_stage_model.transformer.";
-            if (!model_loader.init_from_file(sd_ctx_params->clip_vision_path, prefix)) {
+            if (!model_loader.init_from_file(sd_ctx_params->clip_vision_path, "clip_vision.")) {
                 LOG_WARN("loading clip_vision from '%s' failed", sd_ctx_params->clip_vision_path);
             }
         }
@@ -2100,21 +2120,33 @@ public:
     }
 
     void compute_ip_adapter_tokens(const sd_image_t& image, float strength) {
-        ip_adapter_tokens   = {};
-        ip_adapter_strength = strength;
+        ip_adapter_tokens        = {};
+        ip_adapter_uncond_tokens = {};
+        ip_adapter_strength      = strength;
         if (ip_adapter == nullptr || clip_vision == nullptr || image.data == nullptr) {
             return;
         }
         auto image_tensor = sd_image_to_tensor(image);
-        auto embed        = get_clip_vision_output(image_tensor, true, -1);
+        auto embed        = ip_adapter->is_plus
+                                ? get_clip_vision_output(image_tensor, false, 2)
+                                : get_clip_vision_output(image_tensor, true, -1);
         if (embed.empty()) {
             return;
         }
         ip_adapter_tokens = ip_adapter->compute(n_threads, embed);
-        if (!ip_adapter_tokens.empty()) {
-            LOG_INFO("IP-Adapter: %lld image tokens, strength %.2f",
-                     (long long)ip_adapter_tokens.shape()[1], strength);
+        if (ip_adapter_tokens.empty()) {
+            LOG_ERROR("IP-Adapter conditional image projection failed");
+            return;
         }
+        auto uncond_embed        = sd::Tensor<float>::zeros_like(embed);
+        ip_adapter_uncond_tokens = ip_adapter->compute(n_threads, uncond_embed);
+        if (ip_adapter_uncond_tokens.empty()) {
+            LOG_ERROR("IP-Adapter unconditional image projection failed");
+            ip_adapter_tokens = {};
+            return;
+        }
+        LOG_INFO("IP-Adapter: %lld image tokens, strength %.2f",
+                 (long long)ip_adapter_tokens.shape()[1], strength);
     }
 
     std::vector<float> process_timesteps(const std::vector<float>& timesteps,
@@ -2606,7 +2638,7 @@ public:
                                      const sd::Tensor<float>* c_concat_override                 = nullptr,
                                      const std::vector<int>* local_skip_layers                  = nullptr,
                                      const std::vector<sd::Tensor<float>>* ref_latents_override = nullptr,
-                                     bool apply_ip                                              = true) -> sd::Tensor<float> {
+                                     bool use_uncond_ip                                         = false) -> sd::Tensor<float> {
                 diffusion_params.context     = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn;
                 diffusion_params.c_concat    = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat);
                 diffusion_params.y           = condition.c_vector.empty() ? nullptr : &condition.c_vector;
@@ -2618,8 +2650,9 @@ public:
                         nvf = static_cast<int>(noised_input.shape()[3]);
                     }
                     UNetDiffusionExtra unet_extra{nvf, &controls, control_strength};
-                    if (apply_ip && !ip_adapter_tokens.empty()) {
-                        unet_extra.ip_context = &ip_adapter_tokens;
+                    const auto& ip_tokens = use_uncond_ip ? ip_adapter_uncond_tokens : ip_adapter_tokens;
+                    if (!ip_tokens.empty()) {
+                        unet_extra.ip_context = &ip_tokens;
                         unet_extra.ip_scale   = ip_adapter_strength;
                     }
                     diffusion_params.extra = unet_extra;
@@ -2715,7 +2748,7 @@ public:
                                            uncond.c_concat.empty() ? nullptr : &uncond.c_concat,
                                            uncond_skip_layers,
                                            nullptr,
-                                           false);
+                                           true);
                 if (uncond_out.empty()) {
                     return {};
                 }
@@ -2725,7 +2758,7 @@ public:
                                                img_uncond.c_concat.empty() ? nullptr : &img_uncond.c_concat,
                                                nullptr,
                                                uncond_without_ref_latents ? &empty_ref_latents : nullptr,
-                                               false);
+                                               true);
                 if (img_uncond_out.empty()) {
                     return {};
                 }
@@ -3176,6 +3209,7 @@ const char* sample_method_to_str[] = {
     "euler_ge",
     "dpm++2m_sde",
     "dpm++2m_sde_bt",
+    "lms",
 };
 
 const char* sd_sample_method_name(enum sample_method_t sample_method) {
@@ -3573,21 +3607,22 @@ char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) {
 void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
     *sd_img_gen_params = {};
     sd_sample_params_init(&sd_img_gen_params->sample_params);
-    sd_img_gen_params->clip_skip         = -1;
-    sd_img_gen_params->ref_images_count  = 0;
-    sd_img_gen_params->ref_image_args    = "";
-    sd_img_gen_params->width             = 512;
-    sd_img_gen_params->height            = 512;
-    sd_img_gen_params->strength          = 0.75f;
-    sd_img_gen_params->seed              = -1;
-    sd_img_gen_params->batch_count       = 1;
-    sd_img_gen_params->control_strength  = 0.9f;
-    sd_img_gen_params->qwen_image_layers = 3;
-    sd_img_gen_params->circular_x        = false;
-    sd_img_gen_params->circular_y        = false;
-    sd_img_gen_params->pm_params         = {nullptr, 0, nullptr, 20.f};
-    sd_img_gen_params->pulid_params      = {nullptr, 1.0f};
-    sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
+    sd_img_gen_params->clip_skip           = -1;
+    sd_img_gen_params->ref_images_count    = 0;
+    sd_img_gen_params->ref_image_args      = "";
+    sd_img_gen_params->width               = 512;
+    sd_img_gen_params->height              = 512;
+    sd_img_gen_params->strength            = 0.75f;
+    sd_img_gen_params->seed                = -1;
+    sd_img_gen_params->batch_count         = 1;
+    sd_img_gen_params->control_strength    = 0.9f;
+    sd_img_gen_params->ip_adapter_strength = 1.0f;
+    sd_img_gen_params->qwen_image_layers   = 3;
+    sd_img_gen_params->circular_x          = false;
+    sd_img_gen_params->circular_y          = false;
+    sd_img_gen_params->pm_params           = {nullptr, 0, nullptr, 20.f};
+    sd_img_gen_params->pulid_params        = {nullptr, 1.0f};
+    sd_img_gen_params->vae_tiling_params   = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
     sd_cache_params_init(&sd_img_gen_params->cache);
     sd_hires_params_init(&sd_img_gen_params->hires);
 }
@@ -5776,7 +5811,7 @@ static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd
         auto encode_condition_frame = [&](const sd::Tensor<float>& image,
                                           int64_t latent_frame,
                                           const char* name) -> bool {
-            auto encoded = sd_ctx->sd->encode_first_stage(image);
+            auto encoded = sd_ctx->sd->encode_first_stage(image.unsqueeze(2));
             if (encoded.empty()) {
                 LOG_ERROR("failed to encode Hunyuan Video %s conditioning frame", name);
                 return false;

src/tokenizers/bpe_tokenizer.cpp

@@ -205,7 +205,7 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
         ss << "\"" << token << "\", ";
     }
     ss << "]";
-    LOG_DEBUG("split prompt \"%s\" to tokens %s", text.c_str(), ss.str().c_str());
+    LOG_DEBUG("split prompt \"%s\" to %zu tokens %s", text.c_str(), bpe_tokens.size(), ss.str().c_str());
     return bpe_tokens;
 }