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 @@ public:
         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,

src/model/diffusion/mage_flow.hpp

@@ -0,0 +1,162 @@
+#ifndef __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__
+#define __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__
+
+#include <cmath>
+#include <memory>
+
+#include "model/diffusion/qwen_image.hpp"
+
+namespace MageFlow {
+    constexpr int MAGE_FLOW_GRAPH_SIZE = 20480;
+
+    // Mage-Flow was trained with BF16-rounded timestep frequencies; using Qwen's F32 projection degrades generation quality.
+    struct MageFlowTimestepProjEmbeddings : public Qwen::QwenTimestepProjEmbeddings {
+        static constexpr int TIMESTEP_DIM = 256;
+        static constexpr int HALF_DIM     = TIMESTEP_DIM / 2;
+
+        std::vector<float> frequencies;
+        std::vector<float> timesteps_proj;
+
+        explicit MageFlowTimestepProjEmbeddings(int64_t embedding_dim)
+            : QwenTimestepProjEmbeddings(embedding_dim), frequencies(HALF_DIM) {
+            for (int i = 0; i < HALF_DIM; ++i) {
+                float frequency = std::exp(-std::log(10000.f) * static_cast<float>(i) / HALF_DIM);
+                frequencies[i]  = ggml_bf16_to_fp32(ggml_fp32_to_bf16(frequency));
+            }
+        }
+
+        void prepare(const sd::Tensor<float>& timesteps) {
+            size_t num_timesteps = static_cast<size_t>(timesteps.numel());
+            timesteps_proj.resize(static_cast<size_t>(TIMESTEP_DIM) * num_timesteps);
+            for (size_t b = 0; b < num_timesteps; ++b) {
+                float sigma = ggml_bf16_to_fp32(ggml_fp32_to_bf16(timesteps.values()[b] / 1000.f));
+                for (int i = 0; i < HALF_DIM; ++i) {
+                    float argument = sigma * frequencies[i] * 1000.f;
+                    timesteps_proj[b * TIMESTEP_DIM + i] =
+                        ggml_bf16_to_fp32(ggml_fp32_to_bf16(std::cos(argument)));
+                    timesteps_proj[b * TIMESTEP_DIM + HALF_DIM + i] =
+                        ggml_bf16_to_fp32(ggml_fp32_to_bf16(std::sin(argument)));
+                }
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* timesteps,
+                             ggml_tensor* addition_t_cond = nullptr) override {
+            GGML_ASSERT(addition_t_cond == nullptr);
+            GGML_ASSERT(timesteps_proj.size() ==
+                        static_cast<size_t>(TIMESTEP_DIM * ggml_nelements(timesteps)));
+            auto projection = ggml_new_tensor_2d(ctx->ggml_ctx,
+                                                 GGML_TYPE_F32,
+                                                 TIMESTEP_DIM,
+                                                 ggml_nelements(timesteps));
+            ctx->bind_backend_tensor_data(projection, timesteps_proj.data());
+            auto timestep_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["timestep_embedder"]);
+            return timestep_embedder->forward(ctx, projection);
+        }
+    };
+
+    struct MageFlowRunner : public DiffusionModelRunner {
+    public:
+        Qwen::QwenImageConfig config;
+        Qwen::QwenImageModel mage_flow;
+        std::shared_ptr<MageFlowTimestepProjEmbeddings> time_text_embed;
+        std::vector<float> pe_vec;
+
+        MageFlowRunner(ggml_backend_t backend,
+                       const String2TensorStorage& tensor_storage_map      = {},
+                       const std::string prefix                            = "",
+                       std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager) {
+            config.patch_size          = 1;
+            config.in_channels         = 128;
+            config.out_channels        = 128;
+            config.num_layers          = 12;
+            config.attention_head_dim  = 128;
+            config.num_attention_heads = 24;
+            config.joint_attention_dim = 2560;
+            config.theta               = 10000;
+            config.axes_dim            = {16, 56, 56};
+            config.axes_dim_sum        = 128;
+            time_text_embed            = std::make_shared<MageFlowTimestepProjEmbeddings>(
+                config.num_attention_heads * config.attention_head_dim);
+            mage_flow = Qwen::QwenImageModel(config, time_text_embed);
+            mage_flow.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "mage_flow";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            mage_flow.get_param_tensors(tensors, prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor,
+                                 const std::vector<sd::Tensor<float>>& ref_latents_tensor = {}) {
+            ggml_cgraph* gf        = new_graph_custom(MAGE_FLOW_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
+            GGML_ASSERT(x->ne[3] == 1);
+            GGML_ASSERT(!context_tensor.empty());
+            ggml_tensor* context = make_input(context_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));
+            }
+
+            int batch_size = static_cast<int>(x->ne[3]);
+            pe_vec         = Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
+                                                    static_cast<int>(x->ne[0]),
+                                                    batch_size,
+                                                    static_cast<int>(context->ne[1]),
+                                                    ref_latents,
+                                                    config.theta,
+                                                    config.axes_dim);
+            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());
+
+            time_text_embed->prepare(timesteps_tensor);
+            auto runner_ctx = get_context();
+            auto out        = mage_flow.forward(&runner_ctx,
+                                                x,
+                                                timesteps,
+                                                nullptr,
+                                                context,
+                                                pe,
+                                                ref_latents);
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context,
+                                  const std::vector<sd::Tensor<float>>& ref_latents = {}) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context, ref_latents);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            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),
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents);
+        }
+    };
+}  // namespace MageFlow
+
+#endif  // __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__

