fix: preserve alpha when upscaling RGBA images with ESRGAN (#2029)

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

CMakeLists.txt

@@ -95,6 +95,8 @@ option(SD_MUSA                       "sd: musa backend" OFF)
 option(SD_BUILD_SHARED_LIBS          "sd: build shared libs" OFF)
 option(SD_BUILD_SHARED_GGML_LIB      "sd: build ggml as a separate shared lib" OFF)
 option(SD_USE_SYSTEM_GGML            "sd: use system-installed GGML library" OFF)
+option(SD_USE_UPSTREAM_GGML          "sd: build with upstream GGML instead of the patched GGML extensions" OFF)
+set(SD_GGML_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/ggml" CACHE PATH "sd: ggml source directory (also supplies private headers for system ggml)")
 #option(SD_BUILD_SERVER               "sd: build server example"                           ON)
 
 set(CMAKE_C_STANDARD 11)
@@ -325,18 +327,7 @@ if (NOT SD_USE_SYSTEM_GGML)
 endif()
 
 # deps
-# Only add ggml if it hasn't been added yet
-if (NOT TARGET ggml)
-    if (SD_USE_SYSTEM_GGML)
-        find_package(ggml REQUIRED)
-        if (NOT ggml_FOUND)
-            message(FATAL_ERROR "System-installed GGML library not found.")
-        endif()
-        add_library(ggml ALIAS ggml::ggml)
-    else()
-        add_subdirectory(ggml)
-    endif()
-endif()
+include(cmake/ggml.cmake)
 
 add_subdirectory(thirdparty)
 

CONTRIBUTING.md

@@ -14,6 +14,8 @@ If you want to update a third-party dependency, please open an issue first inste
 
 Keep each PR focused on one clear change. Large or overly complex PRs are harder to review and may not be merged.
 
+Do not include test code or test scripts in commits or PRs. Keep them local and report verification results in the PR description.
+
 Follow Conventional Commit-style subjects seen in history: `feat:`, `fix:`, `refactor:`, `ci:`, `docs:`, `chore:`. Keep subjects imperative and scoped.
 
 PRs should include:
@@ -35,12 +37,11 @@ Naming conventions:
 - In `PascalCase` names, preserve common abbreviations in uppercase, for example `SD`, `API`, `HTTP`, `JSON`, `RGB`, `VAE`, `TAE`, `LoRA`, and `WebP`.
 - Use `snake_case` for functions, methods, variables, and file names unless an existing API requires a different style.
 - Use a trailing underscore for private data member names, for example `hidden_size_` or `tokenizer_`.
-- Use `.h` for C and C++ header files. Do not introduce new `.hpp` headers.
+- Use `.hpp` for model headers under `src/model/`, including new model headers. Do not rename these headers to `.h`. Use `.h` for other C and C++ header files.
 - Use macro-based header include guards instead of `#pragma once`.
 - Format header include guards as `__SD_{PATH}__`, where `{PATH}` is the header path in uppercase snake case without the file extension. For example, `src/sample.h` should use `__SD_SAMPLE_H__`.
 - Do not introduce anonymous namespaces in new or modified code; prefer `static` file-local functions/variables or an explicit named namespace when scoping is needed.
 - In `class`/`struct` definitions, place data members before member functions unless an existing type already clearly follows a different pattern.
-- Keep `test_*.cpp` / `test_*.py` naming for tests.
 
 Some older code in the project may not fully follow the current conventions. Please do not submit PRs that only rewrite existing code to match style rules.
 

README.md

@@ -15,6 +15,7 @@ API and command-line option may change frequently.***
 
 ## 🔥Important News
 
+* **2026/09/20** 🚀 stable-diffusion.cpp adds **Day-0 support for Qwen-Image-2.1**
 * **2026/08/20** 🚀 stable-diffusion.cpp now supports **LTX-2.5**
 * **2026/08/04** 🚀 stable-diffusion.cpp adds **Day-1 support for MiniMax-H3**
 * **2026/06/25** 🚀 stable-diffusion.cpp now supports **Krea2**
@@ -47,6 +48,7 @@ API and command-line option may change frequently.***
     - [Chroma](./docs/chroma.md)
     - [Chroma1-Radiance](./docs/chroma_radiance.md)
     - [Qwen Image](./docs/qwen_image.md)
+    - [Qwen Image 2.1](./docs/qwen_image_2.1.md)
     - [PiD](./docs/pid.md)
     - [LongCat Image](./docs/longcat_image.md)
     - [Z-Image](./docs/z_image.md)
@@ -61,12 +63,14 @@ API and command-line option may change frequently.***
     - [SeFi-Image](./docs/sefi_image.md)
     - [HiDream-O1-Image](./docs/hidream_o1_image.md)
     - [Ideogram4](./docs/ideogram4.md)
+    - [LLaDA-Image](./docs/llada_image.md)
   - [Image Edit Models](./docs/edit.md)
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)
     - [LongCat Image Edit](./docs/longcat_image.md)
     - [Boogu Image Edit](./docs/boogu_image.md)
     - [Mage-Flow-Edit](./docs/mage_flow.md#image-editing)
+    - [LLaDA-Image Edit](./docs/llada_image.md#image-editing)
   - Video Models
     - [Wan2.1/Wan2.2](./docs/wan.md)
     - [MiniMax-H3](./docs/minimax_h3.md)

cmake/ggml.cmake

@@ -0,0 +1,27 @@
+if(NOT TARGET ggml AND NOT TARGET ggml::ggml)
+    if(SD_USE_SYSTEM_GGML)
+        find_package(ggml REQUIRED)
+    else()
+        add_subdirectory("${SD_GGML_SOURCE_DIR}" "${CMAKE_CURRENT_BINARY_DIR}/ggml")
+    endif()
+endif()
+if(NOT TARGET ggml)
+    add_library(ggml ALIAS ggml::ggml)
+endif()
+
+get_target_property(sd_ggml_imported ggml IMPORTED)
+if(sd_ggml_imported)
+    set(sd_ggml_private_include "${SD_GGML_SOURCE_DIR}/src")
+else()
+    get_target_property(sd_ggml_private_include ggml SOURCE_DIR)
+endif()
+if(NOT EXISTS "${sd_ggml_private_include}/ggml-impl.h")
+    message(FATAL_ERROR "Set SD_GGML_SOURCE_DIR to the source tree matching the selected ggml library (ggml-impl.h is required).")
+endif()
+target_include_directories(${SD_LIB} PRIVATE "${sd_ggml_private_include}")
+set_property(TARGET ${SD_LIB} PROPERTY SD_GGML_PRIVATE_INCLUDE_DIR "${sd_ggml_private_include}")
+
+if(SD_USE_UPSTREAM_GGML)
+    target_compile_definitions(${SD_LIB} PUBLIC SD_USE_UPSTREAM_GGML)
+    message(WARNING "Using upstream GGML: INT8 tensorwise/convrot is disabled and FP8 weights are converted to F16 at load time. Some operators may be unsupported and performance may be lower than with patched GGML.")
+endif()

docs/backend.md

@@ -156,8 +156,11 @@ the runner's graph-cut capacity checks.
 
 Runtime capacity checks also leave 512 MiB of currently free device memory for
 backend scratch buffers and pipelines, including with explicit backend assignments.
-They cap stale free-memory reports by the device's total memory minus tracked
-resident allocations and reject reports that exceed the device's total memory.
+They cap free-memory reports by the device's total memory minus tracked
+resident allocations. Vulkan reports exceeding total memory are rejected because
+its heap-budget subtraction can underflow. Other backends use the cap instead of
+treating such reports as zero free memory. Failed checks log the reported free and
+total memory alongside tracked weight and runtime allocations.
 
 Components are considered in `diffusion`, `te`, `vae` order so that repeatedly
 used diffusion weights have priority. Each component's weights use the first

docs/build.md

@@ -16,6 +16,40 @@ git submodule init
 git submodule update
 ```
 
+## Selecting a GGML source tree
+
+By default, sd.cpp builds the patched GGML submodule in `ggml/`. To build with
+an upstream GGML checkout instead, enable `SD_USE_UPSTREAM_GGML` and set
+`SD_GGML_SOURCE_DIR`:
+
+```shell
+cmake -S . -B build-upstream -DSD_USE_UPSTREAM_GGML=ON -DSD_GGML_SOURCE_DIR=../ggml-upstream
+cmake --build build-upstream --config Release
+```
+
+The selected source tree supplies both the library and its private headers.
+Backend options such as `-DSD_CUDA=ON` apply to the selected tree as usual.
+
+`SD_USE_UPSTREAM_GGML` defaults to `OFF`, which enables the patched GGML
+extensions. Set it to `ON` when using upstream GGML; it selects the compatibility
+mode and does not download or replace the GGML source tree. Upstream mode keeps
+the original FP8 safetensors handling: FP8 tensors are converted to F16 at load
+time (one byte per element in the file, two in RAM and VRAM). INT8
+tensorwise/convrot is disabled and its model files are rejected with an explicit
+error. FP8 GGUF files, FP8 weight type requests and tensor type rules are also
+rejected; no automatic conversion is performed.
+
+Upstream GGML may lack some operators and performance optimizations provided by
+the patched version. A warning is emitted during CMake configuration and when
+creating an inference context. Ordinary floating-point and shared GGML
+quantization types remain available, subject to backend operator support.
+
+`SD_USE_SYSTEM_GGML=ON` instead links an installed GGML CMake package, located
+with `ggml_DIR` or `CMAKE_PREFIX_PATH`. In that mode, `SD_GGML_SOURCE_DIR` must
+point to the matching source tree for private headers. The installed library
+must use the same ABI settings as sd.cpp, including `GGML_MAX_NAME`.
+Set `SD_USE_UPSTREAM_GGML=ON` as well if the installed package is upstream GGML.
+
 ## WebP and WebM Support in Examples
 
 The example applications (`examples/cli` and `examples/server`) use `libwebp` to support WebP image I/O, and `examples/cli` can also use `libwebm` for `.webm` video output. Both are enabled by default. WebM output currently reuses `libwebp` to encode each frame as VP8 before muxing with `libwebm`.
@@ -56,6 +90,10 @@ cmake --build . --config Release
 
 ## Build with CUDA
 
+Native SageAttention is included when using CUDA with patched GGML
+(`SD_USE_UPSTREAM_GGML=OFF`).
+See [SageAttention](sage_attention.md) for GPU requirements and `--sage-attn` usage.
+
 This provides GPU acceleration using NVIDIA GPU. Make sure to have the CUDA toolkit installed. You can download it from your Linux distro's package manager (e.g. `apt install nvidia-cuda-toolkit`) or from here: [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads). Recommended to have at least 4 GB of VRAM.
 
 ```shell

docs/edit.md

@@ -17,6 +17,7 @@ Depending on the architecture, different models handle reference images differen
 | [**Boogu Image Edit**](./boogu_image.md) | `z_image_omni` |
 | **Krea2 (Community Edit LoRAs)** | `krea2_ostris_edit` |
 | [**Mage-Flow-Edit**](./mage_flow.md#image-editing) | `mage_flow` |
+| [**LLaDA-Image**](./llada_image.md#image-editing) | `llada_image` |
 | **Anima (Community Edit LoRAs)** | `cosmos_reference` |
 
 Stable-diffusion.spp also supports basic Unet-based editing models like instruct-pix2pix or CosXL-Edit. This document is not about those.
@@ -25,6 +26,9 @@ Stable-diffusion.spp also supports basic Unet-based editing models like instruct
 
 ## Configuring Reference Modes (`--ref-image-args`)
 
+For a one-time input transform before reference presets and model processing,
+including cropping, padding, and resizing algorithms, see [Image preprocessing](./image_preprocessing.md).
+
 Different DiT-based editing models require different configurations to process reference images correctly (e.g., whether to use a Vision Language Model (VLM) encoder or pass VAE-encoded images directly to the DiT).
 
 To simplify this, we provide **Presets**. By default, the system automatically selects the best preset based on the model architecture. However, you can override this using the `--ref-image-args` argument.

docs/esrgan.md

@@ -2,6 +2,8 @@
 
 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.
+
 - Specify the model path using the `--upscale-model PATH` parameter. example:
 
 ```bash

docs/ideogram4.md

@@ -18,6 +18,9 @@ at one byte per element in RAM and VRAM. Backends that cannot multiply FP8
 weights directly cast only the active layer to a temporary BF16 tensor during
 execution; the loader does not expand the entire checkpoint to BF16.
 
+With `SD_USE_UPSTREAM_GGML=ON`, FP8 tensors are converted to F16 at load time
+instead (two bytes per element in RAM and VRAM).
+
 Use `ideogram4_fp8.safetensors` and `ideogram4_uncond_fp8.safetensors` directly
 with `--diffusion-model` and `--uncond-diffusion-model`, respectively.
 

docs/image_preprocessing.md

@@ -0,0 +1,173 @@
+# Image preprocessing
+
+Use `--image-preprocess` to transform each image input once, before generation:
+
+```sh
+sd-cli ... \
+  --image-preprocess "target=init,mode=crop-resize,filter=lanczos,antialias=true" \
+  --image-preprocess "target=mask,filter=nearest-exact" \
+  --image-preprocess "target=ref,index=0,mode=fit-pad,width=768,height=768,filter=bicubic"
+```
+
+CLI and server image loaders decode at the original resolution. The generation
+entry point merges input defaults with user rules and prepares one transformed
+image per input. The original pipeline then consumes those images, including
+its mandatory canvas adaptation, reference resizing, and encoder preprocessing.
+
+```text
+native-resolution image
+  -> input defaults + user overrides
+  -> one input transform
+  -> original generation pipeline and model-specific processing
+```
+
+These rules do not override internal VAE, CLIP/VLM, ControlNet, or pixel-patch preprocessing.
+`--ref-image-args` retains its existing meaning and runs after this input transform.
+
+## Inputs and defaults
+
+| `target` | Input | Default geometry | Indexed? |
+| --- | --- | --- | --- |
+| `init` | img2img image or video first frame | Center crop to the generation aspect ratio, then resize | No |
+| `end` | Video last frame | Center crop, then resize | No |
+| `mask` | Inpainting mask | Inherit init geometry; otherwise center crop, then resize | No |
+| `control` | Control image | Center crop, then resize | No |
+| `ref` | Reference images | Preserve source dimensions | Yes |
+| `ip-adapter` | IP-Adapter image | Preserve source dimensions | No |
+| `id` | PhotoMaker identity images | Preserve source dimensions | Yes |
+| `control-frame` | Control video frames | Center crop, then resize | Yes |
+
+Canvas defaults use the aligned generation dimensions. Reference, IP-Adapter,
+and identity inputs use their original dimensions unless overridden. Default
+resampling is nearest for images and nearest-exact for masks.
+
+These defaults are shared by CLI, server, and C API. Moving geometry out of
+the loaders replaces the previous CLI/server BOX/sRGB resizing, so default
+pixels are not guaranteed to match earlier builds.
+
+Reference video and audio preprocessing are outside these image rules.
+Preprocessing options apply to `img_gen` and `vid_gen`, not standalone upscale
+or ADetailer mode. ADetailer clears the user's rules for its internal crops.
+
+## Rules
+
+Rules are comma-separated `key=value` lists. Repeat the CLI option or separate
+rules with semicolons. Every rule requires a `target` and at least one option.
+Rule syntax and input compatibility are checked when image/video generation
+starts. Unknown keys, invalid values, duplicate keys in a rule, missing images,
+and out-of-range indices cause generation to fail with an error log.
+
+Omit `index` to configure every image of that type; otherwise use a zero-based
+index. CLI directory inputs follow filename order. Indexed rules override
+type-wide rules field by field, regardless of order. At equal specificity,
+the last value for a field wins. `auto` selects the input preset.
+
+| `mode` | Input transform |
+| --- | --- |
+| `auto` | Use the input's default geometry |
+| `none` | Keep source dimensions without resizing, cropping, or padding |
+| `stretch` | Resize to the target dimensions |
+| `crop` | Crop a target-sized rectangle without resizing; fail if the source is too small |
+| `crop-resize` | Crop to the target aspect ratio, then resize |
+| `fit-pad` | Fit the entire image inside the target dimensions, preserving aspect ratio, then pad |
+
+`width` and `height` must be specified together as positive integers. They
+override the input transform's dimensions, not the generation or encoder size.
+For a native-size preset, specifying dimensions without a mode selects stretch.
+`mode=none` with explicit dimensions different from the source is contradictory
+and is rejected.
+
+`anchor=center|top|bottom|left|right` selects crop/padding placement.
+`pad_color=#RRGGBB` or `#RRGGBBAA` selects padding, defaulting to opaque black.
+A grayscale mask uses the first color component.
+
+`filter=auto|nearest|nearest-exact|bilinear|bicubic|lanczos` selects resampling.
+`antialias=auto|true|false` enables antialiasing automatically for filtered
+downscaling; explicit true requires bilinear, bicubic, or Lanczos.
+Filtered RGBA resizing uses premultiplied alpha.
+
+`canny=true|false` enables edge detection for any supported image target,
+defaulting to `false`. It runs once after geometry, before the original
+generation pipeline, including with `mode=none`. Grayscale, grayscale-alpha,
+RGB, and RGBA inputs are supported; alpha is preserved.
+
+Each input has its own Canny setting. Indexed rules can enable or disable it
+for individual references, identity images, or video control frames.
+
+```sh
+--image-preprocess "target=init,mode=fit-pad,canny=true"
+--image-preprocess "target=ref,index=0,mode=none,canny=true"
+--image-preprocess "target=control-frame,index=2,canny=true"
+```
+
+Init and mask sources must have the same dimensions. The mask inherits the
+init crop, resize, and padding coordinates, while retaining its own filter,
+padding value, and Canny setting. Conflicting mask geometry is rejected. An
+omitted mask remains absent until the original pipeline creates its default mask.
+
+## Downstream behavior
+
+`mode=none` only skips the input geometry transform. For example:
+
+```sh
+--image-preprocess "target=init,mode=none" \
+--image-preprocess "target=ref,mode=none"
+```
+
+The init image is still adapted to the generation canvas by the original
+pipeline. Reference images still follow `--ref-image-args` and model-specific
+resizing. CLIP retains its fixed input dimensions and normalization. HiDream-O1
+retains its original pixel-reference and visual preprocessing.
+
+Existing sharing between consumers is preserved: for example, Wan img2video
+uses the same adapted first frame for VAE conditioning and CLIP. High-resolution
+passes reuse the prepared images and apply their original size adaptation;
+they do not apply the user's crop a second time.
+
+To disable reference resizing before VAE encoding, use
+`--ref-image-args "resize_before_vae=false"` or the server field
+`"ref_image_args": "resize_before_vae=false"`. This is separate from
+`target=ref,mode=none`, which only skips input geometry. Model constraints
+still apply.
+
+## Server requests
+
+Native image/video requests and SDAPI accept `image_preprocess` as a string or
+an array of rule strings:
+
+```json
+{
+  "image_preprocess": [
+    "target=init,mode=fit-pad,filter=bicubic",
+    "target=mask,filter=nearest-exact",
+    "target=ref,index=0,mode=none"
+  ]
+}
+```
+
+OpenAI-compatible requests accept it through
+`<sd_cpp_extra_args>{...}</sd_cpp_extra_args>` in the prompt.
+Request rules replace server-default rules. Generation metadata records the
+user rules; image encodings and channel conventions are unchanged.
+
+## C API
+
+Set `image_preprocess` on the existing image/video generation parameters.
+The `generate_image()` and `generate_video()` signatures are unchanged:
+
+```c
+sd_img_gen_params_t params;
+sd_img_gen_params_init(&params);
+/* Set prompt, original-resolution input images, and generation options. */
+params.image_preprocess.rules = "target=init,mode=crop-resize,filter=lanczos;"
+                                "target=mask,filter=nearest-exact";
+bool ok = generate_image(ctx, &params, &images, &count);
+```
+
+Both generation parameter initializers set `image_preprocess.rules` to `NULL`,
+selecting input presets. Rule strings are borrowed for the synchronous call.
+The library owns temporary transformed pixels; caller images and arrays are
+not modified. Add `canny=true` to the desired target's rule in
+`image_preprocess.rules` to enable Canny.
+
+The parameter structs have grown; applications and bindings must be rebuilt.

docs/int8_convrot.md

@@ -2,6 +2,9 @@
 
 sd.cpp can load and execute ComfyUI `int8_tensorwise` safetensors with `convrot` metadata directly. The stored INT8 weights are not converted to another weight type at load time.
 
+This requires the INT8 tensorwise/convrot extensions in the patched GGML.
+Builds with `SD_USE_UPSTREAM_GGML=ON` reject these files during loading.
+
 ## Checkpoint format
 
 Each quantized linear module contains the following tensors:

docs/llada_image.md

@@ -0,0 +1,156 @@
+# How to Use
+
+LLaDA-Image is a 6B text-to-image and instruction-guided editing model. The denoiser is a
+Lumina2/Z-Image-style NextDiT conditioned by a LLaDA2-MoE diffusion-LLM text encoder, and it
+reuses the Flux.2 VAE. Two checkpoints are published: a 50-step base model and
+LLaDA-Image-Turbo, a 4-step distilled model.
+
+## Download weights
+
+Four components are required: a transformer, a text encoder, a VAE, and a connectors file
+holding the QueryFormer, the text projection and, for editing, the SigVQ image encoder.
+
+The two published checkpoints are **not** interchangeable. LLaDA-Image-Turbo and LLaDA-Image
+ship different transformers, text encoders, QueryFormers and text projections; only the VAE,
+the SigVQ encoder and the tokenizer are shared. Mixing the two produces degraded output rather
+than a clean error, so keep each checkpoint's files together.
+
+Both need an external LLaDA2 `tokenizer.json`, which is not embedded in sd.cpp and is the same
+file for either checkpoint. Take `tokenizer/tokenizer.json` from either repository and pass it
+with `--tokenizer`. See [JSON tokenizers](tokenizers.md) for CLI and C API usage.
+
+### LLaDA-Image-Turbo (4 steps)
+
+Converted transformer, text encoder and pre-merged connectors are at
+https://huggingface.co/fszontagh/LLaDA-Image-Turbo-GGUF:
+
+- `llada-image-turbo-f16.gguf`
+- `llada-image-turbo-text_encoder-q8_0.gguf`
+- `llada-image-turbo-connectors.safetensors` for text to image, or
+  `llada-image-turbo-connectors-edit.safetensors`, which also carries the SigVQ encoder that
+  editing needs.
+
+Other quantizations of the transformer and the text encoder are in the same repository.
+
+The VAE comes from the original repository,
+https://huggingface.co/inclusionAI/LLaDA-Image-Turbo: `vae/diffusion_pytorch_model.safetensors`,
+referred to below as `llada_vae.safetensors`.
+
+### LLaDA-Image (50 steps)
+
+Converted transformer, text encoder and pre-merged connectors are at
+https://huggingface.co/fszontagh/LLaDA-Image-GGUF:
+
+- `llada-image-f16.gguf`
+- `llada-image-text_encoder-q8_0.gguf`
+- `llada-image-connectors.safetensors` for text to image, or
+  `llada-image-connectors-edit.safetensors`, which also carries the SigVQ encoder that editing
+  needs.
+
+Other quantizations of the transformer and the text encoder are in the same repository.
+
+The VAE comes from the original repository,
+https://huggingface.co/inclusionAI/LLaDA-Image, and is the same file as the Turbo one.
+
+### Converting the weights yourself
+
+The transformer has to go in through `--diffusion-model` so that its tensor names keep the
+prefix the loader expects, while the text encoder goes in through `-m`:
+
+```bash
+./bin/sd-cli -M convert --diffusion-model transformer/diffusion_pytorch_model.safetensors.index.json \
+  -o llada-image-f16.gguf --type f16
+./bin/sd-cli -M convert -m text_encoder/model.safetensors.index.json \
+  -o llada-image-text_encoder-q8_0.gguf --type q8_0
+```
+
+### Building the connector file yourself
+
+`--embeddings-connectors` takes one file, so the QueryFormer, the text projection and
+(for editing) the SigVQ encoder have to be combined into a single Safetensors file, each
+tensor name prefixed with its component name. Leaving `sigvq` out skips loading the 2.6 GB
+encoder:
+
+```python
+from safetensors.torch import load_file, save_file
+
+merged = {}
+for prefix, path in [
+    ("queryformer", "queryformer/diffusion_pytorch_model.safetensors"),
+    ("text_projection", "text_projection/diffusion_pytorch_model.safetensors"),
+    ("sigvq", "sigvq/diffusion_pytorch_model.safetensors"),
+]:
+    for name, tensor in load_file(path).items():
+        merged[f"{prefix}.{name}"] = tensor
+save_file(merged, "llada_connectors.safetensors")
+```
+
+## Examples
+
+### Text to image
+
+```bash
+./bin/sd-cli \
+  --diffusion-model /path/to/llada-image-turbo-f16.gguf \
+  --llm /path/to/llada-image-turbo-text_encoder-q8_0.gguf \
+  --tokenizer /path/to/tokenizer.json \
+  --vae /path/to/llada_vae.safetensors \
+  --embeddings-connectors /path/to/llada-image-turbo-connectors.safetensors \
+  --prompt "a lovely cat holding a sign says 'llada.cpp'" \
+  --width 1024 \
+  --height 1024 \
+  --steps 4 \
+  --cfg-scale 1.0 \
+  --seed 42 \
+  --output output.png
+```
+
+<img width="256" alt="LLaDA-Image example" src="../assets/llada_image/example.png" />
+
+### Image editing
+
+```bash
+./bin/sd-cli \
+  --diffusion-model /path/to/llada-image-turbo-f16.gguf \
+  --llm /path/to/llada-image-turbo-text_encoder-q8_0.gguf \
+  --tokenizer /path/to/tokenizer.json \
+  --vae /path/to/llada_vae.safetensors \
+  --embeddings-connectors /path/to/llada-image-turbo-connectors-edit.safetensors \
+  --ref-image /path/to/input.png \
+  --prompt "change the sign text to 'sd.cpp'" \
+  --width 1024 \
+  --height 1024 \
+  --steps 4 \
+  --cfg-scale 1.0 \
+  --diffusion-fa \
+  --output output.png
+```
+
+<img width="256" alt="LLaDA-Image edit example" src="../assets/llada_image/edit_example.png" />
+
+See [edit.md](./edit.md) for the shared reference-image options. LLaDA-Image uses the
+`llada_image` preset by default, resizing the reference image to the output width and height
+before VAE encoding. SigVQ uses bilinear resizing to half the output resolution and inputs
+normalized to `[-1, 1]`. CFG keeps the source latent in both branches and uses SigVQ features
+only in the positive branch. Editing requires connectors that include the SigVQ weights.
+
+## Notes
+
+- Use 4 steps and `--cfg-scale 1.0` for LLaDA-Image-Turbo; the guidance is distilled away, so
+  a higher CFG degrades output and doubles the text encoder cost. The 50-step base model uses
+  `--steps 50 --cfg-scale 5`.
+- Width and height are rounded up to a multiple of 16. For editing the reference pipeline
+  requires them to be divisible by 32.
+- Edit the 50-step base model at 1024x1024. At 512x512 it returns the reference image almost
+  unchanged instead of applying the instruction; LLaDA-Image-Turbo edits correctly at both.
+- Editing runs the reference and the target in one sequence, so it needs roughly twice the
+  tokens of text to image at the same size. On 12 GB, editing at 1024x1024 needs
+  `--diffusion-fa`; without it the diffusion graph does not fit.
+- The weights total about 16 GB, but segmented execution streams them, so a much smaller
+  budget works. At 512x512, `--max-vram 6` costs almost nothing over unconstrained execution,
+  and `--max-vram 3` still produces byte-identical output at roughly 2.5x the time.
+- `--scheduler` defaults to `llada_image`, which reproduces the reference Kumaraswamy sigma
+  grid. `--extra-sample-args uniform=1` selects the uniform grid instead.
+- Prompt templating is handled automatically; pass a plain description.
+- VQ-conditioned generation (`generation_mode="vq"`, where the text encoder decodes image
+  tokens before diffusion) is not implemented.

docs/qwen_image_2.1.md

@@ -0,0 +1,59 @@
+# How to Use
+
+Qwen Image 2.1 supports text-to-image generation and image editing, using Qwen3-VL-8B as the text encoder and its own VAE.
+
+## Download weights
+
+- Download Qwen Image 2.1
+    - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image-2.1/tree/main/diffusion_models
+    - gguf: https://huggingface.co/leejet/Qwen-Image-2.1-GGUF/tree/main
+- Download vae
+    - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image-2.1/tree/main/vae
+- Download Qwen3-VL-8B-Instruct
+    - safetensors (BF16 or INT8 convrot): https://huggingface.co/Comfy-Org/Qwen-Image-2.1/tree/main/text_encoders
+    - gguf: https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct-GGUF/tree/main
+    - For image editing with a GGUF text encoder, also download `mmproj-Qwen3VL-8B-Instruct-F16.gguf` from the same repository and pass it with `--llm_vision`.
+
+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.
+
+## Examples
+
+Run the following commands from the build directory. Use image dimensions divisible by 32. The resolution-dependent flow schedule is selected automatically.
+
+### 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
+```
+
+<img alt="Qwen Image 2.1 example" src="../assets/qwen/qwen_image_2.1.png" />
+
+To use GGUF diffusion weights, set `--diffusion-model` to the path of a file such as `qwen_image_2.1-Q4_K.gguf`.
+
+### Image editing
+
+Pass the reference image with `-r` and describe the edit in `-p`. Vision weights are required; the example below loads them separately with `--llm_vision`.
+
+```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
+```
+
+For multiple reference images, repeat `-r` in the desired order, for example `-r first.png -r second.png`.
+
+### Alpha channel
+
+This model supports alpha channel output. As the model determines whether to output a regular image or with transparency through the prompt, according to [official recommendation](https://github.com/QwenLM/Qwen-Image-2.1#transparent-image-generation-rgba), use the following prompt format for better results:
+
+> `This is an RGBA image with transparency. <your description>. The image has alpha channel and the background is transparent.`
+
+Since transparency is decided by the prompt rather than by the input or an explicit switch, the same format applies equally to editing, whether or not the reference image itself has an alpha channel. Note that alpha is kept only in `.png` and `.webp` outputs; saving as `.jpg` drops the transparency.
+
+Here are some examples ran with Q6_K quantization:
+| Input | Prompt | Output |
+| --- | --- | --- |
+| ![Qwen Image 2.1 alpha input example 1](../assets/qwen/qwen-image-2.1-alpha-in1.png) | This is an RGBA image with transparency. Replace the text "BLOOM" with "Qwen Image 2.1", keeping the same font of the original text. The image has alpha channel and the background is transparent. | ![Qwen Image 2.1 alpha output example 1](../assets/qwen/qwen-image-2.1-alpha-out1.png) |
+| ![Qwen Image 2.1 alpha input example 2](../assets/logo.png) | This is an RGBA image with transparency. Remove the background of the image, keeping only the text and cat. The image has alpha channel and the background is transparent. | ![Qwen Image 2.1 alpha output example 2](../assets/qwen/qwen-image-2.1-alpha-out2.png) |
+
+### Other features
+
+Other features of the model could be found on the [model card from QwenLM/Qwen-Image-2.1 repo](https://github.com/QwenLM/Qwen-Image-2.1), including 2 finetuned prompt rewriting Qwen3.5-9B model.

docs/sage_attention.md

@@ -0,0 +1,68 @@
+# SageAttention
+
+`--sage-attn` enables native CUDA SageAttention in the diffusion model, including
+the high-noise diffusion model when present. Python, PyTorch, and Triton are not
+required at build time or runtime.
+
+The CUDA backend automatically selects a kernel supported by both the GPU and
+the compiled CUDA toolkit:
+
+| GPU / toolkit | Implementation |
+| --- | --- |
+| SM89 or newer, CUDA 12.8 or newer (except SM90) | SageAttention2++: per-thread INT8 Q/K, FP8 PV, FP16 instruction accumulation with an FP32 buffer |
+| SM89 or newer, CUDA 12.4 or newer; SM90 also uses this path with newer toolkits | SageAttention2: per-thread INT8 Q/K, FP8 PV, two-level FP32 accumulation |
+| SM80 or newer, CUDA 12.0 or newer | INT8 Q/K, FP16 PV compatibility path |
+
+The FP8 paths smooth K, quantize V per channel, and pad and permute V for FP8
+Tensor Cores. The 2++ path uses the upstream V scale limit of 2.25 to avoid
+overflow in its FP16 instruction accumulator. The public output remains FP32.
+These are the upstream **INT8** SageAttention2/2++ variants; the paper's INT4
+variant and Hopper-specific WGMMA kernel are not implemented here.
+
+## Build
+
+Use the bundled patched GGML, CUDA Toolkit 12.0 or newer, and an NVIDIA GPU with
+compute capability 8.0 or newer. Compile kernels for the GPU being used.
+
+```sh
+cmake -S . -B build -DSD_CUDA=ON -DSD_USE_UPSTREAM_GGML=OFF
+cmake --build build --config Release
+```
+
+No separate SageAttention build option is needed. Upstream GGML builds do not
+support it. A system GGML must include the matching patched API and CUDA
+backend. Enabling `--sage-attn` with an unavailable build or diffusion device
+reports an error. Building with CUDA 12.4 selects SageAttention2 on an RTX 4090;
+rebuild with CUDA 12.8 or newer to use SageAttention2++.
+
+## Use
+
+Replace `--diffusion-fa` with `--sage-attn` in an existing command. For example,
+from the build directory:
+
+```powershell
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model  ..\models\diffusion_models\Wan2.2-T2V-A14B-LowNoise-Q8_0.gguf --high-noise-diffusion-model  ..\models\diffusion_models\Wan2.2-T2V-A14B-HighNoise-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 --offload-to-cpu --video-frames 33 --sage-attn
+```
+
+SageAttention currently handles unmasked attention with head dimensions from
+1 through 128, including grouped-query attention, different query/key lengths,
+and multiple batches. Dimensions below 64 are zero-padded to 64; dimensions
+between 65 and 127 are zero-padded to 128. The original softmax scale is preserved,
+and the output is cropped back to the original dimension. Other attention
+operations fall back to FlashAttention when supported, then ordinary attention.
+SageAttention takes precedence in diffusion
+when combined with `--fa` or `--diffusion-fa`; `--fa` continues to control other
+modules. Existing attention scaling overrides remain effective.
+
+Attention quantization changes numerical results. Compare image quality and
+end-to-end generation time using the same seed, dimensions, and sampling
+settings. Compare sampling steps after the first step for warmed-up inference
+speed, and report model loading and first-step initialization separately.
+Quantization, smoothing, and format conversion costs are included in generation
+time, so short sequences may not benefit.
+
+Library callers set `sd_ctx_params_t.sage_attn = true` before `new_sd_ctx()`,
+like `diffusion_flash_attn`. Context creation fails if the requested feature is
+unavailable. Initialize the parameter structure with `sd_ctx_params_init()`.
+Rebuild library callers against the updated public header.

docs/troubleshooting.md

@@ -1,5 +1,14 @@
 # Troubleshooting
 
+## Video model used in image generation mode
+
+If generation reports that a model cannot be run with `generate_image()`, add
+`--mode vid_gen` to the CLI command. `--video-frames` alone does not select video
+mode. Video models require this mode even when generating a single frame.
+Library callers must use `generate_video()` for these models; use
+`sd_ctx_supports_image_generation()` and `sd_ctx_supports_video_generation()` to
+check the available generation modes.
+
 ## Completely black or white images or videos / NaNs
 
 Some ggml backends can encounter numerical overflow during inference, producing

docs/wan.md

@@ -1,5 +1,7 @@
 # How to Use
 
+Wan models require `-M vid_gen`, including single-frame generation. `--video-frames` alone does not select video mode. Library callers must use `generate_video()` instead of `generate_image()`.
+
 ## Download weights
 
 - Download Wan

examples/cli/README.md

@@ -14,6 +14,12 @@ equivalent to `--log-level verbose`. If repeated, the last logging option wins.
 For direct image repair or automatic post-generation YOLOv8 detection followed by cropped inpainting, see
 [ADetailer](../../docs/adetailer.md).
 
+Use repeatable `--image-preprocess` rules to select resizing, cropping, padding,
+and resampling separately for each image input. Add `canny=true` to any input
+rule for edge detection. See
+[Image preprocessing](../../docs/image_preprocessing.md) for input selectors,
+input defaults, downstream model processing, mask alignment, and examples.
+
 Metadata mode inspects PNG/JPEG container metadata without loading any model:
 
 ```bash

examples/cli/main.cpp

@@ -41,7 +41,6 @@ struct SDCliParams {
     std::string metadata_format = "text";
 
     sd_log_level_t log_level = SD_LOG_INFO;
-    bool canny_preprocess    = false;
     bool convert_name        = false;
 
     preview_t preview_method = PREVIEW_NONE;
@@ -107,10 +106,6 @@ struct SDCliParams {
         };
 
         options.bool_options = {
-            {"",
-             "--canny",
-             "apply canny preprocessor (edge detection)",
-             true, &canny_preprocess},
             {"",
              "--convert-name",
              "convert tensor name (for convert mode)",
@@ -268,7 +263,6 @@ struct SDCliParams {
             << "  metadata_format: \"" << metadata_format << "\",\n"
             << "  log_level: " << log_level_name(log_level) << ",\n"
             << "  color: " << (color ? "true" : "false") << ",\n"
-            << "  canny_preprocess: " << (canny_preprocess ? "true" : "false") << ",\n"
             << "  convert_name: " << (convert_name ? "true" : "false") << ",\n"
             << "  preview_method: " << previews_str[preview_method] << ",\n"
             << "  preview_interval: " << preview_interval << ",\n"
@@ -328,9 +322,7 @@ void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
 
 bool load_images_from_dir(const std::string dir,
                           std::vector<SDImageOwner>& images,
-                          int expected_width  = 0,
-                          int expected_height = 0,
-                          int max_image_num   = 0) {
+                          int max_image_num = 0) {
     if (!fs::exists(dir) || !fs::is_directory(dir)) {
         LOG_ERROR("'%s' is not a valid directory\n", dir.c_str());
         return false;
@@ -357,7 +349,8 @@ bool load_images_from_dir(const std::string dir,
             LOG_VERBOSE("load image %zu from '%s'", images.size(), path.c_str());
             int width             = 0;
             int height            = 0;
-            uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, expected_width, expected_height);
+            int loaded_channel    = 0;
+            uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, loaded_channel, 0, 0);
             if (image_buffer == nullptr) {
                 LOG_ERROR("load image from '%s' failed", path.c_str());
                 return false;
@@ -365,7 +358,7 @@ bool load_images_from_dir(const std::string dir,
 
             images.emplace_back(sd_image_t{(uint32_t)width,
                                            (uint32_t)height,
-                                           3,
+                                           (uint32_t)loaded_channel,
                                            image_buffer});
 
             if (max_image_num > 0 && static_cast<int>(images.size()) >= max_image_num) {
@@ -653,8 +646,8 @@ int main(int argc, const char* argv[]) {
     SDContextParams ctx_params;
     SDGenerationParams gen_params;
 
-    parse_args(argc, argv, cli_params, ctx_params, gen_params);
     sd_set_log_callback(sd_log_cb, (void*)&cli_params);
+    parse_args(argc, argv, cli_params, ctx_params, gen_params);
 
     if (cli_params.mode == METADATA) {
         MetadataReadOptions options;
@@ -750,16 +743,8 @@ int main(int argc, const char* argv[]) {
 
     auto load_image_and_update_size = [&](const std::string& path,
                                           SDImageOwner& image,
-                                          bool resize_image    = true,
                                           int expected_channel = 3) -> bool {
-        int expected_width  = 0;
-        int expected_height = 0;
-        if (resize_image && gen_params.width_and_height_are_set()) {
-            expected_width  = gen_params.width;
-            expected_height = gen_params.height;
-        }
-
-        if (!load_sd_image_from_file(image.put(), path.c_str(), expected_width, expected_height, expected_channel)) {
+        if (!load_sd_image_from_file(image.put(), path.c_str(), 0, 0, expected_channel)) {
             LOG_ERROR("load image from '%s' failed", path.c_str());
             return false;
         }
@@ -781,7 +766,8 @@ int main(int argc, const char* argv[]) {
     };
 
     if (gen_params.init_image_path.size() > 0) {
-        if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image)) {
+        const bool native_init = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER;
+        if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image, native_init ? 0 : 3)) {
             return 1;
         }
     }
@@ -795,8 +781,8 @@ int main(int argc, const char* argv[]) {
     if (gen_params.ref_image_paths.size() > 0) {
         gen_params.ref_images.clear();
         for (auto& path : gen_params.ref_image_paths) {
-            SDImageOwner ref_image({0, 0, 3, nullptr});
-            if (!load_image_and_update_size(path, ref_image, false)) {
+            SDImageOwner ref_image({0, 0, 0, nullptr});
+            if (!load_image_and_update_size(path, ref_image, 0)) {
                 return 1;
             }
             gen_params.ref_images.push_back(std::move(ref_image));
@@ -837,41 +823,22 @@ int main(int argc, const char* argv[]) {
     if (gen_params.mask_image_path.size() > 0) {
         if (!load_sd_image_from_file(gen_params.mask_image.put(),
                                      gen_params.mask_image_path.c_str(),
-                                     gen_params.get_resolved_width(),
-                                     gen_params.get_resolved_height(),
+                                     0,
+                                     0,
                                      1)) {
             LOG_ERROR("load image from '%s' failed", gen_params.mask_image_path.c_str());
             return 1;
         }
-    } else {
-        sd_image_t generated_mask = {0, 0, 1, nullptr};
-        generated_mask.data       = (uint8_t*)malloc(gen_params.get_resolved_width() * gen_params.get_resolved_height());
-        if (generated_mask.data == nullptr) {
-            LOG_ERROR("malloc mask image failed");
-            return 1;
-        }
-        generated_mask.width  = gen_params.get_resolved_width();
-        generated_mask.height = gen_params.get_resolved_height();
-        memset(generated_mask.data, 255, gen_params.get_resolved_width() * gen_params.get_resolved_height());
-        gen_params.mask_image.reset(generated_mask);
     }
 
     if (gen_params.control_image_path.size() > 0) {
         if (!load_sd_image_from_file(gen_params.control_image.put(),
                                      gen_params.control_image_path.c_str(),
-                                     gen_params.get_resolved_width(),
-                                     gen_params.get_resolved_height())) {
+                                     0,
+                                     0)) {
             LOG_ERROR("load image from '%s' failed", gen_params.control_image_path.c_str());
             return 1;
         }
-        if (cli_params.canny_preprocess) {  // apply preprocessor
-            preprocess_canny(gen_params.control_image.get(),
-                             0.08f,
-                             0.08f,
-                             0.8f,
-                             1.0f,
-                             false);
-        }
     }
 
     if (gen_params.ip_adapter_image_path.size() > 0) {
@@ -888,8 +855,6 @@ int main(int argc, const char* argv[]) {
         gen_params.control_frames.clear();
         if (!load_images_from_dir(gen_params.control_video_path,
                                   gen_params.control_frames,
-                                  gen_params.get_resolved_width(),
-                                  gen_params.get_resolved_height(),
                                   gen_params.video_frames)) {
             return 1;
         }
@@ -898,10 +863,7 @@ int main(int argc, const char* argv[]) {
     if (!gen_params.pm_id_images_dir.empty()) {
         gen_params.pm_id_images.clear();
         if (!load_images_from_dir(gen_params.pm_id_images_dir,
-                                  gen_params.pm_id_images,
-                                  0,
-                                  0,
-                                  0)) {
+                                  gen_params.pm_id_images)) {
             return 1;
         }
     }

examples/common/common.cpp

@@ -618,6 +618,10 @@ ArgOptions SDContextParams::get_options() {
          "--diffusion-fa",
          "use flash attention in the diffusion model only",
          true, &diffusion_flash_attn},
+        {"",
+         "--sage-attn",
+         "use native CUDA SageAttention in the diffusion model, with flash/default attention fallback",
+         true, &sage_attn},
         {"",
          "--diffusion-conv-direct",
          "use ggml_conv2d_direct in the diffusion model",
@@ -938,6 +942,7 @@ std::string SDContextParams::to_string() const {
         << "  vae_on_cpu: " << (vae_on_cpu ? "true" : "false") << ",\n"
         << "  flash_attn: " << (flash_attn ? "true" : "false") << ",\n"
         << "  diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n"
+        << "  sage_attn: " << (sage_attn ? "true" : "false") << ",\n"
         << "  linear_scale: " << linear_scale << ",\n"
         << "  attn_scale: " << attn_scale << ",\n"
         << "  diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n"
@@ -995,6 +1000,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.enable_mmap                     = enable_mmap;
     sd_ctx_params.flash_attn                      = flash_attn;
     sd_ctx_params.diffusion_flash_attn            = diffusion_flash_attn;
+    sd_ctx_params.sage_attn                       = sage_attn;
     sd_ctx_params.linear_scale                    = linear_scale;
     sd_ctx_params.attn_scale                      = attn_scale;
     sd_ctx_params.tae_preview_only                = taesd_preview;
@@ -1109,7 +1115,7 @@ ArgOptions SDGenerationParams::get_options() {
          &hires_upscaler},
         {"",
          "--extra-sample-args",
-         "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng",
+         "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; llada_image supports uniform; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng",
          (int)',',
          &extra_sample_args},
         {"",
@@ -1122,6 +1128,9 @@ ArgOptions SDGenerationParams::get_options() {
          "Key-value list to set up the way the reference images are processed (empty = auto-detect from model weigths)",
          (int)',',
          &ref_image_args},
+        {"", "--image-preprocess",
+         "Image preprocessing rule: target=init|end|mask|control|ref|ip-adapter|id|control-frame,index=N,mode=auto|none|stretch|crop|crop-resize|fit-pad,filter=auto|nearest|nearest-exact|bilinear|bicubic|lanczos,antialias=auto|true|false,width=W,height=H,anchor=center|top|bottom|left|right,pad_color=#RRGGBB[AA],canny=true|false. Repeat for multiple rules.",
+         (int)';', &image_preprocess},
     };
 
     options.int_options = {
@@ -1302,11 +1311,6 @@ ArgOptions SDGenerationParams::get_options() {
          "automatically increase the indices of references images based on the order they are listed (starting with 1).",
          true,
          &increase_ref_index},
-        {"",
-         "--disable-auto-resize-ref-image",
-         "disable auto resize of ref images",
-         false,
-         &auto_resize_ref_image},
         {"",
          "--circular",
          "enable circular padding on both axes for tileable output",
@@ -1748,7 +1752,7 @@ ArgOptions SDGenerationParams::get_options() {
          on_scm_policy_arg},
         {"",
          "--vae-tile-size",
-         "tile size for vae tiling, format [X]x[Y] (default: 32x32)",
+         "tile size for vae tiling in latent units, not image pixels, format [X]x[Y] (default: 32x32)",
          on_tile_size_arg},
         {"",
          "--vae-relative-tile-size",
@@ -1842,28 +1846,28 @@ bool decode_base64_image(const std::string& encoded_input,
         return false;
     }
 
-    int decoded_width  = 0;
-    int decoded_height = 0;
-    uint8_t* raw_data  = load_image_from_memory(reinterpret_cast<const char*>(image_bytes.data()),
-                                                static_cast<int>(image_bytes.size()),
-                                                decoded_width,
-                                                decoded_height,
-                                                expected_width,
-                                                expected_height,
-                                                target_channels);
+    int decoded_width    = 0;
+    int decoded_height   = 0;
+    int resolved_channel = target_channels;
+    uint8_t* raw_data    = load_image_from_memory(reinterpret_cast<const char*>(image_bytes.data()),
+                                                  static_cast<int>(image_bytes.size()),
+                                                  decoded_width,
+                                                  decoded_height,
+                                                  resolved_channel,
+                                                  expected_width,
+                                                  expected_height,
+                                                  target_channels);
     if (raw_data == nullptr) {
         return false;
     }
 
-    out_image.reset({(uint32_t)decoded_width, (uint32_t)decoded_height, (uint32_t)target_channels, raw_data});
+    out_image.reset({(uint32_t)decoded_width, (uint32_t)decoded_height, (uint32_t)resolved_channel, raw_data});
     return true;
 }
 
 static bool parse_image_json_field(const json& parent,
                                    const char* key,
                                    int channels,
-                                   int expected_width,
-                                   int expected_height,
                                    SDImageOwner& out_image) {
     if (!parent.contains(key)) {
         return true;
@@ -1875,14 +1879,12 @@ static bool parse_image_json_field(const json& parent,
     if (!parent.at(key).is_string()) {
         return false;
     }
-    return decode_base64_image(parent.at(key).get<std::string>(), channels, expected_width, expected_height, out_image);
+    return decode_base64_image(parent.at(key).get<std::string>(), channels, 0, 0, out_image);
 }
 
 static bool parse_image_array_json_field(const json& parent,
                                          const char* key,
                                          int channels,
-                                         int expected_width,
-                                         int expected_height,
                                          std::vector<SDImageOwner>& out_images) {
     if (!parent.contains(key)) {
         return true;
@@ -1901,7 +1903,7 @@ static bool parse_image_array_json_field(const json& parent,
             return false;
         }
         SDImageOwner image;
-        if (!decode_base64_image(item.get<std::string>(), channels, expected_width, expected_height, image)) {
+        if (!decode_base64_image(item.get<std::string>(), channels, 0, 0, image)) {
             return false;
         }
         out_images.push_back(std::move(image));
@@ -2000,6 +2002,29 @@ static bool resolve_model_file_from_dir(const std::string& model_name,
     return false;
 }
 
+bool SDGenerationParams::parse_image_preprocess_json(const std::string& json_str) {
+    const auto value = json::parse(json_str, nullptr, false);
+    std::string rules;
+    if (value.is_string()) {
+        rules = value.get<std::string>();
+    } else if (value.is_array()) {
+        for (const auto& item : value) {
+            if (!item.is_string()) {
+                LOG_ERROR("image_preprocess must contain rule strings");
+                return false;
+            }
+            if (!rules.empty())
+                rules += ";";
+            rules += item.get<std::string>();
+        }
+    } else {
+        LOG_ERROR("image_preprocess must be a string or array of strings");
+        return false;
+    }
+    image_preprocess = std::move(rules);
+    return true;
+}
+
 bool SDGenerationParams::from_json_str(
     const std::string& json_str,
     const std::function<std::string(const std::string&)>& lora_path_resolver) {
@@ -2011,6 +2036,9 @@ bool SDGenerationParams::from_json_str(
         return false;
     }
 
+    if (j.contains("image_preprocess") && !parse_image_preprocess_json(j["image_preprocess"].dump()))
+        return false;
+
     auto load_if_exists = [&](const char* key, auto& out) {
         if (j.contains(key)) {
             using T = std::decay_t<decltype(out)>;
@@ -2048,6 +2076,7 @@ bool SDGenerationParams::from_json_str(
     load_if_exists("cache_mode", cache_mode);
     load_if_exists("cache_option", cache_option);
     load_if_exists("scm_mask", scm_mask);
+    load_if_exists("ref_image_args", ref_image_args);
 
     load_if_exists("clip_skip", clip_skip);
     load_if_exists("width", width);
@@ -2065,7 +2094,6 @@ bool SDGenerationParams::from_json_str(
     load_if_exists("moe_boundary", moe_boundary);
     load_if_exists("vace_strength", vace_strength);
 
-    load_if_exists("auto_resize_ref_image", auto_resize_ref_image);
     load_if_exists("increase_ref_index", increase_ref_index);
     load_if_exists("embed_image_metadata", embed_image_metadata);
 
@@ -2209,32 +2237,23 @@ bool SDGenerationParams::from_json_str(
         LOG_ERROR("invalid lora");
         return false;
     }
-    if (!parse_image_json_field(j, "init_image", 3, width, height, init_image)) {
-        LOG_ERROR("invalid init_image");
+    auto load_image = [&](const char* key, int channels, SDImageOwner& image) {
+        if (!parse_image_json_field(j, key, channels, image)) {
+            LOG_ERROR("invalid %s", key);
+            return false;
+        }
+        return true;
+    };
+    if (!load_image("init_image", 0, init_image) ||
+        !load_image("end_image", 3, end_image) ||
+        !load_image("mask_image", 1, mask_image) ||
+        !load_image("control_image", 3, control_image) ||
+        !load_image("ip_adapter_image", 3, ip_adapter_image)) {
         return false;
     }
-    if (!parse_image_json_field(j, "end_image", 3, width, height, end_image)) {
-        LOG_ERROR("invalid end_image");
-        return false;
-    }
-    if (!parse_image_array_json_field(j, "ref_images", 3, width, height, ref_images)) {
-        LOG_ERROR("invalid ref_images");
-        return false;
-    }
-    if (!parse_image_array_json_field(j, "control_frames", 3, width, height, control_frames)) {
-        LOG_ERROR("invalid control_frames");
-        return false;
-    }
-    if (!parse_image_json_field(j, "mask_image", 1, width, height, mask_image)) {
-        LOG_ERROR("invalid mask_image");
-        return false;
-    }
-    if (!parse_image_json_field(j, "control_image", 3, width, height, control_image)) {
-        LOG_ERROR("invalid control_image");
-        return false;
-    }
-    if (!parse_image_json_field(j, "ip_adapter_image", 3, width, height, ip_adapter_image)) {
-        LOG_ERROR("invalid ip_adapter_image");
+    if (!parse_image_array_json_field(j, "ref_images", 0, ref_images) ||
+        !parse_image_array_json_field(j, "control_frames", 3, control_frames)) {
+        LOG_ERROR("invalid input image array");
         return false;
     }
 
@@ -2478,6 +2497,10 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s
 }
 
 bool SDGenerationParams::validate(SDMode mode) {
+    if (!image_preprocess.empty() && mode != IMG_GEN && mode != VID_GEN) {
+        LOG_ERROR("--image-preprocess requires img_gen or vid_gen mode");
+        return false;
+    }
     if (batch_count <= 0) {
         LOG_ERROR("error: batch_count must be greater than 0");
         return false;
@@ -2653,14 +2676,6 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
         pulid_id_weight,
     };
 
-    if (!auto_resize_ref_image) {
-        if (!ref_image_args.empty()) {
-            ref_image_args += ",";
-        }
-        ref_image_args += "resize_before_vae=0";
-        LOG_WARN("Notice: --disable-auto-resize-ref-image is deprecated. Use --ref-image-args \"resize_before_vae=off\" instead.");
-    }
-
     if (increase_ref_index) {
         if (!ref_image_args.empty()) {
             ref_image_args += ",";
@@ -2708,6 +2723,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size());
     params.circular_x                = circular || circular_x;
     params.circular_y                = circular || circular_y;
+    params.image_preprocess          = {image_preprocess.c_str()};
     return params;
 }
 
@@ -2810,6 +2826,7 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
     params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size());
     params.circular_x                = circular || circular_x;
     params.circular_y                = circular || circular_y;
+    params.image_preprocess          = {image_preprocess.c_str()};
     return params;
 }
 
@@ -2866,7 +2883,8 @@ std::string SDGenerationParams::to_string() const {
         << "  ref_video_audio_paths: " << vec_str_to_string(ref_video_audio_paths) << ",\n"
         << "  ref_audio_paths: " << vec_str_to_string(ref_audio_paths) << ",\n"
         << "  control_video_path: \"" << control_video_path << "\",\n"
-        << "  auto_resize_ref_image: " << (auto_resize_ref_image ? "true" : "false") << ",\n"
+        << "  image_preprocess: " << image_preprocess << ",\n"
+        << "  ref_image_args: " << ref_image_args << ",\n"
         << "  increase_ref_index: " << (increase_ref_index ? "true" : "false") << ",\n"
         << "  pm_id_images_dir: \"" << pm_id_images_dir << "\",\n"
         << "  pm_id_embed_path: \"" << pm_id_embed_path << "\",\n"
@@ -3017,12 +3035,13 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
     set_json_basename_if_not_empty(models, "control_net", ctx_params.control_net_path);
     root["models"] = std::move(models);
 
-    root["clip_skip"]             = gen_params.clip_skip;
-    root["strength"]              = gen_params.strength;
-    root["control_strength"]      = gen_params.control_strength;
-    root["ip_adapter_strength"]   = gen_params.ip_adapter_strength;
-    root["auto_resize_ref_image"] = gen_params.auto_resize_ref_image;
-    root["increase_ref_index"]    = gen_params.increase_ref_index;
+    root["clip_skip"]           = gen_params.clip_skip;
+    root["strength"]            = gen_params.strength;
+    root["control_strength"]    = gen_params.control_strength;
+    root["ip_adapter_strength"] = gen_params.ip_adapter_strength;
+    root["ref_image_args"]      = gen_params.ref_image_args;
+    root["image_preprocess"]    = gen_params.image_preprocess;
+    root["increase_ref_index"]  = gen_params.increase_ref_index;
     if (mode == VID_GEN) {
         root["video"] = {
             {"frame_count", gen_params.video_frames},

examples/common/common.h

@@ -170,6 +170,7 @@ struct SDContextParams {
     bool vae_on_cpu            = false;
     bool flash_attn            = false;
     bool diffusion_flash_attn  = false;
+    bool sage_attn             = false;
     bool diffusion_conv_direct = false;
     bool vae_conv_direct       = false;
 
@@ -199,18 +200,17 @@ struct SDGenerationParams {
     std::string ad_prompt;
     std::string ad_negative_prompt;
     std::string extra_ad_args;
-    int clip_skip              = -1;  // <= 0 represents unspecified
-    int width                  = -1;
-    int height                 = -1;
-    int batch_count            = 1;
-    int qwen_image_layers      = 3;
-    int64_t seed               = 42;
-    float strength             = 0.75f;
-    float control_strength     = 0.9f;
-    float ip_adapter_strength  = 1.0f;
-    bool auto_resize_ref_image = true;
-    bool increase_ref_index    = false;
-    bool embed_image_metadata  = true;
+    int clip_skip             = -1;  // <= 0 represents unspecified
+    int width                 = -1;
+    int height                = -1;
+    int batch_count           = 1;
+    int qwen_image_layers     = 3;
+    int64_t seed              = 42;
+    float strength            = 0.75f;
+    float control_strength    = 0.9f;
+    float ip_adapter_strength = 1.0f;
+    bool increase_ref_index   = false;
+    bool embed_image_metadata = true;
 
     std::string init_image_path;
     std::string end_image_path;
@@ -246,6 +246,7 @@ struct SDGenerationParams {
     std::string extra_tiling_args;
 
     std::string ref_image_args;
+    std::string image_preprocess;
 
     std::string pm_id_images_dir;
     std::string pm_id_embed_path;
@@ -309,6 +310,7 @@ struct SDGenerationParams {
     ArgOptions get_options();
     bool from_json_str(const std::string& json_str,
                        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 width_and_height_are_set() const;

examples/common/media_io.cpp

@@ -261,6 +261,10 @@ uint8_t* decode_webp_image_to_buffer(const uint8_t* data,
     height               = features.height;
     source_channel_count = features.has_alpha ? 4 : 3;
 
+    if (expected_channel == 0) {
+        expected_channel = source_channel_count;
+    }
+
     const size_t pixel_count = static_cast<size_t>(width) * static_cast<size_t>(height);
 
     if (expected_channel == 1) {
@@ -481,7 +485,8 @@ uint8_t* load_image_common(bool from_memory,
                            int& height,
                            int expected_width,
                            int expected_height,
-                           int expected_channel) {
+                           int expected_channel,
+                           int& out_channel) {
     const char* image_path;
     FreeUniquePtr<uint8_t> image_buffer;
     int source_channel_count = 0;
@@ -538,6 +543,32 @@ uint8_t* load_image_common(bool from_memory,
         LOG_ERROR("load image from '%s' failed", image_path);
         return nullptr;
     }
+    if (expected_channel == 0) {
+        expected_channel = source_channel_count == 2 ? 4 : (source_channel_count == 1 ? 3 : source_channel_count);
+        if (expected_channel != source_channel_count) {
+            FreeUniquePtr<uint8_t> promoted((uint8_t*)malloc((size_t)width * height * expected_channel));
+            if (promoted == nullptr) {
+                LOG_ERROR("error: allocate memory for channel promotion, image_path = %s", image_path);
+                return nullptr;
+            }
+            const size_t pixel_count = (size_t)width * (size_t)height;
+            for (size_t i = 0; i < pixel_count; ++i) {
+                if (source_channel_count == 1) {
+                    promoted.get()[i * 3 + 0] = image_buffer.get()[i];
+                    promoted.get()[i * 3 + 1] = image_buffer.get()[i];
+                    promoted.get()[i * 3 + 2] = image_buffer.get()[i];
+                } else {
+                    promoted.get()[i * 4 + 0] = image_buffer.get()[i * 2];
+                    promoted.get()[i * 4 + 1] = image_buffer.get()[i * 2];
+                    promoted.get()[i * 4 + 2] = image_buffer.get()[i * 2];
+                    promoted.get()[i * 4 + 3] = image_buffer.get()[i * 2 + 1];
+                }
+            }
+            image_buffer         = std::move(promoted);
+            source_channel_count = expected_channel;
+        }
+    }
+    // stb reports the source channel count even when it converts the output.
     if (source_channel_count < expected_channel) {
         fprintf(stderr,
                 "the number of channels for the input image must be >= %d,"
@@ -597,7 +628,7 @@ uint8_t* load_image_common(bool from_memory,
         }
         stbir_resize(image_buffer.get(), width, height, 0,
                      resized_image_buffer.get(), expected_width, expected_height, 0, STBIR_TYPE_UINT8,
-                     expected_channel, STBIR_ALPHA_CHANNEL_NONE, 0,
+                     expected_channel, expected_channel == 4 ? 3 : STBIR_ALPHA_CHANNEL_NONE, 0,
                      STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP,
                      STBIR_FILTER_BOX, STBIR_FILTER_BOX,
                      STBIR_COLORSPACE_SRGB, nullptr);
@@ -605,6 +636,7 @@ uint8_t* load_image_common(bool from_memory,
         height       = expected_height;
         image_buffer = std::move(resized_image_buffer);
     }
+    out_channel = expected_channel;
     return image_buffer.release();
 }
 
@@ -777,10 +809,11 @@ bool write_image_to_file(const std::string& path,
 uint8_t* load_image_from_file(const char* image_path,
                               int& width,
                               int& height,
+                              int& out_channel,
                               int expected_width,
                               int expected_height,
                               int expected_channel) {
-    return load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel);
+    return load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel, out_channel);
 }
 
 bool load_sd_image_from_file(sd_image_t* image,
@@ -790,13 +823,14 @@ bool load_sd_image_from_file(sd_image_t* image,
                              int expected_channel) {
     int width;
     int height;
-    image->data = load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel);
+    int resolved_channel = expected_channel;
+    image->data          = load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel, resolved_channel);
     if (image->data == nullptr) {
         return false;
     }
     image->width   = width;
     image->height  = height;
-    image->channel = expected_channel;
+    image->channel = resolved_channel;
     return true;
 }
 
@@ -804,10 +838,11 @@ uint8_t* load_image_from_memory(const char* image_bytes,
                                 int len,
                                 int& width,
                                 int& height,
+                                int& out_channel,
                                 int expected_width,
                                 int expected_height,
                                 int expected_channel) {
-    return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel);
+    return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel, out_channel);
 }
 
 static void append_avi_metadata(std::vector<uint8_t>& data, const std::string& parameters) {

examples/common/media_io.h

@@ -32,9 +32,12 @@ bool write_image_to_file(const std::string& path,
                          const std::string& parameters = "",
                          int quality                   = 90);
 
+// expected_channel == 0 preserves native channels (grayscale -> RGB, gray+alpha -> RGBA).
+// out_channel receives the output channel count.
 uint8_t* load_image_from_file(const char* image_path,
                               int& width,
                               int& height,
+                              int& out_channel,
                               int expected_width   = 0,
                               int expected_height  = 0,
                               int expected_channel = 3);
@@ -49,6 +52,7 @@ uint8_t* load_image_from_memory(const char* image_bytes,
                                 int len,
                                 int& width,
                                 int& height,
+                                int& out_channel,
                                 int expected_width   = 0,
                                 int expected_height  = 0,
                                 int expected_channel = 3);

examples/server/api.md

@@ -148,6 +148,19 @@ Native extension fields:
 
 - any `sdcpp API` fields embedded through `sd_cpp_extra_args` inside `prompt`
 
+Uploaded images are decoded at their original dimensions. The first decoded
+image establishes the generation dimensions if `size` is omitted. Input
+geometry follows `image_preprocess`: references preserve their dimensions by
+default, while init and mask use the generation canvas preset.
+
+Reference encoding then follows model presets and `ref_image_args`. To skip
+input geometry for references and disable resizing before VAE encoding, include
+this in `prompt`:
+
+```text
+edit this image <sd_cpp_extra_args>{"image_preprocess":"target=ref,mode=none","ref_image_args":"resize_before_vae=false"}</sd_cpp_extra_args>
+```
+
 Response fields:
 
 | Field | Type | Notes |
@@ -526,7 +539,7 @@ LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEH
 | Field | Type |
 | --- | --- |
 | `batch_count` | `integer` |
-| `auto_resize_ref_image` | `boolean` |
+| `ref_image_args` | `string` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
 | `ip_adapter_strength` | `number` |
@@ -653,7 +666,7 @@ Example:
   "strength": 0.75,
   "seed": -1,
   "batch_count": 1,
-  "auto_resize_ref_image": true,
+  "ref_image_args": "",
   "increase_ref_index": false,
   "control_strength": 0.9,
   "ip_adapter_strength": 1.0,
@@ -728,6 +741,17 @@ Example:
 
 ### Image Encoding Rules
 
+Native image/video requests and SDAPI accept `image_preprocess` as a rule string
+or array of rule strings. OpenAI-compatible requests can supply it in
+`sd_cpp_extra_args`. See [Image preprocessing](../../docs/image_preprocessing.md)
+for one-time input geometry, native-resolution decoding, mask alignment, and
+`canny=true` for edge detection on any supported image input.
+
+Image generation also accepts `ref_image_args` as a string (for example,
+`"resize_before_vae=false"`) in native and SDAPI requests, or through
+`sd_cpp_extra_args` in OpenAI-compatible requests. It controls downstream
+reference encoding and is independent of input geometry rules.
+
 Any image field accepts:
 
 - a raw base64 string, or
@@ -735,12 +759,15 @@ Any image field accepts:
 
 Channel expectations:
 
-- `init_image`: 3 channels
-- `ref_images[]`: 3 channels
+- `init_image`: native channels (3 or 4); alpha is preserved and applied per model
+- `ref_images[]`: native channels (3 or 4); alpha is preserved and applied per model
 - `control_image`: 3 channels
 - `ip_adapter_image`: 3 channels
 - `mask_image`: 1 channel
 
+Models that support RGBA (e.g. Qwen-Image 2.1) use the alpha channel of `init_image`
+and `ref_images[]`. RGB-only models drop it, so sending RGBA is safe for every model.
+
 If omitted or null:
 
 - single-image fields map to an empty `sd_image_t`
@@ -760,7 +787,8 @@ Top-level scalar fields:
 | `strength` | `number` |
 | `seed` | `integer` |
 | `batch_count` | `integer` |
-| `auto_resize_ref_image` | `boolean` |
+| `ref_image_args` | `string` |
+| `image_preprocess` | `string \| array<string>` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
 | `ip_adapter_strength` | `number` |

examples/server/main.cpp

@@ -76,9 +76,9 @@ int main(int argc, const char** argv) {
     SDSvrParams svr_params;
     SDContextParams ctx_params;
     SDGenerationParams default_gen_params;
-    parse_args(argc, argv, svr_params, ctx_params, default_gen_params);
 
     sd_set_log_callback(sd_log_cb, (void*)&svr_params);
+    parse_args(argc, argv, svr_params, ctx_params, default_gen_params);
 
     LOG_VERBOSE("version: %s", version_string().c_str());
     LOG_VERBOSE("%s", sd_get_system_info());

examples/server/routes_openai.cpp

@@ -157,42 +157,46 @@ static bool build_openai_edit_request(const httplib::Request& req,
     request.gen_params.height      = height;
     request.gen_params.batch_count = n;
 
+    std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt);
     for (auto& bytes : images_bytes) {
-        int img_w           = 0;
-        int img_h           = 0;
-        uint8_t* raw_pixels = load_image_from_memory(
-            reinterpret_cast<const char*>(bytes.data()),
-            static_cast<int>(bytes.size()),
-            img_w, img_h,
-            width, height, 3);
+        int img_w            = 0;
+        int img_h            = 0;
+        int resolved_channel = 0;
+        uint8_t* raw_pixels  = load_image_from_memory(
+             reinterpret_cast<const char*>(bytes.data()),
+             static_cast<int>(bytes.size()),
+             img_w, img_h, resolved_channel,
+             0, 0,
+             0);
         if (raw_pixels == nullptr) {
             continue;
         }
 
-        SDImageOwner image_owner({(uint32_t)img_w, (uint32_t)img_h, 3, raw_pixels});
+        const bool is_first_ref_image = request.gen_params.ref_images.empty();
+        SDImageOwner image_owner({(uint32_t)img_w, (uint32_t)img_h, (uint32_t)resolved_channel, raw_pixels});
         request.gen_params.set_width_and_height_if_unset(image_owner.get().width, image_owner.get().height);
+
+        if (is_first_ref_image) {
+            request.gen_params.init_image = image_owner;
+            if (request.gen_params.init_image.get().data == nullptr) {
+                error_message = "could not allocate init image";
+                return false;
+            }
+        }
+
         request.gen_params.ref_images.push_back(std::move(image_owner));
     }
 
-    if (!request.gen_params.ref_images.empty()) {
-        request.gen_params.init_image = request.gen_params.ref_images.front();
-    }
-
     if (!mask_bytes.empty()) {
-        int expected_width  = 0;
-        int expected_height = 0;
-        if (request.gen_params.width_and_height_are_set()) {
-            expected_width  = request.gen_params.width;
-            expected_height = request.gen_params.height;
-        }
-        int mask_w = 0;
-        int mask_h = 0;
+        int mask_w       = 0;
+        int mask_h       = 0;
+        int mask_channel = 0;
 
         uint8_t* mask_raw = load_image_from_memory(
             reinterpret_cast<const char*>(mask_bytes.data()),
             static_cast<int>(mask_bytes.size()),
-            mask_w, mask_h,
-            expected_width, expected_height, 1);
+            mask_w, mask_h, mask_channel,
+            0, 0, 1);
         request.gen_params.mask_image.reset({(uint32_t)mask_w, (uint32_t)mask_h, 1, mask_raw});
         const sd_image_t& mask_image = request.gen_params.mask_image.get();
         request.gen_params.set_width_and_height_if_unset(mask_image.width, mask_image.height);
@@ -205,7 +209,6 @@ static bool build_openai_edit_request(const httplib::Request& req,
         });
     }
 
-    std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt);
     if (!sd_cpp_extra_args_str.empty() && !request.gen_params.from_json_str(sd_cpp_extra_args_str)) {
         error_message = "invalid sd_cpp_extra_args";
         return false;

examples/server/routes_sdapi.cpp

@@ -80,17 +80,6 @@ static enum sample_method_t get_sdapi_sample_method(std::string name) {
     return it != hardcoded.end() ? it->second : SAMPLE_METHOD_COUNT;
 }
 
-static void assign_solid_mask(SDImageOwner& mask_owner, int width, int height) {
-    const size_t pixel_count = static_cast<size_t>(width) * static_cast<size_t>(height);
-    uint8_t* raw_mask        = static_cast<uint8_t*>(malloc(pixel_count));
-    if (raw_mask == nullptr) {
-        mask_owner.reset({0, 0, 1, nullptr});
-        return;
-    }
-    std::memset(raw_mask, 255, pixel_count);
-    mask_owner.reset({(uint32_t)width, (uint32_t)height, 1, raw_mask});
-}
-
 static bool build_sdapi_img_gen_request(const json& j,
                                         ServerRuntime& runtime,
                                         bool img2img,
@@ -193,15 +182,25 @@ static bool build_sdapi_img_gen_request(const json& j,
         }
     }
 
-    if (img2img) {
-        const int expected_width  = request.gen_params.width_and_height_are_set() ? request.gen_params.width : 0;
-        const int expected_height = request.gen_params.width_and_height_are_set() ? request.gen_params.height : 0;
+    if (j.contains("ref_image_args")) {
+        if (!j["ref_image_args"].is_string()) {
+            error_message = "ref_image_args must be a string";
+            return false;
+        }
+        request.gen_params.ref_image_args = j["ref_image_args"].get<std::string>();
+    }
 
+    if (j.contains("image_preprocess") && !request.gen_params.parse_image_preprocess_json(j["image_preprocess"].dump())) {
+        error_message = "invalid image_preprocess";
+        return false;
+    }
+
+    if (img2img) {
         if (j.contains("init_images") && j["init_images"].is_array() && !j["init_images"].empty()) {
             if (decode_base64_image(j["init_images"][0].get<std::string>(),
-                                    3,
-                                    expected_width,
-                                    expected_height,
+                                    0,
+                                    0,
+                                    0,
                                     request.gen_params.init_image)) {
                 const sd_image_t& image = request.gen_params.init_image.get();
                 request.gen_params.set_width_and_height_if_unset(image.width, image.height);
@@ -211,8 +210,8 @@ static bool build_sdapi_img_gen_request(const json& j,
         if (j.contains("mask") && j["mask"].is_string()) {
             if (decode_base64_image(j["mask"].get<std::string>(),
                                     1,
-                                    expected_width,
-                                    expected_height,
+                                    0,
+                                    0,
                                     request.gen_params.mask_image)) {
                 const sd_image_t& image = request.gen_params.mask_image.get();
                 request.gen_params.set_width_and_height_if_unset(image.width, image.height);
@@ -225,9 +224,7 @@ static bool build_sdapi_img_gen_request(const json& j,
                 }
             }
         } else {
-            const int resolved_width  = request.gen_params.get_resolved_width();
-            const int resolved_height = request.gen_params.get_resolved_height();
-            assign_solid_mask(request.gen_params.mask_image, resolved_width, resolved_height);
+            request.gen_params.mask_image.reset({0, 0, 1, nullptr});
         }
 
         float denoising_strength = j.value("denoising_strength", -1.f);
@@ -243,9 +240,8 @@ static bool build_sdapi_img_gen_request(const json& j,
             }
             SDImageOwner image_owner;
             if (decode_base64_image(extra_image.get<std::string>(),
-                                    3,
-                                    request.gen_params.width_and_height_are_set() ? request.gen_params.width : 0,
-                                    request.gen_params.width_and_height_are_set() ? request.gen_params.height : 0,
+                                    0,
+                                    0, 0,
                                     image_owner)) {
                 const sd_image_t& image = image_owner.get();
                 request.gen_params.set_width_and_height_if_unset(image.width, image.height);

examples/server/routes_sdcpp.cpp

@@ -127,7 +127,8 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
         {"seed", defaults.seed},
         {"batch_count", defaults.batch_count},
         {"qwen_image_layers", defaults.qwen_image_layers},
-        {"auto_resize_ref_image", defaults.auto_resize_ref_image},
+        {"ref_image_args", defaults.ref_image_args},
+        {"image_preprocess", defaults.image_preprocess},
         {"increase_ref_index", defaults.increase_ref_index},
         {"control_strength", defaults.control_strength},
         {"ip_adapter_strength", defaults.ip_adapter_strength},
@@ -153,6 +154,7 @@ static json make_vid_gen_defaults_json(const SDGenerationParams& defaults, const
         {"strength", defaults.strength},
         {"seed", defaults.seed},
         {"video_frames", defaults.video_frames},
+        {"image_preprocess", defaults.image_preprocess},
         {"fps", defaults.fps},
         {"moe_boundary", defaults.moe_boundary},
         {"vace_strength", defaults.vace_strength},

include/stable-diffusion.h

@@ -79,6 +79,7 @@ enum scheduler_t {
     FLUX2_SCHEDULER,
     FLUX_SCHEDULER,
     BETA_SCHEDULER,
+    LLADA_IMAGE_SCHEDULER,
     SCHEDULER_COUNT
 };
 
@@ -245,6 +246,7 @@ typedef struct {
     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
+    bool sage_attn;
 } sd_ctx_params_t;
 
 typedef struct {
@@ -261,6 +263,11 @@ typedef struct {
     uint8_t* data;
 } sd_image_t;
 
+typedef struct {
+    // Semicolon-separated target=...,key=value rules. NULL preserves defaults.
+    const char* rules;
+} sd_image_preprocess_params_t;
+
 typedef struct {
     sd_image_t* frames;
     int frame_count;
@@ -408,6 +415,7 @@ typedef struct {
     int qwen_image_layers;
     bool circular_x;
     bool circular_y;
+    sd_image_preprocess_params_t image_preprocess;
 } sd_img_gen_params_t;
 
 typedef struct {
@@ -441,6 +449,7 @@ typedef struct {
     sd_hires_params_t hires;
     bool circular_x;
     bool circular_y;
+    sd_image_preprocess_params_t image_preprocess;
 } sd_vid_gen_params_t;
 
 typedef struct sd_ctx_t sd_ctx_t;

src/conditioning/conditioner.hpp

@@ -14,6 +14,7 @@
 #include "core/util.h"
 #include "model/diffusion/model.hpp"
 #include "model/te/clip.hpp"
+#include "model/te/llada_image_te.hpp"
 #include "model/te/llm.hpp"
 #include "model/te/t5.hpp"
 #include "model_loader.h"
@@ -136,6 +137,7 @@ struct ConditionerParams {
     const std::vector<sd::Tensor<float>>* ref_images                    = nullptr;  // for qwen image edit
     const std::vector<MiniMaxH3PresentationItem>* minimax_h3_references = nullptr;
     RefImageParams ref_image_params;
+    bool allow_cache = false;
 };
 
 struct Conditioner {
@@ -1953,6 +1955,10 @@ struct LLMEmbedder : public Conditioner {
     std::shared_ptr<LLM::LLMRunner> llm;
     std::shared_ptr<T5Runner> byt5;
 
+    bool h3_text_cache_valid = false;
+    std::string h3_text_cache_text;
+    SDCondition h3_text_cache;
+
     LLMEmbedder(ggml_backend_t backend,
                 const String2TensorStorage& tensor_storage_map      = {},
                 SDVersion version                                   = VERSION_QWEN_IMAGE,
@@ -1978,7 +1984,8 @@ struct LLMEmbedder : public Conditioner {
             arch = LLM::LLMArch::GPT_OSS_20B;
         } else if (sd_version_is_pid(version)) {
             arch = LLM::LLMArch::GEMMA2_2B;
-        } else if (sd_version_is_lingbot_video(version) ||
+        } else if (version == VERSION_QWEN_IMAGE_2_1 ||
+                   sd_version_is_lingbot_video(version) ||
                    sd_version_is_ideogram4(version) ||
                    sd_version_is_boogu_image(version) ||
                    sd_version_is_sefi_image(version) ||
@@ -2217,7 +2224,10 @@ struct LLMEmbedder : public Conditioner {
                                           false,
                                           deepstack_image_embeds,
                                           image_grids);
-        GGML_ASSERT(!hidden_states.empty());
+        if (hidden_states.empty()) {
+            LOG_ERROR("LLM prompt encoding failed");
+            return {};
+        }
         hidden_states = apply_token_weights(std::move(hidden_states), weights);
         GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
 
@@ -2293,6 +2303,25 @@ struct LLMEmbedder : public Conditioner {
 
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
+        const bool h3_text_cacheable =
+            sd_version_is_minimax_h3(version) &&
+            conditioner_params.allow_cache &&
+            (conditioner_params.minimax_h3_references == nullptr ||
+             conditioner_params.minimax_h3_references->empty()) &&
+            (conditioner_params.ref_images == nullptr ||
+             conditioner_params.ref_images->empty());
+
+        if (sd_version_is_minimax_h3(version) && !h3_text_cacheable) {
+            h3_text_cache_valid = false;
+        }
+
+        if (h3_text_cacheable &&
+            h3_text_cache_valid &&
+            h3_text_cache_text == conditioner_params.text) {
+            LOG_INFO("H3 conditioning cache hit");
+            return h3_text_cache;
+        }
+
         std::string prompt;
         std::pair<int, int> prompt_attn_range;
         std::vector<std::string> extra_prompts;
@@ -2547,6 +2576,67 @@ struct LLMEmbedder : public Conditioner {
             prompt += conditioner_params.text;
             prompt_attn_range = {0, 0};
             prompt += "<|im_end|>\n<|im_start|>assistant\n";
+        } else if (version == VERSION_QWEN_IMAGE_2_1) {
+            if (!llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
+                LOG_ERROR("Qwen Image 2.1 editing requires Qwen3-VL vision weights; provide --llm_vision or a combined encoder");
+                return {};
+            }
+            prompt = "<|im_start|>system\nComprehend and analyze the provided prompt.<|im_end|>\n";
+            std::vector<int> system_tokens;
+            if (!tokenizer->encode(prompt, system_tokens, nullptr)) {
+                return {};
+            }
+            prompt_template_encode_start_idx = static_cast<int>(system_tokens.size());
+            out_layers                       = {static_cast<int>(llm->config.num_layers)};
+            prompt += "<|im_start|>user\n";
+            if (llm->enable_vision && conditioner_params.ref_images != nullptr) {
+                for (size_t i = 0; i < conditioner_params.ref_images->size(); ++i) {
+                    const auto& image = (*conditioner_params.ref_images)[i];
+                    int64_t width     = image.shape()[0];
+                    int64_t height    = image.shape()[1];
+                    int64_t pixels    = width * height;
+                    if (width % 32 != 0 || height % 32 != 0) {
+                        LOG_ERROR("Qwen Image 2.1 reference dimensions must be multiples of 32");
+                        return {};
+                    }
+                    auto rgb = sd::Tensor<float>({width, height, 3, 1});
+                    for (int64_t p = 0; p < pixels; ++p) {
+                        float alpha = image.shape()[2] == 4 ? image[p + 3 * pixels] : 1.f;
+                        for (int c = 0; c < 3; ++c) {
+                            rgb[p + c * pixels] = 2.f * (image[p + c * pixels] * alpha + 1.f - alpha) - 1.f;
+                        }
+                    }
+                    auto outputs = llm->encode_image_outputs(n_threads, rgb, false);
+                    if (outputs.empty()) {
+                        return {};
+                    }
+                    prompt += (i == 0 ? "" : " ") + std::string("<image") + std::to_string(i + 1) + "><|vision_start|>";
+                    std::vector<int> prefix_tokens;
+                    if (!tokenizer->encode(prompt, prefix_tokens, nullptr)) {
+                        return {};
+                    }
+                    int index = static_cast<int>(prefix_tokens.size());
+                    int count = static_cast<int>(outputs[0].shape()[1]);
+                    image_embeds.emplace_back(index, std::move(outputs[0]));
+                    if (deepstack_image_embeds.empty()) {
+                        deepstack_image_embeds.resize(outputs.size() - 1);
+                    }
+                    for (size_t layer = 1; layer < outputs.size(); ++layer) {
+                        deepstack_image_embeds[layer - 1].emplace_back(index, std::move(outputs[layer]));
+                    }
+                    image_grids.push_back({index, count,
+                                           static_cast<int>(height) / llm->config.vision.patch_size,
+                                           static_cast<int>(width) / llm->config.vision.patch_size});
+                    for (int j = 0; j < count; ++j) {
+                        prompt += "<|image_pad|>";
+                    }
+                    prompt += "<|vision_end|>";
+                }
+            }
+            prompt_attn_range.first = static_cast<int>(prompt.size());
+            prompt += conditioner_params.text.empty() ? " " : conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+            prompt += "<|im_end|>\n<|im_start|>assistant\n";
         } else if (sd_version_is_qwen_image(version) || sd_version_is_mage_flow(version)) {
             if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
                 LOG_INFO("%s", sd_version_is_mage_flow(version) ? "MageFlowEditPipeline" : "QwenImageEditPlusPipeline");
@@ -3074,6 +3164,21 @@ struct LLMEmbedder : public Conditioner {
         SDCondition result;
         result.c_crossattn        = std::move(hidden_states);
         result.extra_c_crossattns = std::move(extra_hidden_states_vec);
+        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) {
+                int64_t begin = image_embeds[i].first - prompt_template_encode_start_idx;
+                int64_t end   = begin + image_embeds[i].second.shape()[1];
+                if (begin < 0 || end > slots.numel()) {
+                    LOG_ERROR("Qwen Image 2.1 image slots exceed the encoded prompt");
+                    return {};
+                }
+                for (int64_t j = begin; j < end; ++j) {
+                    slots[j] = static_cast<int32_t>(i + 1);
+                }
+            }
+            result.c_token_types = std::move(slots);
+        }
         if (sd_version_is_minimax_h3(version)) {
             std::vector<int32_t> tags(static_cast<size_t>(result.c_crossattn.shape()[1]), 1);
             for (const auto& [index, image_embed] : image_embeds) {
@@ -3085,6 +3190,14 @@ struct LLMEmbedder : public Conditioner {
             int64_t tag_count    = static_cast<int64_t>(tags.size());
             result.c_token_types = sd::Tensor<int32_t>({tag_count}, std::move(tags));
         }
+
+        if (h3_text_cacheable) {
+            h3_text_cache_text  = conditioner_params.text;
+            h3_text_cache       = result;
+            h3_text_cache_valid = true;
+            LOG_INFO("H3 conditioning cache stored");
+        }
+
         return result;
     }
 };
@@ -3159,6 +3272,214 @@ struct LTXAVTextProjectionRunner : public GGMLRunner {
     }
 };
 
+// LLaDA-Image's text path is a three-stage pipeline rather than a single encoder pass:
+// the token embeddings feed a QueryFormer whose 256 queries are appended to the backbone
+// input, and the backbone's final hidden states are projected to the denoiser's caption dim.
+// Ref: LLaDAImagePipeline._encode_text.
+struct LLaDAImageEmbedder : public Conditioner {
+    std::shared_ptr<Tokenizer> tokenizer;
+    std::shared_ptr<LLM::LLMRunner> llm;
+    std::shared_ptr<LLaDAImageTE::QueryFormerRunner> query_former;
+    std::shared_ptr<LLaDAImageTE::TextProjectionRunner> text_projection;
+    std::shared_ptr<LLaDAImageTE::SigVQRunner> sigvq;
+
+    std::string llm_prefix;
+    std::string query_former_prefix;
+    std::string text_projection_prefix;
+    std::string sigvq_prefix;
+
+    LLaDAImageEmbedder(ggml_backend_t backend,
+                       const String2TensorStorage& tensor_storage_map      = {},
+                       const std::string& llm_prefix                       = "text_encoders.llm",
+                       const std::string& query_former_prefix              = "queryformer",
+                       const std::string& text_projection_prefix           = "text_projection",
+                       const std::string& sigvq_prefix                     = "sigvq",
+                       std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                       const TokenizerConfig& tokenizers                   = {})
+        : llm_prefix(llm_prefix),
+          query_former_prefix(query_former_prefix),
+          text_projection_prefix(text_projection_prefix),
+          sigvq_prefix(sigvq_prefix) {
+        if (!tokenizers.has(TokenizerConfig::MAIN)) {
+            throw std::runtime_error("LLaDA-Image requires an external LLaDA2 tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer");
+        }
+        llm = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::LLADA2_MOE,
+                                               backend,
+                                               tensor_storage_map,
+                                               llm_prefix,
+                                               false,
+                                               weight_manager);
+        // <|endoftext|> doubles as the pad token in LLaDA2's tokenizer.json.
+        tokenizer       = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 156892);
+        query_former    = std::make_shared<LLaDAImageTE::QueryFormerRunner>(backend,
+                                                                         tensor_storage_map,
+                                                                         query_former_prefix,
+                                                                         weight_manager);
+        text_projection = std::make_shared<LLaDAImageTE::TextProjectionRunner>(backend,
+                                                                               tensor_storage_map,
+                                                                               text_projection_prefix,
+                                                                               weight_manager);
+
+        // SigVQ is only present when the user supplies the editing weights.
+        for (const auto& [name, _] : tensor_storage_map) {
+            if (starts_with(name, sigvq_prefix + ".")) {
+                sigvq = std::make_shared<LLaDAImageTE::SigVQRunner>(backend,
+                                                                    tensor_storage_map,
+                                                                    sigvq_prefix,
+                                                                    weight_manager);
+                break;
+            }
+        }
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        llm->get_param_tensors(tensors, llm_prefix);
+        query_former->get_param_tensors(tensors, query_former_prefix);
+        text_projection->get_param_tensors(tensors, text_projection_prefix);
+        if (sigvq != nullptr) {
+            sigvq->get_param_tensors(tensors, sigvq_prefix);
+        }
+    }
+
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
+        llm->get_param_tensor_ops(tensor_ops);
+    }
+
+    void set_flash_attention_enabled(bool enabled) override {
+        llm->set_flash_attention_enabled(enabled);
+        query_former->set_flash_attention_enabled(enabled);
+        text_projection->set_flash_attention_enabled(enabled);
+        if (sigvq != nullptr) {
+            sigvq->set_flash_attention_enabled(enabled);
+        }
+    }
+
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
+        llm->set_max_graph_vram_bytes(max_vram_bytes);
+        query_former->set_max_graph_vram_bytes(max_vram_bytes);
+        text_projection->set_max_graph_vram_bytes(max_vram_bytes);
+        if (sigvq != nullptr) {
+            sigvq->set_max_graph_vram_bytes(max_vram_bytes);
+        }
+    }
+
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        llm->set_runtime_backends(backends);
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        llm->set_graph_cut_layer_split_enabled(enabled);
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        llm->set_graph_cut_layer_split_backend_vram_limits(limits);
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        llm->get_param_tensors(tensors, llm_prefix);
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        llm->set_weight_adapter(adapter);
+        query_former->set_weight_adapter(adapter);
+        text_projection->set_weight_adapter(adapter);
+        if (sigvq != nullptr) {
+            sigvq->set_weight_adapter(adapter);
+        }
+    }
+
+    void runner_end() override {
+        llm->runner_end();
+        query_former->runner_end();
+        text_projection->runner_end();
+        if (sigvq != nullptr) {
+            sigvq->runner_end();
+        }
+    }
+
+    SDCondition get_learned_condition(int n_threads,
+                                      const ConditionerParams& conditioner_params) override {
+        const int64_t num_queries = 256;
+        const bool has_ref_images = conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty();
+        if (has_ref_images && sigvq == nullptr) {
+            LOG_ERROR("LLaDA-Image editing requires connectors with SigVQ weights");
+            return {};
+        }
+
+        std::string text = conditioner_params.text;
+        while (!text.empty() && std::isspace(static_cast<unsigned char>(text.front()))) {
+            text.erase(text.begin());
+        }
+        while (!text.empty() && std::isspace(static_cast<unsigned char>(text.back()))) {
+            text.pop_back();
+        }
+        std::string prompt = text.empty()
+                                 ? "<role>HUMAN</role> Generate an image.\n<role>ASSISTANT</role>\n<IMAGE1>"
+                                 : "<role>HUMAN</role> Generate an image: " + text + "\n<role>ASSISTANT</role>\n<IMAGE1>";
+
+        std::vector<int> tokens;
+        if (!tokenizer->encode(prompt, tokens, nullptr)) {
+            return {};
+        }
+        int64_t n_text = static_cast<int64_t>(tokens.size());
+        GGML_ASSERT(n_text > 0);
+
+        sd::Tensor<int32_t> text_ids({n_text}, std::vector<int32_t>(tokens.begin(), tokens.end()));
+        auto inputs_embeds = llm->compute_input_embeds(n_threads, text_ids);
+        auto query_embeds  = query_former->compute(n_threads, inputs_embeds);
+
+        // splice_image_embeds() replaces tokens in place, so the query slots have to exist in
+        // input_ids; their ids are irrelevant because the embeddings are overwritten.
+        std::vector<int32_t> padded(tokens.begin(), tokens.end());
+        padded.resize(static_cast<size_t>(n_text + num_queries), tokenizer->PAD_TOKEN_ID);
+        int64_t n_total = static_cast<int64_t>(padded.size());
+        sd::Tensor<int32_t> input_ids({n_total}, padded);
+
+        // Bidirectional everywhere except that the text tokens must not see the appended
+        // queries, matching backbone_attention_mask[:, :, :text_length, text_length:] = min.
+        const float mask_min = std::numeric_limits<float>::lowest() / 4.0f;
+        sd::Tensor<float> attention_mask({n_total, n_total});
+        for (int64_t i1 = 0; i1 < n_total; ++i1) {
+            for (int64_t i0 = 0; i0 < n_total; ++i0) {
+                float value                       = (i1 < n_text && i0 >= n_text) ? mask_min : 0.0f;
+                attention_mask[i0 + n_total * i1] = value;
+            }
+        }
+
+        LLM::ImageEmbeds image_embeds;
+        image_embeds.emplace_back(static_cast<int>(n_text), query_embeds);
+
+        std::set<int> out_layers = {static_cast<int>(llm->config.num_layers) + 1};
+        auto hidden_states       = llm->compute(n_threads,
+                                                input_ids,
+                                                attention_mask,
+                                                image_embeds,
+                                                out_layers);
+
+        SDCondition result;
+        result.c_crossattn = text_projection->compute(n_threads, hidden_states);
+
+        // Editing: SigVQ sees the reference at half the output resolution, as in
+        // LLaDAImagePipeline._encode_source_image.
+        if (has_ref_images) {
+            const auto& ref = conditioner_params.ref_images->front();
+            auto resized    = sd::ops::interpolate(ref,
+                                                   {conditioner_params.width / 2,
+                                                    conditioner_params.height / 2,
+                                                    ref.shape()[2],
+                                                    ref.shape()[3]},
+                                                   sd::ops::InterpolateMode::Bilinear);
+            resized         = resized * 2.f - 1.f;
+            auto semantic   = sigvq->compute(n_threads, resized);
+            if (semantic.empty()) {
+                return {};
+            }
+            result.extra_c_crossattns.push_back(std::move(semantic));
+        }
+        return result;
+    }
+};
+
 struct LTXAVEmbedder : public Conditioner {
     static constexpr int64_t kHiddenSize = 3840;
     static constexpr int64_t kNumStates  = 49;

src/convert.cpp

@@ -362,6 +362,9 @@ bool convert_with_components(const char* model_path,
                              const char* tensor_type_rules,
                              bool convert_name,
                              int n_threads) {
+    if (!validate_tensor_types(output_type, tensor_type_rules)) {
+        return false;
+    }
     ModelLoader model_loader;
     bool loaded_any = false;
 

src/core/backend_fit.cpp

@@ -478,7 +478,11 @@ namespace sd::backend_fit {
         return true;
     }
 
-    bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) {
+    bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling, ggml_status status) {
+        // Execution failures can leave the device unusable; tiling only helps with allocation failures.
+        if (status != GGML_STATUS_ALLOC_FAILED) {
+            return false;
+        }
         const char* retry_mode = nullptr;
         if (prefer_temporal_tiling && !tiling_params.temporal_tiling) {
             tiling_params.temporal_tiling = true;
@@ -498,7 +502,7 @@ namespace sd::backend_fit {
             return false;
         }
 
-        LOG_WARN("VAE decode failed (likely out of memory); retrying with %s tiling",
+        LOG_WARN("VAE decode ran out of memory; retrying with %s tiling",
                  retry_mode);
         return true;
     }

src/core/backend_fit.h

@@ -16,7 +16,8 @@ namespace sd::backend_fit {
                               std::string& params_spec);
 
     bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params,
-                                         bool prefer_temporal_tiling);
+                                         bool prefer_temporal_tiling,
+                                         ggml_status status);
 
 }  // namespace sd::backend_fit
 

src/core/compute_workspace.cpp

@@ -11,7 +11,7 @@
 #include "core/ggml_graph_cut.h"
 #include "core/util.h"
 #include "ggml-cpu.h"
-#include "ggml/src/ggml-impl.h"
+#include "ggml-impl.h"
 
 namespace sd {
     ComputeWorkspace::~ComputeWorkspace() {

src/core/ggml_extend.cpp

@@ -1,6 +1,7 @@
 #include "core/ggml_extend.h"
 
 #include <cmath>
+#include <stdexcept>
 #include <utility>
 
 #include "core/ggml_extend_backend.h"
@@ -247,6 +248,7 @@ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
                                            ggml_tensor* b,
                                            int convrot_group_size,
                                            float scale) {
+#ifndef SD_USE_UPSTREAM_GGML
     GGML_ASSERT(x->type == GGML_TYPE_F32 || (x->type == GGML_TYPE_I8 && scale == 1.f));
     if (scale != 1.f) {
         x = ggml_ext_scale(ctx, x, scale);
@@ -270,6 +272,16 @@ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
         }
     }
     return x;
+#else
+    GGML_UNUSED(ctx);
+    GGML_UNUSED(x);
+    GGML_UNUSED(w);
+    GGML_UNUSED(weight_scale);
+    GGML_UNUSED(b);
+    GGML_UNUSED(convrot_group_size);
+    GGML_UNUSED(scale);
+    throw std::runtime_error("INT8 tensorwise/convrot is not supported by this ggml build");
+#endif
 }
 
 ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
@@ -610,7 +622,8 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
                                     ggml_tensor* mask,
                                     bool skip_reshape,
                                     bool flash_attn,
-                                    float kv_scale) {  // avoid overflow
+                                    float kv_scale,
+                                    bool sage_attn) {  // avoid overflow
     int64_t L_q;
     int64_t L_k;
     int64_t C;
@@ -701,7 +714,37 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
         return out;
     };
 
-    if (flash_attn) {
+#ifndef SD_USE_UPSTREAM_GGML
+    if (sage_attn && mask == nullptr && d_head > 0 && d_head <= 128) {
+        auto q_in                 = ggml_reshape_4d(ctx, ggml_ext_cont(ctx, q->type == GGML_TYPE_F32 ? q : ggml_cast(ctx, q, GGML_TYPE_F32)), d_head, L_q, n_head, N);
+        auto k_in                 = ggml_reshape_4d(ctx, ggml_ext_cont(ctx, k->type == GGML_TYPE_F32 ? k : ggml_cast(ctx, k, GGML_TYPE_F32)), d_head, L_k, n_kv_head, N);
+        auto v_in                 = ggml_ext_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3));
+        const int64_t padded_head = d_head <= 64 ? 64 : 128;
+        if ((padded_head != d_head || kv_scale != 1.0f) && v_in->type != GGML_TYPE_F32) {
+            v_in = ggml_cast(ctx, v_in, GGML_TYPE_F32);
+        }
+        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);
+        }
+        if (kv_scale != 1.0f) {
+            k_in = ggml_ext_scale(ctx, k_in, kv_scale);
+            v_in = ggml_ext_scale(ctx, v_in, kv_scale);
+        }
+        v_in     = ggml_cast(ctx, v_in, GGML_TYPE_F16);
+        auto out = ggml_sage_attn(ctx, q_in, k_in, v_in, scale / kv_scale, GGML_SAGE_ATTN_AUTO);
+        if (ggml_backend_supports_op(backend, out)) {
+            kqv = kv_scale != 1.0f ? ggml_ext_scale(ctx, out, 1.0f / kv_scale) : out;
+            if (padded_head != d_head) {
+                kqv = ggml_ext_slice(ctx, kqv, 0, 0, d_head);
+            }
+        }
+    }
+#endif
+
+    if (kqv == nullptr && (flash_attn || sage_attn)) {
         // LOG_VERBOSE("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
         bool can_use_flash_attn = true;
         if (mask != nullptr) {

src/core/ggml_extend.h

@@ -220,7 +220,8 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
                                     ggml_tensor* mask = nullptr,
                                     bool skip_reshape = false,
                                     bool flash_attn   = false,
-                                    float kv_scale    = 1.0f);
+                                    float kv_scale    = 1.0f,
+                                    bool sage_attn    = false);
 
 ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
                                  ggml_tensor* x,

src/core/ggml_extend_backend.cpp

@@ -13,7 +13,7 @@
 #endif
 
 #include "core/util.h"
-#include "ggml/src/ggml-impl.h"
+#include "ggml-impl.h"
 #include "stable-diffusion.h"
 
 static std::string trim_copy(const std::string& value) {

src/core/ggml_graph_cut.cpp

@@ -16,7 +16,7 @@
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 
-#include "ggml/src/ggml-impl.h"
+#include "ggml-impl.h"
 
 namespace sd::ggml_graph_cut {
 

src/core/ggml_runner.cpp

@@ -25,7 +25,7 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
     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);
+    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);
 }
 
 void GGMLRunner::alloc_params_ctx() {
@@ -520,6 +520,7 @@ GGMLRunnerContext GGMLRunner::get_context() {
     runner_ctx.ggml_ctx              = compute_ctx;
     runner_ctx.backend               = runtime_backend;
     runner_ctx.flash_attn_enabled    = flash_attn_enabled;
+    runner_ctx.sage_attn_enabled     = sage_attn_enabled;
     runner_ctx.linear_scale          = linear_scale;
     runner_ctx.attn_scale            = attn_scale;
     runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled;
@@ -589,6 +590,7 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
                                                      bool auto_runner_end,
                                                      bool no_return,
                                                      const std::function<bool()>& read_outputs) {
+    last_compute_status_ = GGML_STATUS_FAILED;
     if (graph_active_) {
         LOG_ERROR("%s does not support reentrant graph execution", get_desc().c_str());
         return std::nullopt;
@@ -612,7 +614,9 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
         GGMLRunner& runner;
         const bool& success;
         ~GraphEndGuard() {
-            runner.workspace_.segment_end();
+            if (!runner.workspace_.segment_end()) {
+                runner.last_compute_status_ = GGML_STATUS_FAILED;
+            }
             runner.cache_.graph_end(false);
             runner.cut_cache_.clear();
             runner.free_compute_ctx();
@@ -641,6 +645,7 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
     try {
         output = execute_graph(graph, n_threads, no_return, read_outputs);
     } catch (const std::exception& error) {
+        last_compute_status_ = GGML_STATUS_FAILED;
         LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(),
                   ggml_backend_name(runtime_backend), error.what());
         return std::nullopt;
@@ -648,6 +653,7 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
     success = output.has_value();
     if (success) {
         cache_.graph_end(true);
+        last_compute_status_ = GGML_STATUS_SUCCESS;
     }
     return output;
 }
@@ -765,6 +771,7 @@ bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) {
     }
     workspace_.synchronize();
     if (status != GGML_STATUS_SUCCESS) {
+        last_compute_status_ = status;
         LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
         return false;
     }
@@ -817,6 +824,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
     const auto& cached_plan     = resolve_graph_cut_plan(graph);
     const auto full_measurement = measure(graph, cached_plan.compute_buffer_size);
     if (full_measurement.buffers.empty()) {
+        last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
         return std::nullopt;
     }
     auto manager         = residency_manager.lock();
@@ -887,7 +895,9 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
             SegmentGraphBindings& bindings;
             ggml_context* context;
             ~SegmentCleanup() {
-                runner.workspace_.segment_end();
+                if (!runner.workspace_.segment_end()) {
+                    runner.last_compute_status_ = GGML_STATUS_FAILED;
+                }
                 bindings.restore();
                 weights.segment_end();
                 ggml_free(context);
@@ -897,6 +907,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
 
         auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement;
         if (!workspace_.prepare(measurement)) {
+            last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
             return fail_segment("workspace preparation");
         }
         const size_t cut_bytes       = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment);
@@ -911,7 +922,11 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
                 sync_runtime_residency();
                 requests = memory_requests(measurement.buffers, new_cache_bytes);
             }
-            return weights.ensure_segment_capacity(index, requests);
+            const bool ready = weights.ensure_segment_capacity(index, requests);
+            if (!ready && manager != nullptr) {
+                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+            }
+            return ready;
         };
         if (!weights.segment_start(index, ensure_capacity)) {
             return fail_segment("weight preparation");
@@ -920,12 +935,17 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
         if (!workspace_.measurement_matches(segment_graph, measurement)) {
             measurement = measure(segment_graph, segment.compute_buffer_size);
         }
-        if (!workspace_.prepare(measurement) || !ensure_capacity()) {
+        if (!workspace_.prepare(measurement)) {
+            last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+            return fail_segment("workspace preparation");
+        }
+        if (!ensure_capacity()) {
             return fail_segment("workspace capacity check");
         }
         if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
                 pin_multi_device_nodes(scheduler, current);
             })) {
+            last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
             return fail_segment("workspace allocation");
         }
         for (const auto& size : measurement.buffers) {
@@ -963,6 +983,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
             }
         }
         if (!workspace_.segment_end()) {
+            last_compute_status_ = GGML_STATUS_FAILED;
             return fail_segment("workspace synchronization");
         }
         // Final outputs and their callbacks may still be views of consumed cuts.

src/core/ggml_runner.h

@@ -68,6 +68,7 @@ struct GGMLRunnerContext {
     ggml_backend_t backend                                           = nullptr;
     ggml_context* ggml_ctx                                           = nullptr;
     bool flash_attn_enabled                                          = false;
+    bool sage_attn_enabled                                           = false;
     float linear_scale                                               = 0.f;
     float attn_scale                                                 = 0.f;
     bool conv2d_direct_enabled                                       = false;
@@ -129,7 +130,8 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
 struct GGMLRunner {
 private:
     std::map<ggml_backend_t, size_t> logged_compute_bytes_;
-    size_t logged_segment_count_ = 0;
+    size_t logged_segment_count_     = 0;
+    ggml_status last_compute_status_ = GGML_STATUS_SUCCESS;
 
     sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes);
     std::vector<DeviceMemoryRequest> memory_requests(const std::vector<sd::BackendBufferSize>& sizes,
@@ -176,6 +178,7 @@ protected:
     const std::string final_result_name = "ggml_runner_final_result_tensor";
 
     bool flash_attn_enabled    = false;
+    bool sage_attn_enabled     = false;
     float linear_scale         = 0.f;
     float attn_scale           = 0.f;
     bool conv2d_direct_enabled = false;
@@ -333,10 +336,20 @@ public:
                                              bool no_return                            = false,
                                              const std::function<bool()>& read_outputs = {});
 
+    ggml_status last_compute_status() const { return last_compute_status_; }
+
     void set_flash_attention_enabled(bool enabled) {
         flash_attn_enabled = enabled;
     }
 
+    void set_sage_attention_enabled(bool enabled) {
+        if (sage_attn_enabled != enabled) {
+            free_cache_ctx_and_buffer();
+            graph_cut_plan_cache_.graph_cut_plans.clear();
+            sage_attn_enabled = enabled;
+        }
+    }
+
     void set_scale_overrides(float linear_scale, float attn_scale) {
         this->linear_scale = linear_scale;
         this->attn_scale   = attn_scale;

src/core/util.cpp

@@ -62,17 +62,34 @@ void replace_all_chars(std::string& str, char target, char replacement) {
     }
 }
 
+static std::string sd_vformat(const char* fmt, va_list ap) {
+    char small[128];
+    va_list ap2;
+    va_copy(ap2, ap);
+    int size = vsnprintf(small, sizeof small, fmt, ap);
+    if (size < 0) {
+        va_end(ap2);
+        return {};
+    }
+    size_t needed = (size_t)size;
+    if (needed < sizeof small) {
+        va_end(ap2);
+        return std::string(small, needed);
+    }
+    std::string out(needed, '\0');
+    int size2 = vsnprintf(out.data(), needed + 1, fmt, ap2);
+    va_end(ap2);
+    if (size2 < 0)
+        out.clear();
+    return out;
+}
+
 std::string sd_format(const char* fmt, ...) {
     va_list ap;
-    va_list ap2;
     va_start(ap, fmt);
-    va_copy(ap2, ap);
-    int size = vsnprintf(nullptr, 0, fmt, ap);
-    std::vector<char> buf(size + 1);
-    int size2 = vsnprintf(buf.data(), size + 1, fmt, ap2);
-    va_end(ap2);
+    std::string result = sd_vformat(fmt, ap);
     va_end(ap);
-    return std::string(buf.data(), size);
+    return result;
 }
 
 int round_up_to(int value, int base) {
@@ -414,14 +431,43 @@ std::vector<std::string> split_string(const std::string& str, char delimiter) {
 }
 
 ggml_type sd_type_to_ggml_type(sd_type_t sdtype) {
+    if (sdtype == SD_TYPE_F8_E4M3 || sdtype == SD_TYPE_F8_E5M2) {
+#ifndef SD_USE_UPSTREAM_GGML
+        return sdtype == SD_TYPE_F8_E4M3 ? GGML_TYPE_F8_E4M3 : GGML_TYPE_F8_E5M2;
+#else
+        return GGML_TYPE_COUNT;
+#endif
+    }
     const int type_value = static_cast<int>(sdtype);
-    if (type_value < std::min<int>(SD_TYPE_COUNT, GGML_TYPE_COUNT)) {
+    if (type_value >= 0 && type_value < std::min<int>(SD_TYPE_COUNT, GGML_TYPE_COUNT)) {
         return static_cast<ggml_type>(type_value);
     } else {
         return GGML_TYPE_COUNT;
     }
 }
 
+bool validate_tensor_types(sd_type_t type, const char* tensor_type_rules) {
+    if (type != SD_TYPE_COUNT && sd_type_to_ggml_type(type) == GGML_TYPE_COUNT) {
+        LOG_ERROR("weight type %s is not supported by this ggml build", sd_type_name(type));
+        return false;
+    }
+#ifdef SD_USE_UPSTREAM_GGML
+    for (const auto& rule : split_string(SAFE_STR(tensor_type_rules), ',')) {
+        const auto pos = rule.find('=');
+        if (pos != std::string::npos) {
+            const auto name = rule.substr(pos + 1);
+            if (name == "f8_e4m3" || name == "f8_e5m2") {
+                LOG_ERROR("FP8 is not supported by this ggml build (tensor type rule '%s')", rule.c_str());
+                return false;
+            }
+        }
+    }
+#else
+    GGML_UNUSED(tensor_type_rules);
+#endif
+    return true;
+}
+
 KeyValueArgs parse_key_value_args(const char* args, const char* context) {
     KeyValueArgs pairs;
 
@@ -595,47 +641,45 @@ std::string trim(const std::string& s) {
 static sd_log_cb_t sd_log_cb = nullptr;
 void* sd_log_cb_data         = nullptr;
 
-#define LOG_BUFFER_SIZE 4096
+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';
+    }
+    sd_log_cb(level, message.c_str(), sd_log_cb_data);
+}
 
 void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...) {
     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 < LOG_BUFFER_SIZE) {
-        vsnprintf(log_buffer + written, LOG_BUFFER_SIZE - written, format, args);
-    }
-    size_t len = strlen(log_buffer);
-    if (log_buffer[len - 1] != '\n') {
-        strncat(log_buffer, "\n", LOG_BUFFER_SIZE - len);
-    }
-
-    if (sd_log_cb) {
-        sd_log_cb(level, log_buffer, sd_log_cb_data);
-    }
-
+    std::string message = sd_vformat(format, args);
     va_end(args);
+    std::string origin = sd_format("%s:%-4d", sd_basename(file).c_str(), line);
+    sd_log_dispatch(level, origin, message);
 }
 
 void sd_ggml_log_callback(ggml_log_level level, const char* text, void*) {
+    sd_log_level_t sd_level = SD_LOG_VERBOSE;
     switch (level) {
         case GGML_LOG_LEVEL_DEBUG:
-            LOG_VERBOSE(text);
+            sd_level = SD_LOG_VERBOSE;
             break;
         case GGML_LOG_LEVEL_INFO:
-            LOG_INFO(text);
+            sd_level = SD_LOG_INFO;
             break;
         case GGML_LOG_LEVEL_WARN:
-            LOG_WARN(text);
+            sd_level = SD_LOG_WARN;
             break;
         case GGML_LOG_LEVEL_ERROR:
-            LOG_ERROR(text);
+            sd_level = SD_LOG_ERROR;
             break;
         default:
-            LOG_VERBOSE(text);
+            sd_level = SD_LOG_VERBOSE;
+            break;
     }
+    sd_log_dispatch(sd_level, "ggml", text);
 }
 
 void sd_set_log_callback(sd_log_cb_t cb, void* data) {

src/core/util.h

@@ -90,6 +90,7 @@ void log_printf(sd_log_level_t level, const char* file, int line, const char* fo
 void sd_ggml_log_callback(ggml_log_level level, const char* text, void*);
 
 ggml_type sd_type_to_ggml_type(sd_type_t sdtype);
+bool validate_tensor_types(sd_type_t type, const char* tensor_type_rules);
 
 std::string trim(const std::string& s);
 

src/detailer.cpp

@@ -970,6 +970,7 @@ bool adetail_image(adetailer_ctx_t* context,
         generation.pm_params           = {};
         generation.pulid_params        = {};
         generation.hires.enabled       = false;
+        generation.image_preprocess    = {};
         if (params.steps > 0) {
             generation.sample_params.sample_steps        = params.steps;
             generation.sample_params.custom_sigmas       = nullptr;

src/model.h

@@ -39,6 +39,7 @@ enum SDVersion {
     VERSION_LINGBOT_VIDEO,
     VERSION_QWEN_IMAGE,
     VERSION_QWEN_IMAGE_LAYERED,
+    VERSION_QWEN_IMAGE_2_1,
     VERSION_HUNYUAN_VIDEO,
     VERSION_ANIMA,
     VERSION_FLUX2,
@@ -59,6 +60,7 @@ enum SDVersion {
     VERSION_KREA2,
     VERSION_MAGE_FLOW,
     VERSION_SENSENOVA_U1_5,
+    VERSION_LLADA_IMAGE,
     VERSION_ESRGAN,
     VERSION_COUNT,
 };
@@ -145,7 +147,7 @@ static inline bool sd_version_is_lingbot_video(SDVersion version) {
 }
 
 static inline bool sd_version_is_qwen_image(SDVersion version) {
-    if (version == VERSION_QWEN_IMAGE || version == VERSION_QWEN_IMAGE_LAYERED) {
+    if (version == VERSION_QWEN_IMAGE || version == VERSION_QWEN_IMAGE_LAYERED || version == VERSION_QWEN_IMAGE_2_1) {
         return true;
     }
     return false;
@@ -172,6 +174,13 @@ static inline bool sd_version_is_z_image(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_llada_image(SDVersion version) {
+    if (version == VERSION_LLADA_IMAGE) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_boogu_image(SDVersion version) {
     if (version == VERSION_BOOGU_IMAGE) {
         return true;
@@ -243,6 +252,14 @@ static inline bool sd_version_is_sensenova_u1(SDVersion version) {
     return version == VERSION_SENSENOVA_U1_5;
 }
 
+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);
+}
+
+static inline bool sd_version_supports_image_generation(SDVersion version) {
+    return !sd_version_supports_video_generation(version);
+}
+
 static inline bool sd_version_uses_flux_vae(SDVersion version) {
     if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) {
         return true;
@@ -251,7 +268,7 @@ static inline bool sd_version_uses_flux_vae(SDVersion version) {
 }
 
 static inline bool sd_version_uses_flux2_vae(SDVersion version) {
-    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version)) {
+    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version) || sd_version_is_llada_image(version)) {
         return true;
     }
     return false;
@@ -292,6 +309,7 @@ static inline bool sd_version_is_dit(SDVersion version) {
         version == VERSION_HIDREAM_O1 ||
         sd_version_is_anima(version) ||
         sd_version_is_z_image(version) ||
+        sd_version_is_llada_image(version) ||
         sd_version_is_boogu_image(version) ||
         sd_version_is_ernie_image(version) ||
         sd_version_is_lens(version) ||

src/model/common/ggml_block.hpp

@@ -208,6 +208,7 @@ public:
         ggml_tensor* w            = params["weight"];
         const float scale         = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale;
         ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr;
+#ifndef SD_USE_UPSTREAM_GGML
         if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) {
             bool supports_fp8_matmul = false;
             if (ctx->backend != nullptr) {
@@ -221,6 +222,7 @@ public:
                 w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16);
             }
         }
+#endif
         ggml_tensor* b = nullptr;
         if (bias) {
             b = params["bias"];
@@ -238,6 +240,7 @@ public:
             if (ctx->weight_adapter && b != nullptr) {
                 b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
             }
+#ifndef SD_USE_UPSTREAM_GGML
             if (int8_convrot && scale == 1.f) {
                 const auto cache_key = std::make_pair(x, int8_convrot_group_size);
                 auto cached          = ctx->int8_convrot_cache.find(cache_key);
@@ -248,6 +251,7 @@ public:
                     x = cached->second;
                 }
             }
+#endif
             out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
                                                 x,
                                                 w,
@@ -835,21 +839,30 @@ class RMSNorm : public UnaryBlock {
 protected:
     int64_t hidden_size;
     float eps;
+    bool elementwise_affine;
     std::string prefix;
 
     void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
-        this->prefix         = prefix;
+        this->prefix = prefix;
+        if (!elementwise_affine) {
+            return;
+        }
         enum ggml_type wtype = GGML_TYPE_F32;
         params["weight"]     = ggml_new_tensor_1d(ctx, wtype, hidden_size);
     }
 
 public:
     RMSNorm(int64_t hidden_size,
-            float eps = 1e-06f)
+            float eps               = 1e-06f,
+            bool elementwise_affine = true)
         : hidden_size(hidden_size),
-          eps(eps) {}
+          eps(eps),
+          elementwise_affine(elementwise_affine) {}
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        if (!elementwise_affine) {
+            return ggml_rms_norm(ctx->ggml_ctx, x, eps);
+        }
         ggml_tensor* w = params["weight"];
         if (ctx->weight_adapter) {
             w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");

src/model/common/rope.hpp

@@ -929,6 +929,145 @@ namespace Rope {
         return ids;
     }
 
+    // LLaDA-Image shares Lumina2/z_image's axes layout, but assigns position (0,0,0) to the
+    // padding slots of the caption stream instead of continuing the caption ramp through them.
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_llada_image_ids(int h,
+                                                                          int w,
+                                                                          int patch_size,
+                                                                          int bs,
+                                                                          int context_len,
+                                                                          int seq_multi_of) {
+        int context_pad_len    = bound_mod(context_len, seq_multi_of);
+        int padded_context_len = context_len + context_pad_len;
+        auto txt_ids           = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
+        for (int i = 0; i < bs * padded_context_len; i++) {
+            int pos = i % padded_context_len;
+            if (pos < context_len) {
+                txt_ids[i][0] = pos + 1.f;
+            }
+        }
+
+        int axes_dim_num = 3;
+        int index        = padded_context_len + 1;
+        auto img_ids     = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index);
+
+        int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
+        if (img_pad_len > 0) {
+            std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
+            img_ids = concat_ids(img_ids, img_pad_ids, bs);
+        }
+
+        return concat_ids(txt_ids, img_ids, bs);
+    }
+
+    // LLaDA-Image editing packs two caption copies (clean and noisy), the source and target
+    // latents anchored at their own caption's end position, and the SigVQ stream after both.
+    // Padding slots keep position (0,0,0), as in the text-only layout.
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_llada_image_edit_ids(int h,
+                                                                               int w,
+                                                                               int patch_size,
+                                                                               int context_len,
+                                                                               int sigvq_len,
+                                                                               int seq_multi_of) {
+        const int context_pad    = bound_mod(context_len, seq_multi_of);
+        const int padded_context = context_len + context_pad;
+        const int h_len          = (h + (patch_size / 2)) / patch_size;
+        const int w_len          = (w + (patch_size / 2)) / patch_size;
+        const int image_len      = h_len * w_len;
+        const int image_pad      = bound_mod(image_len, seq_multi_of);
+        const int padded_image   = image_len + image_pad;
+        const int sigvq_pad      = bound_mod(sigvq_len, seq_multi_of);
+
+        std::vector<std::vector<float>> cap_ids;
+        std::vector<int> cap_end_positions;
+        int cursor = 1;
+        for (int copy = 0; copy < 2; ++copy) {
+            for (int i = 0; i < padded_context; ++i) {
+                std::vector<float> id(3, 0.f);
+                if (i < context_len) {
+                    id[0] = static_cast<float>(cursor + i);
+                }
+                cap_ids.push_back(id);
+            }
+            cursor += context_len;
+            cap_end_positions.push_back(cursor);
+            cursor += 2;
+        }
+
+        std::vector<std::vector<float>> img_ids;
+        for (int copy = 0; copy < 2; ++copy) {
+            auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy]);
+            img_ids.insert(img_ids.end(), ids.begin(), ids.end());
+            img_ids.insert(img_ids.end(), image_pad, std::vector<float>(3, 0.f));
+        }
+
+        const int sigvq_start = static_cast<int>(cap_ids.size() + img_ids.size()) + 1;
+        std::vector<std::vector<float>> sigvq_ids;
+        for (int i = 0; i < sigvq_len + sigvq_pad; ++i) {
+            std::vector<float> id(3, 0.f);
+            if (i < sigvq_len) {
+                id[0] = static_cast<float>(sigvq_start + i);
+            }
+            sigvq_ids.push_back(id);
+        }
+
+        std::vector<std::vector<float>> ids;
+        ids.reserve(cap_ids.size() + img_ids.size() + sigvq_ids.size());
+        ids.insert(ids.end(), cap_ids.begin(), cap_ids.end());
+        ids.insert(ids.end(), img_ids.begin(), img_ids.end());
+        ids.insert(ids.end(), sigvq_ids.begin(), sigvq_ids.end());
+        SD_UNUSED(padded_image);
+        return ids;
+    }
+
+    __STATIC_INLINE__ std::vector<float> gen_llada_image_edit_pe(int h,
+                                                                 int w,
+                                                                 int patch_size,
+                                                                 int context_len,
+                                                                 int sigvq_len,
+                                                                 int seq_multi_of,
+                                                                 int theta,
+                                                                 const std::vector<int>& axes_dim) {
+        auto ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of);
+        return embed_nd(ids, 1, static_cast<float>(theta), axes_dim, {});
+    }
+
+    __STATIC_INLINE__ std::vector<float> gen_llada_image_pe(int h,
+                                                            int w,
+                                                            int patch_size,
+                                                            int bs,
+                                                            int context_len,
+                                                            int seq_multi_of,
+                                                            int theta,
+                                                            bool circular_h,
+                                                            bool circular_w,
+                                                            const std::vector<int>& axes_dim) {
+        std::vector<std::vector<float>> ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of);
+        std::vector<std::vector<int>> wrap_dims;
+        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
+            int pad_h = (patch_size - (h % patch_size)) % patch_size;
+            int pad_w = (patch_size - (w % patch_size)) % patch_size;
+            int h_len = (h + pad_h) / patch_size;
+            int w_len = (w + pad_w) / patch_size;
+            if (h_len > 0 && w_len > 0) {
+                size_t pos_len = ids.size() / bs;
+                wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
+                size_t cursor     = context_len + bound_mod(context_len, seq_multi_of);
+                size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
+                for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
+                    if (circular_h) {
+                        wrap_dims[1][cursor + token_i] = h_len;
+                    }
+                    if (circular_w) {
+                        wrap_dims[2][cursor + token_i] = w_len;
+                    }
+                }
+            }
+        }
+
+        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
+    }
+
     // Generate z_image positional embeddings
     __STATIC_INLINE__ std::vector<float> gen_z_image_pe(int h,
                                                         int w,

src/model/diffusion/llada_image.hpp

@@ -0,0 +1,527 @@
+#ifndef __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__
+#define __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__
+
+#include <algorithm>
+#include <cinttypes>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
+#include "model/diffusion/model.hpp"
+#include "model/diffusion/z_image.hpp"
+#include "model_loader.h"
+
+// Ref: https://github.com/inclusionAI/LLaDA-Image/blob/main/src/models/transformer_llada_image.py
+//
+// The denoiser is Lumina2/z_image's NextDiT with identical hyperparameters, so the blocks are
+// reused from ZImage. Two things differ: every norm here is non-parametric (the checkpoint
+// carries no norm weights at all), and latents arrive already patchified from the Flux2 VAE,
+// so patch_size is 1 over 128 channels.
+
+namespace LLaDAImage {
+    constexpr int LLADA_IMAGE_GRAPH_SIZE = 20480;
+
+    struct LLaDAImageConfig {
+        int patch_size             = 1;
+        int64_t hidden_size        = 3840;
+        int64_t in_channels        = 128;
+        int64_t out_channels       = 128;
+        int64_t num_layers         = 30;
+        int64_t num_refiner_layers = 2;
+        int64_t head_dim           = 128;
+        int64_t num_heads          = 30;
+        int64_t num_kv_heads       = 30;
+        int64_t multiple_of        = 256;
+        float ffn_dim_multiplier   = 8.0f / 3.0f;
+        float norm_eps             = 1e-5f;
+        bool qk_norm               = true;
+        int64_t cap_feat_dim       = 2560;
+        int64_t semantic_feat_dim  = 4096;
+        int theta                  = 256;
+        std::vector<int> axes_dim  = {32, 48, 48};
+        int64_t axes_dim_sum       = 128;
+
+        static int64_t count_blocks(const String2TensorStorage& tensor_storage_map,
+                                    const std::string& prefix,
+                                    const std::string& block_prefix) {
+            int64_t count = 0;
+            for (const auto& [name, _] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                size_t pos = name.find(block_prefix);
+                if (pos == std::string::npos) {
+                    continue;
+                }
+                auto items = split_string(name.substr(pos), '.');
+                if (items.size() > 1) {
+                    count = std::max<int64_t>(count, atoi(items[1].c_str()) + 1);
+                }
+            }
+            return count;
+        }
+
+        static LLaDAImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            LLaDAImageConfig config;
+            int64_t detected_q_dim  = 0;
+            int64_t detected_kv_dim = 0;
+
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
+                    int64_t patch_area = config.patch_size * config.patch_size;
+                    config.in_channels = tensor_storage.ne[0] / patch_area;
+                    config.hidden_size = tensor_storage.ne[1];
+                } else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) {
+                    config.cap_feat_dim = tensor_storage.ne[0];
+                    config.hidden_size  = tensor_storage.ne[1];
+                } else if (ends_with(name, "sigvq_embedder.1.weight") && tensor_storage.n_dims == 2) {
+                    config.semantic_feat_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "layers.0.attention.to_q.weight") && tensor_storage.n_dims == 2) {
+                    detected_q_dim = tensor_storage.ne[1];
+                } else if (ends_with(name, "layers.0.attention.to_k.weight") && tensor_storage.n_dims == 2) {
+                    detected_kv_dim = tensor_storage.ne[1];
+                } else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) {
+                    int64_t patch_area  = config.patch_size * config.patch_size;
+                    config.out_channels = tensor_storage.ne[1] / patch_area;
+                }
+            }
+
+            int64_t detected_layers  = count_blocks(tensor_storage_map, prefix, "layers.");
+            int64_t detected_refiner = std::max(count_blocks(tensor_storage_map, prefix, "noise_refiner."),
+                                                count_blocks(tensor_storage_map, prefix, "context_refiner."));
+            if (detected_layers > 0) {
+                config.num_layers = detected_layers;
+            }
+            if (detected_refiner > 0) {
+                config.num_refiner_layers = detected_refiner;
+            }
+            if (detected_q_dim > 0) {
+                config.num_heads = detected_q_dim / config.head_dim;
+            }
+            if (detected_kv_dim > 0) {
+                config.num_kv_heads = detected_kv_dim / config.head_dim;
+            } else if (detected_q_dim > 0) {
+                config.num_kv_heads = config.num_heads;
+            }
+
+            LOG_VERBOSE("llada_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64 ", cap_feat_dim = %" PRId64 ", semantic_feat_dim = %" PRId64,
+                        config.num_layers,
+                        config.num_refiner_layers,
+                        config.hidden_size,
+                        config.num_heads,
+                        config.num_kv_heads,
+                        config.in_channels,
+                        config.out_channels,
+                        config.cap_feat_dim,
+                        config.semantic_feat_dim);
+            return config;
+        }
+    };
+
+    class LLaDAImageModel : public GGMLBlock {
+    protected:
+        LLaDAImageConfig 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);
+            params["sigvq_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
+        }
+
+        std::shared_ptr<ZImage::JointTransformerBlock> make_block(bool modulation) {
+            return std::make_shared<ZImage::JointTransformerBlock>(0,
+                                                                   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,
+                                                                   false,
+                                                                   true);
+        }
+
+    public:
+        LLaDAImageModel() = default;
+        LLaDAImageModel(LLaDAImageConfig 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, ZImage::ADALN_EMBED_DIM);
+
+            blocks["cap_embedder.0"] = std::make_shared<RMSNorm>(config.cap_feat_dim, config.norm_eps, false);
+            blocks["cap_embedder.1"] = std::make_shared<Linear>(config.cap_feat_dim, config.hidden_size);
+
+            blocks["semantic_embedder.0"] = std::make_shared<RMSNorm>(config.semantic_feat_dim, config.norm_eps, false);
+            blocks["semantic_embedder.1"] = std::make_shared<Linear>(config.semantic_feat_dim, config.hidden_size);
+            blocks["sigvq_embedder.0"]    = std::make_shared<RMSNorm>(config.semantic_feat_dim, config.norm_eps, false);
+            blocks["sigvq_embedder.1"]    = std::make_shared<Linear>(config.semantic_feat_dim, config.hidden_size);
+
+            for (int i = 0; i < config.num_refiner_layers; i++) {
+                blocks["noise_refiner." + std::to_string(i)]   = make_block(true);
+                blocks["context_refiner." + std::to_string(i)] = make_block(false);
+                blocks["sigvq_refiner." + std::to_string(i)]   = make_block(false);
+            }
+            for (int i = 0; i < config.num_layers; i++) {
+                blocks["layers." + std::to_string(i)] = make_block(true);
+            }
+
+            blocks["final_layer"] = std::make_shared<ZImage::FinalLayer>(config.hidden_size, config.patch_size, config.out_channels);
+        }
+
+        ggml_tensor* forward_core(GGMLRunnerContext* ctx,
+                                  ggml_tensor* x,
+                                  ggml_tensor* timestep,
+                                  ggml_tensor* context,
+                                  ggml_tensor* pe) {
+            auto x_embedder     = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
+            auto t_embedder     = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
+            auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
+            auto cap_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"]);
+            auto final_layer    = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"]);
+
+            auto txt_pad_token = params["cap_pad_token"];
+            auto img_pad_token = params["x_pad_token"];
+
+            int64_t N           = x->ne[2];
+            int64_t n_img_token = x->ne[1];
+            int64_t n_txt_token = context->ne[1];
+
+            // sdcpp's flow denoiser already hands over sigma * 1000, which is the range the
+            // reference reaches via its own t_scale, so no further scaling here.
+            auto t_emb = t_embedder->forward(ctx, timestep);
+
+            auto txt = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context));  // [N, n_txt_token, hidden_size]
+            auto img = x_embedder->forward(ctx, x);                                          // [N, n_img_token, hidden_size]
+            sd::ggml_graph_cut::mark_graph_cut(txt, "llada_image.prelude", "txt");
+            sd::ggml_graph_cut::mark_graph_cut(img, "llada_image.prelude", "img");
+            sd::ggml_graph_cut::mark_graph_cut(t_emb, "llada_image.prelude", "t_emb");
+
+            int64_t n_txt_pad_token = Rope::bound_mod(static_cast<int>(n_txt_token), ZImage::SEQ_MULTI_OF);
+            if (n_txt_pad_token > 0) {
+                auto txt_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, txt_pad_token, txt_pad_token->ne[0], n_txt_pad_token, N, 1);
+                txt                 = ggml_concat(ctx->ggml_ctx, txt, txt_pad_tokens, 1);
+            }
+
+            int64_t n_img_pad_token = Rope::bound_mod(static_cast<int>(n_img_token), ZImage::SEQ_MULTI_OF);
+            if (n_img_pad_token > 0) {
+                auto img_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, img_pad_token, img_pad_token->ne[0], n_img_pad_token, N, 1);
+                img                 = ggml_concat(ctx->ggml_ctx, img, img_pad_tokens, 1);
+            }
+
+            GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]);
+
+            auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]);
+            auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]);
+
+            for (int i = 0; i < config.num_refiner_layers; i++) {
+                auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
+
+                txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr);
+                sd::ggml_graph_cut::mark_graph_cut(txt, "llada_image.context_refiner." + std::to_string(i), "txt");
+            }
+
+            for (int i = 0; i < config.num_refiner_layers; i++) {
+                auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
+
+                img = block->forward(ctx, img, img_pe, nullptr, t_emb);
+                sd::ggml_graph_cut::mark_graph_cut(img, "llada_image.noise_refiner." + std::to_string(i), "img");
+            }
+
+            auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1);
+            sd::ggml_graph_cut::mark_graph_cut(txt_img, "llada_image.prelude", "txt_img");
+
+            for (int i = 0; i < config.num_layers; i++) {
+                auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
+
+                txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb);
+                sd::ggml_graph_cut::mark_graph_cut(txt_img, "llada_image.layers." + std::to_string(i), "txt_img");
+            }
+
+            txt_img = final_layer->forward(ctx, txt_img, t_emb);
+
+            return ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token);
+        }
+
+        ggml_tensor* pad_stream(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pad_token) {
+            int64_t n_pad = Rope::bound_mod(static_cast<int>(x->ne[1]), ZImage::SEQ_MULTI_OF);
+            if (n_pad == 0) {
+                return x;
+            }
+            auto pads = ggml_repeat_4d(ctx->ggml_ctx, pad_token, pad_token->ne[0], n_pad, x->ne[2], 1);
+            return ggml_concat(ctx->ggml_ctx, x, pads, 1);
+        }
+
+        // Editing runs one joint sequence carrying two timesteps: the caption and source latent
+        // are clean (t = 0) while the second caption copy and the target latent are noisy. adaLN
+        // is a linear map of the timestep embedding, so feeding a per-token embedding selects the
+        // right modulation exactly, without duplicating the modulation projections.
+        ggml_tensor* forward_editing(GGMLRunnerContext* ctx,
+                                     ggml_tensor* x,
+                                     ggml_tensor* timestep,
+                                     ggml_tensor* context,
+                                     ggml_tensor* semantic,
+                                     ggml_tensor* source_latent,
+                                     ggml_tensor* pe) {
+            ggml_context* gctx = ctx->ggml_ctx;
+
+            auto x_embedder     = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
+            auto t_embedder     = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
+            auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
+            auto cap_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"]);
+            auto sigvq_embed_0  = std::dynamic_pointer_cast<RMSNorm>(blocks["sigvq_embedder.0"]);
+            auto sigvq_embed_1  = std::dynamic_pointer_cast<Linear>(blocks["sigvq_embedder.1"]);
+            auto final_layer    = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"]);
+
+            auto t_noisy = t_embedder->forward(ctx, timestep);
+            auto t_clean = t_embedder->forward(ctx, ggml_scale(gctx, timestep, 0.f));
+
+            auto per_token = [&](ggml_tensor* emb, int64_t n) {
+                return ggml_repeat_4d(gctx, emb, emb->ne[0], n, 1, 1);
+            };
+
+            auto cap        = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context));
+            cap             = pad_stream(ctx, cap, params["cap_pad_token"]);
+            int64_t cap_len = cap->ne[1];
+            cap             = ggml_concat(gctx, cap, cap, 1);
+
+            auto src            = pad_stream(ctx, x_embedder->forward(ctx, source_latent), params["x_pad_token"]);
+            auto tgt_embed      = x_embedder->forward(ctx, x);
+            int64_t n_img_token = tgt_embed->ne[1];
+            auto tgt            = pad_stream(ctx, tgt_embed, params["x_pad_token"]);
+            int64_t img_len     = tgt->ne[1];
+            auto img            = ggml_concat(gctx, src, tgt, 1);
+
+            ggml_tensor* sig = nullptr;
+            int64_t sig_len  = 0;
+            if (semantic != nullptr) {
+                sig     = sigvq_embed_1->forward(ctx, sigvq_embed_0->forward(ctx, semantic));
+                sig     = pad_stream(ctx, sig, params["sigvq_pad_token"]);
+                sig_len = sig->ne[1];
+            }
+
+            GGML_ASSERT(cap_len * 2 + img_len * 2 + sig_len == pe->ne[3]);
+
+            auto cap_pe = ggml_ext_slice(gctx, pe, 3, 0, cap_len * 2);
+            auto img_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2);
+
+            auto img_adaln = ggml_concat(gctx, per_token(t_clean, img_len), per_token(t_noisy, img_len), 1);
+
+            for (int i = 0; i < config.num_refiner_layers; i++) {
+                auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
+                cap        = block->forward(ctx, cap, cap_pe, nullptr, nullptr);
+            }
+            for (int i = 0; i < config.num_refiner_layers; i++) {
+                auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
+                img        = block->forward(ctx, img, img_pe, nullptr, img_adaln);
+            }
+            if (sig != nullptr) {
+                auto sig_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]);
+                for (int i = 0; i < config.num_refiner_layers; i++) {
+                    auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["sigvq_refiner." + std::to_string(i)]);
+                    sig        = block->forward(ctx, sig, sig_pe, nullptr, nullptr);
+                }
+            }
+
+            auto seq = ggml_concat(gctx, cap, img, 1);
+
+            auto cap_adaln = ggml_concat(gctx, per_token(t_clean, cap_len), per_token(t_noisy, cap_len), 1);
+            auto seq_adaln = ggml_concat(gctx, cap_adaln, img_adaln, 1);
+            if (sig != nullptr) {
+                seq       = ggml_concat(gctx, seq, sig, 1);
+                seq_adaln = ggml_concat(gctx, seq_adaln, per_token(t_clean, sig_len), 1);
+            }
+
+            for (int i = 0; i < config.num_layers; i++) {
+                auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
+                seq        = block->forward(ctx, seq, pe, nullptr, seq_adaln);
+                sd::ggml_graph_cut::mark_graph_cut(seq, "llada_image.layers." + std::to_string(i), "seq");
+            }
+
+            seq = final_layer->forward(ctx, seq, seq_adaln);
+
+            // Only the target latent is denoised; the source half of the image stream is context.
+            // The stream is padded to SEQ_MULTI_OF, so drop the pad tokens: they are not part of
+            // the latent grid that unpatchify reconstructs.
+            int64_t target_start = cap_len * 2 + img_len;
+            return ggml_ext_slice(gctx, seq, 1, target_start, target_start + n_img_token);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timestep,
+                             ggml_tensor* context,
+                             ggml_tensor* pe) {
+            // x: [N, C, H, W]
+            // timestep: [N,]
+            // context: [N, L, cap_feat_dim]
+            // pe: [L, d_head/2, 2, 2]
+            // return: [N, C, H, W]
+            int64_t W = x->ne[0];
+            int64_t H = x->ne[1];
+
+            int patch_size = config.patch_size;
+
+            auto img = DiT::pad_and_patchify(ctx, x, patch_size, patch_size, false);
+
+            auto out = forward_core(ctx, img, timestep, context, pe);
+
+            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false);
+
+            // The reference pipeline negates the model output before the scheduler step.
+            return ggml_ext_scale(ctx->ggml_ctx, out, -1.f);
+        }
+    };
+
+    struct LLaDAImageRunner : public DiffusionModelRunner {
+    public:
+        LLaDAImageConfig config;
+        LLaDAImageModel llada_image;
+        std::vector<float> pe_vec;
+
+        LLaDAImageRunner(ggml_backend_t backend,
+                         const String2TensorStorage& tensor_storage_map      = {},
+                         const std::string prefix                            = "",
+                         std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(LLaDAImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
+            llada_image = LLaDAImageModel(config);
+            llada_image.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "llada_image";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            llada_image.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) {
+            ggml_cgraph* gf        = new_graph_custom(LLADA_IMAGE_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);
+
+            pe_vec      = Rope::gen_llada_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,
+                                                   config.theta,
+                                                   circular_y_enabled,
+                                                   circular_x_enabled,
+                                                   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 runner_ctx = get_context();
+
+            ggml_tensor* out = llada_image.forward(&runner_ctx, x, timesteps, context, pe);
+
+            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) {
+            // x: [N, in_channels, h, w]
+            // timesteps: [N, ]
+            // context: [N, max_position, cap_feat_dim]
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context);
+            };
+
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
+        }
+
+        ggml_cgraph* build_edit_graph(const sd::Tensor<float>& x_tensor,
+                                      const sd::Tensor<float>& timesteps_tensor,
+                                      const sd::Tensor<float>& context_tensor,
+                                      const sd::Tensor<float>& semantic_tensor,
+                                      const sd::Tensor<float>& source_tensor) {
+            ggml_cgraph* gf        = new_graph_custom(LLADA_IMAGE_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
+            ggml_tensor* context   = make_input(context_tensor);
+            ggml_tensor* semantic  = make_optional_input(semantic_tensor);
+            ggml_tensor* source    = make_input(source_tensor);
+            GGML_ASSERT(x->ne[3] == 1);
+
+            pe_vec      = Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
+                                                        static_cast<int>(x->ne[0]),
+                                                        config.patch_size,
+                                                        static_cast<int>(context->ne[1]),
+                                                   semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
+                                                        ZImage::SEQ_MULTI_OF,
+                                                        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 runner_ctx = get_context();
+
+            int64_t W   = x->ne[0];
+            int64_t H   = x->ne[1];
+            auto target = DiT::pad_and_patchify(&runner_ctx, x, config.patch_size, config.patch_size, false);
+            auto src    = DiT::pad_and_patchify(&runner_ctx, source, config.patch_size, config.patch_size, false);
+
+            auto out = llada_image.forward_editing(&runner_ctx, target, timesteps, context, semantic, src, pe);
+            out      = DiT::unpatchify_and_crop(runner_ctx.ggml_ctx, out, H, W, config.patch_size, config.patch_size, false);
+            out      = ggml_ext_scale(runner_ctx.ggml_ctx, out, -1.f);
+
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+
+            const auto* extra   = std::get_if<LLaDAImageDiffusionExtra>(&diffusion_params.extra);
+            bool has_semantic   = extra != nullptr && extra->semantic != nullptr && !extra->semantic->empty();
+            bool has_ref_latent = diffusion_params.ref_latents != nullptr && !diffusion_params.ref_latents->empty();
+            if (has_semantic && !has_ref_latent) {
+                LOG_WARN("llada_image: SigVQ features without a reference latent are not supported; falling back to text to image");
+            }
+            if (has_ref_latent) {
+                const auto& source = diffusion_params.ref_latents->front();
+                if (source.shape() != diffusion_params.x->shape()) {
+                    LOG_ERROR("llada_image: reference latent must match the target shape; use resize_vae_to_target=1");
+                    return {};
+                }
+                auto get_graph = [&]() -> ggml_cgraph* {
+                    return build_edit_graph(*diffusion_params.x,
+                                            *diffusion_params.timesteps,
+                                            tensor_or_empty(diffusion_params.context),
+                                            tensor_or_empty(extra != nullptr ? extra->semantic : nullptr),
+                                            source);
+                };
+                return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false),
+                                                       diffusion_params.x->dim());
+            }
+
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context));
+        }
+    };
+
+}  // namespace LLaDAImage
+
+#endif  // __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__

src/model/diffusion/model.hpp

@@ -39,6 +39,9 @@ const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
     {"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
     {"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}},
     {"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}},
+    // pass_to_vlm routes the reference image to the conditioner, which is where LLaDA-Image's
+    // SigVQ encoder lives; it does its own half-resolution resize.
+    {"llada_image", {true, true, Rope::RefIndexMode::FIXED, true, true, -1, RefImageResizeMode::NONE, -1, -1, true}},
     {"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}},
 };
 
@@ -66,6 +69,10 @@ struct AnimaDiffusionExtra {
     const sd::Tensor<float>* t5_weights = nullptr;
 };
 
+struct QwenImage21DiffusionExtra {
+    const sd::Tensor<int32_t>* image_slots = nullptr;
+};
+
 struct WanDiffusionExtra {
     const sd::Tensor<float>* vace_context = nullptr;
     float vace_strength                   = 1.f;
@@ -127,18 +134,25 @@ struct HunyuanVideoDiffusionExtra {
     const sd::Tensor<float>* timestep_r = nullptr;
 };
 
+struct LLaDAImageDiffusionExtra {
+    // SigVQ semantic features of the reference image; present only in editing mode.
+    const sd::Tensor<float>* semantic = nullptr;
+};
+
 using DiffusionExtraParams = std::variant<std::monostate,
                                           UNetDiffusionExtra,
                                           SkipLayerDiffusionExtra,
                                           FluxDiffusionExtra,
                                           AnimaDiffusionExtra,
+                                          QwenImage21DiffusionExtra,
                                           WanDiffusionExtra,
                                           HiDreamO1DiffusionExtra,
                                           LTXAVDiffusionExtra,
                                           MiniMaxH3DiffusionExtra,
                                           MiniT2IDiffusionExtra,
                                           SenseNovaU1DiffusionExtra,
-                                          HunyuanVideoDiffusionExtra>;
+                                          HunyuanVideoDiffusionExtra,
+                                          LLaDAImageDiffusionExtra>;
 
 struct DiffusionParams {
     const sd::Tensor<float>* x                        = nullptr;

src/model/diffusion/qwen_image_2_1.hpp

@@ -0,0 +1,384 @@
+#ifndef __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__
+#define __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__
+
+#include "model/diffusion/qwen_image.hpp"
+
+namespace Qwen {
+
+    struct QwenImage21Config {
+        int64_t in_channels       = 64;
+        int64_t out_channels      = 64;
+        int64_t hidden_size       = 4096;
+        int64_t context_dim       = 4096;
+        int64_t head_dim          = 128;
+        int64_t intermediate_size = 12288;
+        int num_layers            = 32;
+        bool fused_mlp            = false;
+        std::vector<int> axes_dim = {16, 56, 56};
+
+        static QwenImage21Config detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) {
+            QwenImage21Config 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("img_in.weight")) {
+                config.in_channels = w->ne[0];
+                config.hidden_size = w->ne[1];
+            }
+            if (auto w = find("proj_out.weight")) {
+                config.out_channels = w->ne[1];
+            }
+            if (auto w = find("txt_in.in_layer.weight")) {
+                config.context_dim = w->ne[0];
+            }
+            if (auto w = find("transformer_blocks.0.attn.norm_q.weight")) {
+                config.head_dim = w->ne[0];
+            }
+            if (auto w = find("transformer_blocks.0.img_mlp.gate_up.weight")) {
+                config.intermediate_size = w->ne[1] / 2;
+                config.fused_mlp         = true;
+            } else if (auto w = find("transformer_blocks.0.img_mlp.proj.weight")) {
+                config.intermediate_size = w->ne[1];
+            }
+            int layers                     = 0;
+            const std::string block_prefix = prefix + ".transformer_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("qwen_image_2_1: layers = %d, hidden_size = %" PRId64 ", context_dim = %" PRId64,
+                            layers, config.hidden_size, config.context_dim);
+            }
+            return config;
+        }
+    };
+
+    struct QwenImage21Segment {
+        int64_t start;
+        int64_t end;
+        int64_t context_start;
+        int image_index;
+    };
+
+    struct QwenImage21Layout {
+        std::vector<QwenImage21Segment> segments;
+        std::vector<std::vector<float>> positions;
+        int64_t prefix_length = 0;
+
+        static QwenImage21Layout build(int64_t text_length,
+                                       const sd::Tensor<int32_t>& image_slots,
+                                       const std::vector<std::pair<int64_t, int64_t>>& image_shapes) {
+            if (image_shapes.empty() || (!image_slots.empty() && image_slots.numel() != text_length)) {
+                throw std::runtime_error("Qwen Image 2.1: invalid image token layout");
+            }
+            QwenImage21Layout layout;
+            int64_t position  = 0;
+            int next_image    = 0;
+            auto append_image = [&](int index, int64_t context_start) {
+                auto [height, width] = image_shapes[index];
+                int64_t start        = static_cast<int64_t>(layout.positions.size());
+                layout.segments.push_back({start, start + height * width, context_start, index});
+                for (int64_t h = 0; h < height; ++h) {
+                    for (int64_t w = 0; w < width; ++w) {
+                        layout.positions.push_back({static_cast<float>(position),
+                                                    static_cast<float>(h - (height - height / 2)),
+                                                    static_cast<float>(w - (width - width / 2))});
+                    }
+                }
+                position += std::max(height, width);
+            };
+            for (int64_t i = 0; i < text_length;) {
+                int tag       = image_slots.empty() ? 0 : image_slots[i];
+                int64_t begin = i++;
+                while (i < text_length && (image_slots.empty() ? 0 : image_slots[i]) == tag) {
+                    ++i;
+                }
+                if (tag != 0) {
+                    if (tag != next_image + 1 || next_image + 1 >= static_cast<int>(image_shapes.size()) ||
+                        (i - begin) * 4 != image_shapes[next_image].first * image_shapes[next_image].second) {
+                        throw std::runtime_error("Qwen Image 2.1: vision slots and reference latents must have matching sizes");
+                    }
+                    append_image(next_image++, begin);
+                } else {
+                    int64_t start = static_cast<int64_t>(layout.positions.size());
+                    layout.segments.push_back({start, start + i - begin, begin, -1});
+                    for (int64_t j = begin; j < i; ++j, ++position) {
+                        float p = static_cast<float>(position);
+                        layout.positions.push_back({p, p, p});
+                    }
+                }
+            }
+            if (next_image + 1 != static_cast<int>(image_shapes.size())) {
+                throw std::runtime_error("Qwen Image 2.1: missing reference image slots");
+            }
+            layout.prefix_length = static_cast<int64_t>(layout.positions.size());
+            append_image(next_image, text_length);
+            return layout;
+        }
+    };
+
+    class QwenImage21ZeroCenterRMSNorm : public RMSNorm {
+    public:
+        using RMSNorm::RMSNorm;
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto weight = params["weight"];
+            if (ctx->weight_adapter) {
+                weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
+            }
+            weight = ggml_scale_bias(ctx->ggml_ctx, weight, 1.f, 1.f);
+            return ggml_mul(ctx->ggml_ctx, ggml_rms_norm(ctx->ggml_ctx, x, eps), weight);
+        }
+    };
+
+    class QwenImage21TextProjection : public GGMLBlock {
+    public:
+        QwenImage21TextProjection(const QwenImage21Config& config) {
+            blocks["text_norm"] = std::make_shared<QwenImage21ZeroCenterRMSNorm>(config.context_dim, 1e-6f);
+            blocks["in_layer"]  = std::make_shared<Linear>(config.context_dim, config.hidden_size, false);
+            blocks["out_layer"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            x = std::dynamic_pointer_cast<QwenImage21ZeroCenterRMSNorm>(blocks["text_norm"])->forward(ctx, x);
+            x = std::dynamic_pointer_cast<Linear>(blocks["in_layer"])->forward(ctx, x);
+            x = ggml_ext_gelu(ctx->ggml_ctx, x);
+            return std::dynamic_pointer_cast<Linear>(blocks["out_layer"])->forward(ctx, x);
+        }
+    };
+
+    class QwenImage21Attention : public QwenImageAttention {
+    public:
+        QwenImage21Attention(const QwenImage21Config& config)
+            : QwenImageAttention(config.hidden_size, config.head_dim, config.hidden_size / config.head_dim, 0, 0, false, false) {
+            for (const auto* name : {"add_q_proj", "add_k_proj", "add_v_proj", "norm_added_q", "norm_added_k", "to_add_out"}) {
+                blocks.erase(name);
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector<QwenImage21Segment>& segments, const std::vector<ggml_tensor*>& masks) {
+            int64_t heads = x->ne[0] / dim_head;
+            auto project  = [&](const char* name) {
+                auto h = std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x);
+                return ggml_reshape_4d(ctx->ggml_ctx, h, dim_head, heads, x->ne[1], x->ne[2]);
+            };
+            auto q              = project("to_q");
+            auto k              = project("to_k");
+            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);
+            ggml_tensor* result = nullptr;
+            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);
+                result              = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
+            }
+            auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
+            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
+                to_out->set_force_prec_f32(true);
+            }
+            return to_out->forward(ctx, result);
+        }
+    };
+
+    class QwenImage21TransformerBlock : public GGMLBlock {
+    public:
+        QwenImage21TransformerBlock(const QwenImage21Config& config) {
+            blocks["img_norm1"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
+            blocks["img_norm2"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
+            blocks["attn"]      = std::make_shared<QwenImage21Attention>(config);
+            if (config.fused_mlp) {
+                blocks["img_mlp.gate_up"] = std::make_shared<Linear>(config.hidden_size, 2 * config.intermediate_size, false);
+            } else {
+                blocks["img_mlp.proj"]       = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, false);
+                blocks["img_mlp.gate_layer"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, false);
+            }
+            blocks["img_mlp.out"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, false);
+        }
+
+        static ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* params, int64_t prefix_length, bool gate = false) {
+            auto rows  = ggml_ext_chunk(ctx, params, 2, 1);
+            auto apply = [&](ggml_tensor* part, ggml_tensor* row) {
+                row = gate ? ggml_tanh(ctx, row) : ggml_scale_bias(ctx, row, 1.f, 1.f);
+                return ggml_mul(ctx, part, row);
+            };
+            auto target = apply(ggml_ext_slice(ctx, x, 1, prefix_length, x->ne[1]), rows[0]);
+            if (prefix_length == 0) {
+                return target;
+            }
+            auto prefix = apply(ggml_ext_slice(ctx, x, 1, 0, prefix_length), rows[1]);
+            return ggml_concat(ctx, prefix, target, 1);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector<ggml_tensor*>& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks) {
+            auto h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"])->forward(ctx, x);
+            h      = modulate(ctx->ggml_ctx, h, modulation[0], layout.prefix_length);
+            h      = std::dynamic_pointer_cast<QwenImage21Attention>(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks);
+            x      = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], layout.prefix_length, true));
+            h      = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"])->forward(ctx, x);
+            h      = modulate(ctx->ggml_ctx, h, modulation[2], layout.prefix_length);
+            ggml_tensor* gate;
+            auto fused = blocks.find("img_mlp.gate_up");
+            if (fused != blocks.end()) {
+                auto gate_up = std::dynamic_pointer_cast<Linear>(fused->second)->forward(ctx, h);
+                auto parts   = ggml_ext_chunk(ctx->ggml_ctx, gate_up, 2, 0);
+                gate         = parts[0];
+                h            = parts[1];
+            } else {
+                gate = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.gate_layer"])->forward(ctx, h);
+                h    = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.proj"])->forward(ctx, h);
+            }
+            h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate));
+            h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.out"])->forward(ctx, h);
+            return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], layout.prefix_length, true));
+        }
+    };
+
+    class QwenImage21Model : public GGMLBlock {
+        QwenImage21Config config;
+
+    public:
+        QwenImage21Model(const QwenImage21Config& config)
+            : config(config) {
+            blocks["time_text_embed.timestep_embedder"] = std::make_shared<TimestepEmbedding>(256, config.hidden_size, 0, 0, false);
+            blocks["txt_in"]                            = std::make_shared<QwenImage21TextProjection>(config);
+            blocks["img_in"]                            = std::make_shared<Linear>(config.in_channels, config.hidden_size, false);
+            blocks["modulation.1"]                      = std::make_shared<Linear>(config.hidden_size, 4 * config.hidden_size, false);
+            blocks["norm_out.linear"]                   = std::make_shared<Linear>(config.hidden_size, config.hidden_size, false);
+            blocks["norm_out.norm"]                     = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
+            blocks["proj_out"]                          = std::make_shared<Linear>(config.hidden_size, config.out_channels, false);
+            for (int i = 0; i < config.num_layers; ++i) {
+                blocks["transformer_blocks." + std::to_string(i)] = std::make_shared<QwenImage21TransformerBlock>(config);
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector<ggml_tensor*>& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks) {
+            auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), 0);
+            // Runtime flow timesteps already use the [0, 1000] scale.
+            time               = ggml_ext_timestep_embedding(ctx->ggml_ctx, time, 256, 10000, 1.f);
+            time               = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["time_text_embed.timestep_embedder"])->forward(ctx, time);
+            time               = ggml_silu(ctx->ggml_ctx, time);
+            auto modulation    = std::dynamic_pointer_cast<Linear>(blocks["modulation.1"])->forward(ctx, time);
+            auto mod           = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0);
+            auto text          = std::dynamic_pointer_cast<QwenImage21TextProjection>(blocks["txt_in"])->forward(ctx, context);
+            auto img_in        = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
+            ggml_tensor* joint = nullptr;
+            for (const auto& segment : layout.segments) {
+                ggml_tensor* h;
+                if (segment.image_index < 0) {
+                    h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start,
+                                       segment.context_start + segment.end - segment.start);
+                } else {
+                    auto image = segment.image_index == static_cast<int>(refs.size()) ? x : refs[segment.image_index];
+                    h          = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1));
+                }
+                joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1);
+            }
+            sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.prelude", "joint");
+            for (int i = 0; i < config.num_layers; ++i) {
+                auto block = std::dynamic_pointer_cast<QwenImage21TransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
+                joint      = block->forward(ctx, joint, mod, pe, layout, masks);
+                sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.transformer_blocks." + std::to_string(i), "joint");
+            }
+            joint      = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]);
+            auto scale = std::dynamic_pointer_cast<Linear>(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]);
+            joint      = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out.norm"])->forward(ctx, joint);
+            joint      = ggml_mul(ctx->ggml_ctx, joint, ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f));
+            joint      = std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->forward(ctx, joint);
+            return DiT::unpatchify_and_crop(ctx->ggml_ctx, joint, x->ne[1], x->ne[0], 1, 1);
+        }
+    };
+
+    struct QwenImage21Runner : public DiffusionModelRunner {
+        QwenImage21Config config;
+        QwenImage21Model model;
+        std::vector<float> pe_data;
+        std::vector<sd::Tensor<float>> mask_data;
+
+        QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(QwenImage21Config::detect_from_weights(weights, prefix)),
+              model(config) {
+            model.init(params_ctx, weights, prefix);
+        }
+
+        std::string get_desc() override { return "qwen_image_2_1"; }
+
+        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& x       = tensor_or_empty(inputs.x);
+            const auto& context = tensor_or_empty(inputs.context);
+            if (x.empty() || context.empty() || context.dim() < 2 || context.shape()[0] != config.context_dim ||
+                tensor_or_empty(inputs.timesteps).numel() != 1 ||
+                x.dim() != 4 || x.shape()[3] != 1 || x.shape()[2] != config.in_channels) {
+                LOG_ERROR("Qwen Image 2.1 requires an image latent and text conditioning with batch size 1");
+                return {};
+            }
+            static const std::vector<sd::Tensor<float>> empty_refs;
+            const auto& refs = inputs.ref_latents && inputs.ref_image_params.pass_to_dit ? *inputs.ref_latents : empty_refs;
+            std::vector<std::pair<int64_t, int64_t>> shapes;
+            for (const auto& ref : refs) {
+                if (ref.dim() != 4 || ref.shape()[2] != config.in_channels || ref.shape()[3] != 1) {
+                    LOG_ERROR("Qwen Image 2.1: invalid reference latent shape");
+                    return {};
+                }
+                shapes.emplace_back(ref.shape()[1], ref.shape()[0]);
+            }
+            shapes.emplace_back(x.shape()[1], x.shape()[0]);
+            const auto* extra = std::get_if<QwenImage21DiffusionExtra>(&inputs.extra);
+            QwenImage21Layout layout;
+            try {
+                layout = QwenImage21Layout::build(context.shape()[1], tensor_or_empty(extra ? extra->image_slots : nullptr), shapes);
+            } catch (const std::exception& error) {
+                LOG_ERROR("%s", error.what());
+                return {};
+            }
+            pe_data = Rope::embed_nd(layout.positions, 1, 10000.f, config.axes_dim);
+            mask_data.clear();
+            for (const auto& segment : layout.segments) {
+                sd::Tensor<float> mask;
+                if (segment.image_index < 0) {
+                    mask = sd::Tensor<float>::zeros({segment.end, segment.end - segment.start});
+                    for (int64_t q = segment.start; q < segment.end; ++q) {
+                        for (int64_t k = q + 1; k < segment.end; ++k) {
+                            mask[k + segment.end * (q - segment.start)] = -INFINITY;
+                        }
+                    }
+                }
+                mask_data.push_back(std::move(mask));
+            }
+            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());
+                set_backend_tensor_data(pe, pe_data.data());
+                std::vector<ggml_tensor*> masks, ref_inputs;
+                for (const auto& mask : mask_data) {
+                    masks.push_back(mask.empty() ? nullptr : make_input(mask));
+                }
+                for (const auto& ref : refs) {
+                    ref_inputs.push_back(make_input(ref));
+                }
+                auto ctx = get_context();
+                auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), make_input(context),
+                                         ref_inputs, pe, layout, masks);
+                ggml_build_forward_expand(graph, out);
+                return graph;
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim());
+        }
+    };
+}
+
+#endif  // __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__

