feat: add Qwen Image 2.1 Turbo sigma preset (#2130)

.github/workflows/build.yml

@@ -451,8 +451,8 @@ jobs:
     runs-on: windows-2022
 
     env:
-      ROCM_VERSION: "7.14.0"
-      GPU_TARGETS: "gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
+      ROCM_VERSION: "10.1.0"
+      gpu_targets: "gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
 
     steps:
       - uses: actions/checkout@v3
@@ -498,7 +498,7 @@ jobs:
 
           write-host "Installing ROCm wheels for multi-arch support"
           # Install ROCm wheels for multi-arch support (this may take several minutes)
-          python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{env.ROCM_VERSION}}"
+          python -m pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ "rocm[libraries,devel]==${{env.ROCM_VERSION}}"
 
           # Pre-expand the devel tree so it is included in the cache
           write-host "Initializing ROCm devel tree"
@@ -594,9 +594,9 @@ jobs:
     strategy:
       matrix:
         include:
-          - ROCM_VERSION: "7.14.0"
+          - ROCM_VERSION: "10.1.0"
             gpu_targets: "gfx900;gfx906;gfx908;gfx90a;gfx90c;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
-            build: x64
+            build: 'x64'
 
     steps:
       - name: Clone
@@ -610,6 +610,16 @@ jobs:
         with:
           key: ubuntu-rocm-cmake-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
           evict-old-files: 1d
+          max-size: "2G"
+
+      - name: Tune ccache for reinstalled ROCm toolchain
+        run: |
+          # ROCm is pip-installed fresh each run, so the clang binary's mtime
+          # changes every time. With the default compiler_check=mtime that
+          # invalidates the cache; hash compiler contents instead so warm
+          # builds hit.
+          ccache --set-config=compiler_check=content
+          ccache --set-config=sloppiness=time_macros,include_file_mtime,include_file_ctime
 
       - name: Dependencies
         id: depends
@@ -646,7 +656,7 @@ jobs:
           # libraries = HIP runtime and CMake configs needed for linking
           # devel = compilers, headers, static libs
           python -m pip install --upgrade pip
-          python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{matrix.ROCM_VERSION}}"
+          python -m pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ "rocm[libraries,devel]==${{matrix.ROCM_VERSION}}"
 
           # Get ROCm installation paths using the rocm-sdk CLI tool
           ROCM_PATH=$(rocm-sdk path --root)

CONTRIBUTING.md

@@ -29,6 +29,10 @@ PRs should include:
 - Linked issue/PR context when applicable.
 - Screenshots or sample outputs for UI/visual behavior changes.
 