src/model/diffusion/model.hpp

@@ -26,6 +26,7 @@ struct RefImageParams {
     RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA;
     int vlm_min_size                   = -1;
     int vlm_max_size                   = -1;
+    bool resize_vae_to_target          = false;
 };
 
 const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
@@ -34,6 +35,7 @@ const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
     {"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}},
+    {"mage_flow", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, -1, 384, true}},
     {"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}},

src/model/diffusion/qwen_image.hpp

@@ -103,9 +103,9 @@ namespace Qwen {
             }
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* timesteps,
-                             ggml_tensor* addition_t_cond = nullptr) {
+        virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                     ggml_tensor* timesteps,
+                                     ggml_tensor* addition_t_cond = nullptr) {
             // timesteps: [N,]
             // return: [N, embedding_dim]
             auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
@@ -416,10 +416,14 @@ namespace Qwen {
 
     public:
         QwenImageModel() {}
-        QwenImageModel(QwenImageConfig config)
+        QwenImageModel(QwenImageConfig config,
+                       std::shared_ptr<QwenTimestepProjEmbeddings> time_text_embed = nullptr)
             : config(config) {
-            int64_t inner_dim         = config.num_attention_heads * config.attention_head_dim;
-            blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond));
+            int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
+            if (time_text_embed == nullptr) {
+                time_text_embed = std::make_shared<QwenTimestepProjEmbeddings>(inner_dim, config.use_additional_t_cond);
+            }
+            blocks["time_text_embed"] = std::move(time_text_embed);
             blocks["txt_norm"]        = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
             blocks["img_in"]          = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
             blocks["txt_in"]          = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));

src/model/vae/mage_vae.hpp