src/model/diffusion/z_image.hpp

@@ -131,16 +131,30 @@ namespace ZImage {
         int64_t num_heads;
         int64_t num_kv_heads;
         bool qk_norm;
+        bool split_qkv;
 
     public:
-        JointAttention(int64_t hidden_size, int64_t head_dim, int64_t num_heads, int64_t num_kv_heads, bool qk_norm)
-            : head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm) {
-            blocks["qkv"] = std::make_shared<Linear>(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false);
-            float scale   = 1.f;
-            blocks["out"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
+        JointAttention(int64_t hidden_size,
+                       int64_t head_dim,
+                       int64_t num_heads,
+                       int64_t num_kv_heads,
+                       bool qk_norm,
+                       bool norm_elementwise_affine = true,
+                       bool split_qkv               = false)
+            : 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) {
+                blocks["to_q"]     = std::make_shared<Linear>(hidden_size, num_heads * head_dim, false);
+                blocks["to_k"]     = std::make_shared<Linear>(hidden_size, num_kv_heads * head_dim, false);
+                blocks["to_v"]     = std::make_shared<Linear>(hidden_size, num_kv_heads * head_dim, false);
+                blocks["to_out.0"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
+            } else {
+                blocks["qkv"] = std::make_shared<Linear>(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false);
+                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);
-                blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim);
+                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);
             }
         }
 