+## README UI Listings
+
+The [UIs list in the README](README.md#uis) only includes open-source projects that use `stable-diffusion.cpp` as a backend. Do not submit PRs to add closed-source UIs.
+
 ## Code Style
 
 Format code according to the repository style before submitting changes.

README.md

@@ -67,6 +67,7 @@ API and command-line option may change frequently.***
     - [LLaDA-Image](./docs/llada_image.md)
     - [Ming-Image Design](./docs/ming_image.md)
     - [PixArt](./docs/pixart.md)
+    - [Iris-3B](./docs/iris.md)
   - [Image Edit Models](./docs/edit.md)
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)
@@ -182,7 +183,6 @@ These projects wrap `stable-diffusion.cpp` for easier use in other languages/fra
 These projects use `stable-diffusion.cpp` as a backend for their image generation.
 
 - [GIMP Plugins](https://github.com/themanyone/gimp-plugins)
-- [Jellybox](https://jellybox.com)
 - [Stable Diffusion GUI](https://github.com/fszontagh/sd.cpp.gui.wx)
 - [Stable Diffusion CLI-GUI](https://github.com/piallai/stable-diffusion.cpp)
 - [Local Diffusion](https://github.com/rmatif/Local-Diffusion)

docs/esrgan.md

@@ -2,7 +2,9 @@
 
 You can use ESRGAN—such as the model [RealESRGAN_x4plus_anime_6B.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth)—to upscale the generated images and improve their overall resolution and clarity.
 
-RGBA images, including Qwen Image 2.1 output, keep their alpha channel during model upscaling and hires fix. ESRGAN processes the RGB channels; the alpha channel is resized with bilinear interpolation and recombined with the upscaled image.
+MoSR models with the PixelShuffle upsampler, such as [4xNomos2_hq_mosr.pth](https://github.com/Phhofm/models/releases/download/4xNomos2_hq_mosr/4xNomos2_hq_mosr.pth), are also supported. The architecture is detected from the weights, so they are used the same way as ESRGAN models.
+
+RGBA images, including Qwen Image 2.1 output, keep their alpha channel during model upscaling and hires fix. The upscaler processes the RGB channels; the alpha channel is resized with bilinear interpolation and recombined with the upscaled image.
 
 - Specify the model path using the `--upscale-model PATH` parameter. example:
 

docs/iris.md

@@ -0,0 +1,24 @@
+# How to Use
+
+Iris-3B generates images directly in pixel space and uses Qwen3-VL-4B-Instruct as the text encoder. No VAE is required.
+
+## Download weights
+
+- Download Iris-3B
+    - safetensors: https://huggingface.co/speridlabs/iris-3b/tree/main (`model.safetensors` in the root directory)
+- Download Qwen3-VL-4B-Instruct
+    - safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders
+    - gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main
+
+## Examples
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\iris-3b.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf -p "a lovely cat" --cfg-scale 3 -H 1024 -W 1024 --diffusion-fa
+```
+
+<img width="256" alt="iris-3b example" src="../assets/iris/example.png" />
+
+## Notes
+
+- Width and height must be multiples of 16. Do not pass `--vae`.
+- Captions are limited to 300 tokens including the assistant-turn suffix. Positive and negative prompts use the same template.

docs/performance.md

@@ -21,6 +21,16 @@ CPU fallback. It excludes weights and cache buffers. Within a runner lifecycle,
 the summary is printed only on the first graph or when backend capacities or the
 segment count change.
 
+## Run conditional and unconditional CFG in one batched UNet forward.
+
+For UNet models, the conditional and unconditional guidance branches are
+concatenated into a single batch of two and run through one UNet forward per
+step instead of two separate forwards. This is enabled by default whenever the
+run qualifies for it.
+
+Use `--batched-cfg off` to force separate conditional and unconditional
+forwards.
+
 ## Use VAE tiling to reduce encode and decode memory usage.
 
 `--vae-tiling` enables spatial tiling for both VAE encoding and decoding. The

docs/qwen_image_2.1.md

@@ -16,6 +16,8 @@ Qwen Image 2.1 supports text-to-image generation and image editing, using Qwen3-
 
 Use `qwen_image_2.1_vae_bf16.safetensors` with this model. The earlier Qwen Image and Wan 2.2 VAE weights are not interchangeable with the Qwen Image 2.1 VAE weights.
 
+For faster encoding, decoding, or previews, see [TAEQI2.1](taesd.md#qwen-image-21-taeqi21).
+
 ## Examples
 
 Run the following commands from the build directory. Use image dimensions divisible by 32. The resolution-dependent flow schedule is selected automatically.
@@ -30,6 +32,16 @@ Run the following commands from the build directory. Use image dimensions divisi
 
 To use GGUF diffusion weights, set `--diffusion-model` to the path of a file such as `qwen_image_2.1-Q4_K.gguf`.
 
+### Turbo
+
+For [Qwen Image 2.1 Turbo](https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo), use the Turbo diffusion weights and select its fixed 8-step sigma schedule with `--sigmas qwen-image-2.1-turbo`. For example, add these sampling options to the command above:
+
+```text
+--sigmas qwen-image-2.1-turbo --steps 8 --cfg-scale 1.0 --sampling-method euler
+```
+
+The preset is equivalent to `--sigmas 1.0,0.978453,0.95418,0.926626,0.89508,0.845148,0.704534,0.414568,0.0`. It only supplies the sigma values; CFG scale and sampling method remain independently configurable. As with an explicit sigma list, the sigma count determines the actual number of sampling steps. The same preset name is also accepted by `--hires-sigmas` for the highres fix second pass.
+
 ### Image editing
 
 Pass the reference image with `-r` and describe the edit in `-p`. Vision weights are required; the example below loads them separately with `--llm_vision`.

docs/taesd.md

@@ -37,3 +37,15 @@ sd.cpp also supports [TAEHV](https://github.com/madebyollin/taehv) (#937), which
   ```
 
 Then simply replace the `--vae xxx.safetensors` with `--tae xxx.safetensors` in the commands. If it still out of VRAM, add `--vae-conv-direct` to your command though might be slower.
+
+### Qwen Image 2.1 (TAEQI2.1)
+
+For Qwen Image 2.1, use [taeqi2_1](https://github.com/madebyollin/taesd), which supports 64-channel latents, 16x spatial scaling, and RGBA images.
+
+Download the official [safetensors weights](https://huggingface.co/madebyollin/taeqi2_1/blob/main/taeqi2_1.safetensors) directly; no conversion is needed:
+
+```bash
+curl -L -o taeqi2_1.safetensors https://huggingface.co/madebyollin/taeqi2_1/resolve/main/taeqi2_1.safetensors
+```
+
+Replace `--vae PATH` with `--taesd taeqi2_1.safetensors` in the [Qwen Image 2.1 examples](qwen_image_2.1.md) to use it for encoding and decoding. For previews only, keep `--vae PATH` and add `--taesd taeqi2_1.safetensors --taesd-preview-only --preview tae`.

examples/common/common.cpp

@@ -378,6 +378,23 @@ static int parse_scale_override(int argc, const char** argv, int index, float& s
     return 1;
 }
 
+static int parse_on_off_arg(int argc, const char** argv, int index, const char* option, bool& value) {
+    if (++index >= argc) {
+        LOG_ERROR("%s requires 'on' or 'off'", option);
+        return -1;
+    }
+    const std::string arg = argv[index];
+    if (arg == "on") {
+        value = true;
+    } else if (arg == "off") {
+        value = false;
+    } else {
+        LOG_ERROR("invalid %s value '%s'; expected 'on' or 'off'", option, argv[index]);
+        return -1;
+    }
+    return 1;
+}
+
 ArgOptions SDContextParams::get_options() {
     ArgOptions options;
     options.string_options = {
@@ -637,23 +654,6 @@ ArgOptions SDContextParams::get_options() {
          true, &vae_conv_direct},
     };
 
-    auto on_auto_fit_arg = [&](int argc, const char** argv, int index) {
-        if (++index >= argc) {
-            LOG_ERROR("--auto-fit requires 'on' or 'off'");
-            return -1;
-        }
-        const std::string arg = argv[index];
-        if (arg == "on") {
-            auto_fit = true;
-        } else if (arg == "off") {
-            auto_fit = false;
-        } else {
-            LOG_ERROR("invalid --auto-fit value '%s'; expected 'on' or 'off'", argv[index]);
-            return -1;
-        }
-        return 1;
-    };
-
     auto on_type_arg = [&](int argc, const char** argv, int index) {
         if (++index >= argc) {
             return -1;
@@ -742,7 +742,15 @@ ArgOptions SDContextParams::get_options() {
          "on|off (default: on). Preserve --backend (otherwise select one GPU) and place weights on the compute GPU, "
          "RAM, another GPU, or disk in that order, according to available memory (--max-vram limits GPU budgets). "
          "Disabled by explicit --params-backend; uses automatic graph segmentation when needed",
-         on_auto_fit_arg},
+         [this](int argc, const char** argv, int index) {
+             return parse_on_off_arg(argc, argv, index, "--auto-fit", auto_fit);
+         }},
+        {"",
+         "--batched-cfg",
+         "on|off (default: on). Run the conditional and unconditional CFG branches in one batched UNet forward when supported",
+         [this](int argc, const char** argv, int index) {
+             return parse_on_off_arg(argc, argv, index, "--batched-cfg", batched_cfg);
+         }},
         {"",
          "--type",
          "weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K). "
@@ -940,6 +948,7 @@ std::string SDContextParams::to_string() const {
         << "  max_vram: \"" << max_vram << "\",\n"
         << "  disable_prefetch: " << (disable_prefetch ? "true" : "false") << ",\n"
         << "  disable_segmented_compute: " << (disable_segmented_compute ? "true" : "false") << ",\n"
+        << "  batched_cfg: " << (batched_cfg ? "true" : "false") << ",\n"
         << "  eager_load: " << (eager_load ? "true" : "false") << ",\n"
         << "  backend: \"" << backend << "\",\n"
         << "  params_backend: \"" << params_backend << "\",\n"
@@ -1022,6 +1031,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.max_vram                        = max_vram.c_str();
     sd_ctx_params.disable_prefetch                = disable_prefetch;
     sd_ctx_params.disable_segmented_compute       = disable_segmented_compute;
+    sd_ctx_params.batched_cfg                     = batched_cfg;
     sd_ctx_params.eager_load                      = eager_load;
     sd_ctx_params.backend                         = effective_backend.c_str();
     sd_ctx_params.params_backend                  = effective_params_backend.c_str();
@@ -1465,6 +1475,9 @@ ArgOptions SDGenerationParams::get_options() {
             return -1;
         }
         std::string sigmas_str = value;
+        if (sigmas_str == "qwen-image-2.1-turbo") {
+            sigmas_str = "1.0,0.978453,0.95418,0.926626,0.89508,0.845148,0.704534,0.414568,0.0";
+        }
         if (!sigmas_str.empty() && sigmas_str.front() == '[') {
             sigmas_str.erase(0, 1);
         }
@@ -1711,11 +1724,11 @@ ArgOptions SDGenerationParams::get_options() {
          on_scheduler_arg},
         {"",
          "--sigmas",
-         "custom sigma values for the sampler, comma-separated (e.g., \"14.61,7.8,3.5,0.0\").",
+         "custom sigma values for the sampler, comma-separated (e.g., \"14.61,7.8,3.5,0.0\") or preset: qwen-image-2.1-turbo.",
          on_sigmas_arg},
         {"",
          "--hires-sigmas",
-         "custom sigma values for the highres fix second pass, comma-separated (e.g., \"0.85,0.725,0.421875,0.0\").",
+         "custom sigma values for the highres fix second pass, comma-separated (e.g., \"0.85,0.725,0.421875,0.0\") or preset: qwen-image-2.1-turbo.",
          on_hires_sigmas_arg},
         {"",
          "--skip-layers",

examples/common/common.h

@@ -157,6 +157,7 @@ struct SDContextParams {
     bool disable_prefetch          = false;
     bool disable_segmented_compute = false;
     bool eager_load                = false;
+    bool batched_cfg               = true;
     std::string backend;
     std::string params_backend;
     std::string split_mode;

include/stable-diffusion.h

@@ -245,6 +245,7 @@ typedef struct {
     const char* rpc_servers;
     const char* model_args;
     bool disable_segmented_compute;  // Force monolithic graph execution even when automatic graph cutting would fit memory better
+    bool batched_cfg;                // Run the conditional and unconditional CFG branches in one batched UNet forward when supported
     float linear_scale;              // Override linear input scaling; 0 keeps the model default
     float attn_scale;                // Override flash-attention K/V scaling; 0 keeps the model default
     const char* tokenizer;           // tokenizer.json path or main=FILE,clip-l=FILE,clip-g=FILE assignments; required for PiD and Lens

src/conditioning/conditioner.hpp

@@ -2012,7 +2012,8 @@ struct LLMEmbedder : public Conditioner {
                    sd_version_is_sefi_image(version) ||
                    sd_version_is_krea2(version) ||
                    sd_version_is_minimax_h3(version) ||
-                   sd_version_is_mage_flow(version)) {
+                   sd_version_is_mage_flow(version) ||
+                   version == VERSION_IRIS) {
             arch = LLM::LLMArch::QWEN3_VL;
         } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
             arch = LLM::LLMArch::QWEN3;
@@ -2138,9 +2139,10 @@ struct LLMEmbedder : public Conditioner {
 
     std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
                                                                                   const std::pair<int, int>& attn_range,
-                                                                                  size_t min_length = 0,
-                                                                                  size_t max_length = 100000000,
-                                                                                  bool spell_quotes = false) {
+                                                                                  size_t min_length         = 0,
+                                                                                  size_t max_length         = 100000000,
+                                                                                  bool spell_quotes         = false,
+                                                                                  const std::string& suffix = "") {
         std::vector<std::pair<std::string, float>> parsed_attention;
         if (attn_range.first >= 0 && attn_range.second > 0) {
             if (attn_range.first > 0) {
@@ -2185,6 +2187,20 @@ struct LLMEmbedder : public Conditioner {
             weights.insert(weights.end(), curr_tokens.size(), curr_weight);
         }
 
+        if (!suffix.empty()) {
+            std::vector<int> suffix_tokens;
+            if (!tokenizer->encode(suffix, suffix_tokens) || suffix_tokens.size() >= max_length) {
+                return {};
+            }
+            // Reserve the assistant-turn marker before truncating the caption.
+            if (tokens.size() > max_length - suffix_tokens.size()) {
+                tokens.resize(max_length - suffix_tokens.size());
+                weights.resize(tokens.size());
+            }
+            tokens.insert(tokens.end(), suffix_tokens.begin(), suffix_tokens.end());
+            weights.insert(weights.end(), suffix_tokens.size(), 1.f);
+        }
+
         std::vector<float> mask;
         tokenizer->pad_tokens(tokens, &weights, &mask, min_length, max_length);
 
@@ -2207,8 +2223,10 @@ struct LLMEmbedder : public Conditioner {
                                     bool spell_quotes                                       = false,
                                     int max_length                                          = 100000000,
                                     const LLM::DeepStackImageEmbeds& deepstack_image_embeds = {},
-                                    const std::vector<LLM::ImageGrid>& image_grids          = {}) {
-        auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes);
+                                    const std::vector<LLM::ImageGrid>& image_grids          = {},
+                                    const std::string& prompt_suffix                        = "",
+                                    sd::Tensor<float>* output_mask                          = nullptr) {
+        auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes, prompt_suffix);
         auto& tokens             = std::get<0>(tokens_weights_mask);
         auto& weights            = std::get<1>(tokens_weights_mask);
         auto& mask               = std::get<2>(tokens_weights_mask);
@@ -2271,6 +2289,12 @@ struct LLMEmbedder : public Conditioner {
                                                 1);
         }
 
+        if (output_mask != nullptr) {
+            *output_mask = sd::Tensor<float>::zeros({new_hidden_states.shape()[1]});
+            for (size_t i = prompt_template_encode_start_idx; i < mask.size(); ++i) {
+                output_mask->data()[i - prompt_template_encode_start_idx] = mask[i];
+            }
+        }
         return new_hidden_states;
     }
 
@@ -2337,6 +2361,8 @@ struct LLMEmbedder : public Conditioner {
         int hidden_states_min_length         = 0;  // zero pad hidden_states
         bool spell_quotes                    = false;
         std::set<int> out_layers;
+        std::string prompt_suffix;
+        sd::Tensor<float> text_mask;
 
         int64_t t0                     = ggml_time_ms();
         RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode;
@@ -3086,6 +3112,23 @@ struct LLMEmbedder : public Conditioner {
             SDCondition result;
             result.c_crossattn = std::move(hidden_states);
             return result;
+        } else if (version == VERSION_IRIS) {
+            prompt =
+                "<|im_start|>system\n"
+                "Describe the image by detailing the color, shape, size, texture, quantity, text, spatial "
+                "relationships of the objects and background:<|im_end|>\n<|im_start|>user\n";
+            auto prefix_tokens = std::get<0>(tokenize(prompt, {0, 0}));
+            if (prefix_tokens.empty()) {
+                return {};
+            }
+            prompt_template_encode_start_idx = static_cast<int>(prefix_tokens.size());
+            hidden_states_min_length         = 300;
+            max_length                       = prompt_template_encode_start_idx + hidden_states_min_length;
+            out_layers                       = {2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35};
+            prompt_attn_range.first          = static_cast<int>(prompt.size());
+            prompt += conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+            prompt_suffix            = "<|im_end|>\n<|im_start|>assistant\n";
         } else {
             GGML_ABORT("unknown version %d", version);
         }
@@ -3101,7 +3144,9 @@ struct LLMEmbedder : public Conditioner {
                                            spell_quotes,
                                            max_length,
                                            deepstack_image_embeds,
-                                           image_grids);
+                                           image_grids,
+                                           prompt_suffix,
+                                           version == VERSION_IRIS ? &text_mask : nullptr);
         if (hidden_states.empty()) {
             return {};
         }
@@ -3166,6 +3211,7 @@ struct LLMEmbedder : public Conditioner {
         SDCondition result;
         result.c_crossattn        = std::move(hidden_states);
         result.extra_c_crossattns = std::move(extra_hidden_states_vec);
+        result.c_vector           = std::move(text_mask);
         if (version == VERSION_QWEN_IMAGE_2_1) {
             auto slots = sd::Tensor<int32_t>::zeros({result.c_crossattn.shape()[1]});
             for (size_t i = 0; i < image_embeds.size(); ++i) {

src/core/backend_fit.cpp

@@ -497,11 +497,7 @@ namespace sd::backend_fit {
                 return false;
             }
             auto smaller_tile = [&](int size) {
-                int next_size = size / 2;
-                if (!tiling_params.enabled) {
-                    next_size = std::min(next_size, 256 / scale_factor);
-                }
-                return std::min(size, std::max(4, next_size));
+                return std::min(size, std::max(4, size / 2));
             };
             const int tile_size_w = smaller_tile(latent_tile_size_w);
             const int tile_size_h = smaller_tile(latent_tile_size_h);

src/model.h

@@ -65,6 +65,7 @@ enum SDVersion {
     VERSION_PIXART,
     VERSION_MING_IMAGE,
     VERSION_Z_IMAGE_L2P,
+    VERSION_IRIS,
     VERSION_COUNT,
 };
 
@@ -334,7 +335,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_krea2(version) ||
         sd_version_is_mage_flow(version) ||
         sd_version_is_sensenova_u1(version) ||
-        sd_version_is_pixart(version)) {
+        sd_version_is_pixart(version) ||
+        version == VERSION_IRIS) {
         return true;
     }
     return false;

src/model/adapter/lora_ops.cpp

@@ -11,7 +11,9 @@ ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
                                  ggml_tensor* lora_mid) {
     ggml_tensor* updown;
     // flat lora tensors to multiply it
-    int64_t lora_up_rows  = lora_up->ne[ggml_n_dims(lora_up) - 1];
+    // ggml_n_dims drops the output dimension when a linear layer has one output.
+    int lora_up_n_dims    = ggml_n_dims(lora_up);
+    int64_t lora_up_rows  = lora_up_n_dims == 1 ? 1 : lora_up->ne[lora_up_n_dims - 1];
     lora_up               = ggml_reshape_2d(ctx, lora_up, ggml_nelements(lora_up) / lora_up_rows, lora_up_rows);
     auto lora_down_n_dims = ggml_n_dims(lora_down);
     // assume n_dims should always be a multiple of 2 (otherwise rank 1 doesn't work)

src/model/diffusion/iris.hpp

@@ -0,0 +1,450 @@
+#ifndef __SD_MODEL_DIFFUSION_IRIS_HPP__
+#define __SD_MODEL_DIFFUSION_IRIS_HPP__
+
+#include <array>
+#include <stdexcept>
+#include "model/diffusion/pid.hpp"
+
+namespace Iris {
+    struct IrisConfig {
+        int64_t hidden_size    = 2560;
+        int64_t depth          = 24;
+        int64_t dual_depth     = 8;
+        int64_t head_dim       = 128;
+        int64_t num_heads      = 20;
+        int64_t kv_heads       = 5;
+        int64_t mlp_dim        = 6826;
+        int64_t patch_size     = 16;
+        int64_t text_dim       = 2560;
+        int64_t text_len       = 300;
+        int64_t text_layers    = 12;
+        int64_t layer_heads    = 32;
+        int64_t layer_mlp_dim  = 3328;
+        int64_t pixel_dim      = 16;
+        int64_t pixel_attn_dim = 1280;
+        int64_t pixel_heads    = 10;
+        int64_t pixel_depth    = 4;
+
+        static IrisConfig detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) {
+            IrisConfig config;
+            auto find = [&](const std::string& name) -> const TensorStorage& {
+                auto it = weights.find(prefix + "." + name);
+                if (it == weights.end()) {
+                    throw std::runtime_error("Iris-3B: missing weight " + prefix + "." + name);
+                }
+                return it->second;
+            };
+            config.hidden_size    = find("s_embedder.proj.weight").ne[1];
+            config.patch_size     = static_cast<int64_t>(std::sqrt(find("s_embedder.proj.weight").ne[0] / 3));
+            config.head_dim       = find("blocks.0.attn.q_norm_x.weight").ne[0];
+            config.num_heads      = config.hidden_size / config.head_dim;
+            config.kv_heads       = find("blocks.0.attn.k_proj_x.weight").ne[1] / config.head_dim;
+            config.mlp_dim        = find("blocks.0.mlp_x.w1.weight").ne[1];
+            config.text_dim       = find("y_embedder.refiner.proj.weight").ne[0];
+            config.text_len       = find("y_pos_embedding").ne[1];
+            config.text_layers    = find("y_embedder.layer_pool.weight").ne[0];
+            config.layer_mlp_dim  = find("y_embedder.layer_blocks.0.mlp.0.weight").ne[1];
+            config.pixel_dim      = find("pixel_embedder.proj.weight").ne[1];
+            config.pixel_attn_dim = find("pixel_blocks.0.compress.weight").ne[1];
+            config.pixel_heads    = config.pixel_attn_dim / find("pixel_blocks.0.attn.q_norm.weight").ne[0];
+            config.depth = config.dual_depth = config.pixel_depth = 0;
+            for (const auto& [name, tensor] : weights) {
+                if (!starts_with(name, prefix + ".")) {
+                    continue;
+                }
+                auto parts = split_string(name.substr(prefix.size() + 1), '.');
+                if (parts.size() > 2 && parts[0] == "blocks") {
+                    int64_t index = std::stoi(parts[1]);
+                    config.depth  = std::max(config.depth, index + 1);
+                    if (parts[2] == "adaln_img") {
+                        config.dual_depth = std::max(config.dual_depth, index + 1);
+                    }
+                } else if (parts.size() > 2 && parts[0] == "pixel_blocks") {
+                    config.pixel_depth = std::max(config.pixel_depth, int64_t(std::stoi(parts[1]) + 1));
+                }
+            }
+            if (config.text_layers != 12 || config.text_dim != 2560 || config.text_len != 300 ||
+                find("modulation_cores.adaln_img.weight").ne[1] != 6 * config.hidden_size ||
+                find("pixel_blocks.0.adaln.weight").ne[1] != 4 * config.pixel_dim * config.patch_size * config.patch_size) {
+                throw std::runtime_error("Iris-3B: unsupported checkpoint configuration");
+            }
+            LOG_VERBOSE("iris: hidden_size = %" PRId64 ", depth = %" PRId64 ", dual_depth = %" PRId64 ", heads = %" PRId64 "/%" PRId64 ", pixel_depth = %" PRId64,
+                        config.hidden_size, config.depth, config.dual_depth, config.num_heads, config.kv_heads, config.pixel_depth);
+            return config;
+        }
+    };
+
+    struct SelfAttention : public GGMLBlock {
+        int64_t dim;
+        int64_t heads;
+        bool qk_norm;
+
+        SelfAttention(int64_t dim, int64_t heads, bool qk_norm = true)
+            : dim(dim), heads(heads), qk_norm(qk_norm) {
+            blocks["qkv"]  = std::make_shared<Linear>(dim, 3 * dim, false);
+            blocks["proj"] = std::make_shared<Linear>(dim, dim, true);
+            if (qk_norm) {
+                blocks["q_norm"] = std::make_shared<RMSNorm>(dim / heads, 1e-6f);
+                blocks["k_norm"] = std::make_shared<RMSNorm>(dim / heads, 1e-6f);
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe = nullptr, ggml_tensor* mask = nullptr) {
+            auto c   = ctx->ggml_ctx;
+            auto qkv = split_qkv(c, std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x));
+            auto q   = ggml_reshape_4d(c, qkv[0], dim / heads, heads, x->ne[1], x->ne[2]);
+            auto k   = ggml_reshape_4d(c, qkv[1], dim / heads, heads, x->ne[1], x->ne[2]);
+            auto v   = ggml_reshape_4d(c, qkv[2], dim / heads, heads, x->ne[1], x->ne[2]);
+            if (qk_norm) {
+                q = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"])->forward(ctx, q);
+                k = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"])->forward(ctx, k);
+            }
+            if (pe) {
+                x = Rope::attention(ctx, q, k, v, pe, mask);
+            } else {
+                q = ggml_reshape_3d(c, q, dim, x->ne[1], x->ne[2]);
+                k = ggml_reshape_3d(c, k, dim, x->ne[1], x->ne[2]);
+                v = ggml_reshape_3d(c, v, dim, x->ne[1], x->ne[2]);
+                x = ggml_ext_attention_ext(ctx, q, k, v, heads, mask, false, ctx->flash_attn_enabled);
+            }
+            return std::dynamic_pointer_cast<Linear>(blocks["proj"])->forward(ctx, x);
+        }
+    };
+
+    struct TextBlock : public GGMLBlock {
+        bool layerwise;
+
+        TextBlock(int64_t dim, int64_t heads, int64_t mlp_dim, bool layerwise)
+            : layerwise(layerwise) {
+            blocks["norm1"] = std::make_shared<RMSNorm>(dim, 1e-6f);
+            blocks["norm2"] = std::make_shared<RMSNorm>(dim, 1e-6f);
+            blocks["attn"]  = std::make_shared<SelfAttention>(dim, heads, !layerwise);
+            if (layerwise) {
+                blocks["mlp.0"] = std::make_shared<Linear>(dim, mlp_dim, true);
+                blocks["mlp.2"] = std::make_shared<Linear>(mlp_dim, dim, true);
+            } else {
+                blocks["mlp"] = std::make_shared<Pid::FeedForward>(dim, mlp_dim);
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* mask = nullptr) {
+            auto h = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"])->forward(ctx, x);
+            h      = std::dynamic_pointer_cast<SelfAttention>(blocks["attn"])->forward(ctx, h, nullptr, mask);
+            x      = ggml_add(ctx->ggml_ctx, x, h);
+            h      = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"])->forward(ctx, x);
+            if (layerwise) {
+                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);
+            } else {
+                h = std::dynamic_pointer_cast<Pid::FeedForward>(blocks["mlp"])->forward(ctx, h);
+            }
+            return ggml_add(ctx->ggml_ctx, x, h);
+        }
+    };
+
+    struct TextEmbedder : public GGMLBlock {
+        IrisConfig config;
+
+        TextEmbedder(const IrisConfig& config)
+            : config(config) {
+            for (int i = 0; i < 2; ++i) {
+                blocks["layer_blocks." + std::to_string(i)]   = std::make_shared<TextBlock>(config.text_dim, config.layer_heads, config.layer_mlp_dim, true);
+                blocks["refiner.blocks." + std::to_string(i)] = std::make_shared<TextBlock>(config.hidden_size, config.num_heads, config.mlp_dim, false);
+            }
+            blocks["layer_pool"]   = std::make_shared<Linear>(config.text_layers, 1, true);
+            blocks["refiner.proj"] = std::make_shared<Linear>(config.text_dim, config.hidden_size, true);
+            blocks["refiner.norm"] = std::make_shared<RMSNorm>(config.hidden_size, 1e-6f);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* mask) {
+            auto c         = ctx->ggml_ctx;
+            int64_t tokens = x->ne[1];
+            int64_t batch  = x->ne[2];
+            // The encoder concatenates selected layers within each token's feature dimension.
+            x = ggml_reshape_3d(c, x, config.text_dim, config.text_layers, tokens * batch);
+            for (int i = 0; i < 2; ++i) {
+                x = std::dynamic_pointer_cast<TextBlock>(blocks["layer_blocks." + std::to_string(i)])->forward(ctx, x);
+            }
+            x = ggml_cont(c, ggml_permute(c, x, 1, 0, 2, 3));
+            x = std::dynamic_pointer_cast<Linear>(blocks["layer_pool"])->forward(ctx, x);
+            x = ggml_reshape_3d(c, x, config.text_dim, tokens, batch);
+            x = std::dynamic_pointer_cast<Linear>(blocks["refiner.proj"])->forward(ctx, x);
+            for (int i = 0; i < 2; ++i) {
+                x = std::dynamic_pointer_cast<TextBlock>(blocks["refiner.blocks." + std::to_string(i)])->forward(ctx, x, mask);
+            }
+            return std::dynamic_pointer_cast<RMSNorm>(blocks["refiner.norm"])->forward(ctx, x);
+        }
+    };
+
+    struct TrunkBlock : public GGMLBlock {
+        IrisConfig config;
+        bool dual;
+
+        TrunkBlock(const IrisConfig& config, bool dual)
+            : config(config), dual(dual) {
+            int64_t d = config.hidden_size;
+            for (const std::string stream : dual ? std::vector<std::string>{"x", "y"} : std::vector<std::string>{""}) {
+                std::string suffix                                = stream.empty() ? "" : "_" + stream;
+                std::string attn                                  = dual ? "attn." : "";
+                blocks[attn + "q_proj" + suffix]                  = std::make_shared<Linear>(d, d, false);
+                blocks[attn + "k_proj" + suffix]                  = std::make_shared<Linear>(d, config.kv_heads * config.head_dim, false);
+                blocks[attn + "v_proj" + suffix]                  = std::make_shared<Linear>(d, config.kv_heads * config.head_dim, false);
+                blocks[attn + "q_norm" + suffix]                  = std::make_shared<RMSNorm>(config.head_dim, 1e-6f);
+                blocks[attn + "k_norm" + suffix]                  = std::make_shared<RMSNorm>(config.head_dim, 1e-6f);
+                blocks[dual ? "attn.proj" + suffix : "attn_proj"] = std::make_shared<Linear>(d, d, true);
+                blocks["attn_gate" + suffix]                      = std::make_shared<Linear>(d, d, false);
+                blocks["norm" + suffix + "1"]                     = std::make_shared<RMSNorm>(d, 1e-6f);
+                blocks["norm" + suffix + "2"]                     = std::make_shared<RMSNorm>(d, 1e-6f);
+                blocks["attn_post_norm" + suffix]                 = std::make_shared<RMSNorm>(d, 1e-6f);
+                blocks["mlp_post_norm" + suffix]                  = std::make_shared<RMSNorm>(d, 1e-6f);
+                blocks["mlp" + suffix]                            = std::make_shared<Pid::FeedForward>(d, config.mlp_dim);
+            }
+        }
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& = {}, const std::string = "") override {
+            if (dual) {
+                params["adaln_img.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 6 * config.hidden_size);
+                params["adaln_txt.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 6 * config.hidden_size);
+            } else {
+                params["adaln.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 6 * config.hidden_size);
+            }
+        }
+
+        ggml_tensor* norm(GGMLRunnerContext* ctx, const std::string& name, ggml_tensor* x) {
+            return std::dynamic_pointer_cast<RMSNorm>(blocks[name])->forward(ctx, x);
+        }
+
+        ggml_tensor* linear(GGMLRunnerContext* ctx, const std::string& name, ggml_tensor* x) {
+            return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x);
+        }
+
+        std::array<ggml_tensor*, 3> project(GGMLRunnerContext* ctx, ggml_tensor* h, const std::string& suffix) {
+            auto c        = ctx->ggml_ctx;
+            std::string a = dual ? "attn." : "";
+            auto q        = linear(ctx, a + "q_proj" + suffix, h);
+            auto k        = linear(ctx, a + "k_proj" + suffix, h);
+            auto v        = linear(ctx, a + "v_proj" + suffix, h);
+            q             = ggml_reshape_4d(c, q, config.head_dim, config.num_heads, h->ne[1], h->ne[2]);
+            k             = ggml_reshape_4d(c, k, config.head_dim, config.kv_heads, h->ne[1], h->ne[2]);
+            v             = ggml_reshape_4d(c, v, config.head_dim, config.kv_heads, h->ne[1], h->ne[2]);
+            return {norm(ctx, a + "q_norm" + suffix, q), norm(ctx, a + "k_norm" + suffix, k), v};
+        }
+
+        ggml_tensor* finish(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* h, ggml_tensor* out, const std::vector<ggml_tensor*>& mods, const std::string& suffix) {
+            auto c = ctx->ggml_ctx;
+            out    = ggml_mul(c, out, ggml_sigmoid(c, linear(ctx, "attn_gate" + suffix, h)));
+            out    = linear(ctx, dual ? "attn.proj" + suffix : "attn_proj", out);
+            out    = norm(ctx, "attn_post_norm" + suffix, out);
+            x      = ggml_add(c, x, ggml_mul(c, out, mods[2]));
+            h      = Pid::apply_adaln(c, norm(ctx, "norm" + suffix + "2", x), mods[3], mods[4]);
+            h      = std::dynamic_pointer_cast<Pid::FeedForward>(blocks["mlp" + suffix])->forward(ctx, h);
+            h      = norm(ctx, "mlp_post_norm" + suffix, h);
+            return ggml_add(c, x, ggml_mul(c, h, mods[5]));
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* y, ggml_tensor* img_mod, ggml_tensor* txt_mod, ggml_tensor* pe) {
+            auto c     = ctx->ggml_ctx;
+            int64_t nt = y->ne[1];
+            auto mx    = ggml_ext_chunk(c, ggml_add(c, img_mod, params[dual ? "adaln_img.bias" : "adaln.bias"]), 6, 0);
+            if (!dual) {
+                auto tokens = ggml_concat(c, y, x, 1);
+                auto h      = Pid::apply_adaln(c, norm(ctx, "norm1", tokens), mx[0], mx[1]);
+                auto qkv    = project(ctx, h, "");
+                auto out    = Rope::attention(ctx, qkv[0], qkv[1], qkv[2], pe, nullptr);
+                tokens      = finish(ctx, tokens, h, out, mx, "");
+                return {ggml_ext_slice(c, tokens, 1, nt, tokens->ne[1]), ggml_ext_slice(c, tokens, 1, 0, nt)};
+            }
+            auto my  = ggml_ext_chunk(c, ggml_add(c, txt_mod, params["adaln_txt.bias"]), 6, 0);
+            auto hx  = Pid::apply_adaln(c, norm(ctx, "norm_x1", x), mx[0], mx[1]);
+            auto hy  = Pid::apply_adaln(c, norm(ctx, "norm_y1", y), my[0], my[1]);
+            auto qx  = project(ctx, hx, "_x");
+            auto qy  = project(ctx, hy, "_y");
+            auto out = Rope::attention(ctx, ggml_concat(c, qy[0], qx[0], 2), ggml_concat(c, qy[1], qx[1], 2),
+                                       ggml_concat(c, qy[2], qx[2], 2), pe, nullptr);
+            x        = finish(ctx, x, hx, ggml_ext_slice(c, out, 1, nt, out->ne[1]), mx, "_x");
+            y        = finish(ctx, y, hy, ggml_ext_slice(c, out, 1, 0, nt), my, "_y");
+            return {x, y};
+        }
+    };
+
+    struct PixelBlock : public GGMLBlock {
+        IrisConfig config;
+
+        PixelBlock(const IrisConfig& config)
+            : config(config) {
+            auto d             = config.pixel_dim;
+            auto flat          = d * config.patch_size * config.patch_size;
+            blocks["norm1"]    = std::make_shared<RMSNorm>(d, 1e-6f);
+            blocks["norm2"]    = std::make_shared<RMSNorm>(d, 1e-6f);
+            blocks["compress"] = std::make_shared<Linear>(flat, config.pixel_attn_dim, true);
+            blocks["expand"]   = std::make_shared<Linear>(config.pixel_attn_dim, flat, true);
+            blocks["adaln"]    = std::make_shared<Linear>(config.hidden_size, 4 * flat, true);
+            blocks["attn"]     = std::make_shared<SelfAttention>(config.pixel_attn_dim, config.pixel_heads);
+            blocks["mlp.fc1"]  = std::make_shared<Linear>(d, d * 4, true);
+            blocks["mlp.fc2"]  = std::make_shared<Linear>(d * 4, d, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cond, ggml_tensor* pe) {
+            auto c       = ctx->ggml_ctx;
+            auto d       = config.pixel_dim;
+            auto p2      = config.patch_size * config.patch_size;
+            auto patches = cond->ne[1];
+            auto batch   = cond->ne[2];
+            auto mods    = std::dynamic_pointer_cast<Linear>(blocks["adaln"])->forward(ctx, cond);
+            mods         = ggml_reshape_3d(c, mods, 4 * d, p2, patches * batch);
+            auto m       = ggml_ext_chunk(c, mods, 4, 0);
+            auto h       = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"])->forward(ctx, x);
+            h            = ggml_reshape_3d(c, h, d * p2, patches, batch);
+            h            = std::dynamic_pointer_cast<Linear>(blocks["compress"])->forward(ctx, h);
+            h            = std::dynamic_pointer_cast<SelfAttention>(blocks["attn"])->forward(ctx, h, pe);
+            h            = std::dynamic_pointer_cast<Linear>(blocks["expand"])->forward(ctx, h);
+            h            = ggml_reshape_3d(c, h, d, p2, patches * batch);
+            x            = ggml_add(c, x, Pid::apply_adaln(c, h, m[1], m[0]));
+            h            = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"])->forward(ctx, x);
+            h            = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"])->forward(ctx, h);
+            h            = ggml_gelu_erf(c, h);
+            h            = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"])->forward(ctx, h);
+            return ggml_add(c, x, Pid::apply_adaln(c, h, m[3], m[2]));
+        }
+    };
+
+    struct IrisModel : public GGMLBlock {
+        IrisConfig config;
+
+        IrisModel(const IrisConfig& config)
+            : config(config) {
+            blocks["s_embedder"]                 = std::make_shared<Pid::PatchTokenEmbedder>(3 * config.patch_size * config.patch_size, config.hidden_size);
+            blocks["t_embedder"]                 = std::make_shared<Pid::PixelDiTTimestepEmbedder>(config.hidden_size);
+            blocks["y_embedder"]                 = std::make_shared<TextEmbedder>(config);
+            blocks["modulation_cores.adaln_img"] = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size, true);
+            blocks["modulation_cores.adaln_txt"] = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size, true);
+            for (int64_t i = 0; i < config.depth; ++i) {
+                blocks["blocks." + std::to_string(i)] = std::make_shared<TrunkBlock>(config, i < config.dual_depth);
+            }
+            blocks["pixel_embedder"] = std::make_shared<Pid::PixelTokenEmbedder>(3, config.pixel_dim);
+            for (int64_t i = 0; i < config.pixel_depth; ++i) {
+                blocks["pixel_blocks." + std::to_string(i)] = std::make_shared<PixelBlock>(config);
+            }
+            blocks["final_layer"] = std::make_shared<Pid::FinalLayer>(config.pixel_dim, 3);
+        }
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& = {}, const std::string = "") override {
+            params["y_pos_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, config.hidden_size, config.text_len, 1);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image, ggml_tensor* t, ggml_tensor* text, ggml_tensor* mask, ggml_tensor* joint_pe, ggml_tensor* pixel_pe, ggml_tensor* pixel_pos) {
+            auto c    = ctx->ggml_ctx;
+            int p     = static_cast<int>(config.patch_size);
+            auto x    = DiT::patchify(c, image, p, p, true);
+            x         = std::dynamic_pointer_cast<Pid::PatchTokenEmbedder>(blocks["s_embedder"])->forward(ctx, x);
+            auto te   = std::dynamic_pointer_cast<Pid::PixelDiTTimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
+            te        = ggml_reshape_3d(c, te, config.hidden_size, 1, image->ne[3]);
+            auto cond = ggml_silu(c, te);
+            auto im   = std::dynamic_pointer_cast<Linear>(blocks["modulation_cores.adaln_img"])->forward(ctx, cond);
+            auto tm   = std::dynamic_pointer_cast<Linear>(blocks["modulation_cores.adaln_txt"])->forward(ctx, cond);
+            auto y    = std::dynamic_pointer_cast<TextEmbedder>(blocks["y_embedder"])->forward(ctx, text, mask);
+            y         = ggml_add(c, y, params["y_pos_embedding"]);
+            for (int64_t i = 0; i < config.depth; ++i) {
+                auto block     = std::dynamic_pointer_cast<TrunkBlock>(blocks["blocks." + std::to_string(i)]);
+                std::tie(x, y) = block->forward(ctx, x, y, im, tm, joint_pe);
+                sd::ggml_graph_cut::mark_graph_cut(x, "iris.blocks." + std::to_string(i), "x");
+                sd::ggml_graph_cut::mark_graph_cut(y, "iris.blocks." + std::to_string(i), "y");
+            }
+            cond = ggml_silu(c, ggml_add(c, x, te));
+            x    = std::dynamic_pointer_cast<Pid::PixelTokenEmbedder>(blocks["pixel_embedder"])->forward(ctx, image, p, pixel_pos);
+            for (int64_t i = 0; i < config.pixel_depth; ++i) {
+                x = std::dynamic_pointer_cast<PixelBlock>(blocks["pixel_blocks." + std::to_string(i)])->forward(ctx, x, cond, pixel_pe);
+                sd::ggml_graph_cut::mark_graph_cut(x, "iris.pixel_blocks." + std::to_string(i), "x");
+            }
+            x = std::dynamic_pointer_cast<Pid::FinalLayer>(blocks["final_layer"])->forward(ctx, x);
+            x = ggml_reshape_3d(c, x, 3 * p * p, cond->ne[1], image->ne[3]);
+            return DiT::unpatchify(c, x, image->ne[1] / p, image->ne[0] / p, p, p, false);
+        }
+    };
+
+    inline std::vector<float> image_rope(int64_t height, int64_t width, int64_t dim) {
+        std::vector<float> pe;
+        pe.reserve(height * width * dim * 2);
+        float step = 16.f / static_cast<float>(std::max<int64_t>(std::max(height, width) - 1, 1));
+        for (int64_t y = 0; y < height; ++y) {
+            for (int64_t x = 0; x < width; ++x) {
+                for (int64_t i = 0; i < dim / 4; ++i) {
+                    float frequency = std::pow(10000.f, -4.f * static_cast<float>(i) / static_cast<float>(dim));
+                    for (int64_t pos : {x, y}) {
+                        float angle = static_cast<float>(pos) * step * frequency;
+                        pe.insert(pe.end(), {std::cos(angle), -std::sin(angle), std::sin(angle), std::cos(angle)});
+                    }
+                }
+            }
+        }
+        return pe;
+    }
+
+    struct IrisRunner : public DiffusionModelRunner {
+        IrisConfig config;
+        IrisModel model;
+        std::vector<float> joint_pe_data;
+        std::vector<float> pixel_pe_data;
+        std::vector<float> pixel_pos_data;
+        std::vector<float> text_mask_data;
+
+        IrisRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(IrisConfig::detect_from_weights(weights, prefix)),
+              model(config) {
+            model.init(params_ctx, weights, prefix);
+        }
+
+        std::string get_desc() override { return "Iris-3B"; }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            model.get_param_tensors(tensors, prefix);
+        }
+
+        sd::Tensor<float> compute(int n_threads, const DiffusionParams& inputs) override {
+            if (!inputs.x || !inputs.timesteps || !inputs.context || !inputs.y || inputs.y->empty()) {
+                LOG_ERROR("Iris-3B requires text features and their padding mask");
+                return {};
+            }
+            auto get_graph = [&]() {
+                auto gf   = new_graph_custom(196608);
+                auto x    = make_input(*inputs.x);
+                auto t    = make_input(*inputs.timesteps);
+                auto text = make_input(*inputs.context);
+                int h     = static_cast<int>(x->ne[1]);
+                int w     = static_cast<int>(x->ne[0]);
+                int nt    = static_cast<int>(config.text_len);
+                GGML_ASSERT(h % config.patch_size == 0 && w % config.patch_size == 0);
+                GGML_ASSERT(text->ne[0] == config.text_dim * config.text_layers && text->ne[1] == nt);
+                GGML_ASSERT(inputs.y->numel() == nt && x->ne[3] == 1);
+                joint_pe_data = Pid::make_rope_1d(nt, static_cast<int>(config.head_dim), 10000.f);
+                auto img      = image_rope(h / config.patch_size, w / config.patch_size, config.head_dim);
+                joint_pe_data.insert(joint_pe_data.end(), img.begin(), img.end());
+                pixel_pe_data  = image_rope(h / config.patch_size, w / config.patch_size, config.pixel_attn_dim / config.pixel_heads);
+                pixel_pos_data = Pid::make_pixel_abs_pos(h, w, static_cast<int>(config.pixel_dim));
+                text_mask_data.resize(nt * nt);
+                for (int q = 0; q < nt; ++q) {
+                    for (int k = 0; k < nt; ++k) {
+                        text_mask_data[q * nt + k] = inputs.y->values()[k] != 0.f || q == k ? 0.f : -INFINITY;
+                    }
+                }
+                auto pe   = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, nt + img.size() / (2 * config.head_dim));
+                auto pp   = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.pixel_attn_dim / config.pixel_heads / 2, (h / config.patch_size) * (w / config.patch_size));
+                auto pos  = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.pixel_dim, h * w, 1);
+                auto mask = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, nt, nt);
+                set_backend_tensor_data(pe, joint_pe_data.data());
+                set_backend_tensor_data(pp, pixel_pe_data.data());
+                set_backend_tensor_data(pos, pixel_pos_data.data());
+                set_backend_tensor_data(mask, text_mask_data.data());
+                auto ctx = get_context();
+                auto out = model.forward(&ctx, x, t, text, mask, pe, pp, pos);
+                ggml_build_forward_expand(gf, out);
+                return gf;
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), inputs.x->dim());
+        }
+    };
+}  // namespace Iris
+
+#endif  // __SD_MODEL_DIFFUSION_IRIS_HPP__

src/model/diffusion/unet.hpp

@@ -587,7 +587,9 @@ public:
             label_emb      = ggml_silu_inplace(ctx->ggml_ctx, label_emb);
             label_emb      = label_embed_2->forward(ctx, label_emb);  // [N, time_embed_dim]
 
-            emb = ggml_add(ctx->ggml_ctx, emb, label_emb);  // [N, time_embed_dim]
+            emb = label_emb->ne[1] > emb->ne[1]
+                      ? ggml_add(ctx->ggml_ctx, label_emb, emb)
+                      : ggml_add(ctx->ggml_ctx, emb, label_emb);  // [N, time_embed_dim]
         }
         // sd::ggml_graph_cut::mark_graph_cut(emb, "unet.prelude", "emb");
 

src/model/upscaler/mosr.hpp

@@ -0,0 +1,337 @@
+#ifndef __SD_MODEL_UPSCALER_MOSR_HPP__
+#define __SD_MODEL_UPSCALER_MOSR_HPP__
+
+#include <algorithm>
+#include <cmath>
+#include <map>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
+
+/*
+    ===================================    MoSR  ===================================
+    References:
+    https://github.com/umzi2/MoSR
+    https://github.com/chaiNNer-org/spandrel/blob/main/libs/spandrel/spandrel/architectures/MoSR/__arch/mosr_arch.py
+
+    Only the PixelShuffle ("ps") upsampler is supported.
+*/
+
+struct MoSRConfig {
+    int scale       = 4;
+    int n_block     = 24;
+    int in_ch       = 3;
+    int dim         = 64;
+    int hidden      = 96;
+    int conv_ch     = 64;
+    int kernel_size = 7;
+
+    bool is_valid() const {
+        // The upscaler feeds RGB tiles, so only 3-channel models are usable.
+        return in_ch == 3 && scale > 0 && n_block > 0 && conv_ch <= hidden;
+    }
+
+    static MoSRConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                          const std::string& prefix = "") {
+        MoSRConfig config;
+        auto find_weight = [&](const std::string& suffix) -> const TensorStorage* {
+            std::string name = prefix.empty() ? suffix : prefix + "." + suffix;
+            auto iter        = tensor_storage_map.find(name);
+            if (iter == tensor_storage_map.end() || iter->second.n_dims != 4) {
+                return nullptr;
+            }
+            return &iter->second;
+        };
+
+        const TensorStorage* first = find_weight("gblocks.0.weight");
+        const TensorStorage* fc1   = find_weight("gblocks.1.fc1.weight");
+        const TensorStorage* conv  = find_weight("gblocks.1.conv.weight");
+        const TensorStorage* up    = find_weight("upsampler.0.weight");
+        if (first == nullptr || fc1 == nullptr || conv == nullptr || up == nullptr) {
+            // Missing tensors, or the unsupported "gps"/"dys" upsamplers.
+            config.n_block = 0;
+            return config;
+        }
+
+        int seq_len                      = 0;
+        const std::string gblocks_prefix = prefix.empty() ? "gblocks." : prefix + ".gblocks.";
+        for (const auto& [name, _] : tensor_storage_map) {
+            if (!starts_with(name, gblocks_prefix)) {
+                continue;
+            }
+            size_t pos = name.find('.', gblocks_prefix.size());
+            if (pos == std::string::npos) {
+                continue;
+            }
+            try {
+                int idx = std::stoi(name.substr(gblocks_prefix.size(), pos - gblocks_prefix.size()));
+                seq_len = std::max(seq_len, idx + 1);
+            } catch (...) {
+            }
+        }
+
+        // gblocks = [conv_in, n_block x GatedCNNBlock, conv, Mish, conv, Mish, conv_1x1]
+        config.n_block     = seq_len - 6;
+        config.in_ch       = (int)first->ne[2];
+        config.dim         = (int)first->ne[3];
+        config.hidden      = (int)fc1->ne[3] / 2;
+        config.conv_ch     = (int)conv->ne[3];
+        config.kernel_size = (int)conv->ne[0];
+        config.scale       = (int)std::lround(std::sqrt((double)up->ne[3] / config.in_ch));
+        if (config.scale * config.scale * config.in_ch != up->ne[3]) {
+            config.scale = 0;
+        }
+
+        LOG_VERBOSE("mosr: scale = %d, n_block = %d, in_ch = %d, dim = %d, hidden = %d, conv_ch = %d, kernel_size = %d",
+                    config.scale,
+                    config.n_block,
+                    config.in_ch,
+                    config.dim,
+                    config.hidden,
+                    config.conv_ch,
+                    config.kernel_size);
+        return config;
+    }
+};
+
+static inline ggml_tensor* mosr_mish(ggml_context* ctx, ggml_tensor* x) {
+    return ggml_mul(ctx, x, ggml_tanh(ctx, ggml_softplus(ctx, x)));
+}
+
+class MoSRLayerNorm2d : public UnaryBlock {
+protected:
+    int64_t channels;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        params["weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
+        params["bias"]   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
+    }
+
+public:
+    explicit MoSRLayerNorm2d(int64_t channels)
+        : channels(channels) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        // [N, C, H, W] -> [N, H, W, C] so the 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, params["weight"], params["bias"], 1e-6f);
+        return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+    }
+};
+
+class MoSRDepthwiseConv2d : public UnaryBlock {
+protected:
+    int64_t channels;
+    int kernel_size;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        // The direct depthwise kernels of the CUDA and Vulkan backends need an F32 kernel.
+        params["weight"] = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, kernel_size, kernel_size, 1, channels);
+        params["bias"]   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
+    }
+
+public:
+    MoSRDepthwiseConv2d(int64_t channels, int kernel_size)
+        : channels(channels), kernel_size(kernel_size) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        const int pad = kernel_size / 2;
+        x             = ggml_conv_2d_dw_direct(ctx->ggml_ctx, params["weight"], ggml_ext_cont(ctx->ggml_ctx, x), 1, 1, pad, pad, 1, 1);
+        auto b        = ggml_reshape_4d(ctx->ggml_ctx, params["bias"], 1, 1, channels, 1);
+        return ggml_add_inplace(ctx->ggml_ctx, x, b);
+    }
+};
+
+class MoSRGatedCNNBlock : public GGMLBlock {
+protected:
+    int hidden;
+    int conv_ch;
+
+public:
+    MoSRGatedCNNBlock(int dim, int hidden, int conv_ch, int kernel_size)
+        : hidden(hidden), conv_ch(conv_ch) {
+        blocks["norm"] = std::shared_ptr<GGMLBlock>(new MoSRLayerNorm2d(dim));
+        blocks["fc1"]  = std::shared_ptr<GGMLBlock>(new Conv2d(dim, hidden * 2, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv"] = std::shared_ptr<GGMLBlock>(new MoSRDepthwiseConv2d(conv_ch, kernel_size));
+        blocks["fc2"]  = std::shared_ptr<GGMLBlock>(new Conv2d(hidden, dim, {3, 3}, {1, 1}, {1, 1}));
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        // x: [n, dim, h, w]
+        // return: [n, dim, h, w]
+        auto norm = std::dynamic_pointer_cast<MoSRLayerNorm2d>(blocks["norm"]);
+        auto fc1  = std::dynamic_pointer_cast<Conv2d>(blocks["fc1"]);
+        auto conv = std::dynamic_pointer_cast<MoSRDepthwiseConv2d>(blocks["conv"]);
+        auto fc2  = std::dynamic_pointer_cast<Conv2d>(blocks["fc2"]);
+
+        auto h = fc1->forward(ctx, norm->forward(ctx, x));
+
+        auto channel_view = [&](int64_t offset, int64_t count) {
+            return ggml_view_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], count, h->ne[3],
+                                h->nb[1], h->nb[2], h->nb[3], offset * h->nb[2]);
+        };
+        auto g = channel_view(0, hidden);
+        auto i = channel_view(hidden, hidden - conv_ch);
+        auto c = conv->forward(ctx, channel_view(2 * hidden - conv_ch, conv_ch));
+        if (hidden > conv_ch) {
+            c = ggml_concat(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, i), c, 2);
+        }
+
+        h = ggml_mul(ctx->ggml_ctx, c, mosr_mish(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, g)));
+        h = mosr_mish(ctx->ggml_ctx, fc2->forward(ctx, h));
+        return ggml_add(ctx->ggml_ctx, h, ggml_scale_bias(ctx->ggml_ctx, x, 1.f, -0.5f));
+    }
+};
+
+class MoSRConvBlock : public GGMLBlock {
+public:
+    MoSRConvBlock(int in_ch, int out_ch) {
+        blocks["block.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_ch, out_ch, {3, 3}, {1, 1}, {1, 1}));
+        blocks["block.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(out_ch, out_ch, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv11"]  = std::shared_ptr<GGMLBlock>(new Conv2d(in_ch, out_ch, {1, 1}));
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        auto conv0  = std::dynamic_pointer_cast<Conv2d>(blocks["block.0"]);
+        auto conv2  = std::dynamic_pointer_cast<Conv2d>(blocks["block.2"]);
+        auto conv11 = std::dynamic_pointer_cast<Conv2d>(blocks["conv11"]);
+
+        auto out = mosr_mish(ctx->ggml_ctx, conv0->forward(ctx, x));
+        out      = mosr_mish(ctx->ggml_ctx, conv2->forward(ctx, out));
+        return ggml_add(ctx->ggml_ctx, out, conv11->forward(ctx, x));
+    }
+};
+
+class MoSRNet : public GGMLBlock {
+protected:
+    MoSRConfig config;
+
+    static std::string gblock_name(int idx) {
+        return "gblocks." + std::to_string(idx);
+    }
+
+    // torch.nn.PixelShuffle: [n, c*r*r, h, w] -> [n, c, h*r, w*r]
+    static ggml_tensor* pixel_shuffle(ggml_context* ctx, ggml_tensor* x, int r) {
+        const int64_t w = x->ne[0];
+        const int64_t h = x->ne[1];
+        const int64_t c = x->ne[2] / (r * r);
+        const int64_t n = x->ne[3];
+        x               = ggml_reshape_4d(ctx, x, w, h, r, r * c * n);       // [n*c*r, r, h, w]
+        x               = ggml_cont(ctx, ggml_permute(ctx, x, 1, 2, 0, 3));  // [n*c*r, h, w, r]
+        x               = ggml_reshape_4d(ctx, x, r * w, h, r, c * n);       // [n*c, r, h, w*r]
+        x               = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [n*c, h, r, w*r]
+        return ggml_reshape_4d(ctx, x, r * w, r * h, c, n);
+    }
+
+public:
+    explicit MoSRNet(MoSRConfig config)
+        : config(std::move(config)) {
+        const int dim = this->config.dim;
+        const int nb  = this->config.n_block;
+
+        blocks[gblock_name(0)] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.in_ch, dim, {3, 3}, {1, 1}, {1, 1}));
+        for (int i = 1; i <= nb; i++) {
+            blocks[gblock_name(i)] = std::shared_ptr<GGMLBlock>(new MoSRGatedCNNBlock(dim,
+                                                                                      this->config.hidden,
+                                                                                      this->config.conv_ch,
+                                                                                      this->config.kernel_size));
+        }
+        blocks[gblock_name(nb + 1)] = std::shared_ptr<GGMLBlock>(new Conv2d(dim, dim * 2, {3, 3}, {1, 1}, {1, 1}));
+        blocks[gblock_name(nb + 3)] = std::shared_ptr<GGMLBlock>(new Conv2d(dim * 2, dim, {3, 3}, {1, 1}, {1, 1}));
+        blocks[gblock_name(nb + 5)] = std::shared_ptr<GGMLBlock>(new Conv2d(dim, dim, {1, 1}));
+        blocks["shortcut"]          = std::shared_ptr<GGMLBlock>(new MoSRConvBlock(this->config.in_ch, dim));
+        blocks["upsampler.0"]       = std::shared_ptr<GGMLBlock>(new Conv2d(dim,
+                                                                            this->config.in_ch * this->config.scale * this->config.scale,
+                                                                            {3, 3},
+                                                                            {1, 1},
+                                                                            {1, 1}));
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        // x: [n, in_ch, h, w]
+        // return: [n, in_ch, h*scale, w*scale]
+        const int nb   = config.n_block;
+        auto conv_in   = std::dynamic_pointer_cast<Conv2d>(blocks[gblock_name(0)]);
+        auto conv_mid1 = std::dynamic_pointer_cast<Conv2d>(blocks[gblock_name(nb + 1)]);
+        auto conv_mid2 = std::dynamic_pointer_cast<Conv2d>(blocks[gblock_name(nb + 3)]);
+        auto conv_out  = std::dynamic_pointer_cast<Conv2d>(blocks[gblock_name(nb + 5)]);
+        auto shortcut  = std::dynamic_pointer_cast<MoSRConvBlock>(blocks["shortcut"]);
+        auto upsampler = std::dynamic_pointer_cast<Conv2d>(blocks["upsampler.0"]);
+
+        auto skip = shortcut->forward(ctx, x);
+        sd::ggml_graph_cut::mark_graph_cut(skip, "mosr.shortcut", "skip");
+
+        auto feat = conv_in->forward(ctx, x);
+        sd::ggml_graph_cut::mark_graph_cut(feat, "mosr.prelude", "feat");
+        for (int i = 1; i <= nb; i++) {
+            auto block = std::dynamic_pointer_cast<MoSRGatedCNNBlock>(blocks[gblock_name(i)]);
+            feat       = block->forward(ctx, feat);
+            sd::ggml_graph_cut::mark_graph_cut(feat, "mosr.gblocks." + std::to_string(i), "feat");
+        }
+        feat = mosr_mish(ctx->ggml_ctx, conv_mid1->forward(ctx, feat));
+        feat = mosr_mish(ctx->ggml_ctx, conv_mid2->forward(ctx, feat));
+        feat = conv_out->forward(ctx, feat);
+        feat = ggml_add(ctx->ggml_ctx, feat, ggml_scale_bias(ctx->ggml_ctx, skip, 1.f, -0.5f));
+        sd::ggml_graph_cut::mark_graph_cut(feat, "mosr.body.out", "feat");
+
+        auto out = pixel_shuffle(ctx->ggml_ctx, upsampler->forward(ctx, feat), config.scale);
+        sd::ggml_graph_cut::mark_graph_cut(out, "mosr.final", "out");
+        return out;
+    }
+};
+
+struct MoSR : public GGMLRunner {
+    MoSRConfig config;
+    std::unique_ptr<MoSRNet> mosr_net;
+
+    MoSR(ggml_backend_t backend,
+         const String2TensorStorage& tensor_storage_map      = {},
+         std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : GGMLRunner(backend, weight_manager),
+          config(MoSRConfig::detect_from_weights(tensor_storage_map)) {
+        if (!config.is_valid()) {
+            return;
+        }
+        mosr_net = std::make_unique<MoSRNet>(config);
+        mosr_net->init(params_ctx, tensor_storage_map, "");
+    }
+
+    std::string get_desc() override {
+        return "mosr";
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
+        if (!mosr_net) {
+            return;
+        }
+
+        mosr_net->get_param_tensors(tensors);
+    }
+
+    ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
+        if (!mosr_net)
+            return nullptr;
+        constexpr int kGraphNodes = 1 << 16;  // 65k
+        ggml_cgraph* gf           = new_graph_custom(kGraphNodes);
+        ggml_tensor* x            = make_input(x_tensor);
+
+        auto runner_ctx  = get_context();
+        ggml_tensor* out = mosr_net->forward(&runner_ctx, x);
+        ggml_build_forward_expand(gf, out);
+        return gf;
+    }
+
+    sd::Tensor<float> compute(const int n_threads,
+                              const sd::Tensor<float>& x) {
+        auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
+        auto result    = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
+        return result;
+    }
+};
+
+#endif  // __SD_MODEL_UPSCALER_MOSR_HPP__

src/model/vae/tae.hpp

@@ -84,12 +84,17 @@ class TinyEncoder : public UnaryBlock {
     int channels    = 64;
     int z_channels  = 4;
     int num_blocks  = 3;
+    bool f16;
 
 public:
-    TinyEncoder(int z_channels = 4, bool use_midblock_gn = false)
-        : z_channels(z_channels) {
-        int index                       = 0;
+    TinyEncoder(int z_channels = 4, bool use_midblock_gn = false, bool f16 = false)
+        : z_channels(z_channels), f16(f16) {
+        in_channels                     = f16 ? 16 : 3;
+        int index                       = f16 ? 1 : 0;
         blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, channels, {3, 3}, {1, 1}, {1, 1}));
+        if (f16) {
+            index++;  // nn.ReLU()
+        }
         blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
 
         blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
@@ -97,12 +102,14 @@ public:
             blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
         }
 
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, f16 ? channels * 2 : channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
+        channels *= f16 ? 2 : 1;
         for (int i = 0; i < num_blocks; i++) {
             blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
         }
 
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, f16 ? channels * 2 : channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
+        channels *= f16 ? 2 : 1;
         for (int i = 0; i < num_blocks; i++) {
             blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels, use_midblock_gn));
         }
