feat: add MoSR upscaler support (#2128)

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"
@@ -560,6 +560,18 @@ jobs:
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
         uses: pr-mpt/actions-commit-hash@v2
 
+      - name: Bundle HIP runtime libraries
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
+        run: |
+          $ErrorActionPreference = "Stop"
+          $binPath = (rocm-sdk path --bin).Trim()
+          if (-not $binPath) { throw "rocm-sdk path --bin returned empty" }
+          foreach ($dll in @("amdhip64_7.dll", "rocm_kpack.dll", "amd_comgr.dll")) {
+            $f = Get-ChildItem -Path $binPath -Filter $dll -ErrorAction SilentlyContinue
+            if (-not $f) { throw "no match for $dll in $binPath" }
+            Copy-Item $f.FullName -Destination .\build\bin\ -Force
+          }
+
       - name: Pack artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
         run: |
@@ -582,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
@@ -598,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
@@ -634,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)
@@ -64,6 +65,9 @@ API and command-line option may change frequently.***
     - [HiDream-O1-Image](./docs/hidream_o1_image.md)
     - [Ideogram4](./docs/ideogram4.md)
     - [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)
@@ -179,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)

docker/Dockerfile

@@ -40,4 +40,6 @@ RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \
     printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \
     chmod +x /sd-cli /sd-server
 
+WORKDIR /sd.cpp
+
 ENTRYPOINT [ "/sd-cli" ]

docker/Dockerfile.cuda

@@ -57,4 +57,6 @@ RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \
     printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \
     chmod +x /sd-cli /sd-server
 
+WORKDIR /sd.cpp
+
 ENTRYPOINT [ "/sd-cli" ]

docker/Dockerfile.musa

@@ -42,4 +42,6 @@ RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \
     printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \
     chmod +x /sd-cli /sd-server
 
+WORKDIR /sd.cpp
+
 ENTRYPOINT [ "/sd-cli" ]

docker/Dockerfile.sycl

@@ -29,4 +29,6 @@ FROM intel/oneapi-basekit:${SYCL_VERSION}-devel-ubuntu24.04 AS runtime
 COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli
 COPY --from=build /sd.cpp/build/bin/sd-server /sd-server
 
+WORKDIR /sd.cpp
+
 ENTRYPOINT [ "/sd-cli" ]

docker/Dockerfile.vulkan

@@ -41,4 +41,6 @@ RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \
     printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \
     chmod +x /sd-cli /sd-server
 
+WORKDIR /sd.cpp
+
 ENTRYPOINT [ "/sd-cli" ]

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/int8_convrot.md

@@ -84,17 +84,18 @@ The floating-point output is reconstructed as
 Y[r, o] ~= A[r, o] * s_x[r] * s_w[o] + b[o]
 ```
 
-The packed runtime activation tensor contains the I8 activation rows and their floating-point row scales. Linear layers that share the same input and convrot group size reuse this packed tensor, avoiding repeated rotation and activation quantization within the graph.
+The packed runtime activation tensor contains the I8 activation rows and their floating-point row scales. Linear layers with an input scale of `1` that share the same input and convrot group size reuse this packed tensor, avoiding repeated rotation and activation quantization within the graph. For other input scales, activations are scaled before explicit convrot quantization, and the linear output is unscaled before adding bias.
 
 ## Backend support
 
 - CPU provides the portable regular Hadamard, activation quantization, INT8 matrix multiplication, and scale restoration implementations.
 - NVIDIA CUDA devices with compute capability 7.5 or newer use the native accelerated path. For H256, CUDA fuses the rotation, row-wise maximum reduction, and activation quantization. It uses cuBLAS for I8 x I8 to I32 GEMM and a CUDA kernel for scale restoration and bias addition.
-- Vulkan and other GPU backends do not currently have dedicated INT8 convrot kernels. They use the backend scheduler to fall back to CPU, which is expected to be substantially slower than the CUDA path.
+- AMD HIP devices in the CDNA, RDNA3 (including RDNA3.5), and RDNA4 families use the same INT8/H256 kernels with hipBLAS for I8 x I8 to I32 GEMM. Other AMD architectures fall back to CPU for INT8 matrix multiplication.
+- Vulkan provides native H256 activation quantization and INT8 matrix multiplication when the build and device support accelerated packed INT8 dot products. Unsupported configurations and GPU backends without these kernels use the backend scheduler to fall back to CPU.
 
 LoRA adapters are applied at runtime without modifying the INT8 weights. The INT8 convrot path computes the base linear output, while LoRA, LoHa, LoKr, and raw weight-difference adapters compute their output corrections from the original, unrotated activation and add them to the base output. `--lora-apply-mode auto` selects this path for models containing INT8 tensorwise weights. If `immediately` is requested, sd.cpp falls back to runtime application because merging an adapter would require dequantizing and rotating its weight update, then recalculating the per-row scales and requantizing the result.
 
-The dedicated CUDA convrot activation path currently requires a group size of `256`; other supported group sizes use CPU execution.
+The dedicated CUDA and HIP convrot activation paths currently require a group size of `256`; other supported group sizes use CPU execution.
 
 ## 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/mage_flow.md

@@ -4,20 +4,8 @@
 
 ## Download weights
 
-- Download Mage-Flow
-    - safetensors: https://huggingface.co/microsoft/Mage-Flow/tree/main/transformer
-- Download Mage-Flow-Base
-    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Base/tree/main/transformer
-- Download Mage-Flow-Turbo
-    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Turbo/tree/main/transformer
-- Download Mage-Flow-Edit
-    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit/tree/main/transformer
-- Download Mage-Flow-Edit-Turbo
-    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo/tree/main/transformer
-- Download Mage-Flow-Edit-Base
-    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit-Base/tree/main/transformer
-- Download Mage-Flow vae
-    - safetensors: https://huggingface.co/microsoft/Mage-Flow/tree/main/vae
+- Download Mage-Flow diffusion from https://huggingface.co/Comfy-Org/Mage-Flow/tree/main/diffusion_models
+- Download Mage-Flow vae from https://huggingface.co/Comfy-Org/Mage-Flow/tree/main/vae
 - Download Qwen3-VL 4B
     - safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders
     - gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main

docs/ming_image.md

@@ -0,0 +1,24 @@
+# Ming-Image
+
+[Ming-Image](https://github.com/inclusionAI/Ming-Image) 0.1 Design uses a 6B diffusion transformer (DiT), Ling-mini-2.0 for text conditioning, and the Ming-Image VAE. Text-to-image generation with RGBA output is supported.
+
+## Download weights
+
+- Download Ming-Image 0.1 Design DiT
+    - safetensors: https://huggingface.co/Comfy-Org/Ming-Image/tree/main/diffusion_models
+- Download Ling-mini-2.0 BF16
+    - safetensors: https://huggingface.co/Comfy-Org/Ming-Image/tree/main/text_encoders
+- Download Ming-Image VAE
+    - safetensors: https://huggingface.co/Comfy-Org/Ming-Image/tree/main/vae
+- Download Ling tokenizer
+    - tokenizer.json: https://huggingface.co/inclusionAI/Ming-Image-0.1-Design/blob/main/mllm/tokenizer.json
+
+The example below uses `ming_image_0.1_design_bf16.safetensors` for the DiT. You can also use `ming_image_0.1_design_int8_convrot.safetensors` with [INT8 convrot support](int8_convrot.md). Use the BF16 text encoder.
+
+## Text-to-image
+
+Pass the Ling `tokenizer.json` with `--tokenizer` and use the matching Ming-Image VAE. Image dimensions must be multiples of 16. Save the output as PNG to preserve the alpha channel.
+
+```bash
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\ming_image_0.1_design_bf16.safetensors --llm ..\models\text_encoders\ming_image_0.1_ling_mini_2.0_bf16.safetensors --vae ..\models\vae\ming_image_vae_bf16.safetensors --tokenizer ..\models\text_encoders\tokenizer.json -p "A cheerful orange cat sticker, transparent background" --width 1024 --height 1024 --steps 12 --cfg-scale 1 --sampling-method euler --diffusion-fa -v --offload-to-cpu -o ming_image.png
+```

docs/performance.md

@@ -21,6 +21,64 @@ 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
+default tile size is 256x256 **image pixels**, independent of the VAE scale factor:
+
+```shell
+--vae-tiling --vae-tile-size 256x256 --vae-tile-overlap 0.5
+```
+
+`--vae-tile-size` accepts one size or `WIDTHxHEIGHT`. A zero dimension uses the
+256-pixel default. Sizes are rounded down to a multiple of the VAE scale factor
+and capped at the current input dimensions. Explicit sizes below four latent
+pixels per axis (or the full axis when it is smaller) are rejected. Encoding and
+decoding use the same spatial sizes, without an additional encoding multiplier.
+Inputs that fit within a tile are processed as one tile.
+
+For a 512x512 image with the default 50% overlap, both encoding and decoding use
+3x3 tiles. A 256-pixel tile corresponds to 32 latent pixels for an 8x VAE, 16 for
+a 16x VAE, and 8 for a 32x VAE. Smaller tiles reduce each graph's memory demand,
+but overlapping work can increase processing time and tiling can affect image
+quality, especially during encoding. Use larger tiles when more context is needed.
+
+`--vae-relative-tile-size` overrides the absolute size on each axis with a positive
+value. Values up to and including 1 specify a fraction of the current input size;
+values greater than 1 specify a target number of tiles per axis, accounting for
+overlap. For example, `0.5x0.5` uses half the width and height in both encode and
+decode. The target overlap is clamped to 0 through 0.5 and the actual overlap is
+adjusted to fit the image. Size and overlap options require `--vae-tiling`.
+
+**Migration:** `--vae-tile-size` and the C/JSON fields `tile_size_w` and
+`tile_size_h` now use image pixels instead of latent units. The C/JSON fields
+`tile_size_x/y` have been renamed to `tile_size_w/h`, and `rel_size_x/y` to
+`rel_size_w/h`. The command-line option names are unchanged. For example, an old
+decode tile size of 32 corresponds to 256 pixels for an 8x VAE or 512 pixels for a
+16x VAE. Encoding no longer enlarges explicit or relative tile sizes.
+
+The main VAE decode path retries allocation failures with smaller tiles, even
+without `--vae-tiling`. Supported video VAEs first try temporal tiling; spatial
+retries use at most 256-pixel tiles initially and then halve the effective tile
+dimensions down to the minimum size. Each spatial retry must reduce the effective
+tile size. These runtime adjustments do not change the caller's parameters.
+Execution failures are not retried, and encoding has no automatic OOM retry.
+
+`--temporal-tiling` remains independent of spatial tiling. MiniMax H3 always uses
+spatial tiling (256x256 pixels and 25% overlap by default) and its own temporal
+windows. With `--vae-tiling`, its overlap follows `--vae-tile-overlap`; explicit
+spatial sizes are honored.
+
 ## Offload weights to the CPU to save VRAM without reducing generation speed.
 
 Using `--offload-to-cpu` allows you to offload weights to the CPU, saving VRAM without reducing generation speed.
@@ -63,6 +121,12 @@ See [backend selection](./backend.md) for full syntax.
 
 When a graph has cut markers and its missing weights plus incremental compute workspace exceed the available device headroom, it runs its fixed segment list in order. A reusable monolithic compute buffer is not counted as a new allocation. An explicit `--max-vram` budget deducts already-resident managed weights and compute/cache buffers registered by every runner sharing the device, so later graph runs remain segmented when the full graph exceeds the budget. The current segment's weights are pinned during compute, and the next parameter-bearing segment is prefetched when the device supports asynchronous transfer. No opt-in streaming flag is required.
 
+When choosing between monolithic and segmented execution, the runner requires
+an additional 128 MiB of headroom in both available device memory and any explicit
+managed budget. This planning headroom absorbs small allocation estimate changes;
+subsequent capacity checks can consume it while still preserving the 512 MiB device
+scratch reserve and respecting the managed budget.
+
 - `--max-vram <GiB>` optionally lowers the live-memory limit. A positive value is a managed per-device budget, `0` uses the device's current free memory without an explicit budget, and a negative value snapshots free memory at startup while reserving that many GiB (`--max-vram -1` reserves about 1 GiB). Driver contexts and unrelated external allocations remain outside the managed budget.
 - `--disable-prefetch` disables asynchronous next-segment prefetch while retaining synchronous loading, eviction, and segmented execution.
 - `--disable-segmented-compute` forces monolithic graph execution for diagnostics or compatibility, even when the automatic memory check would select segments.

docs/pixart.md

@@ -0,0 +1,46 @@
+# How to Use
+
+You can run PixArt-α / PixArt-Σ with stable-diffusion.cpp.
+
+PixArt is a DiT-based text-to-image model family conditioned by a T5-XXL text
+encoder and a 4-channel VAE: SDXL-style for PixArt-Σ and SD1.x-style for PixArt-α.
+
+## Download weights
+
+- Download the transformer (diffusion model)
+    - PixArt-Σ XL-2 1024-MS: https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS/tree/main/transformer
+    - PixArt-α XL-2 1024-MS: https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS/tree/main/transformer
+- Download the T5-XXL text encoder
+    - safetensors: https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS/tree/main/text_encoder
+- Download the VAE
+    - PixArt-Σ: https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS/tree/main/vae
+    - PixArt-α: https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS/tree/main/vae
+    - Use the VAE matching the checkpoint's latent space. For TAE decoding or
+      preview, use TAESDXL for PixArt-Σ and TAESD for PixArt-α.
+- Tokenizer: the T5 vocabulary is embedded; no extra tokenizer file is needed.
+
+## Examples
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\pixart_sigma_xl2_1024_ms.safetensors --t5xxl ..\models\text_encoders\t5xxl.safetensors --vae ..\models\vae\pixart_vae.safetensors -p "a lovely cat" --cfg-scale 4.5 -W 1024 -H 1024 --steps 20 -v
+```
+
+## Notes
+
+- The VAE scaling factor defaults to `0.13025` for PixArt-Σ. PixArt-α
+  checkpoints with resolution micro-condition weights use `0.18215`.
+  PixArt-α 512 has the same tensor layout as PixArt-Σ, so it requires an
+  explicit override: `--model-args "pixart_vae_scale_factor=0.18215"`.
+  This argument can also override the scale for other compatible checkpoints.
+- PixArt-Σ checkpoints compute 2D sincos positional embeddings at runtime;
+  the trained grid is 64x64 patches with an interpolation scale of 2.
+  For checkpoints trained at a different resolution, the positional embedding
+  parameters can be adjusted via model args:
+  `--model-args "pixart_pos_embed_base_size=<trained grid>,pixart_interpolation_scale=<scale>"`
+  (e.g. `pixart_pos_embed_base_size=32,pixart_interpolation_scale=1,pixart_vae_scale_factor=0.18215` for
+  PixArt-α XL-2 512).
+- Checkpoints carrying resolution/aspect-ratio micro-condition weights are
+  detected but those conditions are not applied yet; a warning is logged and
+  generation proceeds with the timestep embedding only.
+- The transformer predicts 8 channels (noise + learned variance); only the
+  noise half is used for sampling, matching the reference implementation.

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.
@@ -23,7 +25,7 @@ Run the following commands from the build directory. Use image dimensions divisi
 ### Text to image
 
 ```powershell
-.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf -p "a lovely cat holding a sign says 'qwen2.1.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu -o qwen_image_2.1.png
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf -p "a lovely cat holding a sign says 'qwen2.1.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu --fa -o qwen_image_2.1.png
 ```
 
 <img alt="Qwen Image 2.1 example" src="../assets/qwen/qwen_image_2.1.png" />
@@ -35,7 +37,7 @@ To use GGUF diffusion weights, set `--diffusion-model` to the path of a file suc
 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`.
 
 ```powershell
-.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3VL-8B-Instruct-mmproj-BF16.gguf -r ..\assets\qwen\qwen_image_2.1.png -p "change 'qwen2.1.cpp' to 'sd.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu -o qwen_image_2.1_edit.png
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3VL-8B-Instruct-mmproj-BF16.gguf -r ..\assets\qwen\qwen_image_2.1.png -p "change 'qwen2.1.cpp' to 'sd.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu --fa -o qwen_image_2.1_edit.png
 ```
 
 For multiple reference images, repeat `-r` in the desired order, for example `-r first.png -r second.png`.
@@ -44,7 +46,27 @@ For multiple reference images, repeat `-r` in the desired order, for example `-r
 
 By default, the first denoising call for each fixed condition saves the text and reference-image keys and values from every transformer layer. Later calls only compute the target-image tokens. Positive and negative conditions use separate caches, which are released when sampling ends.
 
-The cache uses FP32 on all attention backends. For the default 32-layer model, a prefix of 4096 tokens takes about 4 GiB per condition, in addition to weights and working buffers. The runner accounts for the cache when checking the memory budget. If a cached execution runs out of memory, it releases the prefix caches, disables caching for the rest of that sampling run, and retries the full sequence once. Per-step conditioning extensions currently use the full-sequence path.
+Set `qwen_image_2_1_prefix_cache_type` in `--model-args` to `auto` or a type name using the same parser and case-sensitive names as `--type`:
+
+- `auto` (default): use FP16 only when Flash Attention is enabled, Sage Attention is disabled, the attention scale is unchanged, and every attention operation in the cache-writing or cache-reading graph selects Flash Attention after backend support checks. If an operation falls back, rebuild the prefix in FP32 before executing and keep FP32 for the rest of that sampling run.
+- `f32`: always store FP32 keys and values.
+- `f16`: always store FP16 keys and values, including with ordinary attention or custom attention scaling. This saves cache memory but can introduce additional rounding error.
+- Other types, such as `bf16`, `q4_1`, `q5_0`, `q5_1`, `q8_0`, `q4_K`, `q6_K`, `iq4_nl`, and `iq4_xs`: use the requested storage type if the ggml build provides runtime conversion to and from FP32. Quantization is lossy and must be selected explicitly; `auto` never selects a quantized type.
+
+Cache data is packed into contiguous rows of `hidden_size` elements before conversion, so 256-element quantization blocks work with the model's 128-element attention heads without padding. The type's block size must divide `hidden_size`. Unknown types, types lacking runtime conversion (for example `q8_1` and several IQ formats), and incompatible block sizes are ignored with a warning, leaving the previous setting or the default `auto` unchanged.
+
+For example, use `--model-args qwen_image_2_1_prefix_cache_type=q8_0` to enable 8-bit cache storage. Cached keys and values are converted back to the attention input type before concatenating with the current target tokens. This reduces persistent cache memory; attention working buffers still use floating-point values, and conversion adds work on each step. Backends without the required conversion operations use the existing CPU fallback.
+
+For the default 32-layer model, a prefix of 4096 tokens takes approximately the following memory per condition, excluding weights, working buffers, and allocation overhead:
+
+| Cache type | Memory |
+| --- | ---: |
+| `f32` | 4 GiB |
+| `f16` | 2 GiB |
+| `q8_0` | 1.0625 GiB |
+| `q4_0` | 0.5625 GiB |
+
+The runner accounts for the cache when checking the memory budget. If a cached execution runs out of memory, it releases the prefix caches, disables caching for the rest of that sampling run, and retries the full sequence once. Per-step conditioning extensions currently use the full-sequence path.
 
 Disable this optimization with `--model-args qwen_image_2_1_prefix_cache=false`. It reuses step-independent activations; numerical results can still differ slightly because the matrix sizes change.
 

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/wan.md

@@ -40,6 +40,9 @@ Wan models require `-M vid_gen`, including single-frame generation. `--video-fra
             - safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/tree/main/split_files/diffusion_models
             - gguf: https://huggingface.co/QuantStack/Wan2.2-S2V-14B-GGUF/tree/main
             - int8_convrot safetensors: https://huggingface.co/noctrex/Wan2.2-S2V-14B-int8_convrot
+        - Wan2.2 VACE-Fun A14B
+            - safetensors: https://huggingface.co/alibaba-pai/Wan2.2-VACE-Fun-A14B
+            - gguf: https://huggingface.co/QuantStack/Wan2.2-VACE-Fun-A14B-GGUF/tree/main
 - Download vae
     - wan_2.1_vae (for all the wan model except Wan2.2 TI2V 5B)
         - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors
@@ -256,3 +259,25 @@ ffmpeg -i ..\..\ComfyUI\input\post+depth.mp4 -qscale:v 1 -vf fps=8 post+depth\fr
 ```
 
 <video src=../assets/wan/Wan2.1_14B_vace_v2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>