@@ -151,8 +165,35 @@ namespace ZImage {
             // x: [N, n_token, hidden_size]
             int64_t n_token = x->ne[1];
             int64_t N       = x->ne[2];
-            auto qkv_proj   = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
-            auto out_proj   = std::dynamic_pointer_cast<Linear>(blocks["out"]);
+            auto out_proj   = std::dynamic_pointer_cast<Linear>(blocks[split_qkv ? "to_out.0" : "out"]);
+
+            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);
+                    v_proj->set_force_prec_f32(true);
+                }
+
+                auto q = ggml_reshape_4d(ctx->ggml_ctx, q_proj->forward(ctx, x), head_dim, num_heads, n_token, N);
+                auto k = ggml_reshape_4d(ctx->ggml_ctx, k_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N);
+                auto v = ggml_reshape_4d(ctx->ggml_ctx, v_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N);
+
+                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);
+                }
+
+                auto out = Rope::attention(ctx, q, k, v, pe, mask, 1.f / 128.f);
+                return out_proj->forward(ctx, out);
+            }
+
+            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);
@@ -252,9 +293,12 @@ namespace ZImage {
                                             ggml_tensor* x,
                                             ggml_tensor* scale) {
         // x: [N, L, C]
-        // scale: [N, C]
-        scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]);  // [N, 1, C]
-        x     = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
+        // scale: [N, C], or [N, L, C] when the caller modulates per token (LLaDA-Image editing
+        // feeds a per-token timestep embedding so each segment carries its own modulation).
+        if (scale->ne[1] != x->ne[1]) {
+            scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]);  // [N, 1, C]
+        }
+        x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
         return x;
     }
 