@@ -114,7 +121,11 @@ public:
         // x: [n, in_channels, h, w]
         // return: [n, z_channels, h/8, w/8]
 
-        for (int i = 0; i < num_blocks * 3 + 6; i++) {
+        for (int i = f16 ? 1 : 0; i < num_blocks * 3 + 6 + (f16 ? 2 : 0); i++) {
+            if (f16 && i == 2) {
+                x = ggml_relu_inplace(ctx->ggml_ctx, x);
+                continue;
+            }
             auto block = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(i)]);
 
             x = block->forward(ctx, x);
@@ -131,9 +142,11 @@ class TinyDecoder : public UnaryBlock {
     int num_blocks   = 3;
 
 public:
-    TinyDecoder(int z_channels = 4, bool use_midblock_gn = false)
+    TinyDecoder(int z_channels = 4, bool use_midblock_gn = false, bool f16 = false)
         : z_channels(z_channels) {
-        int index = 0;
+        channels     = f16 ? 256 : 64;
+        out_channels = f16 ? 16 : 3;
+        int index    = 0;
 
         blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(z_channels, channels, {3, 3}, {1, 1}, {1, 1}));
         index++;  // nn.ReLU()
@@ -142,13 +155,15 @@ public:
             blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels, use_midblock_gn));
         }
         index++;  // nn.Upsample()
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, f16 ? channels / 2 : channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
+        channels /= f16 ? 2 : 1;
 
         for (int i = 0; i < num_blocks; i++) {
             blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
         }
         index++;  // nn.Upsample()
