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

Author: stduhpf

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,30 +2423,45 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
         pulid_id_weight,
     };
 
-    params.loras                 = lora_vec.empty() ? nullptr : lora_vec.data();
-    params.lora_count            = static_cast<uint32_t>(lora_vec.size());
-    params.prompt                = prompt.c_str();
-    params.negative_prompt       = negative_prompt.c_str();
-    params.clip_skip             = clip_skip;
-    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.increase_ref_index    = increase_ref_index;
-    params.mask_image            = mask_image.get();
-    params.width                 = get_resolved_width();
-    params.height                = get_resolved_height();
-    params.sample_params         = sample_params;
-    params.strength              = strength;
-    params.seed                  = seed;
-    params.batch_count           = batch_count;
-    params.qwen_image_layers     = qwen_image_layers;
-    params.control_image         = control_image.get();
-    params.control_strength      = control_strength;
-    params.pm_params             = pm_params;
-    params.pulid_params          = pulid_params;
-    params.vae_tiling_params     = vae_tiling_params;
-    params.cache                 = cache_params;
+    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();
+    params.negative_prompt   = negative_prompt.c_str();
+    params.clip_skip         = clip_skip;
+    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.ref_image_args    = ref_image_args.c_str();
+    params.mask_image        = mask_image.get();
+    params.width             = get_resolved_width();
+    params.height            = get_resolved_height();
+    params.sample_params     = sample_params;
+    params.strength          = strength;
+    params.seed              = seed;
+    params.batch_count       = batch_count;
+    params.qwen_image_layers = qwen_image_layers;
+    params.control_image     = control_image.get();
+    params.control_strength  = control_strength;
+    params.pm_params         = pm_params;
+    params.pulid_params      = pulid_params;
+    params.vae_tiling_params = vae_tiling_params;
+    params.cache             = cache_params;
 
     params.hires.enabled             = hires_enabled;
     params.hires.upscaler            = resolved_hires_upscaler;

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;
-    bool increase_ref_index;
+    const char* ref_image_args;
     sd_image_t mask_image;
     int width;
     int height;