feat: add Mage-Flow support

README.md

@@ -54,6 +54,7 @@ API and command-line option may change frequently.***
     - [ERNIE-Image](./docs/ernie_image.md)
     - [Boogu Image](./docs/boogu_image.md)
     - [Krea2](./docs/krea2.md)
+    - [Mage-Flow](./docs/mage_flow.md)
     - [SeFi-Image](./docs/sefi_image.md)
     - [HiDream-O1-Image](./docs/hidream_o1_image.md)
     - [Ideogram4](./docs/ideogram4.md)
@@ -62,6 +63,7 @@ API and command-line option may change frequently.***
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)
     - [LongCat Image Edit](./docs/longcat_image.md)
     - [Boogu Image Edit](./docs/boogu_image.md)
+    - [Mage-Flow-Edit](./docs/mage_flow.md#image-editing)
   - Video Models
     - [Wan2.1/Wan2.2](./docs/wan.md)
     - [LTX-2.3](./docs/ltx2.md)

docs/edit.md

@@ -16,6 +16,7 @@ Depending on the architecture, different models handle reference images differen
 | [**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` |
+| [**Mage-Flow-Edit**](./mage_flow.md#image-editing) | `mage_flow` |
 | **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.
@@ -48,6 +49,7 @@ The `--ref-image-args` argument accepts a comma-separated list of key-value pair
 | `qwen_layered` | Qwen Image Layered |
 | `z_image_omni` | Boogu, Z-Image Omni |
 | `krea2_ostris_edit` | Most Krea2 Community edit LoRAs (trained with Ostris script) |
+| `mage_flow` | Mage-Flow-Edit |
 | `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. |
@@ -83,6 +85,7 @@ For a technical overview of how each preset is configured, see the table below.
 | `flux2` | No | `increase` | `none` | |
 | `qwen` | Yes | `increase` | `area` | |
 | `qwen_layered` | Yes | `decrease` | `area` | |
+| `mage_flow` | Yes | `increase` | `longest` | `vlm_max_size = 384`, VAE input resized to target |
 | `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` |

docs/mage_flow.md

@@ -0,0 +1,45 @@
+# Mage-Flow
+
+[Mage-Flow](https://github.com/microsoft/Mage) uses a 4B native-resolution multimodal diffusion transformer, Qwen3-VL for text and image conditioning, and the 128-channel Mage-VAE. Both text-to-image and instruction-based image editing checkpoints are supported.
+
+## Download weights
+
+- Download Mage-Flow
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow/tree/main/transformer
+- Download Mage-Flow-Base
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Base/tree/main/transformer
+- Download Mage-Flow-Turbo
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Turbo/tree/main/transformer
+- Download Mage-Flow-Edit
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit/tree/main/transformer
+- Download Mage-Flow-Edit-Turbo
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo/tree/main/transformer
+- Download Mage-Flow-Edit-Base
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit-Base/tree/main/transformer
+- Download Mage-Flow vae
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow/tree/main/vae
+- Download Qwen3-VL 4B
+    - safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders
+    - gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main
+
+## Text-to-image
+
+Use 30 steps for Base models and 4 steps with `--cfg-scale 1` for Turbo models. Image dimensions must be multiples of 16; the official checkpoints are trained for native resolutions from 512 to 2048 pixels.
+
+```bash
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\Mage-Flow-Turbo.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\models\vae\mage_vae.safetensors -p "a lovely cat holding a sign says 'mage.cpp'" --cfg-scale 1.0 --steps 4 --diffusion-fa -v --offload-to-cpu
+```
+
+<img width="256" alt="Mage-Flow example" src="../assets/mage_flow/example.png" />
+
+## Image editing
+
+Mage-Flow-Edit accepts one or more reference images. The default `mage_flow` reference preset sends each image to both Qwen3-VL and the diffusion transformer, caps the VLM copy's longest edge at 384 pixels, and keeps the VAE copy at the requested output resolution.
+
+For the Turbo edit checkpoint, use 4 steps and `--cfg-scale 1`.
+
+```bash
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\Mage-Flow-Edit.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3-VL-4B-Instruct-mmproj-BF16.gguf --vae ..\models\vae\mage_vae.safetensors -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'mage.cpp'" --cfg-scale 4.0 --sampling-method euler -v --diffusion-fa --offload-to-cpu
+```
+
+<img width="256" alt="Mage-Flow-Edit example" src="../assets/mage_flow/edit_example.png" />

examples/CMakeLists.txt

@@ -1,4 +1,4 @@
 include_directories(${CMAKE_CURRENT_SOURCE_DIR})
 
 add_subdirectory(cli)
-add_subdirectory(server)
+add_subdirectory(server)

src/conditioning/conditioner.hpp

@@ -1800,7 +1800,12 @@ struct LLMEmbedder : public Conditioner {
             arch = LLM::LLMArch::GPT_OSS_20B;
         } else if (sd_version_is_pid(version)) {
             arch = LLM::LLMArch::GEMMA2_2B;
-        } else if (sd_version_is_lingbot_video(version) || sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) {
+        } else if (sd_version_is_lingbot_video(version) ||
+                   sd_version_is_ideogram4(version) ||
+                   sd_version_is_boogu_image(version) ||
+                   sd_version_is_sefi_image(version) ||
+                   sd_version_is_krea2(version) ||
+                   sd_version_is_mage_flow(version)) {
             arch = LLM::LLMArch::QWEN3_VL;
         } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
             arch = LLM::LLMArch::QWEN3;
@@ -2205,22 +2210,22 @@ struct LLMEmbedder : public Conditioner {
             prompt += conditioner_params.text;
             prompt_attn_range = {0, 0};
             prompt += "<|im_end|>\n<|im_start|>assistant\n";
-        } else if (sd_version_is_qwen_image(version)) {
+        } else if (sd_version_is_qwen_image(version) || sd_version_is_mage_flow(version)) {
             if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
-                LOG_INFO("QwenImageEditPlusPipeline");
+                LOG_INFO("%s", sd_version_is_mage_flow(version) ? "MageFlowEditPipeline" : "QwenImageEditPlusPipeline");
                 prompt_template_encode_start_idx = 64;
                 int image_embed_idx              = 64 + 6;
 
                 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 = sd_version_is_mage_flow(version) ? -1 : 384;
+                    if (min_pixels > 0 && 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;
+                    max_pixels = sd_version_is_mage_flow(version) ? 384 : 560;
                     if (resize_mode == RefImageResizeMode::AREA) {
                         max_pixels *= max_pixels;
                     }
@@ -2248,7 +2253,7 @@ struct LLMEmbedder : public Conditioner {
                     image_embeds.emplace_back(image_embed_idx, image_embed);
                     image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
 
-                    img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";  // [24669, 220, index, 25, 220, 151652]
+                    img_prompt += (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>";
                     int64_t num_image_tokens = image_embed.shape()[1];
                     img_prompt.reserve(num_image_tokens * placeholder.size());
                     for (int j = 0; j < num_image_tokens; j++) {
@@ -2276,6 +2281,9 @@ struct LLMEmbedder : public Conditioner {
 
                 prompt += "<|im_end|>\n<|im_start|>assistant\n";
             }
+            if (sd_version_is_mage_flow(version)) {
+                max_length = 2048 + prompt_template_encode_start_idx;
+            }
         } else if (sd_version_is_boogu_image(version)) {
             prompt_template_encode_start_idx = 0;
 

src/core/ggml_extend.hpp

@@ -3609,18 +3609,19 @@ class Conv2d_grouped : public UnaryBlock {
         if (groups == in_channels && groups == out_channels) {
             ggml_tensor* res;
             if (ctx->conv2d_direct_enabled) {
-                res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, x, w,
+                res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, w, x,
                                              stride.second, stride.first,
                                              padding.second, padding.first,
                                              dilation.second, dilation.first);
             } else {
-                res = ggml_conv_2d_dw(ctx->ggml_ctx, x, w,
+                res = ggml_conv_2d_dw(ctx->ggml_ctx, w, x,
                                       stride.second, stride.first,
                                       padding.second, padding.first,
                                       dilation.second, dilation.first);
             }
             if (b) {
-                res = ggml_add(ctx->ggml_ctx, res, b);
+                b   = ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, b->ne[0], 1);
+                res = ggml_add_inplace(ctx->ggml_ctx, res, b);
             }
             return res;
         }

src/model.h

@@ -55,6 +55,7 @@ enum SDVersion {
     VERSION_IDEOGRAM4,
     VERSION_SEFI_IMAGE,
     VERSION_KREA2,
+    VERSION_MAGE_FLOW,
     VERSION_ESRGAN,
     VERSION_COUNT,
 };
@@ -227,6 +228,10 @@ static inline bool sd_version_is_krea2(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_mage_flow(SDVersion version) {
+    return version == VERSION_MAGE_FLOW;
+}
+
 static inline bool sd_version_uses_flux_vae(SDVersion version) {
     if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) {
         return true;
@@ -283,7 +288,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_pid(version) ||
         sd_version_is_ideogram4(version) ||
         sd_version_is_sefi_image(version) ||
-        sd_version_is_krea2(version)) {
+        sd_version_is_krea2(version) ||
+        sd_version_is_mage_flow(version)) {
         return true;
     }
     return false;

src/model/common/rope.hpp

@@ -654,6 +654,43 @@ namespace Rope {
         return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
     }
 
+    __STATIC_INLINE__ std::vector<float> gen_mage_flow_pe(int h,
+                                                          int w,
+                                                          int bs,
+                                                          int context_len,
+                                                          const std::vector<ggml_tensor*>& ref_latents,
+                                                          int theta,
+                                                          const std::vector<int>& axes_dim) {
+        const int axes_dim_num = static_cast<int>(axes_dim.size());
+        auto make_image_ids    = [=](int image_h, int image_w, int image_index) {
+            std::vector<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w,
+                                                         std::vector<float>(axes_dim_num, 0.f));
+            int h_start = -(image_h - image_h / 2);
+            int w_start = -(image_w - image_w / 2);
+            for (int b = 0; b < bs; ++b) {
+                for (int y = 0; y < image_h; ++y) {
+                    for (int x = 0; x < image_w; ++x) {
+                        auto& id = image_ids[static_cast<size_t>(b) * image_h * image_w + y * image_w + x];
+                        id[0]    = static_cast<float>(image_index);
+                        id[1]    = static_cast<float>(h_start + y);
+                        id[2]    = static_cast<float>(w_start + x);
+                    }
+                }
+            }
+            return image_ids;
+        };
+        auto ids     = gen_flux_txt_ids(bs, context_len, axes_dim_num, {});
+        auto img_ids = make_image_ids(h, w, 0);
+        ids          = concat_ids(ids, img_ids, bs);
+        for (size_t i = 0; i < ref_latents.size(); ++i) {
+            auto ref_ids = make_image_ids(static_cast<int>(ref_latents[i]->ne[1]),
+                                          static_cast<int>(ref_latents[i]->ne[0]),
+                                          static_cast<int>(i + 1));
+            ids          = concat_ids(ids, ref_ids, bs);
+        }
+        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
+    }
+
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
                                                                    int w,
                                                                    int bs,