-        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
+        blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, f16 ? channels / 2 : channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
+        channels /= f16 ? 2 : 1;
 
         for (int i = 0; i < num_blocks; i++) {
             blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
@@ -691,6 +706,7 @@ class TAESD : public GGMLBlock {
 protected:
     bool decode_only;
     bool taef2 = false;
+    bool f16   = false;
 
 public:
     int z_channels = 4;
@@ -708,10 +724,14 @@ public:
             z_channels      = 32;
             use_midblock_gn = true;
         }
-        blocks["decoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyDecoder(z_channels, use_midblock_gn));
+        f16 = version == VERSION_QWEN_IMAGE_2_1;
+        if (f16) {
+            z_channels = 64;
+        }
+        blocks["decoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyDecoder(z_channels, use_midblock_gn, f16));
 
         if (!decode_only) {
-            blocks["encoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyEncoder(z_channels, use_midblock_gn));
+            blocks["encoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyEncoder(z_channels, use_midblock_gn, f16));
         }
     }
 
@@ -720,10 +740,15 @@ public:
         if (taef2) {
             z = unpatchify(ctx->ggml_ctx, z, 2);
         }
-        return decoder->forward(ctx, z);
+        auto x = decoder->forward(ctx, z);
+        return f16 ? unpatchify(ctx->ggml_ctx, x, 2) : x;
     }
 
     ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        if (f16) {
+            GGML_ASSERT(x->ne[0] % 2 == 0 && x->ne[1] % 2 == 0);
+            x = patchify(ctx->ggml_ctx, x, 2);
+        }
         auto encoder = std::dynamic_pointer_cast<TinyEncoder>(blocks["encoder.layers"]);
         auto z       = encoder->forward(ctx, x);
         if (taef2) {

src/model/vae/vae.hpp

@@ -176,11 +176,20 @@ public:
         int scale_factor = 8;
         if (version == VERSION_LTXAV) {
             scale_factor = 32;
-        } else if (version == VERSION_WAN2_2_TI2V || version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
+        } else if (version == VERSION_WAN2_2_TI2V ||
+                   version == VERSION_QWEN_IMAGE_2_1 ||
+                   sd_version_is_hunyuan_video(version) ||
+                   sd_version_is_mage_flow(version) ||
+                   sd_version_is_minimax_h3(version)) {
             scale_factor = 16;
         } else if (sd_version_uses_flux2_vae(version)) {
             scale_factor = 16;
-        } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version) || sd_version_is_sensenova_u1(version) || sd_version_is_z_image_l2p(version)) {
+        } else if (version == VERSION_CHROMA_RADIANCE ||
+                   version == VERSION_HIDREAM_O1 ||
+                   sd_version_is_minit2i(version) ||
+                   sd_version_is_sensenova_u1(version) ||
+                   sd_version_is_z_image_l2p(version) ||
+                   version == VERSION_IRIS) {
             scale_factor = 1;
         }
         return scale_factor;

src/model_loader.cpp

@@ -556,6 +556,10 @@ SDVersion ModelLoader::get_sd_version() const {
             tensor_storage_map.find("model.diffusion_model.audio_patch_proj.weight") != tensor_storage_map.end()) {
             return VERSION_MINIMAX_H3;
         }
+        if (name == "model.diffusion_model.y_embedder.layer_pool.weight" &&
+            tensor_storage_map.count("model.diffusion_model.modulation_cores.adaln_img.weight")) {
+            return VERSION_IRIS;
+        }
         if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
             is_wan = true;
         }

src/model_manager.cpp

@@ -1904,6 +1904,49 @@ bool ModelManager::ensure_compute_backend_capacity(
         }
     }
 
