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

Co-authored-by: leejet <leejet714@gmail.com>

.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

@@ -52,6 +52,7 @@ API and command-line option may change frequently.***
     - [PiD](./docs/pid.md)
     - [LongCat Image](./docs/longcat_image.md)
     - [Z-Image](./docs/z_image.md)
+    - [Z-Image L2P](./docs/z_image_l2p.md)
     - [MiniT2I](./docs/minit2i.md)
     - [SenseNova U1.5](./docs/sensenova_u1.md)
     - [Ovis-Image](./docs/ovis_image.md)
@@ -66,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)
@@ -181,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`.

docs/z_image_l2p.md

@@ -0,0 +1,21 @@
+# Z-Image L2P
+
+[L2P](https://nju-pcalab.github.io/projects/L2P/) (Latent-to-Pixel) transfers Z-Image-Turbo to pixel space. The VAE is replaced by 16x16 patch tokenization on RGB pixels and a small convolutional local decoder that turns the last DiT hidden states back into pixels, so no VAE is needed. The text encoder is the Qwen3-4B used by Z-Image-Turbo.
+
+## Download weights
+
+- Download L2P (1K)
+    - safetensors: https://huggingface.co/zhen-nan/L2P/tree/main
+- Download Qwen3 4b
+    - safetensors: https://huggingface.co/Comfy-Org/z_image_turbo/tree/main/split_files/text_encoders
+    - gguf: https://huggingface.co/unsloth/Qwen3-4B-Instruct-2507-GGUF/tree/main
+
+## Text-to-image
+
+Do not pass `--vae`. Image dimensions must be multiples of 16. The reference pipeline uses 30 steps with a CFG scale of 2.0.
+
+The released checkpoint keeps part of the transformer in F32. Pass `--type bf16` to load it in BF16, as the reference pipeline runs it; this reduces the diffusion model from about 18.6 GB to 11.8 GB.
+
+```bash
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\model-1k-merge.safetensors --llm ..\models\text_encoders\qwen_3_4b.safetensors -p "an origami pig on fire in the middle of a dark room with a pentagram on the floor" --cfg-scale 2.0 --steps 30 --sampling-method euler -W 1024 -H 1024 --diffusion-fa --type bf16 -v -o z_image_l2p.png
+```

examples/cli/main.cpp

@@ -716,6 +716,15 @@ int main(int argc, const char* argv[]) {
     }
 
     if (cli_params.mode == CONVERT) {
+        std::vector<sd_lora_t> loras;
+        loras.reserve(gen_params.lora_map.size() + gen_params.high_noise_lora_map.size());
+        for (const auto& [path, multiplier] : gen_params.lora_map) {
+            loras.push_back({false, multiplier, path.c_str()});
+        }
+        for (const auto& [path, multiplier] : gen_params.high_noise_lora_map) {
+            loras.push_back({true, multiplier, path.c_str()});
+        }
+
         bool success = convert_with_components(ctx_params.model_path.c_str(),
                                                ctx_params.clip_l_path.c_str(),
                                                ctx_params.clip_g_path.c_str(),
@@ -726,7 +735,9 @@ int main(int argc, const char* argv[]) {
                                                ctx_params.wtype,
                                                ctx_params.tensor_type_rules.c_str(),
                                                cli_params.convert_name,
-                                               ctx_params.n_threads);
+                                               ctx_params.n_threads,
+                                               loras.empty() ? nullptr : loras.data(),
+                                               (int)loras.size());
         if (!success) {
             LOG_ERROR("convert '%s'/'%s' to '%s' failed",
                       ctx_params.model_path.c_str(),

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",
@@ -2271,14 +2284,15 @@ bool SDGenerationParams::from_json_str(
     return true;
 }
 
-void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_dir) {
+bool SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_dir) {
     if (lora_model_dir.empty()) {
-        return;
+        return true;
     }
     static const std::regex re(R"(<lora:([^:>]+):([^>]+)>)");
     static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"};
     std::smatch m;
 
+    bool missing    = false;
     std::string tmp = prompt;
 
     while (std::regex_search(tmp, m, re)) {
@@ -2320,8 +2334,9 @@ void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_d
             }
             if (!found) {
                 LOG_WARN("can not found lora %s", final_path.lexically_normal().string().c_str());
-                tmp    = m.suffix().str();
-                prompt = std::regex_replace(prompt, re, "", std::regex_constants::format_first_only);
+                tmp     = m.suffix().str();
+                prompt  = std::regex_replace(prompt, re, "", std::regex_constants::format_first_only);
+                missing = true;
                 continue;
             }
         }
@@ -2337,6 +2352,8 @@ void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_d
 
         tmp = m.suffix().str();
     }
+
+    return !missing;
 }
 
 bool SDGenerationParams::width_and_height_are_set() const {
@@ -2458,7 +2475,7 @@ bool SDGenerationParams::initialize_cache_params() {
     return true;
 }
 
-bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict) {
+bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict, bool validate_missing_loras) {
     vae_tiling_params.extra_tiling_args = extra_tiling_args.empty() ? nullptr : extra_tiling_args.c_str();
 
     if (high_noise_sample_params.sample_steps <= 0) {
@@ -2502,7 +2519,9 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s
 
     prompt_with_lora = prompt;
     if (!lora_model_dir.empty()) {
-        extract_and_remove_lora(lora_model_dir);
+        if (!extract_and_remove_lora(lora_model_dir) && validate_missing_loras) {
+            return false;
+        }
     }
     return true;
 }
@@ -2630,7 +2649,7 @@ bool SDGenerationParams::resolve_and_validate(SDMode mode,
                                               const std::string& lora_model_dir,
                                               const std::string& hires_upscalers_dir,
                                               bool strict) {
-    if (!resolve(lora_model_dir, hires_upscalers_dir, strict)) {
+    if (!resolve(lora_model_dir, hires_upscalers_dir, strict, mode == CONVERT)) {
         return false;
     }
     if (!validate(mode)) {

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;
@@ -313,12 +314,12 @@ struct SDGenerationParams {
                        const std::function<std::string(const std::string&)>& lora_path_resolver = {});
     bool parse_image_preprocess_json(const std::string& json_str);
     bool initialize_cache_params();
-    void extract_and_remove_lora(const std::string& lora_model_dir);
+    bool extract_and_remove_lora(const std::string& lora_model_dir);
     bool width_and_height_are_set() const;
     void set_width_and_height_if_unset(int w, int h);
     int get_resolved_width() const;
     int get_resolved_height() const;
-    bool resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict = false);
+    bool resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict = false, bool validate_missing_loras = false);
     bool validate(SDMode mode);
     bool resolve_and_validate(SDMode mode,
                               const std::string& lora_model_dir,

examples/common/log.cpp

@@ -91,34 +91,34 @@ void log_print(enum sd_log_level_t level, const char* log, sd_log_level_t min_le
     switch (level) {
         case SD_LOG_DEBUG:
             tag_color = 37;
-            level_str = "DEBUG";
+            level_str = "D";
             break;
         case SD_LOG_VERBOSE:
             tag_color = 37;
-            level_str = "VERBOSE";
+            level_str = "V";
             break;
         case SD_LOG_INFO:
             tag_color = 34;
-            level_str = "INFO";
+            level_str = "I";
             break;
         case SD_LOG_WARN:
             tag_color = 35;
-            level_str = "WARN";
+            level_str = "W";
             break;
         case SD_LOG_ERROR:
             tag_color = 31;
-            level_str = "ERROR";
+            level_str = "E";
             break;
         default:
             tag_color = 33;
-            level_str = "?????";
+            level_str = "?";
             break;
     }
 
     if (color) {
-        fprintf(out_stream, "\033[%d;1m[%-7s]\033[0m ", tag_color, level_str);
+        fprintf(out_stream, "\033[%d;1m[%s]\033[0m ", tag_color, level_str);
     } else {
-        fprintf(out_stream, "[%-7s] ", level_str);
+        fprintf(out_stream, "[%s] ", level_str);
     }
     fflush(out_stream);
     print_utf8(out_stream, log);
@@ -131,18 +131,15 @@ void example_log_printf(sd_log_level_t level, const char* file, int line, const
     va_list args;
     va_start(args, format);
 
-    static char log_buffer[LOG_BUFFER_SIZE + 1];
-    int written = snprintf(log_buffer, LOG_BUFFER_SIZE, "%s:%-4d - ", sd_basename(file).c_str(), line);
-
-    if (written >= 0 && written < static_cast<int>(LOG_BUFFER_SIZE)) {
-        vsnprintf(log_buffer + written, LOG_BUFFER_SIZE - written, format, args);
-    }
-    size_t len = strlen(log_buffer);
-    if (len == 0 || log_buffer[len - 1] != '\n') {
-        strncat(log_buffer, "\n", LOG_BUFFER_SIZE - len);
-    }
-
-    log_print(level, log_buffer, log_level, log_color);
-
+    char log_buffer[LOG_BUFFER_SIZE + 1] = {};
+    vsnprintf(log_buffer, sizeof(log_buffer), format, args);
     va_end(args);
+
+    std::string message = log_buffer;
+    while (!message.empty() && (message.back() == '\n' || message.back() == '\r')) {
+        message.pop_back();
+    }
+    message += " --- " + sd_basename(file) + ":" + std::to_string(line) + '\n';
+
+    log_print(level, message.c_str(), log_level, log_color);
 }

examples/server/api.md

@@ -37,6 +37,7 @@ Current generation-related endpoints include:
 
 - `POST /sdapi/v1/txt2img`
 - `POST /sdapi/v1/img2img`
+- `GET /sdapi/v1/progress`
 - `GET /sdapi/v1/loras`
 - `GET /sdapi/v1/upscalers`
 - `GET /sdapi/v1/latent-upscale-modes`
@@ -278,6 +279,26 @@ Response fields:
 | `parameters` | `object` | Echo of the parsed outer request body |
 | `info` | `string` | Currently empty string |
 
+#### `GET /sdapi/v1/progress`
+
+Poll this endpoint while a synchronous SDAPI generation request is running.
+An optional `id_task` query parameter selects the task with the same `id_task`
+provided in the generation request (default: `"sdapi"`). Set
+`skip_current_image=true` or `1` to omit the preview image.
+
+`current_image` contains the latest base64 JPEG latent projection, or `null`
+when unavailable. `state.sampling_step` is the positive logical step and
+`state.sampling_steps` is the actual step count for the current sampling pass,
+including schedule and img2img strength adjustments. `state.job_no` is the
+zero-based pass index; `state.job_count` includes one pass per batch image and
+an additional pass per image when highres fix is enabled. `progress` weights
+these passes equally. Sampling progress can reach `1` before final decoding
+finishes; `eta_relative` is currently always `0`.
+
+When no SDAPI task is active, or `id_task` does not match, the endpoint returns
+`progress=0`, `state.job=""`, zero job/step counts, and `current_image=null`.
+Completed and failed requests do not retain an active preview.
+
 #### Discovery / Compatibility Endpoints
 
 Currently exposed:
@@ -410,6 +431,20 @@ Field types:
 | `queue_position` | `integer` |
 | `result` | `object \| null` |
 | `error` | `object \| null` |
+| `preview` | `object \| null` |
+
+`preview` sub-fields:
+
+| Field | Type | Notes |
+| --- | --- | --- |
+| `preview.pass` | `integer` | Sampling pass number, starting at 1 within the job |
+| `preview.step` | `integer` | Positive logical sample step within this pass |
+| `preview.total_steps` | `integer` | Actual sample steps in this pass, after schedule and strength adjustments |
+| `preview.b64_json` | `string` | Base64-encoded PNG of the current preview frame |
+
+The preview is updated in place as sampling progresses; poll `GET /sdcpp/v1/jobs/{id}` to retrieve the latest frame. Only the most recent preview is retained.
+
+`step / total_steps` describes the current sampling pass, not overall job completion. A new batch item, high/low-noise stage, or highres pass increments `pass` and restarts `step`. Samplers with multiple denoiser evaluations per logical step may update a preview more than once at the same step. The frame and its pass/step metadata are published together.
 
 ### Endpoints
 
@@ -437,6 +472,7 @@ Top-level fields:
 | `loras` | `array<object>` | Available LoRA entries |
 | `upscalers` | `array<object>` | Available highres upscalers, built-in and model-backed |
 | `upscale` | `boolean` | Whether a compatible RGB ESRGAN model is available for `POST /sdcpp/v1/upscale` |
+| `preview_modes` | `array<string>` | Available preview modes, e.g. `["none", "proj", "tae", "vae"]` |
 | `limits` | `object` | Shared queue and size limits |
 
 `model`
@@ -598,6 +634,7 @@ Fields returned in `features_by_mode.img_gen`:
 - `cache`
 - `cancel_queued`
 - `cancel_generating`
+- `preview`
 
 Fields returned in `features_by_mode.vid_gen`:
 
@@ -610,6 +647,7 @@ Fields returned in `features_by_mode.vid_gen`:
 - `cache`
 - `cancel_queued`
 - `cancel_generating`
+- `preview`
 
 #### `POST /sdcpp/v1/img_gen`
 
@@ -782,6 +820,9 @@ Example:
   "scm_mask": "",
   "scm_policy_dynamic": true,
 
+  "preview": "none",
+  "preview_interval": 1,
+
   "output_format": "png",
   "output_compression": 100
 }
@@ -847,6 +888,8 @@ Top-level scalar fields:
 | `control_strength` | `number` |
 | `ip_adapter_strength` | `number` |
 | `embed_image_metadata` | `boolean` |
+| `preview` | `string` |
+| `preview_interval` | `integer` |
 
 Image fields:
 
@@ -934,6 +977,10 @@ When omitted, backend defaults apply to these fields:
 - `sample_params.flow_shift`
 - `sample_params.guidance.img_cfg`
 
+### Preview Interval Semantics
+
+`preview_interval` controls the period (in sample steps) at which preview frames are generated. The default is `1` (every step). Any non-positive values will be clamped to `1`. Note that if `preview` is set to `"none"`, no previews are produced regardless of the interval.
+
 ### Completion Result
 
 Example completed job:
@@ -1048,7 +1095,8 @@ Response fields:
 Compared with `img_gen`, the `vid_gen` request body:
 
 - `vid_gen` is a single video sequence job, so `batch_count` is not part of the request schema
-- `ref_images`, `mask_image`, `control_image`, `control_strength`, `ip_adapter_image`, `ip_adapter_strength`, and `embed_image_metadata` are not part of the request schema
+- `mask_image`, `control_image`, `control_strength`, `ip_adapter_image`, `ip_adapter_strength`, and `embed_image_metadata` are not part of the request schema
+- `ref_images` is accepted for MiniMax-H3 Ref2VA conditioning; other video model families currently ignore it
 - `vid_gen` adds `end_image`, `control_frames`, `high_noise_sample_params`, `video_frames`, `fps`, `moe_boundary`, and `vace_strength`
 
 Example:
@@ -1069,6 +1117,7 @@ Example:
 
   "init_image": null,
   "end_image": null,
+  "ref_images": [],
   "control_frames": [],
 
   "sample_params": {
@@ -1130,11 +1179,25 @@ Example:
   "scm_mask": "",
   "scm_policy_dynamic": true,
 
+  "preview": "none",
+  "preview_interval": 1,
+
   "output_format": "webm",
   "output_compression": 100
 }
 ```
 