+
+### Wan2.2 VACE-Fun A14B
+
+VACE-Fun runs as a MoE pair: `--diffusion-model` takes the low-noise expert and
+`--high-noise-diffusion-model` the high-noise one. Reference-to-video uses `-i`
+for the reference image, same as Wan2.1 VACE.
+
+#### R2V
+
+```
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\Wan2.2-VACE-Fun-A14B-low-noise-Q8_0.gguf --high-noise-diffusion-model ..\models\diffusion_models\Wan2.2-VACE-Fun-A14B-high-noise-Q8_0.gguf --vae ..\models\vae\wan_2.1_vae.safetensors --t5xxl ..\models\text_encoders\umt5-xxl-encoder-Q8_0.gguf -p "a lovely cat" --cfg-scale 3.5 --sampling-method euler --steps 10 --high-noise-cfg-scale 3.5 --high-noise-sampling-method euler --high-noise-steps 8 -v -n "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部， 畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" -W 832 -H 480 --diffusion-fa -i ..\assets\cat_with_sd_cpp_42.png --video-frames 33 --offload-to-cpu
+```
+
+<video src=../assets/wan/Wan2.2_A14B_vace_r2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>
+
+#### T2V
+
+Same command without `-i` (VACE context is synthesized from an empty control
+video, like Wan2.1 VACE t2v).
+
+> On GPUs with ~12 GB VRAM, VACE also needs `--vae-tiling` — the control-video
+> encode can exceed the budget otherwise.

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 = {
@@ -518,7 +535,8 @@ ArgOptions SDContextParams::get_options() {
         {"",
          "--model-args",
          "extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, "
-         "chroma_t5_mask_pad, qwen_image_zero_cond_t, qwen_image_2_1_prefix_cache",
+         "chroma_t5_mask_pad, qwen_image_zero_cond_t, qwen_image_2_1_prefix_cache, "
+         "qwen_image_2_1_prefix_cache_type (auto or a type name from --type)",
          (int)',',
          &model_args},
         {"",
@@ -636,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;
@@ -741,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). "
@@ -939,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"
@@ -1021,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();
@@ -1340,7 +1351,7 @@ ArgOptions SDGenerationParams::get_options() {
          &embed_image_metadata},
         {"",
          "--vae-tiling",
-         "process vae in tiles to reduce memory usage",
+         "process vae encode and decode in spatial tiles to reduce memory usage (default: 256x256 image pixels)",
          true,
          &vae_tiling_params.enabled},
         {"",
@@ -1604,12 +1615,12 @@ ArgOptions SDGenerationParams::get_options() {
         size_t x_pos              = tile_size_str.find('x');
         try {
             if (x_pos != std::string::npos) {
-                std::string tile_x_str        = tile_size_str.substr(0, x_pos);
-                std::string tile_y_str        = tile_size_str.substr(x_pos + 1);
-                vae_tiling_params.tile_size_x = std::stoi(tile_x_str);
-                vae_tiling_params.tile_size_y = std::stoi(tile_y_str);
+                std::string tile_w_str        = tile_size_str.substr(0, x_pos);
+                std::string tile_h_str        = tile_size_str.substr(x_pos + 1);
+                vae_tiling_params.tile_size_w = std::stoi(tile_w_str);
+                vae_tiling_params.tile_size_h = std::stoi(tile_h_str);
             } else {
-                vae_tiling_params.tile_size_x = vae_tiling_params.tile_size_y = std::stoi(tile_size_str);
+                vae_tiling_params.tile_size_w = vae_tiling_params.tile_size_h = std::stoi(tile_size_str);
             }
         } catch (const std::invalid_argument&) {
             return -1;
@@ -1627,12 +1638,12 @@ ArgOptions SDGenerationParams::get_options() {
         size_t x_pos             = rel_size_str.find('x');
         try {
             if (x_pos != std::string::npos) {
-                std::string rel_x_str        = rel_size_str.substr(0, x_pos);
-                std::string rel_y_str        = rel_size_str.substr(x_pos + 1);
-                vae_tiling_params.rel_size_x = std::stof(rel_x_str);
-                vae_tiling_params.rel_size_y = std::stof(rel_y_str);
+                std::string rel_w_str        = rel_size_str.substr(0, x_pos);
+                std::string rel_h_str        = rel_size_str.substr(x_pos + 1);
+                vae_tiling_params.rel_size_w = std::stof(rel_w_str);
+                vae_tiling_params.rel_size_h = std::stof(rel_h_str);
             } else {
-                vae_tiling_params.rel_size_x = vae_tiling_params.rel_size_y = std::stof(rel_size_str);
+                vae_tiling_params.rel_size_w = vae_tiling_params.rel_size_h = std::stof(rel_size_str);
             }
         } catch (const std::invalid_argument&) {
             return -1;
@@ -1762,11 +1773,11 @@ ArgOptions SDGenerationParams::get_options() {
          on_scm_policy_arg},
         {"",
          "--vae-tile-size",
-         "tile size for vae tiling in latent units, not image pixels, format [X]x[Y] (default: 32x32)",
+         "tile size for vae encode and decode in image pixels, format [W]x[H] or [S] (default: 256x256; requires --vae-tiling)",
          on_tile_size_arg},
         {"",
          "--vae-relative-tile-size",
-         "relative tile size for vae tiling, format [X]x[Y], in fraction of image size if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)",
+         "relative tile size for vae encode and decode, format [W]x[H] or [S]: <=1 is a dimension fraction, >1 a target tile count (overrides --vae-tile-size; requires --vae-tiling)",
          on_relative_tile_size_arg},
         {"",
          "--prompt-file",
@@ -2223,20 +2234,20 @@ bool SDGenerationParams::from_json_str(
         if (tiling_json.contains("temporal_tiling") && tiling_json["temporal_tiling"].is_boolean()) {
             vae_tiling_params.temporal_tiling = tiling_json["temporal_tiling"];
         }
-        if (tiling_json.contains("tile_size_x") && tiling_json["tile_size_x"].is_number_integer()) {
-            vae_tiling_params.tile_size_x = tiling_json["tile_size_x"];
+        if (tiling_json.contains("tile_size_w") && tiling_json["tile_size_w"].is_number_integer()) {
+            vae_tiling_params.tile_size_w = tiling_json["tile_size_w"];
         }
-        if (tiling_json.contains("tile_size_y") && tiling_json["tile_size_y"].is_number_integer()) {
-            vae_tiling_params.tile_size_y = tiling_json["tile_size_y"];
+        if (tiling_json.contains("tile_size_h") && tiling_json["tile_size_h"].is_number_integer()) {
+            vae_tiling_params.tile_size_h = tiling_json["tile_size_h"];
         }
         if (tiling_json.contains("target_overlap") && tiling_json["target_overlap"].is_number()) {
             vae_tiling_params.target_overlap = tiling_json["target_overlap"];
         }
-        if (tiling_json.contains("rel_size_x") && tiling_json["rel_size_x"].is_number()) {
-            vae_tiling_params.rel_size_x = tiling_json["rel_size_x"];
+        if (tiling_json.contains("rel_size_w") && tiling_json["rel_size_w"].is_number()) {
+            vae_tiling_params.rel_size_w = tiling_json["rel_size_w"];
         }
-        if (tiling_json.contains("rel_size_y") && tiling_json["rel_size_y"].is_number()) {
-            vae_tiling_params.rel_size_y = tiling_json["rel_size_y"];
+        if (tiling_json.contains("rel_size_h") && tiling_json["rel_size_h"].is_number()) {
+            vae_tiling_params.rel_size_h = tiling_json["rel_size_h"];
         }
         if (tiling_json.contains("extra_tiling_args") && tiling_json["extra_tiling_args"].is_string()) {
             extra_tiling_args = tiling_json["extra_tiling_args"].get<std::string>();
@@ -2270,14 +2281,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)) {
@@ -2319,8 +2331,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;
             }
         }
@@ -2336,6 +2349,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 {
@@ -2457,7 +2472,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) {
@@ -2501,7 +2516,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;
 }
@@ -2629,7 +2646,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)) {
@@ -2933,11 +2950,11 @@ std::string SDGenerationParams::to_string() const {
         << "  vae_tiling_params: { "
         << vae_tiling_params.enabled << ", "
         << vae_tiling_params.temporal_tiling << ", "
-        << vae_tiling_params.tile_size_x << ", "
-        << vae_tiling_params.tile_size_y << ", "
+        << vae_tiling_params.tile_size_w << ", "
+        << vae_tiling_params.tile_size_h << ", "
         << vae_tiling_params.target_overlap << ", "
-        << vae_tiling_params.rel_size_x << ", "
-        << vae_tiling_params.rel_size_y << ", "
+        << vae_tiling_params.rel_size_w << ", "
+        << vae_tiling_params.rel_size_h << ", "
         << "\"" << extra_tiling_args << "\" },\n"
         << "}";
     return oss.str();
@@ -3139,11 +3156,11 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
         root["vae_tiling"] = {
             {"enabled", gen_params.vae_tiling_params.enabled},
             {"temporal_tiling", gen_params.vae_tiling_params.temporal_tiling},
-            {"tile_size_x", gen_params.vae_tiling_params.tile_size_x},
-            {"tile_size_y", gen_params.vae_tiling_params.tile_size_y},
+            {"tile_size_w", gen_params.vae_tiling_params.tile_size_w},
+            {"tile_size_h", gen_params.vae_tiling_params.tile_size_h},
             {"target_overlap", gen_params.vae_tiling_params.target_overlap},
-            {"rel_size_x", gen_params.vae_tiling_params.rel_size_x},
-            {"rel_size_y", gen_params.vae_tiling_params.rel_size_y},
+            {"rel_size_w", gen_params.vae_tiling_params.rel_size_w},
+            {"rel_size_h", gen_params.vae_tiling_params.rel_size_h},
             {"extra_tiling_args", gen_params.extra_tiling_args},
         };
     }

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/CMakeLists.txt

@@ -21,21 +21,21 @@ if(SD_SERVER_BUILD_FRONTEND AND EXISTS "${FRONTEND_DIR}")
         set(HAVE_FRONTEND_BUILD ON)
 
         add_custom_target(${TARGET}_frontend_install
-            COMMAND "${PNPM_EXECUTABLE}" -C "${FRONTEND_DIR}" install
+            COMMAND "${PNPM_EXECUTABLE}" --ignore-workspace -C "${FRONTEND_DIR}" install
             WORKING_DIRECTORY "${FRONTEND_DIR}"
             COMMENT "Installing frontend dependencies"
             VERBATIM
         )
 
         add_custom_target(${TARGET}_frontend_build
-            COMMAND "${PNPM_EXECUTABLE}" -C "${FRONTEND_DIR}" run build
+            COMMAND "${PNPM_EXECUTABLE}" --ignore-workspace -C "${FRONTEND_DIR}" run build
             WORKING_DIRECTORY "${FRONTEND_DIR}"
             COMMENT "Building frontend"
             VERBATIM
         )
 
         add_custom_target(${TARGET}_frontend_header
-            COMMAND "${PNPM_EXECUTABLE}" -C "${FRONTEND_DIR}" run build:header
+            COMMAND "${PNPM_EXECUTABLE}" --ignore-workspace -C "${FRONTEND_DIR}" run build:header
             WORKING_DIRECTORY "${FRONTEND_DIR}"
             COMMENT "Generating gen_index_html.h"
             VERBATIM

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`
@@ -524,11 +560,11 @@ Shared default fields used by both `img_gen` and `vid_gen`:
 | `vae_tiling_params` | `object` |
 | `vae_tiling_params.enabled` | `boolean` |
 | `vae_tiling_params.temporal_tiling` | `boolean` |
-| `vae_tiling_params.tile_size_x` | `integer` |
-| `vae_tiling_params.tile_size_y` | `integer` |
+| `vae_tiling_params.tile_size_w` | `integer` |
+| `vae_tiling_params.tile_size_h` | `integer` |
 | `vae_tiling_params.target_overlap` | `number` |
-| `vae_tiling_params.rel_size_x` | `number` |
-| `vae_tiling_params.rel_size_y` | `number` |
+| `vae_tiling_params.rel_size_w` | `number` |
+| `vae_tiling_params.rel_size_h` | `number` |
 | `vae_tiling_params.extra_tiling_args` | `string` |
 | `cache_mode` | `string` |
 | `cache_option` | `string` |
@@ -537,6 +573,8 @@ Shared default fields used by both `img_gen` and `vid_gen`:
 | `output_format` | `string` |
 | `output_compression` | `integer` |
 
+`vae_tiling_params.tile_size_w` and `tile_size_h` are in **image pixels**, with `0` selecting the 256-pixel default. Both encode and decode use these sizes without an encoding multiplier. Positive `rel_size_w`/`rel_size_h` values override the corresponding absolute size: values up to 1 are dimension fractions, and values greater than 1 are target tile counts. Set `enabled` to use spatial tiling. Sizes are aligned down to the VAE scale factor and capped at the input dimensions; explicit sizes below the minimum supported tile size are rejected. These fields previously used latent units; see [VAE tiling](../../docs/performance.md#use-vae-tiling-to-reduce-encode-and-decode-memory-usage) for migration and OOM retry behavior.
+
 `vae_tiling_params.extra_tiling_args` accepts a key=value list. Supported video VAEs accept `temporal_tile_frames` (alias `temporal_tile_size`, default `4`) and `temporal_tile_overlap` (default `1`).
 LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEHV use overlap blending. MiniMax H3 keeps its model-specific fixed temporal windows because its latent-to-frame mapping is non-linear.
 
@@ -596,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`:
 
@@ -608,6 +647,7 @@ Fields returned in `features_by_mode.vid_gen`:
 - `cache`
 - `cancel_queued`
 - `cancel_generating`
+- `preview`
 
 #### `POST /sdcpp/v1/img_gen`
 
@@ -767,11 +807,11 @@ Example:
   "vae_tiling_params": {
     "enabled": false,
     "temporal_tiling": false,
-    "tile_size_x": 0,
-    "tile_size_y": 0,
+    "tile_size_w": 0,
+    "tile_size_h": 0,
     "target_overlap": 0.5,
-    "rel_size_x": 0.0,
-    "rel_size_y": 0.0,
+    "rel_size_w": 0.0,
+    "rel_size_h": 0.0,
     "extra_tiling_args": ""
   },
 
@@ -780,6 +820,9 @@ Example:
   "scm_mask": "",
   "scm_policy_dynamic": true,
 
+  "preview": "none",
+  "preview_interval": 1,
+
   "output_format": "png",
   "output_compression": 100
 }
@@ -845,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:
 
@@ -900,11 +945,11 @@ Other native fields:
 | `vae_tiling_params` | `object` |
 | `vae_tiling_params.enabled` | `boolean` |
 | `vae_tiling_params.temporal_tiling` | `boolean` |
-| `vae_tiling_params.tile_size_x` | `integer` |
-| `vae_tiling_params.tile_size_y` | `integer` |
+| `vae_tiling_params.tile_size_w` | `integer` |
+| `vae_tiling_params.tile_size_h` | `integer` |
 | `vae_tiling_params.target_overlap` | `number` |
-| `vae_tiling_params.rel_size_x` | `number` |
-| `vae_tiling_params.rel_size_y` | `number` |
+| `vae_tiling_params.rel_size_w` | `number` |
+| `vae_tiling_params.rel_size_h` | `number` |
 | `vae_tiling_params.extra_tiling_args` | `string` |
 | `cache_mode` | `string` |
 | `cache_option` | `string` |
@@ -932,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:
@@ -1046,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:
@@ -1067,6 +1117,7 @@ Example:
 
   "init_image": null,
   "end_image": null,
+  "ref_images": [],
   "control_frames": [],
 
   "sample_params": {
@@ -1115,11 +1166,11 @@ Example:
   "vae_tiling_params": {
     "enabled": false,
     "temporal_tiling": false,
-    "tile_size_x": 0,
-    "tile_size_y": 0,
+    "tile_size_w": 0,
+    "tile_size_h": 0,
     "target_overlap": 0.5,
-    "rel_size_x": 0.0,
-    "rel_size_y": 0.0,
+    "rel_size_w": 0.0,
+    "rel_size_h": 0.0,
     "extra_tiling_args": ""
   },
 
@@ -1128,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.
@@ -1150,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:
@@ -1179,6 +1245,8 @@ Top-level scalar fields:
 | `fps` | `integer` |
 | `moe_boundary` | `number` |
 | `vace_strength` | `number` |
+| `preview` | `string` |
+| `preview_interval` | `integer` |
 
 Image and frame fields:
 
@@ -1186,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:
@@ -1240,11 +1309,11 @@ Other native fields:
 | `vae_tiling_params` | `object` |
 | `vae_tiling_params.enabled` | `boolean` |
 | `vae_tiling_params.temporal_tiling` | `boolean` |
-| `vae_tiling_params.tile_size_x` | `integer` |
-| `vae_tiling_params.tile_size_y` | `integer` |
+| `vae_tiling_params.tile_size_w` | `integer` |
+| `vae_tiling_params.tile_size_h` | `integer` |
 | `vae_tiling_params.target_overlap` | `number` |
-| `vae_tiling_params.rel_size_x` | `number` |
-| `vae_tiling_params.rel_size_y` | `number` |
+| `vae_tiling_params.rel_size_w` | `number` |
+| `vae_tiling_params.rel_size_h` | `number` |
 | `vae_tiling_params.extra_tiling_args` | `string` |
 | `cache_mode` | `string` |
 | `cache_option` | `string` |

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

@@ -78,11 +78,11 @@ static json make_vae_tiling_json(const sd_tiling_params_t& params) {
     return {
         {"enabled", params.enabled},
         {"temporal_tiling", params.temporal_tiling},
-        {"tile_size_x", params.tile_size_x},
-        {"tile_size_y", params.tile_size_y},
+        {"tile_size_w", params.tile_size_w},
+        {"tile_size_h", params.tile_size_h},
         {"target_overlap", params.target_overlap},
-        {"rel_size_x", params.rel_size_x},
-        {"rel_size_y", params.rel_size_y},
+        {"rel_size_w", params.rel_size_w},
+        {"rel_size_h", params.rel_size_h},
         {"extra_tiling_args", params.extra_tiling_args ? params.extra_tiling_args : ""},
     };
 }
@@ -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,28 +254,55 @@ 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;
-    if (fs::exists(lora_dir) && fs::is_directory(lora_dir)) {
-        for (auto& entry : fs::recursive_directory_iterator(lora_dir, fs::directory_options::skip_permission_denied)) {
-            if (!entry.is_regular_file()) {
-                continue;
-            }
-            const fs::path& p = entry.path();
-            if (!is_supported_model_ext(p)) {
-                continue;
-            }
+    std::error_code ec;
+    if (fs::exists(lora_dir, ec) && !ec && fs::is_directory(lora_dir, ec) && !ec) {
+        try {
+            auto it = fs::recursive_directory_iterator(
+                lora_dir,
+                fs::directory_options::skip_permission_denied,
+                ec);
+            auto end = fs::recursive_directory_iterator();
+            while (!ec && it != end) {
+                std::error_code entry_ec;
+                bool is_reg = it->is_regular_file(entry_ec);
+                if (entry_ec || !is_reg) {
+                    it.increment(ec);
+                    continue;
+                }
+                const fs::path& p = it->path();
+                if (!is_supported_model_ext(p)) {
+                    it.increment(ec);
+                    continue;
+                }
 
-            LoraEntry lora_entry;
-            lora_entry.name     = p.stem().u8string();
-            lora_entry.fullpath = p.u8string();
-            std::string rel     = p.lexically_relative(lora_dir).u8string();
-            std::replace(rel.begin(), rel.end(), '\\', '/');
-            lora_entry.path = rel;
+                LoraEntry lora_entry;
+                lora_entry.name     = p.stem().u8string();
+                lora_entry.fullpath = p.u8string();
+                std::string rel     = p.lexically_relative(lora_dir).u8string();
+                std::replace(rel.begin(), rel.end(), '\\', '/');
+                lora_entry.path = rel;
 
-            new_cache.push_back(std::move(lora_entry));
+                new_cache.push_back(std::move(lora_entry));
+                it.increment(ec);
+            }
+        } catch (const std::exception& e) {
+            LOG_WARN("error while scanning lora directory '%s': %s", lora_dir.string().c_str(), e.what());
+            return;
         }
     }
 
+    if (ec) {
+        LOG_WARN("error while scanning lora directory '%s': %s", lora_dir.string().c_str(), ec.message().c_str());
+        return;
+    }
+
     std::sort(new_cache.begin(), new_cache.end(), [](const LoraEntry& a, const LoraEntry& b) {
         return a.path < b.path;
     });
@@ -302,40 +329,71 @@ void refresh_upscaler_cache(ServerRuntime& rt) {
     }
 
     fs::path upscaler_dir = rt.ctx_params->hires_upscalers_dir;
-    if (fs::exists(upscaler_dir) && fs::is_directory(upscaler_dir)) {
-        for (auto& entry : fs::directory_iterator(upscaler_dir)) {
-            if (!entry.is_regular_file()) {
-                continue;
-            }
-            const fs::path& p = entry.path();
-            if (!is_supported_model_ext(p)) {
-                continue;
-            }
+    std::error_code ec;
+    if (fs::exists(upscaler_dir, ec) && !ec && fs::is_directory(upscaler_dir, ec) && !ec) {
+        try {
+            auto it = fs::directory_iterator(
+                upscaler_dir,
+                fs::directory_options::skip_permission_denied,
+                ec);
+            auto end = fs::directory_iterator();
+            while (!ec && it != end) {
+                std::error_code entry_ec;
+                bool is_reg = it->is_regular_file(entry_ec);
+                if (entry_ec || !is_reg) {
+                    it.increment(ec);
+                    continue;
+                }
+                const fs::path& p = it->path();
+                if (!is_supported_model_ext(p)) {
+                    it.increment(ec);
+                    continue;
+                }
 
-            UpscalerEntry upscaler_entry;
-            upscaler_entry.name                 = p.stem().u8string();
-            upscaler_entry.fullpath             = fs::absolute(p).lexically_normal().u8string();
-            upscaler_entry.model_name           = "ESRGAN_4x";
-            upscaler_entry.path                 = p.filename().u8string();
-            upscaler_entry.file_size            = entry.file_size();
-            upscaler_entry.last_modified        = entry.last_write_time();
-            auto previous                       = std::find_if(previous_cache.begin(), previous_cache.end(), [&](const UpscalerEntry& cached) {
-                return cached.fullpath == upscaler_entry.fullpath &&
-                       cached.file_size == upscaler_entry.file_size &&
-                       cached.last_modified == upscaler_entry.last_modified;
-            });
-            upscaler_entry.image_upscale_factor = previous != previous_cache.end()
-                                                      ? previous->image_upscale_factor
-                                                      : get_upscaler_model_scale(upscaler_entry.fullpath.c_str());
-            if (upscaler_entry.image_upscale_factor > 0) {
-                upscaler_entry.scale      = upscaler_entry.image_upscale_factor;
-                upscaler_entry.model_name = "ESRGAN_" + std::to_string(upscaler_entry.scale) + "x";
-            }
+                UpscalerEntry upscaler_entry;
+                upscaler_entry.name = p.stem().u8string();
+                std::error_code abs_ec;
+                fs::path abs_path         = fs::absolute(p, abs_ec);
+                upscaler_entry.fullpath   = (abs_ec ? p : abs_path.lexically_normal()).u8string();
+                upscaler_entry.model_name = "ESRGAN_4x";
+                upscaler_entry.path       = p.filename().u8string();
 
-            new_cache.push_back(std::move(upscaler_entry));
+                std::error_code size_ec;
+                upscaler_entry.file_size = it->file_size(size_ec);
+                if (size_ec) {
+                    it.increment(ec);
+                    continue;
+                }
+
+                std::error_code time_ec;
+                upscaler_entry.last_modified        = it->last_write_time(time_ec);
+                auto previous                       = std::find_if(previous_cache.begin(), previous_cache.end(), [&](const UpscalerEntry& cached) {
+                    return cached.fullpath == upscaler_entry.fullpath &&
+                           cached.file_size == upscaler_entry.file_size &&
+                           (!time_ec && cached.last_modified == upscaler_entry.last_modified);
+                });
+                upscaler_entry.image_upscale_factor = previous != previous_cache.end()
+                                                          ? previous->image_upscale_factor
+                                                          : get_upscaler_model_scale(upscaler_entry.fullpath.c_str());
+                if (upscaler_entry.image_upscale_factor > 0) {
+                    upscaler_entry.scale      = upscaler_entry.image_upscale_factor;
+                    upscaler_entry.model_name = "ESRGAN_" + std::to_string(upscaler_entry.scale) + "x";
+                }
+
+                new_cache.push_back(std::move(upscaler_entry));
+                it.increment(ec);
+            }
+        } catch (const std::exception& e) {
+            LOG_WARN("error while scanning upscalers directory '%s': %s", upscaler_dir.string().c_str(), e.what());
+            return;
         }
     }
 
+    if (ec) {
+        LOG_WARN("error while scanning upscalers directory '%s': %s", upscaler_dir.string().c_str(), ec.message().c_str());
+        return;
+    }
+
     std::sort(new_cache.begin(), new_cache.end(), [](const UpscalerEntry& a, const UpscalerEntry& b) {
         return a.name < b.name;
     });

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

@@ -173,11 +173,13 @@ enum lora_apply_mode_t {
 typedef struct {
     bool enabled;
     bool temporal_tiling;
-    int tile_size_x;
-    int tile_size_y;
+    // Spatial tile dimensions in image pixels for both encode and decode; 0 uses 256.
+    int tile_size_w;
+    int tile_size_h;
     float target_overlap;
-    float rel_size_x;
-    float rel_size_y;
+    // Positive values override tile_size: <= 1 is a dimension fraction, > 1 a target tile count.
+    float rel_size_w;
+    float rel_size_h;
     const char* extra_tiling_args;
 } sd_tiling_params_t;
 
@@ -243,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
@@ -459,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);
@@ -467,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();
@@ -599,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

@@ -16,6 +16,7 @@
 #include "model/te/clip.hpp"
 #include "model/te/llada_image_te.hpp"
 #include "model/te/llm.hpp"
+#include "model/te/ming_image_te.hpp"
 #include "model/te/t5.hpp"
 #include "model_loader.h"
 #include "tokenizers/sensenova_u1_tokenizer.h"
@@ -398,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 {
@@ -816,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 {
@@ -908,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;
@@ -1206,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 {
@@ -1483,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 {
@@ -1872,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;
@@ -1986,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,
@@ -2112,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) {
@@ -2143,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;
@@ -2159,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);
 
@@ -2181,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);
@@ -2245,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;
     }
 
@@ -2311,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;
@@ -2952,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
 
@@ -3060,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);
         }
@@ -3075,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 {};
         }
@@ -3140,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) {
@@ -3171,6 +3243,65 @@ struct LLMEmbedder : public Conditioner {
     }
 };
 
+struct MingImageEmbedder : public Conditioner {
+    std::shared_ptr<Tokenizer> tokenizer;
+    std::shared_ptr<MingImageTE::MingImageTextRunner> text_model;
+    const std::string prefix = "text_encoders.llm";
+
+    MingImageEmbedder(ggml_backend_t backend, const String2TensorStorage& tensors, std::shared_ptr<RunnerWeightManager> weight_manager, const TokenizerConfig& tokenizers) {
+        if (!tokenizers.has(TokenizerConfig::MAIN)) {
+            throw std::runtime_error("Ming-Image requires the Ling tokenizer.json; pass --tokenizer FILE");
+        }
+        text_model = std::make_shared<MingImageTE::MingImageTextRunner>(backend, tensors, prefix, weight_manager);
+        tokenizer  = tokenizers.create(TokenizerConfig::MAIN, text_model->config.backbone.vocab_size, 156895);
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        text_model->get_param_tensors(tensors, prefix);
+    }
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& ops) override {
+        text_model->get_param_tensor_ops(ops);
+    }
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        text_model->get_param_tensors(tensors, prefix);
+    }
+    void set_flash_attention_enabled(bool enabled) override { text_model->set_flash_attention_enabled(enabled); }
+    void set_max_graph_vram_bytes(size_t bytes) override { text_model->set_max_graph_vram_bytes(bytes); }
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override { text_model->set_runtime_backends(backends); }
+    void set_graph_cut_layer_split_enabled(bool enabled) override { text_model->set_graph_cut_layer_split_enabled(enabled); }
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override { text_model->set_graph_cut_layer_split_backend_vram_limits(limits); }
+    void set_scale_overrides(float linear, float attention) override { text_model->set_scale_overrides(linear, attention); }
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override { text_model->set_weight_adapter(adapter); }
+    void runner_end() override { text_model->runner_end(); }
+
+    SDCondition get_learned_condition(int n_threads, const ConditionerParams& input) override {
+        if (input.ref_images != nullptr && !input.ref_images->empty()) {
+            LOG_ERROR("Ming-Image currently supports text-to-image only");
+            return {};
+        }
+        std::string prompt =
+            "<role>SYSTEM</role>你是一个友好的AI助手。\n\ndetailed thinking off<|role_end|>"
+            "<role>HUMAN</role>" +
+            input.text + "<|role_end|><role>ASSISTANT</role>";
+        std::vector<int> tokens;
+        if (!tokenizer->encode(prompt, tokens, nullptr)) {
+            return {};
+        }
+        if (tokens.size() + 258 > 32768) {
+            LOG_ERROR("Ming-Image prompt exceeds the text encoder context length");
+            return {};
+        }
+        auto output = text_model->compute(n_threads, tokens);
+        if (output.empty()) {
+            return {};
+        }
+        SDCondition result;
+        result.c_crossattn = sd::ops::slice(sd::ops::slice(output, 1, 0, 256), 0, 0, text_model->config.caption_dim);
+        result.extra_c_crossattns.push_back(sd::ops::slice(output, 1, 256, output.shape()[1]));
+        return result;
+    }
+};
+
 struct LTXAVTextProjection : public GGMLBlock {
     static constexpr int64_t kHiddenSize = 3840;
     static constexpr int64_t kNumStates  = 49;

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

@@ -478,7 +478,12 @@ namespace sd::backend_fit {
         return true;
     }
 
-    bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling, ggml_status status) {
+    bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params,
+                                         bool prefer_temporal_tiling,
+                                         ggml_status status,
+                                         int latent_tile_size_w,
+                                         int latent_tile_size_h,
+                                         int scale_factor) {
         // Execution failures can leave the device unusable; tiling only helps with allocation failures.
         if (status != GGML_STATUS_ALLOC_FAILED) {
             return false;
@@ -487,19 +492,27 @@ namespace sd::backend_fit {
         if (prefer_temporal_tiling && !tiling_params.temporal_tiling) {
             tiling_params.temporal_tiling = true;
             retry_mode                    = tiling_params.enabled ? "spatial+temporal" : "temporal";
-        } else if (!tiling_params.enabled) {
-            tiling_params.enabled    = true;
-            tiling_params.rel_size_x = 0.5f;
-            tiling_params.rel_size_y = 0.5f;
-            if (tiling_params.tile_size_x <= 0) {
-                tiling_params.tile_size_x = 256;
-            }
-            if (tiling_params.tile_size_y <= 0) {
-                tiling_params.tile_size_y = 256;
-            }
-            retry_mode = tiling_params.temporal_tiling ? "spatial+temporal" : "spatial";
         } else {
-            return false;
+            if (latent_tile_size_w <= 0 || latent_tile_size_h <= 0 || scale_factor <= 0) {
+                return false;
+            }
+            auto smaller_tile = [&](int 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);
+            if (tile_size_w == latent_tile_size_w && tile_size_h == latent_tile_size_h) {
+                return false;
+            }
+            tiling_params.enabled     = true;
+            tiling_params.rel_size_w  = 0.0f;
+            tiling_params.rel_size_h  = 0.0f;
+            tiling_params.tile_size_w = tile_size_w * scale_factor;
+            tiling_params.tile_size_h = tile_size_h * scale_factor;
+            retry_mode                = tiling_params.temporal_tiling ? "spatial+temporal" : "spatial";
+            LOG_WARN("Reducing VAE decode tiles from %dx%d to %dx%d image pixels",
+                     latent_tile_size_w * scale_factor, latent_tile_size_h * scale_factor,
+                     tiling_params.tile_size_w, tiling_params.tile_size_h);
         }
 
         LOG_WARN("VAE decode ran out of memory; retrying with %s tiling",

src/core/backend_fit.h

@@ -17,7 +17,10 @@ namespace sd::backend_fit {
 
     bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params,
                                          bool prefer_temporal_tiling,
-                                         ggml_status status);
+                                         ggml_status status,
+                                         int latent_tile_size_w,
+                                         int latent_tile_size_h,
+                                         int scale_factor);
 
 }  // namespace sd::backend_fit
 

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) {
@@ -255,14 +255,20 @@ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
     }
 
     ggml_tensor* fused_bias = scale == 1.f ? b : nullptr;
+    auto mul_mat            = [&](ggml_tensor* input) {
+        if (input->type == GGML_TYPE_F32 && convrot_group_size > 0) {
+            input = ggml_quantize_i8_convrot(ctx, input, convrot_group_size);
+        }
+        return ggml_mul_mat_i8_tensorwise(ctx, w, input, weight_scale, fused_bias, convrot_group_size);
+    };
     if (x->ne[2] * x->ne[3] > 1024) {
         int64_t ne2 = x->ne[2];
         int64_t ne3 = x->ne[3];
         x           = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]);
-        x           = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size);
+        x           = mul_mat(x);
         x           = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3);
     } else {
-        x = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size);
+        x = mul_mat(x);
     }
 
     if (scale != 1.f) {
@@ -465,7 +471,11 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
                               int d1,
                               int d2,
                               bool force_prec_f32,
-                              bool direct) {
+                              bool direct,
+                              float scale) {
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, scale);
+    }
     if (direct) {
         int64_t OC = w->ne[3] / IC;
         int64_t N  = x->ne[3] / IC;
@@ -478,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);
@@ -502,6 +512,9 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
         }
     }
 
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, 1.f / scale);
+    }
     if (b != nullptr) {
         b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]);  // [OC, 1, 1, 1]
         x = ggml_add_inplace(ctx, x, b);