+    // The loader retains mmap buffers after storage blocks are released, so
+    // demoting mapped states would only force duplicate allocations on reload.
+    std::unordered_set<TensorState*> mapped_states;
+    for (const auto& block : params_storage_blocks_) {
+        if (!block->mmap_tensor_stores.empty()) {
+            mapped_states.insert(block->states.begin(), block->states.end());
+        }
+    }
+    std::unordered_set<TensorState*> required_set(required_states.begin(), required_states.end());
+    std::vector<TensorState*> demoted_states;
+    auto revert_demotions = [&]() {
+        for (TensorState* state : demoted_states) {
+            if (state->loaded_to_params_backend) {
+                state->residency_mode = ResidencyMode::ParamBackend;
+            }
+        }
+    };
+    auto can_demote_to_disk = [&](TensorState* state) {
+        return state != nullptr && state->tensor != nullptr && state->tensor->view_src == nullptr &&
+               state->residency_mode == ResidencyMode::ParamBackend &&
+               state->params_backend != nullptr && state->params_backend == state->compute_backend &&
+               state->compute_backend == compute_backend && !sd_backend_is_cpu(state->params_backend) &&
+               state->pin_count == 0 && !state->staged_to_compute_backend && state->loaded_to_params_backend &&
+               protected_states.find(state) == protected_states.end() &&
+               required_set.find(state) == required_set.end() &&
+               mapped_states.find(state) == mapped_states.end() &&
+               !state->sources.empty() && state->applied_lora_epoch == UINT64_MAX &&
+               state->split_buffer_type == nullptr;
+    };
+    for (TensorState* state : global_candidates) {
+        if (!can_demote_to_disk(state)) {
+            continue;
+        }
+        state->residency_mode = ResidencyMode::Disk;
+        demoted_states.push_back(state);
+        add_evictable_state(state);
+        if (release_eviction_states()) {
+            revert_demotions();
+            return true;
+        }
+    }
+    revert_demotions();
+
     const auto capacity                = check_capacity(request, required_states, true);
     const std::string available_device = capacity.available_device_bytes == SIZE_MAX
                                              ? "unknown"