@@ -0,0 +1,521 @@
+#ifndef __SD_MODEL_VAE_MAGE_VAE_HPP__
+#define __SD_MODEL_VAE_MAGE_VAE_HPP__
+
+#include "model/diffusion/dit.hpp"
+#include "model/vae/vae.hpp"
+
+namespace MageVAE {
+    constexpr int MAGE_VAE_GRAPH_SIZE = 327680;
+    constexpr int HIDDEN_SIZE         = 384;
+    constexpr int LATENT_CHANNELS     = 128;
+    constexpr int PATCH_SIZE          = 16;
+
+    struct LayerNorm2d : public UnaryBlock {
+        int64_t channels;
+        bool affine;
+        std::string prefix;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            this->prefix = prefix;
+            if (affine) {
+                params["weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
+                params["bias"]   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
+            }
+        }
+
+        LayerNorm2d(int64_t channels, bool affine = true)
+            : channels(channels), affine(affine) {}
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            ggml_tensor* weight = affine ? params["weight"] : nullptr;
+            ggml_tensor* bias   = affine ? params["bias"] : nullptr;
+            if (affine && ctx->weight_adapter) {
+                weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
+                bias   = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias");
+            }
+            // [N, C, H, W] -> [N, H, W, C] so layer norm reduces over channels.
+            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
+            x = ggml_ext_layer_norm(ctx->ggml_ctx, x, weight, bias, 1e-6f);
+            return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+        }
+    };
+
+    inline ggml_tensor* modulate_2d(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    ggml_tensor* shift,
+                                    ggml_tensor* scale) {
+        shift = ggml_reshape_4d(ctx, shift, 1, 1, shift->ne[0], shift->ne[1]);
+        scale = ggml_reshape_4d(ctx, scale, 1, 1, scale->ne[0], scale->ne[1]);
+        return ggml_add(ctx, ggml_mul(ctx, x, ggml_add(ctx, scale, ggml_ext_ones(ctx, 1, 1, 1, 1))), shift);
+    }
+
+    inline ggml_tensor* channel_attention(GGMLRunnerContext* ctx,
+                                          ggml_tensor* x,
+                                          Conv2d* projection) {
+        auto pooled = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
+        pooled      = ggml_mean(ctx->ggml_ctx, pooled);
+        pooled      = ggml_reshape_4d(ctx->ggml_ctx, pooled, 1, 1, x->ne[2], x->ne[3]);
+        pooled      = ggml_sigmoid(ctx->ggml_ctx, projection->forward(ctx, pooled));
+        return ggml_mul(ctx->ggml_ctx, x, pooled);
+    }
+
+    struct TimestepEmbedder : public GGMLBlock {
+        TimestepEmbedder() {
+            blocks["mlp.0"] = std::make_shared<Linear>(256, HIDDEN_SIZE);
+            blocks["mlp.2"] = std::make_shared<Linear>(HIDDEN_SIZE, HIDDEN_SIZE);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timestep) {
+            auto linear_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
+            auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
+            auto x        = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f);
+            x             = linear_0->forward(ctx, x);
+            x             = ggml_silu_inplace(ctx->ggml_ctx, x);
+            return linear_2->forward(ctx, x);
+        }
+    };
+
+    struct EncoderDiCoBlock : public UnaryBlock {
+        explicit EncoderDiCoBlock(int64_t channels) {
+            blocks["conv1"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["conv2"] = std::make_shared<Conv2d_grouped>(channels, channels, static_cast<int>(channels), std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+            blocks["conv3"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["ca.1"]  = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["conv4"] = std::make_shared<Conv2d>(channels, channels * 4, std::pair{1, 1});
+            blocks["conv5"] = std::make_shared<Conv2d>(channels * 4, channels, std::pair{1, 1});
+            blocks["norm1"] = std::make_shared<LayerNorm2d>(channels);
+            blocks["norm2"] = std::make_shared<LayerNorm2d>(channels);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
+            auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
+            auto conv2 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv2"]);
+            auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
+            auto ca    = std::dynamic_pointer_cast<Conv2d>(blocks["ca.1"]);
+            auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
+            auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
+            auto norm1 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm1"]);
+            auto norm2 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm2"]);
+
+            auto x = norm1->forward(ctx, input);
+            x      = conv1->forward(ctx, x);
+            x      = conv2->forward(ctx, x);
+            x      = ggml_gelu(ctx->ggml_ctx, x);
+            x      = channel_attention(ctx, x, ca.get());
+            x      = conv3->forward(ctx, x);
+            x      = ggml_add(ctx->ggml_ctx, input, x);
+            auto h = norm2->forward(ctx, x);
+            h      = conv4->forward(ctx, h);
+            h      = ggml_gelu(ctx->ggml_ctx, h);
+            h      = conv5->forward(ctx, h);
+            return ggml_add(ctx->ggml_ctx, x, h);
+        }
+    };
+
+    struct DiCoBlock : public GGMLBlock {
+        explicit DiCoBlock(int64_t channels) {
+            blocks["conv1"]              = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["conv2"]              = std::make_shared<Conv2d_grouped>(channels, channels, static_cast<int>(channels), std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+            blocks["conv3"]              = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["ca.1"]               = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["conv4"]              = std::make_shared<Conv2d>(channels, channels * 4, std::pair{1, 1});
+            blocks["conv5"]              = std::make_shared<Conv2d>(channels * 4, channels, std::pair{1, 1});
+            blocks["norm1"]              = std::make_shared<LayerNorm2d>(channels, false);
+            blocks["norm2"]              = std::make_shared<LayerNorm2d>(channels, false);
+            blocks["adaLN_modulation.1"] = std::make_shared<Linear>(channels, channels * 6);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input, ggml_tensor* condition) {
+            auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
+            auto conv2 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv2"]);
+            auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
+            auto ca    = std::dynamic_pointer_cast<Conv2d>(blocks["ca.1"]);
+            auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
+            auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
+            auto norm1 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm1"]);
+            auto norm2 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm2"]);
+            auto ada   = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
+
+            auto params = ada->forward(ctx, ggml_silu(ctx->ggml_ctx, condition));
+            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 6, 0);
+            auto x      = norm1->forward(ctx, input);
+            x           = modulate_2d(ctx->ggml_ctx, x, chunks[0], chunks[1]);
+            x           = conv1->forward(ctx, x);
+            x           = conv2->forward(ctx, x);
+            x           = ggml_gelu(ctx->ggml_ctx, x);
+            x           = channel_attention(ctx, x, ca.get());
+            x           = conv3->forward(ctx, x);
+            auto gate_1 = ggml_reshape_4d(ctx->ggml_ctx, chunks[2], 1, 1, chunks[2]->ne[0], chunks[2]->ne[1]);
+            x           = ggml_add(ctx->ggml_ctx, input, ggml_mul(ctx->ggml_ctx, x, gate_1));
+
+            auto h      = norm2->forward(ctx, x);
+            h           = modulate_2d(ctx->ggml_ctx, h, chunks[3], chunks[4]);
+            h           = conv4->forward(ctx, h);
+            h           = ggml_gelu(ctx->ggml_ctx, h);
+            h           = conv5->forward(ctx, h);
+            auto gate_2 = ggml_reshape_4d(ctx->ggml_ctx, chunks[5], 1, 1, chunks[5]->ne[0], chunks[5]->ne[1]);
+            return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, h, gate_2));
+        }
+    };
+
+    struct MageResnetBlock : public UnaryBlock {
+        explicit MageResnetBlock(int64_t channels) {
+            blocks["norm1"] = std::make_shared<GroupNorm32>(channels);
+            blocks["conv1"] = std::make_shared<Conv2d>(channels, channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+            blocks["norm2"] = std::make_shared<GroupNorm32>(channels);
+            blocks["conv2"] = std::make_shared<Conv2d>(channels, channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
+            auto norm1 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm1"]);
+            auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
+            auto norm2 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm2"]);
+            auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv2"]);
+            auto x     = conv1->forward(ctx, ggml_silu(ctx->ggml_ctx, norm1->forward(ctx, input)));
+            x          = conv2->forward(ctx, ggml_silu(ctx->ggml_ctx, norm2->forward(ctx, x)));
+            return ggml_add(ctx->ggml_ctx, input, x);
+        }
+    };
+
+    inline ggml_tensor* replicate_pad_right_bottom(ggml_context* ctx,
+                                                   ggml_tensor* x,
+                                                   int pad_w,
+                                                   int pad_h) {
+        if (pad_w > 0) {
+            auto edge = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
+            edge      = ggml_repeat_4d(ctx, edge, pad_w, x->ne[1], x->ne[2], x->ne[3]);
+            x         = ggml_concat(ctx, x, edge, 0);
+        }
+        if (pad_h > 0) {
+            auto edge = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]);
+            edge      = ggml_repeat_4d(ctx, edge, x->ne[0], pad_h, x->ne[2], x->ne[3]);
+            x         = ggml_concat(ctx, x, edge, 1);
+        }
+        return x;
+    }
+
+    struct MageAttnBlock : public UnaryBlock {
+        int64_t channels;
+        int patch_size;
+
+        MageAttnBlock(int64_t channels, int patch_size = 32)
+            : channels(channels), patch_size(patch_size) {
+            blocks["norm"]     = std::make_shared<GroupNorm32>(channels);
+            blocks["q"]        = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["k"]        = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["v"]        = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["proj_out"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+        }
+
+        ggml_tensor* to_patches(ggml_context* ctx, ggml_tensor* x) {
+            x = DiT::patchify(ctx, x, patch_size, patch_size);
+            x = ggml_reshape_4d(ctx, x, patch_size * patch_size, channels, x->ne[1], x->ne[2]);
+            // [N, np, C, P] -> [N, np, P, C] for attention over P pixels.
+            x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
+            return ggml_reshape_3d(ctx, x, channels, patch_size * patch_size, x->ne[2] * x->ne[3]);
+        }
+
+        ggml_tensor* from_patches(ggml_context* ctx,
+                                  ggml_tensor* x,
+                                  int64_t patch_count,
+                                  int64_t batch_size,
+                                  int64_t h_patches,
+                                  int64_t w_patches) {
+            x = ggml_reshape_4d(ctx, x, channels, patch_size * patch_size, patch_count, batch_size);
+            // [N, np, P, C] -> [N, np, C, P] before spatial unpatchify.
+            x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
+            x = ggml_reshape_3d(ctx, x, patch_size * patch_size * channels, patch_count, batch_size);
+            return DiT::unpatchify(ctx, x, h_patches, w_patches, patch_size, patch_size);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
+            auto norm     = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
+            auto q_proj   = std::dynamic_pointer_cast<Conv2d>(blocks["q"]);
+            auto k_proj   = std::dynamic_pointer_cast<Conv2d>(blocks["k"]);
+            auto v_proj   = std::dynamic_pointer_cast<Conv2d>(blocks["v"]);
+            auto proj_out = std::dynamic_pointer_cast<Conv2d>(blocks["proj_out"]);
+
+            int64_t width  = input->ne[0];
+            int64_t height = input->ne[1];
+            int64_t batch  = input->ne[3];
+            int pad_w      = (patch_size - static_cast<int>(width % patch_size)) % patch_size;
+            int pad_h      = (patch_size - static_cast<int>(height % patch_size)) % patch_size;
+            int64_t wp     = (width + pad_w) / patch_size;
+            int64_t hp     = (height + pad_h) / patch_size;
+            int64_t np     = wp * hp;
+
+            auto h = norm->forward(ctx, input);
+            auto q = replicate_pad_right_bottom(ctx->ggml_ctx, q_proj->forward(ctx, h), pad_w, pad_h);
+            auto k = replicate_pad_right_bottom(ctx->ggml_ctx, k_proj->forward(ctx, h), pad_w, pad_h);
+            auto v = replicate_pad_right_bottom(ctx->ggml_ctx, v_proj->forward(ctx, h), pad_w, pad_h);
+            q      = to_patches(ctx->ggml_ctx, q);
+            k      = to_patches(ctx->ggml_ctx, k);
+            v      = to_patches(ctx->ggml_ctx, v);
+            h      = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
+            h      = from_patches(ctx->ggml_ctx, h, np, batch, hp, wp);
+            if (pad_h > 0) {
+                h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, height);
+            }
+            if (pad_w > 0) {
+                h = ggml_ext_slice(ctx->ggml_ctx, h, 0, 0, width);
+            }
+            return ggml_add(ctx->ggml_ctx, input, proj_out->forward(ctx, h));
+        }
+    };
+
+    struct Decoder : public UnaryBlock {
+        Decoder() {
+            blocks["conv_in"]  = std::make_shared<Conv2d>(LATENT_CHANNELS, HIDDEN_SIZE, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+            blocks["block.0"]  = std::make_shared<MageResnetBlock>(HIDDEN_SIZE);
+            blocks["block.1"]  = std::make_shared<MageAttnBlock>(HIDDEN_SIZE);
+            blocks["block.2"]  = std::make_shared<MageResnetBlock>(HIDDEN_SIZE);
+            blocks["block.3"]  = std::make_shared<MageAttnBlock>(HIDDEN_SIZE);
+            blocks["block.4"]  = std::make_shared<MageResnetBlock>(HIDDEN_SIZE);
+            blocks["norm_out"] = std::make_shared<GroupNorm32>(HIDDEN_SIZE);
+            blocks["conv_out"] = std::make_shared<Conv2d>(HIDDEN_SIZE, HIDDEN_SIZE, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            x = std::dynamic_pointer_cast<Conv2d>(blocks["conv_in"])->forward(ctx, x);
+            for (int i = 0; i < 5; ++i) {
+                x = std::dynamic_pointer_cast<UnaryBlock>(blocks["block." + std::to_string(i)])->forward(ctx, x);
+            }
+            x = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm_out"])->forward(ctx, x);
+            x = ggml_silu(ctx->ggml_ctx, x);
+            return std::dynamic_pointer_cast<Conv2d>(blocks["conv_out"])->forward(ctx, x);
+        }
+    };
+
+    struct DConvEncoder : public UnaryBlock {
+        DConvEncoder() {
+            blocks["patch_cond_embed"] = std::make_shared<Conv2d>(3, 768, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{PATCH_SIZE, PATCH_SIZE});
+            for (int i = 0; i < 2; ++i) {
+                blocks["head_blocks." + std::to_string(i)] = std::make_shared<EncoderDiCoBlock>(768);
+            }
+            blocks["proj_down"]  = std::make_shared<Conv2d>(768, HIDDEN_SIZE, std::pair{1, 1});
+            blocks["z_proj"]     = std::make_shared<Conv2d>(LATENT_CHANNELS, HIDDEN_SIZE, std::pair{1, 1});
+            blocks["fuse_proj"]  = std::make_shared<Conv2d>(HIDDEN_SIZE * 2, HIDDEN_SIZE, std::pair{1, 1});
+            blocks["t_embedder"] = std::make_shared<TimestepEmbedder>();
+            for (int i = 0; i < 21; ++i) {
+                blocks["blocks." + std::to_string(i)] = std::make_shared<DiCoBlock>(HIDDEN_SIZE);
+            }
+            blocks["norm_out"] = std::make_shared<LayerNorm2d>(HIDDEN_SIZE);
+            blocks["proj_out"] = std::make_shared<Conv2d>(HIDDEN_SIZE, LATENT_CHANNELS * 2, std::pair{1, 1});
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image) override {
+            auto cond = std::dynamic_pointer_cast<Conv2d>(blocks["patch_cond_embed"])->forward(ctx, image);
+            for (int i = 0; i < 2; ++i) {
+                cond = std::dynamic_pointer_cast<EncoderDiCoBlock>(blocks["head_blocks." + std::to_string(i)])->forward(ctx, cond);
+            }
+            cond   = std::dynamic_pointer_cast<Conv2d>(blocks["proj_down"])->forward(ctx, cond);
+            auto z = ggml_ext_zeros(ctx->ggml_ctx, cond->ne[0], cond->ne[1], LATENT_CHANNELS, cond->ne[3]);
+            z      = std::dynamic_pointer_cast<Conv2d>(blocks["z_proj"])->forward(ctx, z);
+            z      = ggml_concat(ctx->ggml_ctx, cond, z, 2);
+            z      = std::dynamic_pointer_cast<Conv2d>(blocks["fuse_proj"])->forward(ctx, z);
+            auto t = ggml_ext_zeros(ctx->ggml_ctx, image->ne[3], 1, 1, 1);
+            t      = ggml_reshape_1d(ctx->ggml_ctx, t, image->ne[3]);
+            auto c = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
+            for (int i = 0; i < 21; ++i) {
+                z = std::dynamic_pointer_cast<DiCoBlock>(blocks["blocks." + std::to_string(i)])->forward(ctx, z, c);
+            }
+            z = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm_out"])->forward(ctx, z);
+            return std::dynamic_pointer_cast<Conv2d>(blocks["proj_out"])->forward(ctx, z);
+        }
+    };
+
+    struct MLPResBlock : public GGMLBlock {
+        MLPResBlock() {
+            blocks["in_ln"]              = std::make_shared<LayerNorm>(32, 1e-6f);
+            blocks["mlp.0"]              = std::make_shared<Linear>(32, 32);
+            blocks["mlp.2"]              = std::make_shared<Linear>(32, 32);
+            blocks["adaLN_modulation.1"] = std::make_shared<Linear>(32, 96);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* condition) {
+            auto params = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"])->forward(ctx, ggml_silu(ctx->ggml_ctx, condition));
+            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 3, 0);
+            auto h      = std::dynamic_pointer_cast<LayerNorm>(blocks["in_ln"])->forward(ctx, x);
+            h           = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, h, ggml_add(ctx->ggml_ctx, chunks[1], ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1))), chunks[0]);
+            h           = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"])->forward(ctx, h);
+            h           = ggml_silu(ctx->ggml_ctx, h);
+            h           = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"])->forward(ctx, h);
+            return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, chunks[2], h));
+        }
+    };
+
+    struct DConvDenoiser : public GGMLBlock {
+        DConvDenoiser() {
+            blocks["t_embedder"]            = std::make_shared<TimestepEmbedder>();
+            blocks["y_embedder_x"]          = std::make_shared<Conv2d>(HIDDEN_SIZE, 32 * PATCH_SIZE * PATCH_SIZE, std::pair{1, 1});
+            blocks["x_embedder.embedder.0"] = std::make_shared<Linear>(3 + 32 + 64, 32);
+            blocks["s_embedder.proj1"]      = std::make_shared<Conv2d>(3, LATENT_CHANNELS, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{0, 0}, std::pair{1, 1}, false);
+            blocks["s_embedder.proj2"]      = std::make_shared<Conv2d>(LATENT_CHANNELS + HIDDEN_SIZE, HIDDEN_SIZE, std::pair{1, 1});
+            for (int i = 0; i < 21; ++i) {
+                blocks["blocks." + std::to_string(i)] = std::make_shared<DiCoBlock>(HIDDEN_SIZE);
+            }
+            blocks["dec_net.cond_embed"] = std::make_shared<Linear>(HIDDEN_SIZE, PATCH_SIZE * PATCH_SIZE * 32);
+            blocks["dec_net.input_proj"] = std::make_shared<Linear>(32, 32);
+            for (int i = 0; i < 3; ++i) {
+                blocks["dec_net.res_blocks." + std::to_string(i)] = std::make_shared<MLPResBlock>();
+            }
+            blocks["final_layer.norm"]   = std::make_shared<RMSNorm>(32);
+            blocks["final_layer.linear"] = std::make_shared<Linear>(32, 3);
+            blocks["y_embedder.decoder"] = std::make_shared<Decoder>();
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* dct) {
+            auto cond      = std::dynamic_pointer_cast<Decoder>(blocks["y_embedder.decoder"])->forward(ctx, latent);
+            int64_t w      = cond->ne[0];
+            int64_t h      = cond->ne[1];
+            int64_t n      = cond->ne[3];
+            int64_t length = w * h;
+
+            auto image = ggml_ext_zeros(ctx->ggml_ctx, w * PATCH_SIZE, h * PATCH_SIZE, 3, n);
+            auto t     = ggml_ext_zeros(ctx->ggml_ctx, n, 1, 1, 1);
+            t          = ggml_reshape_1d(ctx->ggml_ctx, t, n);
+            auto c     = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
+
+            auto s0 = std::dynamic_pointer_cast<Conv2d>(blocks["s_embedder.proj1"])->forward(ctx, image);
+            s0      = ggml_concat(ctx->ggml_ctx, s0, cond, 2);
+            auto s  = std::dynamic_pointer_cast<Conv2d>(blocks["s_embedder.proj2"])->forward(ctx, s0);
+            for (int i = 0; i < 21; ++i) {
+                s = std::dynamic_pointer_cast<DiCoBlock>(blocks["blocks." + std::to_string(i)])->forward(ctx, s, c);
+            }
+            // [N, C, H, W] -> [N*H*W, C].
+            s = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 2, 0, 3));
+            s = ggml_reshape_2d(ctx->ggml_ctx, s, HIDDEN_SIZE, length * n);
+
+            auto y = std::dynamic_pointer_cast<Conv2d>(blocks["y_embedder_x"])->forward(ctx, cond);
+            // Split 32*P channels as [32, P], then produce [N*L, P, 32].
+            y          = ggml_reshape_4d(ctx->ggml_ctx, y, length, PATCH_SIZE * PATCH_SIZE, 32, n);
+            y          = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, y, 2, 1, 0, 3));
+            y          = ggml_reshape_3d(ctx->ggml_ctx, y, 32, PATCH_SIZE * PATCH_SIZE, length * n);
+            auto zeros = ggml_ext_zeros(ctx->ggml_ctx, 3, PATCH_SIZE * PATCH_SIZE, length * n, 1);
+            dct        = ggml_repeat_4d(ctx->ggml_ctx, dct, 64, PATCH_SIZE * PATCH_SIZE, length * n, 1);
+            auto x     = ggml_concat(ctx->ggml_ctx, zeros, y, 0);
+            x          = ggml_concat(ctx->ggml_ctx, x, dct, 0);
+            x          = std::dynamic_pointer_cast<Linear>(blocks["x_embedder.embedder.0"])->forward(ctx, x);
+            x          = std::dynamic_pointer_cast<Linear>(blocks["dec_net.input_proj"])->forward(ctx, x);
+
+            auto dec_cond = std::dynamic_pointer_cast<Linear>(blocks["dec_net.cond_embed"])->forward(ctx, s);
+            dec_cond      = ggml_reshape_3d(ctx->ggml_ctx, dec_cond, 32, PATCH_SIZE * PATCH_SIZE, length * n);
+            for (int i = 0; i < 3; ++i) {
+                x = std::dynamic_pointer_cast<MLPResBlock>(blocks["dec_net.res_blocks." + std::to_string(i)])->forward(ctx, x, dec_cond);
+            }
+            x = std::dynamic_pointer_cast<RMSNorm>(blocks["final_layer.norm"])->forward(ctx, x);
+            x = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, x);
+            // [N*L, P, 3] -> [N, L, 3*P] for fold/unpatchify.
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, 3, PATCH_SIZE * PATCH_SIZE, length, n);
+            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x = ggml_reshape_3d(ctx->ggml_ctx, x, 3 * PATCH_SIZE * PATCH_SIZE, length, n);
+            return DiT::unpatchify(ctx->ggml_ctx, x, h, w, PATCH_SIZE, PATCH_SIZE);
+        }
+    };
+
+    struct MageVAEModel : public GGMLBlock {
+        MageVAEModel() {
+            blocks["student.dconv_encoder"] = std::make_shared<DConvEncoder>();
+            blocks["pipeline"]              = std::make_shared<DConvDenoiser>();
+        }
+
+        ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* image) {
+            return std::dynamic_pointer_cast<DConvEncoder>(blocks["student.dconv_encoder"])->forward(ctx, image);
+        }
+
+        ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* dct) {
+            return std::dynamic_pointer_cast<DConvDenoiser>(blocks["pipeline"])->forward(ctx, latent, dct);
+        }
+    };
+
+    struct MageVAERunner : public VAE {
+        MageVAEModel model;
+        std::vector<float> dct_vec;
+
+        MageVAERunner(ggml_backend_t backend,
+                      const String2TensorStorage& tensor_storage_map,
+                      const std::string& prefix,
+                      std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : VAE(VERSION_MAGE_FLOW, backend, prefix, weight_manager) {
+            model = MageVAEModel();
+            model.init(params_ctx, tensor_storage_map, prefix);
+            dct_vec.resize(64 * PATCH_SIZE * PATCH_SIZE);
+            constexpr float pi = 3.14159265358979323846f;
+            for (int py = 0; py < PATCH_SIZE; ++py) {
+                float y = static_cast<float>(py) / static_cast<float>(PATCH_SIZE - 1);
+                for (int px = 0; px < PATCH_SIZE; ++px) {
+                    float x = static_cast<float>(px) / static_cast<float>(PATCH_SIZE - 1);
+                    int pos = py * PATCH_SIZE + px;
+                    for (int fy = 0; fy < 8; ++fy) {
+                        for (int fx = 0; fx < 8; ++fx) {
+                            int freq                 = fx * 8 + fy;
+                            float freq_x             = static_cast<float>(fx) * 8.f / 7.f;
+                            float freq_y             = static_cast<float>(fy) * 8.f / 7.f;
+                            float coeff              = 1.f / (1.f + freq_x * freq_y);
+                            dct_vec[freq + 64 * pos] = std::cos(x * freq_x * pi) *
+                                                       std::cos(y * freq_y * pi) * coeff;
+                        }
+                    }
+                }
+            }
+        }
+
+        std::string get_desc() override {
+            return "mage_vae";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+            model.get_param_tensors(tensors, weight_prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& input_tensor, bool decode_graph) {
+            ggml_cgraph* gf  = new_graph_custom(MAGE_VAE_GRAPH_SIZE);
+            auto input       = make_input(input_tensor);
+            auto runner_ctx  = get_context();
+            ggml_tensor* dct = nullptr;
+            if (decode_graph) {
+                dct = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, 64, PATCH_SIZE * PATCH_SIZE, 1);
+                set_backend_tensor_data(dct, dct_vec.data());
+            }
+            auto out = decode_graph ? model.decode(&runner_ctx, input, dct) : model.encode(&runner_ctx, input);
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> _compute(const int n_threads,
+                                   const sd::Tensor<float>& input,
+                                   bool decode_graph) override {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(input, decode_graph);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), input.dim());
+        }
+
+        int get_encoder_output_channels(int input_channels) override {
+            SD_UNUSED(input_channels);
+            return LATENT_CHANNELS * 2;
+        }
+
+        sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
+            const auto chunks         = sd::ops::chunk(vae_output, 2, 2);
+            const auto& mean          = chunks[0];
+            const auto& logvar        = chunks[1];
+            sd::Tensor<float> stddev  = sd::ops::exp(0.5f * sd::ops::clamp(logvar, -20.0f, 10.0f));
+            sd::Tensor<float> noise   = sd::Tensor<float>::randn_like(mean, rng);
+            sd::Tensor<float> latents = mean + stddev * noise;
+            return latents;
+        }
+
+        sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
+            return latents;
+        }
+
+        sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
+            return latents;
+        }
+    };
+}  // namespace MageVAE
+
+#endif  // __SD_MODEL_VAE_MAGE_VAE_HPP__