@@ -272,14 +316,16 @@ namespace ZImage {
                               float ffn_dim_multiplier,
                               float norm_eps,
                               bool qk_norm,
-                              bool modulation = true)
+                              bool modulation              = true,
+                              bool norm_elementwise_affine = true,
+                              bool split_qkv               = false)
             : modulation(modulation) {
-            blocks["attention"]       = std::make_shared<JointAttention>(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm);
+            blocks["attention"]       = std::make_shared<JointAttention>(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm, norm_elementwise_affine, split_qkv);
             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);
-            blocks["ffn_norm1"]       = std::make_shared<RMSNorm>(hidden_size, norm_eps);
-            blocks["attention_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
-            blocks["ffn_norm2"]       = std::make_shared<RMSNorm>(hidden_size, norm_eps);
+            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);
+            blocks["attention_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
+            blocks["ffn_norm2"]       = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
             if (modulation) {
                 blocks["adaLN_modulation.0"] = std::make_shared<Linear>(MIN(hidden_size, ADALN_EMBED_DIM), 4 * hidden_size);
             }

src/model/te/llada_image_te.hpp

@@ -0,0 +1,604 @@
+#ifndef __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__
+#define __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__
+
+#include <algorithm>
+#include <array>
+#include <cmath>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "model/common/ggml_block.hpp"
+#include "model_loader.h"
+
+// The conditioning components LLaDA-Image puts around its LLaDA2-MoE backbone.
+// Ref: LLaDAImageQueryFormerModel / LLaDAImageTextProjectionModel in
+// https://github.com/inclusionAI/LLaDA-Image/blob/main/src/models/transformer_llada_image.py
+//
+// QueryFormer turns the LLaDA token embeddings into 256 learned queries that the pipeline
+// appends to the backbone input; TextProjection maps the backbone hidden states to the
+// denoiser's caption dimension. Neither uses RoPE, and every norm is parameter-free.
+// Both MLPs use the tanh GELU approximation, so ggml_gelu (not ggml_gelu_erf).
+//
+// SigVQ is the editing-only image encoder: a 40-layer ViT whose output is quantized against a
+// 16384-entry codebook, with the resulting ids embedded and projected into the semantic features
+// the denoiser consumes. Its MLP uses the exact erf GELU, unlike the two above.
+
+namespace LLaDAImageTE {
+    constexpr int LLADA_IMAGE_TE_GRAPH_SIZE = 16384;
+
+    struct QueryFormerConfig {
+        int64_t num_queries       = 256;
+        int64_t hidden_size       = 2048;
+        int64_t num_layers        = 1;
+        int64_t num_heads         = 16;
+        int64_t intermediate_size = 8192;
+        float norm_eps            = 1e-6f;
+    };
+
+    struct TextProjectionConfig {
+        int64_t hidden_size       = 2048;
+        int64_t intermediate_size = 8960;
+        int64_t num_layers        = 6;
+        int64_t num_heads         = 32;
+        int64_t projection_dim    = 2560;
+        float norm_eps            = 1e-6f;
+    };
+
+    // Cross-attention with a single fused in_proj over q (from the queries) and k/v (from the
+    // token embeddings). The checkpoint stores in_proj as one [3*hidden, hidden] parameter.
+    struct QueryAttention : public GGMLBlock {
+    protected:
+        int64_t hidden_size;
+        int64_t num_heads;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         std::string prefix                             = "") override {
+            GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
+            enum ggml_type wtype     = get_type(prefix + "in_proj_weight", tensor_storage_map, GGML_TYPE_F32);
+            params["in_proj_weight"] = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size * 3);
+            params["in_proj_bias"]   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size * 3);
+        }
+
+    public:
+        QueryAttention(int64_t hidden_size, int64_t num_heads)
+            : hidden_size(hidden_size), num_heads(num_heads) {
+            blocks["out_proj"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* query,
+                             ggml_tensor* context,
+                             ggml_tensor* mask = nullptr) {
+            // query: [N, num_queries, hidden_size], context: [N, n_token, hidden_size]
+            ggml_context* gctx = ctx->ggml_ctx;
+            auto out_proj      = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
+
+            auto w = params["in_proj_weight"];
+            auto b = params["in_proj_bias"];
+
+            auto slice_w = [&](int64_t index) {
+                return ggml_ext_slice(gctx, w, 1, index * hidden_size, (index + 1) * hidden_size);
+            };
+            auto slice_b = [&](int64_t index) {
+                return ggml_ext_slice(gctx, b, 0, index * hidden_size, (index + 1) * hidden_size);
+            };
+
+            auto q = ggml_ext_linear(gctx, query, slice_w(0), slice_b(0));
+            auto k = ggml_ext_linear(gctx, context, slice_w(1), slice_b(1));
+            auto v = ggml_ext_linear(gctx, context, slice_w(2), slice_b(2));
+
+            auto x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask);  // [N, num_queries, hidden_size]
+            return out_proj->forward(ctx, x);
+        }
+    };
+
+    struct QueryFormerBlock : public GGMLBlock {
+    protected:
+        QueryFormerConfig config;
+
+    public:
+        QueryFormerBlock(const QueryFormerConfig& config)
+            : config(config) {
+            blocks["norm_q"]     = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
+            blocks["norm_k"]     = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
+            blocks["cross_attn"] = std::make_shared<QueryAttention>(config.hidden_size, config.num_heads);
+            blocks["norm1"]      = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
+            blocks["mlp.fc1"]    = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, true);
+            blocks["mlp.fc2"]    = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* query,
+                             ggml_tensor* context,
+                             ggml_tensor* mask = nullptr) {
+            auto norm_q     = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_q"]);
+            auto norm_k     = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_k"]);
+            auto cross_attn = std::dynamic_pointer_cast<QueryAttention>(blocks["cross_attn"]);
+            auto norm1      = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
+            auto fc1        = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
+            auto fc2        = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"]);
+
+            // The reference overwrites query_embeds with its normalized value before the
+            // residual add, so both residuals here are on normalized activations.
+            query      = norm_q->forward(ctx, query);
+            auto ctx_n = norm_k->forward(ctx, context);
+            query      = ggml_add(ctx->ggml_ctx, query, cross_attn->forward(ctx, query, ctx_n, mask));
+            query      = norm1->forward(ctx, query);
+
+            auto h = fc1->forward(ctx, query);
+            h      = ggml_gelu(ctx->ggml_ctx, h);
+            h      = fc2->forward(ctx, h);
+            return ggml_add(ctx->ggml_ctx, query, h);
+        }
+    };
+
+    struct QueryFormerModel : public GGMLBlock {
+    protected:
+        QueryFormerConfig config;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            params["meta_queries"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, config.num_queries);
+        }
+
+    public:
+        QueryFormerModel() = default;
+        QueryFormerModel(const QueryFormerConfig& config)
+            : config(config) {
+            for (int i = 0; i < config.num_layers; i++) {
+                blocks["query_blocks." + std::to_string(i)] = std::make_shared<QueryFormerBlock>(config);
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* inputs_embeds,
+                             ggml_tensor* mask = nullptr) {
+            // inputs_embeds: [N, n_token, hidden_size] -> [N, num_queries, hidden_size]
+            auto query = params["meta_queries"];
+            query      = ggml_reshape_3d(ctx->ggml_ctx, query, config.hidden_size, config.num_queries, 1);
+
+            for (int i = 0; i < config.num_layers; i++) {
+                auto block = std::dynamic_pointer_cast<QueryFormerBlock>(blocks["query_blocks." + std::to_string(i)]);
+                query      = block->forward(ctx, query, inputs_embeds, mask);
+            }
+            return query;
+        }
+    };
+
+    struct TextProjectionAttention : public GGMLBlock {
+    protected:
+        int64_t num_heads;
+        int64_t head_dim;
+
+    public:
+        TextProjectionAttention(const TextProjectionConfig& config)
+            : num_heads(config.num_heads), head_dim(config.hidden_size / config.num_heads) {
+            blocks["q_proj"]   = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
+            blocks["k_proj"]   = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
+            blocks["v_proj"]   = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
+            blocks["out_proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
+            blocks["q_norm"]   = std::make_shared<RMSNorm>(head_dim, config.norm_eps, false);
+            blocks["k_norm"]   = std::make_shared<RMSNorm>(head_dim, config.norm_eps, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            // x: [N, n_token, hidden_size]
+            ggml_context* gctx = ctx->ggml_ctx;
+            int64_t n_token    = x->ne[1];
+            int64_t N          = x->ne[2];
+
+            auto q_proj   = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
+            auto k_proj   = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
+            auto v_proj   = std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
+            auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
+            auto q_norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
+            auto k_norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
+
+            auto q = q_proj->forward(ctx, x);
+            auto k = k_proj->forward(ctx, x);
+            auto v = v_proj->forward(ctx, x);
+
+            q = ggml_reshape_4d(gctx, q, head_dim, num_heads, n_token, N);
+            k = ggml_reshape_4d(gctx, k, head_dim, num_heads, n_token, N);
+            q = q_norm->forward(ctx, q);
+            k = k_norm->forward(ctx, k);
+            q = ggml_reshape_3d(gctx, q, head_dim * num_heads, n_token, N);
+            k = ggml_reshape_3d(gctx, k, head_dim * num_heads, n_token, N);
+
+            auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads);
+            return out_proj->forward(ctx, out);
+        }
+    };
+
+    struct TextProjectionBlock : public GGMLBlock {
+    public:
+        TextProjectionBlock(const TextProjectionConfig& config) {
+            blocks["self_attn"]   = std::make_shared<TextProjectionAttention>(config);
+            blocks["layer_norm1"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps, false);
+            blocks["layer_norm2"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps, false);
+            blocks["mlp.fc1"]     = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, true);
+            blocks["mlp.fc2"]     = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto self_attn   = std::dynamic_pointer_cast<TextProjectionAttention>(blocks["self_attn"]);
+            auto layer_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["layer_norm1"]);
+            auto layer_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["layer_norm2"]);
+            auto fc1         = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
+            auto fc2         = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"]);
+
+            x = ggml_add(ctx->ggml_ctx, x, self_attn->forward(ctx, layer_norm1->forward(ctx, x)));
+
+            auto h = fc1->forward(ctx, layer_norm2->forward(ctx, x));
+            h      = ggml_gelu(ctx->ggml_ctx, h);
+            h      = fc2->forward(ctx, h);
+            return ggml_add(ctx->ggml_ctx, x, h);
+        }
+    };
+
+    struct TextProjectionModel : public GGMLBlock {
+    protected:
+        TextProjectionConfig config;
+
+    public:
+        TextProjectionModel() = default;
+        TextProjectionModel(const TextProjectionConfig& config)
+            : config(config) {
+            for (int i = 0; i < config.num_layers; i++) {
+                blocks["layers." + std::to_string(i)] = std::make_shared<TextProjectionBlock>(config);
+            }
+            blocks["projector"] = std::make_shared<Linear>(config.hidden_size, config.projection_dim, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            // x: [N, n_token, hidden_size] -> [N, n_token, projection_dim]
+            for (int i = 0; i < config.num_layers; i++) {
+                auto block = std::dynamic_pointer_cast<TextProjectionBlock>(blocks["layers." + std::to_string(i)]);
+                x          = block->forward(ctx, x);
+            }
+            auto projector = std::dynamic_pointer_cast<Linear>(blocks["projector"]);
+            return projector->forward(ctx, x);
+        }
+    };
+
+    struct SigVQConfig {
+        int64_t image_size         = 2048;
+        int patch_size             = 16;
+        int64_t in_channels        = 3;
+        int64_t hidden_size        = 1536;
+        int64_t intermediate_size  = 6144;
+        int64_t num_layers         = 40;
+        int64_t num_heads          = 16;
+        int64_t codebook_size      = 16384;
+        int64_t codebook_embed_dim = 2048;
+        int64_t semantic_embed_dim = 4096;
+        float norm_eps             = 1e-6f;
+    };
+
+    struct SigVQAttention : public GGMLBlock {
+    protected:
+        int64_t num_heads;
+        int64_t head_dim;
+
+    public:
+        SigVQAttention(const SigVQConfig& config)
+            : num_heads(config.num_heads), head_dim(config.hidden_size / config.num_heads) {
+            blocks["qkv"]  = std::make_shared<Linear>(config.hidden_size, config.hidden_size * 3, true);
+            blocks["proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            // x: [N, n_token, hidden_size]
+            ggml_context* gctx = ctx->ggml_ctx;
+            auto qkv_proj      = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
+            auto out_proj      = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
+
+            int64_t hidden_size = num_heads * head_dim;
+            auto qkv            = qkv_proj->forward(ctx, x);
+            auto q              = ggml_ext_slice(gctx, qkv, 0, 0, hidden_size);
+            auto k              = ggml_ext_slice(gctx, qkv, 0, hidden_size, hidden_size * 2);
+            auto v              = ggml_ext_slice(gctx, qkv, 0, hidden_size * 2, hidden_size * 3);
+
+            auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads);
+            return out_proj->forward(ctx, out);
+        }
+    };
+
+    struct SigVQBlock : public GGMLBlock {
+    public:
+        SigVQBlock(const SigVQConfig& config) {
+            blocks["norm1"]   = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps);
+            blocks["norm2"]   = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps);
+            blocks["attn"]    = std::make_shared<SigVQAttention>(config);
+            blocks["mlp.fc1"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, true);
+            blocks["mlp.fc2"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
+            auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
+            auto attn  = std::dynamic_pointer_cast<SigVQAttention>(blocks["attn"]);
+            auto fc1   = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
+            auto fc2   = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"]);
+
+            x      = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, norm1->forward(ctx, x)));
+            auto h = fc1->forward(ctx, norm2->forward(ctx, x));
+            h      = ggml_gelu_erf(ctx->ggml_ctx, h);
+            h      = fc2->forward(ctx, h);
+            return ggml_add(ctx->ggml_ctx, x, h);
+        }
+    };
+
+    struct SigVQModel : public GGMLBlock {
+    protected:
+        SigVQConfig config;
+
+    public:
+        SigVQModel() = default;
+        SigVQModel(const SigVQConfig& config)
+            : config(config) {
+            blocks["visual.patch_embed.proj"] = std::make_shared<Conv2d>(config.in_channels,
+                                                                         config.hidden_size,
+                                                                         std::make_pair(config.patch_size, config.patch_size),
+                                                                         std::make_pair(config.patch_size, config.patch_size));
+            for (int i = 0; i < config.num_layers; i++) {
+                blocks["visual.blocks." + std::to_string(i)] = std::make_shared<SigVQBlock>(config);
+            }
+            blocks["vqmodel.quant_conv"]         = std::make_shared<Conv2d>(config.hidden_size,
+                                                                    config.codebook_embed_dim,
+                                                                    std::make_pair(1, 1));
+            blocks["prior_projector.net.0.proj"] = std::make_shared<Linear>(config.semantic_embed_dim, config.semantic_embed_dim, true);
+            blocks["prior_projector.net.2"]      = std::make_shared<Linear>(config.semantic_embed_dim, config.semantic_embed_dim, true);
+        }
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            params["visual.embeddings.position_embedding.weight"] =
+                ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, (config.image_size / config.patch_size) * (config.image_size / config.patch_size));
+            params["vqmodel.quantize.embedding.weight"] =
+                ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.codebook_embed_dim, config.codebook_size);
+            params["prior_token_embedding.weight"] =
+                ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.semantic_embed_dim, config.codebook_size);
+        }
+
+        // Bilinear-resamples the square position-embedding grid onto the image's patch grid.
+        // The reference uses grid_sample(align_corners=False, padding_mode="border"); the source
+        // coordinate for output index j is therefore (j + 0.5) * side / out - 0.5, clamped.
+        ggml_tensor* resample_pos_embed(GGMLRunnerContext* ctx,
+                                        ggml_tensor* pos_idx,
+                                        ggml_tensor* pos_weight) {
+            auto pos_embed = params["visual.embeddings.position_embedding.weight"];
+            auto gathered  = ggml_get_rows(ctx->ggml_ctx, pos_embed, pos_idx);
+            return ggml_mul(ctx->ggml_ctx, gathered, pos_weight);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* pixel_values,
+                             const std::vector<ggml_tensor*>& pos_idx,
+                             const std::vector<ggml_tensor*>& pos_weight) {
+            // pixel_values: [N, in_channels, H, W] -> [N, grid_h * grid_w, semantic_embed_dim]
+            ggml_context* gctx = ctx->ggml_ctx;
+
+            auto patch_embed = std::dynamic_pointer_cast<Conv2d>(blocks["visual.patch_embed.proj"]);
+            auto quant_conv  = std::dynamic_pointer_cast<Conv2d>(blocks["vqmodel.quant_conv"]);
+            auto proj_0      = std::dynamic_pointer_cast<Linear>(blocks["prior_projector.net.0.proj"]);
+            auto proj_2      = std::dynamic_pointer_cast<Linear>(blocks["prior_projector.net.2"]);
+
+            auto x          = patch_embed->forward(ctx, pixel_values);  // [N, hidden_size, grid_h, grid_w]
+            int64_t grid_w  = x->ne[0];
+            int64_t grid_h  = x->ne[1];
+            int64_t n_token = grid_h * grid_w;
+            int64_t N       = x->ne[3];
+
+            x = ggml_reshape_3d(gctx, x, n_token, config.hidden_size, N);
+            x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3));  // [N, n_token, hidden_size]
+
+            ggml_tensor* pos = nullptr;
+            for (size_t i = 0; i < pos_idx.size(); i++) {
+                auto corner = resample_pos_embed(ctx, pos_idx[i], pos_weight[i]);
+                pos         = pos == nullptr ? corner : ggml_add(gctx, pos, corner);
+            }
+            x = ggml_add(gctx, x, ggml_reshape_3d(gctx, pos, config.hidden_size, n_token, N));
+
+            for (int i = 0; i < config.num_layers; i++) {
+                auto block = std::dynamic_pointer_cast<SigVQBlock>(blocks["visual.blocks." + std::to_string(i)]);
+                x          = block->forward(ctx, x);
+            }
+
+            // quant_conv is 1x1, so run it as a per-token projection rather than reshaping to 2-D.
+            x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3));  // [N, hidden_size, n_token]
+            x = ggml_reshape_4d(gctx, x, n_token, 1, config.hidden_size, N);
+            x = quant_conv->forward(ctx, x);  // [N, codebook_embed_dim, 1, n_token]
+            x = ggml_reshape_3d(gctx, x, n_token, config.codebook_embed_dim, N);
+            x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3));  // [N, n_token, codebook_embed_dim]
+
+            // Both sides are L2-normalized, so the nearest codebook entry by euclidean distance
+            // is the one with the largest dot product.
+            auto codebook  = ggml_l2_norm(gctx, params["vqmodel.quantize.embedding.weight"], 1e-12f);
+            auto normed    = ggml_l2_norm(gctx, x, 1e-12f);
+            auto logits    = ggml_mul_mat(gctx, codebook, normed);  // [N, n_token, codebook_size]
+            auto token_ids = ggml_argmax(gctx, ggml_reshape_2d(gctx, logits, config.codebook_size, n_token * N));
+
+            auto semantic = ggml_get_rows(gctx, params["prior_token_embedding.weight"], token_ids);
+            semantic      = ggml_reshape_3d(gctx, semantic, config.semantic_embed_dim, n_token, N);
+
+            auto h = proj_0->forward(ctx, semantic);
+            h      = ggml_silu(gctx, h);
+            return proj_2->forward(ctx, h);
+        }
+    };
+
+    struct QueryFormerRunner : public GGMLRunner {
+    public:
+        QueryFormerConfig config;
+        QueryFormerModel query_former;
+
+        QueryFormerRunner(ggml_backend_t backend,
+                          const String2TensorStorage& tensor_storage_map      = {},
+                          const std::string prefix                            = "",
+                          std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : GGMLRunner(backend, weight_manager) {
+            query_former = QueryFormerModel(config);
+            query_former.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "llada_image_queryformer";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
+            query_former.get_param_tensors(tensors, prefix);
+        }
+
+        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& inputs_embeds) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                ggml_cgraph* gf  = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE);
+                ggml_tensor* x   = make_input(inputs_embeds);
+                auto runner_ctx  = get_context();
+                ggml_tensor* out = query_former.forward(&runner_ctx, x);
+                ggml_build_forward_expand(gf, out);
+                return gf;
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true),
+                                                   inputs_embeds.dim());
+        }
+    };
+
+    struct TextProjectionRunner : public GGMLRunner {
+    public:
+        TextProjectionConfig config;
+        TextProjectionModel text_projection;
+
+        TextProjectionRunner(ggml_backend_t backend,
+                             const String2TensorStorage& tensor_storage_map      = {},
+                             const std::string prefix                            = "",
+                             std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : GGMLRunner(backend, weight_manager) {
+            text_projection = TextProjectionModel(config);
+            text_projection.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "llada_image_text_projection";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
+            text_projection.get_param_tensors(tensors, prefix);
+        }
+
+        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& hidden_states) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                ggml_cgraph* gf  = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE);
+                ggml_tensor* x   = make_input(hidden_states);
+                auto runner_ctx  = get_context();
+                ggml_tensor* out = text_projection.forward(&runner_ctx, x);
+                ggml_build_forward_expand(gf, out);
+                return gf;
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true),
+                                                   hidden_states.dim());
+        }
+    };
+
+    struct SigVQRunner : public GGMLRunner {
+    public:
+        SigVQConfig config;
+        SigVQModel sigvq;
+        std::array<std::vector<int32_t>, 4> pos_idx_data;
+        std::array<std::vector<float>, 4> pos_weight_data;
+
+        SigVQRunner(ggml_backend_t backend,
+                    const String2TensorStorage& tensor_storage_map      = {},
+                    const std::string prefix                            = "",
+                    std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : GGMLRunner(backend, weight_manager) {
+            sigvq = SigVQModel(config);
+            sigvq.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "llada_image_sigvq";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
+            sigvq.get_param_tensors(tensors, prefix);
+        }
+
+        // Precomputes the four bilinear taps that resample the square position-embedding grid
+        // onto a grid_h x grid_w patch grid, matching grid_sample(align_corners=False,
+        // padding_mode="border").
+        void build_pos_embed_taps(int64_t grid_h, int64_t grid_w) {
+            const int64_t side = config.image_size / config.patch_size;
+            for (auto& v : pos_idx_data) {
+                v.clear();
+            }
+            for (auto& v : pos_weight_data) {
+                v.clear();
+            }
+
+            auto clamp_index = [side](int64_t v) {
+                return static_cast<int32_t>(std::min<int64_t>(std::max<int64_t>(v, 0), side - 1));
+            };
+
+            for (int64_t i = 0; i < grid_h; ++i) {
+                double src_h    = (static_cast<double>(i) + 0.5) * side / static_cast<double>(grid_h) - 0.5;
+                int64_t h_floor = static_cast<int64_t>(std::floor(src_h));
+                double dh       = src_h - static_cast<double>(h_floor);
+                for (int64_t j = 0; j < grid_w; ++j) {
+                    double src_w    = (static_cast<double>(j) + 0.5) * side / static_cast<double>(grid_w) - 0.5;
+                    int64_t w_floor = static_cast<int64_t>(std::floor(src_w));
+                    double dw       = src_w - static_cast<double>(w_floor);
+
+                    int32_t h0 = clamp_index(h_floor);
+                    int32_t h1 = clamp_index(h_floor + 1);
+                    int32_t w0 = clamp_index(w_floor);
+                    int32_t w1 = clamp_index(w_floor + 1);
+
+                    pos_idx_data[0].push_back(h0 * static_cast<int32_t>(side) + w0);
+                    pos_idx_data[1].push_back(h0 * static_cast<int32_t>(side) + w1);
+                    pos_idx_data[2].push_back(h1 * static_cast<int32_t>(side) + w0);
+                    pos_idx_data[3].push_back(h1 * static_cast<int32_t>(side) + w1);
+
+                    pos_weight_data[0].push_back(static_cast<float>((1.0 - dh) * (1.0 - dw)));
+                    pos_weight_data[1].push_back(static_cast<float>((1.0 - dh) * dw));
+                    pos_weight_data[2].push_back(static_cast<float>(dh * (1.0 - dw)));
+                    pos_weight_data[3].push_back(static_cast<float>(dh * dw));
+                }
+            }
+        }
+
+        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& pixel_values) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE);
+                ggml_tensor* x  = make_input(pixel_values);
+
+                int64_t grid_h = x->ne[1] / config.patch_size;
+                int64_t grid_w = x->ne[0] / config.patch_size;
+                build_pos_embed_taps(grid_h, grid_w);
+
+                std::vector<ggml_tensor*> pos_idx;
+                std::vector<ggml_tensor*> pos_weight;
+                for (int i = 0; i < 4; i++) {
+                    auto idx = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, static_cast<int64_t>(pos_idx_data[i].size()));
+                    set_backend_tensor_data(idx, pos_idx_data[i].data());
+                    auto w = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, 1, static_cast<int64_t>(pos_weight_data[i].size()));
+                    set_backend_tensor_data(w, pos_weight_data[i].data());
+                    pos_idx.push_back(idx);
+                    pos_weight.push_back(w);
+                }
+
+                auto runner_ctx  = get_context();
+                ggml_tensor* out = sigvq.forward(&runner_ctx, x, pos_idx, pos_weight);
+                ggml_build_forward_expand(gf, out);
+                return gf;
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), 3);
+        }
+    };
+
+}  // namespace LLaDAImageTE
+
+#endif  // __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__