@@ -623,7 +636,11 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
                                     bool skip_reshape,
                                     bool flash_attn,
                                     float kv_scale,
-                                    bool sage_attn) {  // avoid overflow
+                                    bool sage_attn,
+                                    bool* used_flash_attn) {  // avoid overflow
+    if (used_flash_attn != nullptr) {
+        *used_flash_attn = false;
+    }
     int64_t L_q;
     int64_t L_k;
     int64_t C;
@@ -666,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);
@@ -677,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);
@@ -704,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);
         }
@@ -725,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);
@@ -755,6 +772,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
         if (can_use_flash_attn) {
             kqv = build_kqv(q, k, v, mask);
             if (kqv != nullptr) {
+                if (used_flash_attn != nullptr) {
+                    *used_flash_attn = true;
+                }
                 kqv = ggml_view_4d(ctx,
                                    kqv,
                                    d_head,
@@ -777,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/ggml_extend.h

@@ -154,7 +154,8 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
                               int d1              = 1,
                               int d2              = 1,
                               bool force_prec_f32 = false,
-                              bool direct         = false);
+                              bool direct         = false,
+                              float scale         = 1.f);
 
 // w: [OC，IC, KD, 1 * 1]
 // x: [N, IC, ID, IH*IW]
@@ -217,11 +218,12 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
                                     ggml_tensor* k,
                                     ggml_tensor* v,
                                     int64_t n_head,
-                                    ggml_tensor* mask = nullptr,
-                                    bool skip_reshape = false,
-                                    bool flash_attn   = false,
-                                    float kv_scale    = 1.0f,
-                                    bool sage_attn    = false);
+                                    ggml_tensor* mask     = nullptr,
+                                    bool skip_reshape     = false,
+                                    bool flash_attn       = false,
+                                    float kv_scale        = 1.0f,
+                                    bool sage_attn        = false,
+                                    bool* used_flash_attn = nullptr);
 
 ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
                                  ggml_tensor* x,

src/core/ggml_runner.cpp

@@ -21,11 +21,12 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
                                     ggml_tensor* mask,
                                     bool skip_reshape,
                                     bool flash_attn,
-                                    float kv_scale) {
+                                    float kv_scale,
+                                    bool* used_flash_attn) {
     if (ctx->attn_scale > 0.f) {
         kv_scale = ctx->attn_scale;
     }
-    return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled);
+    return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled, used_flash_attn);
 }
 
 void GGMLRunner::alloc_params_ctx() {
@@ -515,9 +516,10 @@ GGMLRunner::~GGMLRunner() {
     free_params_ctx();
 }
 
-GGMLRunnerContext GGMLRunner::get_context() {
+GGMLRunnerContext GGMLRunner::get_context(ggml_cgraph* graph) {
     GGMLRunnerContext runner_ctx;
     runner_ctx.ggml_ctx              = compute_ctx;
+    runner_ctx.graph                 = graph;
     runner_ctx.backend               = runtime_backend;
     runner_ctx.flash_attn_enabled    = flash_attn_enabled;
     runner_ctx.sage_attn_enabled     = sage_attn_enabled;
@@ -532,8 +534,8 @@ GGMLRunnerContext GGMLRunner::get_context() {
     runner_ctx.get_cache_tensor      = [this](const std::string& name) {
         return this->get_cache_tensor_by_name(name);
     };
-    runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) {
-        this->cache(name, tensor);
+    runner_ctx.cache_tensor = [this, graph](const std::string& name, ggml_tensor* tensor) {
+        this->cache(name, tensor, graph);
     };
     runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) {
         this->set_backend_tensor_data(tensor, data);
@@ -575,7 +577,7 @@ ggml_tensor* GGMLRunner::to_backend(ggml_tensor* tensor) {
     }
 }
 
-void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
+void GGMLRunner::cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph) {
     if (tensor != nullptr && tensor->view_src != nullptr) {
         tensor = ggml_cont(compute_ctx, tensor);
     }
@@ -583,6 +585,10 @@ void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
         ggml_set_output(tensor);
     }
     cache_.stage(name, tensor);
+    if (graph != nullptr && tensor != nullptr) {
+        // Schedule the cache output here so its source can be reused before graph end.
+        ggml_build_forward_expand(graph, tensor);
+    }
 }
 
 std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
@@ -831,11 +837,20 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
         last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
         return std::nullopt;
     }
+    auto fits_monolithic = [&]() {
+        // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits.
+        constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL;
+        auto requests                      = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph));
+        for (auto& request : requests) {
+            request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom);
+        }
+        return fits(requests, params);
+    };
     auto manager         = residency_manager.lock();
     const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
                            manager != nullptr && manager->segmented_compute_enabled() &&
                            cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1 &&
-                           !fits(memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)), params);
+                           !fits_monolithic();
     ggml_graph_cut::Plan monolithic_plan;
     if (!segmented) {
         monolithic_plan.segments.emplace_back();

src/core/ggml_runner.h

@@ -67,6 +67,7 @@ struct WeightAdapter {
 struct GGMLRunnerContext {
     ggml_backend_t backend                                           = nullptr;
     ggml_context* ggml_ctx                                           = nullptr;
+    ggml_cgraph* graph                                               = nullptr;
     bool flash_attn_enabled                                          = false;
     bool sage_attn_enabled                                           = false;
     float linear_scale                                               = 0.f;
@@ -102,6 +103,12 @@ struct GGMLRunnerContext {
         return get_cache_tensor(name);
     }
 
+    void expand_graph(ggml_tensor* tensor) const {
+        if (graph != nullptr && tensor != nullptr) {
+            ggml_build_forward_expand(graph, tensor);
+        }
+    }
+
     void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const {
         if (!cache_tensor || tensor == nullptr) {
             return;
@@ -122,10 +129,11 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
                                     ggml_tensor* k,
                                     ggml_tensor* v,
                                     int64_t n_head,
-                                    ggml_tensor* mask = nullptr,
-                                    bool skip_reshape = false,
-                                    bool flash_attn   = false,
-                                    float kv_scale    = 1.f);
+                                    ggml_tensor* mask     = nullptr,
+                                    bool skip_reshape     = false,
+                                    bool flash_attn       = false,
+                                    float kv_scale        = 1.f,
+                                    bool* used_flash_attn = nullptr);
 
 struct GGMLRunner {
 private:
@@ -289,7 +297,8 @@ public:
 
     virtual ~GGMLRunner();
 
-    virtual GGMLRunnerContext get_context();
+    // Binding a graph schedules cache outputs at registration instead of graph end.
+    virtual GGMLRunnerContext get_context(ggml_cgraph* graph = nullptr);
 
     void reset_compute_ctx();
 
@@ -324,7 +333,7 @@ public:
 
     ggml_tensor* to_backend(ggml_tensor* tensor);
 
-    void cache(const std::string name, ggml_tensor* tensor);
+    void cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph = nullptr);
 
     ggml_tensor* get_cache_tensor_by_name(const std::string& name) {
         return cache_.get(name);

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

@@ -62,6 +62,10 @@ enum SDVersion {
     VERSION_SENSENOVA_U1_5,
     VERSION_LLADA_IMAGE,
     VERSION_ESRGAN,
+    VERSION_PIXART,
+    VERSION_MING_IMAGE,
+    VERSION_Z_IMAGE_L2P,
+    VERSION_IRIS,
     VERSION_COUNT,
 };
 
@@ -174,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;
@@ -252,6 +260,10 @@ static inline bool sd_version_is_sensenova_u1(SDVersion version) {
     return version == VERSION_SENSENOVA_U1_5;
 }
 
+static inline bool sd_version_is_pixart(SDVersion version) {
+    return version == VERSION_PIXART;
+}
+
 static inline bool sd_version_supports_video_generation(SDVersion version) {
     return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version);
 }
@@ -275,7 +287,7 @@ static inline bool sd_version_uses_flux2_vae(SDVersion version) {
 }
 
 static inline bool sd_version_uses_wan_vae(SDVersion version) {
-    if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
+    if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version) || version == VERSION_MING_IMAGE) {
         return true;
     }
     return false;
@@ -309,6 +321,8 @@ 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) ||
         sd_version_is_ernie_image(version) ||
@@ -320,7 +334,9 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_sefi_image(version) ||
         sd_version_is_krea2(version) ||
         sd_version_is_mage_flow(version) ||
-        sd_version_is_sensenova_u1(version)) {
+        sd_version_is_sensenova_u1(version) ||
+        sd_version_is_pixart(version) ||
+        version == VERSION_IRIS) {
         return true;
     }
     return false;

src/model/adapter/lora.hpp

@@ -163,6 +163,27 @@ struct LoraModel : public GGMLRunner {
 
             lora_tensors = std::move(new_lora_tensors);
         }
+
+        std::unordered_map<std::string, ggml_tensor*> new_lora_tensors;
+        for (const auto& [old_name, tensor] : lora_tensors) {
+            std::string new_name = old_name;
+            if (starts_with(old_name, "lora.model.diffusion_model.transformer_blocks.")) {
+                // Qwen Image 2.1 stores the gate before the projection in fused MLP weights.
+                for (const auto& suffix : {std::string(".img_mlp.gate_layer.weight."), std::string(".img_mlp.proj.weight.")}) {
+                    size_t pos = old_name.find(suffix);
+                    if (pos == std::string::npos) {
+                        continue;
+                    }
+                    std::string fused_name = old_name.substr(5, pos - 5) + ".img_mlp.gate_up.weight";
+                    if (model_tensor_names.find(fused_name) != model_tensor_names.end()) {
+                        new_name = "lora." + fused_name + (suffix == ".img_mlp.proj.weight." ? ".1." : ".") + old_name.substr(pos + suffix.size());
+                    }
+                    break;
+                }
+            }
+            new_lora_tensors[new_name] = tensor;
+        }
+        lora_tensors = std::move(new_lora_tensors);
     }
 
     ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {

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

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_COMMON_GGML_BLOCK_HPP__
 #define __SD_MODEL_COMMON_GGML_BLOCK_HPP__
 
+#include <cmath>
 #include <cstdint>
 #include <map>
 #include <memory>
@@ -204,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;
@@ -214,7 +220,7 @@ public:
             if (ctx->backend != nullptr) {
                 ggml_tensor* fp8_matmul = ggml_mul_mat(ctx->ggml_ctx, w, x);
                 if (force_prec_f32) {
-                    ggml_mul_mat_set_prec(fp8_matmul, GGML_PREC_F32);
+                    ggml_prec_set_acc(fp8_matmul, GGML_PREC_F32);
                 }
                 supports_fp8_matmul = ggml_backend_supports_op(ctx->backend, fp8_matmul);
             }
@@ -335,15 +341,75 @@ class Embedding : public UnaryBlock {
 protected:
     int64_t embedding_dim;
     int64_t num_embeddings;
+    bool is_int8_tensorwise     = false;
+    int int8_convrot_group_size = 0;
+
+    static void get_rows_i8(ggml_tensor* dst, int ith, int nth, void* userdata) {
+        const auto* embedding     = static_cast<const Embedding*>(userdata);
+        const int group_size      = embedding->int8_convrot_group_size;
+        const auto* weight        = dst->src[0];
+        const auto* input_ids     = static_cast<const int32_t*>(dst->src[1]->data);
+        const auto* scales        = static_cast<const float*>(dst->src[2]->data);
+        const bool scalar_scale   = ggml_nelements(dst->src[2]) == 1;
+        const float normalization = group_size > 0 ? 1.f / std::sqrt(static_cast<float>(group_size)) : 1.f;
+        for (int64_t row = ith; row < dst->ne[1]; row += nth) {
+            const int32_t token = input_ids[row];
+            GGML_ASSERT(token >= 0 && token < weight->ne[1]);
+            const auto* src   = reinterpret_cast<const int8_t*>(static_cast<const char*>(weight->data) + token * weight->nb[1]);
+            auto* out         = reinterpret_cast<float*>(static_cast<char*>(dst->data) + row * dst->nb[1]);
+            const float scale = scales[scalar_scale ? 0 : token] * normalization;
+            for (int64_t i = 0; i < weight->ne[0]; ++i) {
+                out[i] = static_cast<float>(src[i]) * scale;
+            }
+            // The regular Hadamard rotation is symmetric and its own inverse.
+            for (int stride = 1; stride < group_size; stride *= 4) {
+                for (int64_t base = 0; base < weight->ne[0]; base += 4 * stride) {
+                    for (int j = 0; j < stride; ++j) {
+                        float* values      = out + base + j;
+                        const float a      = values[0];
+                        const float b      = values[stride];
+                        const float c      = values[2 * stride];
+                        const float d      = values[3 * stride];
+                        values[0]          = a + b + c - d;
+                        values[stride]     = a + b - c + d;
+                        values[2 * stride] = a - b + c + d;
+                        values[3 * stride] = -a + b + c + d;
+                    }
+                }
+            }
+        }
+    }
+
     void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
-        enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
-        if (!support_get_rows(wtype)) {
+        auto weight_storage     = tensor_storage_map.find(prefix + "weight");
+        is_int8_tensorwise      = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise;
+        int8_convrot_group_size = 0;
+        enum ggml_type wtype    = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
+        if (is_int8_tensorwise) {
+            GGML_ASSERT(wtype == GGML_TYPE_I8);
+            auto scale_storage = tensor_storage_map.find(prefix + "weight_scale");
+            GGML_ASSERT(scale_storage != tensor_storage_map.end());
+            const int64_t scale_nelements = scale_storage->second.nelements();
+            GGML_ASSERT(scale_nelements == 1 || scale_nelements == num_embeddings);
+            params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements);
+            if (weight_storage->second.int8_convrot) {
+                int8_convrot_group_size = weight_storage->second.int8_convrot_group_size;
+                int remainder           = int8_convrot_group_size;
+                while (remainder > 1 && remainder % 4 == 0) {
+                    remainder /= 4;
+                }
+                GGML_ASSERT(remainder == 1 && embedding_dim % int8_convrot_group_size == 0);
+            }
+        } else if (!support_get_rows(wtype)) {
             wtype = GGML_TYPE_F32;
         }
         params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
     }
 
     enum ggml_op param_usage_op(const std::string& name) const override {
+        if (is_int8_tensorwise) {
+            return GGML_OP_CUSTOM;
+        }
         return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
     }
 
@@ -363,9 +429,17 @@ public:
         int64_t n = input_ids->ne[1];
         input_ids = ggml_reshape_1d(ctx->ggml_ctx, input_ids, input_ids->ne[0] * input_ids->ne[1]);
 
-        input_ids      = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]);
-        auto embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids);
-        embedding      = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n);
+        ggml_tensor* embedding;
+        if (is_int8_tensorwise) {
+            GGML_ASSERT(input_ids->type == GGML_TYPE_I32);
+            ggml_tensor* args[] = {weight, input_ids, params["weight_scale"]};
+            embedding           = ggml_custom_4d(ctx->ggml_ctx, GGML_TYPE_F32, embedding_dim, input_ids->ne[0], 1, 1,
+                                                 args, 3, get_rows_i8, GGML_N_TASKS_MAX, this);
+        } else {
+            input_ids = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]);
+            embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids);
+        }
+        embedding = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n);
 
         // [N, n_token, embedding_dim]
         return embedding;

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];
@@ -1082,8 +1094,16 @@ namespace Rope {
             x = ggml_reshape_4d(ctx, x, 2, d_head / 2, L, n_head * N);  // [N * n_head, L, d_head/2, 2]
             x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 0, 1, 2));       // [2, N * n_head, L, d_head/2]
         } else {
-            x = ggml_reshape_4d(ctx, x, d_head / 2, 2, L, n_head * N);       // [N * n_head, L, 2, d_head/2]
-            x = ggml_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 3, 1));  // [2, N * n_head, L, d_head/2]
+            auto x_0 = ggml_view_3d(ctx, x, d_head / 2, L, n_head * N, x->nb[1], x->nb[2], 0);
+            auto x_1 = ggml_view_3d(ctx, x, d_head / 2, L, n_head * N, x->nb[1], x->nb[2], (d_head / 2) * x->nb[0]);
+
+            auto pe_2 = ggml_cont(ctx, ggml_permute(ctx, pe, 2, 3, 0, 1));
+            auto cos  = ggml_view_2d(ctx, pe_2, d_head / 2, L, pe_2->nb[1], 0);
+            auto sin  = ggml_view_2d(ctx, pe_2, d_head / 2, L, pe_2->nb[1], pe_2->nb[3]);
+
+            auto part_0 = ggml_sub(ctx, ggml_mul(ctx, x_0, cos), ggml_mul(ctx, x_1, sin));
+            auto part_1 = ggml_add(ctx, ggml_mul(ctx, x_0, sin), ggml_mul(ctx, x_1, cos));
+            return ggml_concat(ctx, part_0, part_1, 0);
         }
 
         int64_t offset = x->nb[2] * x->ne[2];
@@ -1101,10 +1121,7 @@ namespace Rope {
         auto pe_1 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 1);  // [L, d_head/2, 2]
 
         auto x_out = ggml_add_inplace(ctx, ggml_mul(ctx, x_0, pe_0), ggml_mul(ctx, x_1, pe_1));  // [N * n_head, L, d_head/2, 2]
-        if (!rope_interleaved) {
-            x_out = ggml_cont(ctx, ggml_permute(ctx, x_out, 1, 0, 2, 3));  // [N * n_head, L, x, d_head/2]
-        }
-        x_out = ggml_reshape_3d(ctx, x_out, d_head, L, n_head * N);  // [N*n_head, L, d_head]
+        x_out      = ggml_reshape_3d(ctx, x_out, d_head, L, n_head * N);                         // [N*n_head, L, d_head]
         return x_out;
     }
 
@@ -1121,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/control.hpp

@@ -376,7 +376,7 @@ struct ControlNet : public GGMLRunner {
             hint = make_input(hint_tensor);
         }
 
-        auto runner_ctx = get_context();
+        auto runner_ctx = get_context(gf);
 
         auto outs = control_net.forward(&runner_ctx,
                                         x,
@@ -389,8 +389,7 @@ struct ControlNet : public GGMLRunner {
         if (guided_hint_input == nullptr && !outs.empty()) {
             guided_hint_output_ggml = outs[0];
             ggml_set_output(guided_hint_output_ggml);
-            cache(guided_hint_cache_name(), guided_hint_output_ggml);
-            ggml_build_forward_expand(gf, guided_hint_output_ggml);
+            runner_ctx.persist_cache_tensor(guided_hint_cache_name(), guided_hint_output_ggml);
         }
 
         control_outputs_ggml.reserve(outs.size() > 0 ? outs.size() - 1 : 0);

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

@@ -0,0 +1,135 @@
+#ifndef __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__
+#define __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__
+
+#include "z_image.hpp"
+
+namespace MingImage {
+    struct MingImageConfig : ZImage::ZImageConfig {
+        bool split_qkv = true;
+
+        static MingImageConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) {
+            MingImageConfig config;
+            static_cast<ZImage::ZImageConfig&>(config) = ZImage::ZImageConfig::detect_from_weights(tensors, prefix);
+            config.split_qkv                           = tensors.count(prefix + ".layers.0.attention.qkv.weight") == 0;
+            return config;
+        }
+    };
+
+    class MingImageModel : public GGMLBlock {
+        MingImageConfig config;
+
+    public:
+        explicit MingImageModel(const MingImageConfig& 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>(1024, 256, std::min<int64_t>(config.hidden_size, 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, config.split_qkv, 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["final_layer"] = std::make_shared<ZImage::FinalLayer>(config.hidden_size, config.patch_size, config.out_channels);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, ggml_tensor* direct, ggml_tensor* pe) {
+            auto gctx           = ctx->ggml_ctx;
+            const int64_t width = x->ne[0], height = x->ne[1];
+            auto img            = DiT::pad_and_patchify(ctx, 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);
+            txt                 = ggml_concat(gctx, txt, direct, 1);
+            auto t              = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, timestep);
+            const int64_t n_txt = txt->ne[1], n_img = img->ne[1];
+            auto txt_pe = ggml_ext_slice(gctx, pe, 3, 0, n_txt);
+            auto img_pe = ggml_ext_slice(gctx, pe, 3, n_txt, n_txt + n_img);
+            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, "ming_image.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, "ming_image.noise_refiner." + std::to_string(i), "img");
+            }
+            auto combined = ggml_concat(gctx, txt, img, 1);
+            for (int64_t i = 0; i < config.num_layers; ++i) {
+                combined = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["layers." + std::to_string(i)])->forward(ctx, combined, pe, nullptr, t);
+                sd::ggml_graph_cut::mark_graph_cut(combined, "ming_image.layers." + std::to_string(i), "combined");
+            }
+            img = ggml_ext_slice(gctx, combined, 1, n_txt, n_txt + n_img);
+            img = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"])->forward(ctx, img, t);
+            img = DiT::unpatchify_and_crop(gctx, img, height, width, config.patch_size, config.patch_size, false);
+            return ggml_scale(gctx, img, -1.f);
+        }
+    };
+
+    struct MingImageRunner : DiffusionModelRunner {
+        MingImageConfig config;
+        MingImageModel model;
+        std::vector<float> pe_values;
+
+        MingImageRunner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(MingImageConfig::detect_from_weights(tensors, prefix)),
+              model(config) {
+            model.init(params_ctx, tensors, prefix);
+        }
+
+        std::string get_desc() override { return "ming_image"; }
+
+        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 {
+            const auto* extra = diffusion_extra_as<MingImageDiffusionExtra>(inputs);
+            if (inputs.ref_latents != nullptr && !inputs.ref_latents->empty()) {
+                LOG_ERROR("Ming-Image reference-image conditioning is not supported");
+                return {};
+            }
+            if (inputs.context == nullptr || extra->direct_context == nullptr) {
+                LOG_ERROR("Ming-Image requires both query and direct text conditions");
+                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);
+                auto direct  = make_input(*extra->direct_context);
+                GGML_ASSERT(x->ne[3] == 1);
+                const int64_t n_txt = context->ne[1] + direct->ne[1];
+                const int64_t n_img = ((x->ne[0] + config.patch_size - 1) / config.patch_size) *
+                                      ((x->ne[1] + config.patch_size - 1) / config.patch_size);
+                auto padded = finish_rope_pe(Rope::gen_z_image_pe(
+                    static_cast<int>(x->ne[1]), static_cast<int>(x->ne[0]), config.patch_size, 1,
+                    static_cast<int>(n_txt), ZImage::SEQ_MULTI_OF, {}, Rope::RefIndexMode::FIXED,
+                    config.theta, config.axes_dim));
+                // Zero-masked alignment tokens cannot affect valid queries. Omit them while
+                // retaining the padded caption length used to position image tokens.
+                const size_t stride      = config.axes_dim_sum * 2;
+                const int64_t padded_txt = n_txt + Rope::bound_mod(static_cast<int>(n_txt), ZImage::SEQ_MULTI_OF);
+                pe_values.assign(padded.begin(), padded.begin() + n_txt * stride);
+                pe_values.insert(pe_values.end(), padded.begin() + padded_txt * stride,
+                                 padded.begin() + (padded_txt + n_img) * stride);
+                auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, n_txt + n_img);
+                set_backend_tensor_data(pe, pe_values.data());
+                auto ctx = get_context();
+                auto out = model.forward(&ctx, x, t, context, direct, pe);
+                ggml_build_forward_expand(gf, out);
+                return gf;
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(graph, n_threads, false), inputs.x->dim());
+        }
+    };
+}
+
+#endif  // __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__

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

