feat: add configurable reference image processing for edit models (#1780)

README.md

@@ -57,7 +57,7 @@ API and command-line option may change frequently.***
     - [SeFi-Image](./docs/sefi_image.md)
     - [HiDream-O1-Image](./docs/hidream_o1_image.md)
     - [Ideogram4](./docs/ideogram4.md)
-  - Image Edit Models
+  - [Image Edit Models](./docs/edit.md)
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)
     - [LongCat Image Edit](./docs/longcat_image.md)

docs/edit.md

@@ -0,0 +1,93 @@
+# Image Editing
+
+Image editing in `stable-diffusion.cpp` allows you to use reference images to guide the generation process, enabling tasks like identity preservation, style transfer, or layout modification.
+
+
+## Supported Models
+
+Depending on the architecture, different models handle reference images differently.
+
+| Model | Default Preset |
+| :--- | :--- |
+| [**FLUX.1-Kontext-dev**](./kontext.md) | `flux_kontext` |
+| [**LongCat Image Edit**](./longcat_image.md) | `longcat` |
+| [**Qwen Image Edit**](./qwen_image_edit.md) | `qwen` |
+| **Qwen Image LAYERED** | `qwen_layered` |
+| [**Flux.2 [Dev] / Flux.2 [Klein]**](./flux2.md) | `flux2` |
+| [**Boogu Image Edit**](./boogu_image.md) | `z_image_omni` |
+| **Krea2 (Community Edit LoRAs)** | `krea2_ostris_edit` |
+| **Anima (Community Edit LoRAs)** | `cosmos_reference` |
+
+Stable-diffusion.spp also supports basic Unet-based editing models like instruct-pix2pix or CosXL-Edit. This document is not about those.
+
+---
+
+## Configuring Reference Modes (`--ref-image-args`)
+
+Different DiT-based editing models require different configurations to process reference images correctly (e.g., whether to use a Vision Language Model (VLM) encoder or pass VAE-encoded images directly to the DiT).
+
+To simplify this, we provide **Presets**. By default, the system automatically selects the best preset based on the model architecture. However, you can override this using the `--ref-image-args` argument.
+
+### Usage
+The `--ref-image-args` argument accepts a comma-separated list of key-value pairs:
+
+**Using a preset:**
+`--ref-image-args "preset=qwen_layered"`
+
+**Using a preset with a specific override:**
+`--ref-image-args "preset=krea2_edit,force_ref_timestep_zero=true"`
+
+### Available Presets
+
+| Preset | Primary Use Case |
+| :--- | :--- |
+| `flux_kontext` | FLUX.1 Kontext |
+| `longcat` | LongCat Image Edit |
+| `flux2` | FLUX.2 models |
+| `qwen` | Qwen Image Edit |
+| `qwen_layered` | Qwen Image Layered |
+| `z_image_omni` | Boogu, Z-Image Omni |
+| `krea2_ostris_edit` | Most Krea2 Community edit LoRAs (trained with Ostris script) |
+| `krea2_edit` | Specifically for [lbouaraba/krea2edit](https://huggingface.co/conradlocke/krea2-identity-edit). (or similar) |
+| `cosmos_reference` | For Anima |
+| `default` | Uses the automatic detection based on model architecture. |
+
+---
+
+## Advanced Parameter Reference
+
+If presets are insufficient, you can manually configure the following parameters via `--ref-image-args`:
+
+| Key | Type | Description | Allowed Values |
+| :--- | :--- | :--- | :--- |
+| `preset` | string | Overrides the automatic preset. | (See the Presets table above) |
+| `pass_to_vlm` | bool | Whether reference images are passed to the VLM encoder. | `true`, `false` |
+| `pass_to_dit` | bool | Whether VAE-encoded references are passed directly to the DiT. | `true`, `false` |
+| `ref_index_mode` | string | Behavior of the RoPE index. | `fixed`, `increase`, `decrease` |
+| `force_ref_timestep_zero` | bool | Forces timestep=0 for reference tokens. | `true`, `false` (Krea2 only) |
+| `resize_before_vae` | bool | Whether reference images are resized before VAE encoding. | `true`, `false` |
+| `vae_input_max_pixels` | int | Maximum pixel area for VAE reference inputs. | Integer |
+| `vlm_resize_mode` | string | How to resize VLM reference inputs. | `longest_side`, `area`, `none` |
+| `vlm_max_size` | int | Maximum VLM input size; interpreted according to `vlm_resize_mode`. | Integer |
+| `vlm_min_size` | int | Minimum VLM input size; interpreted according to `vlm_resize_mode`. | Integer |
+| `vlm_size` | int | Shortcut to set both VLM min and max size to the same value. | Integer |
+
+### Preset Default Values
+
+For a technical overview of how each preset is configured, see the table below.
+
+| Preset | VLM | RoPE Index | Cond Resize | Special Notes |
+| :--- | :---: | :---: | :---: | :--- |
+| `flux_kontext` | No | `fixed` | `none` | |
+| `longcat` | Yes | `fixed` | `area` | |
+| `flux2` | No | `increase` | `none` | |
+| `qwen` | Yes | `increase` | `area` | |
+| `qwen_layered` | Yes | `decrease` | `area` | |
+| `z_image_omni` | Yes | `fixed` | `area` | |
+| `krea2_ostris_edit`| Yes | `increase` | `area` | `force_ref_timestep_zero = true` |
+| `krea2_edit` | Yes | `increase` | `longest` | `vlm_size = 768` |
+| `cosmos_reference` | No | `fixed` | `none` | `resize_before_vae = false` |
+
+**Additional Default Notes:**
+- **VLM Input Sizes:** For most presets, `vlm_max_size` and `vlm_min_size` are set to `-1`, meaning the values are model-dependent and handled automatically. In `area` mode they represent pixel area; in `longest_side` mode they represent a side length in pixels.
+- **VAE Input Size:** `vae_input_max_pixels` defaults to $1024 \times 1024$ pixels (`1048576`).

examples/common/common.cpp

@@ -975,6 +975,11 @@ ArgOptions SDGenerationParams::get_options() {
          "extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
          (int)',',
          &extra_tiling_args},
+        {"",
+         "--ref-image-args",
+         "Key-value list to set up the way the reference images are processed (empty = auto-detect from model weigths)",
+         (int)',',
+         &ref_image_args},
     };
 
     options.int_options = {
@@ -2418,6 +2423,22 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
         pulid_id_weight,
     };
 