src/model/vae/vae.hpp

@@ -74,7 +74,7 @@ public:
         int scale_factor = 8;
         if (version == VERSION_LTXAV) {
             scale_factor = 32;
-        } else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version)) {
+        } else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version)) {
             scale_factor = 16;
         } else if (sd_version_uses_flux2_vae(version)) {
             scale_factor = 16;

src/model_loader.cpp

@@ -498,6 +498,10 @@ SDVersion ModelLoader::get_sd_version() {
             return VERSION_MINIT2I;
         }
         if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
+            auto img_in = tensor_storage_map.find("model.diffusion_model.img_in.weight");
+            if (img_in != tensor_storage_map.end() && img_in->second.ne[0] == 128) {
+                return VERSION_MAGE_FLOW;
+            }
             if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) {
                 return VERSION_QWEN_IMAGE_LAYERED;
             }

src/name_conversion.cpp

@@ -1048,7 +1048,7 @@ std::string convert_diffusers_to_original_wan_vae(std::string name) {
 }
 
 std::string convert_first_stage_model_name(std::string name, std::string prefix, SDVersion version) {
-    if (sd_version_is_hunyuan_video(version)) {
+    if (sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version)) {
         return name;
     }
     if (sd_version_uses_wan_vae(version)) {
@@ -1414,8 +1414,14 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
 
     replace_with_prefix_map(name, prefix_map);
 
-    if ((sd_version_is_boogu_image(version) || sd_version_is_krea2(version)) && starts_with(name, "text_encoders.llm.visual.")) {
-        name = convert_qwen3_vl_vision_name(std::move(name));
+    if (sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version)) {
+        const std::string hf_vision_prefix = "text_encoders.llm.model.visual.";
+        if (starts_with(name, hf_vision_prefix)) {
+            name = "text_encoders.llm.visual." + name.substr(hf_vision_prefix.size());
+        }
+        if (starts_with(name, "text_encoders.llm.visual.")) {
+            name = convert_qwen3_vl_vision_name(std::move(name));
+        }
     }
 
     // diffusion model

src/stable-diffusion.cpp

@@ -36,6 +36,7 @@
 #include "model/diffusion/lens.hpp"
 #include "model/diffusion/lingbot_video.hpp"
 #include "model/diffusion/ltxv.hpp"
+#include "model/diffusion/mage_flow.hpp"
 #include "model/diffusion/minit2i.hpp"
 #include "model/diffusion/mmdit.hpp"
 #include "model/diffusion/model.hpp"
@@ -50,6 +51,7 @@
 #include "model/vae/hunyuan_vae.hpp"
 #include "model/vae/ltx_audio_vae.hpp"
 #include "model/vae/ltx_vae.hpp"
+#include "model/vae/mage_vae.hpp"
 #include "model/vae/tae.hpp"
 #include "model/vae/vae.hpp"
 #include "model/vae/wan_vae.hpp"
@@ -115,6 +117,7 @@ const char* model_version_to_str[] = {
     "Ideogram 4",
     "SeFi-Image",
     "Krea2",
+    "Mage Flow",
     "ESRGAN",
 };
 
@@ -145,6 +148,7 @@ static bool sd_version_supports_ref_latent_img_cfg(SDVersion version) {
     return version == VERSION_FLUX ||
            sd_version_is_flux2(version) ||
            sd_version_is_qwen_image(version) ||
+           sd_version_is_mage_flow(version) ||
            sd_version_is_longcat(version) ||
            sd_version_is_z_image(version) ||
            sd_version_is_boogu_image(version);
@@ -1147,6 +1151,17 @@ public:
                                                                           version,
                                                                           model_manager,
                                                                           sd_ctx_params->model_args);
+            } else if (sd_version_is_mage_flow(version)) {
+                cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
+                                                                 tensor_storage_map,
+                                                                 version,
+                                                                 "",
+                                                                 true,
+                                                                 model_manager);
+                diffusion_model  = std::make_shared<MageFlow::MageFlowRunner>(backend_for(SDBackendModule::DIFFUSION),
+                                                                             tensor_storage_map,
+                                                                             "model.diffusion_model",
+                                                                             model_manager);
             } else if (sd_version_is_longcat(version)) {
                 cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
                                                                  tensor_storage_map,
@@ -1321,6 +1336,11 @@ public:
                                                          false,
                                                          version,
                                                          model_manager);