@@ -141,6 +141,10 @@ struct LLaDAImageDiffusionExtra {
     const sd::Tensor<float>* semantic = nullptr;
 };
 
+struct MingImageDiffusionExtra {
+    const sd::Tensor<float>* direct_context = nullptr;
+};
+
 using DiffusionExtraParams = std::variant<std::monostate,
                                           UNetDiffusionExtra,
                                           SkipLayerDiffusionExtra,
@@ -154,7 +158,8 @@ using DiffusionExtraParams = std::variant<std::monostate,
                                           MiniT2IDiffusionExtra,
                                           SenseNovaU1DiffusionExtra,
                                           HunyuanVideoDiffusionExtra,
-                                          LLaDAImageDiffusionExtra>;
+                                          LLaDAImageDiffusionExtra,
+                                          MingImageDiffusionExtra>;
 
 struct DiffusionParams {
     const sd::Tensor<float>* x                        = nullptr;

src/model/diffusion/pixart.hpp

@@ -0,0 +1,399 @@
+#ifndef __SD_MODEL_DIFFUSION_PIXART_HPP__
+#define __SD_MODEL_DIFFUSION_PIXART_HPP__
+
+#include <algorithm>
+#include <cinttypes>
+#include <cmath>
+#include <vector>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/mmdit.hpp"
+#include "model/diffusion/model.hpp"
+#include "model_loader.h"
+
+// Ref: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/pixart_transformer_2d.py
+// Ref: https://github.com/PixArt-alpha/PixArt-sigma
+
+namespace PixArt {
+    constexpr int PIXART_GRAPH_SIZE = 20480;
+    constexpr int ADALN_EMBED_DIM   = 256;
+
+    struct PixArtConfig {
+        int64_t in_channels         = 4;
+        int64_t out_channels        = 8;  // learn_sigma: noise prediction + learned variance
+        int64_t hidden_size         = 1152;
+        int64_t cross_attention_dim = 1152;
+        int64_t caption_channels    = 4096;
+        int64_t num_heads           = 16;
+        int64_t patch_size          = 2;
+        int64_t ffn_dim             = 4608;
+        int64_t pos_embed_base_size = 64;
+        float interpolation_scale   = 2.f;
+        int num_layers              = 28;
+
+        static PixArtConfig detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) {
+            PixArtConfig config;
+            auto find = [&](const std::string& suffix) -> const TensorStorage* {
+                auto it = weights.find(prefix + "." + suffix);
+                return it == weights.end() ? nullptr : &it->second;
+            };
+            if (auto w = find("x_embedder.proj.weight")) {
+                config.hidden_size = w->ne[3];
+                config.in_channels = w->ne[2];
+                config.patch_size  = w->ne[0];
+            }
+            if (auto w = find("final_layer.linear.weight")) {
+                config.out_channels = w->ne[1] / (config.patch_size * config.patch_size);
+            }
+            if (auto w = find("y_embedder.y_proj.fc1.weight")) {
+                config.caption_channels = w->ne[0];
+            }
+            if (auto w = find("blocks.0.cross_attn.kv_linear.weight")) {
+                config.cross_attention_dim = w->ne[0];
+            }
+            if (auto w = find("blocks.0.mlp.fc1.weight")) {
+                config.ffn_dim = w->ne[1];
+            }
+            if (find("csize_embedder.mlp.0.weight") != nullptr) {
+                LOG_WARN("pixart: resolution/aspect-ratio micro conditions are not supported; output may differ from the reference");
+            }
+            int layers                     = 0;
+            const std::string block_prefix = prefix + ".blocks.";
+            for (const auto& [name, _] : weights) {
+                if (starts_with(name, block_prefix)) {
+                    layers = std::max(layers, atoi(name.substr(block_prefix.size()).c_str()) + 1);
+                }
+            }
+            if (layers > 0) {
+                config.num_layers = layers;
+                LOG_VERBOSE("pixart: layers = %d, hidden_size = %" PRId64,
+                            layers, config.hidden_size);
+            }
+            return config;
+        }
+    };
+
+    // Mirrors diffusers get_2d_sincos_pos_embed for a (gh, gw) patch grid.
+    static std::vector<float> gen_2d_sincos_pos_embed(int64_t dim,
+                                                      int64_t gh,
+                                                      int64_t gw,
+                                                      int64_t base_size,
+                                                      float interpolation_scale) {
+        // diffusers: meshgrid(grid_w, grid_h, indexing="xy") -> grid[0]=w, grid[1]=h,
+        // embedding = concat(sincos(w), sincos(h))
+        std::vector<float> out(static_cast<size_t>(gh) * gw * dim);
+        int64_t quarter = dim / 4;
+        for (int64_t h = 0; h < gh; ++h) {
+            float pos_h = static_cast<float>(h) / (static_cast<float>(gh) / base_size) / interpolation_scale;
+            for (int64_t w = 0; w < gw; ++w) {
+                float pos_w  = static_cast<float>(w) / (static_cast<float>(gw) / base_size) / interpolation_scale;
+                float* dst_w = out.data() + (h * gw + w) * dim;
+                float* dst_h = dst_w + dim / 2;
+                for (int64_t i = 0; i < quarter; ++i) {
+                    float omega        = 1.f / powf(10000.f, static_cast<float>(i) / quarter);
+                    dst_w[i]           = sinf(pos_w * omega);
+                    dst_w[i + quarter] = cosf(pos_w * omega);
+                    dst_h[i]           = sinf(pos_h * omega);
+                    dst_h[i + quarter] = cosf(pos_h * omega);
+                }
+            }
+        }
+        return out;
+    }
+
+    class PixArtTimestepEmbedding : public GGMLBlock {
+    public:
+        PixArtTimestepEmbedding(int64_t in_channels, int64_t out_dim) {
+            blocks["mlp.0"] = std::make_shared<Linear>(in_channels, out_dim);
+            blocks["mlp.2"] = std::make_shared<Linear>(out_dim, out_dim);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            x = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"])->forward(ctx, x);
+            x = ggml_silu(ctx->ggml_ctx, x);
+            return std::dynamic_pointer_cast<Linear>(blocks["mlp.2"])->forward(ctx, x);
+        }
+    };
+
+    class PixArtAttention : public GGMLBlock {
+        int64_t num_heads;
+        bool self_attention;
+
+    public:
+        PixArtAttention(int64_t dim, int64_t num_heads, int64_t context_dim, bool self_attention)
+            : num_heads(num_heads), self_attention(self_attention) {
+            if (self_attention) {
+                blocks["qkv"] = std::make_shared<Linear>(dim, 3 * dim);
+            } else {
+                blocks["q_linear"]  = std::make_shared<Linear>(dim, dim);
+                blocks["kv_linear"] = std::make_shared<Linear>(context_dim, 2 * dim);
+            }
+            blocks["proj"] = std::make_shared<Linear>(dim, dim);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* context, ggml_tensor* mask = nullptr) {
+            std::vector<ggml_tensor*> qkv;
+            if (self_attention) {
+                auto projected = std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x);
+                qkv            = ggml_ext_chunk(ctx->ggml_ctx, projected, 3, 0);
+            } else {
+                auto q     = std::dynamic_pointer_cast<Linear>(blocks["q_linear"])->forward(ctx, x);
+                auto kv    = std::dynamic_pointer_cast<Linear>(blocks["kv_linear"])->forward(ctx, context);
+                auto parts = ggml_ext_chunk(ctx->ggml_ctx, kv, 2, 0);
+                qkv        = {q, parts[0], parts[1]};
+            }
+            auto out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, mask, false, ctx->flash_attn_enabled);
+            return std::dynamic_pointer_cast<Linear>(blocks["proj"])->forward(ctx, out);
+        }
+    };
+
+    class PixArtBlock : public GGMLBlock {
+        int64_t dim;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            ggml_type wtype             = get_type(prefix + "scale_shift_table", tensor_storage_map, GGML_TYPE_F32);
+            params["scale_shift_table"] = ggml_new_tensor_2d(ctx, wtype, dim, 6);
+        }
+
+    public:
+        PixArtBlock(int64_t dim, int64_t num_heads, int64_t context_dim, int64_t ffn_dim)
+            : dim(dim) {
+            blocks["attn"]       = std::make_shared<PixArtAttention>(dim, num_heads, dim, true);
+            blocks["cross_attn"] = std::make_shared<PixArtAttention>(dim, num_heads, context_dim, false);
+            blocks["mlp.fc1"]    = std::make_shared<Linear>(dim, ffn_dim);
+            blocks["mlp.fc2"]    = std::make_shared<Linear>(ffn_dim, dim);
+        }
+
+        static ggml_tensor* norm(ggml_context* ctx, ggml_tensor* x) {
+            return ggml_norm(ctx, x, 1e-6f);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* mod, ggml_tensor* context, ggml_tensor* context_mask) {
+            // x: [N, n_token, dim]
+            // mod: [N, 6 * dim], shared adaLN-single output
+            int64_t N = x->ne[2];
+
+            auto table = params["scale_shift_table"];
+            if (table->type != GGML_TYPE_F32) {
+                table = ggml_cast(ctx->ggml_ctx, table, GGML_TYPE_F32);
+            }
+            table   = ggml_reshape_3d(ctx->ggml_ctx, table, dim, 6, 1);
+            auto m  = ggml_add(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, mod, dim, 6, N), table);
+            auto mv = ggml_ext_chunk(ctx->ggml_ctx, ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, m), dim * 6, N), 6, 0);
+
+            auto attn1 = std::dynamic_pointer_cast<PixArtAttention>(blocks["attn"]);
+            auto attn2 = std::dynamic_pointer_cast<PixArtAttention>(blocks["cross_attn"]);
+            auto proj  = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
+            auto fc2   = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"]);
+
+            auto gate = [&](ggml_tensor* y, ggml_tensor* g) {
+                g = ggml_reshape_3d(ctx->ggml_ctx, g, dim, 1, N);
+                return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, y, g));
+            };
+
+            auto h = modulate(ctx->ggml_ctx, norm(ctx->ggml_ctx, x), mv[0], mv[1]);
+            x      = gate(attn1->forward(ctx, h, h), mv[2]);
+            // ada_norm_single: no norm before cross-attention (PixArtMS.py)
+            x = ggml_add(ctx->ggml_ctx, x, attn2->forward(ctx, x, context, context_mask));
+            h = modulate(ctx->ggml_ctx, norm(ctx->ggml_ctx, x), mv[3], mv[4]);
+            h = proj->forward(ctx, h);
+            h = ggml_ext_gelu(ctx->ggml_ctx, h, true);
+            h = fc2->forward(ctx, h);
+            return gate(h, mv[5]);
+        }
+    };
+
+    class PixArtModel : public GGMLBlock {
+        PixArtConfig config;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            ggml_type wtype                         = get_type(prefix + "final_layer.scale_shift_table", tensor_storage_map, GGML_TYPE_F32);
+            params["final_layer.scale_shift_table"] = ggml_new_tensor_2d(ctx, wtype, config.hidden_size, 2);
+        }
+
+    public:
+        PixArtModel() = default;
+        PixArtModel(const PixArtConfig& config)
+            : config(config) {
+            blocks["x_embedder.proj"]       = std::make_shared<Conv2d>(config.in_channels,
+                                                                 config.hidden_size,
+                                                                 std::pair<int, int>{static_cast<int>(config.patch_size), static_cast<int>(config.patch_size)},
+                                                                 std::pair<int, int>{static_cast<int>(config.patch_size), static_cast<int>(config.patch_size)});
+            blocks["t_embedder"]            = std::make_shared<PixArtTimestepEmbedding>(ADALN_EMBED_DIM, config.hidden_size);
+            blocks["t_block.1"]             = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size);
+            blocks["y_embedder.y_proj.fc1"] = std::make_shared<Linear>(config.caption_channels, config.hidden_size);
+            blocks["y_embedder.y_proj.fc2"] = std::make_shared<Linear>(config.hidden_size, config.cross_attention_dim);
+            for (int i = 0; i < config.num_layers; ++i) {
+                blocks["blocks." + std::to_string(i)] =
+                    std::make_shared<PixArtBlock>(config.hidden_size, config.num_heads, config.cross_attention_dim, config.ffn_dim);
+            }
+            blocks["final_layer.norm_final"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
+            blocks["final_layer.linear"]     = std::make_shared<Linear>(config.hidden_size,
+                                                                    config.patch_size * config.patch_size * config.out_channels);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timesteps,
+                             ggml_tensor* context,
+                             ggml_tensor* pos_embed,
+                             ggml_tensor* context_mask) {
+            // x: [N, C, H, W] latent, context: [N, n_ctx, caption_channels]
+            int64_t W  = x->ne[0];
+            int64_t H  = x->ne[1];
+            int64_t N  = x->ne[3];
+            int64_t p  = config.patch_size;
+            int64_t wp = W / p;
+            int64_t hp = H / p;
+
+            auto h = std::dynamic_pointer_cast<Conv2d>(blocks["x_embedder.proj"])->forward(ctx, x);  // [N, hidden, hp, wp]
+            h      = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 1, 2, 0, 3));       // [N, hp, wp, hidden] -> [N, hp*wp, hidden]
+            h      = ggml_reshape_3d(ctx->ggml_ctx, h, config.hidden_size, wp * hp, N);              // [N, hp*wp, hidden]
+            h      = ggml_add(ctx->ggml_ctx, h, pos_embed);
+
+            auto t   = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, ADALN_EMBED_DIM, 10000);
+            auto emb = std::dynamic_pointer_cast<PixArtTimestepEmbedding>(blocks["t_embedder"])->forward(ctx, t);
+
+            auto mod = std::dynamic_pointer_cast<Linear>(blocks["t_block.1"])
+                           ->forward(ctx, ggml_silu(ctx->ggml_ctx, emb));  // [N, 6 * hidden]
+
+            auto ctx_emb = std::dynamic_pointer_cast<Linear>(blocks["y_embedder.y_proj.fc1"])->forward(ctx, context);
+            ctx_emb      = ggml_ext_gelu(ctx->ggml_ctx, ctx_emb, true);
+            ctx_emb      = std::dynamic_pointer_cast<Linear>(blocks["y_embedder.y_proj.fc2"])->forward(ctx, ctx_emb);
+
+            for (int i = 0; i < config.num_layers; ++i) {
+                auto block = std::dynamic_pointer_cast<PixArtBlock>(blocks["blocks." + std::to_string(i)]);
+                h          = block->forward(ctx, h, mod, ctx_emb, context_mask);
+                sd::ggml_graph_cut::mark_graph_cut(h, "pixart.blocks." + std::to_string(i), "h");
+            }
+
+            // scale_shift_table + emb -> (shift, scale) for the affine-free final norm
+            auto tail_table = params["final_layer.scale_shift_table"];
+            if (tail_table->type != GGML_TYPE_F32) {
+                tail_table = ggml_cast(ctx->ggml_ctx, tail_table, GGML_TYPE_F32);
+            }
+            auto ss    = ggml_add(ctx->ggml_ctx,
+                                  ggml_reshape_3d(ctx->ggml_ctx, tail_table, config.hidden_size, 2, 1),
+                                  ggml_reshape_3d(ctx->ggml_ctx, emb, config.hidden_size, 1, N));  // [2, hidden, N]
+            auto parts = ggml_ext_chunk(ctx->ggml_ctx,
+                                        ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, ss), config.hidden_size * 2, N),
+                                        2, 0);
+            h          = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer.norm_final"])->forward(ctx, h);
+            h          = modulate(ctx->ggml_ctx, h, parts[0], parts[1]);
+            h          = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, h);  // [N, hp*wp, p*p*out_ch]
+            h          = DiT::unpatchify(ctx->ggml_ctx, h, hp, wp, static_cast<int>(p), static_cast<int>(p), false);
+            return h;  // [N, out_channels, H, W]
+        }
+    };
+
+    struct PixArtRunner : public DiffusionModelRunner {
+        PixArtConfig config;
+        PixArtModel model;
+        std::vector<float> pos_vec;
+
+        PixArtRunner(ggml_backend_t backend,
+                     const String2TensorStorage& tensor_storage_map      = {},
+                     const std::string prefix                            = "",
+                     std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                     const char* model_args                              = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(PixArtConfig::detect_from_weights(tensor_storage_map, prefix)) {
+            for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
+                if (key == "pixart_pos_embed_base_size") {
+                    int parsed = 0;
+                    if (parse_strict_int(value, parsed)) {
+                        config.pos_embed_base_size = parsed;
+                    } else {
+                        LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                } else if (key == "pixart_interpolation_scale") {
+                    float parsed = 0.f;
+                    if (parse_strict_float(value, parsed)) {
+                        config.interpolation_scale = parsed;
+                    } else {
+                        LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                }
+            }
+            model = PixArtModel(config);
+            model.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "pixart";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            model.get_param_tensors(tensors, prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor,
+                                 const sd::Tensor<float>& mask_tensor) {
+            ggml_cgraph* gf        = new_graph_custom(PIXART_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
+            GGML_ASSERT(x->ne[3] == 1);
+            GGML_ASSERT(!context_tensor.empty());
+            ggml_tensor* context = make_input(context_tensor);
+
+            ggml_tensor* context_mask = nullptr;
+            if (!mask_tensor.empty()) {
+                // additive attention bias over context tokens: 0 keep / -inf discard
+                context_mask = ggml_reshape_4d(compute_ctx, make_input(mask_tensor), mask_tensor.shape()[0], 1, 1, 1);
+            }
+
+            int64_t W  = x->ne[0];
+            int64_t H  = x->ne[1];
+            int64_t wp = W / config.patch_size;
+            int64_t hp = H / config.patch_size;
+
+            pos_vec  = gen_2d_sincos_pos_embed(config.hidden_size, hp, wp,
+                                               config.pos_embed_base_size, config.interpolation_scale);
+            auto pos = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.hidden_size, wp * hp, 1);
+            set_backend_tensor_data(pos, pos_vec.data());
+
+            auto runner_ctx  = get_context();
+            ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pos, context_mask);
+            // learn_sigma: keep the noise prediction half of the output channels
+            out = ggml_ext_slice(compute_ctx, out, 2, 0, config.in_channels);
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context,
+                                  const sd::Tensor<float>& context_mask) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context, context_mask);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            auto context     = tensor_or_empty(diffusion_params.context);
+            auto context_msk = tensor_or_empty(diffusion_params.y);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           context,
+                           context_msk);
+        }
+    };
+}  // namespace PixArt
+
+#endif  // __SD_MODEL_DIFFUSION_PIXART_HPP__

src/model/diffusion/qwen_image_2_1.hpp

@@ -134,6 +134,8 @@ namespace Qwen {
         std::string name;
         std::string cut_group;
         int64_t prefix_length = 0;
+        ggml_type type        = GGML_TYPE_F32;
+        bool* flash_attn_used = nullptr;
     };
 
     class QwenImage21ZeroCenterRMSNorm : public RMSNorm {
@@ -186,13 +188,18 @@ 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.
+                ctx->expand_graph(q);
                 auto persist = [&](ggml_tensor* tensor, int axis, const char* name) {
                     auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length);
-                    auto copy = ggml_new_tensor(ctx->ggml_ctx, GGML_TYPE_F32, 4, part->ne);
-                    copy      = ggml_cpy(ctx->ggml_ctx, part, copy);
+                    // Pack the contiguous data into wider rows so quantization blocks
+                    // can exceed head_dim without padding or changing element order.
+                    part      = ggml_reshape_2d(ctx->ggml_ctx, part, x->ne[0], cache.prefix_length);
+                    auto copy = ggml_cast(ctx->ggml_ctx, part, cache.type);
                     // Keep the copy in this layer's segment so graph cuts do not
                     // retain or recompute the full-sequence K/V in the final segment.
                     sd::ggml_graph_cut::mark_graph_cut(copy, cache.cut_group, name);
@@ -201,21 +208,37 @@ namespace Qwen {
                 persist(k, 1, "k");
                 persist(v, 2, "v");
             }
+            auto attend = [&](ggml_tensor* aq, ggml_tensor* ak, ggml_tensor* av, ggml_tensor* mask) {
+                bool used_flash_attn = false;
+                auto out             = ggml_ext_attention_ext(ctx, aq, ak, av, heads, mask, true, ctx->flash_attn_enabled, 1.f, &used_flash_attn);
+                if (cache.flash_attn_used != nullptr) {
+                    *cache.flash_attn_used &= used_flash_attn;
+                }
+                return out;
+            };
             ggml_tensor* result = nullptr;
             if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
                 auto prefix_k = ctx->load_cache_tensor(cache.name + ".k");
                 auto prefix_v = ctx->load_cache_tensor(cache.name + ".v");
                 GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr);
-                k      = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1);
-                v      = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
-                result = ggml_ext_attention_ext(ctx, q, k, v, heads, nullptr, true, ctx->flash_attn_enabled);
+                if (prefix_k->type != k->type) {
+                    prefix_k = ggml_cast(ctx->ggml_ctx, prefix_k, k->type);
+                }
+                if (prefix_v->type != v->type) {
+                    prefix_v = ggml_cast(ctx->ggml_ctx, prefix_v, v->type);
+                }
+                prefix_k = ggml_reshape_4d(ctx->ggml_ctx, prefix_k, dim_head, cache.prefix_length, heads, k->ne[3]);
+                prefix_v = ggml_reshape_4d(ctx->ggml_ctx, prefix_v, dim_head, heads, cache.prefix_length, v->ne[3]);
+                k        = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1);
+                v        = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
+                result   = attend(q, k, v, nullptr);
             } else {
                 for (size_t i = 0; i < segments.size(); ++i) {
                     const auto& segment = segments[i];
                     auto sq             = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end);
                     auto sk             = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end);
                     auto sv             = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end);
-                    auto out            = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled);
+                    auto out            = attend(sq, sk, sv, masks[i]);
                     result              = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
                 }
             }
@@ -351,8 +374,35 @@ namespace Qwen {
         QwenImage21Model model;
         std::vector<float> pe_data;
         std::vector<sd::Tensor<float>> mask_data;
-        bool prefix_cache_enabled  = true;
-        bool prefix_cache_disabled = false;
+        ggml_type prefix_cache_type = GGML_TYPE_COUNT;
+        bool prefix_cache_enabled   = true;
+        bool prefix_cache_disabled  = false;
+        bool prefix_cache_auto_f32  = false;
+
+        static bool supports_prefix_cache_type(ggml_type type) {
+            if (type == GGML_TYPE_F32) {
+                return true;
+            }
+            const auto* traits = ggml_get_type_traits(type);
+            if (traits->from_float_ref == nullptr || traits->to_float == nullptr) {
+                return false;
+            }
+            auto cpu = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+            if (cpu == nullptr) {
+                return false;
+            }
+            auto ctx = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
+                ggml_init({3 * ggml_tensor_overhead(), nullptr, true}), ggml_free);
+            if (ctx == nullptr) {
+                return false;
+            }
+            // Some reference quantizers have no runtime copy support. Query the
+            // device through the registry so dynamically loaded CPU backends work.
+            auto source  = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, ggml_blck_size(type));
+            auto encoded = ggml_cast(ctx.get(), source, type);
+            auto decoded = ggml_cast(ctx.get(), encoded, GGML_TYPE_F32);
+            return ggml_backend_dev_supports_op(cpu, encoded) && ggml_backend_dev_supports_op(cpu, decoded);
+        }
 
         QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr, const char* model_args = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager),
@@ -361,6 +411,22 @@ namespace Qwen {
             for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
                 if (key == "qwen_image_2_1_prefix_cache" && !parse_strict_bool(value, prefix_cache_enabled)) {
                     LOG_WARN("ignoring invalid Qwen Image 2.1 model arg '%s=%s'", key.c_str(), value.c_str());
+                } else if (key == "qwen_image_2_1_prefix_cache_type") {
+                    if (value == "auto") {
+                        prefix_cache_type = GGML_TYPE_COUNT;
+                        continue;
+                    }
+                    const auto type = sd_type_to_ggml_type(str_to_sd_type(value.c_str()));
+                    if (type == GGML_TYPE_COUNT) {
+                        LOG_WARN("ignoring unknown Qwen Image 2.1 cache type '%s'", value.c_str());
+                    } else if (!supports_prefix_cache_type(type)) {
+                        LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': runtime conversion to and from F32 is unavailable", value.c_str());
+                    } else if (config.hidden_size % ggml_blck_size(type) != 0) {
+                        LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': block size %" PRId64 " does not divide hidden size %" PRId64,
+                                 value.c_str(), ggml_blck_size(type), config.hidden_size);
+                    } else {
+                        prefix_cache_type = type;
+                    }
                 }
             }
             model.init(params_ctx, weights, prefix);