+    if (!auto_resize_ref_image) {
+        if (!ref_image_args.empty()) {
+            ref_image_args += ",";
+        }
+        ref_image_args += "resize_before_vae=0";
+        LOG_WARN("Notice: --disable-auto-resize-ref-image is deprecated. Use --ref-image-args \"resize_before_vae=off\" instead.");
+    }
+
+    if (increase_ref_index) {
+        if (!ref_image_args.empty()) {
+            ref_image_args += ",";
+        }
+        ref_image_args += "ref_index_mode=increase";
+        LOG_WARN("Notice: --increase-ref-index is deprecated. Use --ref-image-args \"ref_index_mode=increase\" instead.");
+    }
+
     params.loras             = lora_vec.empty() ? nullptr : lora_vec.data();
     params.lora_count        = static_cast<uint32_t>(lora_vec.size());
     params.prompt            = prompt.c_str();
@@ -2426,8 +2447,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     params.init_image        = init_image.get();
     params.ref_images        = ref_image_views.empty() ? nullptr : ref_image_views.data();
     params.ref_images_count  = static_cast<int>(ref_image_views.size());
-    params.auto_resize_ref_image = auto_resize_ref_image;
+    params.ref_image_args    = ref_image_args.c_str();
-    params.increase_ref_index    = increase_ref_index;
     params.mask_image        = mask_image.get();
     params.width             = get_resolved_width();
     params.height            = get_resolved_height();

examples/common/common.h