+                } else if (sd_version_is_mage_flow(vae_version)) {
+                    return std::make_shared<MageVAE::MageVAERunner>(backend_for(SDBackendModule::VAE),
+                                                                    tensor_storage_map,
+                                                                    "first_stage_model",
+                                                                    model_manager);
                 } else if (sd_version_uses_hunyuan_video_vae(vae_version)) {
                     return std::make_shared<Hunyuan::HunyuanVideoVAERunner>(backend_for(SDBackendModule::VAE),
                                                                             tensor_storage_map,
@@ -1642,6 +1662,7 @@ public:
                            sd_version_is_hunyuan_video(version) ||
                            sd_version_is_lingbot_video(version) ||
                            sd_version_is_qwen_image(version) ||
+                           sd_version_is_mage_flow(version) ||
                            version == VERSION_HIDREAM_O1 ||
                            sd_version_is_anima(version) ||
                            sd_version_is_ernie_image(version) ||
@@ -1662,6 +1683,8 @@ public:
                         default_flow_shift = 1.0f;
                     } else if (sd_version_is_boogu_image(version)) {
                         default_flow_shift = 3.16f;
+                    } else if (sd_version_is_mage_flow(version)) {
+                        default_flow_shift = 6.f;
                     } else {
                         default_flow_shift = 3.f;
                     }
@@ -2762,6 +2785,8 @@ public:
                 latent_channel = 144;
             } else if (sd_version_uses_flux2_vae(version)) {
                 latent_channel = 128;
+            } else if (sd_version_is_mage_flow(version)) {
+                latent_channel = 128;
             } else {
                 latent_channel = 16;
             }
@@ -2910,6 +2935,8 @@ public:
             return "qwen_layered";
         } else if (sd_version_is_qwen_image(version)) {
             return "qwen";
+        } else if (sd_version_is_mage_flow(version)) {
+            return "mage_flow";
         } 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)) {
@@ -4773,10 +4800,17 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
         sd::Tensor<float> ref_latent;
         if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd_ctx->sd->version)) {
             LOG_DEBUG("auto resize ref images");
-            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];
+            double vae_width;
+            double vae_height;
+            if (ref_image_params.resize_vae_to_target) {
+                vae_width  = request->width;
+                vae_height = request->height;
+            } else {
+                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);
+                vae_width          = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]);
+                vae_height         = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0];
+            }
 
             int factor = sd_version_is_qwen_image(sd_ctx->sd->version) ? 32 : 16;
             vae_height = round(vae_height / factor) * factor;