@@ -377,11 +443,9 @@ namespace Qwen {
                 const auto name = cache.name + "." + std::to_string(i);
                 auto k          = get_cache_tensor_by_name(name + ".k");
                 auto v          = get_cache_tensor_by_name(name + ".v");
-                if (k == nullptr || v == nullptr || k->type != GGML_TYPE_F32 || v->type != GGML_TYPE_F32 ||
-                    k->ne[0] != config.head_dim || k->ne[1] != cache.prefix_length ||
-                    k->ne[2] != config.hidden_size / config.head_dim || k->ne[3] != 1 ||
-                    v->ne[0] != config.head_dim || v->ne[1] != config.hidden_size / config.head_dim ||
-                    v->ne[2] != cache.prefix_length || v->ne[3] != 1) {
+                if (k == nullptr || v == nullptr || k->type != cache.type || v->type != cache.type ||
+                    k->ne[0] != config.hidden_size || k->ne[1] != cache.prefix_length || k->ne[2] != 1 || k->ne[3] != 1 ||
+                    v->ne[0] != config.hidden_size || v->ne[1] != cache.prefix_length || v->ne[2] != 1 || v->ne[3] != 1) {
                     return false;
                 }
             }
@@ -418,15 +482,28 @@ namespace Qwen {
             }
             if (!runner_started()) {
                 prefix_cache_disabled = false;
+                prefix_cache_auto_f32 = false;
             }
             QwenImage21PrefixCache cache;
             if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) {
                 cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) +
                              ".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled);
                 cache.prefix_length = layout.prefix_length;
-                cache.mode          = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
+                if (prefix_cache_type != GGML_TYPE_COUNT) {
+                    cache.type = prefix_cache_type;
+                } else if (!prefix_cache_auto_f32 && flash_attn_enabled && !sage_attn_enabled &&
+                           (attn_scale <= 0.f || attn_scale == 1.f)) {
+                    cache.type = GGML_TYPE_F16;
+                }
+                cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
             }
-            auto run = [&](const QwenImage21PrefixCache& active_cache) {
+            bool flash_attn_used = true;
+            auto run             = [&](const QwenImage21PrefixCache& active_cache) {
+                flash_attn_used    = true;
+                auto checked_cache = active_cache;
+                if (prefix_cache_type == GGML_TYPE_COUNT && active_cache.type == GGML_TYPE_F16) {
+                    checked_cache.flash_attn_used = &flash_attn_used;
+                }
                 const bool cached         = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
                 const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0);
                 Rope::Embedding embedding;
@@ -459,7 +536,7 @@ namespace Qwen {
                 auto build = [&]() {
                     auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2);
                     auto pe    = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2,
-                                                    layout.positions.size() - (cached ? layout.prefix_length : 0));
+                                                                layout.positions.size() - (cached ? layout.prefix_length : 0));
                     set_backend_tensor_data(pe, pe_data.data());
                     std::vector<ggml_tensor*> masks, ref_inputs;
                     for (const auto& mask : mask_data) {
@@ -470,15 +547,28 @@ namespace Qwen {
                             ref_inputs.push_back(make_input(ref));
                         }
                     }
-                    auto ctx = get_context();
+                    auto ctx = get_context(graph);
                     auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), cached ? nullptr : make_input(context),
-                                             ref_inputs, pe, layout, masks, active_cache);
+                                                         ref_inputs, pe, layout, masks, checked_cache);
+                    if (!flash_attn_used) {
+                        return static_cast<ggml_cgraph*>(nullptr);
+                    }
                     ggml_build_forward_expand(graph, out);
                     return graph;
                 };
                 return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim());
             };
             auto result = run(cache);
+            if (result.empty() && !flash_attn_used) {
+                // Casting an F16 cache back to F32 cannot recover its original values.
+                // Recompute the prefix before executing a graph that falls back from FA.
+                free_cache_ctx_and_buffer();
+                prefix_cache_auto_f32 = true;
+                cache.type            = GGML_TYPE_F32;
+                cache.mode            = QwenImage21PrefixCache::Mode::STORE;
+                LOG_DEBUG("Qwen Image 2.1: Flash Attention unavailable; using F32 prefix caching for this sampling run");
+                result = run(cache);
+            }
             if (result.empty() && last_compute_status() == GGML_STATUS_ALLOC_FAILED &&
                 (cache.mode != QwenImage21PrefixCache::Mode::NONE || !cache_.empty())) {
                 // The failed graph has ended before persistent inputs are released.
@@ -493,7 +583,7 @@ namespace Qwen {
                     prefix_cache_disabled = true;
                     LOG_WARN("Qwen Image 2.1: incomplete prefix cache; disabling it for this sampling run");
                 } else {
-                    LOG_DEBUG("Qwen Image 2.1: cached prefix %" PRIu64 " (%" PRId64 " tokens)", extra->prefix_id, layout.prefix_length);
+                    LOG_DEBUG("Qwen Image 2.1: cached prefix %" PRIu64 " (%" PRId64 " tokens, %s)", extra->prefix_id, layout.prefix_length, ggml_type_name(cache.type));
                 }
             }
             return result;

src/model/diffusion/sensenova_u1.h

@@ -442,16 +442,9 @@ namespace SenseNovaU1 {
                 k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 2);
                 v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
             } else {
-                // Keep dedicated graph outputs alive until the runner copies them
-                // into its persistent cache buffer after graph execution.
-                auto cache_k = ggml_dup_tensor(ctx->ggml_ctx, k);
-                cache_k      = ggml_cpy(ctx->ggml_ctx, k, cache_k);
-                ggml_set_output(cache_k);
-                auto cache_v = ggml_dup_tensor(ctx->ggml_ctx, v);
-                cache_v      = ggml_cpy(ctx->ggml_ctx, v, cache_v);
-                ggml_set_output(cache_v);
-                ctx->persist_cache_tensor(layer_cache + ".k", cache_k);
-                ctx->persist_cache_tensor(layer_cache + ".v", cache_v);
+                ctx->expand_graph(q);
+                ctx->persist_cache_tensor(layer_cache + ".k", k);
+                ctx->persist_cache_tensor(layer_cache + ".v", v);
             }
 
             q = ggml_cont(ctx->ggml_ctx,
@@ -687,7 +680,7 @@ namespace SenseNovaU1 {
             ggml_set_name(attention_mask, "snu15.prefix.attention_mask");
             set_backend_tensor_data(attention_mask, attention_mask_vec.data());
 
-            auto runner_ctx = get_context();
+            auto runner_ctx = get_context(graph);
             auto text_model = model.text_model();
             auto hidden     = text_model->embed(&runner_ctx, ids);
             hidden          = text_model->forward(&runner_ctx,

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

@@ -140,7 +140,8 @@ namespace ZImage {
                        int64_t num_kv_heads,
                        bool qk_norm,
                        bool norm_elementwise_affine = true,
-                       bool split_qkv               = false)
+                       bool split_qkv               = false,
+                       float qk_norm_eps            = 1e-6f)
             : head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm), split_qkv(split_qkv) {
             float scale = 1.f;
             if (split_qkv) {
@@ -153,8 +154,8 @@ namespace ZImage {
                 blocks["out"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
             }
             if (qk_norm) {
-                blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-06f, norm_elementwise_affine);
-                blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-06f, norm_elementwise_affine);
+                blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, qk_norm_eps, norm_elementwise_affine);
+                blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, qk_norm_eps, norm_elementwise_affine);
             }
         }
 
@@ -167,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);
@@ -196,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);
             }