src/model/te/llm.hpp

@@ -49,6 +49,7 @@ namespace LLM {
         GEMMA2_2B,
         GEMMA4_12B,
         GPT_OSS_20B,
+        LLADA2_MOE,
         ARCH_COUNT,
     };
 
@@ -62,6 +63,7 @@ namespace LLM {
         "gemma2_2b",
         "gemma4_12b",
         "gpt_oss_20b",
+        "llada2_moe",
     };
 
     enum class MLPActivation {
@@ -125,6 +127,17 @@ namespace LLM {
         std::vector<int> sliding_attention;
         int64_t num_experts         = 0;
         int64_t num_experts_per_tok = 0;
+        bool qkv_fused              = false;
+        bool bidirectional          = false;
+        float partial_rotary        = 1.f;
+
+        // DeepSeek-V3-style grouped-sigmoid MoE routing (LLaDA2)
+        int64_t moe_intermediate_size = 0;
+        int64_t num_shared_experts    = 0;
+        int64_t first_k_dense_replace = 0;
+        int64_t n_group               = 0;
+        int64_t topk_group            = 0;
+        float routed_scaling_factor   = 1.f;
         LLMVisionConfig vision;
         bool have_vision_weight = false;
         bool llama_cpp_style    = false;
@@ -212,6 +225,31 @@ namespace LLM {
                 config.intermediate_size       = 9216;
                 config.num_layers              = 26;
                 config.vocab_size              = 256000;
+            } else if (arch == LLMArch::LLADA2_MOE) {
+                config.head_dim                = 128;
+                config.num_heads               = 16;
+                config.num_kv_heads            = 4;
+                config.qkv_bias                = false;
+                config.attention_out_bias      = false;
+                config.qk_norm                 = true;
+                config.rms_norm_eps            = 1e-6f;
+                config.hidden_size             = 2048;
+                config.intermediate_size       = 5120;
+                config.num_layers              = 20;
+                config.vocab_size              = 173568;
+                config.max_position_embeddings = 16384;
+                config.rope_thetas             = {600000.f};
+                config.qkv_fused               = true;
+                config.bidirectional           = true;
+                config.partial_rotary          = 0.5f;
+                config.num_experts             = 256;
+                config.num_experts_per_tok     = 8;
+                config.moe_intermediate_size   = 512;
+                config.num_shared_experts      = 1;
+                config.first_k_dense_replace   = 1;
+                config.n_group                 = 8;
+                config.topk_group              = 4;
+                config.routed_scaling_factor   = 2.5f;
             } else if (arch == LLMArch::GPT_OSS_20B) {
                 config.head_dim                = 64;
                 config.num_heads               = 64;
@@ -419,6 +457,195 @@ namespace LLM {
         }
     };
 
+    // LLaDA2's MoE differs from GPT-OSS's in three ways that all change the result:
+    // routing scores are sigmoid (not softmax over the selected logits), expert selection is
+    // group-limited and uses a bias term that the returned weights do NOT include, and the
+    // experts carry no biases. Ref: LLaDA2MoeGate / LLaDA2MoeSparseMoeBlock in
+    // modeling_llada2uni_moe.py.
+    struct LLaDA2MoEMLP : public GGMLBlock {
+    protected:
+        int64_t hidden_size;
+        int64_t moe_intermediate_size;
+        int64_t num_experts;
+        int64_t num_experts_per_tok;
+        int64_t n_group;
+        int64_t topk_group;
+        float routed_scaling_factor;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         std::string prefix                             = "") override {
+            GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
+
+            auto supported_type = [](ggml_type wtype, int64_t in_features) {
+                if (in_features % ggml_blck_size(wtype) != 0) {
+                    return GGML_TYPE_F32;
+                }
+                return wtype;
+            };
+
+            // The reference runs the router in fp32; keep the weight in fp32 so the sigmoid
+            // 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);
+
+            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);
+            ggml_type down_type = supported_type(get_type(prefix + "experts.down_proj.weight", tensor_storage_map, GGML_TYPE_F32), moe_intermediate_size);
+
+            // HF ships the stacked experts as 3-D nn.Parameters, while the ComfyUI GGUF repack
+            // flattens the expert axis into ne[1]. Declare whichever the file holds - the two are
+            // bit-identical, and forward() reshapes to 3-D for ggml_mul_mat_id either way.
+            auto declare_experts = [&](const std::string& name, ggml_type type, int64_t in_dim, int64_t out_dim) {
+                auto storage = tensor_storage_map.find(prefix + name);
+                if (storage != tensor_storage_map.end() && storage->second.n_dims == 2) {
+                    GGML_ASSERT(storage->second.nelements() == in_dim * out_dim * num_experts);
+                    params[name] = ggml_new_tensor_2d(ctx, type, in_dim, out_dim * num_experts);
+                } else {
+                    params[name] = ggml_new_tensor_3d(ctx, type, in_dim, out_dim, num_experts);
+                }
+            };
+
+            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);
+        }
+
+    public:
+        LLaDA2MoEMLP(const LLMConfig& config)
+            : hidden_size(config.hidden_size),
+              moe_intermediate_size(config.moe_intermediate_size),
+              num_experts(config.num_experts),
+              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) {
+            if (config.num_shared_experts > 0) {
+                blocks["shared_experts"] = std::make_shared<MLP>(config.hidden_size,
+                                                                 config.moe_intermediate_size * config.num_shared_experts,
+                                                                 false,
+                                                                 config.mlp_activation);
+            }
+        }
+
+        // Reproduces group_limited_topk(): keep the topk_group groups with the highest
+        // "sum of the two best scores in the group", then take the global top-k among them.
+        ggml_tensor* group_limited_mask(GGMLRunnerContext* ctx,
+                                        ggml_tensor* routing_scores,
+                                        int64_t n_token_total) {
+            ggml_context* gctx      = ctx->ggml_ctx;
+            const int64_t per_group = num_experts / n_group;
+
+            // [experts_per_group, n_group * tokens] so top-2 runs per (group, token) row.
+            auto grouped     = ggml_reshape_2d(gctx, routing_scores, per_group, n_group * n_token_total);
+            auto best2_idx   = ggml_argsort_top_k(gctx, grouped, 2);  // [2, n_group * tokens]
+            auto grouped_val = ggml_reshape_3d(gctx, grouped, 1, per_group, n_group * n_token_total);
+            auto best2       = ggml_get_rows(gctx, grouped_val, best2_idx);  // [1, 2, n_group * tokens]
+            best2            = ggml_reshape_2d(gctx, best2, 2, n_group * n_token_total);
+            auto group_score = ggml_reshape_2d(gctx, ggml_sum_rows(gctx, best2), n_group, n_token_total);  // [n_group, tokens]
+
+            // Threshold = the topk_group-th largest group score, taken from the sorted top-k.
+            auto top_groups = ggml_argsort_top_k(gctx, group_score, (int)topk_group);  // [topk_group, tokens]
+            auto group_val  = ggml_reshape_3d(gctx, group_score, 1, n_group, n_token_total);
+            auto top_scores = ggml_get_rows(gctx, group_val, top_groups);  // [1, topk_group, tokens]
+            top_scores      = ggml_reshape_2d(gctx, top_scores, topk_group, n_token_total);
+            auto threshold  = ggml_view_2d(gctx,
+                                           top_scores,
+                                           1,
+                                           n_token_total,
+                                           top_scores->nb[1],
+                                           (topk_group - 1) * top_scores->nb[0]);  // [1, tokens]
+            threshold       = ggml_cont(gctx, threshold);
+
+            // keep = 1 - step(threshold - score). step(0) == 0, so the group sitting exactly on
+            // the threshold is kept without needing an epsilon.
+            auto diff = ggml_sub(gctx, ggml_repeat(gctx, threshold, group_score), group_score);
+            auto keep = ggml_scale_bias(gctx, ggml_step(gctx, diff), -1.f, 1.f);  // [n_group, tokens]
+
+            // 0 for kept groups, a large negative for dropped ones, broadcast over the group.
+            auto additive = ggml_scale_bias(gctx, keep, 1e30f, -1e30f);
+            additive      = ggml_reshape_3d(gctx, additive, 1, n_group, n_token_total);
+            auto expanded = ggml_repeat_4d(gctx, additive, per_group, n_group, n_token_total, 1);
+            return ggml_reshape_2d(gctx, expanded, num_experts, n_token_total);
+        }
+
+        ggml_tensor* expert_linear(GGMLRunnerContext* ctx,
+                                   const std::string& weight_name,
+                                   ggml_tensor* x,
+                                   ggml_tensor* selected_experts) {
+            ggml_tensor* w = params[weight_name];
+            if (w->ne[2] != num_experts) {
+                // Flattened layout: split the expert axis back out. ne[0] is untouched, so this
+                // stays valid for quantized types.
+                w = ggml_reshape_3d(ctx->ggml_ctx, w, w->ne[0], w->ne[1] / num_experts, num_experts);
+            }
+            return ggml_mul_mat_id(ctx->ggml_ctx, w, x, selected_experts);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            // 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);
+
+            ggml_context* gctx          = ctx->ggml_ctx;
+            const int64_t n_token       = x->ne[1];
+            const int64_t N             = x->ne[2];
+            const int64_t n_token_total = n_token * N;
+
+            auto identity = x;
+
+            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]
+
+            // The bias steers selection only; the combine weights come from the unbiased scores.
+            auto routing = ggml_add(gctx, scores, params["gate.expert_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]
+            auto score_rows       = ggml_reshape_3d(gctx, scores, 1, num_experts, n_token_total);
+            auto weights          = ggml_get_rows(gctx, score_rows, selected_experts);  // [1, top_k, tokens]
+            weights               = ggml_reshape_2d(gctx, weights, num_experts_per_tok, n_token_total);
+
+            if (num_experts_per_tok > 1) {
+                auto denom = ggml_scale_bias(gctx, ggml_sum_rows(gctx, weights), 1.f, 1e-20f);  // [1, tokens]
+                weights    = ggml_div(gctx, weights, ggml_repeat(gctx, denom, weights));
+            }
+            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);
+
+            ggml_tensor* out = nullptr;
+            for (int64_t i = 0; i < num_experts_per_tok; ++i) {
+                auto expert_out = ggml_view_2d(gctx,
+                                               experts,
+                                               hidden_size,
+                                               n_token_total,
+                                               experts->nb[2],
+                                               i * experts->nb[1]);
+                out             = out == nullptr ? expert_out : ggml_add(gctx, out, expert_out);
+            }
+            if (num_experts_per_tok == 1) {
+                out = ggml_cont(gctx, out);
+            }
+            out = ggml_reshape_3d(gctx, out, hidden_size, n_token, N);
+
+            auto shared_it = blocks.find("shared_experts");
+            if (shared_it != blocks.end()) {
+                auto shared_experts = std::dynamic_pointer_cast<MLP>(shared_it->second);
+                out                 = ggml_add(gctx, out, shared_experts->forward(ctx, identity));
+            }
+
+            return out;
+        }
+    };
+
     struct GPTOSSMLP : public GGMLBlock {
     protected:
         int64_t hidden_size;
@@ -605,21 +832,31 @@ namespace LLM {
             }
             txt_token_end = image_embeds[i].first;
 
-            auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
-            if (input_embed == nullptr) {
-                input_embed = txt_embed;
-            } else {
-                input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1);
+            // An embed can sit flush against the previous one or at the very start/end of the
+            // sequence, leaving no text tokens to splice around it.
+            if (txt_token_end > txt_token_start) {
+                auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
+                if (input_embed == nullptr) {
+                    input_embed = txt_embed;
+                } else {
+                    input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1);
+                }
             }
 