+### Reference Image Rules
+
+- `ref_images` contains reference images for MiniMax-H3 Ref2VA conditioning.
+- Images retain request order and correspond to `<Picture 1>`, `<Picture 2>`, and so on in the prompt.
+- For MiniMax-H3, non-empty `ref_images` cannot be combined with `init_image` or `end_image`. A conflicting API request fails during generation; the WebUI checks this before submission.
+- MiniMax-H3 does not support `control_frames`; leave that array empty.
+- Other video model families currently ignore `ref_images`.
+- `features_by_mode.vid_gen` does not currently advertise `ref_images`; its absence is not an indication that MiniMax-H3 reference images are unsupported.
+
+See [MiniMax-H3 reference conditioning](../../docs/minimax_h3.md#reference-to-audio-video-conditioning) for model requirements and prompt examples.
+
 ### LoRA Rules
 
 - The server only accepts explicit LoRA entries from the `lora` field.
@@ -1152,6 +1215,7 @@ Channel expectations:
 
 - `init_image`: 3 channels
 - `end_image`: 3 channels
+- `ref_images[]`: decoded with native channels, then converted to RGB by MiniMax-H3
 - `control_frames[]`: 3 channels
 
 Frame ordering rules:
@@ -1181,6 +1245,8 @@ Top-level scalar fields:
 | `fps` | `integer` |
 | `moe_boundary` | `number` |
 | `vace_strength` | `number` |
+| `preview` | `string` |
+| `preview_interval` | `integer` |
 
 Image and frame fields:
 
@@ -1188,6 +1254,7 @@ Image and frame fields:
 | --- | --- |
 | `init_image` | `string \| null` |
 | `end_image` | `string \| null` |
+| `ref_images` | `array<string>` |
 | `control_frames` | `array<string>` |
 
 LoRA fields:

examples/server/async_jobs.cpp

@@ -10,6 +10,55 @@
 #include "common/media_io.h"
 #include "common/resource_owners.hpp"
 
+static void preview_callback(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data) {
+    (void)is_noisy;
+    if (frame_count <= 0 || frames == nullptr || frames[0].data == nullptr) {
+        return;
+    }
+
+    AsyncGenerationJob* job      = static_cast<AsyncGenerationJob*>(data);
+    const sd_image_t& frame      = frames[0];
+    const sd_preview_info_t info = sd_get_preview_info();
+
+    auto image_bytes = encode_image_to_vector(EncodedImageFormat::PNG,
+                                              frame.data,
+                                              frame.width,
+                                              frame.height,
+                                              frame.channel,
+                                              "",
+                                              80);
+    if (image_bytes.empty()) {
+        return;
+    }
+
+    std::string b64 = base64_encode(image_bytes);
+    std::lock_guard<std::mutex> lock(job->preview_mutex);
+    job->preview_b64         = std::move(b64);
+    job->preview_step        = step < 0 ? -step : step;
+    job->preview_pass        = info.sample_pass;
+    job->preview_total_steps = info.total_steps;
+}
+
+static void clear_preview_callback() {
+    sd_set_preview_callback(nullptr, PREVIEW_NONE, 1, false, false, nullptr);
+}
+
+static void set_preview_callback(AsyncGenerationJob& job) {
+    preview_t mode = str_to_preview(job.kind == AsyncJobKind::ImgGen
+                                        ? job.img_gen.preview_mode.c_str()
+                                        : job.vid_gen.preview_mode.c_str());
+    if (mode == PREVIEW_COUNT || mode == PREVIEW_NONE) {
+        return;
+    }
+    int interval = job.kind == AsyncJobKind::ImgGen
+                       ? job.img_gen.preview_interval
+                       : job.vid_gen.preview_interval;
+    if (interval <= 0) {
+        interval = 1;
+    }
+    sd_set_preview_callback(preview_callback, mode, interval, true, false, &job);
+}
+
 const char* async_job_kind_name(AsyncJobKind kind) {
     switch (kind) {
         case AsyncJobKind::ImgGen:
@@ -161,6 +210,18 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo
         result["error"]  = nullptr;
     }
 
+    if (job.status == AsyncJobStatus::Generating) {
+        std::lock_guard<std::mutex> lock(job.preview_mutex);
+        if (!job.preview_b64.empty()) {
+            result["preview"] = {
+                {"pass", job.preview_pass},
+                {"step", job.preview_step},
+                {"total_steps", job.preview_total_steps},
+                {"b64_json", job.preview_b64},
+            };
+        }
+    }
+
     return result;
 }
 
@@ -174,12 +235,14 @@ bool execute_img_gen_job(ServerRuntime& runtime,
 
     {
         std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
+        set_preview_callback(job);
         sd_image_t* raw_results = nullptr;
         int num_results         = 0;
         if (!generate_image(runtime.sd_ctx, &params, &raw_results, &num_results)) {
             raw_results = nullptr;
             num_results = 0;
         }
+        clear_preview_callback();
         results.adopt(raw_results, num_results);
     }
 
@@ -247,10 +310,12 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
 
     {
         std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
+        set_preview_callback(job);
         sd_image_t* raw_results = nullptr;
         if (!generate_video(runtime.sd_ctx, &params, &raw_results, &num_results, &generated_audio, &output_fps)) {
             raw_results = nullptr;
         }
+        clear_preview_callback();
         results.adopt(raw_results, num_results);
     }
 

examples/server/async_jobs.h

@@ -44,6 +44,11 @@ struct AsyncGenerationJob {
     int result_fps         = 0;
     std::string error_code;
     std::string error_message;
+    mutable std::mutex preview_mutex;
+    std::string preview_b64;
+    int preview_pass          = 0;
+    int preview_step          = 0;
+    int preview_total_steps   = 0;
 };
 
 struct AsyncJobManager {

examples/server/main.cpp

@@ -75,6 +75,8 @@ int main(int argc, const char** argv) {
     }
     SDSvrParams svr_params;
     SDContextParams ctx_params;
+    // Preserve the CLI default while disabling cwd scans for sd-server.
+    ctx_params.lora_model_dir.clear();
     SDGenerationParams default_gen_params;
 
     sd_set_log_callback(sd_log_cb, (void*)&svr_params);

examples/server/routes_sdapi.cpp

@@ -301,6 +301,76 @@ static nlohmann::json prepare_info_field(const SDContextParams& ctx_params,
     return jsoninfo;
 }
 
+struct SDAPIPreviewState {
+    std::mutex mutex;
+    std::string b64;
+    std::string id_task;
+    std::string job_type;
+    int step          = 0;
+    int total         = 0;
+    bool active       = false;
+    int64_t job_count = 0;
+    int job_no        = 0;
+};
+
+static SDAPIPreviewState g_preview;
+
+static void preview_callback(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data) {
+    (void)is_noisy;
+    (void)data;
+    if (frame_count <= 0 || frames == nullptr || frames[0].data == nullptr) {
+        return;
+    }
+
+    const sd_image_t& frame      = frames[0];
+    const sd_preview_info_t info = sd_get_preview_info();
+
+    auto image_bytes = encode_image_to_vector(EncodedImageFormat::JPEG,
+                                              frame.data,
+                                              frame.width,
+                                              frame.height,
+                                              frame.channel,
+                                              "",
+                                              80);
+    if (image_bytes.empty()) {
+        return;
+    }
+
+    std::string b64 = base64_encode(image_bytes);
+    std::lock_guard<std::mutex> lock(g_preview.mutex);
+    g_preview.b64    = std::move(b64);
+    g_preview.step   = static_cast<int>(std::min<int64_t>(step < 0 ? -static_cast<int64_t>(step) : step,
+                                                        std::max(0, info.total_steps)));
+    g_preview.total  = std::max(0, info.total_steps);
+    g_preview.job_no = std::max(0, info.sample_pass - 1);
+}
+
+struct SDAPIPreviewGuard {
+    SDAPIPreviewGuard(const std::string& id_task, const sd_img_gen_params_t& params, const std::string& job_type) {
+        std::lock_guard<std::mutex> lock(g_preview.mutex);
+        g_preview.id_task  = id_task;
+        g_preview.total    = 0;
+        g_preview.job_type = job_type;
+        g_preview.b64.clear();
+        g_preview.step   = 0;
+        g_preview.active = true;
+        // Each batch item has a base pass and, when enabled, a hires pass.
+        g_preview.job_count = static_cast<int64_t>(std::max(1, params.batch_count)) * (params.hires.enabled ? 2 : 1);
+        g_preview.job_no    = 0;
+        sd_set_preview_callback(preview_callback, PREVIEW_PROJ, 1, true, false, nullptr);
+    }
+
+    ~SDAPIPreviewGuard() {
+        sd_set_preview_callback(nullptr, PREVIEW_NONE, 1, false, false, nullptr);
+        std::lock_guard<std::mutex> lock(g_preview.mutex);
+        g_preview.active = false;
+        g_preview.b64.clear();
+    }
+
+    SDAPIPreviewGuard(const SDAPIPreviewGuard&)            = delete;
+    SDAPIPreviewGuard& operator=(const SDAPIPreviewGuard&) = delete;
+};
+
 void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
     ServerRuntime* runtime = &rt;
 
@@ -317,7 +387,9 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
                 return;
             }
 
-            json j = json::parse(req.body);
+            std::string id_task = "sdapi";
+            json j              = json::parse(req.body);
+            id_task             = j.value("id_task", id_task);
             ImgGenJobRequest request;
             std::string error_message;
             if (!build_sdapi_img_gen_request(j, *runtime, img2img, request, error_message)) {
@@ -334,6 +406,7 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
 
             {
                 std::lock_guard<std::mutex> lock(*runtime->sd_ctx_mutex);
+                SDAPIPreviewGuard preview_guard(id_task, img_gen_params, img2img ? "img2img" : "txt2img");
                 sd_image_t* raw_results = nullptr;
                 if (!generate_image(runtime->sd_ctx, &img_gen_params, &raw_results, &num_results)) {
                     raw_results = nullptr;
@@ -528,4 +601,62 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
         r["sd_model_checkpoint"] = model_path.stem();
         res.set_content(r.dump(), "application/json");
     });
+
+    svr.Get("/sdapi/v1/progress", [](const httplib::Request& req, httplib::Response& res) {
+        bool skip_current_image = false;
+        bool has_id_param       = req.has_param("id_task");
+        std::string req_id_task = has_id_param ? req.get_param_value("id_task") : "";
+
+        if (req.has_param("skip_current_image")) {
+            std::string skip_str = req.get_param_value("skip_current_image");
+            skip_current_image   = (skip_str == "true" || skip_str == "1");
+        }
+
+        json r;
+        r["eta_relative"] = 0.0f;
+        json state;
+        state["skipped"]       = false;
+        state["interrupted"]   = false;
+        state["job_timestamp"] = "0";
+        r["textinfo"]          = nullptr;
+
+        {
+            std::lock_guard<std::mutex> lock(g_preview.mutex);
+
+            bool is_matching_task = g_preview.active && (!has_id_param || req_id_task == g_preview.id_task);
+
+            int step             = is_matching_task ? g_preview.step : 0;
+            int steps            = is_matching_task ? g_preview.total : 0;
+            std::string b64      = is_matching_task ? g_preview.b64 : "";
+            std::string job_type = is_matching_task ? g_preview.job_type : "";
+            int64_t job_count    = is_matching_task ? g_preview.job_count : 0;
+            int job_no           = is_matching_task ? g_preview.job_no : 0;
+
+            float batch_progress = (steps > 0) ? (static_cast<float>(step) / static_cast<float>(steps)) : 0.0f;
+            if (batch_progress > 1.0f)
+                batch_progress = 1.0f;
+
+            float overall_progress = job_count > 0
+                                         ? (static_cast<float>(job_no) + batch_progress) / static_cast<float>(job_count)
+                                         : 0.0f;
+            if (overall_progress > 1.0f)
+                overall_progress = 1.0f;
+
+            r["progress"]           = overall_progress;
+            state["job"]            = job_type;
+            state["sampling_step"]  = step;
+            state["sampling_steps"] = steps;
+            state["job_count"]      = job_count;
+            state["job_no"]         = job_no;
+            r["state"]              = state;
+
+            if (!skip_current_image && !b64.empty()) {
+                r["current_image"] = b64;
+            } else {
+                r["current_image"] = nullptr;
+            }
+        }
+
+        res.set_content(r.dump(), "application/json");
+    });
 }

examples/server/routes_sdcpp.cpp

@@ -205,6 +205,7 @@ static json make_img_gen_features_json() {
         {"cache", true},
         {"cancel_queued", true},
         {"cancel_generating", false},
+        {"preview", true},
     };
 }
 
@@ -219,9 +220,18 @@ static json make_vid_gen_features_json() {
         {"cache", true},
         {"cancel_queued", true},
         {"cancel_generating", false},
+        {"preview", true},
     };
 }
 