@@ -318,9 +321,10 @@ namespace ZImage {
                               bool qk_norm,
                               bool modulation              = true,
                               bool norm_elementwise_affine = true,
-                              bool split_qkv               = false)
+                              bool split_qkv               = false,
+                              float qk_norm_eps            = 1e-6f)
             : modulation(modulation) {
-            blocks["attention"]       = std::make_shared<JointAttention>(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm, norm_elementwise_affine, split_qkv);
+            blocks["attention"]       = std::make_shared<JointAttention>(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm, norm_elementwise_affine, split_qkv, qk_norm_eps);
             blocks["feed_forward"]    = std::make_shared<FeedForward>(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier);
             blocks["attention_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
             blocks["ffn_norm1"]       = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);

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/te/llm.hpp

@@ -50,6 +50,8 @@ namespace LLM {
         GEMMA4_12B,
         GPT_OSS_20B,
         LLADA2_MOE,
+        BAILING_MOE,
+        QWEN2,
         ARCH_COUNT,
     };
 
@@ -64,6 +66,8 @@ namespace LLM {
         "gemma4_12b",
         "gpt_oss_20b",
         "llada2_moe",
+        "bailing_moe",
+        "qwen2",
     };
 
     enum class MLPActivation {
@@ -225,7 +229,7 @@ namespace LLM {
                 config.intermediate_size       = 9216;
                 config.num_layers              = 26;
                 config.vocab_size              = 256000;
-            } else if (arch == LLMArch::LLADA2_MOE) {
+            } else if (arch == LLMArch::LLADA2_MOE || arch == LLMArch::BAILING_MOE) {
                 config.head_dim                = 128;
                 config.num_heads               = 16;
                 config.num_kv_heads            = 4;
@@ -240,7 +244,7 @@ namespace LLM {
                 config.max_position_embeddings = 16384;
                 config.rope_thetas             = {600000.f};
                 config.qkv_fused               = true;
-                config.bidirectional           = true;
+                config.bidirectional           = arch == LLMArch::LLADA2_MOE;
                 config.partial_rotary          = 0.5f;
                 config.num_experts             = 256;
                 config.num_experts_per_tok     = 8;
@@ -250,6 +254,18 @@ namespace LLM {
                 config.n_group                 = 8;
                 config.topk_group              = 4;
                 config.routed_scaling_factor   = 2.5f;
+                if (arch == LLMArch::BAILING_MOE) {
+                    config.vocab_size              = 157184;
+                    config.max_position_embeddings = 32768;
+                }
+            } else if (arch == LLMArch::QWEN2) {
+                config.hidden_size             = 1536;
+                config.intermediate_size       = 8960;
+                config.num_heads               = 12;
+                config.num_kv_heads            = 2;
+                config.vocab_size              = 151936;
+                config.max_position_embeddings = 32768;
+                config.bidirectional           = true;
             } else if (arch == LLMArch::GPT_OSS_20B) {
                 config.head_dim                = 64;
                 config.num_heads               = 64;
@@ -283,22 +299,22 @@ namespace LLM {
                     if (contains(name, "attn.q_proj")) {
                         config.llama_cpp_style = true;
                     }
-                    if (contains(name, "visual.patch_embed.proj.1.weight")) {
+                    if (ends_with(name, "visual.patch_embed.proj.1.weight")) {
                         config.vision.split_patch_embed = true;
                     }
-                    if (contains(name, "visual.patch_embed.proj.0.weight")) {
+                    if (ends_with(name, "visual.patch_embed.proj.0.weight")) {
                         config.vision.patch_size  = static_cast<int>(tensor_storage.ne[0]);
                         config.vision.in_channels = tensor_storage.ne[2];
                         config.vision.hidden_size = tensor_storage.ne[3];
                     }
                     // HF-format checkpoints keep the patch embed unsplit under a single name.
-                    if (contains(name, "visual.patch_embed.proj.weight")) {
+                    if (ends_with(name, "visual.patch_embed.proj.weight")) {
                         config.vision.patch_size = static_cast<int>(tensor_storage.ne[0]);
                     }
                     if (contains(name, "visual.patch_embed.bias") || contains(name, "visual.patch_embed.proj.bias")) {
                         config.vision.hidden_size = tensor_storage.ne[0];
                     }
-                    if (contains(name, "visual.pos_embed.weight")) {
+                    if (ends_with(name, "visual.pos_embed.weight") && tensor_storage.n_dims == 2) {
                         config.vision.hidden_size             = tensor_storage.ne[0];
                         config.vision.num_position_embeddings = static_cast<int>(tensor_storage.ne[1]);
                     }
@@ -332,7 +348,7 @@ namespace LLM {
                         }
                     }
                 }
-                if (contains(name, "embed_tokens.weight")) {
+                if (ends_with(name, "embed_tokens.weight") && tensor_storage.n_dims == 2) {
                     config.hidden_size = tensor_storage.ne[0];
                     config.vocab_size  = tensor_storage.ne[1];
                 }
@@ -471,6 +487,8 @@ namespace LLM {
         int64_t n_group;
         int64_t topk_group;
         float routed_scaling_factor;
+        bool image_router;
+        bool fused_experts = false;
 
         void init_params(ggml_context* ctx,
                          const String2TensorStorage& tensor_storage_map = {},
@@ -488,6 +506,10 @@ namespace LLM {
             // scores and the group sums match.
             params["gate.weight"]      = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_experts);
             params["gate.expert_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts);
+            if (image_router) {
+                params["image_gate.weight"]      = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_experts);
+                params["image_gate.expert_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts);
+            }
 
             ggml_type gate_type = supported_type(get_type(prefix + "experts.gate_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
             ggml_type up_type   = supported_type(get_type(prefix + "experts.up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
@@ -506,8 +528,14 @@ namespace LLM {
                 }
             };
 
-            declare_experts("experts.gate_proj.weight", gate_type, hidden_size, moe_intermediate_size);
-            declare_experts("experts.up_proj.weight", up_type, hidden_size, moe_intermediate_size);
+            fused_experts = tensor_storage_map.count(prefix + "experts.gate_up_proj.weight") != 0;
+            if (fused_experts) {
+                auto type = supported_type(get_type(prefix + "experts.gate_up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
+                declare_experts("experts.gate_up_proj.weight", type, hidden_size, 2 * moe_intermediate_size);
+            } else {
+                declare_experts("experts.gate_proj.weight", gate_type, hidden_size, moe_intermediate_size);
+                declare_experts("experts.up_proj.weight", up_type, hidden_size, moe_intermediate_size);
+            }
             declare_experts("experts.down_proj.weight", down_type, moe_intermediate_size, hidden_size);
         }
 
@@ -519,7 +547,8 @@ namespace LLM {
               num_experts_per_tok(config.num_experts_per_tok),
               n_group(config.n_group),
               topk_group(config.topk_group),
-              routed_scaling_factor(config.routed_scaling_factor) {
+              routed_scaling_factor(config.routed_scaling_factor),
+              image_router(config.arch == LLMArch::BAILING_MOE) {
             if (config.num_shared_experts > 0) {
                 blocks["shared_experts"] = std::make_shared<MLP>(config.hidden_size,
                                                                  config.moe_intermediate_size * config.num_shared_experts,
@@ -582,7 +611,7 @@ namespace LLM {
             return ggml_mul_mat_id(ctx->ggml_ctx, w, x, selected_experts);
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* image_mask = nullptr) {
             // x: [N, n_token, hidden_size]
             GGML_ASSERT(num_experts > 0 && num_experts_per_tok > 0);
             GGML_ASSERT(n_group > 0 && topk_group > 0 && num_experts % n_group == 0);
@@ -596,10 +625,21 @@ namespace LLM {
 
             auto logits = ggml_mul_mat(gctx, params["gate.weight"], x);
             logits      = ggml_reshape_2d(gctx, logits, num_experts, n_token_total);
-            auto scores = ggml_sigmoid(gctx, logits);  // [num_experts, tokens]
+            auto bias   = params["gate.expert_bias"];
+            if (image_router) {
+                GGML_ASSERT(image_mask != nullptr);
+                auto mask         = ggml_reshape_2d(gctx, image_mask, 1, n_token_total);
+                auto inverse_mask = ggml_scale_bias(gctx, mask, -1.f, 1.f);
+                auto image_logits = ggml_reshape_2d(gctx, ggml_mul_mat(gctx, params["image_gate.weight"], x), num_experts, n_token_total);
+                logits            = ggml_add(gctx, ggml_mul(gctx, logits, inverse_mask), ggml_mul(gctx, image_logits, mask));
+                bias              = ggml_add(gctx,
+                                             ggml_mul(gctx, ggml_repeat_4d(gctx, bias, num_experts, n_token_total, 1, 1), inverse_mask),
+                                             ggml_mul(gctx, ggml_repeat_4d(gctx, params["image_gate.expert_bias"], num_experts, n_token_total, 1, 1), mask));
+            }
+            auto scores = ggml_sigmoid(gctx, logits);
 
             // The bias steers selection only; the combine weights come from the unbiased scores.
-            auto routing = ggml_add(gctx, scores, params["gate.expert_bias"]);
+            auto routing = ggml_add(gctx, scores, bias);
             routing      = ggml_add(gctx, routing, group_limited_mask(ctx, routing, n_token_total));
 
             auto selected_experts = ggml_argsort_top_k(gctx, routing, (int)num_experts_per_tok);  // [top_k, tokens]
@@ -614,12 +654,17 @@ namespace LLM {
             weights = ggml_scale(gctx, weights, routed_scaling_factor);
             weights = ggml_reshape_3d(gctx, weights, 1, num_experts_per_tok, n_token_total);
 
-            auto xf        = ggml_reshape_3d(gctx, x, hidden_size, 1, n_token_total);
-            auto gate      = expert_linear(ctx, "experts.gate_proj.weight", xf, selected_experts);
-            auto up        = expert_linear(ctx, "experts.up_proj.weight", xf, selected_experts);
-            auto activated = ggml_swiglu_split(gctx, gate, up);
-            auto experts   = expert_linear(ctx, "experts.down_proj.weight", activated, selected_experts);
-            experts        = ggml_mul(gctx, experts, weights);
+            auto xf = ggml_reshape_3d(gctx, x, hidden_size, 1, n_token_total);
+            ggml_tensor* activated;
+            if (fused_experts) {
+                activated = ggml_swiglu(gctx, expert_linear(ctx, "experts.gate_up_proj.weight", xf, selected_experts));
+            } else {
+                auto gate = expert_linear(ctx, "experts.gate_proj.weight", xf, selected_experts);
+                auto up   = expert_linear(ctx, "experts.up_proj.weight", xf, selected_experts);
+                activated = ggml_swiglu_split(gctx, gate, up);
+            }
+            auto experts = expert_linear(ctx, "experts.down_proj.weight", activated, selected_experts);
+            experts      = ggml_mul(gctx, experts, weights);
 
             ggml_tensor* out = nullptr;
             for (int64_t i = 0; i < num_experts_per_tok; ++i) {
@@ -1428,7 +1473,9 @@ namespace LLM {
                              ggml_tensor* x,
                              ggml_tensor* input_pos,
                              ggml_tensor* attention_mask = nullptr,
-                             int rope_index              = 0) {
+                             int rope_index              = 0,
+                             ggml_tensor* rope_cos       = nullptr,
+                             ggml_tensor* rope_sin       = nullptr) {
             // x: [N, n_token, hidden_size]
             int64_t n_token = x->ne[1];
             int64_t N       = x->ne[2];
@@ -1471,13 +1518,28 @@ namespace LLM {
                 v = ggml_rms_norm(ctx->ggml_ctx, v, rms_norm_eps);
             }
 
-            if (arch == LLMArch::MISTRAL_SMALL_3_2) {
+            if (rope_cos != nullptr) {
+                GGML_ASSERT(rope_sin != nullptr);
+                // Bailing video RoPE interleaves spatial frequencies within a partial NEOX head.
+                auto rotate = [&](ggml_tensor* input) {
+                    auto gctx    = ctx->ggml_ctx;
+                    int64_t half = rope_cos->ne[0];
+                    auto first   = ggml_ext_slice(gctx, input, 0, 0, half);
+                    auto second  = ggml_ext_slice(gctx, input, 0, half, 2 * half);
+                    auto left    = ggml_sub(gctx, ggml_mul(gctx, first, rope_cos), ggml_mul(gctx, second, rope_sin));
+                    auto right   = ggml_add(gctx, ggml_mul(gctx, second, rope_cos), ggml_mul(gctx, first, rope_sin));
+                    auto rotated = ggml_concat(gctx, left, right, 0);
+                    return ggml_concat(gctx, rotated, ggml_ext_slice(gctx, input, 0, 2 * half, input->ne[0]), 0);
+                };
+                q = rotate(q);
+                k = rotate(k);
+            } else if (arch == LLMArch::MISTRAL_SMALL_3_2) {
                 q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
                 k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
             } else if (arch == LLMArch::MINISTRAL_3_3B) {
                 q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 262144, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
                 k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 262144, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
-            } else if (arch == LLMArch::QWEN3) {
+            } else if (arch == LLMArch::QWEN3 || arch == LLMArch::QWEN2) {
                 q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
                 k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
             } else if (arch == LLMArch::GPT_OSS_20B) {
@@ -1655,7 +1717,7 @@ namespace LLM {
                 v_attn      = ggml_reshape_3d(ctx->ggml_ctx, v_attn, n_token, head_dim, num_kv_heads * N);
 
                 auto kq = ggml_mul_mat(ctx->ggml_ctx, k, q);
-                ggml_mul_mat_set_prec(kq, GGML_PREC_F32);
+                ggml_prec_set_acc(kq, GGML_PREC_F32);
                 kq = ggml_scale_inplace(ctx->ggml_ctx, kq, 1.0f / std::sqrt(static_cast<float>(head_dim)));
                 if (attention_mask != nullptr) {
                     kq = ggml_add_inplace(ctx->ggml_ctx, kq, attention_mask);
@@ -1724,7 +1786,7 @@ namespace LLM {
             blocks["self_attn"] = std::make_shared<Attention>(config, sliding_attention == 0);
             if (config.arch == LLMArch::GPT_OSS_20B) {
                 blocks["mlp"] = std::make_shared<GPTOSSMLP>(config);
-            } else if (config.arch == LLMArch::LLADA2_MOE && layer_index >= config.first_k_dense_replace) {
+            } else if ((config.arch == LLMArch::LLADA2_MOE || config.arch == LLMArch::BAILING_MOE) && layer_index >= config.first_k_dense_replace) {
                 blocks["mlp"] = std::make_shared<LLaDA2MoEMLP>(config);
             } else {
                 blocks["mlp"] = std::make_shared<MLP>(config.hidden_size,
@@ -1746,7 +1808,10 @@ namespace LLM {
                              ggml_tensor* x,
                              ggml_tensor* input_pos,
                              ggml_tensor* attention_mask         = nullptr,
-                             ggml_tensor* sliding_attention_mask = nullptr) {
+                             ggml_tensor* sliding_attention_mask = nullptr,
+                             ggml_tensor* image_mask             = nullptr,
+                             ggml_tensor* rope_cos               = nullptr,
+                             ggml_tensor* rope_sin               = nullptr) {
             // x: [N, n_token, hidden_size]
             auto self_attn                                  = std::dynamic_pointer_cast<Attention>(blocks["self_attn"]);
             auto input_layernorm                            = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["input_layernorm"]);
@@ -1768,7 +1833,7 @@ namespace LLM {
 
             auto residual = x;
             x             = input_layernorm->forward(ctx, x);
-            x             = self_attn->forward(ctx, x, input_pos, block_attention_mask, rope_index);
+            x             = self_attn->forward(ctx, x, input_pos, block_attention_mask, rope_index, rope_cos, rope_sin);
             if (post_attention_norm != nullptr) {
                 x = post_attention_norm->forward(ctx, x);
             }
@@ -1782,7 +1847,7 @@ namespace LLM {
             } else if (auto moe_mlp = std::dynamic_pointer_cast<LLaDA2MoEMLP>(blocks["mlp"])) {
                 // LLaDA2 is dense for the first first_k_dense_replace layers and MoE afterwards,
                 // so the block type varies per layer rather than per arch.
-                x = moe_mlp->forward(ctx, x);
+                x = moe_mlp->forward(ctx, x, image_mask);
             } else {
                 auto mlp = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
                 x        = mlp->forward(ctx, x);
@@ -1804,10 +1869,11 @@ namespace LLM {
     protected:
         int64_t num_layers;
         LLMConfig config;
+        std::string graph_cut_prefix;
 
     public:
-        TextModel(const LLMConfig& config)
-            : num_layers(config.num_layers), config(config) {
+        TextModel(const LLMConfig& config, const std::string& graph_cut_prefix = "llm.text")
+            : num_layers(config.num_layers), config(config), graph_cut_prefix(graph_cut_prefix) {
             blocks["embed_tokens"] = std::shared_ptr<GGMLBlock>(new Embedding(config.vocab_size, config.hidden_size));
             for (int i = 0; i < num_layers; i++) {
                 blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new TransformerBlock(config, i));
@@ -1831,7 +1897,10 @@ namespace LLM {
                                     std::set<int> out_layers,
                                     const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds = {},
                                     ggml_tensor* sliding_attention_mask                                                  = nullptr,
-                                    bool return_all_hidden_states                                                        = false) {
+                                    bool return_all_hidden_states                                                        = false,
+                                    ggml_tensor* image_mask                                                              = nullptr,
+                                    ggml_tensor* rope_cos                                                                = nullptr,
+                                    ggml_tensor* rope_sin                                                                = nullptr) {
             auto norm = config.final_norm ? std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"])
                                           : nullptr;
             std::vector<ggml_tensor*> intermediate_outputs;
@@ -1843,18 +1912,18 @@ namespace LLM {
                 intermediate_outputs.push_back(x);
             }
 
-            sd::ggml_graph_cut::mark_graph_cut(x, "llm.text.prelude", "x");
+            sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".prelude", "x");
             for (int i = 0; i < num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["layers." + std::to_string(i)]);
 
-                x = block->forward(ctx, x, input_pos, attention_mask, sliding_attention_mask);
+                x = block->forward(ctx, x, input_pos, attention_mask, sliding_attention_mask, image_mask, rope_cos, rope_sin);
                 if (i < static_cast<int>(deepstack_image_embeds.size())) {
                     x = add_deepstack_image_embeds(ctx, x, deepstack_image_embeds[static_cast<size_t>(i)]);
                 }
                 if (return_all_hidden_states || out_layers.size() > 1) {
                     x = ggml_cont(ctx->ggml_ctx, x);
                 }
-                sd::ggml_graph_cut::mark_graph_cut(x, "llm.text.layers." + std::to_string(i), "x");
+                sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".layers." + std::to_string(i), "x");
                 if (return_all_hidden_states) {
                     if (i + 1 < num_layers) {
                         intermediate_outputs.push_back(x);

src/model/te/ming_image_te.hpp

@@ -0,0 +1,159 @@
+#ifndef __SD_MODEL_TE_MING_IMAGE_TE_HPP__
+#define __SD_MODEL_TE_MING_IMAGE_TE_HPP__
+
+#include "llm.hpp"
+
+namespace MingImageTE {
+    struct MingImageTEConfig {
+        LLM::LLMConfig backbone;
+        LLM::LLMConfig connector;
+        int64_t num_queries   = 256;
+        int64_t caption_dim   = 2560;
+        int64_t diffusion_dim = 3840;
+
+        static MingImageTEConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) {
+            MingImageTEConfig config;
+            for (const auto& entry : tensors) {
+                if (starts_with(entry.first, prefix + ".backbone.") &&
+                    contains(entry.first, ".mlp.experts.") && entry.second.type == GGML_TYPE_I8) {
+                    throw std::runtime_error("Ming-Image INT8/W4A8 text encoder experts are not supported; use the BF16 text encoder");
+                }
+            }
+            bool vision           = false;
+            config.backbone       = LLM::LLMConfig::detect_from_weights(tensors, prefix + ".backbone.", LLM::LLMArch::BAILING_MOE, vision);
+            config.connector      = LLM::LLMConfig::detect_from_weights(tensors, prefix + ".connector.", LLM::LLMArch::QWEN2, vision);
+            const auto query      = tensors.find(prefix + ".query_tokens_dict.16x16");
+            const auto projection = tensors.find(prefix + ".proj_out.weight");
+            const auto direct     = tensors.find(prefix + ".proj_directvlm.1.weight");
+            if (query == tensors.end() || projection == tensors.end() || direct == tensors.end()) {
+                throw std::runtime_error("Ming-Image requires the learned queries, connector and both condition projections");
+            }
+            config.num_queries   = query->second.ne[1];
+            config.caption_dim   = projection->second.ne[1];
+            config.diffusion_dim = direct->second.ne[1];
+            if (config.num_queries != 256 || config.caption_dim != 2560 || config.diffusion_dim != 3840 ||
+                config.backbone.num_layers != 20 || config.backbone.hidden_size != 2048 ||
+                config.connector.num_layers != 28 || config.connector.hidden_size != 1536) {
+                throw std::runtime_error("unsupported Ming-Image text encoder configuration");
+            }
+            LOG_VERBOSE("ming_image_te: queries = %" PRId64 ", caption_dim = %" PRId64 ", diffusion_dim = %" PRId64,
+                        config.num_queries, config.caption_dim, config.diffusion_dim);
+            return config;
+        }
+    };
+
+    struct ConnectorModel : LLM::TextModel {
+        explicit ConnectorModel(const LLM::LLMConfig& config)
+            : LLM::TextModel(config, "ming_image.connector") {
+            blocks.erase("embed_tokens");
+        }
+    };
+
+    class MingImageTextModel : public GGMLBlock {
+        MingImageTEConfig config;
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensors = {}, const std::string prefix = "") override {
+            params["query_tokens_dict.16x16"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.backbone.hidden_size, config.num_queries);
+        }
+
+    public:
+        explicit MingImageTextModel(const MingImageTEConfig& config)
+            : config(config) {
+            blocks["backbone"]         = std::make_shared<LLM::TextModel>(config.backbone, "ming_image.backbone");
+            blocks["connector"]        = std::make_shared<ConnectorModel>(config.connector);
+            blocks["proj_in"]          = std::make_shared<Linear>(config.backbone.hidden_size, config.connector.hidden_size);
+            blocks["proj_out"]         = std::make_shared<Linear>(config.connector.hidden_size, config.caption_dim);
+            blocks["proj_directvlm.0"] = std::make_shared<RMSNorm>(config.backbone.hidden_size * 3, 1e-5f);
+            blocks["proj_directvlm.1"] = std::make_shared<Linear>(config.backbone.hidden_size * 3, config.diffusion_dim);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* ids, int64_t prompt_length, ggml_tensor* mask, ggml_tensor* image_mask, ggml_tensor* cos, ggml_tensor* sin, ggml_tensor* connector_positions) {
+            auto gctx                 = ctx->ggml_ctx;
+            auto backbone             = std::dynamic_pointer_cast<LLM::TextModel>(blocks["backbone"]);
+            auto connector            = std::dynamic_pointer_cast<ConnectorModel>(blocks["connector"]);
+            auto x                    = backbone->embed(ctx, ids);
+            const int64_t query_start = prompt_length + 1;
+            auto before               = ggml_ext_slice(gctx, x, 1, 0, query_start);
+            auto after                = ggml_ext_slice(gctx, x, 1, query_start + config.num_queries, x->ne[1]);
+            x                         = ggml_concat(gctx, ggml_concat(gctx, before, params["query_tokens_dict.16x16"], 1), after, 1);
+            // HF hidden_states[20] includes final RMSNorm; sd.cpp selects it as num_layers + 1.
+            x            = backbone->forward_embeds(ctx, x, nullptr, mask, {5, 12, 21}, {}, nullptr, false, image_mask, cos, sin);
+            auto direct  = ggml_ext_slice(gctx, x, 1, 0, prompt_length);
+            direct       = std::dynamic_pointer_cast<RMSNorm>(blocks["proj_directvlm.0"])->forward(ctx, direct);
+            direct       = std::dynamic_pointer_cast<Linear>(blocks["proj_directvlm.1"])->forward(ctx, direct);
+            auto queries = ggml_ext_slice(gctx, x, 0, config.backbone.hidden_size * 2, config.backbone.hidden_size * 3);
+            queries      = ggml_ext_slice(gctx, queries, 1, query_start, query_start + config.num_queries);
+            queries      = std::dynamic_pointer_cast<Linear>(blocks["proj_in"])->forward(ctx, queries);
+            queries      = connector->forward_embeds(ctx, queries, connector_positions, nullptr, {});
+            queries      = std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->forward(ctx, queries);
+            queries      = ggml_pad(gctx, queries, static_cast<int>(config.diffusion_dim - config.caption_dim), 0, 0, 0);
+            return ggml_concat(gctx, queries, direct, 1);
+        }
+    };
+
+    struct MingImageTextRunner : GGMLRunner {
+        MingImageTEConfig config;
+        MingImageTextModel model;
+
+        MingImageTextRunner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : GGMLRunner(backend, weight_manager), config(MingImageTEConfig::detect_from_weights(tensors, prefix)), model(config) {
+            model.init(params_ctx, tensors, prefix);
+        }
+
+        std::string get_desc() override { return "ming_image_text"; }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
+            model.get_param_tensors(tensors, prefix);
+        }
+
+        void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& ops) {
+            model.get_param_tensor_ops(ops);
+        }
+
+        sd::Tensor<float> compute(int n_threads, const std::vector<int>& tokens) {
+            const int64_t prompt_length = tokens.size();
+            const int64_t query_start   = prompt_length + 1;
+            const int64_t total         = prompt_length + config.num_queries + 2;
+            std::vector<int32_t> ids(tokens.begin(), tokens.end());
+            ids.push_back(157158);
+            ids.insert(ids.end(), config.num_queries, 157157);
+            ids.push_back(157159);
+            auto input = sd::Tensor<int32_t>({total}, std::move(ids));
+            sd::Tensor<float> attention_mask({total, total});
+            sd::Tensor<float> image_mask({1, total});
+            sd::Tensor<float> cos({32, 1, total}), sin({32, 1, total});
+            std::vector<int32_t> positions(config.num_queries);
+            std::iota(positions.begin(), positions.end(), 0);
+            auto connector_positions = sd::Tensor<int32_t>({config.num_queries}, std::move(positions));
+            for (int64_t token = 0; token < total; ++token) {
+                const bool query  = token >= query_start && token < query_start + config.num_queries;
+                image_mask[token] = query ? 1.f : 0.f;
+                for (int64_t key = 0; key < total; ++key) {
+                    attention_mask[key + token * total] = key > token ? -INFINITY : 0.f;
+                }
+                // A 16x16 query bank is represented upstream as a [1, 2, 512] image grid,
+                // then spatially merged and centered to [1, 1, 256].
+                int64_t temporal = query ? query_start : (token == total - 1 ? query_start + 1 : token);
+                int64_t width    = query ? token - 127 : temporal;
+                for (int j = 0; j < 32; ++j) {
+                    int64_t position    = query && j < 24 && j % 2 ? width : temporal;
+                    float frequency     = 1.f / std::pow(600000.f, static_cast<float>(2 * j) / 64.f);
+                    float angle         = static_cast<float>(position) * frequency;
+                    cos[token * 32 + j] = std::cos(angle);
+                    sin[token * 32 + j] = std::sin(angle);
+                }
+            }
+            auto graph = [&]() {
+                auto gf  = new_graph_custom(LLM::LLM_GRAPH_SIZE);
+                auto ctx = get_context();
+                auto out = model.forward(&ctx, make_input(input), prompt_length, make_input(attention_mask),
+                                         make_input(image_mask), make_input(cos), make_input(sin), make_input(connector_positions));
+                ggml_build_forward_expand(gf, out);
+                return gf;
+            };
+            return take_or_empty(GGMLRunner::compute(graph, n_threads));
+        }
+    };
+}
+
+#endif  // __SD_MODEL_TE_MING_IMAGE_TE_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/auto_encoder_kl.hpp

@@ -544,7 +544,7 @@ public:
                        const String2TensorStorage& tensor_storage_map = {},
                        const std::string& prefix                      = "")
         : version(version), decode_only(decode_only), use_video_decoder(use_video_decoder) {
-        if (sd_version_is_dit(version)) {
+        if (sd_version_is_dit(version) && version != VERSION_PIXART) {
             if (sd_version_uses_flux2_vae(version)) {
                 dd_config.z_channels = 32;
                 embed_dim            = 32;
@@ -678,7 +678,7 @@ struct AutoEncoderKL : public VAE {
         if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
             scale_factor = 0.18215f;
             shift_factor = 0.f;
-        } else if (sd_version_is_sdxl(version)) {
+        } else if (sd_version_is_sdxl(version) || sd_version_is_pixart(version)) {
             scale_factor = 0.13025f;
             shift_factor = 0.f;
         } else if (sd_version_is_sd3(version)) {

src/model/vae/ltx_vae.hpp

@@ -1300,7 +1300,7 @@ struct LTXVideoVAE : public VAE {
             feat_map[feat_idx] = get_cache_tensor_by_name(temporal_feat_cache_name(feat_idx));
         }
 
-        auto runner_ctx  = get_context();
+        auto runner_ctx  = get_context(gf);
         int feat_count   = 0;
         ggml_tensor* out = vae.decode_tiled_chunk(&runner_ctx,
                                                   z,
@@ -1313,8 +1313,7 @@ struct LTXVideoVAE : public VAE {
         for (int feat_idx = 0; feat_idx < feat_count && feat_idx < static_cast<int>(feat_map.size()); ++feat_idx) {
             ggml_tensor* feat_cache = feat_map[static_cast<size_t>(feat_idx)];
             if (feat_cache != nullptr) {
-                cache(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
-                ggml_build_forward_expand(gf, feat_cache);
+                runner_ctx.persist_cache_tensor(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
             }
         }
 

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,
@@ -556,12 +557,18 @@ namespace MiniMaxH3VAE {
                                    tensor.shape()[3]});
         }
 
-        static sd_tiling_params_t h3_tiling(sd_tiling_params_t params) {
+        sd_tiling_params_t resolve_tiling_params(sd_tiling_params_t params) const override {
+            if (!params.enabled) {
+                params.target_overlap = 0.25f;
+            }
+            if (params.tile_size_w == 0 && params.rel_size_w == 0.f) {
+                params.tile_size_w = 256;
+            }
+            if (params.tile_size_h == 0 && params.rel_size_h == 0.f) {
+                params.tile_size_h = 256;
+            }
             params.enabled         = true;
             params.temporal_tiling = false;
-            params.tile_size_x     = 16;
-            params.tile_size_y     = 16;
-            params.target_overlap  = 0.25f;
             return params;
         }
 
@@ -605,7 +612,7 @@ namespace MiniMaxH3VAE {
                                  bool circular_x = false,
                                  bool circular_y = false) override {
             auto input  = ensure_video_shape(x);
-            auto tiling = h3_tiling(tiling_params);
+            auto tiling = resolve_tiling_params(tiling_params);
             if (input.shape()[2] == 1) {
                 auto encoded = VAE::encode(n_threads, input, tiling, circular_x, circular_y);
                 if (!encoded.empty() && encoded.shape()[2] > 1) {
@@ -627,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 {};
             }
@@ -646,7 +656,7 @@ namespace MiniMaxH3VAE {
                                  bool circular_y   = false,
                                  bool silent       = false) override {
             auto input  = ensure_video_shape(x);
-            auto tiling = h3_tiling(tiling_params);
+            auto tiling = resolve_tiling_params(tiling_params);
             if (input.shape()[2] == 1) {
                 auto decoded = VAE::decode(n_threads,
                                            input,
@@ -689,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>();
                 }
@@ -722,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;
@@ -701,17 +717,21 @@ public:
         bool use_midblock_gn = false;
         taef2                = sd_version_uses_flux2_vae(version);
 
-        if (sd_version_is_dit(version)) {
+        if (sd_version_is_dit(version) && !sd_version_is_pixart(version)) {
             z_channels = 16;
         }
         if (taef2) {
             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

@@ -1,6 +1,9 @@
 #ifndef __SD_MODEL_VAE_VAE_HPP__
 #define __SD_MODEL_VAE_VAE_HPP__
 
+#include <cmath>
+#include <limits>
+
 #include "core/tensor_ggml.hpp"
 #include "model/common/block.hpp"
 #include "model/vae/vae_tiling.hpp"
@@ -117,8 +120,8 @@ protected:
                                     int output_width,
                                     int output_height,
                                     int scale,
-                                    int p_tile_size_x,
-                                    int p_tile_size_y,
+                                    int p_tile_size_w,
+                                    int p_tile_size_h,
                                     float tile_overlap_factor,
                                     bool circular_x,
                                     bool circular_y,
@@ -138,17 +141,28 @@ protected:
             }
             return output_tile;
         };
-        return ::process_tiles_2d(input,
-                                  output_width,
-                                  output_height,
-                                  scale,
-                                  p_tile_size_x,
-                                  p_tile_size_y,
-                                  tile_overlap_factor,
-                                  circular_x,
-                                  circular_y,
-                                  on_processing,
-                                  silent);
+        const bool original_circular_x = circular_x_enabled;
+        const bool original_circular_y = circular_y_enabled;
+        const int64_t latent_width     = decode_graph ? input.shape()[0] : output_width;
+        const int64_t latent_height    = decode_graph ? input.shape()[1] : output_height;
+        circular_x                     = circular_x || original_circular_x;
+        circular_y                     = circular_y || original_circular_y;
+        // Full-width axes wrap in convolutions; split axes wrap between tiles.
+        set_circular_axes(circular_x && p_tile_size_w >= latent_width,
+                          circular_y && p_tile_size_h >= latent_height);
+        auto output = ::process_tiles_2d(input,
+                                         output_width,
+                                         output_height,
+                                         scale,
+                                         p_tile_size_w,
+                                         p_tile_size_h,
+                                         tile_overlap_factor,
+                                         circular_x && p_tile_size_w < latent_width,
+                                         circular_y && p_tile_size_h < latent_height,
+                                         on_processing,
+                                         silent);
+        set_circular_axes(original_circular_x, original_circular_y);
+        return output;
     }
 
 public:
@@ -162,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;
@@ -178,33 +201,48 @@ public:
         return supports_temporal_tiling(VAETemporalDirection::DECODE);
     }
 
-    void get_tile_sizes(int& tile_size_x,
-                        int& tile_size_y,
+    virtual sd_tiling_params_t resolve_tiling_params(sd_tiling_params_t params) const {
+        return params;
+    }
+
+    bool get_tile_sizes(int& tile_size_w,
+                        int& tile_size_h,
                         float& tile_overlap,
                         const sd_tiling_params_t& params,
-                        int64_t latent_x,
-                        int64_t latent_y,
-                        float encoding_factor = 1.0f) {
-        tile_overlap       = std::max(std::min(params.target_overlap, 0.5f), 0.0f);
-        auto get_tile_size = [&](int requested_size, float factor, int64_t latent_size) {
-            const int default_tile_size  = 32;
-            const int min_tile_dimension = 4;
-            int tile_size                = default_tile_size;
-            // factor <= 1 means simple fraction of the latent dimension
-            // factor > 1 means number of tiles across that dimension
-            if (factor > 0.f) {
-                if (factor > 1.0)
-                    factor = 1 / (factor - factor * tile_overlap + tile_overlap);
-                tile_size = static_cast<int>(std::round(latent_size * factor));
-            } else if (requested_size >= min_tile_dimension) {
-                tile_size = requested_size;
+                        int64_t latent_w,
+                        int64_t latent_h) {
+        const auto tiling = resolve_tiling_params(params);
+        if (latent_w <= 0 || latent_h <= 0 ||
+            latent_w > std::numeric_limits<int>::max() || latent_h > std::numeric_limits<int>::max() ||
+            !std::isfinite(tiling.target_overlap)) {
+            LOG_ERROR("invalid VAE tiling dimensions or overlap");
+            return false;
+        }
+        const int scale_factor = get_scale_factor();
+        tile_overlap           = std::max(std::min(tiling.target_overlap, 0.5f), 0.0f);
+        auto get_tile_size     = [&](int requested_size, double factor, int64_t latent_size, int& tile_size) {
+            if (requested_size < 0 || !std::isfinite(factor) || factor < 0.0) {
+                LOG_ERROR("VAE tile sizes and relative sizes must be finite and non-negative");
+                return false;
             }
-            tile_size = static_cast<int>(tile_size * encoding_factor);
-            return std::max(std::min(tile_size, static_cast<int>(latent_size)), min_tile_dimension);
+            const int min_tile_dimension = std::min(4, static_cast<int>(latent_size));
+            double size                  = (requested_size > 0 ? requested_size : 256) / scale_factor;
+            if (factor > 0.0) {
+                if (factor > 1.0) {
+                    factor = 1.0 / (factor * (1.0 - tile_overlap) + tile_overlap);
+                }
+                size = std::floor(static_cast<double>(latent_size) * factor);
+            }
+            if (size < min_tile_dimension && (requested_size > 0 || factor > 0.0)) {
+                LOG_ERROR("VAE tile size must be at least %d image pixels on this axis", min_tile_dimension * scale_factor);
+                return false;
+            }
+            tile_size = static_cast<int>(std::min(static_cast<double>(latent_size), std::max<double>(min_tile_dimension, size)));
+            return true;
         };
 
-        tile_size_x = get_tile_size(params.tile_size_x, params.rel_size_x, latent_x);
-        tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
+        return get_tile_size(tiling.tile_size_w, tiling.rel_size_w, latent_w, tile_size_w) &&
+               get_tile_size(tiling.tile_size_h, tiling.rel_size_h, latent_h, tile_size_h);
     }
 
     virtual sd::Tensor<float> encode(int n_threads,
@@ -212,7 +250,17 @@ 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;
         sd::Tensor<float> output;
         if (scale_input) {
@@ -224,21 +272,19 @@ public:
             int64_t W              = input.shape()[0] / scale_factor;
             int64_t H              = input.shape()[1] / scale_factor;
             float tile_overlap;
-            int tile_size_x, tile_size_y;
-            // Image VAE encode is more sensitive to tile boundary context than decode.
-            // Keep the smaller legacy factor for video VAEs, but default image encode
-            // tiles to 64 latent pixels so a 512px SD image is encoded as one tile.
-            const float encode_tile_factor = sd_version_is_minimax_h3(version) ? 1.f : (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f
-                                                                                                                                                                                            : 2.0f;
-            get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor);
-            LOG_VERBOSE("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
+            int tile_size_w, tile_size_h;
+            if (!get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params, W, H)) {
+                return {};
+            }
+            LOG_VERBOSE("VAE encode tile size: %dx%d pixels (%dx%d latent)",
+                        tile_size_w * scale_factor, tile_size_h * scale_factor, tile_size_w, tile_size_h);
             output = tiled_compute(input,
                                    n_threads,
                                    static_cast<int>(W),
                                    static_cast<int>(H),
                                    scale_factor,
-                                   tile_size_x,
-                                   tile_size_y,
+                                   tile_size_w,
+                                   tile_size_h,
                                    tile_overlap,
                                    circular_x,
                                    circular_y,
@@ -252,7 +298,9 @@ public:
                                                   tiling_params);
         }
 
-        runner_end();
+        if (auto_end) {
+            runner_end();
+        }
 
         if (output.empty()) {
             LOG_ERROR("vae encode compute failed");
@@ -270,7 +318,19 @@ 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;
         sd::Tensor<float> output;
 
@@ -279,10 +339,13 @@ public:
             int64_t W              = input.shape()[0] * scale_factor;
             int64_t H              = input.shape()[1] * scale_factor;
             float tile_overlap;
-            int tile_size_x, tile_size_y;
-            get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, input.shape()[0], input.shape()[1]);
+            int tile_size_w, tile_size_h;
+            if (!get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params, input.shape()[0], input.shape()[1])) {
+                return {};
+            }
             if (!silent) {
-                LOG_VERBOSE("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
+                LOG_VERBOSE("VAE decode tile size: %dx%d pixels (%dx%d latent)",
+                            tile_size_w * scale_factor, tile_size_h * scale_factor, tile_size_w, tile_size_h);
             }
             output = tiled_compute(
                 input,
@@ -290,8 +353,8 @@ public:
                 static_cast<int>(W),
                 static_cast<int>(H),
                 scale_factor,
-                tile_size_x,
-                tile_size_y,
+                tile_size_w,
+                tile_size_h,
                 tile_overlap,
                 circular_x,
                 circular_y,
@@ -306,7 +369,9 @@ public:
                                                   tiling_params);
         }
 
-        runner_end();
+        if (auto_end) {
+            runner_end();
+        }
 
         if (output.empty()) {
             LOG_ERROR("vae decode compute failed");

src/model/vae/wan_vae.hpp

@@ -24,6 +24,7 @@ namespace WAN {
         std::tuple<int, int, int> padding;
         std::tuple<int, int, int> dilation;
         bool bias;
+        float scale = 1.f;
 
         void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
             auto weight = tensor_storage_map.find(prefix + "weight");
@@ -60,6 +61,10 @@ namespace WAN {
               dilation(std::move(dilation)),
               bias(bias) {}
 
+        void set_scale(float scale_value) {
+            scale = scale_value;
+        }
+
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cache_x = nullptr) {
             // x: [N*IC, ID, IH, IW]
             // result: x: [N*OC, ID, IH, IW]
@@ -93,14 +98,14 @@ namespace WAN {
                 x2              = ggml_ext_conv_2d(ctx->ggml_ctx, x2, w2, b,
                                                    std::get<2>(stride), std::get<1>(stride), 0, 0,
                                                    std::get<2>(dilation), std::get<1>(dilation),
-                                                   ctx->conv2d_direct_enabled);
+                                                   ctx->conv2d_direct_enabled, false, false, scale);
                 return ggml_reshape_4d(ctx->ggml_ctx, x2, x2->ne[0], x2->ne[1], 1, out_channels);
             }
             return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
                                     std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                     0, 0, 0,
                                     std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
-                                    false, ctx->conv3d_direct_enabled);
+                                    false, ctx->conv3d_direct_enabled, scale);
         }
     };
 
@@ -1079,7 +1084,7 @@ namespace WAN {
 
                 _conv_num     = 34;
                 _enc_conv_num = 26;
-            } else if (version == VERSION_QWEN_IMAGE_LAYERED) {
+            } else if (version == VERSION_QWEN_IMAGE_LAYERED || version == VERSION_MING_IMAGE) {
                 input_channels = 4;
             }
 
@@ -1117,6 +1122,19 @@ namespace WAN {
             } else {
                 blocks["conv2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1}));
             }
+            if (version == VERSION_QWEN_IMAGE_2_1) {
+                // Keep large VAE activations within the FP16 convolution range.
+                const float conv_scale = 1.f / 128.f;
+                std::vector<GGMLBlock*> all_blocks;
+                get_all_blocks(all_blocks);
+                for (auto block : all_blocks) {
+                    if (auto conv = dynamic_cast<Conv2d*>(block)) {
+                        conv->set_scale(conv_scale);
+                    } else if (auto conv = dynamic_cast<CausalConv3d*>(block)) {
+                        conv->set_scale(conv_scale);
+                    }
+                }
+            }
         }
 
         static ggml_tensor* patchify(ggml_context* ctx,
@@ -1405,11 +1423,17 @@ namespace WAN {
         }
 
         sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
+            if (version == VERSION_MING_IMAGE) {
+                return latents / 8.0064f;
+            }
             auto [mean_tensor, std_tensor] = get_latents_mean_std(latents);
             return (latents * std_tensor) / scale_factor + mean_tensor;
         }
 
         sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
+            if (version == VERSION_MING_IMAGE) {
+                return latents * 8.0064f;
+            }
             auto [mean_tensor, std_tensor] = get_latents_mean_std(latents);
             return ((latents - mean_tensor) * scale_factor) / std_tensor;
         }
@@ -1443,15 +1467,14 @@ namespace WAN {
 
             ggml_tensor* z = make_input(z_tensor);
 
-            auto runner_ctx = get_context();
+            auto runner_ctx = get_context(gf);
 
             ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx);
 
             for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) {
                 ggml_tensor* feat_cache = ae._feat_map[feat_idx];
                 if (feat_cache != nullptr) {
-                    cache("feat_idx:" + std::to_string(feat_idx), feat_cache);
-                    ggml_build_forward_expand(gf, feat_cache);
+                    runner_ctx.persist_cache_tensor("feat_idx:" + std::to_string(feat_idx), feat_cache);
                 }
             }
 

src/model_io/gguf_io.cpp

@@ -47,10 +47,14 @@ bool read_gguf_file(const std::string& file_path,
     gguf_context* ctx_gguf_ = nullptr;
     ggml_context* ctx_meta_ = nullptr;
 
-    ctx_gguf_ = gguf_init_from_file(file_path.c_str(), {true, &ctx_meta_});
+    GGUFReader gguf_reader;
+    bool probe_ok = gguf_reader.load(file_path);
+
+    if (!probe_ok || !gguf_reader.has_tensors_beyond_ggml_limits()) {
+        ctx_gguf_ = gguf_init_from_file(file_path.c_str(), {true, &ctx_meta_});
+    }
     if (!ctx_gguf_) {
-        GGUFReader gguf_reader;
-        if (!gguf_reader.load(file_path)) {
+        if (!probe_ok && !gguf_reader.load(file_path)) {
             set_error(error, "failed to open '" + file_path + "' with GGUFReader");
             return false;
         }

src/model_io/gguf_reader_ext.h

@@ -35,23 +35,72 @@ enum class GGUFMetadataType : uint32_t {
 class GGUFReader {
 private:
     std::vector<GGUFTensorInfo> tensors_;
+    bool has_wide_tensors_    = false;
+    uint64_t remaining_bytes_ = 0;
     size_t data_offset_;
     size_t alignment_ = 32;  // default alignment is 32
 
     template <typename T>
     bool safe_read(std::ifstream& fin, T& value) {
-        fin.read(reinterpret_cast<char*>(&value), sizeof(T));
-        return fin.good();
+        return safe_read(fin, reinterpret_cast<char*>(&value), sizeof(T));
     }
 
     bool safe_read(std::ifstream& fin, char* buffer, size_t size) {
+        if (size > remaining_bytes_)
+            return false;
         fin.read(buffer, size);
-        return fin.good();
+        if (!fin.good())
+            return false;
+        remaining_bytes_ -= size;
+        return true;
     }
 
-    bool safe_seek(std::ifstream& fin, std::streamoff offset, std::ios::seekdir dir) {
-        fin.seekg(offset, dir);
-        return fin.good();
+    bool safe_skip(std::ifstream& fin, uint64_t count, uint64_t element_size = 1) {
+        if (count > remaining_bytes_ / element_size)
+            return false;
+        uint64_t size = count * element_size;
+        fin.seekg(static_cast<std::streamoff>(size), std::ios::cur);
+        if (!fin.good())
+            return false;
+        remaining_bytes_ -= size;
+        return true;
+    }
+
+    bool skip_metadata_values(std::ifstream& fin, GGUFMetadataType type, uint64_t count) {
+        switch (type) {
+            case GGUFMetadataType::UINT8:
+            case GGUFMetadataType::INT8:
+            case GGUFMetadataType::BOOL:
+                return safe_skip(fin, count);
+
+            case GGUFMetadataType::UINT16:
+            case GGUFMetadataType::INT16:
+                return safe_skip(fin, count, 2);
+
+            case GGUFMetadataType::UINT32:
+            case GGUFMetadataType::INT32:
+            case GGUFMetadataType::FLOAT32:
+                return safe_skip(fin, count, 4);
+
+            case GGUFMetadataType::UINT64:
+            case GGUFMetadataType::INT64:
+            case GGUFMetadataType::FLOAT64:
+                return safe_skip(fin, count, 8);
+
+            case GGUFMetadataType::STRING:
+                if (count > remaining_bytes_ / sizeof(uint64_t))
+                    return false;
+                for (uint64_t i = 0; i < count; i++) {
+                    uint64_t len = 0;
+                    if (!safe_read(fin, len) || !safe_skip(fin, len))
+                        return false;
+                }
+                return true;
+
+            default:
+                LOG_ERROR("Unknown metadata type=%u", static_cast<uint32_t>(type));
+                return false;
+        }
     }
 
     bool read_metadata(std::ifstream& fin) {
@@ -84,52 +133,12 @@ private:
             return true;
         }
 
-        switch (static_cast<GGUFMetadataType>(type)) {
-            case GGUFMetadataType::UINT8:
-            case GGUFMetadataType::INT8:
-            case GGUFMetadataType::BOOL:
-                return safe_seek(fin, 1, std::ios::cur);
-
-            case GGUFMetadataType::UINT16:
-            case GGUFMetadataType::INT16:
-                return safe_seek(fin, 2, std::ios::cur);
-
-            case GGUFMetadataType::UINT32:
-            case GGUFMetadataType::INT32:
-            case GGUFMetadataType::FLOAT32:
-                return safe_seek(fin, 4, std::ios::cur);
-
-            case GGUFMetadataType::UINT64:
-            case GGUFMetadataType::INT64:
-            case GGUFMetadataType::FLOAT64:
-                return safe_seek(fin, 8, std::ios::cur);
-
-            case GGUFMetadataType::STRING: {
-                uint64_t len = 0;
-                if (!safe_read(fin, len))
-                    return false;
-                return safe_seek(fin, len, std::ios::cur);
-            }
-
-            case GGUFMetadataType::ARRAY: {
-                uint32_t elem_type = 0;
-                uint64_t len       = 0;
-                if (!safe_read(fin, elem_type))
-                    return false;
-                if (!safe_read(fin, len))
-                    return false;
-
-                for (uint64_t i = 0; i < len; i++) {
-                    if (!read_metadata(fin))
-                        return false;
-                }
-                return true;
-            }
-
-            default:
-                LOG_ERROR("Unknown metadata type=%u", type);
+        uint64_t count = 1;
+        if (type == static_cast<uint32_t>(GGUFMetadataType::ARRAY)) {
+            if (!safe_read(fin, type) || !safe_read(fin, count))
                 return false;
         }
+        return skip_metadata_values(fin, static_cast<GGUFMetadataType>(type), count);
     }
 
     GGUFTensorInfo read_tensor_info(std::ifstream& fin) {
@@ -154,6 +163,7 @@ private:
         }
 
         if (n_dims > GGML_MAX_DIMS) {
+            has_wide_tensors_ = true;
             for (uint32_t i = GGML_MAX_DIMS; i < n_dims; i++) {
                 info.shape[GGML_MAX_DIMS - 1] *= info.shape[i];  // stack to last dim;
             }
@@ -174,12 +184,20 @@ private:
 
 public:
     bool load(const std::string& file_path) {
-        std::ifstream fin(file_path, std::ios::binary);
+        std::ifstream fin(file_path, std::ios::binary | std::ios::ate);
         if (!fin) {
             LOG_ERROR("failed to open '%s'", file_path.c_str());
             return false;
         }
 
+        std::streamoff file_size = fin.tellg();
+        if (file_size < 0)
+            return false;
+        remaining_bytes_ = static_cast<uint64_t>(file_size);
+        fin.seekg(0, std::ios::beg);
+        if (!fin.good())
+            return false;
+
         // --- Header ---
         char magic[4];
         if (!safe_read(fin, magic, 4) || strncmp(magic, "GGUF", 4) != 0) {
@@ -228,6 +246,8 @@ public:
     }
 
     const std::vector<GGUFTensorInfo>& tensors() const { return tensors_; }
+
+    bool has_tensors_beyond_ggml_limits() const { return has_wide_tensors_; }
     size_t data_offset() const { return data_offset_; }
 };
 

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);
@@ -322,7 +322,13 @@ bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path,
     }
 
     for (const std::string& shard_path : shard_paths) {
-        if (!parse_file(shard_path, prefix)) {
+        FileStamp stamp;
+        if (!read_file_stamp(shard_path, stamp)) {
+            return false;
+        }
+        parsed_dependencies_.push_back(stamp);
+        LOG_INFO("load %s using safetensors format", shard_path.c_str());
+        if (!init_from_safetensors_file(shard_path, prefix)) {
             return false;
         }
     }
@@ -518,6 +524,12 @@ SDVersion ModelLoader::get_sd_version() const {
             return VERSION_LLADA_IMAGE;
         }
         if (tensor_storage.name.find("model.diffusion_model.cap_embedder.0.weight") != std::string::npos) {
+            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) {
@@ -526,13 +538,28 @@ SDVersion ModelLoader::get_sd_version() const {
         if (tensor_storage.name.find("model.diffusion_model.layers.0.adaLN_sa_ln.weight") != std::string::npos) {
             return VERSION_ERNIE_IMAGE;
         }
+        if (tensor_storage.name.find("model.diffusion_model.t_block.1.weight") != std::string::npos &&
+            tensor_storage_map.find("model.diffusion_model.x_embedder.proj.weight") != tensor_storage_map.end() &&
+            tensor_storage_map.find("model.diffusion_model.audio_patchify_proj.weight") == tensor_storage_map.end()) {
+            return VERSION_PIXART;
+        }
         if (tensor_storage.name.find("model.diffusion_model.adaln_single.emb.timestep_embedder.linear_1.bias") != std::string::npos) {
+            // PixArt shares this timestep embedding with LTX-AV.
+            if (tensor_storage_map.find("model.diffusion_model.pos_embed.proj.weight") != tensor_storage_map.end() &&
+                tensor_storage_map.find("model.diffusion_model.adaln_single.linear.weight") != tensor_storage_map.end() &&
+                tensor_storage_map.find("model.diffusion_model.audio_patchify_proj.weight") == tensor_storage_map.end()) {
+                return VERSION_PIXART;
+            }
             return VERSION_LTXAV;
         }
         if (tensor_storage.name.find("model.diffusion_model.video_patch_proj.weight") != std::string::npos &&
             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;
         }
@@ -1084,10 +1111,12 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
         if (tensors_to_process.empty()) {
             continue;
         }
-        LOG_VERBOSE("loading %zu/%zu tensors from %s",
-                    tensors_to_process.size(),
-                    file_tensors.size(),
-                    file_path.c_str());
+        if (log_progress) {
+            LOG_VERBOSE("loading %zu/%zu tensors from %s",
+                        tensors_to_process.size(),
+                        file_tensors.size(),
+                        file_path.c_str());
+        }
 
         bool is_zip = fdata.is_zip;
 
@@ -1591,6 +1620,9 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage
             // Pass, do not convert. For Unet
         } else if (contains(name, "embedding")) {
             // Pass, do not convert embedding
+        } else if (contains(name, "scale_shift_table")) {
+            // Pass, do not convert. adaLN modulation tables (PixArt, LTXV) are sliced
+            // element-wise, which is invalid on quantized block layouts.
         } else if (ends_with(name, "_pad_token")) {
             // Pass, do not convert. LLaDA-Image stores its pad tokens far outside the f16
             // range, so any format with an f16 scale or payload turns them into inf.

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

@@ -79,6 +79,10 @@ static bool device_supports_param_op(ggml_backend_dev_t device,
     if (op == GGML_OP_GET_ROWS) {
         ggml_tensor* indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
         op_tensor            = ggml_get_rows(ctx, weight, indices);
+    } else if (op == GGML_OP_CUSTOM) {
+        op_tensor         = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+        op_tensor->op     = op;
+        op_tensor->src[0] = weight;
     }
     if (op_tensor == nullptr) {
         ggml_free(ctx);
@@ -246,17 +250,13 @@ bool ModelManager::register_param_tensors(ModelComponent component,
         }
         ggml_set_name(tensor, name.c_str());
 
-        auto state            = std::make_unique<TensorState>();
-        state->name           = name;
-        state->tensor         = tensor;
-        state->component      = component;
-        state->source_file    = source_file;
-        state->source_version = source_version;
-        auto source           = sources.find(name);
-        if (source != sources.end()) {
-            state->source     = source->second;
-            state->has_source = true;
-        }
+        auto state             = std::make_unique<TensorState>();
+        state->name            = name;
+        state->tensor          = tensor;
+        state->component       = component;
+        state->source_file     = source_file;
+        state->source_version  = source_version;
+        state->sources         = find_tensor_sources(*state, sources);
         state->residency_mode  = residency_mode;
         state->compute_backend = compute_backend;
         state->params_backend  = params_backend;
@@ -511,7 +511,9 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
             LOG_ERROR("model manager params backend is null for tensor '%s'", state->name.c_str());
             return false;
         }
-        if (state->compute_backend == state->params_backend || state->staged_to_compute_backend) {
+        // Custom CPU operators must retain host weights even when the runner uses a GPU.
+        if (state->usage_op == GGML_OP_CUSTOM ||
+            state->compute_backend == state->params_backend || state->staged_to_compute_backend) {
             continue;
         }
         if (!state->loaded_to_params_backend || state->tensor == nullptr || state->tensor->data == nullptr) {
@@ -757,12 +759,33 @@ bool ModelManager::validate_tensor(const TensorState& state) const {
         return true;
     }
 
-    if (!state.has_source) {
+    if (state.sources.empty()) {
         LOG_ERROR("%s tensor '%s' not in model metadata", model_component_name(state.component), state.name.c_str());
         return false;
     }
 
-    const TensorStorage& tensor_storage = state.source;
+    TensorStorage tensor_storage = state.sources.front();
+    if (state.sources.size() > 1) {
+        const int dim = tensor_storage.n_dims - 1;
+        if (dim < 0 || dim >= GGML_MAX_DIMS) {
+            return false;
+        }
+        tensor_storage.ne[dim] = 0;
+        for (const auto& part : state.sources) {
+            if (part.n_dims != tensor_storage.n_dims || part.ne[dim] < 0 ||
+                part.ne[dim] > state.tensor->ne[dim] - tensor_storage.ne[dim]) {
+                LOG_ERROR("invalid tensor part '%s' for '%s'", part.name.c_str(), state.name.c_str());
+                return false;
+            }
+            for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+                if (i != dim && part.ne[i] != tensor_storage.ne[i]) {
+                    LOG_ERROR("incompatible tensor part '%s' for '%s'", part.name.c_str(), state.name.c_str());
+                    return false;
+                }
+            }
+            tensor_storage.ne[dim] += part.ne[dim];
+        }
+    }
     if (state.tensor->ne[0] != tensor_storage.ne[0] ||
         state.tensor->ne[1] != tensor_storage.ne[1] ||
         state.tensor->ne[2] != tensor_storage.ne[2] ||
@@ -831,7 +854,7 @@ bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
 }
 
 bool ModelManager::can_mmap_storage(const TensorState& state) const {
-    if (state.source_file != 0 || !enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
+    if (state.sources.size() > 1 || state.source_file != 0 || !enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
         return false;
     }
     if (state.compute_backend == nullptr || state.params_backend == nullptr) {
@@ -941,6 +964,62 @@ bool ModelManager::alloc_params_buffers(const std::vector<TensorState*>& states,
     return true;
 }
 
+bool ModelManager::load_tensor_parts(TensorState& state) {
+    auto ctx = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
+        ggml_init({state.sources.size() * ggml_tensor_overhead(), nullptr, true}), ggml_free);
+    if (!ctx) {
+        return false;
+    }
+    const size_t size = ggml_nbytes(state.tensor);
+    std::vector<uint8_t> buffer;
+    void* data = state.tensor->data;
+    if (!ggml_backend_buffer_is_host(state.tensor->buffer)) {
+        buffer.resize(size);
+        data = buffer.data();
+    }
+    std::map<std::string, ggml_tensor*> parts;
+    std::set<std::string> names;
+    size_t offset = 0;
+    for (const auto& source : state.sources) {
+        auto part              = ggml_new_tensor(ctx.get(), state.tensor->type, source.n_dims, source.ne);
+        const size_t part_size = ggml_nbytes(part);
+        if (part_size > size - offset) {
+            return false;
+        }
+        part->data         = static_cast<uint8_t*>(data) + offset;
+        parts[source.name] = part;
+        names.insert(source.name);
+        offset += part_size;
+    }
+    if (offset != size) {
+        return false;
+    }
+    std::set<std::string> loaded;
+    std::mutex mutex;
+    auto callback = [&](const TensorStorage& source, ggml_tensor** dst) {
+        *dst      = nullptr;
+        auto part = parts.find(source.name);
+        if (part != parts.end()) {
+            *dst = part->second;
+            std::lock_guard<std::mutex> lock(mutex);
+            loaded.insert(source.name);
+        }
+        return true;
+    };
+    const bool success = state.source_file == 0
+                             ? model_loader_.load_tensors(callback, enable_mmap_, &names, false)
+                             : model_loader_.load_file_tensors(state.source_file, state.source_version, callback, names, enable_mmap_);
+    if (!success || loaded != names) {
+        return false;
+    }
+    if (!buffer.empty()) {
+        // Upload the assembled tensor once, including for row-split backend buffers.
+        ggml_backend_tensor_set(state.tensor, buffer.data(), 0, size);
+    }
+    state.loaded_to_params_backend = true;
+    return true;
+}
+
 bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
     using ReadGroup = std::pair<ModelLoader::FileId, SDVersion>;
     using ReadBatch = std::map<std::string, std::vector<TensorState*>>;
@@ -948,6 +1027,12 @@ bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
     for (auto* state : states) {
         if (state == nullptr)
             continue;
+        if (state->sources.size() > 1) {
+            if (!load_tensor_parts(*state)) {
+                return false;
+            }
+            continue;
+        }
         auto& batches = groups[{state->source_file, state->source_version}];
         // The loader supplies one destination per name; only conflicting types need another batch.
         auto batch = std::find_if(batches.begin(), batches.end(), [&](const ReadBatch& candidate) {
@@ -1323,6 +1408,9 @@ size_t ModelManager::compute_backend_alloc_size(const std::vector<TensorState*>&
         if (state == nullptr || state->tensor == nullptr) {
             continue;
         }
+        if (state->usage_op == GGML_OP_CUSTOM && !sd_backend_is_cpu(state->compute_backend)) {
+            continue;
+        }
         const bool compute_resident =
             state->compute_backend == state->params_backend
                 ? state->loaded_to_params_backend
@@ -1816,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/model_manager.h

@@ -38,8 +38,7 @@ private:
         std::string name;
         ggml_tensor* tensor      = nullptr;
         ModelComponent component = ModelComponent::Count;
-        TensorStorage source;
-        bool has_source                 = false;
+        std::vector<TensorStorage> sources;
         ModelLoader::FileId source_file = 0;
         SDVersion source_version        = VERSION_COUNT;
 
@@ -138,6 +137,8 @@ private:
     bool apply_loras_to_params(const std::vector<TensorState*>& states);
     bool mmap_params(const std::vector<TensorState*>& states,
                      std::vector<ParamsStorageBlock*>& created_storage_blocks);
+    static std::vector<TensorStorage> find_tensor_sources(const TensorState& state, const String2TensorStorage& sources);
+    bool load_tensor_parts(TensorState& state);
     bool can_mmap_storage(const TensorState& state) const;
     bool alloc_params_buffers(const std::vector<TensorState*>& states,
                               std::vector<ParamsStorageBlock*>& created_storage_blocks);

src/model_manager_files.cpp

@@ -15,6 +15,26 @@ static bool same_tensor_source(const TensorStorage& a, const TensorStorage& b) {
            a.int8_convrot_group_size == b.int8_convrot_group_size;
 }
 
+std::vector<TensorStorage> ModelManager::find_tensor_sources(const TensorState& state, const String2TensorStorage& sources) {
+    auto first = sources.find(state.name);
+    if (first == sources.end()) {
+        return {};
+    }
+    std::vector<TensorStorage> result{first->second};
+    if (state.component == ModelComponent::LoRA ||
+        std::equal(first->second.ne, first->second.ne + GGML_MAX_DIMS, state.tensor->ne)) {
+        return result;
+    }
+    for (size_t i = 1;; ++i) {
+        auto part = sources.find(state.name + "." + std::to_string(i));
+        if (part == sources.end()) {
+            break;
+        }
+        result.push_back(part->second);
+    }
+    return result;
+}
+
 void ModelManager::invalidate_sources(const std::unordered_set<TensorState*>& states) {
     auto affected = states;
     for (const auto& block : params_storage_blocks_) {
@@ -76,20 +96,16 @@ bool ModelManager::set_loader(ModelLoader loader) {
     }
     std::unordered_set<TensorState*> changed;
     for (const auto& state : tensor_states_) {
-        const auto& sources = sources_for(*state);
-        auto source         = sources.find(state->name);
-        const bool found    = source != sources.end();
-        if (found != state->has_source || (found && !same_tensor_source(state->source, source->second)) ||
+        const auto sources = find_tensor_sources(*state, sources_for(*state));
+        if (sources.size() != state->sources.size() ||
+            !std::equal(sources.begin(), sources.end(), state->sources.begin(), same_tensor_source) ||
             (lora_changed && state->component != ModelComponent::LoRA && state->applied_lora_epoch != UINT64_MAX)) {
             changed.insert(state.get());
         }
     }
     invalidate_sources(changed);
     for (auto* state : changed) {
-        const auto& sources = sources_for(*state);
-        auto source         = sources.find(state->name);
-        state->has_source   = source != sources.end();
-        state->source       = state->has_source ? source->second : TensorStorage{};
+        state->sources = find_tensor_sources(*state, sources_for(*state));
     }
     if (lora_changed) {
         ++current_lora_epoch_;
@@ -134,9 +150,8 @@ ModelLoader::FileVersions ModelManager::source_versions(const std::set<ModelComp
             versions[state->source_file] = loader.file_revision(state->source_file);
             continue;
         }
-        auto source = sources.find(state->name);
-        if (source != sources.end()) {
-            versions[source->second.file_id] = source->second.file_revision;
+        for (const auto& source : find_tensor_sources(*state, sources)) {
+            versions[source.file_id] = source.file_revision;
         }
     }
     return versions;

src/model_manager_prefetch.cpp

@@ -210,7 +210,7 @@ WeightPrefetchResult ModelManager::prefetch_params(
             is_optional_missing_tensor(state->name)) {
             continue;
         }
-        if (state->compute_backend == state->params_backend) {
+        if (state->usage_op == GGML_OP_CUSTOM || state->compute_backend == state->params_backend) {
             needs_synchronous_load = needs_synchronous_load ||
                                      !state->loaded_to_params_backend;
             continue;
@@ -275,6 +275,7 @@ bool ModelManager::activate_prefetched_params(
         [&](TensorState* state) {
             return state == nullptr || should_ignore(*state) ||
                    is_optional_missing_tensor(state->name) ||
+                   state->usage_op == GGML_OP_CUSTOM ||
                    state->compute_backend == state->params_backend ||
                    state->staged_to_compute_backend;
         });

src/name_conversion.cpp

@@ -105,7 +105,25 @@ std::string convert_open_clip_to_hf_clip_name(std::string name) {
 
 std::string convert_llada2_moe_te_name(std::string name);
 
-std::string convert_cond_stage_model_name(std::string name, std::string prefix) {
+static std::string convert_ming_image_te_name(std::string name) {
+    if (name == "llm.query_tokens") {
+        name = "llm.query_tokens_dict.16x16";
+    }
+    static const std::vector<std::pair<std::string, std::string>> name_map = {
+        {"thinker.", "backbone."},
+        {"attention.", "self_attn."},
+        {"self_attn.dense.", "self_attn.o_proj."},
+        {"gate.proj.", "gate."},
+    };
+    replace_with_name_map(name, name_map);
+    return name;
+}
+
+std::string convert_cond_stage_model_name(std::string name, std::string prefix, SDVersion version) {
+    if (version == VERSION_MING_IMAGE && prefix == "text_encoders." && starts_with(name, "llm.")) {
+        name = convert_ming_image_te_name(name);
+    }
+
     static const std::vector<std::pair<std::string, std::string>> clip_name_map{
         {"transformer.text_projection.weight", "transformer.text_model.text_projection"},
         {"model.text_projection.weight", "transformer.text_model.text_projection"},
@@ -806,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) {
@@ -932,6 +951,79 @@ static bool is_diffusers_controlnet_name(const std::string& name) {
     return false;
 }
 
+static std::string convert_diffusers_dit_to_original_pixart(std::string name) {
+    static const std::vector<std::pair<std::string, std::string>> prefix_map = {
+        {"pos_embed.proj.", "x_embedder.proj."},
+        {"adaln_single.emb.timestep_embedder.linear_1.", "t_embedder.mlp.0."},
+        {"adaln_single.emb.timestep_embedder.linear_2.", "t_embedder.mlp.2."},
+        {"adaln_single.emb.resolution_embedder.linear_1.", "csize_embedder.mlp.0."},
+        {"adaln_single.emb.resolution_embedder.linear_2.", "csize_embedder.mlp.2."},
+        {"adaln_single.emb.aspect_ratio_embedder.linear_1.", "ar_embedder.mlp.0."},
+        {"adaln_single.emb.aspect_ratio_embedder.linear_2.", "ar_embedder.mlp.2."},
+        {"adaln_single.linear.", "t_block.1."},
+        {"caption_projection.linear_1.", "y_embedder.y_proj.fc1."},
+        {"caption_projection.linear_2.", "y_embedder.y_proj.fc2."},
+        {"proj_out.", "final_layer.linear."},
+    };
+    for (const auto& entry : prefix_map) {
+        if (starts_with(name, entry.first)) {
+            return entry.second + name.substr(entry.first.size());
+        }
+    }
+    if (name == "scale_shift_table") {
+        return "final_layer.scale_shift_table";
+    }
+    const std::string block_prefix = "transformer_blocks.";
+    if (!starts_with(name, block_prefix)) {
+        return name;
+    }
+    const size_t block_end = name.find('.', block_prefix.size());
+    if (block_end == std::string::npos) {
+        return name;
+    }
+    const std::string prefix                                                = "blocks." + name.substr(block_prefix.size(), block_end - block_prefix.size()) + ".";
+    name                                                                    = name.substr(block_end + 1);
+    static const std::vector<std::pair<std::string, std::string>> block_map = {
+        {"attn1.to_q.", "attn.qkv."},
+        {"attn1.to_out.0.", "attn.proj."},
+        {"attn2.to_q.", "cross_attn.q_linear."},
+        {"attn2.to_k.", "cross_attn.kv_linear."},
+        {"attn2.to_out.0.", "cross_attn.proj."},
+        {"ff.net.0.proj.", "mlp.fc1."},
+        {"ff.net.2.", "mlp.fc2."},
+    };
+    for (const auto& entry : block_map) {
+        if (starts_with(name, entry.first)) {
+            return prefix + entry.second + name.substr(entry.first.size());
+        }
+    }
+    static const std::vector<std::pair<std::string, std::string>> part_map = {
+        {"attn1.to_k.weight", "attn.qkv.weight.1"},
+        {"attn1.to_k.bias", "attn.qkv.bias.1"},
+        {"attn1.to_v.weight", "attn.qkv.weight.2"},
+        {"attn1.to_v.bias", "attn.qkv.bias.2"},
+        {"attn2.to_v.weight", "cross_attn.kv_linear.weight.1"},
+        {"attn2.to_v.bias", "cross_attn.kv_linear.bias.1"},
+    };
+    for (const auto& entry : part_map) {
+        if (name == entry.first || starts_with(name, entry.first + ".")) {
+            return prefix + entry.second + name.substr(entry.first.size());
+        }
+    }
+    return prefix + name;
+}
+
+static std::string convert_ming_image_dit_name(std::string name) {
+    static const std::vector<std::pair<std::string, std::string>> name_map = {
+        {"all_x_embedder.2-1.", "x_embedder."},
+        {"all_final_layer.2-1.", "final_layer."},
+        {"attention.norm_q.", "attention.q_norm."},
+        {"attention.norm_k.", "attention.k_norm."},
+    };
+    replace_with_name_map(name, name_map);
+    return name;
+}
+
 std::string convert_diffusion_model_name(std::string name, std::string prefix, SDVersion version) {
     if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
         name = convert_diffusers_unet_to_original_sd1(name);
@@ -943,7 +1035,9 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
         name = convert_diffusers_dit_to_original_flux(name);
     } else if (sd_version_is_hunyuan_video(version)) {
         name = convert_hunyuan_video_to_original_flux(name);
-    } else if (sd_version_is_z_image(version)) {
+    } else if (version == VERSION_MING_IMAGE) {
+        name = convert_ming_image_dit_name(name);
+    } 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);
@@ -951,6 +1045,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
         name = convert_other_dit_to_original_anima(name);
     } else if (sd_version_is_krea2(version)) {
         name = convert_diffusers_dit_to_original_krea2(name);
+    } else if (sd_version_is_pixart(version)) {
+        name = convert_diffusers_dit_to_original_pixart(name);
     }
     return name;
 }
@@ -1604,7 +1700,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
     {
         for (const auto& prefix : cond_stage_model_prefix_vec) {
             if (starts_with(name, prefix)) {
-                name = convert_cond_stage_model_name(name.substr(prefix.size()), prefix);
+                name = convert_cond_stage_model_name(name.substr(prefix.size()), prefix, version);
                 name = prefix + name;
                 break;
             }

src/pipeline/diffusion_engine.cpp

@@ -105,6 +105,10 @@ const char* model_version_to_str[] = {
     "SenseNova U1.5",
     "LLaDA-Image",
     "ESRGAN",
+    "PixArt",
+    "Ming-Image",
+    "Z-Image L2P",
+    "Iris-3B",
 };
 
 static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]),
@@ -125,6 +129,18 @@ void calculate_alphas_cumprod(float* alphas_cumprod,
     }
 }
 
+void calculate_alphas_cumprod_linear_beta(float* alphas_cumprod,
+                                          float beta_start,
+                                          float beta_end,
+                                          int timesteps = TIMESTEPS) {
+    float product = 1.0f;
+    for (int i = 0; i < timesteps; i++) {
+        float beta = beta_start + (beta_end - beta_start) * ((float)i / (timesteps - 1));
+        product *= 1.0f - beta;
+        alphas_cumprod[i] = product;
+    }
+}
+
 template <typename T, typename = void>
 struct has_set_runtime_backends : std::false_type {};
 template <typename T>
@@ -666,6 +682,10 @@ void StableDiffusionGGML::refresh_compvis_denoiser_sigmas() {
     std::vector<float> alphas_cumprod(TIMESTEPS);
     if (file_alphas_cumprod.size() == TIMESTEPS) {
         alphas_cumprod = file_alphas_cumprod;
+    } else if (sd_version_is_pixart(version)) {
+        // PixArt checkpoints train with a linear beta schedule (0.0001 -> 0.02)
+        // instead of the scaled_linear schedule used by SD1.x/SDXL.
+        calculate_alphas_cumprod_linear_beta(alphas_cumprod.data(), 0.0001f, 0.02f);
     } else {
         calculate_alphas_cumprod(alphas_cumprod.data());
     }
@@ -1193,6 +1213,11 @@ bool StableDiffusionGGML::validate_and_load_runners() {
     LOG_VERBOSE("validating model metadata");
 
     std::set<std::string> ignore_tensors;
+    if (version == VERSION_MING_IMAGE) {
+        ignore_tensors.insert("text_encoders.llm.vision.");
+        ignore_tensors.insert("text_encoders.llm.linear_proj.");
+        ignore_tensors.insert("text_encoders.llm.backbone.lm_head.");
+    }
     if (use_tae && !tae_preview_only) {
         ignore_tensors.insert("first_stage_model.");
     }
@@ -1353,10 +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;
@@ -1364,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;
@@ -1374,6 +1402,8 @@ bool StableDiffusionGGML::build_denoiser() {
                 default_flow_shift = 3.16f;
             } else if (sd_version_is_mage_flow(version)) {
                 default_flow_shift = 6.f;
+            } else if (version == VERSION_MING_IMAGE) {
+                default_flow_shift = INFINITY;
             } else if (sd_version_is_llada_image(version)) {
                 default_flow_shift = 1.0f;  // unused: LLADA_IMAGE_SCHEDULER builds a fixed grid
             } else {
@@ -1434,6 +1464,8 @@ bool StableDiffusionGGML::build_denoiser() {
             } else if (sd_version_is_minimax_h3(version)) {
                 LOG_INFO("running in MiniMax H3 AV FLOW mode");
                 denoiser = std::make_shared<H3AVFlowDenoiser>(default_flow_shift, 3.f, get_latent_channel());
+            } else if (version == VERSION_MING_IMAGE) {
+                denoiser = std::make_shared<MingImageFlowDenoiser>();
             } else {
                 LOG_INFO("running in FLOW mode");
                 denoiser = std::make_shared<DiscreteFlowDenoiser>();
@@ -2144,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)) {
+    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};
@@ -2195,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,
@@ -2343,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);
@@ -2497,6 +2551,9 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
                     condition.c_token_types.empty() ? nullptr : &condition.c_token_types,
                     condition.c_vinput_mask.empty() ? nullptr : &condition.c_vinput_mask,
                     condition.c_image_embeds.empty() ? nullptr : &condition.c_image_embeds};
+            } else if (version == VERSION_MING_IMAGE) {
+                diffusion_params.extra = MingImageDiffusionExtra{
+                    condition.extra_c_crossattns.empty() ? nullptr : &condition.extra_c_crossattns[0]};
             } else if (sd_version_is_llada_image(version)) {
                 diffusion_params.extra = LLaDAImageDiffusionExtra{
                     condition.extra_c_crossattns.empty() ? nullptr : &condition.extra_c_crossattns[0]};
@@ -2568,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) {
@@ -2607,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) {
@@ -2731,7 +2867,9 @@ 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 (version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version)) {
+        } else if (sd_version_is_z_image_l2p(version) || version == VERSION_IRIS) {
+            down_factor = 16;
+        } else if (version == VERSION_QWEN_IMAGE_2_1 || version == VERSION_MING_IMAGE || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version) || sd_version_is_pixart(version)) {
             down_factor = 2;
         } else {
             down_factor = 1;
@@ -2761,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;
@@ -2769,6 +2909,8 @@ int StableDiffusionGGML::get_latent_channel() {
             latent_channel = 128;
         } else if (sd_version_is_mage_flow(version)) {
             latent_channel = 128;
+        } else if (sd_version_is_pixart(version)) {
+            latent_channel = 4;
         } else {
             latent_channel = 16;
         }
@@ -2777,7 +2919,7 @@ int StableDiffusionGGML::get_latent_channel() {
 }
 
 int StableDiffusionGGML::get_image_channels() const {
-    return version == VERSION_QWEN_IMAGE_LAYERED || version == VERSION_QWEN_IMAGE_2_1 ? 4 : 3;
+    return version == VERSION_QWEN_IMAGE_LAYERED || version == VERSION_QWEN_IMAGE_2_1 || version == VERSION_MING_IMAGE ? 4 : 3;
 }
 
 int StableDiffusionGGML::get_image_seq_len(int h, int w) {
@@ -2860,18 +3002,30 @@ 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);
-    auto decoded                      = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
-    const bool prefer_temporal_tiling = decode_video && first_stage_model->can_temporal_tile_decode();
-    while (decoded.empty() &&
-           sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling,
-                                                            first_stage_model->last_compute_status())) {
-        decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
+    auto tiling_params                = first_stage_model->resolve_tiling_params(vae_tiling_params);
+    const bool prefer_temporal_tiling = decode_video && latents.dim() == 5 && latents.shape()[2] > 1 &&
+                                        first_stage_model->can_temporal_tile_decode();
+    for (;;) {
+        int tile_size_w = static_cast<int>(latents.shape()[0]);
+        int tile_size_h = static_cast<int>(latents.shape()[1]);
+        float tile_overlap;
+        if (tiling_params.enabled &&
+            !first_stage_model->get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params,
+                                               latents.shape()[0], latents.shape()[1])) {
+            return {};
+        }
+        auto decoded = first_stage_model->decode(n_threads, latents, tiling_params, decode_video, circular_x, circular_y);
+        if (!decoded.empty() ||
+            !sd::backend_fit::prepare_vae_decode_retry_tiling(tiling_params, prefer_temporal_tiling,
+                                                              first_stage_model->last_compute_status(),
+                                                              tile_size_w, tile_size_h, first_stage_model->get_scale_factor())) {
+            return decoded;
+        }
     }
-    return decoded;
 }
 
 sd::Tensor<float> StableDiffusionGGML::normalize_ltx_video_latents(const sd::Tensor<float>& x) {

src/pipeline/image.cpp

@@ -35,19 +35,21 @@ namespace sd::pipeline {
             return original_axes;
         }
 
-        int tile_size_x, tile_size_y;
+        int tile_size_w, tile_size_h;
         float overlap;
-        int latent_size_x = request.width / request.vae_scale_factor;
-        int latent_size_y = request.height / request.vae_scale_factor;
-        sd->first_stage_model->get_tile_sizes(tile_size_x,
-                                              tile_size_y,
-                                              overlap,
-                                              sd_img_gen_params->vae_tiling_params,
-                                              latent_size_x,
-                                              latent_size_y);
+        int latent_size_w = request.width / request.vae_scale_factor;
+        int latent_size_h = request.height / request.vae_scale_factor;
+        if (!sd->first_stage_model->get_tile_sizes(tile_size_w,
+                                                   tile_size_h,
+                                                   overlap,
+                                                   sd_img_gen_params->vae_tiling_params,
+                                                   latent_size_w,
+                                                   latent_size_h)) {
+            return original_axes;
+        }
 
-        sd->circular_x = sd->circular_x && (tile_size_x >= latent_size_x);
-        sd->circular_y = sd->circular_y && (tile_size_y >= latent_size_y);
+        sd->circular_x = sd->circular_x && (tile_size_w >= latent_size_w);
+        sd->circular_y = sd->circular_y && (tile_size_h >= latent_size_h);
 
         if (sd->first_stage_model) {
             sd->first_stage_model->set_circular_axes(sd->circular_x, sd->circular_y);
@@ -56,8 +58,8 @@ namespace sd::pipeline {
             sd->preview_vae->set_circular_axes(sd->circular_x, sd->circular_y);
         }
 
-        sd->circular_x = original_axes.circular_x && (tile_size_x < latent_size_x);
-        sd->circular_y = original_axes.circular_y && (tile_size_y < latent_size_y);
+        sd->circular_x = original_axes.circular_x && (tile_size_w < latent_size_w);
+        sd->circular_y = original_axes.circular_y && (tile_size_h < latent_size_h);
 
         return original_axes;
     }
@@ -463,6 +465,11 @@ namespace sd::pipeline {
                 // states with a zeroed prompt mask, so no extra text encode is needed.
                 uncond.c_crossattn = cond.c_crossattn;
                 uncond.c_vector    = sd::Tensor<float>::zeros_like(cond.c_vector);
+            } else if (sd->version == VERSION_MING_IMAGE) {
+                uncond.c_crossattn = sd::Tensor<float>::zeros_like(cond.c_crossattn);
+                for (const auto& extra : cond.extra_c_crossattns) {
+                    uncond.extra_c_crossattns.push_back(sd::Tensor<float>::zeros_like(extra));
+                }
             } else if (sd_version_is_sensenova_u1(sd->version)) {
                 auto* sensenova_conditioner = static_cast<SenseNovaU1Conditioner*>(sd->cond_stage_model.get());
                 uncond                      = sensenova_conditioner->get_unconditional_condition(request->negative_prompt);

src/pipeline/model_builders.cpp

@@ -1,5 +1,6 @@
 #include "model_builders.h"
 
+#include <cmath>
 #include <cstring>
 #include <utility>
 
@@ -17,23 +18,27 @@
 #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"
 #include "model/diffusion/llada_image.hpp"
 #include "model/diffusion/ltxv.hpp"
 #include "model/diffusion/mage_flow.hpp"
+#include "model/diffusion/ming_image.hpp"
 #include "model/diffusion/minimax_h3.hpp"
 #include "model/diffusion/minit2i.hpp"
 #include "model/diffusion/mmdit.hpp"
 #include "model/diffusion/model.hpp"
 #include "model/diffusion/pid.hpp"
+#include "model/diffusion/pixart.hpp"
 #include "model/diffusion/qwen_image.hpp"
 #include "model/diffusion/qwen_image_2_1.hpp"
 #include "model/diffusion/sensenova_u1.h"
 #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"
@@ -301,6 +306,19 @@ namespace sd::model_builders {
                                                                            weight_manager,
                                                                            sd_ctx_params->model_args);
             }
+        } else if (version == VERSION_PIXART) {
+            result.conditioner = std::make_shared<T5CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                  tensor_storage_map,
+                                                                  true,
+                                                                  0,
+                                                                  false,
+                                                                  weight_manager,
+                                                                  sd_ctx_params->model_args);
+            result.diffusion   = std::make_shared<PixArt::PixArtRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                      tensor_storage_map,
+                                                                      "model.diffusion_model",
+                                                                      weight_manager,
+                                                                      sd_ctx_params->model_args);
         } else if (sd_version_is_mage_flow(version)) {
             result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                tensor_storage_map,
@@ -359,6 +377,11 @@ namespace sd::model_builders {
                                                                     tensor_storage_map,
                                                                     "model.diffusion_model",
                                                                     weight_manager);
+        } else if (version == VERSION_MING_IMAGE) {
+            result.conditioner = std::make_shared<MingImageEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                     tensor_storage_map, weight_manager, tokenizers);
+            result.diffusion   = std::make_shared<MingImage::MingImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                            tensor_storage_map, "model.diffusion_model", weight_manager);
         } else if (sd_version_is_z_image(version)) {
             result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                tensor_storage_map,
@@ -372,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,
@@ -422,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++) {
@@ -493,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",
@@ -562,6 +602,23 @@ namespace sd::model_builders {
                                                              false,
                                                              vae_version,
                                                              weight_manager);
+                if (sd_version_is_pixart(version)) {
+                    // Alpha-512 and Sigma share tensor layouts; Alpha-512 needs an explicit scale override.
+                    if (tensor_storage_map.count("model.diffusion_model.csize_embedder.mlp.0.weight") != 0) {
+                        model->scale_factor = 0.18215f;
+                    }
+                    for (const auto& [key, value] : parse_key_value_args(sd_ctx_params->model_args, "model arg")) {
+                        if (key == "pixart_vae_scale_factor") {
+                            float parsed = 0.f;
+                            if (parse_strict_float(value, parsed) && std::isfinite(parsed) && parsed > 0.f) {
+                                model->scale_factor = parsed;
+                            } else {
+                                LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str());
+                            }
+                        }
+                    }
+                    LOG_VERBOSE("pixart: VAE scale factor = %.5f", model->scale_factor);
+                }
                 if (sd_version_is_sdxl(version) &&
                     (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || options.external_vae_is_invalid)) {
                     float vae_conv_2d_scale = 1.f / 32.f;
@@ -575,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/denoiser.hpp

@@ -728,12 +728,20 @@ inline float flux_time_shift(float mu, float sigma, float t) {
 
 // https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py#L289
 struct FluxScheduler : SigmaScheduler {
-    int image_seq_len = 0;
-    float base_shift  = 0.5f;
-    float max_shift   = 1.15f;
+    int image_seq_len      = 0;
+    int base_image_seq_len = 256;
+    int max_image_seq_len  = 4096;
+    float base_shift       = 0.5f;
+    float max_shift        = 1.15f;
+    float shift_terminal   = 0.0f;
 
-    explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr)
+    FluxScheduler(int image_seq_len, SDVersion version, const char* extra_sample_args = nullptr)
         : image_seq_len(image_seq_len) {
+        if (version == VERSION_QWEN_IMAGE_2_1) {
+            max_image_seq_len = 8192;
+            max_shift         = 0.9f;
+            shift_terminal    = 0.02f;
+        }
         parse_extra_sample_args(extra_sample_args);
     }
 
@@ -752,10 +760,8 @@ struct FluxScheduler : SigmaScheduler {
     }
 
     float compute_mu() const {
-        constexpr float base_shift_anchor = 256.0f;
-        constexpr float max_shift_anchor  = 4096.0f;
-        float m                           = (max_shift - base_shift) / (max_shift_anchor - base_shift_anchor);
-        float b                           = base_shift - m * base_shift_anchor;
+        float m = (max_shift - base_shift) / static_cast<float>(max_image_seq_len - base_image_seq_len);
+        float b = base_shift - m * static_cast<float>(base_image_seq_len);
         return static_cast<float>(image_seq_len) * m + b;
     }
 
@@ -764,7 +770,7 @@ struct FluxScheduler : SigmaScheduler {
         sigmas.reserve(n + 1);
 
         float mu = compute_mu();
-        LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu);
+        LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f, shift_terminal=%.3f", image_seq_len, n, mu, shift_terminal);
 
         if (n == 0) {
             sigmas.push_back(1.0f);
@@ -780,6 +786,16 @@ struct FluxScheduler : SigmaScheduler {
             }
         }
 
+        if (shift_terminal > 0.0f && n > 1) {
+            // The terminal shift applies to the last model evaluation, not the final zero sigma.
+            float scale_factor = (1.0f - sigmas[n - 1]) / (1.0f - shift_terminal);
+            if (std::isfinite(scale_factor) && scale_factor > 0.0f) {
+                for (uint32_t i = 0; i < n; ++i) {
+                    sigmas[i] = 1.0f - (1.0f - sigmas[i]) / scale_factor;
+                }
+            }
+        }
+
         sigmas[n] = 0.0f;
         return sigmas;
     }
@@ -1178,7 +1194,7 @@ struct Denoiser {
             }
             case FLUX_SCHEDULER: {
                 LOG_INFO("get_sigmas with Flux scheduler");
-                scheduler = std::make_shared<FluxScheduler>(image_seq_len, extra_sample_args);
+                scheduler = std::make_shared<FluxScheduler>(image_seq_len, version, extra_sample_args);
                 break;
             }
             default:
@@ -1353,6 +1369,42 @@ struct DiscreteFlowDenoiser : public Denoiser {
     }
 };
 
+struct MingImageFlowDenoiser : DiscreteFlowDenoiser {
+    float resolution_shift = std::exp(1.35f);
+
+    MingImageFlowDenoiser()
+        : DiscreteFlowDenoiser(INFINITY) {}
+
+    float sigma_min() override { return 0.f; }
+    float sigma_max() override { return 1.f; }
+
+    float t_to_sigma(float t) override {
+        return time_snr_shift(std::isfinite(shift) ? shift : resolution_shift, (t + 1.f) / 1000.f);
+    }
+
+    std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler, SDVersion version, const char* extra_sample_args = nullptr) override {
+        const float tokens = static_cast<float>(image_seq_len) / 4.f;
+        const float mu     = tokens >= 4096.f ? 1.35f : 0.5f + (tokens - 256.f) * (1.15f - 0.5f) / (4096.f - 256.f);
+        resolution_shift   = std::exp(mu);
+        if (scheduler != DISCRETE_SCHEDULER) {
+            return DiscreteFlowDenoiser::get_sigmas(n, image_seq_len, scheduler, version, extra_sample_args);
+        }
+        if (n == 0) {
+            return {};
+        }
+        const float factor = std::isfinite(shift) ? shift : resolution_shift;
+        std::vector<float> sigmas;
+        sigmas.reserve(n + 1);
+        for (uint32_t i = 0; i < n; ++i) {
+            const float t = n == 1 ? 1.f : 1.f - static_cast<float>(i) / static_cast<float>(n - 1);
+            sigmas.push_back(time_snr_shift(factor, t));
+        }
+        // Upstream sets sigma_min to zero before linspace, then appends the terminal zero.
+        sigmas.push_back(0.f);
+        return sigmas;
+    }
+};
+
 struct H3AVFlowDenoiser : public DiscreteFlowDenoiser {
     int64_t video_channels;
     float audio_shift;
@@ -1767,7 +1819,10 @@ static sd::Tensor<float> sample_euler(denoise_cb_t model,
                                       const std::vector<float>& sigmas) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        float sigma       = sigmas[i];
+        float sigma = sigmas[i];
+        if (sigma == sigmas[i + 1]) {
+            continue;
+        }
         auto denoised_opt = model(x, sigma, i + 1);
         if (denoised_opt.pred.empty()) {
             return {};

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];
             }
         }
     }
@@ -142,8 +140,8 @@ sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
                                    int output_width,
                                    int output_height,
                                    int scale,
-                                   int p_tile_size_x,
-                                   int p_tile_size_y,
+                                   int p_tile_size_w,
+                                   int p_tile_size_h,
                                    float tile_overlap_factor,
                                    bool circular_x,
                                    bool circular_y,
@@ -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_x, 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_y, 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_x * tile_overlap_factor_x);
-    int non_tile_overlap_x = p_tile_size_x - tile_overlap_x;
-    int tile_overlap_y     = static_cast<int32_t>(p_tile_size_y * tile_overlap_factor_y);
-    int non_tile_overlap_y = p_tile_size_y - tile_overlap_y;
-    int tile_size_x        = p_tile_size_x < small_width ? p_tile_size_x : small_width;
-    int tile_size_y        = p_tile_size_y < small_height ? p_tile_size_y : small_height;
-    int input_tile_size_x  = tile_size_x;
-    int input_tile_size_y  = tile_size_y;
-    int output_tile_size_x = tile_size_x;
-    int output_tile_size_y = tile_size_y;
-    if (decode) {
-        output_tile_size_x *= scale;
-        output_tile_size_y *= scale;
-    } else {
-        input_tile_size_x *= scale;
-        input_tile_size_y *= 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_y >= small_height) {
-            int original_y = y;
-            y              = small_height - tile_size_y;
-            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_x >= small_width) {
-                int original_x = x;
-                x              = small_width - tile_size_x;
-                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_x, input_tile_size_y, 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_x && output_tile.shape()[1] == output_tile_size_y);
+            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/runtime/tiling.h

@@ -11,8 +11,8 @@ sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
                                    int output_width,
                                    int output_height,
                                    int scale,
-                                   int p_tile_size_x,
-                                   int p_tile_size_y,
+                                   int p_tile_size_w,
+                                   int p_tile_size_h,
                                    float tile_overlap_factor,
                                    bool circular_x,
                                    bool circular_y,

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

@@ -295,7 +295,13 @@ bool T5UniGramTokenizer::encode(const std::string& input, std::vector<int>& resu
     std::vector<int32_t> tokens;
     std::vector<std::string> token_strs;
     std::string normalized = normalize(input);
-    auto splited_texts     = split_with_special_tokens(normalized, special_tokens);
+    if (normalized.empty()) {
+        // HF reference tokenizers emit no pieces for empty input; pad_tokens
+        // still appends EOS so the sequence becomes [EOS] + padding.
+        result = std::move(tokens);
+        return true;
+    }
+    auto splited_texts = split_with_special_tokens(normalized, special_tokens);
     if (splited_texts.empty()) {
         splited_texts.push_back(normalized);  // for empty string
     }
@@ -337,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);
 };