-            input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1);
+            if (input_embed == nullptr) {
+                input_embed = image_embeds[i].second;
+            } else {
+                input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1);
+            }
         }
 
         txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1];
         txt_token_end   = raw_x->ne[1];
 
-        auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
-        input_embed          = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
+        if (txt_token_end > txt_token_start) {
+            auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
+            input_embed          = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
+        }
         GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]);
         return input_embed;
     }
@@ -1122,6 +1359,7 @@ namespace LLM {
         bool k_eq_v;
         bool v_norm;
         bool unscaled_attention;
+        bool qkv_fused;
         float rms_norm_eps;
         int rope_pairs;
 
@@ -1147,12 +1385,20 @@ namespace LLM {
               k_eq_v(global_layer && config.global_k_eq_v),
               v_norm(config.v_norm),
               unscaled_attention(config.unscaled_attention),
+              qkv_fused(config.qkv_fused),
               rms_norm_eps(config.rms_norm_eps),
               rope_pairs(0) {
-            blocks["q_proj"] = std::make_shared<Linear>(config.hidden_size, num_heads * head_dim, config.qkv_bias);
-            blocks["k_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
-            if (!k_eq_v) {
-                blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
+            if (qkv_fused) {
+                // The checkpoint ships q, k and v as one tensor and the loader cannot split a
+                // source tensor, so keep it fused and slice it in forward().
+                GGML_ASSERT(!k_eq_v);
+                blocks["query_key_value"] = std::make_shared<Linear>(config.hidden_size, (num_heads + num_kv_heads * 2) * head_dim, config.qkv_bias);
+            } else {
+                blocks["q_proj"] = std::make_shared<Linear>(config.hidden_size, num_heads * head_dim, config.qkv_bias);
+                blocks["k_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
+                if (!k_eq_v) {
+                    blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
+                }
             }
             blocks["o_proj"] = std::make_shared<Linear>(num_heads * head_dim, config.hidden_size, config.attention_out_bias);
             if (config.qk_norm) {
@@ -1161,7 +1407,7 @@ namespace LLM {
             }
             // Proportional RoPE rotates only the leading `rope_pairs` dimension pairs of the head;
             // the rest are left unrotated through freq_factors (see rope_freq_factors()).
-            float partial = global_layer ? config.global_partial_rotary : 1.f;
+            float partial = global_layer && config.global_partial_rotary != 1.f ? config.global_partial_rotary : config.partial_rotary;
             rope_pairs    = static_cast<int>(partial * head_dim / 2.f);
         }
 
@@ -1186,14 +1432,28 @@ namespace LLM {
             // x: [N, n_token, hidden_size]
             int64_t n_token = x->ne[1];
             int64_t N       = x->ne[2];
-            auto q_proj     = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
-            auto k_proj     = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
-            auto v_proj     = k_eq_v ? nullptr : std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
             auto out_proj   = std::dynamic_pointer_cast<Linear>(blocks["o_proj"]);
 
-            auto q = q_proj->forward(ctx, x);               // [N, n_token, num_heads*head_dim]
-            auto k = k_proj->forward(ctx, x);               // [N, n_token, num_kv_heads*head_dim]
-            auto v = k_eq_v ? k : v_proj->forward(ctx, x);  // [N, n_token, num_kv_heads*head_dim]
+            ggml_tensor* q = nullptr;
+            ggml_tensor* k = nullptr;
+            ggml_tensor* v = nullptr;
+            if (qkv_fused) {
+                auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["query_key_value"]);
+                auto qkv      = qkv_proj->forward(ctx, x);  // [N, n_token, (num_heads + num_kv_heads*2)*head_dim]
+                int64_t q_len = num_heads * head_dim;
+                int64_t k_len = num_kv_heads * head_dim;
+                q             = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, 0, q_len);
+                k             = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len, q_len + k_len);
+                v             = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len + k_len, q_len + k_len * 2);
+            } else {
+                auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
+                auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
+                auto v_proj = k_eq_v ? nullptr : std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
+
+                q = q_proj->forward(ctx, x);               // [N, n_token, num_heads*head_dim]
+                k = k_proj->forward(ctx, x);               // [N, n_token, num_kv_heads*head_dim]
+                v = k_eq_v ? k : v_proj->forward(ctx, x);  // [N, n_token, num_kv_heads*head_dim]
+            }
 
             q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, n_token, N);     // [N, n_token, num_heads, head_dim]
             k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_kv_heads, n_token, N);  // [N, n_token, num_kv_heads, head_dim]
@@ -1336,6 +1596,38 @@ namespace LLM {
                                   1.f,
                                   32.f,
                                   1.f);
+            } else if (arch == LLMArch::LLADA2_MOE) {
+                // LLaDA2 slices the head (query[..., :rotary_dim]) instead of zero-padding
+                // inv_freq like gemma does, so rotate_half pairs i with i + rotary_dim/2 and the
+                // frequencies use rotary_dim as the exponent denominator. Passing n_dims =
+                // rotary_dim reproduces both; freq_factors would give the wrong pairing.
+                int rotary_dim = rope_pairs * 2;
+                q              = ggml_rope_ext(ctx->ggml_ctx,
+                                               q,
+                                               input_pos,
+                                               nullptr,
+                                               rotary_dim,
+                                               GGML_ROPE_TYPE_NEOX,
+                                               static_cast<int>(max_position_embeddings),
+                                               rope_thetas[0],
+                                               1.f,
+                                               0.f,
+                                               1.f,
+                                               32.f,
+                                               1.f);
+                k              = ggml_rope_ext(ctx->ggml_ctx,
+                                               k,
+                                               input_pos,
+                                               nullptr,
+                                               rotary_dim,
+                                               GGML_ROPE_TYPE_NEOX,
+                                               static_cast<int>(max_position_embeddings),
+                                               rope_thetas[0],
+                                               1.f,
+                                               0.f,
+                                               1.f,
+                                               32.f,
+                                               1.f);
             } else if (arch == LLMArch::QWEN3_VL) {
                 int sections[4] = {24, 20, 20, 0};
                 q               = ggml_rope_multi(ctx->ggml_ctx, q, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_IMROPE, 262144, 5000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
@@ -1432,6 +1724,8 @@ 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) {
+                blocks["mlp"] = std::make_shared<LLaDA2MoEMLP>(config);
             } else {
                 blocks["mlp"] = std::make_shared<MLP>(config.hidden_size,
                                                       config.intermediate_size,
@@ -1485,6 +1779,10 @@ namespace LLM {
             if (arch == LLMArch::GPT_OSS_20B) {
                 auto mlp = std::dynamic_pointer_cast<GPTOSSMLP>(blocks["mlp"]);
                 x        = mlp->forward(ctx, x);
+            } 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);
             } else {
                 auto mlp = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
                 x        = mlp->forward(ctx, x);
@@ -1650,6 +1948,11 @@ namespace LLM {
             return x;
         }
 
+        ggml_tensor* embed(GGMLRunnerContext* ctx, ggml_tensor* input_ids) {
+            auto model = std::dynamic_pointer_cast<TextModel>(blocks["model"]);
+            return model->embed(ctx, input_ids);
+        }
+
         std::shared_ptr<VisionModel> vision_model() {
             GGML_ASSERT(enable_vision);
             return std::dynamic_pointer_cast<VisionModel>(blocks["visual"]);
@@ -1990,7 +2293,8 @@ namespace LLM {
                 config.arch == LLMArch::GEMMA3_12B ||
                 config.arch == LLMArch::GEMMA4_12B ||
                 config.arch == LLMArch::GEMMA2_2B ||
-                config.arch == LLMArch::GPT_OSS_20B) {
+                config.arch == LLMArch::GPT_OSS_20B ||
+                config.arch == LLMArch::LLADA2_MOE) {
                 input_pos_vec.resize(n_tokens);
                 for (int i = 0; i < n_tokens; ++i) {
                     input_pos_vec[i] = i;
@@ -2042,8 +2346,9 @@ namespace LLM {
                 attention_mask_vec.resize(n_tokens * n_tokens);
                 for (int i0 = 0; i0 < n_tokens; i0++) {
                     for (int i1 = 0; i1 < n_tokens; i1++) {
+                        // Diffusion LLMs attend in both directions; only causal LMs get the triangle.
                         float value = 0.f;
-                        if (i0 > i1) {
+                        if (!config.bidirectional && i0 > i1) {
                             value = -INFINITY;
                         }
                         attention_mask_vec[i1 * n_tokens + i0] = value;
@@ -2115,6 +2420,22 @@ namespace LLM {
                                                    input_ids.dim() + 1);
         }
 
+        // LLaDA-Image's QueryFormer consumes the raw token embeddings before the backbone runs,
+        // so it needs the embedding lookup on its own.
+        sd::Tensor<float> compute_input_embeds(const int n_threads,
+                                               const sd::Tensor<int32_t>& input_ids) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                ggml_cgraph* gf  = new_graph_custom(LLM_GRAPH_SIZE);
+                ggml_tensor* ids = make_input(input_ids);
+                auto runner_ctx  = get_context();
+                ggml_tensor* out = model.embed(&runner_ctx, ids);
+                ggml_build_forward_expand(gf, out);
+                return gf;
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true),
+                                                   input_ids.dim() + 1);
+        }
+
         int64_t get_num_image_tokens(int64_t t, int64_t h, int64_t w) {
             int64_t grid_t     = 1;
             int64_t grid_h     = h / config.vision.patch_size;
@@ -2370,11 +2691,13 @@ namespace LLM {
                 pad_id = 199999;
             } else if (arch == LLMArch::GEMMA2_2B) {
                 pad_id = 0;
+            } else if (arch == LLMArch::LLADA2_MOE) {
+                pad_id = 156892;
             }
             tokenizer = tokenizers.create(TokenizerConfig::MAIN, model.config.vocab_size, pad_id);
             if (!tokenizer) {
-                if (arch == LLMArch::GPT_OSS_20B || arch == LLMArch::GEMMA2_2B) {
-                    throw std::runtime_error("GPT-OSS and Gemma 2 require an external tokenizer.json in the main tokenizer slot");
+                if (arch == LLMArch::GPT_OSS_20B || arch == LLMArch::GEMMA2_2B || arch == LLMArch::LLADA2_MOE) {
+                    throw std::runtime_error("GPT-OSS, Gemma 2 and LLaDA2 require an external tokenizer.json in the main tokenizer slot");
                 }
                 if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
                     tokenizer = std::make_shared<MistralTokenizer>();

src/model/vae/vae.hpp

@@ -162,7 +162,7 @@ public:
         int scale_factor = 8;
         if (version == VERSION_LTXAV) {
             scale_factor = 32;
-        } else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version) || 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;

src/model/vae/wan_vae.hpp

@@ -26,6 +26,13 @@ namespace WAN {
         bool bias;
 
         void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            auto weight = tensor_storage_map.find(prefix + "weight");
+            if (weight != tensor_storage_map.end() && weight->second.ne[2] == 1 &&
+                weight->second.ne[3] == in_channels * out_channels) {
+                // Image VAE exports may retain Conv3d weights with a singleton temporal kernel.
+                std::get<0>(kernel_size) = 1;
+                std::get<0>(padding)     = 0;
+            }
             params["weight"] = ggml_new_tensor_4d(ctx,
                                                   GGML_TYPE_F16,
                                                   std::get<2>(kernel_size),
@@ -140,7 +147,7 @@ namespace WAN {
         std::string mode;
 
     public:
-        Resample(int64_t dim, const std::string& mode, bool wan2_2 = false)
+        Resample(int64_t dim, const std::string& mode, bool wan2_2 = false, bool is_2D = false)
             : dim(dim), mode(mode) {
             if (mode == "upsample2d") {
                 if (wan2_2) {
@@ -154,12 +161,20 @@ namespace WAN {
                 } else {
                     blocks["resample.1"] = std::shared_ptr<GGMLBlock>(new Conv2d(dim, dim / 2, {3, 3}, {1, 1}, {1, 1}));
                 }
-                blocks["time_conv"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(dim, dim * 2, {3, 1, 1}, {1, 1, 1}, {1, 0, 0}));
+                if (is_2D) {
+                    blocks["time_conv"] = std::make_shared<Conv2dBut3d>(dim, dim * 2, std::pair<int, int>{1, 1});
+                } else {
+                    blocks["time_conv"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(dim, dim * 2, {3, 1, 1}, {1, 1, 1}, {1, 0, 0}));
+                }
             } else if (mode == "downsample2d") {
                 blocks["resample.1"] = std::shared_ptr<GGMLBlock>(new Conv2d(dim, dim, {3, 3}, {2, 2}));
             } else if (mode == "downsample3d") {
                 blocks["resample.1"] = std::shared_ptr<GGMLBlock>(new Conv2d(dim, dim, {3, 3}, {2, 2}));
-                blocks["time_conv"]  = std::shared_ptr<GGMLBlock>(new CausalConv3d(dim, dim, {3, 1, 1}, {2, 1, 1}, {0, 0, 0}));
+                if (is_2D) {
+                    blocks["time_conv"] = std::make_shared<Conv2dBut3d>(dim, dim, std::pair<int, int>{1, 1});
+                } else {
+                    blocks["time_conv"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(dim, dim, {3, 1, 1}, {2, 1, 1}, {0, 0, 0}));
+                }
             } else if (mode == "none") {
                 // nn.Identity()
             } else {
@@ -469,7 +484,7 @@ namespace WAN {
             }
             if (down_flag) {
                 std::string mode                           = temperal_downsample ? "downsample3d" : "downsample2d";
-                blocks["downsamples." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new Resample(out_dim, mode, true));
+                blocks["downsamples." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new Resample(out_dim, mode, true, is_2D));
                 i++;
             }
         }
@@ -532,7 +547,7 @@ namespace WAN {
             }
             if (up_flag) {
                 std::string mode                         = temperal_upsample ? "upsample3d" : "upsample2d";
-                blocks["upsamples." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new Resample(out_dim, mode, true));
+                blocks["upsamples." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new Resample(out_dim, mode, true, is_2D));
                 i++;
             }
         }
@@ -1054,9 +1069,24 @@ namespace WAN {
                 input_channels = 4;
             }
 
+            if (version == VERSION_QWEN_IMAGE_2_1) {
+                wan2_2         = true;
+                dec_dim        = 144;
+                z_dim          = 64;
+                input_channels = 4;
+                dim_mult       = {1, 2, 4, 8, 8};
+            }
+
             if (is_2D) {
-                temperal_upsample   = {false, false, false};
-                temperal_downsample = {false, false, false};
+                temperal_upsample.assign(dim_mult.size() - 1, false);
+                temperal_downsample.assign(dim_mult.size() - 1, false);
+            }
+            if (version == VERSION_QWEN_IMAGE_2_1) {
+                // Temporal shortcut factors still affect single-frame channel grouping.
+                temperal_upsample   = {true, true, true, false};
+                temperal_downsample = {false, true, true, true};
+                _conv_num           = 2 * (2 + static_cast<int>(dim_mult.size()) * (num_res_blocks + 1)) + 3 + (is_2D ? 0 : 2);
+                _enc_conv_num       = 2 * (2 + static_cast<int>(dim_mult.size()) * num_res_blocks) + 3 + (is_2D ? 0 : 2);
             }
 
             if (!decode_only) {
@@ -1271,18 +1301,9 @@ namespace WAN {
                      SDVersion version                                   = VERSION_WAN2,
                      std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
             : VAE(version, backend, prefix, weight_manager), decode_only(decode_only) {
-            bool is_2D = false;
-            for (const auto& [name, tensor_storage] : tensor_storage_map) {
-                if (ends_with(name, "decoder.conv1.weight")) {
-                    if (tensor_storage.ne[2] > 3) {
-                        is_2D = true;
-                    }
-                    break;
-                }
-            }
-            if (is_2D) {
-                LOG_VERBOSE("USING 2D VAE");
-            }
+            const auto conv_in = tensor_storage_map.find((prefix.empty() ? "" : prefix + ".") + "decoder.conv1.weight");
+            const bool is_2D   = conv_in != tensor_storage_map.end() && conv_in->second.ne[2] > 3;
+            LOG_VERBOSE("Wan VAE convolution type: %s", is_2D ? "2D" : "3D");
             ae = WanVAE(decode_only, version, is_2D);
             ae.init(params_ctx, tensor_storage_map, prefix);
         }
@@ -1342,6 +1363,28 @@ namespace WAN {
                 std_tensor.reshape_(stats_shape);
                 return {std::move(mean_tensor), std::move(std_tensor)};
             }
+            if (version == VERSION_QWEN_IMAGE_2_1 && latents.shape()[channel_dim] == 64) {
+                stats_shape[static_cast<size_t>(channel_dim)] = 64;
+                auto mean_tensor                              = sd::Tensor<float>::from_vector({0.5126f, 0.7721f, -0.0631f, 1.3506f, -0.7855f, -2.1025f, -0.3458f, 1.3722f,
+                                                                                                1.8873f, -1.7177f, -0.6510f, 0.2732f, 0.7562f, -0.6163f, -1.0277f, 3.8363f,
+                                                                                                2.0210f, 0.0472f, 0.9320f, 2.0087f, 2.4954f, -0.1391f, -1.4249f, 1.8464f,
+                                                                                                -0.5236f, 1.2826f, 3.7046f, -1.3035f, 2.7286f, -1.4518f, -1.9036f, -1.9955f,
+                                                                                                -0.0342f, -1.0265f, -0.7636f, 3.0555f, 0.0746f, -3.0751f, -0.1076f, 1.7376f,
+                                                                                                -1.0914f, -1.9435f, -0.2784f, -1.3680f, 0.4809f, -0.4433f, 0.3764f, 0.5729f,
+                                                                                                -2.0595f, 1.0960f, -1.3260f, -2.0211f, -5.0179f, 0.5275f, 4.0162f, 1.8505f,
+                                                                                                0.3026f, 1.9373f, 1.4937f, 0.2632f, 0.5547f, -1.7121f, -0.1562f, 0.0304f});
+                auto std_tensor                               = sd::Tensor<float>::from_vector({3.2001f, 3.2936f, 3.4321f, 3.0091f, 3.1061f, 4.0379f, 4.0705f, 3.7910f,
+                                                                                                3.0785f, 3.6500f, 3.9308f, 3.0904f, 2.8778f, 3.7675f, 3.7320f, 5.0756f,
+                                                                                                3.2864f, 4.0397f, 3.1317f, 4.0443f, 2.9249f, 3.9454f, 3.0988f, 4.2489f,
+                                                                                                3.4896f, 3.8513f, 3.9323f, 3.4719f, 3.7498f, 4.2830f, 3.5694f, 4.2467f,
+                                                                                                3.9037f, 3.2947f, 5.0770f, 3.5075f, 3.2700f, 3.4767f, 2.8063f, 5.1125f,
+                                                                                                3.5327f, 4.7833f, 3.1286f, 4.1819f, 3.8527f, 3.8312f, 3.5605f, 4.3875f,
+                                                                                                3.9624f, 4.0168f, 3.5643f, 4.0550f, 5.5614f, 4.2963f, 4.4080f, 3.4959f,
+                                                                                                3.8747f, 3.7608f, 3.5735f, 3.1490f, 3.7662f, 3.6746f, 3.4563f, 3.8161f});
+                mean_tensor.reshape_(stats_shape);
+                std_tensor.reshape_(stats_shape);
+                return {std::move(mean_tensor), std::move(std_tensor)};
+            }
             GGML_ABORT("unexpected latent channel dimension %lld for version %d",
                        (long long)latents.shape()[channel_dim],
                        version);

src/model_io/gguf_io.cpp

@@ -57,6 +57,18 @@ bool read_gguf_file(const std::string& file_path,
 
         size_t data_offset = gguf_reader.data_offset();
         for (const auto& gguf_tensor_info : gguf_reader.tensors()) {
+#ifdef SD_USE_UPSTREAM_GGML
+            if (static_cast<int>(gguf_tensor_info.type) == SD_TYPE_F8_E4M3 ||
+                static_cast<int>(gguf_tensor_info.type) == SD_TYPE_F8_E5M2) {
+                set_error(error, "FP8 is not supported by this ggml build (tensor '" + gguf_tensor_info.name + "')");
+                return false;
+            }
+#endif
+            if (static_cast<unsigned>(gguf_tensor_info.type) >= GGML_TYPE_COUNT ||
+                ggml_get_type_traits(gguf_tensor_info.type)->type_size == 0) {
+                set_error(error, "unsupported GGUF tensor type (tensor '" + gguf_tensor_info.name + "')");
+                return false;
+            }
             TensorStorage tensor_storage(
                 gguf_tensor_info.name,
                 gguf_tensor_info.type,

src/model_io/safetensors_io.cpp

@@ -86,10 +86,17 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) {
         ttype = GGML_TYPE_F32;
     } else if (dtype == "F64") {
         ttype = GGML_TYPE_F32;
+#ifdef SD_USE_UPSTREAM_GGML
+    } else if (dtype == "F8_E4M3") {
+        ttype = GGML_TYPE_F16;
+    } else if (dtype == "F8_E5M2") {
+        ttype = GGML_TYPE_F16;
+#else
     } else if (dtype == "F8_E4M3") {
         ttype = GGML_TYPE_F8_E4M3;
     } else if (dtype == "F8_E5M2") {
         ttype = GGML_TYPE_F8_E5M2;
+#endif
     } else if (dtype == "I32") {
         ttype = GGML_TYPE_I32;
     } else if (dtype == "I64") {
@@ -230,6 +237,12 @@ bool read_safetensors_file(const std::string& file_path,
         if (!read_comfy_quant_config(file, file_path, name, data_start + begin, end - begin, config, error)) {
             return false;
         }
+#ifdef SD_USE_UPSTREAM_GGML
+        if (config.format == "int8_tensorwise") {
+            set_error(error, "INT8 tensorwise/convrot is not supported by this ggml build (tensor '" + name + "')");
+            return false;
+        }
+#endif
         const std::string module_name = name.substr(0, name.size() - std::string(".comfy_quant").size());
         comfy_quant_configs.emplace(module_name, std::move(config));
     }
@@ -362,10 +375,20 @@ bool read_safetensors_file(const std::string& file_path,
         bool tensor_size_ok;
         if (dtype == "F8_E4M3") {
             tensor_storage.is_f8_e4m3 = true;
-            tensor_size_ok            = (tensor_storage.nbytes() == tensor_data_size);
+#ifdef SD_USE_UPSTREAM_GGML
+            // f8 -> f16
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
+#else
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
+#endif
         } else if (dtype == "F8_E5M2") {
             tensor_storage.is_f8_e5m2 = true;
-            tensor_size_ok            = (tensor_storage.nbytes() == tensor_data_size);
+#ifdef SD_USE_UPSTREAM_GGML
+            // f8 -> f16
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
+#else
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
+#endif
         } else if (dtype == "F64") {
             tensor_storage.is_f64 = true;
             // f64 -> f32

src/model_io/tensor_storage.h

@@ -58,6 +58,10 @@ struct TensorStorage {
     int64_t nbytes_to_read() const {
         if (is_f64 || is_i64) {
             return nbytes() * 2;
+#ifdef SD_USE_UPSTREAM_GGML
+        } else if (is_f8_e4m3 || is_f8_e5m2) {
+            return nbytes() / 2;
+#endif
         } else {
             return nbytes();
         }

src/model_loader.cpp

@@ -35,6 +35,66 @@
 
 /*================================================= Preprocess ==================================================*/
 
+#ifdef SD_USE_UPSTREAM_GGML
+uint16_t f8_e4m3_to_f16(uint8_t f8) {
+    const uint32_t exponent_bias = 7;
+    if (f8 == 0xff) {
+        return ggml_fp32_to_fp16(-NAN);
+    } else if (f8 == 0x7f) {
+        return ggml_fp32_to_fp16(NAN);
+    }
+
+    uint32_t sign     = f8 & 0x80;
+    uint32_t exponent = (f8 & 0x78) >> 3;
+    uint32_t mantissa = f8 & 0x07;
+    uint32_t result   = sign << 24;
+    if (exponent == 0) {
+        if (mantissa > 0) {
+            exponent = 0x7f - exponent_bias;
+
+            // yes, 2 times
+            if ((mantissa & 0x04) == 0) {
+                mantissa &= 0x03;
+                mantissa <<= 1;
+                exponent -= 1;
+            }
+            if ((mantissa & 0x04) == 0) {
+                mantissa &= 0x03;
+                mantissa <<= 1;
+                exponent -= 1;
+            }
+
+            result |= (mantissa & 0x03) << 21;
+            result |= exponent << 23;
+        }
+    } else {
+        result |= mantissa << 20;
+        exponent += 0x7f - exponent_bias;
+        result |= exponent << 23;
+    }
+
+    return ggml_fp32_to_fp16(*reinterpret_cast<const float*>(&result));
+}
+
+uint16_t f8_e5m2_to_f16(uint8_t fp8) {
+    return static_cast<uint16_t>(fp8) << 8;
+}
+
+void f8_e4m3_to_f16_vec(uint8_t* src, uint16_t* dst, int64_t n) {
+    // support inplace op
+    for (int64_t i = n - 1; i >= 0; i--) {
+        dst[i] = f8_e4m3_to_f16(src[i]);
+    }
+}
+
+void f8_e5m2_to_f16_vec(uint8_t* src, uint16_t* dst, int64_t n) {
+    // support inplace op
+    for (int64_t i = n - 1; i >= 0; i--) {
+        dst[i] = f8_e5m2_to_f16(src[i]);
+    }
+}
+#endif
+
 void f64_to_f32_vec(double* src, float* dst, int64_t n) {
     // support inplace op
     for (int64_t i = 0; i < n; i++) {
@@ -423,6 +483,9 @@ SDVersion ModelLoader::get_sd_version() const {
         if (tensor_storage.name.find("language_model.model.layers.0.self_attn.q_proj_mot_gen.weight") != std::string::npos) {
             return VERSION_SENSENOVA_U1_5;
         }
+        if (tensor_storage.name == "model.diffusion_model.txt_in.text_norm.weight") {
+            return VERSION_QWEN_IMAGE_2_1;
+        }
         if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
             auto img_in = tensor_storage_map.find("model.diffusion_model.img_in.weight");
             if (img_in != tensor_storage_map.end() && img_in->second.ne[0] == 128) {
@@ -451,6 +514,9 @@ SDVersion ModelLoader::get_sd_version() const {
         if (tensor_storage.name.find("model.diffusion_model.double_blocks.0.img_mlp.gate_proj.weight") != std::string::npos) {
             return VERSION_OVIS_IMAGE;
         }
+        if (tensor_storage.name.find("model.diffusion_model.sigvq_embedder.1.weight") != std::string::npos) {
+            return VERSION_LLADA_IMAGE;
+        }
         if (tensor_storage.name.find("model.diffusion_model.cap_embedder.0.weight") != std::string::npos) {
             return VERSION_Z_IMAGE;
         }
@@ -858,6 +924,10 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
 
             if (tensor_storage.is_f64 ||
                 tensor_storage.is_i64 ||
+#ifdef SD_USE_UPSTREAM_GGML
+                tensor_storage.is_f8_e4m3 ||
+                tensor_storage.is_f8_e5m2 ||
+#endif
                 tensor_storage.type != dst_tensor->type) {
                 continue;
             }
@@ -1147,6 +1217,12 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                         f64_to_f32_vec((double*)read_buf, (float*)target_buf, tensor_storage.nelements());
                     } else if (tensor_storage.is_i64) {
                         i64_to_i32_vec((int64_t*)read_buf, (int32_t*)target_buf, tensor_storage.nelements());
+#ifdef SD_USE_UPSTREAM_GGML
+                    } else if (tensor_storage.is_f8_e4m3) {
+                        f8_e4m3_to_f16_vec((uint8_t*)read_buf, (uint16_t*)target_buf, tensor_storage.nelements());
+                    } else if (tensor_storage.is_f8_e5m2) {
+                        f8_e5m2_to_f16_vec((uint8_t*)read_buf, (uint16_t*)target_buf, tensor_storage.nelements());
+#endif
                     }
                     if (tensor_storage.type != dst_tensor->type) {
                         if (convert_buf == nullptr) {
@@ -1463,6 +1539,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 (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.
         } else {
             return true;
         }

src/model_loader_files.cpp

@@ -147,7 +147,7 @@ bool ModelLoader::add_file_impl(const std::string& path, const std::string& pref
         }
         if (tensor.index_in_zip < 0) {
             const auto& stamp = physical_files[tensor.file_index];
-            if (tensor.offset > stamp.size || elements / block_size * type_size > stamp.size - tensor.offset) {
+            if (tensor.offset > stamp.size || static_cast<uint64_t>(tensor.nbytes_to_read()) > stamp.size - tensor.offset) {  // kcpp int8 fp8
                 LOG_ERROR("tensor '%s' extends beyond its model file", tensor.name.c_str());
                 return false;
             }

src/model_manager.cpp

@@ -1613,7 +1613,8 @@ void ModelManager::remove_runtime_owner(uintptr_t owner_id) {
 
 ModelManager::CapacityCheck ModelManager::check_capacity(
     const DeviceMemoryRequest& request,
-    const std::vector<TensorState*>& states) const {
+    const std::vector<TensorState*>& states,
+    bool log_details) const {
     CapacityCheck result;
     if (request.compute_backend == nullptr || sd_backend_is_cpu(request.compute_backend)) {
         return result;
@@ -1631,16 +1632,23 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
         }
         size_t free_bytes = 0, total_bytes = 0;
         ggml_backend_dev_memory(device, &free_bytes, &total_bytes);
+        const size_t weights_resident = compute_backend_resident_bytes(backend);
+        const size_t other_runtime    = other_runtime_resident_bytes(request.owner_id, backend);
+        const size_t resident         = add(weights_resident, add(other_runtime, request.runtime_resident_bytes));
+        if (log_details) {
+            LOG_WARN("model manager memory on %s: reported free %.2f MB / total %.2f MB, tracked weights %.2f MB / other runtime %.2f MB / current runtime %.2f MB",
+                        ggml_backend_name(backend),
+                        free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0),
+                        weights_resident / (1024.0 * 1024.0), other_runtime / (1024.0 * 1024.0),
+                        request.runtime_resident_bytes / (1024.0 * 1024.0));
+        }
         if (free_bytes == 0 && total_bytes == 0) {
             return SIZE_MAX;
         }
         // Vulkan's heap budget subtraction can underflow when usage exceeds the budget.
-        if (total_bytes > 0 && free_bytes > total_bytes) {
+        if (total_bytes > 0 && free_bytes > total_bytes && sd_backend_is(backend, "Vulkan")) {
             return size_t{0};
         }
-        const size_t resident = add(compute_backend_resident_bytes(backend),
-                                       add(other_runtime_resident_bytes(request.owner_id, backend),
-                                           request.runtime_resident_bytes));
         if (total_bytes > 0) {
             free_bytes = std::min(free_bytes, resident < total_bytes ? total_bytes - resident : 0);
         }
@@ -1786,7 +1794,7 @@ bool ModelManager::ensure_compute_backend_capacity(
         }
     }
 
-    const auto capacity                = check_capacity(request, required_states);
+    const auto capacity                = check_capacity(request, required_states, true);
     const std::string available_device = capacity.available_device_bytes == SIZE_MAX
                                              ? "unknown"
                                              : sd_format("%.2f MB", capacity.available_device_bytes / (1024.0 * 1024.0));

src/model_manager.h

@@ -157,7 +157,8 @@ private:
         }
     };
     CapacityCheck check_capacity(const DeviceMemoryRequest& request,
-                                 const std::vector<TensorState*>& states) const;
+                                 const std::vector<TensorState*>& states,
+                                 bool log_details = false) const;
 
     ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const;
     ggml_backend_buffer_type_t split_buffer_type_for(const TensorState& state) const;

src/name_conversion.cpp

@@ -103,6 +103,8 @@ std::string convert_open_clip_to_hf_clip_name(std::string name) {
     return 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 const std::vector<std::pair<std::string, std::string>> clip_name_map{
         {"transformer.text_projection.weight", "transformer.text_model.text_projection"},
@@ -177,6 +179,7 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix)
             replace_with_name_map(name, llm_vision_name_map);
         } else {
             replace_with_name_map(name, llm_name_map);
+            name = convert_llada2_moe_te_name(name);
         }
     } else {
         name = convert_open_clip_to_hf_clip_name(name);
@@ -749,6 +752,52 @@ std::string convert_hunyuan_video_to_original_flux(std::string name) {
     return name;
 }
 
+// LLaDA-Image's LLaDA2-MoE text encoder. Both published layouts use these names; the ComfyUI
+// GGUF repack differs only by appending ".weight" to the bare 3-D expert parameters.
+// Called with the "text_encoders." prefix already stripped, so the name still carries "llm.".
+std::string convert_llada2_moe_te_name(std::string name) {
+    static const std::vector<std::pair<std::string, std::string>> name_map = {
+        {"model.language_model.word_embeddings.", "model.embed_tokens."},
+        {"model.language_model.norm.", "model.norm."},
+        {"model.language_model.lm_head.", "lm_head."},
+        {"model.language_model.layers.", "model.layers."},
+        {"attention.query_key_value.", "self_attn.query_key_value."},
+        {"attention.dense.", "self_attn.o_proj."},
+        {"attention.query_layernorm.", "self_attn.q_norm."},
+        {"attention.key_layernorm.", "self_attn.k_norm."},
+    };
+    replace_with_name_map(name, name_map);
+
+    // The HF checkpoint stores the stacked experts as bare nn.Parameters with no ".weight".
+    static const std::vector<std::string> bare_expert_params = {
+        "mlp.experts.gate_proj",
+        "mlp.experts.up_proj",
+        "mlp.experts.down_proj",
+    };
+    for (const auto& suffix : bare_expert_params) {
+        if (ends_with(name, suffix)) {
+            name += ".weight";
+            break;
+        }
+    }
+
+    return name;
+}
+
+// The attention projections keep their diffusers names (JointAttention's split_qkv mode), so
+// only the patch-size-keyed dicts need flattening. Latents arrive already patchified from the
+// Flux2 VAE, so the only patch key is 1-1.
+std::string convert_diffusers_dit_to_original_llada_image(std::string name) {
+    static const std::vector<std::pair<std::string, std::string>> prefix_map = {
+        {"all_x_embedder.1-1.", "x_embedder."},
+        {"all_final_layer.1-1.", "final_layer."},
+    };
+
+    replace_with_prefix_map(name, prefix_map);
+
+    return name;
+}
+
 std::string convert_diffusers_dit_to_original_lumina2(std::string name) {
     int num_layers         = 30;
     int num_refiner_layers = 2;
@@ -896,6 +945,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
         name = convert_hunyuan_video_to_original_flux(name);
     } else if (sd_version_is_z_image(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);
     } else if (sd_version_is_anima(version)) {
         name = convert_other_dit_to_original_anima(name);
     } else if (sd_version_is_krea2(version)) {
@@ -999,7 +1050,30 @@ std::string convert_diffusers_vae_to_original_sd1(std::string name) {
     return result;
 }
 
-std::string convert_diffusers_to_original_wan_vae(std::string name) {
+std::string convert_diffusers_to_original_wan_vae(std::string name, bool qwen_image_2_1 = false) {
+    if (qwen_image_2_1) {
+        for (int i = 0; i < 5; ++i) {
+            const auto index = std::to_string(i);
+            for (const auto& side : {std::string("encoder"), std::string("decoder")}) {
+                const bool encoder           = side == "encoder";
+                const std::string old_prefix = side + (encoder ? ".down_blocks." : ".up_blocks.") + index + ".";
+                const std::string new_prefix = side + (encoder ? ".downsamples." : ".upsamples.") + index + ".";
+                if (!starts_with(name, old_prefix)) {
+                    continue;
+                }
+                name.replace(0, old_prefix.size(), new_prefix);
+                const std::string layers = encoder ? "downsamples." : "upsamples.";
+                for (int j = 0; j < (encoder ? 2 : 3); ++j) {
+                    const auto old_resnet = new_prefix + "resnets." + std::to_string(j) + ".";
+                    const auto new_resnet = new_prefix + layers + std::to_string(j) + ".";
+                    replace_with_prefix_map(name, std::vector<std::pair<std::string, std::string>>{{old_resnet + "conv_shortcut.", new_resnet + "shortcut."},
+                                                                                                   {old_resnet, new_resnet + "residual."}});
+                }
+                replace_with_prefix_map(name, std::vector<std::pair<std::string, std::string>>{{new_prefix + (encoder ? "downsampler." : "upsampler."),
+                                                                                                new_prefix + layers + (encoder ? "2." : "3.")}});
+            }
+        }
+    }
     static const std::vector<std::pair<std::string, std::string>> prefix_map = {
         {"quant_conv.", "conv1."},
         {"post_quant_conv.", "conv2."},
@@ -1055,7 +1129,11 @@ std::string convert_diffusers_to_original_wan_vae(std::string name) {
     };
 
     replace_with_name_map(name, shared_name_map);
-    replace_with_prefix_map(name, prefix_map);
+    if (qwen_image_2_1) {
+        replace_with_prefix_map(name, std::vector<std::pair<std::string, std::string>>{{"quant_conv.", "conv1."}, {"post_quant_conv.", "conv2."}});
+    } else {
+        replace_with_prefix_map(name, prefix_map);
+    }
 
     // Only apply the ResNet-specific renaming if the tensor belongs to a ResNet block.
     // This prevents generic ".conv1." or ".conv2." matching on top-level encoder/decoder convolutions.
@@ -1071,7 +1149,7 @@ std::string convert_first_stage_model_name(std::string name, std::string prefix,
         return name;
     }
     if (sd_version_uses_wan_vae(version)) {
-        return convert_diffusers_to_original_wan_vae(name);
+        return convert_diffusers_to_original_wan_vae(name, version == VERSION_QWEN_IMAGE_2_1);
     }
     static std::unordered_map<std::string, std::string> vae_name_map = {
         {"decoder.post_quant_conv.", "post_quant_conv."},
@@ -1489,7 +1567,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
 
     replace_with_prefix_map(name, prefix_map);
 
-    if (sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
+    if (version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
         const std::string hf_vision_prefix = "text_encoders.llm.model.visual.";
         if (starts_with(name, hf_vision_prefix)) {
             name = "text_encoders.llm.visual." + name.substr(hf_vision_prefix.size());

src/pipeline/diffusion_engine.cpp

@@ -80,6 +80,7 @@ const char* model_version_to_str[] = {
     "LingBot Video",
     "Qwen Image",
     "Qwen Image Layered",
+    "Qwen Image 2.1",
     "Hunyuan Video",
     "Anima",
     "Flux.2",
@@ -100,6 +101,7 @@ const char* model_version_to_str[] = {
     "Krea2",
     "Mage Flow",
     "SenseNova U1.5",
+    "LLaDA-Image",
     "ESRGAN",
 };
 
@@ -856,7 +858,57 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf
     return true;
 }
 
+bool StableDiffusionGGML::set_sage_attention_enabled(bool enabled) {
+    if (!diffusion_model) {
+        return false;
+    }
+    if (enabled) {
+#ifndef SD_USE_UPSTREAM_GGML
+        auto* ctx = ggml_init({4 * ggml_tensor_overhead(), nullptr, true});
+        if (ctx == nullptr) {
+            return false;
+        }
+        auto* q        = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1);
+        auto* k        = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1);
+        auto* v        = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 128, 128, 1, 1);
+        auto* op       = ggml_sage_attn(ctx, q, k, v, 1.f / sqrtf(128.f), GGML_SAGE_ATTN_AUTO);
+        bool supported = true;
+        for (auto backend : backend_manager.runtime_backends(SDBackendModule::DIFFUSION)) {
+            if (!ggml_backend_supports_op(backend, op)) {
+                LOG_ERROR("SageAttention is unavailable on %s; it requires patched GGML, CUDA Toolkit 12.0 or newer, and SM80 or newer kernels",
+                          ggml_backend_name(backend));
+                supported = false;
+            }
+        }
+        ggml_free(ctx);
+        if (!supported) {
+            return false;
+        }
+#else
+        LOG_ERROR("SageAttention requires -DSD_USE_UPSTREAM_GGML=OFF and a CUDA backend");
+        return false;
+#endif
+    }
+    diffusion_model->set_sage_attention_enabled(enabled);
+    if (high_noise_diffusion_model) {
+        high_noise_diffusion_model->set_sage_attention_enabled(enabled);
+    }
+    if (enabled) {
+        LOG_INFO("Using SageAttention in the diffusion model; CUDA selects the supported kernel, unsupported layers use flash/default attention");
+    }
+    return true;
+}
+
 bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) {
+#ifdef SD_USE_UPSTREAM_GGML
+    LOG_WARN(
+        "Using upstream GGML: INT8 tensorwise/convrot is disabled and FP8 weights are "
+        "converted to F16 at load time. Some operators may be unsupported and performance "
+        "may be lower than with patched GGML.");
+#endif
+    if (!validate_tensor_types(sd_ctx_params->wtype, sd_ctx_params->tensor_type_rules)) {
+        return false;
+    }
     for (float scale : {sd_ctx_params->linear_scale, sd_ctx_params->attn_scale}) {
         if (!std::isfinite(scale) || scale < 0.f || (scale > 0.f && !std::isfinite(1.f / scale))) {
             LOG_ERROR("scale overrides must be finite positive values, or 0 to keep model defaults");
@@ -1125,6 +1177,9 @@ bool StableDiffusionGGML::validate_and_load_runners() {
             high_noise_diffusion_model->set_flash_attention_enabled(true);
         }
     }
+    if (sd_ctx_params->sage_attn && !set_sage_attention_enabled(true)) {
+        return false;
+    }
     LOG_VERBOSE("validating model metadata");
 
     std::set<std::string> ignore_tensors;
@@ -1288,6 +1343,7 @@ bool StableDiffusionGGML::build_denoiser() {
                    sd_version_is_anima(version) ||
                    sd_version_is_ernie_image(version) ||
                    sd_version_is_z_image(version) ||
+                   sd_version_is_llada_image(version) ||
                    sd_version_is_boogu_image(version) ||
                    sd_version_is_pid(version) ||
                    sd_version_is_ideogram4(version)) {
@@ -1308,6 +1364,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 (sd_version_is_llada_image(version)) {
+                default_flow_shift = 1.0f;  // unused: LLADA_IMAGE_SCHEDULER builds a fixed grid
             } else {
                 default_flow_shift = 3.f;
             }
@@ -1704,6 +1762,8 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun
         extension->collect_loras(all_loras);
     }
 
+    conditioning_cache_allowed_ = all_loras.empty();
+
     int64_t t0 = ggml_time_ms();
     end_runners();
     clear_lora_adapters();
@@ -2367,6 +2427,8 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
             } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
                 diffusion_params.extra = FluxDiffusionExtra{&guidance_tensor,
                                                             local_skip_layers};
+            } else if (version == VERSION_QWEN_IMAGE_2_1) {
+                diffusion_params.extra = QwenImage21DiffusionExtra{&condition.c_token_types};
             } else if (sd_version_is_anima(version)) {
                 diffusion_params.extra = AnimaDiffusionExtra{condition.c_t5_ids.empty() ? nullptr : &condition.c_t5_ids,
                                                              condition.c_t5_weights.empty() ? nullptr : &condition.c_t5_weights};
@@ -2387,6 +2449,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 (sd_version_is_llada_image(version)) {
+                diffusion_params.extra = LLaDAImageDiffusionExtra{
+                    condition.extra_c_crossattns.empty() ? nullptr : &condition.extra_c_crossattns[0]};
             } else if (sd_version_is_minimax_h3(version)) {
                 diffusion_params.extra = MiniMaxH3DiffusionExtra{
                     condition.c_token_types.empty() ? nullptr : &condition.c_token_types,
@@ -2567,7 +2632,7 @@ int StableDiffusionGGML::get_diffusion_model_down_factor() {
     if (sd_version_is_dit(version)) {
         if (sd_version_is_sensenova_u1(version)) {
             down_factor = 32;
-        } else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version)) {
+        } 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)) {
             down_factor = 2;
         } else {
             down_factor = 1;
@@ -2583,6 +2648,8 @@ int StableDiffusionGGML::get_latent_channel() {
             latent_channel = 128;
         } else if (sd_version_is_minimax_h3(version)) {
             latent_channel = 24;
+        } else if (version == VERSION_QWEN_IMAGE_2_1) {
+            latent_channel = 64;
         } else if (version == VERSION_WAN2_2_TI2V) {
             latent_channel = 48;
         } else if (sd_version_is_hunyuan_video(version)) {
@@ -2611,7 +2678,7 @@ int StableDiffusionGGML::get_latent_channel() {
 }
 
 int StableDiffusionGGML::get_image_channels() const {
-    return version == VERSION_QWEN_IMAGE_LAYERED ? 4 : 3;
+    return version == VERSION_QWEN_IMAGE_LAYERED || version == VERSION_QWEN_IMAGE_2_1 ? 4 : 3;
 }
 
 int StableDiffusionGGML::get_image_seq_len(int h, int w) {
@@ -2701,7 +2768,8 @@ sd::Tensor<float> StableDiffusionGGML::decode_first_stage(const sd::Tensor<float
     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)) {
+           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);
     }
     return decoded;
@@ -2764,6 +2832,8 @@ std::string StableDiffusionGGML::get_default_ref_image_preset(SDVersion version)
         return "mage_flow";
     } else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) {
         return "z_image_omni";
+    } else if (sd_version_is_llada_image(version)) {
+        return "llada_image";
     } else if (sd_version_is_krea2(version)) {
         // have to make a choice between "krea2_edit" mode (for lbouaraba/krea2edit)
         // and "krea2_ostris_edit" (for krea2 ostris edit)

src/pipeline/diffusion_engine.h

@@ -37,7 +37,7 @@ extern const char* model_version_to_str[];
 static inline bool sd_version_supports_ref_latent_img_cfg(SDVersion version) {
     return version == VERSION_FLUX ||
            sd_version_is_flux2(version) ||
-           sd_version_is_qwen_image(version) ||
+           (sd_version_is_qwen_image(version) && version != VERSION_QWEN_IMAGE_2_1) ||
            sd_version_is_mage_flow(version) ||
            sd_version_is_longcat(version) ||
            sd_version_is_z_image(version) ||
@@ -178,6 +178,7 @@ public:
     std::recursive_mutex execution_mutex;
     std::unique_ptr<ModelConfig> config_;
     RunnerState runner_state_;
+    bool conditioning_cache_allowed_ = false;
     bool executing_ = false;
 
     std::shared_ptr<Denoiser> denoiser;
@@ -312,6 +313,7 @@ public:
     bool init_model_loader(ModelLoader& model_loader, ModelConfig& configuration);
 
     bool init(const sd_ctx_params_t* sd_ctx_params);
+    bool set_sage_attention_enabled(bool enabled);
 
     bool uses_tae() const;
 

src/pipeline/image.cpp

@@ -10,6 +10,7 @@
 #include "model/vae/vae.hpp"
 #include "request.h"
 #include "runtime/denoiser.hpp"
+#include "runtime/image_preprocess.h"
 #include "upscaler.h"
 
 namespace sd::pipeline {
@@ -285,7 +286,8 @@ namespace sd::pipeline {
                     vae_width  = request->width;
                     vae_height = request->height;
                 } else {
-                    int target_pixels  = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : 1024 * 1024;
+                    int default_pixels = sd->version == VERSION_QWEN_IMAGE_2_1 ? request->width * request->height : 1024 * 1024;
+                    int target_pixels  = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : default_pixels;
                     int vae_image_size = std::min(target_pixels, request->width * request->height);
                     vae_width          = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]);
                     vae_height         = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0];
@@ -309,6 +311,9 @@ namespace sd::pipeline {
                             resized_ref_img.shape()[0]);
 
                 ref_latent = sd->encode_first_stage(resized_ref_img);
+                if (sd->version == VERSION_QWEN_IMAGE_2_1) {
+                    ref_images[i] = std::move(resized_ref_img);
+                }
             } else {
                 ref_latent = sd->encode_first_stage(ref_images[i]);
             }
@@ -471,8 +476,12 @@ namespace sd::pipeline {
                 }
                 condition_params.text            = request->negative_prompt;
                 condition_params.zero_out_masked = zero_out_masked;
-                uncond                           = sd->cond_stage_model->get_learned_condition(sd->n_threads,
-                                                                                               condition_params);
+                if (sd_version_is_llada_image(sd->version)) {
+                    // LLaDA-Image CFG keeps the source latent but drops its SigVQ features.
+                    condition_params.ref_images = nullptr;
+                }
+                uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
+                                                                     condition_params);
                 if (uncond.empty()) {
                     LOG_ERROR("failed to encode negative prompt");
                     return std::nullopt;
@@ -781,15 +790,9 @@ namespace sd::pipeline {
             return false;
         }
 
-        // MiniMax-H3 is video-only. Its denoiser always splits the packed latent into a video and an
-        // audio half, and only generate_video ever computes the audio length, so reaching this
-        // function with an H3 checkpoint is guaranteed to die on
-        // GGML_ASSERT(!audio_input_cache.empty()) with a core dump, after the several minutes it
-        // takes to load the weights, and with nothing in the output pointing at the missing --mode.
-        // (The AnimateDiff path below routes vid_gen back through here, but that is SD1.5 plus a
-        // motion module, never H3.)
-        if (sd_version_is_minimax_h3(sd->version)) {
-            LOG_ERROR("MiniMax-H3 is a video model and cannot be run in img_gen mode; use --mode vid_gen");
+        if (!sd_version_supports_image_generation(sd->version)) {
+            LOG_ERROR("%s cannot be run with generate_image(); use generate_video() or --mode vid_gen in the CLI",
+                      model_version_to_str[sd->version]);
             return false;
         }
 
@@ -798,6 +801,12 @@ namespace sd::pipeline {
         int64_t t0            = ggml_time_ms();
         sd->vae_tiling_params = sd_img_gen_params->vae_tiling_params;
         GenerationRequest request(sd, sd_img_gen_params);
+        sd::ImagePreprocessor preprocessing(sd_img_gen_params->image_preprocess.rules);
+        sd_img_gen_params_t processed_params = *sd_img_gen_params;
+        if (!preprocessing.prepare_inputs(processed_params, request.width, request.height))
+            return false;
+        sd_img_gen_params = &processed_params;
+        request.pm_params = processed_params.pm_params;
         LOG_INFO("generate_image %dx%d", request.width, request.height);
 
         sd->rng->manual_seed(request.seed);

src/pipeline/model_builders.cpp

@@ -20,6 +20,7 @@
 #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/minimax_h3.hpp"
@@ -28,6 +29,7 @@
 #include "model/diffusion/model.hpp"
 #include "model/diffusion/pid.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"
@@ -285,12 +287,19 @@ namespace sd::model_builders {
                                                                enable_vision,
                                                                weight_manager,
                                                                tokenizers);
-            result.diffusion   = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
-                                                                       tensor_storage_map,
-                                                                       "model.diffusion_model",
-                                                                       version,
-                                                                       weight_manager,
-                                                                       sd_ctx_params->model_args);
+            if (version == VERSION_QWEN_IMAGE_2_1) {
+                result.diffusion = std::make_shared<Qwen::QwenImage21Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                             tensor_storage_map,
+                                                                             "model.diffusion_model",
+                                                                             weight_manager);
+            } else {
+                result.diffusion = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                           tensor_storage_map,
+                                                                           "model.diffusion_model",
+                                                                           version,
+                                                                           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,
@@ -362,6 +371,19 @@ namespace sd::model_builders {
                                                                       "model.diffusion_model",
                                                                       version,
                                                                       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,
+                                                                      "text_encoders.llm",
+                                                                      "queryformer",
+                                                                      "text_projection",
+                                                                      "sigvq",
+                                                                      weight_manager,
+                                                                      tokenizers);
+            result.diffusion   = std::make_shared<LLaDAImage::LLaDAImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                              tensor_storage_map,
+                                                                              "model.diffusion_model",
+                                                                              weight_manager);
         } else if (sd_version_is_boogu_image(version)) {
             result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                tensor_storage_map,

src/pipeline/request.cpp

@@ -65,7 +65,7 @@ namespace sd::pipeline {
             return LCM_SCHEDULER;
         } else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) {
             return SIMPLE_SCHEDULER;
-        } else if (sd != nullptr && sd_version_is_flux(sd->version)) {
+        } else if (sd != nullptr && (sd_version_is_flux(sd->version) || sd->version == VERSION_QWEN_IMAGE_2_1)) {
             return FLUX_SCHEDULER;
         } else if (sd != nullptr && sd_version_is_flux2(sd->version)) {
             return FLUX2_SCHEDULER;
@@ -73,6 +73,8 @@ namespace sd::pipeline {
             return LTX2_SCHEDULER;
         } else if (sd != nullptr && sd_version_is_ideogram4(sd->version)) {
             return LOGIT_NORMAL_SCHEDULER;
+        } else if (sd != nullptr && sd_version_is_llada_image(sd->version)) {
+            return LLADA_IMAGE_SCHEDULER;
         }
         return DISCRETE_SCHEDULER;
     }

src/pipeline/video.cpp

@@ -15,6 +15,7 @@
 #include "model/vae/vae.hpp"
 #include "request.h"
 #include "runtime/denoiser.hpp"
+#include "runtime/image_preprocess.h"
 
 namespace sd::pipeline {
 
@@ -470,11 +471,15 @@ namespace sd::pipeline {
         sd::Tensor<float> end_image;
 
         if (sd_vid_gen_params->init_image.data) {
-            start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, request->width, request->height);
+            start_image = ensure_image_tensor_channels(
+                sd_image_to_tensor(sd_vid_gen_params->init_image, request->width, request->height),
+                sd->get_image_channels());
         }
 
         if (sd_vid_gen_params->end_image.data) {
-            end_image = sd_image_to_tensor(sd_vid_gen_params->end_image, request->width, request->height);
+            end_image = ensure_image_tensor_channels(
+                sd_image_to_tensor(sd_vid_gen_params->end_image, request->width, request->height),
+                sd->get_image_channels());
         }
 
         if (sd_version_is_minimax_h3(sd->version)) {
@@ -1152,6 +1157,10 @@ namespace sd::pipeline {
         condition_params.zero_out_masked       = true;
         condition_params.ref_images            = &latents.ref_images;
         condition_params.minimax_h3_references = &latents.minimax_presentation_refs;
+        condition_params.allow_cache =
+            sd_version_is_minimax_h3(sd->version) &&
+            sd->conditioning_cache_allowed_ &&
+            !request.use_uncond;
         if (sd_version_is_lingbot_video(sd->version) || sd_version_is_minimax_h3(sd->version)) {
             condition_params.ref_image_params.vlm_resize_mode = RefImageResizeMode::AREA;
         }
@@ -1416,7 +1425,9 @@ namespace sd::pipeline {
         sd::Tensor<float> video_mask = make_ltxav_video_denoise_mask(video_latent, 1.f);
 
         if (sd_vid_gen_params->init_image.data != nullptr) {
-            sd::Tensor<float> start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, image_width, image_height);
+            sd::Tensor<float> start_image = ensure_image_tensor_channels(
+                sd_image_to_tensor(sd_vid_gen_params->init_image, image_width, image_height),
+                sd->get_image_channels());
             if (!apply_ltxav_condition_image_by_latent_index(sd,
                                                              start_image,
                                                              &video_latent,
@@ -1429,7 +1440,9 @@ namespace sd::pipeline {
         }
 
         if (sd_vid_gen_params->end_image.data != nullptr) {
-            sd::Tensor<float> end_image        = sd_image_to_tensor(sd_vid_gen_params->end_image, image_width, image_height);
+            sd::Tensor<float> end_image = ensure_image_tensor_channels(
+                sd_image_to_tensor(sd_vid_gen_params->end_image, image_width, image_height),
+                sd->get_image_channels());
             sd::Tensor<float> end_image_latent = encode_ltxav_condition_image(sd, end_image, "end");
             if (end_image_latent.empty()) {
                 return false;
@@ -1518,6 +1531,7 @@ namespace sd::pipeline {
         img_gen_params.qwen_image_layers = 0;
         img_gen_params.circular_x        = sd_vid_gen_params->circular_x;
         img_gen_params.circular_y        = sd_vid_gen_params->circular_y;
+        img_gen_params.image_preprocess  = sd_vid_gen_params->image_preprocess;
 
         sd->animatediff_num_frames = n_frames;
         bool ok                    = generate_image(sd, &img_gen_params, frames_out, num_frames_out);
@@ -1548,6 +1562,11 @@ namespace sd::pipeline {
         sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params;
         sd->apply_circular_axes(sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y);
         GenerationRequest request(sd, sd_vid_gen_params);
+        sd::ImagePreprocessor preprocessing(sd_vid_gen_params->image_preprocess.rules);
+        sd_vid_gen_params_t processed_params = *sd_vid_gen_params;
+        if (!preprocessing.prepare_inputs(processed_params, request.width, request.height))
+            return false;
+        sd_vid_gen_params = &processed_params;
         if (fps_out != nullptr) {
             *fps_out = request.fps;
         }

src/runtime/denoiser.hpp

@@ -786,6 +786,54 @@ struct FluxScheduler : SigmaScheduler {
 };
 
 // https://github.com/black-forest-labs/flux2/blob/main/src/flux2/sampling.py#L244
+// LLaDA-Image does not use a shift-based flow schedule. The reference pipeline builds a
+// Kumaraswamy-shaped grid over t = linspace(0.001, 1, n + 1)[:-1]:
+//     schedule = (1 - (1 - t^1.17)^0.8)^1.1
+//     sigma    = 1 - schedule
+// Its scheduler config can also set use_uniform_sigmas, which replaces the whole curve with a
+// plain linspace(1, 0, n + 1)[:-1] pre-shift grid.
+struct LLaDAImageScheduler : SigmaScheduler {
+    bool uniform_sigmas = false;
+
+    explicit LLaDAImageScheduler(const char* extra_sample_args = nullptr) {
+        parse_extra_sample_args(extra_sample_args);
+    }
+
+    void parse_extra_sample_args(const char* extra_sample_args) {
+        for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "llada_image scheduler arg")) {
+            if (key == "uniform") {
+                if (!parse_strict_bool(value, uniform_sigmas)) {
+                    LOG_WARN("ignoring invalid llada_image scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            }
+        }
+    }
+
+    std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override {
+        std::vector<float> sigmas;
+        sigmas.reserve(n + 1);
+
+        if (n == 0) {
+            sigmas.push_back(1.0f);
+            return sigmas;
+        }
+
+        for (uint32_t i = 0; i < n; ++i) {
+            float progress = static_cast<float>(i) / static_cast<float>(n);
+            if (uniform_sigmas) {
+                sigmas.push_back(1.0f - progress);
+            } else {
+                float t        = 0.001f + progress * (1.0f - 0.001f);
+                float schedule = powf(1.0f - powf(1.0f - powf(t, 1.17f), 0.8f), 1.1f);
+                sigmas.push_back(1.0f - schedule);
+            }
+        }
+
+        sigmas.push_back(0.0f);
+        return sigmas;
+    }
+};
+
 struct Flux2Scheduler : SigmaScheduler {
     int image_seq_len = 0;
 
@@ -1123,6 +1171,11 @@ struct Denoiser {
                 scheduler = std::make_shared<Flux2Scheduler>(image_seq_len);
                 break;
             }
+            case LLADA_IMAGE_SCHEDULER: {
+                LOG_INFO("get_sigmas with LLaDA-Image scheduler");
+                scheduler = std::make_shared<LLaDAImageScheduler>(extra_sample_args);
+                break;
+            }
             case FLUX_SCHEDULER: {
                 LOG_INFO("get_sigmas with Flux scheduler");
                 scheduler = std::make_shared<FluxScheduler>(image_seq_len, extra_sample_args);

src/runtime/image_preprocess.cpp

@@ -0,0 +1,447 @@
+#include "image_preprocess.h"
+#include "core/util.h"
+
+#include <climits>
+#include <set>
+
+namespace sd {
+
+    static constexpr std::pair<const char*, ImageTarget> image_targets[] = {
+        {"init", ImageTarget::Init},
+        {"end", ImageTarget::End},
+        {"mask", ImageTarget::Mask},
+        {"control", ImageTarget::Control},
+        {"ref", ImageTarget::Ref},
+        {"ip-adapter", ImageTarget::IPAdapter},
+        {"id", ImageTarget::ID},
+        {"control-frame", ImageTarget::ControlFrame},
+    };
+
+    static constexpr std::pair<const char*, ImageResizeMode> image_resize_modes[] = {
+        {"auto", ImageResizeMode::Auto},
+        {"none", ImageResizeMode::None},
+        {"stretch", ImageResizeMode::Stretch},
+        {"crop", ImageResizeMode::Crop},
+        {"crop-resize", ImageResizeMode::CropResize},
+        {"fit-pad", ImageResizeMode::FitPad},
+    };
+
+    template <typename T, size_t N>
+    static bool parse_enum(const std::string& text, const std::pair<const char*, T> (&names)[N], T& value) {
+        for (const auto& entry : names) {
+            if (text == entry.first) {
+                value = entry.second;
+                return true;
+            }
+        }
+        return false;
+    }
+
+    template <typename T, size_t N>
+    static const char* enum_name(T value, const std::pair<const char*, T> (&names)[N]) {
+        for (const auto& entry : names) {
+            if (value == entry.second)
+                return entry.first;
+        }
+        return "unknown";
+    }
+
+    template <typename T>
+    static bool one_of(T value, std::initializer_list<T> choices) {
+        return std::find(choices.begin(), choices.end(), value) != choices.end();
+    }
+
+    static bool one_of(const std::string& value, std::initializer_list<const char*> choices) {
+        for (const char* choice : choices) {
+            if (value == choice)
+                return true;
+        }
+        return false;
+    }
+
+    static ImageResizeMode resolve_mode(const std::map<std::string, std::string>& options, ImageResizeMode default_mode) {
+        auto it              = options.find("mode");
+        ImageResizeMode mode = ImageResizeMode::Auto;
+        if (it != options.end())
+            parse_enum(it->second, image_resize_modes, mode);
+        if (mode != ImageResizeMode::Auto)
+            return mode;
+        return options.count("width") && default_mode == ImageResizeMode::None ? ImageResizeMode::Stretch : default_mode;
+    }
+
+    bool ImagePreprocessor::fail(const std::string& message) const {
+        LOG_ERROR("image preprocessing: %s", message.c_str());
+        valid_ = false;
+        return false;
+    }
+
+    ImagePreprocessor::ImagePreprocessor(const char* text) {
+        if (text == nullptr || trim(text).empty())
+            return;
+        for (const auto& part : split_string(text, ';')) {
+            ImagePreprocessRule rule;
+            std::set<std::string> keys;
+            if (trim(part).empty()) {
+                fail("empty rule");
+                return;
+            }
+            for (const auto& entry : split_string(part, ',')) {
+                size_t equal = entry.find('=');
+                if (equal == std::string::npos) {
+                    fail("expected key=value: " + entry);
+                    return;
+                }
+                std::string key   = trim(entry.substr(0, equal));
+                std::string value = trim(entry.substr(equal + 1));
+                bool ok           = !value.empty() && keys.insert(key).second;
+                int number        = 0;
+                if (key == "target") {
+                    ok &= parse_enum(value, image_targets, rule.target);
+                } else if (key == "index") {
+                    ok &= parse_strict_int(value, rule.index) && rule.index >= 0;
+                } else {
+                    if (key == "mode") {
+                        ImageResizeMode mode;
+                        ok &= parse_enum(value, image_resize_modes, mode);
+                    } else if (key == "filter") {
+                        ok &= one_of(value, {"auto", "nearest", "nearest-exact", "bilinear", "bicubic", "lanczos"});
+                    } else if (key == "antialias") {
+                        ok &= one_of(value, {"auto", "true", "false"});
+                    } else if (key == "canny") {
+                        ok &= one_of(value, {"true", "false"});
+                    } else if (key == "anchor") {
+                        ok &= one_of(value, {"center", "top", "bottom", "left", "right"});
+                    } else if (key == "width" || key == "height") {
+                        ok &= parse_strict_int(value, number) && number > 0;
+                    } else if (key == "pad_color") {
+                        ok &= value.size() == 7 || value.size() == 9;
+                        ok &= !value.empty() && value[0] == '#';
+                        for (size_t i = 1; i < value.size(); ++i)
+                            ok &= std::isxdigit(static_cast<unsigned char>(value[i])) != 0;
+                    } else {
+                        ok = false;
+                    }
+                    rule.options[key] = value;
+                }
+                if (!ok) {
+                    fail("invalid or duplicate option: " + entry);
+                    return;
+                }
+            }
+            if (!keys.count("target") || rule.options.empty() ||
+                rule.options.count("width") != rule.options.count("height") ||
+                (rule.index >= 0 && !one_of(rule.target, {ImageTarget::Ref, ImageTarget::ID, ImageTarget::ControlFrame}))) {
+                fail("invalid target, index, or incomplete dimensions: " + part);
+                return;
+            }
+            rules_.push_back(std::move(rule));
+        }
+        for (const auto& rule : rules_) {
+            const auto options = resolve_options(rule.target, std::max(0, rule.index));
+            if (options.count("antialias") && options.at("antialias") == "true" && options.count("filter") &&
+                one_of(options.at("filter"), {"nearest", "nearest-exact"})) {
+                fail("antialias requires bilinear, bicubic, or lanczos");
+                return;
+            }
+        }
+    }
+
+    std::map<std::string, std::string> ImagePreprocessor::resolve_options(ImageTarget target, int index) const {
+        std::map<std::string, std::string> options;
+        for (int specificity = 0; specificity < 2; ++specificity) {
+            for (const auto& rule : rules_) {
+                if (rule.target == target &&
+                    rule.index == (specificity == 0 ? -1 : index)) {
+                    for (const auto& entry : rule.options)
+                        options[entry.first] = entry.second;
+                }
+            }
+        }
+        return options;
+    }
+
+    bool ImagePreprocessor::validate_inputs(const sd_img_gen_params_t& params) const {
+        const std::map<ImageTarget, int> counts = {
+            {ImageTarget::Init, params.init_image.data != nullptr},
+            {ImageTarget::Mask, params.mask_image.data != nullptr},
+            {ImageTarget::Control, params.control_image.data != nullptr},
+            {ImageTarget::IPAdapter, params.ip_adapter_image.data != nullptr},
+            {ImageTarget::Ref, params.ref_images != nullptr ? params.ref_images_count : 0},
+            {ImageTarget::ID, params.pm_params.id_images != nullptr ? params.pm_params.id_images_count : 0},
+        };
+        for (const auto& rule : rules_) {
+            auto it   = counts.find(rule.target);
+            int count = it == counts.end() ? 0 : it->second;
+            if (count <= 0 || rule.index >= count) {
+                return fail(std::string("rule targets an unavailable image: ") + enum_name(rule.target, image_targets));
+            }
+        }
+        return valid_;
+    }
+
+    bool ImagePreprocessor::validate_inputs(const sd_vid_gen_params_t& params) const {
+        const std::map<ImageTarget, int> counts = {
+            {ImageTarget::Init, params.init_image.data != nullptr},
+            {ImageTarget::End, params.end_image.data != nullptr},
+            {ImageTarget::Ref, params.ref_images != nullptr ? params.ref_images_count : 0},
+            {ImageTarget::ControlFrame, params.control_frames != nullptr ? params.control_frames_size : 0},
+        };
+        for (const auto& rule : rules_) {
+            auto it   = counts.find(rule.target);
+            int count = it == counts.end() ? 0 : it->second;
+            if (count <= 0 || rule.index >= count)
+                return fail(std::string("rule targets an unavailable video input: ") + enum_name(rule.target, image_targets));
+        }
+        return valid_;
+    }
+
+    static int anchor_offset(int remaining, const std::string& anchor, bool horizontal) {
+        if (anchor == (horizontal ? "left" : "top"))
+            return 0;
+        if (anchor == (horizontal ? "right" : "bottom"))
+            return remaining;
+        return remaining / 2;
+    }
+
+    Tensor<float> ImagePreprocessor::apply_transform(const Tensor<float>& image, const std::map<std::string, std::string>& options, ImageTransform p, const std::string& label, ops::InterpolateMode default_filter) const {
+        auto value = [&](const char* key, const char* fallback) {
+            auto it = options.find(key);
+            return it == options.end() ? std::string(fallback) : it->second;
+        };
+        std::string filter        = value("filter", "auto");
+        ops::InterpolateMode mode = default_filter;
+        if (filter == "nearest")
+            mode = ops::InterpolateMode::Nearest;
+        if (filter == "nearest-exact")
+            mode = ops::InterpolateMode::NearestExact;
+        if (filter == "bilinear")
+            mode = ops::InterpolateMode::Bilinear;
+        if (filter == "bicubic")
+            mode = ops::InterpolateMode::Bicubic;
+        if (filter == "lanczos")
+            mode = ops::InterpolateMode::Lanczos;
+        bool filtered  = ops::is_2d_filter_interpolate_mode(mode);
+        bool antialias = value("antialias", "auto") == "true" ||
+                         (value("antialias", "auto") == "auto" && filtered &&
+                          (p.resize_width < p.crop_width || p.resize_height < p.crop_height));
+        if (antialias && !filtered) {
+            fail(label + ": antialias requires bilinear, bicubic, or lanczos");
+            return {};
+        }
+        auto cropped = ops::slice(ops::slice(image, 0, p.x, p.x + p.crop_width), 1, p.y, p.y + p.crop_height);
+        int channels = static_cast<int>(image.shape()[2]);
+        bool resize  = p.resize_width != p.crop_width || p.resize_height != p.crop_height;
+        if (resize && channels == 4 && filtered) {
+            for (int64_t i = 0, pixels = cropped.shape()[0] * cropped.shape()[1]; i < pixels; ++i) {
+                for (int c = 0; c < 3; ++c)
+                    cropped[i + c * pixels] *= cropped[i + 3 * pixels];
+            }
+        }
+        auto resized = ops::interpolate(cropped, {p.resize_width, p.resize_height, channels, 1}, mode, false, antialias);
+        if (resize && channels == 4 && filtered) {
+            for (int64_t i = 0, pixels = resized.shape()[0] * resized.shape()[1]; i < pixels; ++i) {
+                float alpha = std::clamp(resized[i + 3 * pixels], 0.f, 1.f);
+                for (int c = 0; c < 3; ++c)
+                    resized[i + c * pixels] = alpha > 1e-6f ? resized[i + c * pixels] / alpha : 0.f;
+            }
+        }
+        resized = ops::clamp(resized, 0.f, 1.f);
+        Tensor<float> output({p.width, p.height, channels, 1});
+        std::string color = value("pad_color", "#000000ff");
+        if (color.size() == 7)
+            color += "ff";
+        uint8_t rgba[4];
+        for (int c = 0; c < 4; ++c)
+            rgba[c] = static_cast<uint8_t>(std::strtoul(color.substr(1 + c * 2, 2).c_str(), nullptr, 16));
+        for (int c = 0; c < channels; ++c) {
+            float fill = rgba[channels == 1 ? 0 : c] / 255.f;
+            for (int y = 0; y < p.height; ++y) {
+                for (int x = 0; x < p.width; ++x) {
+                    output.index(x, y, c, 0) = x >= p.pad_x && x < p.pad_x + p.resize_width && y >= p.pad_y && y < p.pad_y + p.resize_height
+                                                   ? resized.index(x - p.pad_x, y - p.pad_y, c, 0)
+                                                   : fill;
+                }
+            }
+        }
+        LOG_INFO("preprocess %s: %dx%d crop=(%d,%d,%d,%d) resize=%dx%d pad=(%d,%d) output=%dx%d filter=%s(%d) antialias=%s",
+                 label.c_str(), p.source_width, p.source_height, p.x, p.y, p.crop_width, p.crop_height,
+                 p.resize_width, p.resize_height, p.pad_x, p.pad_y, p.width, p.height, filter.c_str(), static_cast<int>(mode), BOOL_STR(antialias));
+        return output;
+    }
+
+    Tensor<float> ImagePreprocessor::apply_geometry(const Tensor<float>& image, ImageTarget target, int index, int width, int height, ImageResizeMode default_mode, ops::InterpolateMode default_filter, ImageTransform* plan_out) const {
+        if (!valid_ || image.empty())
+            return {};
+        const std::string label = std::string(enum_name(target, image_targets)) + "[" + std::to_string(index) + "]";
+        auto options            = resolve_options(target, index);
+        if (image.dim() != 4 || image.shape()[3] != 1 || image.shape()[2] < 1 || image.shape()[2] > 4) {
+            fail(label + ": expected one image with 1 to 4 channels");
+            return {};
+        }
+        ImageTransform p;
+        p.source_width = p.crop_width = static_cast<int>(image.shape()[0]);
+        p.source_height = p.crop_height = static_cast<int>(image.shape()[1]);
+        int target_width                = width > 0 ? width : p.source_width;
+        int target_height               = height > 0 ? height : p.source_height;
+        if (options.count("width")) {
+            parse_strict_int(options.at("width"), target_width);
+            parse_strict_int(options.at("height"), target_height);
+        }
+        ImageResizeMode mode = resolve_mode(options, default_mode);
+        std::string anchor   = options.count("anchor") ? options.at("anchor") : "center";
+        p.width = p.resize_width = target_width;
+        p.height = p.resize_height = target_height;
+        if (mode == ImageResizeMode::None) {
+            if (options.count("width") && (target_width != p.source_width || target_height != p.source_height)) {
+                fail(label + ": mode=none conflicts with requested dimensions");
+                return {};
+            }
+            p.width = p.resize_width = p.source_width;
+            p.height = p.resize_height = p.source_height;
+        } else if (mode == ImageResizeMode::Crop || mode == ImageResizeMode::CropResize) {
+            if (mode == ImageResizeMode::Crop) {
+                p.crop_width  = target_width;
+                p.crop_height = target_height;
+            } else if (int64_t(p.source_width) * target_height > int64_t(p.source_height) * target_width) {
+                p.crop_width = std::max(1, static_cast<int>(int64_t(p.source_height) * target_width / target_height));
+            } else {
+                p.crop_height = std::max(1, static_cast<int>(int64_t(p.source_width) * target_height / target_width));
+            }
+            if (p.crop_width > p.source_width || p.crop_height > p.source_height) {
+                fail(label + ": crop exceeds source dimensions");
+                return {};
+            }
+            p.x = anchor_offset(p.source_width - p.crop_width, anchor, true);
+            p.y = anchor_offset(p.source_height - p.crop_height, anchor, false);
+        } else if (mode == ImageResizeMode::FitPad) {
+            double scale    = std::min(double(target_width) / p.source_width, double(target_height) / p.source_height);
+            p.resize_width  = std::max(1, std::min(target_width, static_cast<int>(std::round(p.source_width * scale))));
+            p.resize_height = std::max(1, std::min(target_height, static_cast<int>(std::round(p.source_height * scale))));
+            p.pad_x         = anchor_offset(target_width - p.resize_width, anchor, true);
+            p.pad_y         = anchor_offset(target_height - p.resize_height, anchor, false);
+        }
+        if (p.width <= 0 || p.height <= 0) {
+            fail(label + ": invalid output dimensions");
+            return {};
+        }
+        uint64_t max_pixels = std::min<uint64_t>(INT64_MAX, SIZE_MAX / sizeof(float)) / static_cast<uint64_t>(image.shape()[2]);
+        if (uint64_t(p.width) * p.height > max_pixels || uint64_t(p.resize_width) * p.resize_height > max_pixels) {
+            fail(label + ": image allocation size overflows");
+            return {};
+        }
+        if (plan_out != nullptr)
+            *plan_out = p;
+        return apply_transform(image, options, p, label, default_filter);
+    }
+
+    Tensor<float> ImagePreprocessor::preprocess_input(sd_image_t image, ImageTarget target, int index, int width, int height) {
+        if (image.data == nullptr || image.width == 0 || image.height == 0 || image.width > INT_MAX || image.height > INT_MAX || image.channel < 1 || image.channel > 4) {
+            fail(std::string(enum_name(target, image_targets)) + ": invalid input image");
+            return {};
+        }
+        auto tensor = sd_image_to_tensor(image);
+        if (target == ImageTarget::Mask && has_init_transform_) {
+            auto options = resolve_options(target, index);
+            if (image.width != init_transform_.source_width || image.height != init_transform_.source_height) {
+                fail("mask and init source dimensions must match");
+                return {};
+            }
+            bool geometry_override = options.count("width") || options.count("anchor") ||
+                                     (options.count("mode") && options.at("mode") != "auto");
+            if (geometry_override) {
+                ImageTransform p;
+                auto init_options            = resolve_options(ImageTarget::Init, 0);
+                ImageResizeMode default_mode = resolve_mode(init_options, ImageResizeMode::CropResize);
+                auto result                  = apply_geometry(tensor, target, index, init_transform_.width, init_transform_.height, default_mode, ops::InterpolateMode::NearestExact, &p);
+                if (result.empty())
+                    return {};
+                const auto& q = init_transform_;
+                if (p.x != q.x || p.y != q.y || p.crop_width != q.crop_width || p.crop_height != q.crop_height ||
+                    p.resize_width != q.resize_width || p.resize_height != q.resize_height || p.pad_x != q.pad_x || p.pad_y != q.pad_y || p.width != q.width || p.height != q.height) {
+                    fail("mask geometry conflicts with init; configure geometry on init and filter on mask");
+                    return {};
+                }
+                return result;
+            }
+            return apply_transform(tensor, options, init_transform_, "mask[0]", ops::InterpolateMode::NearestExact);
+        }
+        auto result = apply_geometry(tensor, target, index, width, height, width > 0 ? ImageResizeMode::CropResize : ImageResizeMode::None,
+                                     target == ImageTarget::Mask ? ops::InterpolateMode::NearestExact : ops::InterpolateMode::Nearest,
+                                     target == ImageTarget::Init ? &init_transform_ : nullptr);
+        if (target == ImageTarget::Init)
+            has_init_transform_ = !result.empty();
+        return result;
+    }
+
+    ImagePreprocessor::~ImagePreprocessor() {
+        for (const auto& image : owned_images_)
+            std::free(image.data);
+    }
+
+    bool ImagePreprocessor::prepare_image(sd_image_t& image, ImageTarget target, int index, int width, int height) {
+        if (image.data == nullptr)
+            return true;
+        auto options = resolve_options(target, index);
+        bool canny   = options.count("canny") && options.at("canny") == "true";
+        auto tensor  = preprocess_input(image, target, index, width, height);
+        if (tensor.empty())
+            return false;
+        auto output = tensor_to_sd_image(tensor);
+        if (output.data == nullptr)
+            return fail("could not allocate input preprocessing buffer");
+        owned_images_.push_back(output);
+        if (canny && !preprocess_canny(output, 0.08f, 0.08f, 0.8f, 1.f, false))
+            return fail("Canny preprocessing failed");
+        image = output;
+        return true;
+    }
+
+    bool ImagePreprocessor::prepare_array(sd_image_t*& images, int count, ImageTarget target, std::vector<sd_image_t>& storage, int width, int height) {
+        if (count < 0 || (count > 0 && images == nullptr))
+            return fail(std::string("invalid image array: ") + enum_name(target, image_targets));
+        if (count == 0)
+            return true;
+        storage.assign(images, images + count);
+        for (int i = 0; i < count; ++i) {
+            if (storage[i].data == nullptr)
+                return fail(std::string("empty image in array: ") + enum_name(target, image_targets));
+            if (!prepare_image(storage[i], target, i, width, height))
+                return false;
+        }
+        images = storage.data();
+        return true;
+    }
+
+    bool ImagePreprocessor::prepare_inputs(sd_img_gen_params_t& params, int width, int height) {
+        if (prepared_)
+            return fail("inputs have already been prepared");
+        prepared_ = true;
+        if (!valid_ || !validate_inputs(params))
+            return false;
+        if (!prepare_image(params.init_image, ImageTarget::Init, 0, width, height) ||
+            !prepare_image(params.mask_image, ImageTarget::Mask, 0, width, height) ||
+            !prepare_image(params.control_image, ImageTarget::Control, 0, width, height) ||
+            !prepare_image(params.ip_adapter_image, ImageTarget::IPAdapter, 0, -1, -1) ||
+            !prepare_array(params.ref_images, params.ref_images_count, ImageTarget::Ref, ref_images_) ||
+            !prepare_array(params.pm_params.id_images, params.pm_params.id_images_count, ImageTarget::ID, id_images_))
+            return false;
+        params.image_preprocess = {};
+        return true;
+    }
+
+    bool ImagePreprocessor::prepare_inputs(sd_vid_gen_params_t& params, int width, int height) {
+        if (prepared_)
+            return fail("inputs have already been prepared");
+        prepared_ = true;
+        if (!valid_ || !validate_inputs(params))
+            return false;
+        if (!prepare_image(params.init_image, ImageTarget::Init, 0, width, height) ||
+            !prepare_image(params.end_image, ImageTarget::End, 0, width, height) ||
+            !prepare_array(params.ref_images, params.ref_images_count, ImageTarget::Ref, ref_images_) ||
+            !prepare_array(params.control_frames, params.control_frames_size, ImageTarget::ControlFrame, control_frames_, width, height))
+            return false;
+        params.image_preprocess = {};
+        return true;
+    }
+
+}  // namespace sd

src/runtime/image_preprocess.h

@@ -0,0 +1,88 @@
+#ifndef __SD_RUNTIME_IMAGE_PREPROCESS_H__
+#define __SD_RUNTIME_IMAGE_PREPROCESS_H__
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include "core/tensor.hpp"
+#include "stable-diffusion.h"
+
+namespace sd {
+
+    enum class ImageTarget {
+        Init,
+        End,
+        Mask,
+        Control,
+        Ref,
+        IPAdapter,
+        ID,
+        ControlFrame,
+    };
+
+    enum class ImageResizeMode {
+        Auto,
+        None,
+        Stretch,
+        Crop,
+        CropResize,
+        FitPad,
+    };
+
+    struct ImageTransform {
+        int source_width  = 0;
+        int source_height = 0;
+        int x             = 0;
+        int y             = 0;
+        int crop_width    = 0;
+        int crop_height   = 0;
+        int resize_width  = 0;
+        int resize_height = 0;
+        int width         = 0;
+        int height        = 0;
+        int pad_x         = 0;
+        int pad_y         = 0;
+    };
+
+    struct ImagePreprocessRule {
+        ImageTarget target = ImageTarget::Init;
+        int index          = -1;
+        std::map<std::string, std::string> options;
+    };
+
+    class ImagePreprocessor {
+        std::vector<ImagePreprocessRule> rules_;
+        mutable bool valid_ = true;
+        ImageTransform init_transform_;
+        bool has_init_transform_ = false;
+        bool prepared_           = false;
+        std::vector<sd_image_t> owned_images_;
+        std::vector<sd_image_t> ref_images_;
+        std::vector<sd_image_t> id_images_;
+        std::vector<sd_image_t> control_frames_;
+
+        bool fail(const std::string& message) const;
+        std::map<std::string, std::string> resolve_options(ImageTarget target, int index) const;
+        Tensor<float> apply_transform(const Tensor<float>& image, const std::map<std::string, std::string>& options, ImageTransform plan, const std::string& label, ops::InterpolateMode default_filter) const;
+
+        bool prepare_image(sd_image_t& image, ImageTarget target, int index, int width, int height);
+        bool prepare_array(sd_image_t*& images, int count, ImageTarget target, std::vector<sd_image_t>& storage, int width = -1, int height = -1);
+
+    public:
+        explicit ImagePreprocessor(const char* rules = nullptr);
+        ~ImagePreprocessor();
+        ImagePreprocessor(const ImagePreprocessor&)            = delete;
+        ImagePreprocessor& operator=(const ImagePreprocessor&) = delete;
+        bool prepare_inputs(sd_img_gen_params_t& params, int width, int height);
+        bool prepare_inputs(sd_vid_gen_params_t& params, int width, int height);
+        bool is_valid() const { return valid_; }
+        bool validate_inputs(const sd_img_gen_params_t& params) const;
+        bool validate_inputs(const sd_vid_gen_params_t& params) const;
+        Tensor<float> apply_geometry(const Tensor<float>& image, ImageTarget target, int index, int width, int height, ImageResizeMode default_mode = ImageResizeMode::Stretch, ops::InterpolateMode default_filter = ops::InterpolateMode::Nearest, ImageTransform* plan_out = nullptr) const;
+        Tensor<float> preprocess_input(sd_image_t image, ImageTarget target, int index = 0, int width = -1, int height = -1);
+    };
+
+}  // namespace sd
+
+#endif  // __SD_RUNTIME_IMAGE_PREPROCESS_H__

src/runtime/preprocessing.hpp

@@ -165,16 +165,18 @@ static inline sd::Tensor<float> convolve_tensor(const sd::Tensor<float>& input,
     return output;
 }
 
-static inline sd::Tensor<float> grayscale_tensor(const sd::Tensor<float>& rgb_img) {
-    GGML_ASSERT(rgb_img.dim() == 4);
-    GGML_ASSERT(rgb_img.shape()[2] >= 3);
-    sd::Tensor<float> grayscale({rgb_img.shape()[0], rgb_img.shape()[1], 1, rgb_img.shape()[3]});
-    for (int64_t iy = 0; iy < rgb_img.shape()[1]; ++iy) {
-        for (int64_t ix = 0; ix < rgb_img.shape()[0]; ++ix) {
-            float r    = preprocessing_get_4d(rgb_img, ix, iy, 0, 0);
-            float g    = preprocessing_get_4d(rgb_img, ix, iy, 1, 0);
-            float b    = preprocessing_get_4d(rgb_img, ix, iy, 2, 0);
-            float gray = 0.2989f * r + 0.5870f * g + 0.1140f * b;
+static inline sd::Tensor<float> grayscale_tensor(const sd::Tensor<float>& image) {
+    GGML_ASSERT(image.dim() == 4);
+    GGML_ASSERT(image.shape()[2] >= 1);
+    sd::Tensor<float> grayscale({image.shape()[0], image.shape()[1], 1, image.shape()[3]});
+    for (int64_t iy = 0; iy < image.shape()[1]; ++iy) {
+        for (int64_t ix = 0; ix < image.shape()[0]; ++ix) {
+            float gray = preprocessing_get_4d(image, ix, iy, 0, 0);
+            if (image.shape()[2] >= 3) {
+                float g = preprocessing_get_4d(image, ix, iy, 1, 0);
+                float b = preprocessing_get_4d(image, ix, iy, 2, 0);
+                gray    = 0.2989f * gray + 0.5870f * g + 0.1140f * b;
+            }
             preprocessing_set_4d(grayscale, gray, ix, iy, 0, 0);
         }
     }
@@ -317,11 +319,12 @@ bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold,
     image_gray              = non_max_supression(G, theta);
     threshold_hystersis(&image_gray, high_threshold, low_threshold, weak, strong);
 
+    const uint32_t color_channels = img.channel == 2 || img.channel == 4 ? img.channel - 1 : img.channel;
     for (uint32_t iy = 0; iy < img.height; ++iy) {
         for (uint32_t ix = 0; ix < img.width; ++ix) {
             float gray = preprocessing_get_4d(image_gray, ix, iy, 0, 0);
             gray       = inverse ? 1.0f - gray : gray;
-            for (uint32_t c = 0; c < img.channel; ++c) {
+            for (uint32_t c = 0; c < color_channels; ++c) {
                 preprocessing_set_4d(image, gray, ix, iy, c, 0);
             }
         }

src/stable-diffusion.cpp

@@ -26,13 +26,26 @@ static float get_cache_reuse_threshold(const sd_cache_params_t& params) {
 }
 
 const char* sd_type_name(enum sd_type_t type) {
-    if ((int)type < std::min<int>(SD_TYPE_COUNT, GGML_TYPE_COUNT)) {
-        return ggml_type_name((ggml_type)type);
+    if (type == SD_TYPE_F8_E4M3) {
+        return "f8_e4m3";
+    }
+    if (type == SD_TYPE_F8_E5M2) {
+        return "f8_e5m2";
+    }
+    const auto ggml_type = sd_type_to_ggml_type(type);
+    if (ggml_type != GGML_TYPE_COUNT) {
+        return ggml_type_name(ggml_type);
     }
     return NONE_STR;
 }
 
 enum sd_type_t str_to_sd_type(const char* str) {
+    if (!strcmp(str, "f8_e4m3")) {
+        return SD_TYPE_F8_E4M3;
+    }
+    if (!strcmp(str, "f8_e5m2")) {
+        return SD_TYPE_F8_E5M2;
+    }
     for (int i = 0; i < std::min<int>(SD_TYPE_COUNT, GGML_TYPE_COUNT); i++) {
         auto trait = ggml_get_type_traits((ggml_type)i);
         if (!strcmp(str, trait->type_name)) {
@@ -124,6 +137,7 @@ const char* scheduler_to_str[] = {
     "flux2",
     "flux",
     "beta",
+    "llada_image",
 };
 
 static_assert(SCHEDULER_COUNT == sizeof(scheduler_to_str) / sizeof(scheduler_to_str[0]),
@@ -323,6 +337,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_params->eager_load                = false;
     sd_ctx_params->enable_mmap               = false;
     sd_ctx_params->diffusion_flash_attn      = false;
+    sd_ctx_params->sage_attn                 = false;
     sd_ctx_params->linear_scale              = 0.f;
     sd_ctx_params->attn_scale                = 0.f;
     sd_ctx_params->vae_format                = SD_VAE_FORMAT_AUTO;
@@ -378,6 +393,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "auto_fit: %s\n"
              "flash_attn: %s\n"
              "diffusion_flash_attn: %s\n"
+             "sage_attn: %s\n"
              "linear_scale: %g\n"
              "attn_scale: %g\n"
              "vae_format: %s\n",
@@ -417,6 +433,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              BOOL_STR(sd_ctx_params->auto_fit),
              BOOL_STR(sd_ctx_params->flash_attn),
              BOOL_STR(sd_ctx_params->diffusion_flash_attn),
+             BOOL_STR(sd_ctx_params->sage_attn),
              sd_ctx_params->linear_scale,
              sd_ctx_params->attn_scale,
              sd_vae_format_name(sd_ctx_params->vae_format));
@@ -613,14 +630,6 @@ struct sd_ctx_t {
     StableDiffusionGGML* sd = nullptr;
 };
 
-static 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);
-}
-
-static bool sd_version_supports_image_generation(SDVersion version) {
-    return !sd_version_supports_video_generation(version);
-}
-
 sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t));
     if (sd_ctx == nullptr) {
@@ -742,28 +751,20 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                            int* num_frames_out,
                            sd_audio_t** audio_out,
                            int* fps_out) {
-    if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) {
-        if (fps_out != nullptr) {
-            *fps_out = 0;
-        }
-        return false;
-    }
-
-    if (frames_out != nullptr) {
+    if (frames_out != nullptr)
         *frames_out = nullptr;
-    }
-    if (audio_out != nullptr) {
+    if (audio_out != nullptr)
         *audio_out = nullptr;
-    }
-    if (num_frames_out != nullptr) {
+    if (num_frames_out != nullptr)
         *num_frames_out = 0;
+    if (fps_out != nullptr)
+        *fps_out = 0;
+    if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) {
+        return false;
     }
 
     StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd);
     if (!execution.ready) {
-        if (fps_out != nullptr) {
-            *fps_out = 0;
-        }
         return false;
     }
 

src/upscaler.cpp

@@ -111,10 +111,22 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
 
 sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_tensor) {
     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");
+        return {};
+    }
+
+    const bool has_alpha = input_tensor.shape()[2] == 4;
+    sd::Tensor<float> rgb;
+    if (has_alpha) {
+        rgb = sd::ops::slice(input_tensor, 2, 0, 3);
+    }
+    const sd::Tensor<float>& model_input = has_alpha ? rgb : input_tensor;
     sd::Tensor<float> upscaled;
     const int scale = esrgan_upscaler->config.scale;
     if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) {
-        upscaled = esrgan_upscaler->compute(n_threads, input_tensor);
+        upscaled = esrgan_upscaler->compute(n_threads, 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);
@@ -125,7 +137,7 @@ sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
             return output_tile;
         };
 
-        upscaled = process_tiles_2d(input_tensor,
+        upscaled = process_tiles_2d(model_input,
                                     static_cast<int>(input_tensor.shape()[0] * scale),
                                     static_cast<int>(input_tensor.shape()[1] * scale),
                                     scale,
@@ -141,6 +153,14 @@ sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
         LOG_ERROR("esrgan compute failed");
         return {};
     }
+    if (has_alpha) {
+        auto alpha       = sd::ops::slice(input_tensor, 2, 3, 4);
+        auto alpha_shape = alpha.shape();
+        alpha_shape[0]   = upscaled.shape()[0];
+        alpha_shape[1]   = upscaled.shape()[1];
+        alpha            = sd::ops::interpolate(alpha, alpha_shape, sd::ops::InterpolateMode::Bilinear);
+        upscaled         = sd::ops::concat(upscaled, alpha, 2);
+    }
     return upscaled;
 }