@@ -227,6 +227,8 @@ struct SDGenerationParams {
     sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
     std::string extra_tiling_args;
 
+    std::string ref_image_args;
+
     std::string pm_id_images_dir;
     std::string pm_id_embed_path;
     float pm_style_strength = 20.f;

include/stable-diffusion.h

@@ -363,8 +363,7 @@ typedef struct {
     sd_image_t init_image;
     sd_image_t* ref_images;
     int ref_images_count;
-    bool auto_resize_ref_image;
+    const char* ref_image_args;
-    bool increase_ref_index;
     sd_image_t mask_image;
     int width;
     int height;

src/conditioning/conditioner.hpp

@@ -7,6 +7,7 @@
 
 #include "core/tensor_ggml.hpp"
 #include "core/util.h"
+#include "model/diffusion/model.hpp"
 #include "model/te/clip.hpp"
 #include "model/te/llm.hpp"
 #include "model/te/t5.hpp"
@@ -106,6 +107,7 @@ struct ConditionerParams {
     int height                                       = -1;
     bool zero_out_masked                             = false;
     const std::vector<sd::Tensor<float>>* ref_images = nullptr;  // for qwen image edit
+    RefImageParams ref_image_params;
 };
 
 struct Conditioner {
@@ -1993,6 +1995,54 @@ struct LLMEmbedder : public Conditioner {
         return new_hidden_states;
     }
 
+    void resize_image_dims(int height, int width, int& h_bar, int& w_bar, int factor, int min_size, int max_size, RefImageResizeMode mode) {
+        if (min_size > 0 && min_size == max_size) {
+            if (mode == RefImageResizeMode::AREA) {
+                double beta = std::sqrt(static_cast<double>(min_size) / (static_cast<double>(height) * width));
+                h_bar       = std::max(static_cast<int>(factor),
+                                       static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
+                w_bar       = std::max(static_cast<int>(factor),
+                                       static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
+            } else if (mode == RefImageResizeMode::LONGEST_SIDE) {
+                int current_max_side = std::max(height, width);
+                double beta          = static_cast<double>(min_size) / current_max_side;
+                h_bar                = std::max(static_cast<int>(factor),
+                                                static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
+                w_bar                = std::max(static_cast<int>(factor),
+                                                static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
+            }
+            return;
+        }
+
+        if (mode == RefImageResizeMode::AREA) {
+            double current_area = static_cast<double>(h_bar) * w_bar;
+            if (max_size > 0 && current_area > max_size) {
+                double beta = std::sqrt((static_cast<double>(height) * width) / static_cast<double>(max_size));
+                h_bar       = std::max(static_cast<int>(factor),
+                                       static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
+                w_bar       = std::max(static_cast<int>(factor),
+                                       static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
+            } else if (min_size > 0 && current_area < min_size) {
+                double beta = std::sqrt(static_cast<double>(min_size) / (static_cast<double>(height) * width));
+                h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
+                w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
+            }
+        } else if (mode == RefImageResizeMode::LONGEST_SIDE) {
+            int current_max_side = std::max(height, width);
+            if (max_size > 0 && current_max_side > max_size) {
+                double beta = static_cast<double>(max_size) / current_max_side;
+                h_bar       = std::max(static_cast<int>(factor),
+                                       static_cast<int>(std::floor(height * beta / factor)) * static_cast<int>(factor));
+                w_bar       = std::max(static_cast<int>(factor),
+                                       static_cast<int>(std::floor(width * beta / factor)) * static_cast<int>(factor));
+            } else if (min_size > 0 && current_max_side < min_size) {
+                double beta = static_cast<double>(min_size) / current_max_side;
+                h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
+                w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
+            }
+        }
+    }
+
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
         std::string prompt;
@@ -2008,6 +2058,7 @@ struct LLMEmbedder : public Conditioner {
         std::set<int> out_layers;
 
         int64_t t0                     = ggml_time_ms();
+        RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode;
 
         if (sd_version_is_lingbot_video(version)) {
             const int pad_token = 151643;
@@ -2040,28 +2091,35 @@ struct LLMEmbedder : public Conditioner {
 
                 for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
                     const auto& image = (*conditioner_params.ref_images)[i];
-                    double factor     = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
+                    const int factor  = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
                     int height        = static_cast<int>(image.shape()[1]);
                     int width         = static_cast<int>(image.shape()[0]);
-                    int min_pixels    = static_cast<int>(4 * factor * factor);
+
-                    int max_pixels    = static_cast<int>(16384 * factor * factor);
+                    int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
-                    int h_bar         = std::max(static_cast<int>(factor), static_cast<int>(std::round(height / factor) * factor));
+                    if (min_pixels <= 0) {
-                    int w_bar         = std::max(static_cast<int>(factor), static_cast<int>(std::round(width / factor) * factor));
+                        if (resize_mode == RefImageResizeMode::AREA) {
+                            min_pixels = static_cast<int>(4 * factor * factor);
+                        } else {
+                            min_pixels = static_cast<int>(2 * factor);
+                        }
+                    }
+                    int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
+                    if (max_pixels <= 0) {
+                        if (resize_mode == RefImageResizeMode::AREA) {
+                            max_pixels = static_cast<int>(16384 * factor * factor);
+                        } else {
+                            max_pixels = static_cast<int>(128 * factor);
+                        }
+                    }
+
+                    int h_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(height) / factor) * factor));
+                    int w_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(width) / factor) * factor));
 
                     if (std::max(height, width) > 200 * std::min(height, width)) {
                         LOG_WARN("LingBotVideo image aspect ratio is very large: %dx%d", width, height);
                     }
-                    if (h_bar * w_bar > max_pixels) {
+
-                        double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
+                    resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
-                        h_bar       = std::max(static_cast<int>(factor),
-                                               static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
-                        w_bar       = std::max(static_cast<int>(factor),
-                                               static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
-                    } else if (h_bar * w_bar < min_pixels) {
-                        double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
-                        h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
-                        w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
-                    }
 
                     LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
                     auto resized_image = clip_preprocess(image, w_bar, h_bar);
@@ -2092,30 +2150,33 @@ struct LLMEmbedder : public Conditioner {
                 prompt_template_encode_start_idx = 64;
                 int image_embed_idx              = 64 + 6;
 
-                int min_pixels          = 384 * 384;
+                int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
-                int max_pixels          = 560 * 560;
+                if (min_pixels <= 0) {
+                    min_pixels = 384;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        min_pixels *= min_pixels;
+                    }
+                }
+                int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
+                if (max_pixels <= 0) {
+                    max_pixels = 560;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        max_pixels *= max_pixels;
+                    }
+                }
+
                 std::string placeholder = "<|image_pad|>";
                 std::string img_prompt;
 
                 for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
                     const auto& image = (*conditioner_params.ref_images)[i];
-                    double factor     = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
+                    const int factor  = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
                     int height        = static_cast<int>(image.shape()[1]);
                     int width         = static_cast<int>(image.shape()[0]);
-                    int h_bar         = static_cast<int>(std::round(height / factor) * factor);
+                    int h_bar         = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
-                    int w_bar         = static_cast<int>(std::round(width / factor) * factor);
+                    int w_bar         = static_cast<int>(std::round(static_cast<double>(width) / factor) * factor);
 
-                    if (static_cast<double>(h_bar) * w_bar > max_pixels) {
+                    resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
-                        double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
-                        h_bar       = std::max(static_cast<int>(factor),
-                                               static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
-                        w_bar       = std::max(static_cast<int>(factor),
-                                               static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
-                    } else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
-                        double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
-                        h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
-                        w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
-                    }
 
                     LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
 
@@ -2170,16 +2231,33 @@ struct LLMEmbedder : public Conditioner {
                 std::string img_prompt;
                 const std::string placeholder = "<|image_pad|>";
 
+                int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
+                if (min_pixels <= 0) {
+                    min_pixels = 384;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        min_pixels *= min_pixels;
+                    }
+                }
+                int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
+                if (max_pixels <= 0) {
+                    max_pixels = 384;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        max_pixels *= max_pixels;
+                    }
+                }
+
                 for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
                     const auto& image = (*conditioner_params.ref_images)[i];
-                    double factor     = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
+                    const int factor  = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
                     int height        = static_cast<int>(image.shape()[1]);
                     int width         = static_cast<int>(image.shape()[0]);
-                    double beta       = std::sqrt((384.0 * 384.0) / (static_cast<double>(height) * static_cast<double>(width)));
+
-                    int h_bar         = std::max(static_cast<int>(factor),
+                    int h_bar = std::max(factor,
-                                                 static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
+                                         static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
-                    int w_bar         = std::max(static_cast<int>(factor),
+                    int w_bar = std::max(factor,
-                                                 static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
+                                         static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
+
+                    resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
 
                     LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
 
@@ -2221,17 +2299,33 @@ struct LLMEmbedder : public Conditioner {
             if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
                 std::string img_prompt        = "";
                 const std::string placeholder = "<|image_pad|>";
+                int min_pixels                = conditioner_params.ref_image_params.vlm_min_size;
+                if (min_pixels <= 0) {
+                    min_pixels = 384;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        min_pixels *= min_pixels;
+                    }
+                }
+                int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
+                if (max_pixels <= 0) {
+                    max_pixels = 1024;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        max_pixels *= max_pixels;
+                    }
+                }
 
                 for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
                     const auto& image = (*conditioner_params.ref_images)[i];
-                    double factor     = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
+                    const int factor  = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
                     int height        = static_cast<int>(image.shape()[1]);
                     int width         = static_cast<int>(image.shape()[0]);
-                    double beta       = std::sqrt((384.0 * 384.0) / (static_cast<double>(height) * static_cast<double>(width)));
+
-                    int h_bar         = std::max(static_cast<int>(factor),
+                    int h_bar = std::max(factor,
-                                                 static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
+                                         static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
-                    int w_bar         = std::max(static_cast<int>(factor),
+                    int w_bar = std::max(factor,
-                                                 static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
+                                         static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
+
+                    resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
 
                     LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
 
@@ -2268,30 +2362,33 @@ struct LLMEmbedder : public Conditioner {
                 min_length                       = 512 + prompt_template_encode_start_idx;
                 int image_embed_idx              = 36 + 6;
 
-                int min_pixels          = 384 * 384;
+                int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
-                int max_pixels          = 560 * 560;
+                if (min_pixels <= 0) {
+                    min_pixels = 384;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        min_pixels *= min_pixels;
+                    }
+                }
+                int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
+                if (max_pixels <= 0) {
+                    max_pixels = 560;
+                    if (resize_mode == RefImageResizeMode::AREA) {
+                        max_pixels *= max_pixels;
+                    }
+                }
+
                 std::string placeholder = "<|image_pad|>";
                 std::string img_prompt;
 
                 for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
                     const auto& image = (*conditioner_params.ref_images)[i];
-                    double factor     = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
+                    const int factor  = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
                     int height        = static_cast<int>(image.shape()[1]);
                     int width         = static_cast<int>(image.shape()[0]);
-                    int h_bar         = static_cast<int>(std::round(height / factor) * factor);
+                    int h_bar         = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
-                    int w_bar         = static_cast<int>(std::round(width / factor) * factor);
+                    int w_bar         = static_cast<int>(std::round(static_cast<double>(width) / factor) * factor);
 
-                    if (static_cast<double>(h_bar) * w_bar > max_pixels) {
+                    resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
-                        double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
-                        h_bar       = std::max(static_cast<int>(factor),
-                                               static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
-                        w_bar       = std::max(static_cast<int>(factor),
-                                               static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
-                    } else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
-                        double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
-                        h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
-                        w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
-                    }
 
                     LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
 

src/model/common/rope.hpp

@@ -18,10 +18,6 @@ namespace Rope {
         DECREASE,
     };
 
-    __STATIC_INLINE__ RefIndexMode ref_index_mode_from_bool(bool increase_ref_index) {
-        return increase_ref_index ? RefIndexMode::INCREASE : RefIndexMode::FIXED;
-    }
-
     template <class T>
     __STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
         std::vector<T> result(num);

src/model/diffusion/anima.hpp

@@ -487,7 +487,8 @@ namespace Anima {
                              ggml_tensor* t5_ids                   = nullptr,
                              ggml_tensor* t5_weights               = nullptr,
                              ggml_tensor* adapter_q_pe             = nullptr,
-                             ggml_tensor* adapter_k_pe = nullptr) {
+                             ggml_tensor* adapter_k_pe             = nullptr,
+                             std::vector<ggml_tensor*> ref_latents = {}) {
             GGML_ASSERT(x->ne[3] == 1);
 
             auto x_embedder       = std::dynamic_pointer_cast<XEmbedder>(blocks["x_embedder"]);
@@ -503,7 +504,15 @@ namespace Anima {
             x                 = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2);  // [N, C + 1, H, W]
 
             x               = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);  // [N, h*w, (C+1)*ph*pw]
-
+            int64_t img_len = x->ne[1];
+            if (ref_latents.size() > 0) {
+                for (ggml_tensor* ref : ref_latents) {
+                    auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, ref->ne[0], ref->ne[1], 1, ref->ne[3]);
+                    ref               = ggml_concat(ctx->ggml_ctx, ref, padding_mask, 2);  // [N, C + 1, H, W]
+                    ref               = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
+                    x                 = ggml_concat(ctx->ggml_ctx, x, ref, 1);
+                }
+            }
             x = x_embedder->forward(ctx, x);
 
             auto timestep_proj     = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
@@ -543,6 +552,7 @@ namespace Anima {
                 x          = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
                 sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x");
             }
+            x = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, img_len);
 
             x = final_layer->forward(ctx, x, embedded_timestep, temb);  // [N, h*w, ph*pw*C]
 
@@ -602,8 +612,8 @@ namespace Anima {
                                                          const std::vector<int>& axes_dim,
                                                          float h_extrapolation_ratio,
                                                          float w_extrapolation_ratio,
-                                                         float t_extrapolation_ratio) {
+                                                         float t_extrapolation_ratio,
-            static const std::vector<ggml_tensor*> empty_ref_latents;
+                                                         const std::vector<ggml_tensor*>& ref_latents) {
             auto ids = Rope::gen_flux_ids(h,
                                           w,
                                           patch_size,
@@ -611,7 +621,7 @@ namespace Anima {
                                           static_cast<int>(axes_dim.size()),
                                           0,
                                           {},
-                                          empty_ref_latents,
+                                          ref_latents,
                                           Rope::RefIndexMode::FIXED,
                                           1.0f,
                                           false);
@@ -628,12 +638,18 @@ namespace Anima {
                                  const sd::Tensor<float>& timesteps_tensor,
                                  const sd::Tensor<float>& context_tensor                  = {},
                                  const sd::Tensor<int32_t>& t5_ids_tensor                 = {},
-                                 const sd::Tensor<float>& t5_weights_tensor = {}) {
+                                 const sd::Tensor<float>& t5_weights_tensor               = {},
+                                 const std::vector<sd::Tensor<float>>& ref_latents_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* t5_ids     = make_optional_input(t5_ids_tensor);
             ggml_tensor* t5_weights = make_optional_input(t5_weights_tensor);
+            std::vector<ggml_tensor*> ref_latents;
+            ref_latents.reserve(ref_latents_tensor.size());
+            for (const auto& ref_latent_tensor : ref_latents_tensor) {
+                ref_latents.push_back(make_input(ref_latent_tensor));
+            }
             GGML_ASSERT(x->ne[3] == 1);
             ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
 
@@ -650,7 +666,8 @@ namespace Anima {
                                                            config.axes_dim,
                                                            4.0f,
                                                            4.0f,
-                                                           1.0f);
+                                                           1.0f,
+                                                           ref_latents);
             int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
             auto image_pe         = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
             set_backend_tensor_data(image_pe, image_pe_vec.data());
@@ -682,7 +699,8 @@ namespace Anima {
                                           t5_ids,
                                           t5_weights,
                                           adapter_q_pe,
-                                          adapter_k_pe);
+                                          adapter_k_pe,
+                                          ref_latents);
 
             ggml_build_forward_expand(gf, out);
             return gf;
@@ -693,9 +711,11 @@ namespace Anima {
                                   const sd::Tensor<float>& timesteps,
                                   const sd::Tensor<float>& context                  = {},
                                   const sd::Tensor<int32_t>& t5_ids                 = {},
-                                  const sd::Tensor<float>& t5_weights = {}) {
+                                  const sd::Tensor<float>& t5_weights               = {},
+                                  const std::vector<sd::Tensor<float>>& ref_latents = {},
+                                  const RefImageParams& ref_image_params            = REF_IMAGE_PRESETS.at("cosmos_reference")) {
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, t5_ids, t5_weights);
+                return build_graph(x, timesteps, context, t5_ids, t5_weights, ref_latents);
             };
             return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
@@ -705,12 +725,15 @@ namespace Anima {
             GGML_ASSERT(diffusion_params.x != nullptr);
             GGML_ASSERT(diffusion_params.timesteps != nullptr);
             const auto* extra = diffusion_extra_as<AnimaDiffusionExtra>(diffusion_params);
+            static const std::vector<sd::Tensor<float>> empty_ref_latents;
             return compute(n_threads,
                            *diffusion_params.x,
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
                            tensor_or_empty(extra->t5_ids),
-                           tensor_or_empty(extra->t5_weights));
+                           tensor_or_empty(extra->t5_weights),
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.ref_image_params);
         }
     };
 }  // namespace Anima

src/model/diffusion/boogu.hpp

@@ -827,7 +827,7 @@ namespace Boogu {
                            *diffusion_params.x,
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
-                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents);
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents);
         }
     };
 }  // namespace Boogu

src/model/diffusion/flux.hpp

@@ -1642,8 +1642,8 @@ namespace Flux {
                            tensor_or_empty(diffusion_params.c_concat),
                            tensor_or_empty(diffusion_params.y),
                            tensor_or_empty(extra->guidance),
-                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.ref_index_mode,
+                           diffusion_params.ref_image_params.ref_index_mode,
                            extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
                            tensor_or_empty(extra->pulid_id),
                            extra->pulid_id_weight);

src/model/diffusion/krea2.hpp

@@ -615,7 +615,8 @@ namespace Krea2 {
                              ggml_tensor* timestep,
                              ggml_tensor* context,
                              ggml_tensor* pe,
-                             std::vector<ggml_tensor*> ref_latents = {}) {
+                             std::vector<ggml_tensor*> ref_latents = {},
+                             bool zero_timestep_refs               = false) {
             int64_t W = x->ne[0];
             int64_t H = x->ne[1];
             int64_t N = x->ne[3];
@@ -645,7 +646,7 @@ namespace Krea2 {
             auto tvec = tproj->forward(ctx, t);
 
             ggml_tensor* tvec_0 = nullptr;
-            if (ref_latents.size() > 0) {
+            if (ref_latents.size() > 0 && zero_timestep_refs) {
                 // "index_timestep_zero" mode: use timestep = 0 for ref latents
                 auto timestep_0 = ggml_scale(ctx->ggml_ctx, timestep, 0.0f);
                 auto t_0        = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_0, static_cast<int>(config.timestep_dim), 10000, 1000.f);
@@ -719,7 +720,7 @@ namespace Krea2 {
                                  const sd::Tensor<float>& timesteps_tensor,
                                  const sd::Tensor<float>& context_tensor,
                                  const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
-                                 Rope::RefIndexMode ref_index_mode                        = Rope::RefIndexMode::FIXED) {
+                                 const RefImageParams& ref_image_params                   = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) {
             ggml_cgraph* gf        = new_graph_custom(KREA2_GRAPH_SIZE);
             ggml_tensor* x         = make_input(x_tensor);
             ggml_tensor* timesteps = make_input(timesteps_tensor);
@@ -741,13 +742,13 @@ namespace Krea2 {
                                        config.theta,
                                        config.axes_dim,
                                        ref_latents,
-                                       ref_index_mode);
+                                       ref_image_params.ref_index_mode);
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
             set_backend_tensor_data(pe, pe_vec.data());
 
             auto runner_ctx  = get_context();
-            ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents);
+            ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents, ref_image_params.force_ref_timestep_zero);
             ggml_build_forward_expand(gf, out);
             return gf;
         }
@@ -757,9 +758,9 @@ namespace Krea2 {
                                   const sd::Tensor<float>& timesteps,
                                   const sd::Tensor<float>& context,
                                   const std::vector<sd::Tensor<float>>& ref_latents = {},
-                                  Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::FIXED) {
+                                  const RefImageParams& ref_image_params            = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) {
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
+                return build_graph(x, timesteps, context, ref_latents, ref_image_params);
             };
             return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
@@ -773,8 +774,8 @@ namespace Krea2 {
                            *diffusion_params.x,
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
-                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.ref_index_mode);
+                           diffusion_params.ref_image_params);
         }
     };
 }  // namespace Krea2

src/model/diffusion/model.hpp

@@ -10,6 +10,36 @@
 #include "model/common/rope.hpp"
 #include "model_manager.h"
 
+enum class RefImageResizeMode {
+    NONE,
+    LONGEST_SIDE,
+    AREA,
+};
+
+struct RefImageParams {
+    bool pass_to_vlm                   = false;
+    bool pass_to_dit                   = true;
+    Rope::RefIndexMode ref_index_mode  = Rope::RefIndexMode::FIXED;
+    bool force_ref_timestep_zero       = false;
+    bool resize_before_vae             = true;
+    int vae_input_max_pixels           = -1;
+    RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA;
+    int vlm_min_size                   = -1;
+    int vlm_max_size                   = -1;
+};
+
+const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
+    {"flux_kontext", {false, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
+    {"longcat", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
+    {"flux2", {false, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
+    {"qwen", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
+    {"qwen_layered", {true, true, Rope::RefIndexMode::DECREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
+    {"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
+    {"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}},
+    {"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}},
+    {"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}},
+};
+
 struct UNetDiffusionExtra {
     int num_video_frames                           = -1;
     const std::vector<sd::Tensor<float>>* controls = nullptr;
@@ -74,7 +104,7 @@ struct DiffusionParams {
     const sd::Tensor<float>* c_concat                 = nullptr;
     const sd::Tensor<float>* y                        = nullptr;
     const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
-    Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::FIXED;
+    RefImageParams ref_image_params                   = {false, false, Rope::RefIndexMode::FIXED, false};
     DiffusionExtraParams extra                        = std::monostate{};
 };
 

src/model/diffusion/qwen_image.hpp

@@ -715,8 +715,8 @@ namespace Qwen {
                            *diffusion_params.x,
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
-                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.ref_index_mode);
+                           diffusion_params.ref_image_params.ref_index_mode);
         }
 
         void test() {

src/model/diffusion/z_image.hpp

@@ -648,8 +648,8 @@ namespace ZImage {
                            *diffusion_params.x,
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
-                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.ref_index_mode);
+                           diffusion_params.ref_image_params.ref_index_mode);
         }
 
         void test() {

src/stable-diffusion.cpp

@@ -2297,7 +2297,7 @@ public:
                              const char* extra_sample_args,
                              const std::vector<float>& sigmas,
                              const std::vector<sd::Tensor<float>>& ref_latents,
-                             bool increase_ref_index,
+                             const RefImageParams& ref_image_params,
                              const sd::Tensor<float>& denoise_mask,
                              const sd::Tensor<float>& vace_context,
                              float vace_strength,
@@ -2460,7 +2460,7 @@ public:
             DiffusionParams diffusion_params;
             diffusion_params.x                = &noised_input;
             diffusion_params.timesteps        = &timesteps_tensor;
-            diffusion_params.ref_index_mode = Rope::ref_index_mode_from_bool(increase_ref_index);
+            diffusion_params.ref_image_params = ref_image_params;
             sd::guidance::GuidanceInput step_guidance_input;
             step_guidance_input.step          = step;
             step_guidance_input.schedule_size = sigmas.size();
@@ -2840,6 +2840,126 @@ public:
         auto flow_denoiser = std::dynamic_pointer_cast<DiscreteFlowDenoiser>(denoiser);
         return !!flow_denoiser;
     }
+
+    std::string get_default_ref_image_preset(SDVersion version) const {
+        if (sd_version_is_longcat(version)) {
+            return "longcat";
+        } else if (sd_version_is_flux(version)) {
+            return "flux_kontext";
+        } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
+            return "flux2";
+        } else if (version == VERSION_QWEN_IMAGE_LAYERED) {
+            return "qwen_layered";
+        } else if (sd_version_is_qwen_image(version)) {
+            return "qwen";
+        } else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) {
+            return "z_image_omni";
+        } else if (sd_version_is_krea2(version)) {
+            // have to make a choice between "krea2_edit" mode (for lbouaraba/krea2edit)
+            // and "krea2_ostris_edit" (for krea2 ostris edit)
+            // since krea2 ostris edit support predates, it should probably be default
+            return "krea2_ostris_edit";
+        } else if (sd_version_is_anima(version)) {
+            return "cosmos_reference";
+        }
+        return "default";
+    }
+
+    RefImageParams resolve_ref_image_params(const char* ref_image_args) const {
+        RefImageParams params;
+        std::string preset_name = get_default_ref_image_preset(version);
+
+        for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) {
+            if (key == "preset") {
+                std::string requested_preset_name = value;
+                if (REF_IMAGE_PRESETS.count(requested_preset_name)) {
+                    preset_name = requested_preset_name;
+                } else if (value != "default") {
+                    std::string valid_list;
+                    for (auto const& [name, _] : REF_IMAGE_PRESETS) {
+                        valid_list += (valid_list.empty() ? "" : ", ") + name;
+                    }
+                    LOG_WARN("ignoring invalid reference image preset '%s'. Valid options: [%s]", value.c_str(), valid_list.c_str());
+                }
+                break;
+            }
+        }
+        if (preset_name != "default") {
+            LOG_INFO("Using '%s' preset for reference images", preset_name.c_str());
+            params = REF_IMAGE_PRESETS.at(preset_name);
+        }
+
+        for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) {
+            if (key == "pass_to_vlm") {
+                if (!parse_strict_bool(value, params.pass_to_vlm)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "pass_to_dit") {
+                if (!parse_strict_bool(value, params.pass_to_dit)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "ref_index_mode") {
+                if (value == "fixed") {
+                    params.ref_index_mode = Rope::RefIndexMode::FIXED;
+                } else if (value == "increase") {
+                    params.ref_index_mode = Rope::RefIndexMode::INCREASE;
+                } else if (value == "decrease") {
+                    params.ref_index_mode = Rope::RefIndexMode::DECREASE;
+                } else {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "force_ref_timestep_zero") {
+                if (!parse_strict_bool(value, params.force_ref_timestep_zero)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "resize_before_vae") {
+                if (!parse_strict_bool(value, params.resize_before_vae)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "vae_input_max_pixels") {
+                if (!parse_strict_int(value, params.vae_input_max_pixels)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "vlm_resize_mode") {
+                if (value == "longest_side") {
+                    params.vlm_resize_mode = RefImageResizeMode::LONGEST_SIDE;
+                } else if (value == "area") {
+                    params.vlm_resize_mode = RefImageResizeMode::AREA;
+                } else if (value == "none") {
+                    params.vlm_resize_mode = RefImageResizeMode::NONE;
+                } else {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "vlm_max_size") {
+                if (!parse_strict_int(value, params.vlm_max_size)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "vlm_min_size") {
+                if (!parse_strict_int(value, params.vlm_min_size)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key != "preset" && key != "vlm_size") {
+                LOG_WARN("ignoring unknown reference image arg '%s'", key.c_str());
+            }
+        }
+        for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) {
+            if (key == "vlm_size") {
+                int vlm_size;
+                if (!parse_strict_int(value, vlm_size)) {
+                    LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str());
+                } else {
+                    LOG_INFO("vlm_size override: setting both min and max size to %ld", (long)vlm_size);
+                    params.vlm_min_size = vlm_size;
+                    params.vlm_max_size = vlm_size;
+                }
+                break;
+            }
+        }
+        if (params.force_ref_timestep_zero && !sd_version_is_krea2(version)) {
+            LOG_WARN("force_ref_timestep_zero is only supported by Krea2 architecture for now");
+        }
+        return params;
+    }
 };
 
 /*================================================= SD API ==================================================*/
@@ -3301,6 +3421,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* 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;
@@ -3338,8 +3459,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              "batch_count: %d\n"
              "qwen_image_layers: %d\n"
              "ref_images_count: %d\n"
-             "auto_resize_ref_image: %s\n"
+             "ref_image_args: %s\n"
-             "increase_ref_index: %s\n"
              "control_strength: %.2f\n"
              "photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n"
              "VAE tiling: %s (temporal=%s, extra_tiling_args=%s)\n"
@@ -3357,8 +3477,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              sd_img_gen_params->batch_count,
              sd_img_gen_params->qwen_image_layers,
              sd_img_gen_params->ref_images_count,
-             BOOL_STR(sd_img_gen_params->auto_resize_ref_image),
+             SAFE_STR(sd_img_gen_params->ref_image_args),
-             BOOL_STR(sd_img_gen_params->increase_ref_index),
              sd_img_gen_params->control_strength,
              sd_img_gen_params->pm_params.style_strength,
              sd_img_gen_params->pm_params.id_images_count,
@@ -3655,8 +3774,6 @@ struct GenerationRequest {
     float strength                           = 1.f;
     float control_strength                   = 0.f;
     float eta                                = 0.f;
-    bool increase_ref_index                  = false;
-    bool auto_resize_ref_image               = false;
     sd_guidance_params_t guidance            = {};
     sd_guidance_params_t high_noise_guidance = {};
     sd_pm_params_t pm_params                 = {};
@@ -3682,8 +3799,6 @@ struct GenerationRequest {
         strength                    = sd_img_gen_params->strength;
         control_strength            = sd_img_gen_params->control_strength;
         eta                         = sd_img_gen_params->sample_params.eta;
-        increase_ref_index          = sd_img_gen_params->increase_ref_index;
-        auto_resize_ref_image       = sd_img_gen_params->auto_resize_ref_image;
         has_ref_images              = sd_img_gen_params->ref_images_count > 0;
         guidance                    = sd_img_gen_params->sample_params.guidance;
         pm_params                   = sd_img_gen_params->pm_params;
@@ -4431,7 +4546,8 @@ static sd::Tensor<float> ensure_image_tensor_channels(sd::Tensor<float> image, i
 static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd_ctx_t* sd_ctx,
                                                                               const sd_img_gen_params_t* sd_img_gen_params,
                                                                               GenerationRequest* request,
-                                                                              SamplePlan* plan) {
+                                                                              SamplePlan* plan,
+                                                                              const RefImageParams& ref_image_params) {
     int64_t prepare_start_ms = ggml_time_ms();
 
     sd::Tensor<float> init_image_tensor;
@@ -4574,9 +4690,10 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
             continue;
         }
         sd::Tensor<float> ref_latent;
-        if (request->auto_resize_ref_image && !sd_version_is_pid(sd_ctx->sd->version)) {
+        if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd_ctx->sd->version)) {
             LOG_DEBUG("auto resize ref images");
-            int vae_image_size = std::min(1024 * 1024, request->width * request->height);
+            int target_pixels  = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : 1024 * 1024;
+            int vae_image_size = std::min(target_pixels, request->width * request->height);
             double vae_width   = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]);
             double vae_height  = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0];
 
@@ -4703,7 +4820,8 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
                                                                             const sd_img_gen_params_t* sd_img_gen_params,
                                                                             GenerationRequest* request,
                                                                             SamplePlan* plan,
-                                                                            ImageGenerationLatents* latents) {
+                                                                            ImageGenerationLatents* latents,
+                                                                            const RefImageParams& ref_image_params) {
     ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()};
 
     ConditionerParams condition_params;
@@ -4711,7 +4829,11 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
     condition_params.clip_skip = request->clip_skip;
     condition_params.width     = request->width;
     condition_params.height    = request->height;
+    if (ref_image_params.pass_to_vlm) {
         condition_params.ref_images = &latents->ref_images;
+    }
+
+    condition_params.ref_image_params = ref_image_params;
 
     sd_ctx->sd->prepare_generation_extensions(request->pm_params,
                                               request->pulid_params,
@@ -4721,7 +4843,7 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
     condition_params.zero_out_masked = false;
     auto cond                        = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
                                                                                            condition_params);
-    if (cond.c_concat.empty()) {
+    if (cond.c_concat.empty() && ref_image_params.pass_to_dit) {
         cond.c_concat = latents->concat_latent;  // TODO: optimize
     }
 
@@ -4750,7 +4872,7 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
             uncond                           = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
                                                                                                    condition_params);
         }
-        if (uncond.c_concat.empty()) {
+        if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
             uncond.c_concat = latents->concat_latent;  // TODO: optimize
         }
     }
@@ -4774,7 +4896,7 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
             }
             img_uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
                                                                              condition_params);
-            if (img_uncond.c_concat.empty()) {
+            if (img_uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
                 img_uncond.c_concat = latents->img_uncond_concat_latent;  // TODO: optimize
             }
         }
@@ -5091,13 +5213,16 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
     sd_ctx->sd->apply_loras(sd_img_gen_params->loras, sd_img_gen_params->lora_count);
     apply_circular_axes_to_diffusion(sd_ctx, sd_img_gen_params->circular_x, sd_img_gen_params->circular_y);
 
+    const RefImageParams ref_image_params = sd_ctx->sd->resolve_ref_image_params(sd_img_gen_params->ref_image_args);
+
     ImageVaeAxesGuard axes_guard(sd_ctx, sd_img_gen_params, request);
 
     SamplePlan plan(sd_ctx, sd_img_gen_params, request);
     auto latents_opt = prepare_image_generation_latents(sd_ctx,
                                                         sd_img_gen_params,
                                                         &request,
-                                                        &plan);
+                                                        &plan,
+                                                        ref_image_params);
     if (!latents_opt.has_value()) {
         return false;
     }
@@ -5107,7 +5232,8 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
                                                       sd_img_gen_params,
                                                       &request,
                                                       &plan,
-                                                      &latents);
+                                                      &latents,
+                                                      ref_image_params);
     if (!embeds_opt.has_value()) {
         return false;
     }
@@ -5153,7 +5279,7 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
                                                    plan.extra_sample_args,
                                                    plan.sigmas,
                                                    latents.ref_latents,
-                                                   request.increase_ref_index,
+                                                   ref_image_params,
                                                    latents.denoise_mask,
                                                    sd::Tensor<float>(),
                                                    1.f,
@@ -5274,7 +5400,7 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
                                                             plan.extra_sample_args,
                                                             hires_sigma_sched,
                                                             latents.ref_latents,
-                                                            request.increase_ref_index,
+                                                            ref_image_params,
                                                             hires_denoise_mask,
                                                             sd::Tensor<float>(),
                                                             1.f,
@@ -5665,6 +5791,9 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx,
     condition_params.text            = request.prompt;
     condition_params.zero_out_masked = true;
     condition_params.ref_images      = &latents.ref_images;
+    if (sd_version_is_lingbot_video(sd_ctx->sd->version)) {
+        condition_params.ref_image_params.vlm_resize_mode = RefImageResizeMode::AREA;
+    }
 
     int64_t prepare_start_ms = ggml_time_ms();
     embeds.cond              = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
@@ -5945,6 +6074,8 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
 
     sd_ctx->sd->reset_cancel_flag();
 
+    const RefImageParams ref_image_params;
+
     if (num_frames_out != nullptr) {
         *num_frames_out = 0;
     }
@@ -6034,7 +6165,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                                                            plan.high_noise_extra_sample_args,
                                                            high_noise_sigmas,
                                                            std::vector<sd::Tensor<float>>{},
-                                                           false,
+                                                           ref_image_params,
                                                            latents.denoise_mask,
                                                            latents.vace_context,
                                                            request.vace_strength,
@@ -6076,7 +6207,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                                                         plan.extra_sample_args,
                                                         plan.sigmas,
                                                         std::vector<sd::Tensor<float>>{},
-                                                        false,
+                                                        ref_image_params,
                                                         latents.denoise_mask,
                                                         latents.vace_context,
                                                         request.vace_strength,
@@ -6214,7 +6345,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                                             plan.extra_sample_args,
                                             hires_sigma_sched,
                                             std::vector<sd::Tensor<float>>{},
-                                            false,
+                                            ref_image_params,
                                             hires_denoise_mask,
                                             sd::Tensor<float>(),
                                             hires_request.vace_strength,