+static json make_preview_modes_json() {
+    json modes = json::array();
+    for (int i = 0; i < PREVIEW_COUNT; ++i) {
+        modes.push_back(sd_preview_name((preview_t)i));
+    }
+    return modes;
+}
+
 static json make_capabilities_json(ServerRuntime& runtime) {
     refresh_lora_cache(runtime);
     refresh_upscaler_cache(runtime);
@@ -396,6 +406,7 @@ static json make_capabilities_json(ServerRuntime& runtime) {
     result["loras"]                  = available_loras;
     result["upscalers"]              = available_upscalers;
     result["upscale"]                = have_upscaler_models;
+    result["preview_modes"]          = make_preview_modes_json();
     return result;
 }
 
@@ -418,6 +429,8 @@ static bool parse_img_gen_request(const json& body,
     if (!assign_output_options(request, output_format, output_compression, true, error_message)) {
         return false;
     }
+    request.preview_mode     = body.value("preview", std::string("none"));
+    request.preview_interval = body.value("preview_interval", 1);
     // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
     if (!request.gen_params.resolve_and_validate(IMG_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) {
         error_message = "invalid generation parameters";
@@ -445,6 +458,8 @@ static bool parse_vid_gen_request(const json& body,
     if (!assign_output_options(request, output_format, output_compression, error_message)) {
         return false;
     }
+    request.preview_mode     = body.value("preview", std::string("none"));
+    request.preview_interval = body.value("preview_interval", 1);
     // Intentionally disable prompt-embedded LoRA tag parsing for server APIs.
     if (!request.gen_params.resolve_and_validate(VID_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) {
         error_message = "invalid generation parameters";

examples/server/runtime.cpp

@@ -254,6 +254,12 @@ std::string SDSvrParams::to_string() const {
 void refresh_lora_cache(ServerRuntime& rt) {
     std::vector<LoraEntry> new_cache;
 
+    if (rt.ctx_params->lora_model_dir.empty()) {
+        std::lock_guard<std::mutex> lock(*rt.lora_mutex);
+        rt.lora_cache->clear();
+        return;
+    }
+
     fs::path lora_dir = rt.ctx_params->lora_model_dir;
     std::error_code ec;
     if (fs::exists(lora_dir, ec) && !ec && fs::is_directory(lora_dir, ec) && !ec) {

examples/server/runtime.h

@@ -66,6 +66,8 @@ struct ImgGenJobRequest {
     SDGenerationParams gen_params;
     std::string output_format = "png";
     int output_compression    = 100;
+    std::string preview_mode  = "none";
+    int preview_interval      = 1;
 
     sd_img_gen_params_t to_sd_img_gen_params_t() {
         return gen_params.to_sd_img_gen_params_t();
@@ -76,6 +78,8 @@ struct VidGenJobRequest {
     SDGenerationParams gen_params;
     std::string output_format = "webm";
     int output_compression    = 100;
+    std::string preview_mode  = "none";
+    int preview_interval      = 1;
 
     sd_vid_gen_params_t to_sd_vid_gen_params_t() {
         return gen_params.to_sd_vid_gen_params_t();

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
@@ -461,6 +462,11 @@ struct ggml_tensor;
 typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data);
 typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data);
 typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data);
+
+typedef struct {
+    int sample_pass;
+    int total_steps;
+} sd_preview_info_t;
 typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void* user_data);
 
 SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data);
@@ -469,6 +475,10 @@ SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data);
 // negative interval previews only completed logical step -interval. Zero previews the final
 // completed step of the first sampling pass (base-resolution or high-noise).
 SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data);
+// Query from a preview callback. Passes are numbered from 1 since the last
+// sd_set_preview_callback call; total_steps is the actual count in that pass.
+// Both fields are zero before sampling starts.
+SD_API sd_preview_info_t sd_get_preview_info();
 SD_API void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data);
 SD_API int32_t sd_get_num_physical_cores();
 SD_API const char* sd_get_system_info();
@@ -601,7 +611,9 @@ SD_API bool convert_with_components(const char* model_path,
                                     enum sd_type_t output_type,
                                     const char* tensor_type_rules,
                                     bool convert_name,
-                                    int n_threads);
+                                    int n_threads,
+                                    const sd_lora_t* loras,
+                                    int lora_count);
 
 SD_API bool preprocess_canny(sd_image_t image,
                              float high_threshold,

src/conditioning/conditioner.hpp

@@ -399,7 +399,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                 ss << "['" << item.first << "', " << item.second << "], ";
             }
             ss << "]";
-            LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
+            LOG_VERBOSE("parse '%s' to %s", escape_newlines(text).c_str(), escape_newlines(ss.str()).c_str());
         }
 
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
@@ -817,7 +817,7 @@ struct SD3CLIPEmbedder : public Conditioner {
                 ss << "['" << item.first << "', " << item.second << "], ";
             }
             ss << "]";
-            LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
+            LOG_VERBOSE("parse '%s' to %s", escape_newlines(text).c_str(), escape_newlines(ss.str()).c_str());
         }
 
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
@@ -909,6 +909,31 @@ struct SD3CLIPEmbedder : public Conditioner {
 
         size_t chunk_count = std::max(std::max(clip_l_tokens.size(), clip_g_tokens.size()), t5_tokens.size()) / chunk_len;
 
+        // Each active encoder must cover every chunk of the longest token sequence.
+        Tokenizer* tokenizers[] = {
+            clip_l ? clip_l_tokenizer.get() : nullptr,
+            clip_g ? clip_g_tokenizer.get() : nullptr,
+            t5 ? &t5_tokenizer : nullptr,
+        };
+        const size_t target_length = chunk_count * chunk_len;
+
+        for (size_t encoder_idx = 0; encoder_idx < 3; encoder_idx++) {
+            auto& tokens  = token_and_weights[encoder_idx].first;
+            auto& weights = token_and_weights[encoder_idx].second;
+            if (!tokenizers[encoder_idx] || tokens.size() >= target_length) {
+                continue;
+            }
+
+            std::vector<int> empty_tokens;
+            std::vector<float> empty_weights;
+            tokenizers[encoder_idx]->pad_tokens(empty_tokens, &empty_weights, nullptr, chunk_len, chunk_len, false);
+
+            while (tokens.size() < target_length) {
+                tokens.insert(tokens.end(), empty_tokens.begin(), empty_tokens.end());
+                weights.insert(weights.end(), empty_weights.begin(), empty_weights.end());
+            }
+        }
+
         for (int chunk_idx = 0; chunk_idx < chunk_count; chunk_idx++) {
             // clip_l
             sd::Tensor<float> chunk_hidden_states_l;
@@ -1207,7 +1232,7 @@ struct FluxCLIPEmbedder : public Conditioner {
                 ss << "['" << item.first << "', " << item.second << "], ";
             }
             ss << "]";
-            LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
+            LOG_VERBOSE("parse '%s' to %s", escape_newlines(text).c_str(), escape_newlines(ss.str()).c_str());
         }
 
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
@@ -1484,7 +1509,7 @@ struct T5CLIPEmbedder : public Conditioner {
                 ss << "['" << item.first << "', " << item.second << "], ";
             }
             ss << "]";
-            LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
+            LOG_VERBOSE("parse '%s' to %s", escape_newlines(text).c_str(), escape_newlines(ss.str()).c_str());
         }
 
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
@@ -1873,7 +1898,7 @@ struct AnimaConditioner : public Conditioner {
                 ss << "['" << item.first << "', " << item.second << "], ";
             }
             ss << "]";
-            LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
+            LOG_VERBOSE("parse '%s' to %s", escape_newlines(text).c_str(), escape_newlines(ss.str()).c_str());
         }
 
         std::vector<int> qwen_tokens;
@@ -1987,9 +2012,10 @@ 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) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
+        } 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;
         }
         llm        = std::make_shared<LLM::LLMRunner>(arch,
@@ -2113,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) {
@@ -2144,7 +2171,7 @@ struct LLMEmbedder : public Conditioner {
                 ss << "['" << item.first << "', " << item.second << "], ";
             }
             ss << "]";
-            LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
+            LOG_VERBOSE("parse '%s' to %s", escape_newlines(text).c_str(), escape_newlines(ss.str()).c_str());
         }
 
         std::vector<int> tokens;
@@ -2160,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);
 
@@ -2182,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);
@@ -2246,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;
     }
 
@@ -2312,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;
@@ -2953,7 +3004,7 @@ struct LLMEmbedder : public Conditioner {
             prompt_attn_range.second = static_cast<int>(prompt.size());
 
             prompt += "<|end|><|start|>assistant<|channel|>analysis<|message|>Need to generate one image according to the description.<|end|><|start|>assistant<|channel|>final<|message|>";
-        } else if (sd_version_is_z_image(version)) {
+        } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version)) {
             prompt_template_encode_start_idx = 0;
             out_layers                       = {35};  // -2
 
@@ -3061,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);
         }
@@ -3076,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 {};
         }
@@ -3141,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/convert.cpp

@@ -4,6 +4,7 @@
 #include <cstring>
 #include <exception>
 #include <fstream>
+#include <map>
 #include <memory>
 #include <mutex>
 #include <regex>
@@ -11,11 +12,15 @@
 #include <thread>
 #include <vector>
 
+#include "core/ggml_extend_backend.h"
 #include "core/util.h"
 #include "model_io/gguf_io.h"
 #include "model_io/safetensors_io.h"
 #include "model_io/streaming_writer.h"
 #include "model_loader.h"