src/pipeline/diffusion_engine.cpp

@@ -108,6 +108,7 @@ const char* model_version_to_str[] = {
     "PixArt",
     "Ming-Image",
     "Z-Image L2P",
+    "Iris-3B",
 };
 
 static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]),
@@ -1382,7 +1383,8 @@ bool StableDiffusionGGML::build_denoiser() {
                    sd_version_is_llada_image(version) ||
                    sd_version_is_boogu_image(version) ||
                    sd_version_is_pid(version) ||
-                   sd_version_is_ideogram4(version)) {
+                   sd_version_is_ideogram4(version) ||
+                   version == VERSION_IRIS) {
             pred_type = FLOW_PRED;
             if (sd_version_is_wan(version)) {
                 default_flow_shift = 5.f;
@@ -1390,7 +1392,7 @@ bool StableDiffusionGGML::build_denoiser() {
                 default_flow_shift = 7.f;
             } else if (sd_version_is_minimax_h3(version)) {
                 default_flow_shift = 12.f;
-            } else if (sd_version_is_ernie_image(version)) {
+            } else if (sd_version_is_ernie_image(version) || version == VERSION_IRIS) {
                 default_flow_shift = 4.f;
             } else if (sd_version_is_pid(version)) {
                 default_flow_shift = 1.5f;
@@ -2225,6 +2227,27 @@ void StableDiffusionGGML::report_sample_progress(int step,
     }
 }
 
+static sd::Tensor<float> batch_two_condition_tensors(const sd::Tensor<float>& a, const sd::Tensor<float>& b) {
+    if (a.empty() || b.empty() || a.dim() != b.dim()) {
+        return {};
+    }
+    if (a.dim() == 1) {
+        if (a.shape() != b.shape()) {
+            return {};
+        }
+        auto batched = sd::ops::concat(a, b, 0);
+        batched.reshape_({a.shape()[0], 2});
+        return batched;
+    }
+    const int64_t batch_dim = a.dim() - 1;
+    for (int64_t d = 0; d < batch_dim; d++) {
+        if (a.shape()[d] != b.shape()[d]) {
+            return {};
+        }
+    }
+    return sd::ops::concat(a, b, static_cast<size_t>(batch_dim));
+}
+
 void StableDiffusionGGML::compute_sample_controls(const sd::Tensor<float>& control_image,
                                                   const sd::Tensor<float>& noised_input,
                                                   const sd::Tensor<float>& timesteps_tensor,
@@ -2602,6 +2625,57 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
             return output_opt;
         };
 
+        auto run_batched_condition = [&](const SDCondition& condition,
+                                         const sd::Tensor<float>* c_concat_override) -> sd::Tensor<float> {
+            const sd::Tensor<float>& condition_concat =
+                c_concat_override != nullptr ? *c_concat_override : condition.c_concat;
+
+            sd::Tensor<float> batched_context = batch_two_condition_tensors(condition.c_crossattn, uncond.c_crossattn);
+            sd::Tensor<float> batched_y       = batch_two_condition_tensors(condition.c_vector, uncond.c_vector);
+            sd::Tensor<float> batched_concat  = batch_two_condition_tensors(condition_concat, uncond.c_concat);
+            if (!condition.c_crossattn.empty() && batched_context.empty()) {
+                return {};
+            }
+            if ((!condition.c_vector.empty() || !uncond.c_vector.empty()) && batched_y.empty()) {
+                return {};
+            }
+            if ((!condition_concat.empty() || !uncond.c_concat.empty()) && batched_concat.empty()) {
+                return {};
+            }
+
+            std::vector<sd::Tensor<float>> uncond_controls;
+            compute_sample_controls(control_image, noised_input, timesteps_tensor, uncond, &uncond_controls);
+            if (controls.size() != uncond_controls.size()) {
+                return {};
+            }
+            std::vector<sd::Tensor<float>> batched_controls;
+            batched_controls.reserve(controls.size());
+            for (size_t i = 0; i < controls.size(); i++) {
+                sd::Tensor<float> batched_control = batch_two_condition_tensors(controls[i], uncond_controls[i]);
+                if (batched_control.empty()) {
+                    return {};
+                }
+                batched_controls.push_back(std::move(batched_control));
+            }
+
+            sd::Tensor<float> batched_x =
+                sd::ops::concat(noised_input, noised_input, static_cast<size_t>(noised_input.dim() - 1));
+
+            DiffusionParams batched_params = diffusion_params;
+            batched_params.x               = &batched_x;
+            batched_params.context         = batched_context.empty() ? nullptr : &batched_context;
+            batched_params.c_concat        = batched_concat.empty() ? nullptr : &batched_concat;
+            batched_params.y               = batched_y.empty() ? nullptr : &batched_y;
+            batched_params.ref_latents     = nullptr;
+            batched_params.extra           = UNetDiffusionExtra{1, &batched_controls, control_strength};
+
+            sd::Tensor<float> output = work_diffusion_model->compute(n_threads, batched_params);
+            if (output.empty()) {
+                LOG_ERROR("batched diffusion model compute failed");
+            }
+            return output;
+        };
+
         const SDCondition* positive_condition      = &cond;
         const sd::Tensor<float>* c_concat_override = nullptr;
         for (const auto& extension : generation_extensions) {
@@ -2641,12 +2715,40 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
             }
         }
 
-        cond_out = run_condition(*positive_condition, c_concat_override);
-        if (cond_out.empty()) {
-            return {};
+        const bool batch_cfg_ok = config_->params.batched_cfg &&
+                                  sd_version_is_unet(version) &&
+                                  !uncond.empty() &&
+                                  img_uncond.empty() &&
+                                  !skip_uncond &&
+                                  !cache_runtime.ucache_enabled() &&
+                                  !(is_skiplayer_step && slg_uncond) &&
+                                  ip_adapter_tokens.empty() &&
+                                  ip_adapter_uncond_tokens.empty() &&
+                                  !config_->animatediff_loaded &&
+                                  (noised_input.dim() < 4 || noised_input.shape()[3] <= 1) &&
+                                  std::none_of(generation_extensions.begin(),
+                                               generation_extensions.end(),
+                                               [](const std::shared_ptr<GenerationExtension>& extension) {
+                                                   return extension->is_enabled();
+                                               });
+
+        if (batch_cfg_ok) {
+            sd::Tensor<float> batched_out = run_batched_condition(*positive_condition, c_concat_override);
+            if (!batched_out.empty() && batched_out.dim() >= 4 && batched_out.shape()[3] == 2) {
+                auto parts = sd::ops::chunk(batched_out, 2, 3);
+                cond_out   = std::move(parts[0]);
+                uncond_out = std::move(parts[1]);
+            }
         }
 
-        if (!uncond.empty()) {
+        if (cond_out.empty()) {
+            cond_out = run_condition(*positive_condition, c_concat_override);
+            if (cond_out.empty()) {
+                return {};
+            }
+        }
+
+        if (uncond_out.empty() && !uncond.empty()) {
             if (!skip_uncond) {
                 const std::vector<int>* uncond_skip_layers = nullptr;
                 if (is_skiplayer_step && slg_uncond) {
@@ -2765,7 +2867,7 @@ int StableDiffusionGGML::get_diffusion_model_down_factor() {
     if (sd_version_is_dit(version)) {
         if (sd_version_is_sensenova_u1(version)) {
             down_factor = 32;
-        } else if (sd_version_is_z_image_l2p(version)) {
+        } else if (sd_version_is_z_image_l2p(version) || version == VERSION_IRIS) {
             down_factor = 16;
         } else if (version == VERSION_QWEN_IMAGE_2_1 || version == VERSION_MING_IMAGE || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version) || sd_version_is_pixart(version)) {
             down_factor = 2;
@@ -2799,7 +2901,7 @@ int StableDiffusionGGML::get_latent_channel() {
             latent_channel = 3;
         } else if (sd_version_is_z_image_l2p(version)) {
             latent_channel = 3;
-        } else if (sd_version_is_pid(version)) {
+        } else if (sd_version_is_pid(version) || version == VERSION_IRIS) {
             latent_channel = 3;
         } else if (sd_version_is_sefi_image(version)) {
             latent_channel = 144;
@@ -2900,7 +3002,7 @@ sd::Tensor<float> StableDiffusionGGML::encode_first_stage(const sd::Tensor<float
 }
 
 sd::Tensor<float> StableDiffusionGGML::decode_first_stage(const sd::Tensor<float>& x, bool decode_video) {
-    if (sd_version_is_pid(version) || sd_version_is_minit2i(version)) {
+    if (sd_version_is_pid(version) || sd_version_is_minit2i(version) || version == VERSION_IRIS) {
         return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f);
     }
     auto latents                      = first_stage_model->diffusion_to_vae_latents(x);

src/pipeline/model_builders.cpp

@@ -18,6 +18,7 @@
 #include "model/diffusion/hidream_o1.hpp"
 #include "model/diffusion/hunyuan.hpp"
 #include "model/diffusion/ideogram4.hpp"
+#include "model/diffusion/iris.hpp"
 #include "model/diffusion/krea2.hpp"
 #include "model/diffusion/lens.hpp"
 #include "model/diffusion/lingbot_video.hpp"
@@ -456,6 +457,11 @@ namespace sd::model_builders {
                                                                   tensor_storage_map,
                                                                   "model.diffusion_model",
                                                                   weight_manager);
+        } else if (version == VERSION_IRIS) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map, version, "", false, weight_manager, tokenizers);
+            result.diffusion   = std::make_shared<Iris::IrisRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                  tensor_storage_map, "model.diffusion_model", weight_manager);
         } else {  // SD1.x SD2.x SDXL
             std::map<std::string, std::string> embbeding_map;
             for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) {
@@ -527,7 +533,7 @@ namespace sd::model_builders {
         }
 
         auto create_tae = [&](bool decode_only) -> std::shared_ptr<VAE> {
-            if (sd_version_uses_wan_vae(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) {
+            if ((sd_version_uses_wan_vae(version) && version != VERSION_QWEN_IMAGE_2_1) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) {
                 return std::make_shared<TinyVideoAutoEncoder>(ctx.backends.runtime_backend(SDBackendModule::VAE),
                                                               tensor_storage_map,
                                                               "decoder",
@@ -626,7 +632,12 @@ namespace sd::model_builders {
             }
         };
 
-        if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version) || sd_version_is_sensenova_u1(version) || sd_version_is_z_image_l2p(version)) {
+        if (version == VERSION_CHROMA_RADIANCE ||
+            version == VERSION_HIDREAM_O1 ||
+            sd_version_is_minit2i(version) ||
+            sd_version_is_sensenova_u1(version) ||
+            sd_version_is_z_image_l2p(version) ||
+            version == VERSION_IRIS) {
             LOG_INFO("using FakeVAE");
             result.vae = std::make_shared<FakeVAE>(version,
                                                    ctx.backends.runtime_backend(SDBackendModule::VAE),

src/stable-diffusion.cpp

@@ -335,6 +335,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_params->max_vram                  = nullptr;
     sd_ctx_params->disable_prefetch          = false;
     sd_ctx_params->disable_segmented_compute = false;
+    sd_ctx_params->batched_cfg               = true;
     sd_ctx_params->eager_load                = false;
     sd_ctx_params->enable_mmap               = false;
     sd_ctx_params->diffusion_flash_attn      = false;

src/upscaler.cpp

@@ -25,15 +25,40 @@ UpscalerGGML::~UpscalerGGML() {
     // ModelManager holds raw ggml tensor pointers owned by the runner context.
     model_manager.reset();
     esrgan_upscaler.reset();
+    mosr_upscaler.reset();
 }
 
 void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) {
     max_graph_vram_bytes = max_vram_bytes;
-    if (esrgan_upscaler) {
-        esrgan_upscaler->set_max_graph_vram_bytes(max_vram_bytes);
+    if (runner() != nullptr) {
+        runner()->set_max_graph_vram_bytes(max_vram_bytes);
     }
 }
 
+GGMLRunner* UpscalerGGML::runner() {
+    if (mosr_upscaler) {
+        return mosr_upscaler.get();
+    }
+    return esrgan_upscaler.get();
+}
+
+int UpscalerGGML::get_scale() {
+    if (mosr_upscaler) {
+        return mosr_upscaler->config.scale;
+    }
+    if (esrgan_upscaler) {
+        return esrgan_upscaler->config.scale;
+    }
+    return 1;
+}
+
+sd::Tensor<float> UpscalerGGML::compute(const sd::Tensor<float>& input_tensor) {
+    if (mosr_upscaler) {
+        return mosr_upscaler->compute(n_threads, input_tensor);
+    }
+    return esrgan_upscaler->compute(n_threads, input_tensor);
+}
+
 bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
                                   int n_threads) {
     sd::ParallelScope tensor_scope(&tensor_executor);
@@ -82,20 +107,33 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
     }
     model_loader.set_wtype_override(model_data_type);
     LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type));
-    esrgan_upscaler = std::make_shared<ESRGAN>(backend_for(SDBackendModule::UPSCALER),
-                                               model_loader.get_tensor_storage_map(),
+    const auto& tensor_storage_map = model_loader.get_tensor_storage_map();
+    std::map<std::string, ggml_tensor*> tensors;
+    if (tensor_storage_map.count("gblocks.0.weight") > 0) {
+        mosr_upscaler = std::make_shared<MoSR>(backend_for(SDBackendModule::UPSCALER),
+                                               tensor_storage_map,
                                                model_manager);
-    if (esrgan_upscaler == nullptr || esrgan_upscaler->rrdb_net == nullptr) {
-        LOG_ERROR("init esrgan model from metadata failed: '%s'", esrgan_path.c_str());
-        return false;
+        if (mosr_upscaler->mosr_net == nullptr) {
+            LOG_ERROR("init mosr model from metadata failed: '%s'", esrgan_path.c_str());
+            return false;
+        }
+        mosr_upscaler->get_param_tensors(tensors);
+    } else {
+        esrgan_upscaler = std::make_shared<ESRGAN>(backend_for(SDBackendModule::UPSCALER),
+                                                   tensor_storage_map,
+                                                   model_manager);
+        if (esrgan_upscaler->rrdb_net == nullptr) {
+            LOG_ERROR("init esrgan model from metadata failed: '%s'", esrgan_path.c_str());
+            return false;
+        }
+        esrgan_upscaler->get_param_tensors(tensors);
     }
-    esrgan_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes);
+    LOG_INFO("Upscaler architecture: %s", runner()->get_desc().c_str());
+    runner()->set_max_graph_vram_bytes(max_graph_vram_bytes);
     if (direct) {
-        esrgan_upscaler->set_conv2d_direct_enabled(true);
+        runner()->set_conv2d_direct_enabled(true);
     }
 
-    std::map<std::string, ggml_tensor*> tensors;
-    esrgan_upscaler->get_param_tensors(tensors);
     if (!model_manager->set_loader(model_loader) ||
         !model_manager->register_param_tensors(ModelComponent::Upscaler,
                                                std::move(tensors),
@@ -103,7 +141,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
                                                backend_for(SDBackendModule::UPSCALER),
                                                params_backend_for(SDBackendModule::UPSCALER)) ||
         !model_manager->validate_registered_tensors()) {
-        LOG_ERROR("register esrgan tensors with model manager failed");
+        LOG_ERROR("register upscaler tensors with model manager failed");
         return false;
     }
     return true;
@@ -113,7 +151,7 @@ sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
     sd::ParallelScope tensor_scope(&tensor_executor);
     if (input_tensor.empty() || input_tensor.dim() != 4 ||
         (input_tensor.shape()[2] != 3 && input_tensor.shape()[2] != 4)) {
-        LOG_ERROR("esrgan expects a 4D RGB or RGBA image tensor");
+        LOG_ERROR("upscaler expects a 4D RGB or RGBA image tensor");
         return {};
     }
 
@@ -124,14 +162,14 @@ sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
     }
     const sd::Tensor<float>& model_input = has_alpha ? rgb : input_tensor;
     sd::Tensor<float> upscaled;
-    const int scale = esrgan_upscaler->config.scale;
+    const int scale = get_scale();
     if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) {
-        upscaled = esrgan_upscaler->compute(n_threads, model_input);
+        upscaled = compute(model_input);
     } else {
         auto on_processing = [&](const sd::Tensor<float>& input_tile) -> sd::Tensor<float> {
-            auto output_tile = esrgan_upscaler->compute(n_threads, input_tile);
+            auto output_tile = compute(input_tile);
             if (output_tile.empty()) {
-                LOG_ERROR("esrgan compute failed while processing a tile");
+                LOG_ERROR("upscaler compute failed while processing a tile");
                 return {};
             }
             return output_tile;
@@ -148,9 +186,9 @@ sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
                                     false,
                                     on_processing);
     }
-    esrgan_upscaler->runner_end();
+    runner()->runner_end();
     if (upscaled.empty()) {
-        LOG_ERROR("esrgan compute failed");
+        LOG_ERROR("upscaler compute failed");
         return {};
     }
     if (has_alpha) {
@@ -168,7 +206,7 @@ sd_image_t UpscalerGGML::upscale(sd_image_t input_image, uint32_t upscale_factor
     sd::ParallelScope tensor_scope(&tensor_executor);
     // upscale_factor, unused for RealESRGAN_x4plus_anime_6B.pth
     sd_image_t upscaled_image = {0, 0, 0, nullptr};
-    const int scale           = esrgan_upscaler->config.scale;
+    const int scale           = get_scale();
     int output_width          = (int)input_image.width * scale;
     int output_height         = (int)input_image.height * scale;
     LOG_INFO("upscaling from (%i x %i) to (%i x %i)",
@@ -256,10 +294,10 @@ bool upscale(upscaler_ctx_t* upscaler_ctx,
 }
 
 int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) {
-    if (upscaler_ctx == nullptr || upscaler_ctx->upscaler == nullptr || upscaler_ctx->upscaler->esrgan_upscaler == nullptr) {
+    if (upscaler_ctx == nullptr || upscaler_ctx->upscaler == nullptr) {
         return 1;
     }
-    return upscaler_ctx->upscaler->esrgan_upscaler->config.scale;
+    return upscaler_ctx->upscaler->get_scale();
 }
 
 int get_upscaler_model_scale(const char* model_path) {
@@ -272,8 +310,12 @@ int get_upscaler_model_scale(const char* model_path) {
             return 0;
         }
         const auto& tensors = loader.get_tensor_storage_map();
-        auto first          = tensors.find("conv_first.weight");
-        auto last           = tensors.find("conv_last.weight");
+        if (tensors.count("gblocks.0.weight") > 0) {
+            MoSRConfig config = MoSRConfig::detect_from_weights(tensors);
+            return config.is_valid() ? config.scale : 0;
+        }
+        auto first = tensors.find("conv_first.weight");
+        auto last  = tensors.find("conv_last.weight");
         if (first == tensors.end() || last == tensors.end() ||
             tensors.count("body.0.rdb1.conv1.weight") == 0 ||
             first->second.n_dims != 4 || last->second.n_dims != 4 ||

src/upscaler.h

@@ -4,6 +4,7 @@
 #include "core/ggml_extend_backend.h"
 #include "core/tensor.hpp"
 #include "model/upscaler/esrgan.hpp"
+#include "model/upscaler/mosr.hpp"
 #include "model_manager.h"
 #include "stable-diffusion.h"
 
@@ -15,6 +16,7 @@ struct UpscalerGGML {
     std::shared_ptr<ModelManager> model_manager;
     ggml_type model_data_type = GGML_TYPE_F16;
     std::shared_ptr<ESRGAN> esrgan_upscaler;
+    std::shared_ptr<MoSR> mosr_upscaler;
     std::string esrgan_path;
     int n_threads;
     sd::ParallelExecutor tensor_executor;
@@ -34,6 +36,9 @@ struct UpscalerGGML {
     bool load_from_file(const std::string& esrgan_path,
                         int n_threads);
     void set_max_graph_vram_bytes(size_t max_vram_bytes);
+    GGMLRunner* runner();
+    int get_scale();
+    sd::Tensor<float> compute(const sd::Tensor<float>& input_tensor);
     sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
     sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
 };