+#include "model_manager.h"
+#include "name_conversion.h"
+#include "runtime/imatrix.h"
 
 struct TensorExportInfo {
     TensorStorage storage;
@@ -75,9 +80,9 @@ static size_t export_tensor_nbytes(const TensorExportInfo& info) {
     return static_cast<size_t>(output_storage.nbytes());
 }
 
-static TensorWritePlan tensor_write_plan_from_export_info(const TensorExportInfo& info) {
+static TensorWritePlan tensor_write_plan_from_export_info(const TensorExportInfo& info, const std::string& name) {
     TensorWritePlan plan;
-    plan.name   = info.storage.name;
+    plan.name   = name;
     plan.type   = info.type;
     plan.n_dims = info.storage.n_dims;
     for (int i = 0; i < SD_MAX_DIMS; i++) {
@@ -86,11 +91,14 @@ static TensorWritePlan tensor_write_plan_from_export_info(const TensorExportInfo
     return plan;
 }
 
-static std::vector<TensorWritePlan> tensor_write_plans_from_export_infos(const std::vector<TensorExportInfo>& tensors) {
+static std::vector<TensorWritePlan> tensor_write_plans_from_export_infos(const std::vector<TensorExportInfo>& tensors,
+                                                                         bool convert_name,
+                                                                         SDVersion version) {
     std::vector<TensorWritePlan> plans;
     plans.reserve(tensors.size());
     for (const TensorExportInfo& info : tensors) {
-        plans.push_back(tensor_write_plan_from_export_info(info));
+        const std::string name = convert_name ? convert_tensor_name(info.storage.name, version) : info.storage.name;
+        plans.push_back(tensor_write_plan_from_export_info(info, name));
     }
     return plans;
 }
@@ -158,16 +166,42 @@ static bool load_tensor_for_export(ModelLoader& model_loader, TensorExportJob& j
     return true;
 }
 
+static bool export_tensor_from_memory(const TensorExportInfo& info,
+                                      const std::map<std::string, ggml_tensor*>& mem,
+                                      SDVersion version,
+                                      std::vector<uint8_t>& out) {
+    auto found = mem.find(info.storage.name);
+    if (found == mem.end() || found->second == nullptr || found->second->data == nullptr) {
+        return false;
+    }
+    const TensorStorage& storage = info.storage;
+    out.resize(export_tensor_nbytes(info));
+    if (out.size() > 0) {
+        std::vector<float> imatrix = get_imatrix_collector().get_values(convert_tensor_name(info.storage.name, version));
+        convert_tensor(found->second->data,
+                       found->second->type,
+                       out.data(),
+                       info.type,
+                       (int)(storage.nelements() / storage.ne[0]),
+                       (int)storage.ne[0],
+                       std::move(imatrix));
+    }
+    return true;
+}
+
 static bool stream_tensor_data(ModelLoader& model_loader,
                                const std::string& output_path,
                                const std::vector<TensorExportInfo>& tensors,
                                const StreamingModelWriter& writer,
                                int n_threads,
+                               const std::map<std::string, ggml_tensor*>* mem,
                                std::string* error) {
     n_threads = n_threads > 0 ? n_threads : sd_get_num_physical_cores();
     n_threads = std::max(1, n_threads);
     LOG_INFO("streaming convert with %d threads", n_threads);
 
+    const SDVersion version = model_loader.get_sd_version();
+
     int64_t start_time       = ggml_time_ms();
     uint64_t bytes_written   = 0;
     size_t tensors_written   = 0;
@@ -241,7 +275,10 @@ static bool stream_tensor_data(ModelLoader& model_loader,
                 TensorExportJob job;
                 job.info = tensors[tensor_index];
                 try {
-                    job.success = load_tensor_for_export(model_loader, job);
+                    job.success = (mem != nullptr) && export_tensor_from_memory(job.info, *mem, version, job.data);
+                    if (!job.success) {
+                        job.success = load_tensor_for_export(model_loader, job);
+                    }
                 } catch (const std::exception& e) {
                     job.error   = e.what();
                     job.success = false;
@@ -298,8 +335,10 @@ static bool write_model_file_streaming(ModelLoader& model_loader,
                                        const std::vector<TensorExportInfo>& tensors,
                                        StreamingModelWriter& writer,
                                        int n_threads,
+                                       bool convert_name,
+                                       const std::map<std::string, ggml_tensor*>* mem,
                                        std::string* error) {
-    std::vector<TensorWritePlan> plans = tensor_write_plans_from_export_infos(tensors);
+    std::vector<TensorWritePlan> plans = tensor_write_plans_from_export_infos(tensors, convert_name, model_loader.get_sd_version());
     if (!writer.write_metadata(output_path, plans, error)) {
         return false;
     }
@@ -307,7 +346,7 @@ static bool write_model_file_streaming(ModelLoader& model_loader,
         return false;
     }
     model_loader.process_model_files(false, false);
-    return stream_tensor_data(model_loader, output_path, tensors, writer, n_threads, error);
+    return stream_tensor_data(model_loader, output_path, tensors, writer, n_threads, mem, error);
 }
 
 static bool init_convert_path(ModelLoader& model_loader, const char* path, const char* prefix, bool& loaded_any) {
@@ -326,7 +365,9 @@ static bool export_loaded_model(ModelLoader& model_loader,
                                 const char* output_path,
                                 sd_type_t output_type,
                                 const char* tensor_type_rules,
-                                int n_threads) {
+                                int n_threads,
+                                bool convert_name,
+                                const std::map<std::string, ggml_tensor*>* mem = nullptr) {
     ggml_type type             = sd_type_to_ggml_type(output_type);
     bool output_is_safetensors = ends_with(output_path, ".safetensors");
     TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
@@ -341,7 +382,7 @@ static bool export_loaded_model(ModelLoader& model_loader,
         } else {
             writer = std::make_unique<GGUFStreamingWriter>();
         }
-        success = write_model_file_streaming(model_loader, output_path, tensors, *writer, n_threads, &error);
+        success = write_model_file_streaming(model_loader, output_path, tensors, *writer, n_threads, convert_name, mem, &error);
     }
 
     if (!success && !error.empty()) {
@@ -351,6 +392,174 @@ static bool export_loaded_model(ModelLoader& model_loader,
     return success;
 }
 
+static bool has_active_loras(const sd_lora_t* loras, int lora_count) {
+    if (loras == nullptr) {
+        return false;
+    }
+    for (int i = 0; i < lora_count; i++) {
+        if (loras[i].multiplier != 0.0f) {
+            return true;
+        }
+    }
+    return false;
+}
+
+static std::vector<ModelManager::LoraSpec> build_lora_specs(const sd_lora_t* loras, int lora_count) {
+    std::vector<ModelManager::LoraSpec> specs;
+    specs.reserve(lora_count);
+    for (int i = 0; i < lora_count; i++) {
+        ModelManager::LoraSpec spec;
+        spec.path          = loras[i].path != nullptr ? loras[i].path : "";
+        spec.multiplier    = loras[i].multiplier;
+        spec.is_high_noise = loras[i].is_high_noise;
+        spec.required      = true;
+        specs.push_back(std::move(spec));
+    }
+    return specs;
+}
+
+static ModelComponent component_for_tensor_name(const std::string& name) {
+    if (starts_with(name, "text_encoders.") || starts_with(name, "cond_stage_model.")) {
+        return ModelComponent::Conditioner;
+    }
+    if (starts_with(name, "model.high_noise_diffusion_model.")) {
+        return ModelComponent::HighNoiseDiffusion;
+    }
+    if (starts_with(name, "model.diffusion_model.")) {
+        return ModelComponent::Diffusion;
+    }
+    if (starts_with(name, "vae.") || starts_with(name, "first_stage_model.")) {
+        return ModelComponent::VAE;
+    }
+    return ModelComponent::Diffusion;
+}
+
+static bool load_model_into_memory(ModelLoader& model_loader,
+                                   const sd_lora_t* loras,
+                                   int lora_count,
+                                   int n_threads,
+                                   ModelManager& manager,
+                                   ggml_backend_t cpu,
+                                   ggml_context* ctx,
+                                   std::map<std::string, ggml_tensor*>& mem,
+                                   std::vector<ggml_tensor*>& all_tensors,
+                                   ggml_type output_type,
+                                   const TensorTypeRules& type_rules) {
+    const String2TensorStorage& storage_map = model_loader.get_tensor_storage_map();
+    if (storage_map.empty()) {
+        LOG_ERROR("no tensors to load into memory for convert");
+        return false;
+    }
+
+    const SDVersion version  = model_loader.get_sd_version();
+    ModelLoader merge_loader = model_loader;
+    merge_loader.convert_tensors_name();
+
+    std::map<ModelComponent, std::map<std::string, ggml_tensor*>> groups;
+    size_t total_bytes = 0;
+    for (const auto& [name, storage] : storage_map) {
+        ggml_type load_type         = storage.type;
+        const ggml_type export_type = get_export_tensor_type(model_loader, storage, output_type, type_rules);
+        const bool quant_src        = ggml_is_quantized(storage.type);
+        const bool quant_dst        = ggml_is_quantized(export_type);
+        if (quant_src && quant_dst) {
+            load_type = GGML_TYPE_F16;
+        } else if (quant_src || quant_dst) {
+            load_type = quant_src ? export_type : storage.type;
+        } else if (storage.type == GGML_TYPE_F32 || export_type == GGML_TYPE_F32) {
+            load_type = GGML_TYPE_F32;
+        }
+        ggml_tensor* tensor = ggml_new_tensor(ctx, load_type, storage.n_dims, storage.ne);
+        if (tensor == nullptr) {
+            LOG_ERROR("failed to create tensor '%s' for in-memory convert", name.c_str());
+            return false;
+        }
+        const std::string merge_name = convert_tensor_name(name, version);
+        auto& group                  = groups[component_for_tensor_name(merge_name)];
+        if (!group.emplace(merge_name, tensor).second) {
+            LOG_ERROR("duplicate tensor name '%s' after name conversion", merge_name.c_str());
+            return false;
+        }
+        mem[name] = tensor;
+        all_tensors.push_back(tensor);
+        total_bytes += ggml_nbytes(tensor);
+    }
+    LOG_INFO("loading %zu tensors (%.2f GB) into memory for convert",
+             all_tensors.size(), total_bytes / (1024.0 * 1024.0 * 1024.0));
+
+    manager.set_n_threads(n_threads);
+    manager.set_enable_mmap(false);
+    if (!manager.set_loader(std::move(merge_loader))) {
+        LOG_ERROR("failed to set model loader for in-memory convert");
+        return false;
+    }
+
+    for (const auto& [component, tensors] : groups) {
+        if (!manager.register_param_tensors(component, tensors, ModelManager::ResidencyMode::ParamBackend, cpu, cpu)) {
+            LOG_ERROR("failed to register %s tensors for in-memory convert", model_component_name(component));
+            return false;
+        }
+    }
+
+    std::vector<ModelManager::LoraSpec> lora_specs = build_lora_specs(loras, lora_count);
+    if (!manager.prepare_lora_sources(lora_specs)) {
+        LOG_ERROR("failed to prepare LoRA sources for convert");
+        return false;
+    }
+    if (!manager.set_loras(lora_specs, version)) {
+        LOG_ERROR("failed to set LoRAs for convert");
+        return false;
+    }
+    if (!manager.prepare_params(all_tensors)) {
+        LOG_ERROR("failed to load model into memory for convert");
+        return false;
+    }
+
+    return true;
+}
+
+static bool convert_model_in_memory(ModelLoader& model_loader,
+                                    const char* output_path,
+                                    sd_type_t output_type,
+                                    const char* tensor_type_rules,
+                                    bool convert_name,
+                                    int n_threads,
+                                    const sd_lora_t* loras,
+                                    int lora_count) {
+    auto cpu = std::unique_ptr<ggml_backend, decltype(&ggml_backend_free)>(sd_backend_cpu_init(), ggml_backend_free);
+    if (cpu == nullptr) {
+        LOG_ERROR("failed to init CPU backend for in-memory convert");
+        return false;
+    }
+
+    ggml_init_params ctx_params;
+    ctx_params.mem_size   = model_loader.get_tensor_storage_map().size() * ggml_tensor_overhead();
+    ctx_params.mem_buffer = nullptr;
+    ctx_params.no_alloc   = true;
+    ggml_context* ctx     = ggml_init(ctx_params);
+    if (ctx == nullptr) {
+        LOG_ERROR("ggml_init failed for in-memory convert");
+        return false;
+    }
+
+    ggml_type type             = sd_type_to_ggml_type(output_type);
+    TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
+
+    std::vector<ggml_tensor*> all_tensors;
+    std::map<std::string, ggml_tensor*> mem;
+    bool success = false;
+    {
+        // The manager owns the buffers backing `mem` and must outlive the export; it is
+        // destroyed (freeing those buffers) before the tensor context is freed below.
+        ModelManager manager;
+        if (load_model_into_memory(model_loader, loras, lora_count, n_threads, manager, cpu.get(), ctx, mem, all_tensors, type, type_rules)) {
+            success = export_loaded_model(model_loader, output_path, output_type, tensor_type_rules, n_threads, convert_name, &mem);
+        }
+    }
+    ggml_free(ctx);
+    return success;
+}
+
 bool convert_with_components(const char* model_path,
                              const char* clip_l_path,
                              const char* clip_g_path,
@@ -361,10 +570,22 @@ bool convert_with_components(const char* model_path,
                              sd_type_t output_type,
                              const char* tensor_type_rules,
                              bool convert_name,
-                             int n_threads) {
+                             int n_threads,
+                             const sd_lora_t* loras,
+                             int lora_count) {
     if (!validate_tensor_types(output_type, tensor_type_rules)) {
         return false;
     }
+
+    if (loras != nullptr) {
+        for (int i = 0; i < lora_count; i++) {
+            LOG_INFO("lora %d: '%s' (multiplier %.2f%s)",
+                     i + 1, loras[i].path != nullptr ? loras[i].path : "",
+                     loras[i].multiplier,
+                     loras[i].is_high_noise ? ", high noise" : "");
+        }
+    }
+
     ModelLoader model_loader;
     bool loaded_any = false;
 
@@ -382,11 +603,17 @@ bool convert_with_components(const char* model_path,
         return false;
     }
 
+    if (has_active_loras(loras, lora_count)) {
+        LOG_INFO("loading the full model into memory to apply LoRAs");
+        return convert_model_in_memory(model_loader, output_path, output_type, tensor_type_rules, convert_name, n_threads,
+                                       loras, lora_count);
+    }
+
     if (convert_name) {
         model_loader.convert_tensors_name();
     }
 
-    return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules, n_threads);
+    return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules, n_threads, false, nullptr);
 }
 
 bool convert(const char* input_path,
@@ -405,5 +632,7 @@ bool convert(const char* input_path,
                                    output_type,
                                    tensor_type_rules,
                                    convert_name,
+                                   0,
+                                   nullptr,
                                    0);
 }

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/core/ggml_extend.cpp

@@ -223,13 +223,13 @@ ggml_tensor* ggml_ext_linear(ggml_context* ctx,
         x           = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]);
         x           = ggml_mul_mat(ctx, w, x);
         if (force_prec_f32) {
-            ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+            ggml_prec_set_acc(x, GGML_PREC_F32);
         }
         x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3);
     } else {
         x = ggml_mul_mat(ctx, w, x);
         if (force_prec_f32) {
-            ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+            ggml_prec_set_acc(x, GGML_PREC_F32);
         }
     }
     if (scale != 1.f) {
@@ -488,7 +488,7 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
         x          = ggml_mul_mat(ctx,
                                   ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]),
                                   ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1] * w->ne[2] * IC, OC));
-        ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+        ggml_prec_set_acc(x, GGML_PREC_F32);
 
         int64_t OD = im2col->ne[3] / N;
         x          = ggml_reshape_4d(ctx, x, im2col->ne[1] * im2col->ne[2], OD, N, OC);
@@ -683,8 +683,8 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
                               v_in->type == GGML_TYPE_F32 && sd_backend_supports_cuda_mma(backend);
         if (pad_head) {
             // CUDA FA MMA starts at 64 channels; keep the original head's attention scale.
-            q_in = ggml_pad(ctx, q_in, 64 - d_head, 0, 0, 0);
-            k_in = ggml_pad(ctx, k_in, 64 - d_head, 0, 0, 0);
+            q_in = ggml_pad(ctx, q_in, static_cast<int>(64 - d_head), 0, 0, 0);
+            k_in = ggml_pad(ctx, k_in, static_cast<int>(64 - d_head), 0, 0, 0);
         }
         if (kv_scale != 1.0f) {
             k_in = ggml_ext_scale(ctx, k_in, kv_scale);
@@ -694,7 +694,7 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
         v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3));
         v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N);
         if (pad_head) {
-            v_in = ggml_pad(ctx, v_in, 64 - d_head, 0, 0, 0);
+            v_in = ggml_pad(ctx, v_in, static_cast<int>(64 - d_head), 0, 0, 0);
         }
         if (kv_scale != 1.0f) {
             v_in = ggml_ext_scale(ctx, v_in, kv_scale);
@@ -721,7 +721,7 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
         if (!ggml_backend_supports_op(backend, out)) {
             return nullptr;
         }
-        ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32);
+        ggml_prec_set_acc(out, GGML_PREC_F32);
         if (kv_scale != 1.0f) {
             out = ggml_ext_scale(ctx, out, 1.0f / kv_scale);
         }
@@ -742,9 +742,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
         }
         if (padded_head != d_head) {
             // Keep the original head's softmax scale when padding for the CUDA kernel.
-            q_in = ggml_pad(ctx, q_in, padded_head - d_head, 0, 0, 0);
-            k_in = ggml_pad(ctx, k_in, padded_head - d_head, 0, 0, 0);
-            v_in = ggml_pad(ctx, v_in, padded_head - d_head, 0, 0, 0);
+            q_in = ggml_pad(ctx, q_in, static_cast<int>(padded_head - d_head), 0, 0, 0);
+            k_in = ggml_pad(ctx, k_in, static_cast<int>(padded_head - d_head), 0, 0, 0);
+            v_in = ggml_pad(ctx, v_in, static_cast<int>(padded_head - d_head), 0, 0, 0);
         }
         if (kv_scale != 1.0f) {
             k_in = ggml_ext_scale(ctx, k_in, kv_scale);
@@ -797,7 +797,7 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
         v = ggml_reshape_3d(ctx, v, L_k, d_head, n_kv_head * N);   // [N * n_kv_head, d_head, L_k]
 
         auto kq = ggml_mul_mat(ctx, k, q);  // [N * n_head, L_q, L_k]
-        ggml_mul_mat_set_prec(kq, GGML_PREC_F32);
+        ggml_prec_set_acc(kq, GGML_PREC_F32);
         kq = ggml_scale_inplace(ctx, kq, scale);
         if (mask) {
             kq = ggml_add_inplace(ctx, kq, mask);

src/core/util.cpp

@@ -62,6 +62,21 @@ void replace_all_chars(std::string& str, char target, char replacement) {
     }
 }
 
+std::string escape_newlines(const std::string& text) {
+    std::string escaped;
+    escaped.reserve(text.size());
+    for (char ch : text) {
+        if (ch == '\n') {
+            escaped += "\\n";
+        } else if (ch == '\r') {
+            escaped += "\\r";
+        } else {
+            escaped += ch;
+        }
+    }
+    return escaped;
+}
+
 static std::string sd_vformat(const char* fmt, va_list ap) {
     char small[128];
     va_list ap2;
@@ -360,12 +375,13 @@ int32_t sd_get_num_physical_cores() {
 static sd_progress_cb_t sd_progress_cb = nullptr;
 void* sd_progress_cb_data              = nullptr;
 
-static sd_preview_cb_t sd_preview_cb = nullptr;
-static void* sd_preview_cb_data      = nullptr;
-preview_t sd_preview_mode            = PREVIEW_NONE;
-int sd_preview_interval              = 1;
-bool sd_preview_denoised             = true;
-bool sd_preview_noisy                = false;
+static sd_preview_info_t sd_preview_info = {};
+static sd_preview_cb_t sd_preview_cb     = nullptr;
+static void* sd_preview_cb_data          = nullptr;
+preview_t sd_preview_mode                = PREVIEW_NONE;
+int sd_preview_interval                  = 1;
+bool sd_preview_denoised                 = true;
+bool sd_preview_noisy                    = false;
 
 static sd_graph_eval_callback_t sd_backend_eval_cb = nullptr;
 static void* sd_backend_eval_cb_data               = nullptr;
@@ -644,9 +660,14 @@ void* sd_log_cb_data         = nullptr;
 static void sd_log_dispatch(sd_log_level_t level, const std::string& origin, const std::string& text) {
     if (sd_log_cb == nullptr)
         return;
-    std::string message = origin + " - " + text;
-    if (message.back() != '\n') {
-        message += '\n';
+    std::string message = text;
+    while (!message.empty() && (message.back() == '\n' || message.back() == '\r')) {
+        message.pop_back();
+    }
+    if (origin == "ggml") {
+        message = origin + ": " + message + '\n';
+    } else {
+        message += " --- " + origin + '\n';
     }
     sd_log_cb(level, message.c_str(), sd_log_cb_data);
 }
@@ -656,7 +677,7 @@ void log_printf(sd_log_level_t level, const char* file, int line, const char* fo
     va_start(args, format);
     std::string message = sd_vformat(format, args);
     va_end(args);
-    std::string origin = sd_format("%s:%-4d", sd_basename(file).c_str(), line);
+    std::string origin = sd_format("%s:%d", sd_basename(file).c_str(), line);
     sd_log_dispatch(level, origin, message);
 }
 
@@ -691,6 +712,7 @@ void sd_set_progress_callback(sd_progress_cb_t cb, void* data) {
     sd_progress_cb_data = data;
 }
 void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode, int interval, bool denoised, bool noisy, void* data) {
+    sd_preview_info     = {};
     sd_preview_cb       = cb;
     sd_preview_cb_data  = data;
     sd_preview_mode     = mode;
@@ -704,6 +726,17 @@ void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data) {
     sd_backend_eval_cb_data = data;
 }
 
+void sd_begin_preview_pass(int total_steps) {
+    if (sd_preview_cb != nullptr) {
+        ++sd_preview_info.sample_pass;
+        sd_preview_info.total_steps = total_steps;
+    }
+}
+
+sd_preview_info_t sd_get_preview_info() {
+    return sd_preview_info;
+}
+
 sd_preview_cb_t sd_get_preview_callback() {
     return sd_preview_cb;
 }

src/core/util.h

@@ -29,6 +29,7 @@ bool contains(const std::string& str, const std::string& substr);
 std::string sd_format(const char* fmt, ...);
 
 void replace_all_chars(std::string& str, char target, char replacement);
+std::string escape_newlines(const std::string& text);
 
 int round_up_to(int value, int base);
 
@@ -101,6 +102,7 @@ std::vector<std::pair<std::string, float>> split_quotation_attention(
 sd_progress_cb_t sd_get_progress_callback();
 void* sd_get_progress_callback_data();
 
+void sd_begin_preview_pass(int total_steps);
 sd_preview_cb_t sd_get_preview_callback();
 void* sd_get_preview_callback_data();
 preview_t sd_get_preview_mode();

src/model.h

@@ -64,6 +64,8 @@ enum SDVersion {
     VERSION_ESRGAN,
     VERSION_PIXART,
     VERSION_MING_IMAGE,
+    VERSION_Z_IMAGE_L2P,
+    VERSION_IRIS,
     VERSION_COUNT,
 };
 
@@ -176,6 +178,10 @@ static inline bool sd_version_is_z_image(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_z_image_l2p(SDVersion version) {
+    return version == VERSION_Z_IMAGE_L2P;
+}
+
 static inline bool sd_version_is_llada_image(SDVersion version) {
     if (version == VERSION_LLADA_IMAGE) {
         return true;
@@ -315,6 +321,7 @@ static inline bool sd_version_is_dit(SDVersion version) {
         version == VERSION_HIDREAM_O1 ||
         sd_version_is_anima(version) ||
         sd_version_is_z_image(version) ||
+        sd_version_is_z_image_l2p(version) ||
         version == VERSION_MING_IMAGE ||
         sd_version_is_llada_image(version) ||
         sd_version_is_boogu_image(version) ||
@@ -328,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/common/ggml_block.hpp

@@ -205,6 +205,11 @@ public:
         force_prec_f32 = force_prec_f32_;
     }
 
+    bool has_quantized_weight() const {
+        auto it = params.find("weight");
+        return it != params.end() && it->second != nullptr && ggml_is_quantized(it->second->type);
+    }
+
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         ggml_tensor* w            = params["weight"];
         const float scale         = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale;

src/model/common/rope.hpp

@@ -1067,10 +1067,22 @@ namespace Rope {
         return result;
     }
 
-    __STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx,
+    __STATIC_INLINE__ ggml_tensor* apply_rope(GGMLRunnerContext* runner_ctx,
                                               ggml_tensor* x,
                                               ggml_tensor* pe,
                                               bool rope_interleaved = true) {
+        auto ctx = runner_ctx->ggml_ctx;
+#ifndef SD_USE_UPSTREAM_GGML
+        auto backend = runner_ctx->backend;
+        if (backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 &&
+            x->ne[0] % 2 == 0 && pe->ne[0] == 2 && pe->ne[1] == 2 &&
+            pe->ne[2] == x->ne[0] / 2 && pe->ne[3] == x->ne[2]) {
+            auto out = ggml_rope_apply(ctx, x, pe, rope_interleaved);
+            if (ggml_backend_supports_op(backend, out)) {
+                return out;
+            }
+        }
+#endif
         // x: [N, L, n_head, d_head]
         // pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]]
         int64_t d_head = x->ne[0];
@@ -1126,8 +1138,8 @@ namespace Rope {
         // return: [N, L, n_head*d_head]
         int64_t n_head = q->ne[1];
 
-        q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved);  // [N*n_head, L, d_head]
-        k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved);  // [N*n_head, L, d_head]
+        q = apply_rope(ctx, q, pe, rope_interleaved);  // [N*n_head, L, d_head]
+        k = apply_rope(ctx, k, pe, rope_interleaved);  // [N*n_head, L, d_head]
 
         auto x = ggml_ext_attention_ext(ctx, q, k, v, n_head, mask, true, ctx->flash_attn_enabled, kv_scale);  // [N, L, n_head*d_head]
         return x;

src/model/diffusion/anima.hpp

@@ -235,8 +235,8 @@ namespace Anima {
                 if (pe_k == nullptr) {
                     pe_k = pe_q;
                 }
-                auto q_rope = Rope::apply_rope(ctx->ggml_ctx, q4, pe_q, false);
-                auto k_rope = Rope::apply_rope(ctx->ggml_ctx, k4, pe_k, false);
+                auto q_rope = Rope::apply_rope(ctx, q4, pe_q, false);
+                auto k_rope = Rope::apply_rope(ctx, k4, pe_k, false);
                 attn_out    = ggml_ext_attention_ext(ctx,
                                                      q_rope,
                                                      k_rope,

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/ltxv.hpp

@@ -548,21 +548,22 @@ namespace LTXV {
         return build_rope_matrix_from_frequencies(freqs, dim);
     }
 
-    __STATIC_INLINE__ ggml_tensor* apply_hidden_rope(ggml_context* ctx,
+    __STATIC_INLINE__ ggml_tensor* apply_hidden_rope(GGMLRunnerContext* runner_ctx,
                                                      ggml_tensor* x,
                                                      ggml_tensor* pe,
                                                      int64_t heads,
                                                      int64_t dim_head,
                                                      bool rope_interleaved) {
+        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(x->ne[0] == heads * dim_head);
         auto x4 = ggml_reshape_4d(ctx, x, dim_head, heads, x->ne[1], x->ne[2]);
         if (pe != nullptr && pe->ne[3] == x->ne[1] * heads) {
             auto x_flat   = ggml_reshape_4d(ctx, x4, dim_head, 1, x->ne[1] * heads, x->ne[2]);
-            auto out_flat = Rope::apply_rope(ctx, x_flat, pe, rope_interleaved);
+            auto out_flat = Rope::apply_rope(runner_ctx, x_flat, pe, rope_interleaved);
             auto out4     = ggml_reshape_4d(ctx, out_flat, dim_head, heads, x->ne[1], x->ne[2]);
             return ggml_reshape_3d(ctx, out4, heads * dim_head, x->ne[1], x->ne[2]);
         }
-        return Rope::apply_rope(ctx, x4, pe, rope_interleaved);
+        return Rope::apply_rope(runner_ctx, x4, pe, rope_interleaved);
     }
 
     struct TimestepEmbedder : public GGMLBlock {
@@ -705,8 +706,8 @@ namespace LTXV {
                 if (k_pe == nullptr) {
                     k_pe = pe;
                 }
-                q = apply_hidden_rope(ctx->ggml_ctx, q, pe, heads, dim_head, rope_interleaved);
-                k = apply_hidden_rope(ctx->ggml_ctx, k, k_pe, heads, dim_head, rope_interleaved);
+                q = apply_hidden_rope(ctx, q, pe, heads, dim_head, rope_interleaved);
+                k = apply_hidden_rope(ctx, k, k_pe, heads, dim_head, rope_interleaved);
             }
 
             auto out = ggml_ext_attention_ext(ctx,

src/model/diffusion/minimax_h3.hpp

@@ -160,12 +160,13 @@ namespace MiniMaxH3 {
         return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]);
     }
 
-    static ggml_tensor* apply_partial_rope(ggml_context* ctx,
+    static ggml_tensor* apply_partial_rope(GGMLRunnerContext* runner_ctx,
                                            ggml_tensor* x,
                                            ggml_tensor* pe) {
+        auto ctx        = runner_ctx->ggml_ctx;
         int64_t rot_dim = pe->ne[2] * 2;
         GGML_ASSERT(rot_dim <= x->ne[0]);
-        auto rotated = Rope::apply_rope(ctx,
+        auto rotated = Rope::apply_rope(runner_ctx,
                                         ggml_ext_slice(ctx, x, 0, 0, rot_dim),
                                         pe,
                                         false);
@@ -209,8 +210,8 @@ namespace MiniMaxH3 {
             q                = q_norm->forward(ctx, q);
             k                = k_norm->forward(ctx, k);
             if (pe != nullptr) {
-                q = apply_partial_rope(ctx->ggml_ctx, q, pe);
-                k = apply_partial_rope(ctx->ggml_ctx, k, pe);
+                q = apply_partial_rope(ctx, q, pe);
+                k = apply_partial_rope(ctx, k, pe);
             } else {
                 q = attention_layout(ctx->ggml_ctx, q);
                 k = attention_layout(ctx->ggml_ctx, k);

src/model/diffusion/qwen_image_2_1.hpp

@@ -188,8 +188,8 @@ namespace Qwen {
             auto v = project("to_v");
             q      = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, q);
             k      = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k);
-            q      = Rope::apply_rope(ctx->ggml_ctx, q, pe);
-            k      = Rope::apply_rope(ctx->ggml_ctx, k, pe);
+            q      = Rope::apply_rope(ctx, q, pe);
+            k      = Rope::apply_rope(ctx, k, pe);
             if (cache.mode == QwenImage21PrefixCache::Mode::STORE) {
                 // Preserve query-first attention evaluation while writing each layer's
                 // prefix before its full-sequence K/V can accumulate across layers.

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/diffusion/z_image.hpp

@@ -168,13 +168,16 @@ namespace ZImage {
             int64_t N       = x->ne[2];
             auto out_proj   = std::dynamic_pointer_cast<Linear>(blocks[split_qkv ? "to_out.0" : "out"]);
 
+            if (out_proj->has_quantized_weight()) {
+                out_proj->set_scale(1.f / 16.f);
+            }
+
             if (split_qkv) {
                 auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
                 auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
                 auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
 
                 if (sd_backend_is(ctx->backend, "ROCm")) {
-                    out_proj->set_scale(1.f / 16.f);
                     out_proj->set_force_prec_f32(true);
                     q_proj->set_force_prec_f32(true);
                     k_proj->set_force_prec_f32(true);
@@ -197,7 +200,6 @@ namespace ZImage {
             auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
 
             if (sd_backend_is(ctx->backend, "ROCm")) {
-                out_proj->set_scale(1.f / 16.f);
                 out_proj->set_force_prec_f32(true);
                 qkv_proj->set_force_prec_f32(true);
             }

src/model/diffusion/z_image_l2p.hpp

@@ -0,0 +1,247 @@
+#ifndef __SD_MODEL_DIFFUSION_Z_IMAGE_L2P_HPP__
+#define __SD_MODEL_DIFFUSION_Z_IMAGE_L2P_HPP__
+
+#include <cmath>
+
+#include "z_image.hpp"
+
+// Ref: https://github.com/TencentYoutuResearch/T2I-L2P/blob/main/diffsynth/models/z_image_dit_L2P.py
+
+namespace ZImageL2P {
+    struct ZImageL2PConfig : ZImage::ZImageConfig {
+        static ZImageL2PConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) {
+            ZImageL2PConfig config;
+            static_cast<ZImage::ZImageConfig&>(config) = ZImage::ZImageConfig::detect_from_weights(tensors, prefix);
+            config.in_channels                         = 3;
+            config.out_channels                        = 3;
+            config.patch_size                          = 16;
+            auto x_embedder                            = tensors.find(prefix + ".x_embedder.weight");
+            if (x_embedder != tensors.end()) {
+                config.patch_size = static_cast<int>(std::lround(std::sqrt(static_cast<double>(x_embedder->second.ne[0] / config.in_channels))));
+            }
+            // L2P ships split to_q/to_k/to_v, which name conversion maps to qkv.weight plus
+            // .1/.2 parts; the base detection only sees the q part and undercounts kv heads.
+            auto k_part = tensors.find(prefix + ".layers.0.attention.qkv.weight.1");
+            if (k_part != tensors.end()) {
+                config.num_kv_heads = k_part->second.ne[1] / config.head_dim;
+            }
+            LOG_VERBOSE("z_image_l2p: patch_size = %d, in_channels = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64,
+                        config.patch_size,
+                        config.in_channels,
+                        config.num_heads,
+                        config.num_kv_heads);
+            return config;
+        }
+    };
+
+    // U-Net over the full-resolution noisy image. Its bottleneck sits at the DiT token
+    // grid, so the number of pooling levels is tied to patch_size == 16.
+    class LocalDecoder : public GGMLBlock {
+        static constexpr int LEVELS                  = 4;
+        static constexpr int64_t CHANNELS[LEVELS]    = {64, 128, 256, 512};
+        static constexpr int64_t BOTTLENECK_CHANNELS = 512;
+
+    public:
+        LocalDecoder(int64_t in_channels, int64_t cond_channels) {
+            int64_t prev = in_channels;
+            for (int i = 0; i < LEVELS; ++i) {
+                blocks["enc" + std::to_string(i + 1) + ".0"] = std::make_shared<Conv2d>(prev, CHANNELS[i], std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+                prev                                         = CHANNELS[i];
+            }
+            blocks["bottleneck.0"] = std::make_shared<Conv2d>(prev + cond_channels, BOTTLENECK_CHANNELS, std::pair{1, 1});
+            prev                   = BOTTLENECK_CHANNELS;
+            for (int i = LEVELS - 1; i >= 0; --i) {
+                const std::string level      = std::to_string(i + 1);
+                const int64_t out            = i == 0 ? CHANNELS[0] : CHANNELS[i - 1];
+                blocks["up" + level + ".1"]  = std::make_shared<Conv2d>(prev, prev, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+                blocks["dec" + level + ".0"] = std::make_shared<Conv2d>(prev + CHANNELS[i], out, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+                prev                         = out;
+            }
+            blocks["out_conv"] = std::make_shared<Conv2d>(prev, in_channels, std::pair{1, 1});
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cond) {
+            // x: [N, C, H, W]
+            // cond: [N, cond_channels, H / 16, W / 16]
+            // return: [N, C, H, W]
+            auto gctx = ctx->ggml_ctx;
+            auto conv = [&](const std::string& name, ggml_tensor* h) {
+                return std::dynamic_pointer_cast<Conv2d>(blocks[name])->forward(ctx, h);
+            };
+
+            std::vector<ggml_tensor*> skips;
+            auto h = x;
+            for (int i = 0; i < LEVELS; ++i) {
+                h = ggml_silu(gctx, conv("enc" + std::to_string(i + 1) + ".0", h));
+                skips.push_back(h);
+                h = ggml_pool_2d(gctx, h, GGML_OP_POOL_MAX, 2, 2, 2, 2, 0, 0);
+            }
+
+            h = ggml_concat(gctx, h, cond, 2);
+            h = ggml_silu(gctx, conv("bottleneck.0", h));
+
+            for (int i = LEVELS - 1; i >= 0; --i) {
+                const std::string level = std::to_string(i + 1);
+                h                       = ggml_upscale(gctx, h, 2, GGML_SCALE_MODE_NEAREST);
+                h                       = conv("up" + level + ".1", h);
+                h                       = ggml_concat(gctx, h, skips[i], 2);
+                h                       = ggml_silu(gctx, conv("dec" + level + ".0", h));
+            }
+            return conv("out_conv", h);
+        }
+    };
+
+    class ZImageL2PModel : public GGMLBlock {
+        ZImageL2PConfig config;
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
+            params["x_pad_token"]   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
+        }
+
+    public:
+        explicit ZImageL2PModel(const ZImageL2PConfig& config)
+            : config(config) {
+            blocks["x_embedder"]     = std::make_shared<Linear>(config.patch_size * config.patch_size * config.in_channels, config.hidden_size);
+            blocks["t_embedder"]     = std::make_shared<TimestepEmbedder>(MIN(config.hidden_size, 1024), 256, 256);
+            blocks["cap_embedder.0"] = std::make_shared<RMSNorm>(config.cap_feat_dim, config.norm_eps);
+            blocks["cap_embedder.1"] = std::make_shared<Linear>(config.cap_feat_dim, config.hidden_size);
+            auto add_blocks          = [&](const std::string& prefix, int64_t count, bool modulation) {
+                for (int64_t i = 0; i < count; ++i) {
+                    blocks[prefix + std::to_string(i)] = std::make_shared<ZImage::JointTransformerBlock>(
+                        static_cast<int>(i), config.hidden_size, config.head_dim, config.num_heads,
+                        config.num_kv_heads, config.multiple_of, config.ffn_dim_multiplier,
+                        config.norm_eps, config.qk_norm, modulation, true, false, 1e-5f);
+                }
+            };
+            add_blocks("noise_refiner.", config.num_refiner_layers, true);
+            add_blocks("context_refiner.", config.num_refiner_layers, false);
+            add_blocks("layers.", config.num_layers, true);
+            blocks["local_decoder"] = std::make_shared<LocalDecoder>(config.in_channels, config.hidden_size);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, ggml_tensor* pe) {
+            // x: [N, C, H, W], H and W are multiples of patch_size
+            // context: [N, L, cap_feat_dim]
+            // return: [N, C, H, W]
+            auto gctx           = ctx->ggml_ctx;
+            const int64_t W     = x->ne[0];
+            const int64_t H     = x->ne[1];
+            const int64_t N     = x->ne[3];
+            const int64_t w     = W / config.patch_size;
+            const int64_t h     = H / config.patch_size;
+            const int64_t n_img = w * h;
+            const int64_t n_txt = context->ne[1];
+
+            auto img = DiT::patchify(gctx, x, config.patch_size, config.patch_size, false);
+            img      = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"])->forward(ctx, img);
+            auto txt = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"])->forward(ctx, context);
+            txt      = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"])->forward(ctx, txt);
+            auto t   = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, timestep);
+            sd::ggml_graph_cut::mark_graph_cut(txt, "z_image_l2p.prelude", "txt");
+            sd::ggml_graph_cut::mark_graph_cut(img, "z_image_l2p.prelude", "img");
+            sd::ggml_graph_cut::mark_graph_cut(t, "z_image_l2p.prelude", "t_emb");
+
+            const int64_t n_txt_pad = Rope::bound_mod(static_cast<int>(n_txt), ZImage::SEQ_MULTI_OF);
+            if (n_txt_pad > 0) {
+                auto pad = params["cap_pad_token"];
+                txt      = ggml_concat(gctx, txt, ggml_repeat_4d(gctx, pad, pad->ne[0], n_txt_pad, N, 1), 1);
+            }
+            const int64_t n_img_pad = Rope::bound_mod(static_cast<int>(n_img), ZImage::SEQ_MULTI_OF);
+            if (n_img_pad > 0) {
+                auto pad = params["x_pad_token"];
+                img      = ggml_concat(gctx, img, ggml_repeat_4d(gctx, pad, pad->ne[0], n_img_pad, N, 1), 1);
+            }
+            GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]);
+
+            auto txt_pe = ggml_ext_slice(gctx, pe, 3, 0, txt->ne[1]);
+            auto img_pe = ggml_ext_slice(gctx, pe, 3, txt->ne[1], pe->ne[3]);
+            for (int64_t i = 0; i < config.num_refiner_layers; ++i) {
+                txt = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)])->forward(ctx, txt, txt_pe);
+                sd::ggml_graph_cut::mark_graph_cut(txt, "z_image_l2p.context_refiner." + std::to_string(i), "txt");
+            }
+            for (int64_t i = 0; i < config.num_refiner_layers; ++i) {
+                img = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)])->forward(ctx, img, img_pe, nullptr, t);
+                sd::ggml_graph_cut::mark_graph_cut(img, "z_image_l2p.noise_refiner." + std::to_string(i), "img");
+            }
+
+            auto txt_img = ggml_concat(gctx, txt, img, 1);
+            for (int64_t i = 0; i < config.num_layers; ++i) {
+                txt_img = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["layers." + std::to_string(i)])->forward(ctx, txt_img, pe, nullptr, t);
+                sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image_l2p.layers." + std::to_string(i), "txt_img");
+            }
+
+            // The local decoder consumes the raw last-layer hidden states: there is no final
+            // norm or adaLN modulation in front of it, unlike Z-Image's final_layer.
+            const int64_t img_start = n_txt + n_txt_pad;
+            auto feat               = ggml_ext_slice(gctx, txt_img, 1, img_start, img_start + n_img);                           // [N, h*w, hidden_size]
+            feat                    = ggml_reshape_4d(gctx, feat, config.hidden_size, w, h, N);                                 // [N, h, w, hidden_size]
+            feat                    = ggml_cont(gctx, ggml_permute(gctx, feat, 2, 0, 1, 3));                                    // [N, hidden_size, h, w]
+            auto out                = std::dynamic_pointer_cast<LocalDecoder>(blocks["local_decoder"])->forward(ctx, x, feat);  // [N, C, H, W]
+            return ggml_ext_scale(gctx, out, -1.f);
+        }
+    };
+
+    struct ZImageL2PRunner : public DiffusionModelRunner {
+        ZImageL2PConfig config;
+        ZImageL2PModel model;
+        std::vector<float> pe_vec;
+
+        ZImageL2PRunner(ggml_backend_t backend,
+                        const String2TensorStorage& tensors,
+                        const std::string& prefix,
+                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(ZImageL2PConfig::detect_from_weights(tensors, prefix)),
+              model(config) {
+            model.init(params_ctx, tensors, prefix);
+        }
+
+        std::string get_desc() override { return "z_image_l2p"; }
+
+        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 {
+            GGML_ASSERT(inputs.x != nullptr);
+            GGML_ASSERT(inputs.timesteps != nullptr);
+            if (inputs.ref_latents != nullptr && !inputs.ref_latents->empty()) {
+                LOG_ERROR("Z-Image L2P reference-image conditioning is not supported");
+                return {};
+            }
+            if (inputs.context == nullptr || inputs.context->empty()) {
+                LOG_ERROR("Z-Image L2P requires a text condition");
+                return {};
+            }
+            auto graph = [&]() {
+                auto gf      = new_graph_custom(ZImage::Z_IMAGE_GRAPH_SIZE);
+                auto x       = make_input(*inputs.x);
+                auto t       = make_input(*inputs.timesteps);
+                auto context = make_input(*inputs.context);
+                GGML_ASSERT(x->ne[3] == 1);
+                GGML_ASSERT(x->ne[0] % config.patch_size == 0 && x->ne[1] % config.patch_size == 0);
+                pe_vec      = finish_rope_pe(Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
+                                                                  static_cast<int>(x->ne[0]),
+                                                                  config.patch_size,
+                                                                  static_cast<int>(x->ne[3]),
+                                                                  static_cast<int>(context->ne[1]),
+                                                                  ZImage::SEQ_MULTI_OF,
+                                                                  {},
+                                                                  Rope::RefIndexMode::FIXED,
+                                                                  config.theta,
+                                                                  config.axes_dim));
+                int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+                auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
+                set_backend_tensor_data(pe, pe_vec.data());
+                auto ctx = get_context();
+                auto out = model.forward(&ctx, x, t, context, pe);
+                ggml_build_forward_expand(gf, out);
+                return gf;
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(graph, n_threads, false), inputs.x->dim());
+        }
+    };
+}  // namespace ZImageL2P
+
+#endif  // __SD_MODEL_DIFFUSION_Z_IMAGE_L2P_HPP__

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/minimax_h3_vae.hpp

@@ -228,11 +228,12 @@ namespace MiniMaxH3VAE {
         return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]);
     }
 
-    static ggml_tensor* apply_partial_rope(ggml_context* ctx,
+    static ggml_tensor* apply_partial_rope(GGMLRunnerContext* runner_ctx,
                                            ggml_tensor* x,
                                            ggml_tensor* pe) {
+        auto ctx        = runner_ctx->ggml_ctx;
         int64_t rot_dim = pe->ne[2] * 2;
-        auto rotated    = Rope::apply_rope(ctx,
+        auto rotated    = Rope::apply_rope(runner_ctx,
                                            ggml_ext_slice(ctx, x, 0, 0, rot_dim),
                                            pe,
                                            false);
@@ -289,8 +290,8 @@ namespace MiniMaxH3VAE {
                                                   batch_size);
             q                   = ggml_rms_norm(ctx->ggml_ctx, q, 1e-5f);
             k                   = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f);
-            q                   = apply_partial_rope(ctx->ggml_ctx, q, pe);
-            k                   = apply_partial_rope(ctx->ggml_ctx, k, pe);
+            q                   = apply_partial_rope(ctx, q, pe);
+            k                   = apply_partial_rope(ctx, k, pe);
             auto out            = ggml_ext_attention_ext(ctx,
                                                          q,
                                                          k,
@@ -633,8 +634,11 @@ namespace MiniMaxH3VAE {
             auto plan   = make_vae_temporal_tile_plan(input.shape()[2], {17, 0});
             auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& chunk, const VAETemporalTile& tile) {
                 SD_UNUSED(tile);
-                return VAE::encode(n_threads, chunk, tiling, circular_x, circular_y);
+                // keep the runner alive across chunks; ending it here would
+                // evict and reload the encoder weights every chunk
+                return VAE::encode(n_threads, chunk, tiling, circular_x, circular_y, false);
             });
+            runner_end();
             if (result.empty()) {
                 return {};
             }
@@ -695,13 +699,16 @@ namespace MiniMaxH3VAE {
                 {static_cast<int>(tokens_per_chunk + token_overlap), static_cast<int>(token_overlap)});
             GGML_ASSERT(plan.tiles.size() == static_cast<size_t>(num_chunks));
             auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& chunk, const VAETemporalTile& tile) {
+                // keep the runner alive across chunks; ending it here would
+                // evict and reload the decoder weights every chunk
                 auto decoded = VAE::decode(n_threads,
                                            chunk,
                                            tiling,
                                            true,
                                            circular_x,
                                            circular_y,
-                                           silent);
+                                           silent,
+                                           false);
                 if (decoded.empty()) {
                     return sd::Tensor<float>();
                 }
@@ -728,6 +735,7 @@ namespace MiniMaxH3VAE {
                 }
                 return first;
             });
+            runner_end();
             if (result.empty()) {
                 return {};
             }

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)) {
+        } 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;
@@ -241,6 +250,15 @@ public:
                                      sd_tiling_params_t tiling_params,
                                      bool circular_x = false,
                                      bool circular_y = false) {
+        return encode(n_threads, x, tiling_params, circular_x, circular_y, true);
+    }
+
+    sd::Tensor<float> encode(int n_threads,
+                             const sd::Tensor<float>& x,
+                             sd_tiling_params_t tiling_params,
+                             bool circular_x,
+                             bool circular_y,
+                             bool auto_end) {
         int64_t t0              = ggml_time_ms();
         tiling_params           = resolve_tiling_params(tiling_params);
         sd::Tensor<float> input = x;
@@ -280,7 +298,9 @@ public:
                                                   tiling_params);
         }
 
-        runner_end();
+        if (auto_end) {
+            runner_end();
+        }
 
         if (output.empty()) {
             LOG_ERROR("vae encode compute failed");
@@ -298,6 +318,17 @@ public:
                                      bool circular_x   = false,
                                      bool circular_y   = false,
                                      bool silent       = false) {
+        return decode(n_threads, x, tiling_params, decode_video, circular_x, circular_y, silent, true);
+    }
+
+    sd::Tensor<float> decode(int n_threads,
+                             const sd::Tensor<float>& x,
+                             sd_tiling_params_t tiling_params,
+                             bool decode_video,
+                             bool circular_x,
+                             bool circular_y,
+                             bool silent,
+                             bool auto_end) {
         int64_t t0              = ggml_time_ms();
         tiling_params           = resolve_tiling_params(tiling_params);
         sd::Tensor<float> input = x;
@@ -338,7 +369,9 @@ public:
                                                   tiling_params);
         }
 
-        runner_end();
+        if (auto_end) {
+            runner_end();
+        }
 
         if (output.empty()) {
             LOG_ERROR("vae decode compute failed");

src/model_loader.cpp

@@ -115,7 +115,7 @@ void convert_tensor(void* src,
                     ggml_type dst_type,
                     int nrows,
                     int n_per_row,
-                    std::vector<float> imatrix = {}) {
+                    std::vector<float> imatrix) {
     int n = nrows * n_per_row;
     if (src_type == dst_type) {
         size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type);
@@ -527,6 +527,9 @@ SDVersion ModelLoader::get_sd_version() const {
             if (tensor_storage_map.find("text_encoders.llm.connector.layers.0.self_attn.q_proj.weight") != tensor_storage_map.end()) {
                 return VERSION_MING_IMAGE;
             }
+            if (tensor_storage_map.find("model.diffusion_model.local_decoder.out_conv.weight") != tensor_storage_map.end()) {
+                return VERSION_Z_IMAGE_L2P;
+            }
             return VERSION_Z_IMAGE;
         }
         if (tensor_storage.name.find("double_stream_layers.0.img_instruct_attn.processor.img_to_q.weight") != std::string::npos) {
@@ -553,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_loader.h

@@ -13,6 +13,14 @@
 
 TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules);
 
+void convert_tensor(void* src,
+                    ggml_type src_type,
+                    void* dst,
+                    ggml_type dst_type,
+                    int nrows,
+                    int n_per_row,
+                    std::vector<float> imatrix = {});
+
 class MmapWrapper;
 
 struct ModelFileData {

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/name_conversion.cpp

@@ -824,6 +824,7 @@ std::string convert_diffusers_dit_to_original_lumina2(std::string name) {
     if (z_image_name_map.empty()) {
         z_image_name_map["all_x_embedder.2-1."]  = "x_embedder.";
         z_image_name_map["all_final_layer.2-1."] = "final_layer.";
+        z_image_name_map["all_x_embedder.16-1."] = "x_embedder.";
 
         // --- transformer blocks ---
         auto add_attention_map = [&](const std::string& prefix, int num) {
@@ -1036,7 +1037,7 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
         name = convert_hunyuan_video_to_original_flux(name);
     } else if (version == VERSION_MING_IMAGE) {
         name = convert_ming_image_dit_name(name);
-    } else if (sd_version_is_z_image(version)) {
+    } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version)) {
         name = convert_diffusers_dit_to_original_lumina2(name);
     } else if (sd_version_is_llada_image(version)) {
         name = convert_diffusers_dit_to_original_llada_image(name);

src/pipeline/diffusion_engine.cpp

@@ -107,6 +107,8 @@ const char* model_version_to_str[] = {
     "ESRGAN",
     "PixArt",
     "Ming-Image",
+    "Z-Image L2P",
+    "Iris-3B",
 };
 
 static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]),
@@ -1376,11 +1378,13 @@ bool StableDiffusionGGML::build_denoiser() {
                    sd_version_is_anima(version) ||
                    sd_version_is_ernie_image(version) ||
                    sd_version_is_z_image(version) ||
+                   sd_version_is_z_image_l2p(version) ||
                    version == VERSION_MING_IMAGE ||
                    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;
@@ -1388,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;
@@ -2172,7 +2176,7 @@ std::vector<float> StableDiffusionGGML::prepare_sample_timesteps(float sigma,
     if (version == VERSION_HIDREAM_O1) {
         return std::vector<float>{1.0f - (t / static_cast<float>(TIMESTEPS))};
     }
-    if (sd_version_is_z_image(version) || sd_version_is_ideogram4(version) || version == VERSION_MING_IMAGE) {
+    if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version) || sd_version_is_ideogram4(version) || version == VERSION_MING_IMAGE) {
         return std::vector<float>{1000.f - t};
     }
     return std::vector<float>{t};
@@ -2223,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,
@@ -2371,7 +2396,8 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
         noise *= eta;
     }
 
-    int64_t last_progress_us     = ggml_time_us();
+    int64_t last_progress_us = ggml_time_us();
+    sd_begin_preview_pass(static_cast<int>(steps));
     SamplePreviewContext preview = prepare_sample_preview_context();
 
     sd::Tensor<float> processed_init_latent       = denoiser->process_latent_in(init_latent);
@@ -2599,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) {
@@ -2638,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) {
@@ -2762,6 +2867,8 @@ 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) || 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;
         } else {
@@ -2792,7 +2899,9 @@ int StableDiffusionGGML::get_latent_channel() {
             latent_channel = 3;
         } else if (sd_version_is_sensenova_u1(version)) {
             latent_channel = 3;
-        } else if (sd_version_is_pid(version)) {
+        } else if (sd_version_is_z_image_l2p(version)) {
+            latent_channel = 3;
+        } else if (sd_version_is_pid(version) || version == VERSION_IRIS) {
             latent_channel = 3;
         } else if (sd_version_is_sefi_image(version)) {
             latent_channel = 144;
@@ -2893,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"
@@ -37,6 +38,7 @@
 #include "model/diffusion/unet.hpp"
 #include "model/diffusion/wan.hpp"
 #include "model/diffusion/z_image.hpp"
+#include "model/diffusion/z_image_l2p.hpp"
 #include "model/vae/auto_encoder_kl.hpp"
 #include "model/vae/hunyuan_vae.hpp"
 #include "model/vae/ltx_audio_vae.hpp"
@@ -393,6 +395,18 @@ namespace sd::model_builders {
                                                                       "model.diffusion_model",
                                                                       version,
                                                                       weight_manager);
+        } else if (sd_version_is_z_image_l2p(version)) {
+            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<ZImageL2P::ZImageL2PRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                            tensor_storage_map,
+                                                                            "model.diffusion_model",
+                                                                            weight_manager);
         } else if (sd_version_is_llada_image(version)) {
             result.conditioner = std::make_shared<LLaDAImageEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                       tensor_storage_map,
@@ -443,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++) {
@@ -514,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",
@@ -613,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)) {
+        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/pipeline/request.cpp

@@ -375,17 +375,19 @@ namespace sd::pipeline {
         total_steps = sample_steps + std::max(0, high_noise_sample_steps);
 
         if (sample_params->custom_sigmas_count > 0) {
-            sigmas      = std::vector<float>(sample_params->custom_sigmas,
+            sigmas = std::vector<float>(sample_params->custom_sigmas,
                                         sample_params->custom_sigmas + sample_params->custom_sigmas_count);
-            total_steps = static_cast<int>(sigmas.size()) - 1;
-            LOG_WARN("total_steps != custom_sigmas_count - 1, set total_steps to %d", total_steps);
-            if (sample_steps >= total_steps) {
-                sample_steps = total_steps;
-                LOG_WARN("total_steps != custom_sigmas_count - 1, set sample_steps to %d", sample_steps);
-            }
-            if (high_noise_sample_steps > 0) {
-                high_noise_sample_steps = total_steps - sample_steps;
-                LOG_WARN("total_steps != custom_sigmas_count - 1, set high_noise_sample_steps to %d", high_noise_sample_steps);
+            if (total_steps != sample_params->custom_sigmas_count - 1) {
+                total_steps = sample_params->custom_sigmas_count - 1;
+                LOG_WARN("total_steps != custom_sigmas_count - 1, set total_steps to %d", total_steps);
+                if ((sample_steps > total_steps) || (high_noise_sample_steps <= 0 && sample_steps != total_steps)) {
+                    sample_steps = total_steps;
+                    LOG_WARN("total_steps != custom_sigmas_count - 1, set sample_steps to %d", sample_steps);
+                }
+                if (high_noise_sample_steps > 0) {
+                    high_noise_sample_steps = total_steps - sample_steps;
+                    LOG_WARN("total_steps != custom_sigmas_count - 1, set high_noise_sample_steps to %d", high_noise_sample_steps);
+                }
             }
         } else {
             scheduler_t scheduler = resolve_scheduler(sd,

src/runtime/tiling.cpp

@@ -8,54 +8,51 @@
 #include "core/util.h"
 #include "ggml.h"
 
-static void sd_tiling_calc_tiles(int& num_tiles_dim,
-                                 float& tile_overlap_factor_dim,
-                                 int small_dim,
-                                 int tile_size,
-                                 const float tile_overlap_factor,
-                                 bool circular) {
-    int tile_overlap     = static_cast<int>(tile_size * tile_overlap_factor);
-    int non_tile_overlap = tile_size - tile_overlap;
+struct TileSpan {
+    int offset;
+    int overlap_before;
+    int overlap_after;
+};
 
-    if (circular) {
-        // circular means the last and first tile are overlapping (wraping around)
-        num_tiles_dim = small_dim / non_tile_overlap;
-
-        if (num_tiles_dim < 1) {
-            num_tiles_dim = 1;
-        }
-
-        tile_overlap_factor_dim = (tile_size - small_dim / num_tiles_dim) / (float)tile_size;
-
-        // if single tile and tile_overlap_factor is not 0, add one to ensure we have at least two overlapping tiles
-        if (num_tiles_dim == 1 && tile_overlap_factor_dim > 0) {
-            num_tiles_dim++;
-            tile_overlap_factor_dim = 0.5;
-        }
-
-        return;
-    }
-    // else, non-circular means the last and first tile are not overlapping
-
-    num_tiles_dim     = (small_dim - tile_overlap) / non_tile_overlap;
-    int overshoot_dim = ((num_tiles_dim + 1) * non_tile_overlap + tile_overlap) % small_dim;
-
-    if ((overshoot_dim != non_tile_overlap) && (overshoot_dim <= num_tiles_dim * (tile_size / 2 - tile_overlap))) {
-        // if tiles don't fit perfectly using the desired overlap
-        // and there is enough room to squeeze an extra tile without overlap becoming >0.5
-        num_tiles_dim++;
+static std::vector<TileSpan> sd_tiling_plan_axis(int dimension,
+                                                 int tile_size,
+                                                 float target_overlap,
+                                                 bool circular) {
+    if (tile_size >= dimension) {
+        return {{0, 0, 0}};
     }
 
-    tile_overlap_factor_dim = (float)(tile_size * num_tiles_dim - small_dim) / (float)(tile_size * (num_tiles_dim - 1));
-    if (num_tiles_dim <= 2) {
-        if (small_dim <= tile_size) {
-            num_tiles_dim           = 1;
-            tile_overlap_factor_dim = 0;
-        } else {
-            num_tiles_dim           = 2;
-            tile_overlap_factor_dim = (2 * tile_size - small_dim) / (float)tile_size;
-        }
+    const int extent = circular ? dimension : dimension - tile_size;
+    int intervals    = extent;
+    if (tile_size > 1) {
+        const float target = extent / ((1.0f - target_overlap) * tile_size);
+        // Adjacent tiles must overlap, but three tiles must not cover the same point.
+        const int min_intervals = 1 + (extent - 1) / (tile_size - 1);
+        const int max_intervals = std::max(1, static_cast<int>(2LL * extent / tile_size));
+        const int lower         = std::clamp(static_cast<int>(std::floor(target)), min_intervals, max_intervals);
+        const int upper         = std::clamp(static_cast<int>(std::ceil(target)), min_intervals, max_intervals);
+        auto error              = [&](int count) {
+            return std::abs(1.0f / count - 1.0f / target);
+        };
+        intervals = error(upper) < error(lower) ? upper : lower;
     }
+
+    const int num_tiles = circular ? intervals : intervals + 1;
+    auto position       = [&](int index) -> int {
+        return static_cast<int>(int64_t(index) * extent / intervals);
+    };
+    std::vector<TileSpan> tiles;
+    tiles.reserve(num_tiles);
+    for (int i = 0; i < num_tiles; ++i) {
+        const int offset = position(i);
+        // Keep wrapped neighbors in the same unwrapped coordinate space.
+        const int previous = i > 0 ? position(i - 1) : position(num_tiles - 1) - dimension;
+        const int next     = i + 1 < num_tiles ? position(i + 1) : dimension;
+        const int before   = i > 0 || circular ? previous + tile_size - offset : 0;
+        const int after    = i + 1 < num_tiles || circular ? offset + tile_size - next : 0;
+        tiles.push_back({offset, before, after});
+    }
+    return tiles;
 }
 
 static int64_t sd_tensor_plane_size(const sd::Tensor<float>& tensor) {
@@ -90,49 +87,50 @@ static void sd_tensor_merge_2d(const sd::Tensor<float>& input,
                                sd::Tensor<float>* output,
                                int x,
                                int y,
-                               int overlap_x,
-                               int overlap_y,
-                               bool circular_x,
-                               bool circular_y,
-                               int x_skip,
-                               int y_skip) {
+                               int overlap_left,
+                               int overlap_right,
+                               int overlap_top,
+                               int overlap_bottom) {
     GGML_ASSERT(output != nullptr);
-    int64_t width        = input.shape()[0];
-    int64_t height       = input.shape()[1];
-    int64_t img_width    = output->shape()[0];
-    int64_t img_height   = output->shape()[1];
-    int64_t input_plane  = sd_tensor_plane_size(input);
-    int64_t output_plane = sd_tensor_plane_size(*output);
-    int64_t plane_count  = input.numel() / input_plane;
-    GGML_ASSERT(output->numel() / output_plane == plane_count);
 
-    // unclamped -> expects x in the range [0-1]
+    int64_t in_width    = input.shape()[0];
+    int64_t in_height   = input.shape()[1];
+    int64_t out_width   = output->shape()[0];
+    int64_t out_height  = output->shape()[1];
+    int64_t in_size     = sd_tensor_plane_size(input);
+    int64_t out_size    = sd_tensor_plane_size(*output);
+    int64_t plane_count = input.numel() / in_size;
+
+    GGML_ASSERT(output->numel() == plane_count * out_size);
+    GGML_ASSERT(x >= 0 && y >= 0);
+    GGML_ASSERT(x < out_width && in_width <= out_width);
+    GGML_ASSERT(y < out_height && in_height <= out_height);
+    GGML_ASSERT(overlap_left >= 0 && overlap_right >= 0);
+    GGML_ASSERT(overlap_top >= 0 && overlap_bottom >= 0);
+
     auto smootherstep_f32 = [](const float x) -> float {
-        GGML_ASSERT(x >= 0.f && x <= 1.f);
         return x * x * x * (x * (6.0f * x - 15.0f) + 10.0f);
     };
-
-    for (int iy = y_skip; iy < height; iy++) {
-        for (int ix = x_skip; ix < width; ix++) {
-            int64_t src_xy = ix + width * iy;
-            int64_t ox     = (x + ix) % img_width;
-            int64_t oy     = (y + iy) % img_height;
-            int64_t dst_xy = ox + img_width * oy;
-            for (int64_t plane = 0; plane < plane_count; ++plane) {
-                float new_value = input[plane * input_plane + src_xy];
-                if (overlap_x > 0 || overlap_y > 0) {
-                    float old_value   = (*output)[plane * output_plane + dst_xy];
-                    const float x_f_0 = (circular_x || (overlap_x > 0 && x > 0)) ? (ix - x_skip) / float(overlap_x) : 1.f;
-                    const float x_f_1 = (circular_x || (overlap_x > 0 && x < (img_width - width))) ? (width - ix) / float(overlap_x) : 1.f;
-                    const float y_f_0 = (circular_y || (overlap_y > 0 && y > 0)) ? (iy - y_skip) / float(overlap_y) : 1.f;
-                    const float y_f_1 = (circular_y || (overlap_y > 0 && y < (img_height - height))) ? (height - iy) / float(overlap_y) : 1.f;
-                    const float x_f   = std::min(std::min(x_f_0, x_f_1), 1.f);
-                    const float y_f   = std::min(std::min(y_f_0, y_f_1), 1.f);
-                    (*output)[plane * output_plane + dst_xy] =
-                        old_value + new_value * smootherstep_f32(y_f) * smootherstep_f32(x_f);
-                } else {
-                    (*output)[plane * output_plane + dst_xy] = new_value;
+    for (int64_t plane = 0; plane < plane_count; ++plane) {
+        for (int iy = 0; iy < in_height; ++iy) {
+            float y_f = 1.0f;
+            if (iy < overlap_top) {
+                y_f = static_cast<float>(iy) / overlap_top;
+            }
+            if (iy >= in_height - overlap_bottom) {
+                y_f = static_cast<float>(in_height - iy) / overlap_bottom;
+            }
+            const float y_weight = smootherstep_f32(std::clamp(y_f, 0.0f, 1.0f));
+            for (int ix = 0; ix < in_width; ++ix) {
+                float x_f = 1.0f;
+                if (ix < overlap_left) {
+                    x_f = static_cast<float>(ix) / overlap_left;
                 }
+                if (ix >= in_width - overlap_right) {
+                    x_f = static_cast<float>(in_width - ix) / overlap_right;
+                }
+                float x_weight = smootherstep_f32(std::clamp(x_f, 0.0f, 1.0f));
+                (*output)[plane * out_size + out_width * ((y + iy) % out_height) + ((x + ix) % out_width)] += x_weight * y_weight * input[plane * in_size + in_width * iy + ix];
             }
         }
     }
@@ -158,111 +156,53 @@ sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
     GGML_ASSERT(((input_width / output_width) == scale) ||
                 ((output_width / input_width) == scale));
 
-    int small_width  = output_width;
-    int small_height = output_height;
-    bool decode      = output_width > input_width;
-    if (decode) {
-        small_width  = input_width;
-        small_height = input_height;
-    }
+    bool decode      = output_width > input_width;  // scale up
+    int small_width  = decode ? input_width : output_width;
+    int small_height = decode ? input_height : output_height;
+    int scale_in     = decode ? 1 : scale;
+    int scale_out    = decode ? scale : 1;
 
-    int num_tiles_x;
-    float tile_overlap_factor_x;
-    sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_w, tile_overlap_factor, circular_x);
+    const auto tiles_x = sd_tiling_plan_axis(small_width, p_tile_size_w, tile_overlap_factor, circular_x);
+    const auto tiles_y = sd_tiling_plan_axis(small_height, p_tile_size_h, tile_overlap_factor, circular_y);
 
-    int num_tiles_y;
-    float tile_overlap_factor_y;
-    sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_h, tile_overlap_factor, circular_y);
+    int tile_width         = std::min(p_tile_size_w, small_width);
+    int tile_height        = std::min(p_tile_size_h, small_height);
+    int input_tile_width   = tile_width * scale_in;
+    int input_tile_height  = tile_height * scale_in;
+    int output_tile_width  = tile_width * scale_out;
+    int output_tile_height = tile_height * scale_out;
 
-    int tile_overlap_x     = static_cast<int32_t>(p_tile_size_w * tile_overlap_factor_x);
-    int non_tile_overlap_x = p_tile_size_w - tile_overlap_x;
-    int tile_overlap_y     = static_cast<int32_t>(p_tile_size_h * tile_overlap_factor_y);
-    int non_tile_overlap_y = p_tile_size_h - tile_overlap_y;
-    int tile_size_w        = p_tile_size_w < small_width ? p_tile_size_w : small_width;
-    int tile_size_h        = p_tile_size_h < small_height ? p_tile_size_h : small_height;
-    int input_tile_size_w  = tile_size_w;
-    int input_tile_size_h  = tile_size_h;
-    int output_tile_size_w = tile_size_w;
-    int output_tile_size_h = tile_size_h;
-    if (decode) {
-        output_tile_size_w *= scale;
-        output_tile_size_h *= scale;
-    } else {
-        input_tile_size_w *= scale;
-        input_tile_size_h *= scale;
-    }
-
-    int num_tiles   = num_tiles_x * num_tiles_y;
-    int tile_count  = 1;
-    bool last_y     = false;
-    bool last_x     = false;
-    float last_time = 0.0f;
+    const int num_tiles = static_cast<int>(tiles_x.size() * tiles_y.size());
+    int tile_count      = 0;
     if (!silent) {
-        LOG_VERBOSE("num tiles : %d, %d ", num_tiles_x, num_tiles_y);
-        LOG_VERBOSE("optimal overlap : %f, %f (targeting %f)", tile_overlap_factor_x, tile_overlap_factor_y, tile_overlap_factor);
+        LOG_VERBOSE("num tiles : %d, %d ", static_cast<int>(tiles_x.size()), static_cast<int>(tiles_y.size()));
         LOG_VERBOSE("processing %i tiles", num_tiles);
         pretty_progress(0, num_tiles, 0.0f);
     }
-    for (int y = 0; y < small_height && !last_y; y += non_tile_overlap_y) {
-        int dy = 0;
-        if (!circular_y && y + tile_size_h >= small_height) {
-            int original_y = y;
-            y              = small_height - tile_size_h;
-            dy             = original_y - y;
-            if (decode) {
-                dy *= scale;
-            }
-            last_y = true;
-        }
-        for (int x = 0; x < small_width && !last_x; x += non_tile_overlap_x) {
-            int dx = 0;
-            if (!circular_x && x + tile_size_w >= small_width) {
-                int original_x = x;
-                x              = small_width - tile_size_w;
-                dx             = original_x - x;
-                if (decode) {
-                    dx *= scale;
-                }
-                last_x = true;
-            }
-
-            int x_in  = decode ? x : scale * x;
-            int y_in  = decode ? y : scale * y;
-            int x_out = decode ? x * scale : x;
-            int y_out = decode ? y * scale : y;
-
-            int overlap_x_out = decode ? tile_overlap_x * scale : tile_overlap_x;
-            int overlap_y_out = decode ? tile_overlap_y * scale : tile_overlap_y;
-
+    for (const auto& y : tiles_y) {
+        for (const auto& x : tiles_x) {
             int64_t t1       = ggml_time_ms();
-            auto input_tile  = sd_tensor_split_2d(input, input_tile_size_w, input_tile_size_h, x_in, y_in);
+            auto input_tile  = sd_tensor_split_2d(input, input_tile_width, input_tile_height, x.offset * scale_in, y.offset * scale_in);
             auto output_tile = on_processing(input_tile);
             if (output_tile.empty()) {
                 return {};
             }
-            GGML_ASSERT(output_tile.shape()[0] == output_tile_size_w && output_tile.shape()[1] == output_tile_size_h);
+            GGML_ASSERT(output_tile.shape()[0] == output_tile_width && output_tile.shape()[1] == output_tile_height);
             if (output.empty()) {
                 std::vector<int64_t> output_shape = output_tile.shape();
                 output_shape[0]                   = output_width;
                 output_shape[1]                   = output_height;
                 output                            = sd::Tensor<float>::zeros(std::move(output_shape));
             }
-            sd_tensor_merge_2d(output_tile, &output, x_out, y_out, overlap_x_out, overlap_y_out, circular_x, circular_y, dx, dy);
-
+            sd_tensor_merge_2d(output_tile, &output,
+                               x.offset * scale_out, y.offset * scale_out,
+                               x.overlap_before * scale_out, x.overlap_after * scale_out,
+                               y.overlap_before * scale_out, y.overlap_after * scale_out);
+            ++tile_count;
             if (!silent) {
-                int64_t t2 = ggml_time_ms();
-                last_time  = (t2 - t1) / 1000.0f;
-                pretty_progress(tile_count, num_tiles, last_time);
+                pretty_progress(tile_count, num_tiles, (ggml_time_ms() - t1) / 1000.0f);
             }
-            tile_count++;
         }
-        last_x = false;
-    }
-    if (!silent && tile_count < num_tiles) {
-        pretty_progress(num_tiles, num_tiles, last_time);
-    }
-    if (output.empty()) {
-        return {};
     }
     return output;
 }

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/tokenizers/bpe_tokenizer.cpp

@@ -242,7 +242,7 @@ bool BPETokenizer::encode(const std::string& text, std::vector<int>& result, on_
         ss << "\"" << token << "\", ";
     }
     ss << "]";
-    LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", text.c_str(), bpe_tokens.size(), ss.str().c_str());
+    LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", escape_newlines(text).c_str(), bpe_tokens.size(), escape_newlines(ss.str()).c_str());
     result = std::move(bpe_tokens);
     return true;
 }

src/tokenizers/hf_tokenizer.cpp

@@ -735,7 +735,7 @@ bool HFTokenizer::encode(const std::string& text, std::vector<int>& tokens, on_n
         }
     }
     ss << "]";
-    LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", text.c_str(), result.size(), ss.str().c_str());
+    LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", escape_newlines(text).c_str(), result.size(), escape_newlines(ss.str()).c_str());
     tokens = std::move(result);
     return true;
 }

src/tokenizers/t5_unigram_tokenizer.cpp

@@ -343,7 +343,7 @@ bool T5UniGramTokenizer::encode(const std::string& input, std::vector<int>& resu
         ss << "\"" << token_str << "\", ";
     }
     ss << "]";
-    LOG_VERBOSE("split prompt \"%s\" to tokens %s", input.c_str(), ss.str().c_str());
+    LOG_VERBOSE("split prompt \"%s\" to tokens %s", escape_newlines(input).c_str(), escape_newlines(ss.str()).c_str());
 
     result = std::move(tokens);
     return true;

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);
 };