fix: handle GPU memory reports and LLM encoding failures (#2020)

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

AGENTS.md

@@ -118,7 +118,8 @@ documentation.
 6. Run the narrowest useful build, test, or inspection command available.
 
 Follow `CONTRIBUTING.md` for formatting, naming, PR expectations, dependency
-update policy, and security rules.
+update policy, and security rules. For tokenizer additions, follow its embedded-data
+allowlist and default to an external `tokenizer.json`.
 
 ---
 

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,23 +327,21 @@ 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)
 
 target_sources(${SD_LIB} PRIVATE $<TARGET_OBJECTS:zip>)
 target_link_libraries(${SD_LIB} PUBLIC ggml)
+find_package(Threads REQUIRED)
+target_link_libraries(${SD_LIB} PRIVATE Threads::Threads)
+target_link_libraries(${SD_LIB} PRIVATE onig sd-utf8proc)
+if (SD_CUDA)
+    find_package(CUDAToolkit REQUIRED)
+    # Keep the driver stub on downstream link lines when no driver is installed.
+    target_link_libraries(${SD_LIB} PUBLIC CUDA::cuda_driver)
+    set_property(SOURCE src/core/ggml_extend_backend.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_USE_CUDA)
+endif()
 target_include_directories(${SD_LIB} PUBLIC . src include)
 target_include_directories(${SD_LIB} PRIVATE src/core)
 target_include_directories(${SD_LIB} PUBLIC . 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,17 +37,31 @@ 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.
 
 When adding or modifying model implementations, follow the model config and weight detection conventions in [docs/model_config.md](docs/model_config.md).
 
+## Tokenizer Data
+
+New model integrations must use an external `tokenizer.json` by default. Do not
+embed new vocabularies or merge tables solely for less widely used models;
+these tables increase the binary size for every user.
+
+The embedded-data allowlist is CLIP, T5/UMT5, Qwen 2/3, Mistral, and Gemma 3/4.
+Models may reuse an existing embedded tokenizer when its vocabulary and behavior
+match their text encoder. Gemma 2 and GPT-OSS require external JSON files.
+
+Adding to this allowlist requires maintainer approval, supported by the model's
+usage, reuse across models, and measured binary-size cost. Document the matching
+JSON and CLI option for models that require an external tokenizer, and fail
+initialization clearly when it is missing.
+
 ## AI-Assisted Contributions
 
 AI tools may be used to assist development, but contributors are responsible for the quality and correctness of the submitted code.

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()

cmake/stable-diffusion-config.cmake.in

@@ -10,6 +10,10 @@ set(SD_BIN_DIR "@PACKAGE_SD_BIN_INSTALL_DIR@")
 
 include(CMakeFindDependencyMacro)
 find_dependency(ggml REQUIRED HINTS "${SD_LIB_DIR}/cmake")
+find_dependency(Threads REQUIRED)
+if(@SD_CUDA@)
+    find_dependency(CUDAToolkit REQUIRED)
+endif()
 
 if(NOT TARGET stable-diffusion)
     find_library(stable-diffusion_LIBRARY stable-diffusion
@@ -22,12 +26,16 @@ if(NOT TARGET stable-diffusion)
     set_target_properties(stable-diffusion
         PROPERTIES
             INTERFACE_INCLUDE_DIRECTORIES "${SD_INCLUDE_DIR}"
-            INTERFACE_LINK_LIBRARIES "ggml::ggml"
+            INTERFACE_LINK_LIBRARIES "ggml::ggml;Threads::Threads"
             IMPORTED_LINK_INTERFACE_LANGUAGES "CXX"
             IMPORTED_LOCATION "${stable-diffusion_LIBRARY}"
             INTERFACE_COMPILE_FEATURES "c_std_11;cxx_std_17"
             POSITION_INDEPENDENT_CODE ON)
 
+    if(@SD_CUDA@)
+        set_property(TARGET stable-diffusion APPEND PROPERTY INTERFACE_LINK_LIBRARIES CUDA::cuda_driver)
+    endif()
+
     if(SD_SHARED_LIB)
         target_compile_definitions(stable-diffusion
             INTERFACE SD_BUILD_SHARED_LIB)

cmake/stable-diffusion.pc.in

@@ -7,5 +7,5 @@ Name: stable-diffusion
 Description: Diffusion model(SD,Flux,Wan,Qwen Image,Z-Image,...) inference in pure C/C++
 Version: @SDCPP_BUILD_VERSION@
 Libs: -L${libdir} -lstable-diffusion
-Libs.private: -lggml -lggml-base
+Libs.private: -lggml -lggml-base @CMAKE_THREAD_LIBS_INIT@
 Cflags: -I${includedir}

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.

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

@@ -12,13 +12,17 @@ Lens uses a Lens diffusion transformer, the FLUX.2 VAE, and GPT-OSS-20B as the L
     - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main
 - Download GPT-OSS-20B
     - gguf: https://huggingface.co/unsloth/gpt-oss-20b-GGUF/tree/main
+- Download GPT-OSS-20B tokenizer.json
+    - https://huggingface.co/openai/gpt-oss-20b/tree/main
+
+Lens and Lens Turbo require an external GPT-OSS `tokenizer.json` matching the text encoder checkpoint. Save it as `tokenizer_gpt_oss.json` and pass it with `--tokenizer`; the tokenizer is not embedded in sd.cpp. See [JSON tokenizers](tokenizers.md) for CLI and C API usage.
 
 ## Examples
 
 ### Lens
 
 ```
-.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\lens_bf16.safetensors --llm "..\models\text_encoders\gpt-oss-20b-UD-Q8_K_XL.gguf" --vae ..\models\vae\flux2_ae.safetensors --cfg-scale 5.0  -p "A crystal dragon soaring through an aurora borealis sky, its entire body made of transparent faceted crystal refracting the green and purple aurora light into rainbow spectra, ice particles trailing from its wings, high fantasy digital art" --diffusion-fa -v
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\lens_bf16.safetensors --llm "..\models\text_encoders\gpt-oss-20b-UD-Q8_K_XL.gguf" --tokenizer ..\models\tokenizers\tokenizer_gpt_oss.json --vae ..\models\vae\flux2_ae.safetensors --cfg-scale 5.0  -p "A crystal dragon soaring through an aurora borealis sky, its entire body made of transparent faceted crystal refracting the green and purple aurora light into rainbow spectra, ice particles trailing from its wings, high fantasy digital art" --diffusion-fa -v
 ```
 
 <img width="256" alt="Lens example" src="../assets/lens/example.png" />
@@ -26,7 +30,7 @@ Lens uses a Lens diffusion transformer, the FLUX.2 VAE, and GPT-OSS-20B as the L
 ### Lens Turbo
 
 ```
-.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\lens_turbo_bf16.safetensors --llm "..\models\text_encoders\gpt-oss-20b-UD-Q8_K_XL.gguf" --vae ..\models\vae\flux2_ae.safetensors --cfg-scale 1.0  -p "A crystal dragon soaring through an aurora borealis sky, its entire body made of transparent faceted crystal refracting the green and purple aurora light into rainbow spectra, ice particles trailing from its wings, high fantasy digital art" --diffusion-fa -v --steps 4
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\lens_turbo_bf16.safetensors --llm "..\models\text_encoders\gpt-oss-20b-UD-Q8_K_XL.gguf" --tokenizer ..\models\tokenizers\tokenizer_gpt_oss.json --vae ..\models\vae\flux2_ae.safetensors --cfg-scale 1.0  -p "A crystal dragon soaring through an aurora borealis sky, its entire body made of transparent faceted crystal refracting the green and purple aurora light into rainbow spectra, ice particles trailing from its wings, high fantasy digital art" --diffusion-fa -v --steps 4
 ```
 
 <img width="256" alt="Lens Turbo example" src="../assets/lens/turbo_example.png" />

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

@@ -11,6 +11,8 @@ In stable-diffusion.cpp, PiD currently runs as an image edit pipeline: provide a
     - safetensors: https://huggingface.co/Comfy-Org/PixelDiT/tree/main/diffusion_models
 - Download Gemma 2 2B
     - safetensors: https://huggingface.co/Comfy-Org/PixelDiT/tree/main/text_encoders
+- Download Gemma 2 2B tokenizer.json
+    - https://huggingface.co/google/gemma-2-2b/tree/main
 - Download the VAE that matches the PiD checkpoint backbone
     - safetensors: https://huggingface.co/nvidia/PiD/tree/main/checkpoints
     - Flux / Z-Image PiD: use the Flux VAE and pass `--vae-format flux`
@@ -20,10 +22,12 @@ In stable-diffusion.cpp, PiD currently runs as an image edit pipeline: provide a
 
 The official PiD model card should be checked before use. At the time of the initial PiD release, the official weights are under the NSCLv1 non-commercial license.
 
+PiD and PiD 1.5 require an external Gemma 2 `tokenizer.json` matching the text encoder checkpoint. Save it as `tokenizer_gemma2.json` and pass it with `--tokenizer`; the tokenizer is not embedded in sd.cpp. See [JSON tokenizers](tokenizers.md) for CLI and C API usage.
+
 ## Examples
 
 ```
-.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\pid_flux1_512_to_2048_4step_bf16.safetensors --llm "..\models\text_encoders\gemma_2_2b_it_elm_bf16.safetensors" --vae ..\models\vae\ae.sft --vae-format flux --cfg-scale 1.0  -p "a lovely cat" -r ..\assets\ernie_image\turbo_example.png --diffusion-fa -v --steps 4 -H 2048 -W 2048 --rng cpu
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\pid_flux1_512_to_2048_4step_bf16.safetensors --llm "..\models\text_encoders\gemma_2_2b_it_elm_bf16.safetensors" --tokenizer ..\models\tokenizers\tokenizer_gemma2.json --vae ..\models\vae\ae.sft --vae-format flux --cfg-scale 1.0  -p "a lovely cat" -r ..\assets\ernie_image\turbo_example.png --diffusion-fa -v --steps 4 -H 2048 -W 2048 --rng cpu
 ```
 
 Before:

docs/qwen_image_2.1.md

@@ -0,0 +1,41 @@
+# 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`.

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

@@ -0,0 +1,107 @@
+# JSON tokenizers
+
+Use a Hugging Face `tokenizer.json` to supply the tokenizer vocabulary, merges,
+added tokens, and processing stages. **PiD (including PiD 1.5) and Lens (including
+Lens Turbo) require an external JSON**; their Gemma 2 and GPT-OSS tokenizers are
+not embedded. Initialization fails if the main tokenizer is missing. Other
+models keep their embedded tokenizer when this option is omitted.
+
+```shell
+sd-cli --diffusion-model model.gguf --llm text_encoder.gguf \
+  --tokenizer tokenizer_gemma2.json --vae vae.safetensors -p "a cat"
+```
+
+Choose the JSON belonging to the text encoder checkpoint. Checking that IDs fit
+the embedding table does not establish that two vocabularies have the same
+meaning. The JSON file is loaded when the text encoder is created; its embedded
+vocabulary is not loaded in this case.
+
+| Model | Required text encoder tokenizer | Example |
+| --- | --- | --- |
+| PiD / PiD 1.5 | Gemma 2 matching the text encoder checkpoint | `--tokenizer tokenizer_gemma2.json` |
+| Lens / Lens Turbo | GPT-OSS matching the text encoder checkpoint | `--tokenizer tokenizer_gpt_oss.json` |
+
+The Gemma 3/4 tokenizer used by LTX-2 remains embedded.
+
+| Option | Encoder |
+| --- | --- |
+| `--tokenizer FILE` | Main LLM/BPE encoder: Gemma 2, Gemma 3, Qwen 2/3, Mistral, GPT-OSS; also Anima and HiDream-O1 |
+| `--tokenizer FILE` | Shared CLIP tokenizer in SD1/SD2/SDXL, or CLIP-L in Flux |
+| `--tokenizer clip-l=FILE` | Separate CLIP-L in SD3 or Flux |
+| `--tokenizer clip-g=FILE` | Separate CLIP-G in SD3 |
+
+Use comma-separated assignments to configure multiple slots, for example
+`--tokenizer main=main.json,clip-l=clip_l.json,clip-g=clip_g.json`.
+A plain file path is equivalent to `main=FILE`. You may also repeat `--tokenizer`
+with explicit assignments, such as `--tokenizer main=main.json --tokenizer clip-l=clip.json`.
+Empty assignment paths, unknown keys, malformed assignments and
+duplicate slots are rejected. Commas separate entries in the assignment form;
+quote the complete argument when paths contain spaces.
+
+SD3 overrides must name the `clip-l` or `clip-g` slot. SDXL uses one shared
+tokenizer for both CLIP encoders. Do not supply both `main` and `clip-l` for Flux.
+A slot targeting an absent or unsupported encoder fails initialization.
+T5/SentencePiece Unigram tokenizers are outside this implementation's scope.
+
+For example, SD3 can load the same CLIP JSON into both slots:
+
+```shell
+sd-cli --diffusion-model sd3.gguf --clip_l clip_l.safetensors \
+  --clip_g clip_g.safetensors --t5xxl t5xxl.gguf --vae vae.safetensors \
+  --tokenizer clip-l=tokenizer_clip.json,clip-g=tokenizer_clip.json \
+  -p "a cat"
+```
+
+The C API accepts the same string in `sd_ctx_params_t::tokenizer`. A null or
+empty value keeps an embedded tokenizer where available; PiD and Lens require
+a nonempty main tokenizer path. The CLI passes the string through;
+`TokenizerConfig` parses and validates it when text encoders are initialized.
+
+```c
+sd_ctx_params_t params;
+sd_ctx_params_init(&params);
+params.tokenizer = "clip-l=tokenizer_clip.json,clip-g=tokenizer_clip.json";
+```
+
+Rebuild applications against the updated public header when using the updated
+library.
+
+## Supported components
+
+| Stage | Supported configurations |
+| --- | --- |
+| Normalizer | `Sequence`, `NFC`, `Lowercase`, `Replace` with String/Regex patterns |
+| PreTokenizer | `Sequence`, `Split` with String/Regex patterns, all five delimiter behaviors and `invert`; `ByteLevel` with `add_prefix_space` and `use_regex` |
+| Model | Deterministic `BPE`, string or array-pair merges, `unk_token`, `fuse_unk`, `byte_fallback`, `ignore_merges`, `end_of_word_suffix` |
+| PostProcessor | Single-sequence `TemplateProcessing` with at most one prefix and one suffix token, `RobertaProcessing`, `ByteLevel` |
+| Decoder | `Sequence`, `Replace`, `ByteLevel`, `ByteFallback`, `Fuse` |
+| AddedToken | Special and ordinary added tokens, original IDs, raw or normalized matching, leftmost-longest matching |
+
+`ByteLevel.use_regex` defaults to true when omitted. ByteLevel postprocessing
+changes offsets only and adds no tokens. Added tokens with `single_word`,
+`lstrip`, or `rstrip` enabled, nonzero BPE dropout, nonempty
+`continuing_subword_prefix`, and unsupported component types fail loading.
+New added-token IDs must follow the model vocabulary consecutively; configurations
+whose IDs Hugging Face would reassign are rejected.
+JSON `padding` and `truncation` must be null. This API returns IDs, not offsets,
+type IDs, or paired-input encodings; the pair template is not used.
+
+The pipeline covers the CLIP, Gemma 2, Gemma 3, GPT-OSS, Mistral 3, Qwen 2 and
+Qwen 3 JSON configurations used by the differential test. It does not imply
+support for every tokenizer published under those model names.
+
+## Prompt integration
+
+Prompt attention parsing and model-specific chat/image templates remain in the
+conditioner. Raw `encode()` does not add BOS/EOS. The conditioner concatenates
+weighted prompt fragments, then the existing padding/chunking step applies the
+JSON single-sequence template once per sequence or CLIP chunk. Padding ID,
+direction, length limits and attention masks remain text encoder policies.
+CLIP requires both BOS and EOS because its chunking reserves those positions.
+
+The internal `encode()`, `tokenize()`, and `decode()` interfaces return a success
+flag and write to an output parameter. A successful result may be empty; a failed
+call clears its output. JSON tokenizer input, normalization, and regex failures
+return `false` with diagnostic information instead of throwing. Invalid
+JSON, unsupported stages, conflicting IDs and IDs outside the encoder embedding
+table fail initialization instead of falling back to the embedded tokenizer.

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

@@ -410,6 +410,11 @@ ArgOptions SDContextParams::get_options() {
          "path to the llm text encoder. For example: (qwenvl2.5 for qwen-image, mistral-small3.2 for flux2, ...)",
          0,
          &llm_path},
+        {"",
+         "--tokenizer",
+         "tokenizer.json path, or comma-separated main=FILE,clip-l=FILE,clip-g=FILE assignments; required for PiD and Lens",
+         (int)',',
+         &tokenizer},
         {"",
          "--llm_vision",
          "path to the llm vit",
@@ -613,13 +618,17 @@ 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",
          true, &diffusion_conv_direct},
         {"",
          "--vae-conv-direct",
-         "use ggml_conv2d_direct in the vae model",
+         "use direct 2D and 3D convolutions in the vae model",
          true, &vae_conv_direct},
     };
 
@@ -896,6 +905,7 @@ std::string SDContextParams::to_string() const {
         << "  t5xxl_path: \"" << t5xxl_path << "\",\n"
         << "  llm_path: \"" << llm_path << "\",\n"
         << "  llm_vision_path: \"" << llm_vision_path << "\",\n"
+        << "  tokenizer: \"" << tokenizer << "\",\n"
         << "  diffusion_model_path: \"" << diffusion_model_path << "\",\n"
         << "  high_noise_diffusion_model_path: \"" << high_noise_diffusion_model_path << "\",\n"
         << "  uncond_diffusion_model_path: \"" << uncond_diffusion_model_path << "\",\n"
@@ -932,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"
@@ -963,6 +974,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.t5xxl_path                      = t5xxl_path.c_str();
     sd_ctx_params.llm_path                        = llm_path.c_str();
     sd_ctx_params.llm_vision_path                 = llm_vision_path.c_str();
+    sd_ctx_params.tokenizer                       = tokenizer.c_str();
     sd_ctx_params.diffusion_model_path            = diffusion_model_path.c_str();
     sd_ctx_params.high_noise_diffusion_model_path = high_noise_diffusion_model_path.c_str();
     sd_ctx_params.uncond_diffusion_model_path     = uncond_diffusion_model_path.c_str();
@@ -988,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;
@@ -1102,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",
+         "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},
         {"",
@@ -1741,7 +1754,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",
@@ -2210,7 +2223,12 @@ bool SDGenerationParams::from_json_str(
         LOG_ERROR("invalid end_image");
         return false;
     }
-    if (!parse_image_array_json_field(j, "ref_images", 3, width, height, ref_images)) {
+    if (!parse_image_array_json_field(j,
+                                      "ref_images",
+                                      3,
+                                      auto_resize_ref_image ? width : 0,
+                                      auto_resize_ref_image ? height : 0,
+                                      ref_images)) {
         LOG_ERROR("invalid ref_images");
         return false;
     }

examples/common/common.h

@@ -124,6 +124,7 @@ struct SDContextParams {
     std::string t5xxl_path;
     std::string llm_path;
     std::string llm_vision_path;
+    std::string tokenizer;
     std::string diffusion_model_path;
     std::string high_noise_diffusion_model_path;
     std::string uncond_diffusion_model_path;
@@ -169,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;
 

examples/server/routes_sdapi.cpp

@@ -244,8 +244,12 @@ 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,
+                                    request.gen_params.auto_resize_ref_image && request.gen_params.width_and_height_are_set()
+                                        ? request.gen_params.width
+                                        : 0,
+                                    request.gen_params.auto_resize_ref_image && request.gen_params.width_and_height_are_set()
+                                        ? request.gen_params.height
+                                        : 0,
                                     image_owner)) {
                 const sd_image_t& image = image_owner.get();
                 request.gen_params.set_width_and_height_if_unset(image.width, image.height);

include/stable-diffusion.h

@@ -79,6 +79,7 @@ enum scheduler_t {
     FLUX2_SCHEDULER,
     FLUX_SCHEDULER,
     BETA_SCHEDULER,
+    LLADA_IMAGE_SCHEDULER,
     SCHEDULER_COUNT
 };
 
@@ -244,6 +245,8 @@ typedef struct {
     bool disable_segmented_compute;  // Force monolithic graph execution even when automatic graph cutting would fit memory better
     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 {

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

@@ -102,7 +102,7 @@ namespace sd::backend_fit {
         for (const auto& [name, stored_tensor] : loader.get_tensor_storage_map()) {
             TensorStorage ts = stored_tensor;
             ComponentKind kind;
-            if (is_unused_tensor(ts.name) || !classify_tensor(ts.name, kind)) {
+            if (!classify_tensor(ts.name, kind)) {
                 continue;
             }
             if (ts.expected_type != GGML_TYPE_COUNT) {
@@ -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,
@@ -452,8 +464,13 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
                               int d0,
                               int d1,
                               int d2,
-                              bool force_prec_f32) {
-    if (force_prec_f32) {
+                              bool force_prec_f32,
+                              bool direct) {
+    if (direct) {
+        int64_t OC = w->ne[3] / IC;
+        int64_t N  = x->ne[3] / IC;
+        x          = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC);
+    } else if (force_prec_f32) {
         ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
 
         int64_t OC = w->ne[3] / IC;
@@ -605,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;
@@ -643,6 +661,14 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
     ggml_tensor* kqv = nullptr;
 
     auto build_kqv = [&](ggml_tensor* q_in, ggml_tensor* k_in, ggml_tensor* v_in, ggml_tensor* mask_in) -> ggml_tensor* {
+        const bool pad_head = d_head > 0 && d_head < 64 && q_in->ne[0] == d_head && k_in->ne[0] == d_head &&
+                              q_in->type == GGML_TYPE_F32 && k_in->type == GGML_TYPE_F32 &&
+                              v_in->type == GGML_TYPE_F32 && sd_backend_supports_cuda_mma(backend);
+        if (pad_head) {
+            // CUDA FA MMA starts at 64 channels; keep the original head's attention scale.
+            q_in = ggml_pad(ctx, q_in, 64 - d_head, 0, 0, 0);
+            k_in = ggml_pad(ctx, k_in, 64 - d_head, 0, 0, 0);
+        }
         if (kv_scale != 1.0f) {
             k_in = ggml_ext_scale(ctx, k_in, kv_scale);
         }
@@ -650,6 +676,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
 
         v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3));
         v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N);
+        if (pad_head) {
+            v_in = ggml_pad(ctx, v_in, 64 - d_head, 0, 0, 0);
+        }
         if (kv_scale != 1.0f) {
             v_in = ggml_ext_scale(ctx, v_in, kv_scale);
         }
@@ -679,10 +708,43 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
         if (kv_scale != 1.0f) {
             out = ggml_ext_scale(ctx, out, 1.0f / kv_scale);
         }
+        if (pad_head) {
+            out = ggml_ext_slice(ctx, out, 0, 0, d_head);
+        }
         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

@@ -153,7 +153,8 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
                               int d0              = 1,
                               int d1              = 1,
                               int d2              = 1,
-                              bool force_prec_f32 = false);
+                              bool force_prec_f32 = false,
+                              bool direct         = false);
 
 // w: [OC，IC, KD, 1 * 1]
 // x: [N, IC, ID, IH*IW]
@@ -219,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

@@ -8,8 +8,12 @@
 #include <stdexcept>
 #include <vector>
 
+#ifdef SD_USE_CUDA
+#include <cuda.h>
+#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) {
@@ -429,6 +433,70 @@ bool sd_backend_is_cpu(ggml_backend_t backend) {
     return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU;
 }
 
+bool sd_backend_supports_cuda_mma(ggml_backend_t backend) {
+#ifdef SD_USE_CUDA
+    if (!sd_backend_is(backend, "CUDA")) {
+        return false;
+    }
+    auto dev = ggml_backend_get_device(backend);
+    if (dev == nullptr) {
+        return false;
+    }
+    static std::mutex mutex;
+    static std::unordered_map<ggml_backend_dev_t, bool> cache;
+    std::lock_guard<std::mutex> lock(mutex);
+    auto it = cache.find(dev);
+    if (it != cache.end()) {
+        return it->second;
+    }
+    const bool supported = [&]() {
+        ggml_backend_dev_props props{};
+        ggml_backend_dev_get_props(dev, &props);
+        CUdevice device;
+        int major = 0, minor = 0;
+        if (props.device_id == nullptr || cuInit(0) != CUDA_SUCCESS ||
+            cuDeviceGetByPCIBusId(&device, props.device_id) != CUDA_SUCCESS ||
+            cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device) != CUDA_SUCCESS ||
+            cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device) != CUDA_SUCCESS) {
+            return false;
+        }
+        auto reg          = ggml_backend_dev_backend_reg(dev);
+        auto get_features = reinterpret_cast<ggml_backend_get_features_t>(
+            ggml_backend_reg_get_proc_address(reg, "ggml_backend_get_features"));
+        if (get_features == nullptr) {
+            return false;
+        }
+        // Match ggml's highest compiled architecture for this device, including PTX fallback.
+        const int cc      = 100 * major + 10 * minor;
+        int compiled_arch = 0;
+        for (auto feature = get_features(reg); feature != nullptr && feature->name != nullptr; ++feature) {
+            if (std::strcmp(feature->name, "ARCHS") != 0 || feature->value == nullptr) {
+                continue;
+            }
+            const char* arch = feature->value;
+            while (*arch != '\0') {
+                char* end        = nullptr;
+                const long value = std::strtol(arch, &end, 10);
+                if (end == arch) {
+                    ++arch;
+                    continue;
+                }
+                if (value <= cc && value > compiled_arch) {
+                    compiled_arch = static_cast<int>(value);
+                }
+                arch = end;
+            }
+        }
+        return compiled_arch == 700 || compiled_arch >= 750;
+    }();
+    cache.emplace(dev, supported);
+    return supported;
+#else
+    (void)backend;
+    return false;
+#endif
+}
+
 ggml_backend_t sd_backend_cpu_init() {
     ggml_backend_load_all_once();
     return ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);

src/core/ggml_extend_backend.h

@@ -87,6 +87,7 @@ private:
 
 bool sd_backend_is(ggml_backend_t backend, const std::string& name);
 bool sd_backend_is_cpu(ggml_backend_t backend);
+bool sd_backend_supports_cuda_mma(ggml_backend_t backend);
 ggml_backend_t sd_backend_cpu_init();
 bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
 ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,

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 {
 
@@ -426,8 +426,8 @@ namespace sd::ggml_graph_cut {
         if (tensor == nullptr || tensor->name[0] == '\0') {
             return false;
         }
-        return starts_with(tensor->name, GGML_RUNNER_CUT_PREFIX) &&
-               ends_with(tensor->name, GGML_RUNNER_CUT_SUFFIX);
+        return std::strncmp(tensor->name, GGML_RUNNER_CUT_PREFIX, std::strlen(GGML_RUNNER_CUT_PREFIX)) == 0 &&
+               tensor->name[std::strlen(tensor->name) - 1] == GGML_RUNNER_CUT_SUFFIX[0];
     }
 
     std::string make_graph_cut_name(const std::string& group, const std::string& output) {
@@ -492,35 +492,88 @@ namespace sd::ggml_graph_cut {
         }
     }
 
-    std::vector<uint64_t> graph_layout(ggml_cgraph* graph, bool include_bindings) {
-        std::vector<const ggml_tensor*> tensors;
-        std::unordered_map<const ggml_tensor*, size_t> indices;
-        auto add = [&](const ggml_tensor* tensor) {
-            if (tensor != nullptr && indices.emplace(tensor, tensors.size() + 1).second) {
-                tensors.push_back(tensor);
-            }
+    struct GraphLayoutTensors {
+        struct Entry {
+            const ggml_tensor* tensor = nullptr;
+            size_t index              = 0;
         };
+
+        std::vector<ggml_tensor*> tensors;
+        std::vector<Entry> entries;
+
+        explicit GraphLayoutTensors(size_t graph_size) {
+            tensors.reserve(graph_size);
+            size_t capacity = 2;
+            while (capacity < 2 * graph_size) {
+                capacity *= 2;
+            }
+            entries.resize(capacity);
+        }
+
+        size_t find(const ggml_tensor* tensor) const {
+            size_t hash = reinterpret_cast<uintptr_t>(tensor) >> 4;
+            hash ^= hash >> 16;
+            const size_t mask = entries.size() - 1;
+            size_t slot       = hash & mask;
+            while (entries[slot].tensor != nullptr && entries[slot].tensor != tensor) {
+                slot = (slot + 1) & mask;
+            }
+            return slot;
+        }
+
+        void add(ggml_tensor* tensor) {
+            if (tensor == nullptr) {
+                return;
+            }
+            size_t slot = find(tensor);
+            if (entries[slot].tensor != nullptr) {
+                return;
+            }
+            if (2 * (tensors.size() + 1) > entries.size()) {
+                // Segment graphs can reference tensors outside their node and leaf arrays.
+                std::vector<Entry> next(2 * entries.size());
+                entries.swap(next);
+                for (size_t i = 0; i < tensors.size(); ++i) {
+                    entries[find(tensors[i])] = {tensors[i], i + 1};
+                }
+                slot = find(tensor);
+            }
+            entries[slot] = {tensor, tensors.size() + 1};
+            tensors.push_back(tensor);
+        }
+
+        size_t index(const ggml_tensor* tensor) const {
+            return tensor == nullptr ? 0 : entries[find(tensor)].index;
+        }
+    };
+
+    std::vector<uint64_t> graph_layout(ggml_cgraph* graph, bool include_bindings) {
+        const size_t graph_size = static_cast<size_t>(graph->n_leafs) + graph->n_nodes;
+        GraphLayoutTensors layout_tensors(graph_size);
+        const auto& tensors = layout_tensors.tensors;
         for (int i = 0; i < graph->n_leafs; ++i) {
-            add(graph->leafs[i]);
+            layout_tensors.add(graph->leafs[i]);
         }
         for (int i = 0; i < graph->n_nodes; ++i) {
-            add(graph->nodes[i]);
+            layout_tensors.add(graph->nodes[i]);
         }
         for (size_t i = 0; i < tensors.size(); ++i) {
-            add(tensors[i]->view_src);
+            layout_tensors.add(tensors[i]->view_src);
             for (auto source : tensors[i]->src) {
-                add(source);
+                layout_tensors.add(source);
             }
         }
         std::vector<uint64_t> signature;
-        signature.reserve(tensors.size() * 24);
+        const size_t tensor_fields = 5 + 2 * GGML_MAX_DIMS + GGML_MAX_SRC +
+                                     GGML_MAX_OP_PARAMS / sizeof(int32_t) + (include_bindings ? 2 : 0);
+        signature.reserve(2 + graph_size + tensors.size() * tensor_fields);
         signature.push_back(graph->n_nodes);
         signature.push_back(graph->n_leafs);
         for (int i = 0; i < graph->n_leafs; ++i) {
-            signature.push_back(indices.at(graph->leafs[i]));
+            signature.push_back(layout_tensors.index(graph->leafs[i]));
         }
         for (int i = 0; i < graph->n_nodes; ++i) {
-            signature.push_back(indices.at(graph->nodes[i]));
+            signature.push_back(layout_tensors.index(graph->nodes[i]));
         }
         for (auto tensor : tensors) {
             signature.push_back(tensor->op);
@@ -536,9 +589,9 @@ namespace sd::ggml_graph_cut {
                 signature.push_back(tensor->ne[d]);
                 signature.push_back(tensor->nb[d]);
             }
-            signature.push_back(tensor->view_src == nullptr ? 0 : indices.at(tensor->view_src));
+            signature.push_back(layout_tensors.index(tensor->view_src));
             for (auto source : tensor->src) {
-                signature.push_back(source == nullptr ? 0 : indices.at(source));
+                signature.push_back(layout_tensors.index(source));
             }
             if (!can_ignore_op_params(tensor->op)) {
                 for (int value : tensor->op_params) {
@@ -562,14 +615,18 @@ namespace sd::ggml_graph_cut {
                 return false;
             }
         }
-        std::vector<std::pair<int, std::string>> cut_markers;
-        for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) {
-            auto node = ggml_graph_node(gf, i);
+        size_t cut_index = 0;
+        for (int i = 0; i < gf->n_nodes; ++i) {
+            auto node = gf->nodes[i];
             if (is_graph_cut_tensor(node)) {
-                cut_markers.emplace_back(i, node->name);
+                if (cut_index >= plan.cut_markers.size() ||
+                    plan.cut_markers[cut_index].first != i || plan.cut_markers[cut_index].second != node->name) {
+                    return false;
+                }
+                ++cut_index;
             }
         }
-        return cut_markers == plan.cut_markers;
+        return cut_index == plan.cut_markers.size();
     }
 
     bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) {
@@ -948,11 +1005,11 @@ namespace sd::ggml_graph_cut {
         return plan;
     }
 
-    Plan resolve_plan(ggml_backend_t backend,
-                      ggml_cgraph* gf,
-                      PlanCache* cache,
-                      const std::unordered_set<const ggml_tensor*>& params_tensor_set,
-                      const char* log_desc) {
+    const Plan& resolve_plan(ggml_backend_t backend,
+                             ggml_cgraph* gf,
+                             PlanCache* cache,
+                             const std::unordered_set<const ggml_tensor*>& params_tensor_set,
+                             const char* log_desc) {
         GGML_ASSERT(backend != nullptr);
         GGML_ASSERT(gf != nullptr);
         GGML_ASSERT(cache != nullptr);

src/core/ggml_graph_cut.h

@@ -94,11 +94,12 @@ namespace sd::ggml_graph_cut {
                     ggml_cgraph* gf,
                     const std::unordered_set<const ggml_tensor*>& params_tensor_set,
                     const char* log_desc);
-    Plan resolve_plan(ggml_backend_t backend,
-                      ggml_cgraph* gf,
-                      PlanCache* cache,
-                      const std::unordered_set<const ggml_tensor*>& params_tensor_set,
-                      const char* log_desc);
+    // The returned reference is valid until its cache entry is evicted or the cache is destroyed.
+    const Plan& resolve_plan(ggml_backend_t backend,
+                             ggml_cgraph* gf,
+                             PlanCache* cache,
+                             const std::unordered_set<const ggml_tensor*>& params_tensor_set,
+                             const char* log_desc);
 
 }  // 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() {
@@ -342,21 +342,17 @@ void GGMLRunner::copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_c
     }
 }
 
-bool GGMLRunner::resolve_graph_cut_plan(ggml_cgraph* gf,
-                                        GraphCutPlan* plan_out) {
-    GGML_ASSERT(plan_out != nullptr);
+const GGMLRunner::GraphCutPlan& GGMLRunner::resolve_graph_cut_plan(ggml_cgraph* gf) {
     GGML_ASSERT(gf != nullptr);
-    *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
-                                                 gf,
-                                                 &graph_cut_plan_cache_,
-                                                 params_tensor_set_,
-                                                 get_desc().c_str());
-    return true;
+    return sd::ggml_graph_cut::resolve_plan(runtime_backend,
+                                            gf,
+                                            &graph_cut_plan_cache_,
+                                            params_tensor_set_,
+                                            get_desc().c_str());
 }
 
-bool GGMLRunner::resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
-                                                    GraphCutPlan* plan_out) {
-    return resolve_graph_cut_plan(gf, plan_out);
+const GGMLRunner::GraphCutPlan& GGMLRunner::resolve_graph_cut_layer_split_plan(ggml_cgraph* gf) {
+    return resolve_graph_cut_plan(gf);
 }
 
 bool GGMLRunner::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
@@ -369,10 +365,7 @@ bool GGMLRunner::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
         return false;
     }
 
-    GraphCutPlan plan;
-    if (!resolve_graph_cut_layer_split_plan(gf, &plan)) {
-        return false;
-    }
+    const auto& plan = resolve_graph_cut_layer_split_plan(gf);
     if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
         auto manager = residency_manager.lock();
         if (manager == nullptr) {
@@ -527,9 +520,11 @@ 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;
+    runner_ctx.conv3d_direct_enabled = conv3d_direct_enabled;
     runner_ctx.circular_x_enabled    = circular_x_enabled;
     runner_ctx.circular_y_enabled    = circular_y_enabled;
     runner_ctx.weight_adapter        = weight_adapter;
@@ -595,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;
@@ -618,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();
@@ -647,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;
@@ -654,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;
 }
@@ -771,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;
     }
@@ -819,24 +820,26 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
     if (!assign_graph_cut_layer_split_backends(graph)) {
         return std::nullopt;
     }
-    const auto params = collect_used_param_tensors(graph);
-    ggml_graph_cut::Plan plan;
-    if (!resolve_graph_cut_plan(graph, &plan)) {
-        return std::nullopt;
-    }
-    const auto full_measurement = measure(graph, plan.compute_buffer_size);
+    const auto params           = collect_used_param_tensors(graph);
+    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();
     const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
                            manager != nullptr && manager->segmented_compute_enabled() &&
-                           plan.valid && plan.has_cuts && plan.segments.size() > 1 &&
+                           cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1 &&
                            !fits(memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)), params);
+    ggml_graph_cut::Plan monolithic_plan;
     if (!segmented) {
-        ggml_graph_cut::Segment segment;
+        monolithic_plan.segments.emplace_back();
+        auto& segment               = monolithic_plan.segments.back();
         segment.group_name          = "graph";
-        segment.compute_buffer_size = plan.compute_buffer_size;
+        segment.compute_buffer_size = cached_plan.compute_buffer_size;
+        segment.internal_node_indices.reserve(ggml_graph_n_nodes(graph));
+        segment.input_refs.reserve(ggml_graph_cut::leaf_count(graph));
         for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
             segment.internal_node_indices.push_back(i);
         }
@@ -849,8 +852,8 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
                                    : ggml_graph_cut::Segment::INPUT_EXTERNAL;
             segment.input_refs.push_back(input);
         }
-        plan.segments = {std::move(segment)};
     }
+    const auto& plan            = segmented ? cached_plan : monolithic_plan;
     const bool segments_changed = plan.segments.size() != logged_segment_count_;
     if (segments_changed && (segmented || logged_segment_count_ > 1)) {
         LOG_VERBOSE("%s using %zu segment%s", get_desc().c_str(),
@@ -892,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);
@@ -902,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);
@@ -909,11 +915,18 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
         auto ensure_capacity         = [&]() {
             sync_runtime_residency();
             auto requests = memory_requests(measurement.buffers, new_cache_bytes);
-            if (!fits(requests, weights.params(index)) && workspace_.release_excess(measurement)) {
+            if (fits(requests, weights.params(index))) {
+                return true;
+            }
+            if (workspace_.release_excess(measurement)) {
                 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");
@@ -922,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) {
@@ -965,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,9 +68,11 @@ 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;
+    bool conv3d_direct_enabled                                       = false;
     bool circular_x_enabled                                          = false;
     bool circular_y_enabled                                          = false;
     ggml_tensor* ip_context                                          = nullptr;
@@ -128,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,
@@ -175,9 +178,11 @@ 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;
+    bool conv3d_direct_enabled = false;
     bool circular_x_enabled    = false;
     bool circular_y_enabled    = false;
 
@@ -263,11 +268,9 @@ protected:
 
     void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true);
 
-    bool resolve_graph_cut_plan(ggml_cgraph* gf,
-                                GraphCutPlan* plan_out);
+    const GraphCutPlan& resolve_graph_cut_plan(ggml_cgraph* gf);
 
-    bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
-                                            GraphCutPlan* plan_out);
+    const GraphCutPlan& resolve_graph_cut_layer_split_plan(ggml_cgraph* gf);
 
     bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf);
 
@@ -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;
@@ -346,6 +359,10 @@ public:
         conv2d_direct_enabled = enabled;
     }
 
+    void set_conv3d_direct_enabled(bool enabled) {
+        conv3d_direct_enabled = enabled;
+    }
+
     void set_circular_axes(bool circular_x, bool circular_y) {
         circular_x_enabled = circular_x;
         circular_y_enabled = circular_y;

src/core/parallel.cpp

@@ -0,0 +1,143 @@
+#include "core/parallel.h"
+
+#include <algorithm>
+#include <condition_variable>
+#include <exception>
+#include <mutex>
+#include <thread>
+#include <vector>
+
+namespace sd {
+
+    namespace parallel_detail {
+        thread_local ParallelExecutor* executor = nullptr;
+        thread_local bool active                = false;
+    }
+
+    struct ParallelExecutor::Impl {
+        struct Worker {
+            std::condition_variable wake;
+            std::thread thread;
+            bool ready = false;
+        };
+
+        ParallelExecutor* owner;
+        std::mutex invocation_mutex;
+        std::mutex mutex;
+        std::condition_variable finished;
+        std::vector<std::unique_ptr<Worker>> workers;
+        bool stopping                                     = false;
+        int pending                                       = 0;
+        int participants                                  = 1;
+        int64_t begin                                     = 0;
+        int64_t count                                     = 0;
+        const std::function<void(int64_t, int64_t)>* task = nullptr;
+        std::exception_ptr error;
+
+        explicit Impl(ParallelExecutor* owner)
+            : owner(owner) {}
+
+        ~Impl() {
+            {
+                std::lock_guard<std::mutex> lock(mutex);
+                stopping = true;
+            }
+            for (auto& worker : workers) {
+                worker->wake.notify_one();
+            }
+            for (auto& worker : workers) {
+                worker->thread.join();
+            }
+        }
+
+        void execute(int index) {
+            ParallelScope scope(owner);
+            parallel_detail::Region region;
+            const int64_t size  = count / participants;
+            const int64_t extra = count % participants;
+            const int64_t first = begin + index * size + std::min<int64_t>(index, extra);
+            const int64_t last  = first + size + (index < extra ? 1 : 0);
+            try {
+                (*task)(first, last);
+            } catch (...) {
+                std::lock_guard<std::mutex> lock(mutex);
+                if (!error) {
+                    error = std::current_exception();
+                }
+            }
+        }
+
+        void worker_loop(Worker* worker, int index) {
+            std::unique_lock<std::mutex> lock(mutex);
+            for (;;) {
+                worker->wake.wait(lock, [&] { return stopping || worker->ready; });
+                if (stopping) {
+                    return;
+                }
+                worker->ready = false;
+                lock.unlock();
+                execute(index);
+                lock.lock();
+                if (--pending == 0) {
+                    finished.notify_one();
+                }
+            }
+        }
+
+        void run(int64_t first, int64_t last, int n_tasks, const std::function<void(int64_t, int64_t)>& callback) {
+            std::lock_guard<std::mutex> invocation_lock(invocation_mutex);
+            std::unique_lock<std::mutex> lock(mutex);
+            while (static_cast<int>(workers.size()) < n_tasks - 1) {
+                workers.push_back(std::make_unique<Worker>());
+                auto* worker    = workers.back().get();
+                const int index = static_cast<int>(workers.size());
+                try {
+                    worker->thread = std::thread([this, worker, index] { worker_loop(worker, index); });
+                } catch (...) {
+                    workers.pop_back();
+                    throw;
+                }
+            }
+            begin        = first;
+            count        = last - first;
+            participants = n_tasks;
+            pending      = n_tasks - 1;
+            task         = &callback;
+            error        = nullptr;
+            for (int i = 0; i < pending; ++i) {
+                workers[i]->ready = true;
+                workers[i]->wake.notify_one();
+            }
+            lock.unlock();
+            execute(0);
+            lock.lock();
+            finished.wait(lock, [&] { return pending == 0; });
+            task = nullptr;
+            if (error) {
+                std::rethrow_exception(error);
+            }
+        }
+    };
+
+    ParallelExecutor::ParallelExecutor(int n_threads)
+        : n_threads_(std::max(1, n_threads)), impl_(std::make_unique<Impl>(this)) {}
+
+    ParallelExecutor::~ParallelExecutor() = default;
+
+    void ParallelExecutor::run(int64_t begin, int64_t end, int64_t grain_size, const std::function<void(int64_t, int64_t)>& task) {
+        if (begin < 0 || grain_size <= 0) {
+            throw std::invalid_argument("parallel_for requires begin >= 0 and grain_size > 0");
+        }
+        if (end <= begin) {
+            return;
+        }
+        const int n_tasks = static_cast<int>(std::min<int64_t>(n_threads_, (end - begin) / grain_size));
+        if (parallel_detail::active || n_tasks <= 1) {
+            parallel_detail::Region region;
+            task(begin, end);
+            return;
+        }
+        impl_->run(begin, end, n_tasks, task);
+    }
+
+}

src/core/parallel.h

@@ -0,0 +1,77 @@
+#ifndef __SD_CORE_PARALLEL_H__
+#define __SD_CORE_PARALLEL_H__
+
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <stdexcept>
+#include <utility>
+
+namespace sd {
+
+    class ParallelExecutor {
+        struct Impl;
+        int n_threads_;
+        std::unique_ptr<Impl> impl_;
+
+    public:
+        explicit ParallelExecutor(int n_threads);
+        ~ParallelExecutor();
+        ParallelExecutor(const ParallelExecutor&)            = delete;
+        ParallelExecutor& operator=(const ParallelExecutor&) = delete;
+
+        int num_threads() const { return n_threads_; }
+        void run(int64_t begin, int64_t end, int64_t grain_size, const std::function<void(int64_t, int64_t)>& task);
+    };
+
+    namespace parallel_detail {
+        extern thread_local ParallelExecutor* executor;
+        extern thread_local bool active;
+
+        class Region {
+            bool previous_;
+
+        public:
+            Region()
+                : previous_(active) { active = true; }
+            ~Region() { active = previous_; }
+            Region(const Region&)            = delete;
+            Region& operator=(const Region&) = delete;
+        };
+    }
+
+    class ParallelScope {
+        ParallelExecutor* previous_;
+
+    public:
+        explicit ParallelScope(ParallelExecutor* executor)
+            : previous_(parallel_detail::executor) {
+            parallel_detail::executor = executor;
+        }
+        ~ParallelScope() { parallel_detail::executor = previous_; }
+        ParallelScope(const ParallelScope&)            = delete;
+        ParallelScope& operator=(const ParallelScope&) = delete;
+    };
+
+    // Ranges are non-negative. The callback may run concurrently and must own its writes.
+    template <typename F>
+    inline void parallel_for(int64_t begin, int64_t end, int64_t grain_size, F&& task) {
+        if (begin < 0 || grain_size <= 0) {
+            throw std::invalid_argument("parallel_for requires begin >= 0 and grain_size > 0");
+        }
+        if (end <= begin) {
+            return;
+        }
+        auto* executor = parallel_detail::executor;
+        if (parallel_detail::active || executor == nullptr || executor->num_threads() <= 1 ||
+            (end - begin) / grain_size < 2) {
+            parallel_detail::Region region;
+            task(begin, end);
+            return;
+        }
+        executor->run(begin, end, grain_size, std::forward<F>(task));
+    }
+
+}
+
+#endif  // __SD_CORE_PARALLEL_H__

src/core/regex.cpp

@@ -0,0 +1,171 @@
+#include "regex.h"
+
+#include <cstdint>
+#include <limits>
+#include <mutex>
+
+#define ONIG_ESCAPE_UCHAR_COLLISION
+#define ONIG_ESCAPE_REGEX_T_COLLISION
+#include <oniguruma.h>
+
+namespace sd {
+
+    struct Regex::Impl {
+        OnigRegex regex = nullptr;
+
+        ~Impl() {
+            if (regex) {
+                onig_free(regex);
+            }
+        }
+    };
+
+    struct RegexRegionDeleter {
+        void operator()(OnigRegion* region) const {
+            onig_region_free(region, 1);
+        }
+    };
+
+    static bool regex_error(std::string* error, const std::string& message) {
+        if (error) {
+            *error = message;
+        }
+        return false;
+    }
+
+    static bool regex_onig_error(std::string* error, int code, OnigErrorInfo* info = nullptr) {
+        OnigUChar buffer[ONIG_MAX_ERROR_MESSAGE_LEN];
+        onig_error_code_to_str(buffer, code, info);
+        return regex_error(error, reinterpret_cast<const char*>(buffer));
+    }
+
+    static int regex_initialize() {
+        static std::once_flag once;
+        static int result = ONIG_NORMAL;
+        std::call_once(once, [] {
+            OnigEncoding encodings[] = {ONIG_ENCODING_UTF8};
+            result                   = onig_initialize(encodings, 1);
+        });
+        // onig_end() would invalidate expressions held by other Regex instances.
+        return result;
+    }
+
+    // Rust str excludes overlong encodings, surrogates and extended UTF-8 accepted by Oniguruma.
+    static bool regex_valid_utf8(const std::string& text) {
+        size_t position = 0;
+        while (position < text.size()) {
+            const auto lead = static_cast<unsigned char>(text[position++]);
+            if (lead < 0x80) {
+                continue;
+            }
+            int count = 0;
+            if (lead >= 0xC2 && lead <= 0xDF) {
+                count = 1;
+            } else if (lead >= 0xE0 && lead <= 0xEF) {
+                count = 2;
+            } else if (lead >= 0xF0 && lead <= 0xF4) {
+                count = 3;
+            }
+            if (count == 0 || text.size() - position < static_cast<size_t>(count)) {
+                return false;
+            }
+            uint32_t codepoint = lead & (0x7F >> count);
+            for (int i = 0; i < count; ++i) {
+                const auto byte = static_cast<unsigned char>(text[position++]);
+                if ((byte & 0xC0) != 0x80) {
+                    return false;
+                }
+                codepoint = (codepoint << 6) | (byte & 0x3F);
+            }
+            constexpr uint32_t minimum[] = {0, 0x80, 0x800, 0x10000};
+            if (codepoint < minimum[count] || codepoint > 0x10FFFF || (codepoint >= 0xD800 && codepoint <= 0xDFFF)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    Regex::Regex()                            = default;
+    Regex::~Regex()                           = default;
+    Regex::Regex(Regex&&) noexcept            = default;
+    Regex& Regex::operator=(Regex&&) noexcept = default;
+
+    bool Regex::compile(const std::string& pattern, std::string* error) {
+        if (error) {
+            error->clear();
+        }
+        const int initialized = regex_initialize();
+        if (initialized != ONIG_NORMAL) {
+            return regex_onig_error(error, initialized);
+        }
+        if (pattern.size() > static_cast<size_t>(std::numeric_limits<int>::max())) {
+            return regex_error(error, "regex pattern exceeds Oniguruma's offset range");
+        }
+        const auto* begin = reinterpret_cast<const OnigUChar*>(pattern.data());
+        const auto* end   = begin + pattern.size();
+        if (!regex_valid_utf8(pattern)) {
+            return regex_error(error, "regex pattern is not valid UTF-8");
+        }
+
+        auto next = std::make_unique<Impl>();
+        OnigErrorInfo info{};
+        static std::mutex compile_mutex;
+        std::lock_guard<std::mutex> lock(compile_mutex);
+        const int result = onig_new(&next->regex, begin, end, ONIG_OPTION_NONE,
+                                    ONIG_ENCODING_UTF8, ONIG_SYNTAX_ONIGURUMA, &info);
+        if (result != ONIG_NORMAL) {
+            return regex_onig_error(error, result, &info);
+        }
+        impl_ = std::move(next);
+        return true;
+    }
+
+    bool Regex::find_matches(const std::string& text, std::vector<Match>& matches, std::string* error) const {
+        matches.clear();
+        if (error) {
+            error->clear();
+        }
+        if (!impl_) {
+            return regex_error(error, "regex has not been compiled");
+        }
+        if (text.size() > static_cast<size_t>(std::numeric_limits<int>::max())) {
+            return regex_error(error, "regex input exceeds Oniguruma's offset range");
+        }
+        const auto* begin = reinterpret_cast<const OnigUChar*>(text.data());
+        const auto* end   = begin + text.size();
+        if (!regex_valid_utf8(text)) {
+            return regex_error(error, "regex input is not valid UTF-8");
+        }
+        std::unique_ptr<OnigRegion, RegexRegionDeleter> region(onig_region_new());
+        if (!region) {
+            return regex_error(error, "failed to allocate regex match region");
+        }
+
+        size_t position = 0;
+        while (position <= text.size()) {
+            const int result = onig_search(impl_->regex, begin, end, begin + position, end,
+                                           region.get(), ONIG_OPTION_NONE);
+            if (result == ONIG_MISMATCH) {
+                break;
+            }
+            if (result < 0) {
+                matches.clear();
+                return regex_onig_error(error, result);
+            }
+            const size_t match_begin = static_cast<size_t>(region->beg[0]);
+            const size_t match_end   = static_cast<size_t>(region->end[0]);
+            // Match rust-onig's find_iter: suppress an empty match at the previous match's end.
+            if (match_begin == match_end && !matches.empty() && matches.back().second == match_end) {
+                if (position == text.size()) {
+                    break;
+                }
+                position += static_cast<size_t>(ONIGENC_MBC_ENC_LEN(ONIG_ENCODING_UTF8, begin + position));
+                continue;
+            }
+            matches.emplace_back(match_begin, match_end);
+            position = match_end;
+        }
+        return true;
+    }
+
+}  // namespace sd

src/core/regex.h

@@ -0,0 +1,32 @@
+#ifndef __SD_CORE_REGEX_H__
+#define __SD_CORE_REGEX_H__
+
+#include <cstddef>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace sd {
+
+    class Regex {
+        struct Impl;
+        std::unique_ptr<Impl> impl_;
+
+    public:
+        using Match = std::pair<size_t, size_t>;
+
+        Regex();
+        ~Regex();
+        Regex(Regex&&) noexcept;
+        Regex& operator=(Regex&&) noexcept;
+
+        // Failed compilation leaves the previous expression intact.
+        bool compile(const std::string& pattern, std::string* error = nullptr);
+        // Matches are non-overlapping UTF-8 byte ranges; each call owns its search state.
+        bool find_matches(const std::string& text, std::vector<Match>& matches, std::string* error = nullptr) const;
+    };
+
+}  // namespace sd
+
+#endif  // __SD_CORE_REGEX_H__

src/core/rng.hpp

@@ -1,6 +1,8 @@
 #ifndef __SD_CORE_RNG_HPP__
 #define __SD_CORE_RNG_HPP__
 
+#include <cstdint>
+#include <memory>
 #include <random>
 #include <vector>
 
@@ -8,6 +10,7 @@ class RNG {
 public:
     virtual void manual_seed(uint64_t seed)      = 0;
     virtual std::vector<float> randn(uint32_t n) = 0;
+    virtual std::shared_ptr<RNG> clone() const   = 0;
 };
 
 class STDDefaultRNG : public RNG {
@@ -15,6 +18,10 @@ private:
     std::default_random_engine generator;
 
 public:
+    std::shared_ptr<RNG> clone() const override {
+        return std::make_shared<STDDefaultRNG>(*this);
+    }
+
     void manual_seed(uint64_t seed) override {
         generator.seed((unsigned int)seed);
     }

src/core/rng_mt19937.hpp

@@ -1,7 +1,9 @@
 #ifndef __SD_CORE_RNG_MT19937_HPP__
 #define __SD_CORE_RNG_MT19937_HPP__
 
+#include <array>
 #include <cmath>
+#include <limits>
 #include <vector>
 
 #include "core/rng.hpp"
@@ -123,6 +125,10 @@ class MT19937RNG : public RNG {
 public:
     MT19937RNG(uint64_t seed = 0) { manual_seed(seed); }
 
+    std::shared_ptr<RNG> clone() const override {
+        return std::make_shared<MT19937RNG>(*this);
+    }
+
     void manual_seed(uint64_t seed) override {
         s.seed_     = seed;
         s.seeded_   = true;

src/core/rng_philox.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_CORE_RNG_PHILOX_HPP__
 #define __SD_CORE_RNG_PHILOX_HPP__
 
+#include <array>
 #include <cmath>
 #include <vector>
 
@@ -14,67 +15,35 @@ private:
     uint32_t offset;
 
 private:
-    std::vector<uint32_t> philox_m = {0xD2511F53, 0xCD9E8D57};
-    std::vector<uint32_t> philox_w = {0x9E3779B9, 0xBB67AE85};
-    float two_pow32_inv            = 2.3283064e-10f;
-    float two_pow32_inv_2pi        = 2.3283064e-10f * 6.2831855f;
+    using Counter = std::array<std::vector<uint32_t>, 4>;
 
-    std::vector<uint32_t> uint32(uint64_t x) {
-        std::vector<uint32_t> result(2);
-        result[0] = static_cast<uint32_t>(x & 0xFFFFFFFF);
-        result[1] = static_cast<uint32_t>(x >> 32);
-        return result;
-    }
-
-    std::vector<std::vector<uint32_t>> uint32(const std::vector<uint64_t>& x) {
-        uint32_t N = (uint32_t)x.size();
-        std::vector<std::vector<uint32_t>> result(2, std::vector<uint32_t>(N));
-
-        for (uint32_t i = 0; i < N; ++i) {
-            result[0][i] = static_cast<uint32_t>(x[i] & 0xFFFFFFFF);
-            result[1][i] = static_cast<uint32_t>(x[i] >> 32);
-        }
-
-        return result;
-    }
+    static constexpr uint32_t philox_m[2] = {0xD2511F53, 0xCD9E8D57};
+    static constexpr uint32_t philox_w[2] = {0x9E3779B9, 0xBB67AE85};
+    float two_pow32_inv                   = 2.3283064e-10f;
+    float two_pow32_inv_2pi               = 2.3283064e-10f * 6.2831855f;
 
     //  A single round of the Philox 4x32 random number generator.
-    void philox4_round(std::vector<std::vector<uint32_t>>& counter,
-                       const std::vector<std::vector<uint32_t>>& key) {
+    void philox4_round(Counter& counter, uint32_t key0, uint32_t key1) {
         uint32_t N = (uint32_t)counter[0].size();
         for (uint32_t i = 0; i < N; i++) {
-            std::vector<uint32_t> v1 = uint32(static_cast<uint64_t>(counter[0][i]) * static_cast<uint64_t>(philox_m[0]));
-            std::vector<uint32_t> v2 = uint32(static_cast<uint64_t>(counter[2][i]) * static_cast<uint64_t>(philox_m[1]));
+            const uint64_t v1 = static_cast<uint64_t>(counter[0][i]) * static_cast<uint64_t>(philox_m[0]);
+            const uint64_t v2 = static_cast<uint64_t>(counter[2][i]) * static_cast<uint64_t>(philox_m[1]);
 
-            counter[0][i] = v2[1] ^ counter[1][i] ^ key[0][i];
-            counter[1][i] = v2[0];
-            counter[2][i] = v1[1] ^ counter[3][i] ^ key[1][i];
-            counter[3][i] = v1[0];
+            counter[0][i] = static_cast<uint32_t>(v2 >> 32) ^ counter[1][i] ^ key0;
+            counter[1][i] = static_cast<uint32_t>(v2);
+            counter[2][i] = static_cast<uint32_t>(v1 >> 32) ^ counter[3][i] ^ key1;
+            counter[3][i] = static_cast<uint32_t>(v1);
         }
     }
 
-    // Generates 32-bit random numbers using the Philox 4x32 random number generator.
-    // Parameters:
-    //     counter : A 4xN array of 32-bit integers representing the counter values (offset into generation).
-    //     key : A 2xN array of 32-bit integers representing the key values (seed).
-    //     rounds : The number of rounds to perform.
-    // Returns:
-    //     std::vector<std::vector<uint32_t>>: A 4xN array of 32-bit integers containing the generated random numbers.
-    std::vector<std::vector<uint32_t>> philox4_32(std::vector<std::vector<uint32_t>>& counter,
-                                                  std::vector<std::vector<uint32_t>>& key,
-                                                  int rounds = 10) {
-        uint32_t N = (uint32_t)counter[0].size();
+    void philox4_32(Counter& counter, uint32_t key0, uint32_t key1, int rounds = 10) {
         for (int i = 0; i < rounds - 1; ++i) {
-            philox4_round(counter, key);
-
-            for (uint32_t j = 0; j < N; ++j) {
-                key[0][j] += philox_w[0];
-                key[1][j] += philox_w[1];
-            }
+            philox4_round(counter, key0, key1);
+            key0 += philox_w[0];
+            key1 += philox_w[1];
         }
 
-        philox4_round(counter, key);
-        return counter;
+        philox4_round(counter, key0, key1);
     }
 
     float box_muller(float x, float y) {
@@ -93,33 +62,35 @@ public:
         this->offset = 0;
     }
 
+    std::shared_ptr<RNG> clone() const override {
+        return std::make_shared<PhiloxRNG>(*this);
+    }
+
     void manual_seed(uint64_t seed) override {
         this->seed   = seed;
         this->offset = 0;
     }
 
     std::vector<float> randn(uint32_t n) override {
-        std::vector<std::vector<uint32_t>> counter(4, std::vector<uint32_t>(n, 0));
-        for (uint32_t i = 0; i < n; i++) {
-            counter[0][i] = this->offset;
-        }
+        Counter counter;
+        counter[0].resize(n, this->offset);
+        counter[1].resize(n);
+        counter[2].resize(n);
+        counter[3].resize(n);
 
         for (uint32_t i = 0; i < n; i++) {
             counter[2][i] = i;
         }
         this->offset += 1;
 
-        std::vector<uint64_t> key(n, this->seed);
-        std::vector<std::vector<uint32_t>> key_uint32 = uint32(key);
+        philox4_32(counter, static_cast<uint32_t>(this->seed), static_cast<uint32_t>(this->seed >> 32));
 
-        std::vector<std::vector<uint32_t>> g = philox4_32(counter, key_uint32);
-
-        std::vector<float> result;
+        std::vector<float> result(n);
         for (uint32_t i = 0; i < n; ++i) {
-            result.push_back(box_muller((float)g[0][i], (float)g[1][i]));
+            result[i] = box_muller((float)counter[0][i], (float)counter[1][i]);
         }
         return result;
     }
 };
 
-#endif  // __SD_CORE_RNG_PHILOX_HPP__
+#endif  // __SD_CORE_RNG_PHILOX_HPP__

src/core/tensor.hpp

@@ -16,6 +16,7 @@
 #include <utility>
 #include <vector>
 
+#include "core/parallel.h"
 #include "core/rng.hpp"
 
 namespace sd {
@@ -59,6 +60,15 @@ namespace sd {
         return numel;
     }
 
+    template <typename F>
+    inline void tensor_for_each(int64_t count, F&& fn, int64_t grain_size = 65536) {
+        parallel_for(0, count, grain_size, [&](int64_t begin, int64_t end) {
+            for (int64_t i = begin; i < end; ++i) {
+                fn(i);
+            }
+        });
+    }
+
     template <typename T>
     class Tensor {
     public:
@@ -230,7 +240,10 @@ namespace sd {
         }
 
         void fill_(const T& value) {
-            std::fill(data_.begin(), data_.end(), value);
+            const T fill_value = value;
+            parallel_for(0, numel(), 65536, [&](int64_t begin, int64_t end) {
+                std::fill_n(data_.data() + begin, end - begin, fill_value);
+            });
         }
 
         Tensor& masked_fill_(const Tensor<uint8_t>& mask, const T& value);
@@ -390,7 +403,7 @@ namespace sd {
                                               tensor_shape_to_string(lhs) + ", rhs_shape=" +
                                               tensor_shape_to_string(rhs));
             }
-            shape[i] = std::max(lhs_dim, rhs_dim);
+            shape[i] = lhs_dim == 1 ? rhs_dim : lhs_dim;
         }
         return shape;
     }
@@ -425,39 +438,55 @@ namespace sd {
                                                  const std::vector<int64_t>& rhs_shape_raw,
                                                  const std::vector<int64_t>& rhs_strides_raw,
                                                  F&& fn) {
-        const size_t ndim                = out_shape.size();
-        std::vector<int64_t> out_strides = tensor_compute_strides(out_shape);
-        std::vector<int64_t> lhs_shape(ndim, 1);
-        std::vector<int64_t> lhs_strides(ndim, 0);
-        std::vector<int64_t> rhs_shape(ndim, 1);
-        std::vector<int64_t> rhs_strides(ndim, 0);
-
-        for (size_t i = 0; i < lhs_shape_raw.size(); ++i) {
-            lhs_shape[i]   = lhs_shape_raw[i];
-            lhs_strides[i] = lhs_strides_raw[i];
-        }
-        for (size_t i = 0; i < rhs_shape_raw.size(); ++i) {
-            rhs_shape[i]   = rhs_shape_raw[i];
-            rhs_strides[i] = rhs_strides_raw[i];
-        }
-
-        const int64_t numel = tensor_numel(out_shape);
-        for (int64_t flat = 0; flat < numel; ++flat) {
-            int64_t remaining  = flat;
-            int64_t lhs_offset = 0;
-            int64_t rhs_offset = 0;
-            for (size_t i = ndim; i-- > 0;) {
-                int64_t coord = remaining / out_strides[i];
-                remaining %= out_strides[i];
-                if (lhs_shape[i] != 1) {
-                    lhs_offset += coord * lhs_strides[i];
+        const int64_t numel    = tensor_numel(out_shape);
+        const size_t ndim      = out_shape.size();
+        auto broadcast_strides = [&](const std::vector<int64_t>& shape,
+                                     const std::vector<int64_t>& strides) {
+            if ((numel != 0 && tensor_numel(shape) == 0) || strides.size() != shape.size()) {
+                tensor_throw_invalid_argument("Tensor broadcast requires non-empty inputs and matching strides");
+            }
+            std::vector<int64_t> result(ndim, 0);
+            for (size_t i = 0; i < std::max(ndim, shape.size()); ++i) {
+                const int64_t input_dim  = i < shape.size() ? shape[i] : 1;
+                const int64_t output_dim = i < ndim ? out_shape[i] : 1;
+                if (input_dim != 1 && input_dim != output_dim) {
+                    tensor_throw_invalid_argument("Tensor broadcast cannot expand the destination: input_shape=" +
+                                                  tensor_shape_to_string(shape) + ", output_shape=" +
+                                                  tensor_shape_to_string(out_shape));
                 }
-                if (rhs_shape[i] != 1) {
-                    rhs_offset += coord * rhs_strides[i];
+                if (i < ndim && input_dim != 1) {
+                    result[i] = strides[i];
                 }
             }
-            fn(flat, lhs_offset, rhs_offset);
+            return result;
+        };
+        const auto lhs_strides = broadcast_strides(lhs_shape_raw, lhs_strides_raw);
+        const auto rhs_strides = broadcast_strides(rhs_shape_raw, rhs_strides_raw);
+        if (numel == 0) {
+            return;
         }
+        parallel_for(0, numel, 16384, [&](int64_t begin, int64_t end) {
+            auto coord         = tensor_unravel_index(begin, out_shape);
+            int64_t lhs_offset = 0;
+            int64_t rhs_offset = 0;
+            for (size_t i = 0; i < ndim; ++i) {
+                lhs_offset += coord[i] * lhs_strides[i];
+                rhs_offset += coord[i] * rhs_strides[i];
+            }
+            for (int64_t flat = begin; flat < end; ++flat) {
+                fn(flat, lhs_offset, rhs_offset);
+                for (size_t i = 0; i < ndim; ++i) {
+                    lhs_offset += lhs_strides[i];
+                    rhs_offset += rhs_strides[i];
+                    if (++coord[i] < out_shape[i]) {
+                        break;
+                    }
+                    coord[i] = 0;
+                    lhs_offset -= out_shape[i] * lhs_strides[i];
+                    rhs_offset -= out_shape[i] * rhs_strides[i];
+                }
+            }
+        });
     }
 
     template <typename T>
@@ -469,6 +498,7 @@ namespace sd {
         const std::vector<int64_t> data_strides = tensor_compute_strides(shape_);
         const std::vector<int64_t> mask_strides = tensor_compute_strides(mask.shape());
         const uint8_t* mask_data                = mask.data();
+        const T fill_value                      = value;
         tensor_for_each_broadcast_offset(shape_,
                                          shape_,
                                          data_strides,
@@ -476,7 +506,7 @@ namespace sd {
                                          mask_strides,
                                          [&](int64_t, int64_t data_offset, int64_t mask_offset) {
                                              if (mask_data[mask_offset] != 0) {
-                                                 data_[static_cast<size_t>(data_offset)] = value;
+                                                 data_[static_cast<size_t>(data_offset)] = fill_value;
                                              }
                                          });
         return *this;
@@ -486,9 +516,9 @@ namespace sd {
     inline Tensor<uint8_t> operator<(const Tensor<T>& lhs, Scalar rhs) {
         Tensor<uint8_t> result(lhs.shape());
         const T value = static_cast<T>(rhs);
-        for (int64_t i = 0; i < lhs.numel(); ++i) {
-            result[i] = lhs[i] < value ? 1 : 0;
-        }
+        tensor_for_each(lhs.numel(), [&](int64_t i) {
+            result.data()[i] = lhs.data()[i] < value ? 1 : 0;
+        });
         return result;
     }
 
@@ -496,9 +526,9 @@ namespace sd {
     inline Tensor<uint8_t> operator<(Scalar lhs, const Tensor<T>& rhs) {
         Tensor<uint8_t> result(rhs.shape());
         const T value = static_cast<T>(lhs);
-        for (int64_t i = 0; i < rhs.numel(); ++i) {
-            result[i] = value < rhs[i] ? 1 : 0;
-        }
+        tensor_for_each(rhs.numel(), [&](int64_t i) {
+            result.data()[i] = value < rhs.data()[i] ? 1 : 0;
+        });
         return result;
     }
 
@@ -516,7 +546,7 @@ namespace sd {
                                          rhs.shape(),
                                          rhs_strides,
                                          [&](int64_t flat, int64_t lhs_offset, int64_t rhs_offset) {
-                                             result[flat] = lhs_data[lhs_offset] < rhs_data[rhs_offset] ? 1 : 0;
+                                             result.data()[flat] = lhs_data[lhs_offset] < rhs_data[rhs_offset] ? 1 : 0;
                                          });
         return result;
     }
@@ -524,9 +554,9 @@ namespace sd {
     template <typename T>
     inline Tensor<T>& operator+=(Tensor<T>& lhs, const Tensor<T>& rhs) {
         if (lhs.shape() == rhs.shape()) {
-            for (int64_t i = 0; i < lhs.numel(); ++i) {
-                lhs[i] += rhs[i];
-            }
+            tensor_for_each(lhs.numel(), [&](int64_t i) {
+                lhs.data()[i] += rhs.data()[i];
+            });
             return lhs;
         }
         tensor_broadcast_shape(lhs.shape(), rhs.shape());
@@ -539,7 +569,7 @@ namespace sd {
                                          rhs.shape(),
                                          rhs_strides,
                                          [&](int64_t, int64_t lhs_offset, int64_t rhs_offset) {
-                                             lhs[static_cast<int64_t>(lhs_offset)] += rhs_data[rhs_offset];
+                                             lhs.data()[lhs_offset] += rhs_data[rhs_offset];
                                          });
         return lhs;
     }
@@ -547,18 +577,18 @@ namespace sd {
     template <typename T, typename Scalar, typename = std::enable_if_t<std::is_arithmetic<Scalar>::value>>
     inline Tensor<T>& operator+=(Tensor<T>& lhs, Scalar rhs) {
         const T value = static_cast<T>(rhs);
-        for (int64_t i = 0; i < lhs.numel(); ++i) {
-            lhs[i] += value;
-        }
+        tensor_for_each(lhs.numel(), [&](int64_t i) {
+            lhs.data()[i] += value;
+        });
         return lhs;
     }
 
     template <typename T>
     inline Tensor<T>& operator-=(Tensor<T>& lhs, const Tensor<T>& rhs) {
         if (lhs.shape() == rhs.shape()) {
-            for (int64_t i = 0; i < lhs.numel(); ++i) {
-                lhs[i] -= rhs[i];
-            }
+            tensor_for_each(lhs.numel(), [&](int64_t i) {
+                lhs.data()[i] -= rhs.data()[i];
+            });
             return lhs;
         }
         tensor_broadcast_shape(lhs.shape(), rhs.shape());
@@ -571,7 +601,7 @@ namespace sd {
                                          rhs.shape(),
                                          rhs_strides,
                                          [&](int64_t, int64_t lhs_offset, int64_t rhs_offset) {
-                                             lhs[static_cast<int64_t>(lhs_offset)] -= rhs_data[rhs_offset];
+                                             lhs.data()[lhs_offset] -= rhs_data[rhs_offset];
                                          });
         return lhs;
     }
@@ -579,18 +609,18 @@ namespace sd {
     template <typename T, typename Scalar, typename = std::enable_if_t<std::is_arithmetic<Scalar>::value>>
     inline Tensor<T>& operator-=(Tensor<T>& lhs, Scalar rhs) {
         const T value = static_cast<T>(rhs);
-        for (int64_t i = 0; i < lhs.numel(); ++i) {
-            lhs[i] -= value;
-        }
+        tensor_for_each(lhs.numel(), [&](int64_t i) {
+            lhs.data()[i] -= value;
+        });
         return lhs;
     }
 
     template <typename T>
     inline Tensor<T>& operator*=(Tensor<T>& lhs, const Tensor<T>& rhs) {
         if (lhs.shape() == rhs.shape()) {
-            for (int64_t i = 0; i < lhs.numel(); ++i) {
-                lhs[i] *= rhs[i];
-            }
+            tensor_for_each(lhs.numel(), [&](int64_t i) {
+                lhs.data()[i] *= rhs.data()[i];
+            });
             return lhs;
         }
         tensor_broadcast_shape(lhs.shape(), rhs.shape());
@@ -603,7 +633,7 @@ namespace sd {
                                          rhs.shape(),
                                          rhs_strides,
                                          [&](int64_t, int64_t lhs_offset, int64_t rhs_offset) {
-                                             lhs[static_cast<int64_t>(lhs_offset)] *= rhs_data[rhs_offset];
+                                             lhs.data()[lhs_offset] *= rhs_data[rhs_offset];
                                          });
         return lhs;
     }
@@ -611,18 +641,18 @@ namespace sd {
     template <typename T, typename Scalar, typename = std::enable_if_t<std::is_arithmetic<Scalar>::value>>
     inline Tensor<T>& operator*=(Tensor<T>& lhs, Scalar rhs) {
         const T value = static_cast<T>(rhs);
-        for (int64_t i = 0; i < lhs.numel(); ++i) {
-            lhs[i] *= value;
-        }
+        tensor_for_each(lhs.numel(), [&](int64_t i) {
+            lhs.data()[i] *= value;
+        });
         return lhs;
     }
 
     template <typename T>
     inline Tensor<T>& operator/=(Tensor<T>& lhs, const Tensor<T>& rhs) {
         if (lhs.shape() == rhs.shape()) {
-            for (int64_t i = 0; i < lhs.numel(); ++i) {
-                lhs[i] /= rhs[i];
-            }
+            tensor_for_each(lhs.numel(), [&](int64_t i) {
+                lhs.data()[i] /= rhs.data()[i];
+            });
             return lhs;
         }
         tensor_broadcast_shape(lhs.shape(), rhs.shape());
@@ -635,7 +665,7 @@ namespace sd {
                                          rhs.shape(),
                                          rhs_strides,
                                          [&](int64_t, int64_t lhs_offset, int64_t rhs_offset) {
-                                             lhs[static_cast<int64_t>(lhs_offset)] /= rhs_data[rhs_offset];
+                                             lhs.data()[lhs_offset] /= rhs_data[rhs_offset];
                                          });
         return lhs;
     }
@@ -643,9 +673,9 @@ namespace sd {
     template <typename T, typename Scalar, typename = std::enable_if_t<std::is_arithmetic<Scalar>::value>>
     inline Tensor<T>& operator/=(Tensor<T>& lhs, Scalar rhs) {
         const T value = static_cast<T>(rhs);
-        for (int64_t i = 0; i < lhs.numel(); ++i) {
-            lhs[i] /= value;
-        }
+        tensor_for_each(lhs.numel(), [&](int64_t i) {
+            lhs.data()[i] /= value;
+        });
         return lhs;
     }
 
@@ -664,7 +694,7 @@ namespace sd {
                                              rhs.shape(),
                                              rhs_strides,
                                              [&](int64_t flat, int64_t lhs_offset, int64_t rhs_offset) {
-                                                 result[flat] = lhs_data[lhs_offset] + rhs_data[rhs_offset];
+                                                 result.data()[flat] = lhs_data[lhs_offset] + rhs_data[rhs_offset];
                                              });
             return result;
         }
@@ -699,7 +729,7 @@ namespace sd {
                                              rhs.shape(),
                                              rhs_strides,
                                              [&](int64_t flat, int64_t lhs_offset, int64_t rhs_offset) {
-                                                 result[flat] = lhs_data[lhs_offset] - rhs_data[rhs_offset];
+                                                 result.data()[flat] = lhs_data[lhs_offset] - rhs_data[rhs_offset];
                                              });
             return result;
         }
@@ -717,9 +747,9 @@ namespace sd {
     inline Tensor<T> operator-(Scalar lhs, const Tensor<T>& rhs) {
         Tensor<T> result = rhs;
         const T value    = static_cast<T>(lhs);
-        for (int64_t i = 0; i < result.numel(); ++i) {
-            result[i] = value - result[i];
-        }
+        tensor_for_each(result.numel(), [&](int64_t i) {
+            result.data()[i] = value - result.data()[i];
+        });
         return result;
     }
 
@@ -738,7 +768,7 @@ namespace sd {
                                              rhs.shape(),
                                              rhs_strides,
                                              [&](int64_t flat, int64_t lhs_offset, int64_t rhs_offset) {
-                                                 result[flat] = lhs_data[lhs_offset] * rhs_data[rhs_offset];
+                                                 result.data()[flat] = lhs_data[lhs_offset] * rhs_data[rhs_offset];
                                              });
             return result;
         }
@@ -773,7 +803,7 @@ namespace sd {
                                              rhs.shape(),
                                              rhs_strides,
                                              [&](int64_t flat, int64_t lhs_offset, int64_t rhs_offset) {
-                                                 result[flat] = lhs_data[lhs_offset] / rhs_data[rhs_offset];
+                                                 result.data()[flat] = lhs_data[lhs_offset] / rhs_data[rhs_offset];
                                              });
             return result;
         }
@@ -791,18 +821,18 @@ namespace sd {
     inline Tensor<T> operator/(Scalar lhs, const Tensor<T>& rhs) {
         Tensor<T> result = rhs;
         const T value    = static_cast<T>(lhs);
-        for (int64_t i = 0; i < result.numel(); ++i) {
-            result[i] = value / result[i];
-        }
+        tensor_for_each(result.numel(), [&](int64_t i) {
+            result.data()[i] = value / result.data()[i];
+        });
         return result;
     }
 
     template <typename T>
     inline Tensor<T> operator-(const Tensor<T>& tensor) {
         Tensor<T> result = tensor;
-        for (int64_t i = 0; i < result.numel(); ++i) {
-            result[i] = -result[i];
-        }
+        tensor_for_each(result.numel(), [&](int64_t i) {
+            result.data()[i] = -result.data()[i];
+        });
         return result;
     }
 
@@ -1067,9 +1097,11 @@ namespace sd {
         template <typename T>
         inline Tensor<T> exp(const Tensor<T>& input) {
             Tensor<T> output(input.shape());
-            for (int64_t i = 0; i < input.numel(); ++i) {
-                output[i] = static_cast<T>(std::exp(static_cast<double>(input[i])));
-            }
+            tensor_for_each(
+                input.numel(), [&](int64_t i) {
+                    output.data()[i] = static_cast<T>(std::exp(static_cast<double>(input.data()[i])));
+                },
+                4096);
             return output;
         }
 
@@ -1079,18 +1111,18 @@ namespace sd {
                 tensor_throw_invalid_argument("Tensor clamp requires min_value <= max_value");
             }
             Tensor<T> output(input.shape());
-            for (int64_t i = 0; i < input.numel(); ++i) {
-                output[i] = std::clamp(input[i], min_value, max_value);
-            }
+            tensor_for_each(input.numel(), [&](int64_t i) {
+                output.data()[i] = std::clamp(input.data()[i], min_value, max_value);
+            });
             return output;
         }
 
         template <typename T>
         inline Tensor<T> round(const Tensor<T>& input) {
             Tensor<T> output(input.shape());
-            for (int64_t i = 0; i < input.numel(); ++i) {
-                output[i] = static_cast<T>(std::round(static_cast<double>(input[i])));
-            }
+            tensor_for_each(input.numel(), [&](int64_t i) {
+                output.data()[i] = static_cast<T>(std::round(static_cast<double>(input.data()[i])));
+            });
             return output;
         }
 

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) {
@@ -755,6 +799,13 @@ sd::Tensor<float> clip_preprocess(const sd::Tensor<float>& image, int target_wid
     int64_t resized_width  = static_cast<int64_t>(scale * static_cast<float>(image.shape()[0]));
     int64_t resized_height = static_cast<int64_t>(scale * static_cast<float>(image.shape()[1]));
 
+    // The resized image must cover the crop window. Floating-point rounding can
+    // leave a side one pixel short of the crop target (e.g. 730 -> 735.999...
+    // -> 735 after truncation), so clamp to keep the center crop in bounds.
+    // Truncation is otherwise preserved to avoid changing existing results.
+    resized_width  = std::max<int64_t>(resized_width, target_width);
+    resized_height = std::max<int64_t>(resized_height, target_height);
+
     sd::Tensor<float> resized = sd::ops::interpolate(
         image,
         {resized_width, resized_height, image.shape()[2], image.shape()[3]});

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/extensions/photomaker_extension.cpp

@@ -18,12 +18,15 @@ tokenize_photomaker_trigger(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
     auto tokens_and_weights           = clip_conditioner.tokenize(text);
     std::vector<int> source_tokens    = std::move(tokens_and_weights.first);
     std::vector<float> source_weights = std::move(tokens_and_weights.second);
+    if (source_tokens.empty()) {
+        return {};
+    }
 
-    if (!source_tokens.empty() && source_tokens.front() == clip_conditioner.tokenizer.BOS_TOKEN_ID) {
+    if (!source_tokens.empty() && source_tokens.front() == clip_conditioner.tokenizer->BOS_TOKEN_ID) {
         source_tokens.erase(source_tokens.begin());
         source_weights.erase(source_weights.begin());
     }
-    if (!source_tokens.empty() && source_tokens.back() == clip_conditioner.tokenizer.EOS_TOKEN_ID) {
+    if (!source_tokens.empty() && source_tokens.back() == clip_conditioner.tokenizer->EOS_TOKEN_ID) {
         source_tokens.pop_back();
         source_weights.pop_back();
     }
@@ -49,12 +52,12 @@ tokenize_photomaker_trigger(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
         weights.push_back(source_weights[i]);
     }
 
-    clip_conditioner.tokenizer.pad_tokens(tokens,
-                                          &weights,
-                                          nullptr,
-                                          clip_conditioner.text_model->model.n_token,
-                                          clip_conditioner.text_model->model.n_token,
-                                          true);
+    clip_conditioner.tokenizer->pad_tokens(tokens,
+                                           &weights,
+                                           nullptr,
+                                           clip_conditioner.text_model->model.n_token,
+                                           clip_conditioner.text_model->model.n_token,
+                                           true);
     std::vector<bool> class_token_mask;
     for (int i = 0; i < tokens.size(); i++) {
         class_token_mask.push_back(class_idx >= 0 && class_idx + 1 <= i && i < class_idx + 1 + trigger_token_count);
@@ -69,8 +72,14 @@ get_photomaker_condition_with_trigger(FrozenCLIPEmbedderWithCustomWords& clip_co
                                       const ConditionerParams& conditioner_params,
                                       const std::string& trigger_word,
                                       int trigger_token_count) {
-    auto image_tokens = clip_conditioner.convert_token_to_id(trigger_word);
-    GGML_ASSERT(image_tokens.size() == 1);
+    std::vector<int> image_tokens;
+    if (!clip_conditioner.convert_token_to_id(trigger_word, image_tokens)) {
+        return {};
+    }
+    if (image_tokens.size() != 1) {
+        LOG_ERROR("PhotoMaker trigger word must encode to one token");
+        return {};
+    }
     auto tokens_and_weights         = tokenize_photomaker_trigger(clip_conditioner,
                                                                   conditioner_params.text,
                                                                   trigger_token_count,
@@ -78,27 +87,43 @@ get_photomaker_condition_with_trigger(FrozenCLIPEmbedderWithCustomWords& clip_co
     std::vector<int>& tokens        = std::get<0>(tokens_and_weights);
     std::vector<float>& weights     = std::get<1>(tokens_and_weights);
     std::vector<bool>& trigger_mask = std::get<2>(tokens_and_weights);
-    auto cond                       = clip_conditioner.get_learned_condition_common(n_threads,
-                                                                                    tokens,
-                                                                                    weights,
-                                                                                    conditioner_params.clip_skip,
-                                                                                    conditioner_params.width,
-                                                                                    conditioner_params.height,
-                                                                                    conditioner_params.zero_out_masked);
+    if (tokens.empty()) {
+        return {};
+    }
+    auto cond = clip_conditioner.get_learned_condition_common(n_threads,
+                                                              tokens,
+                                                              weights,
+                                                              conditioner_params.clip_skip,
+                                                              conditioner_params.width,
+                                                              conditioner_params.height,
+                                                              conditioner_params.zero_out_masked);
     return std::make_tuple(std::move(cond), trigger_mask);
 }
 
-static std::string remove_photomaker_trigger_from_prompt(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
-                                                         const std::string& prompt,
-                                                         const std::string& trigger_word) {
-    auto image_tokens = clip_conditioner.convert_token_to_id(trigger_word);
-    GGML_ASSERT(image_tokens.size() == 1);
+static bool remove_photomaker_trigger_from_prompt(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
+                                                  const std::string& prompt,
+                                                  const std::string& trigger_word,
+                                                  std::string& result) {
+    std::vector<int> image_tokens;
+    if (!clip_conditioner.convert_token_to_id(trigger_word, image_tokens)) {
+        return false;
+    }
+    if (image_tokens.size() != 1) {
+        LOG_ERROR("PhotoMaker trigger word must encode to one token");
+        return false;
+    }
     auto tokens_and_weights  = clip_conditioner.tokenize(prompt);
     std::vector<int>& tokens = tokens_and_weights.first;
-    auto it                  = std::find(tokens.begin(), tokens.end(), image_tokens[0]);
-    GGML_ASSERT(it != tokens.end());
+    if (tokens.empty()) {
+        return false;
+    }
+    auto it = std::find(tokens.begin(), tokens.end(), image_tokens[0]);
+    if (it == tokens.end()) {
+        LOG_ERROR("PhotoMaker trigger word was not found in tokenized prompt");
+        return false;
+    }
     tokens.erase(it);
-    return clip_conditioner.decode(tokens);
+    return clip_conditioner.decode(tokens, result);
 }
 
 struct PhotoMakerExtension : public GenerationExtension {
@@ -223,6 +248,10 @@ struct PhotoMakerExtension : public GenerationExtension {
                                                                                   trigger_token_count);
         SDCondition prepared_id_condition = std::get<0>(cond_tup);
         auto class_tokens_mask            = std::get<1>(cond_tup);
+        if (prepared_id_condition.empty()) {
+            LOG_ERROR("failed to encode PhotoMaker prompt");
+            return false;
+        }
         if (std::find(class_tokens_mask.begin(), class_tokens_mask.end(), true) == class_tokens_mask.end()) {
             LOG_WARN("PhotoMaker trigger word '%s' was not found in prompt", trigger_word.c_str());
             LOG_WARN("Turn off PhotoMaker for this request");
@@ -263,11 +292,16 @@ struct PhotoMakerExtension : public GenerationExtension {
 
         prepared_id_condition.c_crossattn = std::move(res);
         int64_t t1                        = ggml_time_ms();
-        id_condition                      = std::move(prepared_id_condition);
-        start_merge_step                  = int(ctx.pm_params.style_strength / 100.f * ctx.total_steps);
-        ctx.condition_params.text         = remove_photomaker_trigger_from_prompt(*clip_conditioner,
-                                                                                  ctx.condition_params.text,
-                                                                                  trigger_word);
+        std::string prompt;
+        if (!remove_photomaker_trigger_from_prompt(*clip_conditioner,
+                                                   ctx.condition_params.text,
+                                                   trigger_word,
+                                                   prompt)) {
+            return false;
+        }
+        id_condition              = std::move(prepared_id_condition);
+        start_merge_step          = int(ctx.pm_params.style_strength / 100.f * ctx.total_steps);
+        ctx.condition_params.text = std::move(prompt);
         LOG_INFO("Photomaker ID Stacking, taking %" PRId64 " ms", t1 - t0);
         LOG_INFO("PHOTOMAKER: start_merge_step: %d", start_merge_step);
 

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/adapter/lora.hpp

@@ -67,7 +67,7 @@ struct LoraModel : public GGMLRunner {
         std::map<std::string, ggml_tensor*> scalars;
         std::set<std::string> scalar_names;
         for (const auto& [name, source] : sources) {
-            if (is_unused_tensor(name) || (filter && !filter(name)))
+            if (filter && !filter(name))
                 continue;
             const bool scalar  = source.nelements() == 1 && (ends_with(name, ".alpha") || ends_with(name, ".scale"));
             auto* tensor       = ggml_new_tensor(params_ctx, scalar ? GGML_TYPE_F32 : source.type, source.n_dims, source.ne);

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,
@@ -728,7 +732,7 @@ public:
                                 std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                 std::get<2>(padding), std::get<1>(padding), std::get<0>(padding),
                                 std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
-                                force_prec_f32);
+                                force_prec_f32, ctx->conv3d_direct_enabled);
     }
 };
 
@@ -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/hidream_o1.hpp

@@ -484,13 +484,19 @@ namespace HiDreamO1 {
     };
 
     struct HiDreamO1Conditioner : public Conditioner {
-        Qwen2Tokenizer tokenizer;
+        std::shared_ptr<Tokenizer> tokenizer;
         std::shared_ptr<HiDreamO1VisionRunner> vision_runner;
 
         HiDreamO1Conditioner(ggml_backend_t backend,
                              const String2TensorStorage& tensor_storage_map      = {},
-                             std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
-            : vision_runner(std::make_shared<HiDreamO1VisionRunner>(backend, tensor_storage_map, "model.visual", weight_manager)) {}
+                             std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                             const TokenizerConfig& tokenizers                   = {})
+            : vision_runner(std::make_shared<HiDreamO1VisionRunner>(backend, tensor_storage_map, "model.visual", weight_manager)) {
+            tokenizer = tokenizers.create(TokenizerConfig::MAIN, HiDreamO1Config::detect_from_weights(tensor_storage_map, "").llm.vocab_size, 151643);
+            if (!tokenizer) {
+                tokenizer = std::make_shared<Qwen2Tokenizer>();
+            }
+        }
 
         void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
             vision_runner->get_param_tensors(tensors);
@@ -538,7 +544,10 @@ namespace HiDreamO1 {
             if (ref_images.empty()) {
                 prompt += conditioner_params.text;
                 prompt += "<|im_end|>\n<|im_start|>assistant\n<|boi_token|><|tms_token|>";
-                auto input_ids = tokenizer.encode(prompt, nullptr);
+                std::vector<int> input_ids;
+                if (!tokenizer->encode(prompt, input_ids, nullptr)) {
+                    return {};
+                }
 
                 std::vector<int32_t> input_ids_pad = input_ids;
                 input_ids_pad.push_back(VISION_START_TOKEN_ID);
@@ -612,7 +621,11 @@ namespace HiDreamO1 {
 
                 auto patch_img = resized_ref * 2.0f - 1.0f;
                 result.c_ref_images.push_back(std::move(patch_img));
-                int64_t prompt_start = static_cast<int64_t>(tokenizer.encode(prompt + "<|vision_start|>", nullptr).size());
+                std::vector<int> prefix_tokens;
+                if (!tokenizer->encode(prompt + "<|vision_start|>", prefix_tokens, nullptr)) {
+                    return {};
+                }
+                int64_t prompt_start = static_cast<int64_t>(prefix_tokens.size());
                 prompt += "<|vision_start|>";
                 prompt += repeat_special_token("<|image_pad|>", image_tokens);
                 prompt += "<|vision_end|>";
@@ -623,7 +636,10 @@ namespace HiDreamO1 {
 
             prompt += conditioner_params.text;
             prompt += "<|im_end|>\n<|im_start|>assistant\n<|boi_token|><|tms_token|>";
-            auto input_ids = tokenizer.encode(prompt, nullptr);
+            std::vector<int> input_ids;
+            if (!tokenizer->encode(prompt, input_ids, nullptr)) {
+                return {};
+            }
 
             std::vector<int32_t> input_ids_pad = input_ids;
             input_ids_pad.push_back(VISION_START_TOKEN_ID);

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

@@ -15,6 +15,7 @@
 #include <regex>
 #include <set>
 #include <sstream>
+#include <stdexcept>
 #include <string>
 #include <utility>
 #include <vector>
@@ -31,9 +32,9 @@
 #include "model_manager.h"
 #include "tokenizers/bpe_tokenizer.h"
 #include "tokenizers/gemma_tokenizer.h"
-#include "tokenizers/gpt_oss_tokenizer.h"
 #include "tokenizers/mistral_tokenizer.h"
 #include "tokenizers/qwen2_tokenizer.h"
+#include "tokenizers/tokenizer_config.h"
 
 namespace LLM {
     constexpr int LLM_GRAPH_SIZE = 65536;
@@ -48,6 +49,7 @@ namespace LLM {
         GEMMA2_2B,
         GEMMA4_12B,
         GPT_OSS_20B,
+        LLADA2_MOE,
         ARCH_COUNT,
     };
 
@@ -61,6 +63,7 @@ namespace LLM {
         "gemma2_2b",
         "gemma4_12b",
         "gpt_oss_20b",
+        "llada2_moe",
     };
 
     enum class MLPActivation {
@@ -124,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;
@@ -211,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;
@@ -418,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;
@@ -604,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;
     }
@@ -1121,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;
 
@@ -1146,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) {
@@ -1160,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);
         }
 
@@ -1185,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]
@@ -1335,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);
@@ -1376,7 +1669,7 @@ namespace LLM {
                 x        = ggml_ext_cont(ctx->ggml_ctx, kqv);
                 x        = ggml_reshape_3d(ctx->ggml_ctx, x, head_dim * num_heads, n_token, N);
             } else {
-                x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, false);  // [N, n_token, hidden_size]
+                x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled);  // [N, n_token, hidden_size]
             }
 
             x = out_proj->forward(ctx, x);  // [N, n_token, hidden_size]
@@ -1431,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,
@@ -1484,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);
@@ -1649,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"]);
@@ -1989,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;
@@ -2041,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;
@@ -2114,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;
@@ -2351,7 +2673,7 @@ namespace LLM {
     };
 
     struct LLMEmbedder {
-        std::shared_ptr<BPETokenizer> tokenizer;
+        std::shared_ptr<Tokenizer> tokenizer;
         LLMRunner model;
 
         LLMEmbedder(LLMArch arch,
@@ -2359,14 +2681,29 @@ namespace LLM {
                     const String2TensorStorage& tensor_storage_map      = {},
                     const std::string prefix                            = "",
                     bool enable_vision                                  = false,
-                    std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+                    std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                    const TokenizerConfig& tokenizers                   = {})
             : model(arch, backend, tensor_storage_map, prefix, enable_vision, weight_manager) {
+            int pad_id = 151643;
             if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
-                tokenizer = std::make_shared<MistralTokenizer>();
+                pad_id = 11;
             } else if (arch == LLMArch::GPT_OSS_20B) {
-                tokenizer = std::make_shared<GPTOSSTokenizer>();
-            } else {
-                tokenizer = std::make_shared<Qwen2Tokenizer>();
+                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 || 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>();
+                } else {
+                    tokenizer = std::make_shared<Qwen2Tokenizer>();
+                }
             }
         }
 
@@ -2402,7 +2739,10 @@ namespace LLM {
             for (const auto& item : parsed_attention) {
                 const std::string& curr_text = item.first;
                 float curr_weight            = item.second;
-                std::vector<int> curr_tokens = tokenizer->tokenize(curr_text, nullptr);
+                std::vector<int> curr_tokens;
+                if (!tokenizer->tokenize(curr_text, curr_tokens, nullptr)) {
+                    return {};
+                }
                 tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 weights.insert(weights.end(), curr_tokens.size(), curr_weight);
             }

src/model/te/t5.hpp

@@ -567,7 +567,10 @@ struct T5Embedder {
         for (const auto& item : parsed_attention) {
             const std::string& curr_text = item.first;
             float curr_weight            = item.second;
-            std::vector<int> curr_tokens = tokenizer.encode(curr_text);
+            std::vector<int> curr_tokens;
+            if (!tokenizer.encode(curr_text, curr_tokens)) {
+                return {};
+            }
             tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
             weights.insert(weights.end(), curr_tokens.size(), curr_weight);
         }

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),
@@ -78,7 +85,8 @@ namespace WAN {
             return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
                                     std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                     0, 0, 0,
-                                    std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation));
+                                    std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
+                                    false, ctx->conv3d_direct_enabled);
         }
     };
 
@@ -139,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) {
@@ -153,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 {
@@ -468,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++;
             }
         }
@@ -531,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++;
             }
         }
@@ -1053,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) {
@@ -1270,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);
         }
@@ -1341,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,53 +35,66 @@
 
 /*================================================= Preprocess ==================================================*/
 
-const char* unused_tensors[] = {
-    "betas",
-    "alphas_cumprod_prev",
-    "sqrt_alphas_cumprod",
-    "sqrt_one_minus_alphas_cumprod",
-    "log_one_minus_alphas_cumprod",
-    "sqrt_recip_alphas_cumprod",
-    "sqrt_recipm1_alphas_cumprod",
-    "posterior_variance",
-    "posterior_log_variance_clipped",
-    "posterior_mean_coef1",
-    "posterior_mean_coef2",
-    "cond_stage_model.transformer.text_model.embeddings.position_ids",
-    "cond_stage_model.1.model.text_model.embeddings.position_ids",
-    "cond_stage_model.transformer.vision_model.embeddings.position_ids",
-    "cond_stage_model.model.logit_scale",
-    "conditioner.embedders.0.transformer.text_model.embeddings.position_ids",
-    "conditioner.embedders.0.model.logit_scale",
-    "conditioner.embedders.1.model.logit_scale",
-    "model.diffusion_model.time_embedding.cond_proj.weight",
-    "unet.time_embedding.cond_proj.weight",
-    "model_ema.decay",
-    "model_ema.num_updates",
-    "model_ema.diffusion_model",
-    "embedding_manager",
-    "denoiser.sigmas",
-    "text_encoders.t5xxl.transformer.encoder.embed_tokens.weight",  // only used during training
-    "ztsnr",                                                        // Found in some SDXL vpred models
-    "edm_vpred.sigma_min",                                          // Found in CosXL
-    // TODO: find another way to avoid the "unknown tensor" for these two
-    // "edm_vpred.sigma_max", // Used to detect CosXL
-    // "v_pred", // Used to detect SDXL vpred models
-    "text_encoders.llm.output.weight",
-    "text_encoders.llm.lm_head.",
-    "language_model.lm_head.",
-    "vision_model.",
-};
-
-bool is_unused_tensor(const std::string& name) {
-    for (size_t i = 0; i < sizeof(unused_tensors) / sizeof(const char*); i++) {
-        if (starts_with(name, unused_tensors[i])) {
-            return true;
-        }
+#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);
     }
-    return false;
+
+    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++) {
@@ -285,10 +298,6 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
     size_t file_index = add_file_path(file_path);
 
     for (auto& tensor_storage : tensor_storages) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
         if (!starts_with(tensor_storage.name, prefix)) {
             tensor_storage.name = prefix + tensor_storage.name;
         }
@@ -357,10 +366,6 @@ bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, cons
     size_t file_index = add_file_path(file_path);
 
     for (auto& tensor_storage : tensor_storages) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
         if (!starts_with(tensor_storage.name, prefix)) {
             tensor_storage.name = prefix + tensor_storage.name;
         }
@@ -478,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) {
@@ -506,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;
         }
@@ -675,10 +686,6 @@ SDVersion ModelLoader::get_sd_version() const {
 std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() const {
     std::map<ggml_type, uint32_t> wtype_stat;
     for (auto& [name, tensor_storage] : tensor_storage_map) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
         auto iter = wtype_stat.find(tensor_storage.type);
         if (iter != wtype_stat.end()) {
             iter->second++;
@@ -692,10 +699,6 @@ std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() const {
 std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() const {
     std::map<ggml_type, uint32_t> wtype_stat;
     for (auto& [name, tensor_storage] : tensor_storage_map) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
         if ((tensor_storage.name.find("text_encoders") == std::string::npos &&
              tensor_storage.name.find("cond_stage_model") == std::string::npos &&
              tensor_storage.name.find("te.text_model.") == std::string::npos &&
@@ -716,10 +719,6 @@ std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() const {
 std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() const {
     std::map<ggml_type, uint32_t> wtype_stat;
     for (auto& [name, tensor_storage] : tensor_storage_map) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
         if (tensor_storage.name.find("model.diffusion_model.") == std::string::npos && tensor_storage.name.find("unet.") == std::string::npos) {
             continue;
         }
@@ -737,10 +736,6 @@ std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() cons
 std::map<ggml_type, uint32_t> ModelLoader::get_vae_wtype_stat() const {
     std::map<ggml_type, uint32_t> wtype_stat;
     for (auto& [name, tensor_storage] : tensor_storage_map) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
         if (tensor_storage.name.find("vae.") == std::string::npos &&
             tensor_storage.name.find("first_stage_model") == std::string::npos) {
             continue;
@@ -824,9 +819,6 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
 
     std::vector<TensorStorage> processed_tensor_storages;
     for (const auto& [name, tensor_storage] : tensor_storage_map) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
         processed_tensor_storages.push_back(tensor_storage);
     }
 
@@ -932,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;
             }
@@ -1221,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) {
@@ -1537,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;
         }
@@ -1552,9 +1557,6 @@ int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type)
     int64_t mem_size = 0;
     std::vector<TensorStorage> processed_tensor_storages;
     for (auto [name, tensor_storage] : tensor_storage_map) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
         if (tensor_should_be_converted(tensor_storage, type)) {
             tensor_storage.type = type;
         }

src/model_loader.h

@@ -28,8 +28,6 @@ struct MmapTensorStore {
     std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
 };
 
-bool is_unused_tensor(const std::string& name);
-
 class ModelLoader {
 public:
     using FileId       = uint64_t;

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

@@ -274,6 +274,7 @@ bool ModelManager::register_param_tensors(ModelComponent component,
         new_states.push_back(std::move(state));
     }
 
+    resolved_tensor_states_.clear();
     for (auto& state : new_states) {
         TensorState* registered_state                      = state.get();
         tensor_states_by_tensor_[registered_state->tensor] = registered_state;
@@ -369,6 +370,7 @@ bool ModelManager::unregister_tensor_states(const std::unordered_set<TensorState
         }
     }
 
+    resolved_tensor_states_.clear();
     for (auto it = tensor_states_by_tensor_.begin(); it != tensor_states_by_tensor_.end();) {
         if (target_states.count(it->second) > 0) {
             it = tensor_states_by_tensor_.erase(it);
@@ -1199,22 +1201,52 @@ bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*
                                                   std::vector<TensorState*>& required_states,
                                                   ggml_backend_t compute_backend) const {
     required_states.clear();
+    required_states.reserve(tensors.size());
+    auto append_states = [&](const std::vector<TensorState*>& states) {
+        for (TensorState* state : states) {
+            if (compute_backend == nullptr || state->compute_backend == nullptr ||
+                state->compute_backend == compute_backend) {
+                required_states.push_back(state);
+            }
+        }
+    };
+    for (auto it = resolved_tensor_states_.begin(); it != resolved_tensor_states_.end(); ++it) {
+        if (it->tensors == tensors) {
+            append_states(it->states);
+            resolved_tensor_states_.splice(resolved_tensor_states_.begin(), resolved_tensor_states_, it);
+            return true;
+        }
+    }
+    std::vector<TensorState*> states;
+    states.reserve(tensors.size());
     std::unordered_set<TensorState*> seen;
+    seen.reserve(tensors.size());
+    bool cacheable = true;
     for (ggml_tensor* tensor : tensors) {
         if (tensor == nullptr) {
             continue;
         }
-        auto param = resolve_param_tensor(tensor);
-        auto found = tensor_states_by_tensor_.find(param);
+        auto found = tensor_states_by_tensor_.find(tensor);
+        // Unregistered views can be rebound without changing the parameter list.
+        cacheable &= found != tensor_states_by_tensor_.end();
+        for (auto view = tensor->view_src; found == tensor_states_by_tensor_.end() && view != nullptr; view = view->view_src) {
+            found = tensor_states_by_tensor_.find(view);
+        }
         if (found == tensor_states_by_tensor_.end()) {
             LOG_ERROR("model manager tensor '%s' is not registered", ggml_get_name(tensor));
             return false;
         }
         TensorState* state = found->second;
-        if ((compute_backend == nullptr || state->compute_backend == nullptr ||
-             state->compute_backend == compute_backend) &&
-            seen.insert(state).second) {
-            required_states.push_back(state);
+        if (seen.insert(state).second) {
+            states.push_back(state);
+        }
+    }
+    append_states(states);
+    if (cacheable && !tensors.empty()) {
+        static constexpr size_t MAX_RESOLVED_LISTS = 4;
+        resolved_tensor_states_.push_front({tensors, std::move(states)});
+        if (resolved_tensor_states_.size() > MAX_RESOLVED_LISTS) {
+            resolved_tensor_states_.pop_back();
         }
     }
     return true;
@@ -1274,8 +1306,7 @@ size_t ModelManager::compute_backend_alloc_size(const std::vector<TensorState*>&
     size_t total_size = 0;
     std::unordered_set<TensorState*> seen;
     for (TensorState* state : states) {
-        if (state == nullptr || state->tensor == nullptr || !seen.insert(state).second ||
-            should_ignore(*state) || is_optional_missing_tensor(state->name)) {
+        if (state == nullptr || state->tensor == nullptr) {
             continue;
         }
         const bool compute_resident =
@@ -1285,6 +1316,9 @@ size_t ModelManager::compute_backend_alloc_size(const std::vector<TensorState*>&
         if (missing_only && compute_resident) {
             continue;
         }
+        if (!seen.insert(state).second || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
+            continue;
+        }
 
         ggml_backend_buffer_type_t buffer_type = nullptr;
         if (state->compute_backend == state->params_backend) {
@@ -1579,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;
@@ -1597,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);
         }
@@ -1752,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

@@ -2,6 +2,7 @@
 #define __MODEL_MANAGER_H__
 
 #include <cstdint>
+#include <list>
 #include <map>
 #include <memory>
 #include <set>
@@ -84,9 +85,15 @@ private:
         size_t resident_bytes          = 0;
     };
 
+    struct ResolvedTensorStates {
+        std::vector<ggml_tensor*> tensors;
+        std::vector<TensorState*> states;
+    };
+
     ModelLoader model_loader_;
     std::vector<std::unique_ptr<TensorState>> tensor_states_;
     std::map<const ggml_tensor*, TensorState*> tensor_states_by_tensor_;
+    mutable std::list<ResolvedTensorStates> resolved_tensor_states_;
     std::vector<std::unique_ptr<ParamsStorageBlock>> params_storage_blocks_;
     std::vector<std::unique_ptr<ComputeStagingBlock>> compute_staging_blocks_;
     std::map<ggml_backend_t, ggml_backend_buffer_type_t> split_buffer_types_;
@@ -150,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,15 +858,67 @@ 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");
             return false;
         }
     }
-    auto configuration        = std::make_unique<ModelConfig>(*sd_ctx_params);
-    n_threads                 = sd_ctx_params->n_threads;
+    auto configuration = std::make_unique<ModelConfig>(*sd_ctx_params);
+    n_threads          = sd_ctx_params->n_threads;
+    tensor_executor    = std::make_unique<sd::ParallelExecutor>(n_threads > 0 ? n_threads : sd_get_num_physical_cores());
+    sd::ParallelScope tensor_scope(tensor_executor.get());
     enable_mmap               = sd_ctx_params->enable_mmap;
     disable_prefetch          = sd_ctx_params->disable_prefetch;
     disable_segmented_compute = sd_ctx_params->disable_segmented_compute;
@@ -1123,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;
@@ -1286,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)) {
@@ -1306,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;
             }
@@ -2158,6 +2218,10 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
 
     RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr};
 
+    const bool apply_denoise_mask = !denoise_mask.empty() &&
+                                    std::any_of(denoise_mask.values().begin(), denoise_mask.values().end(),
+                                                [](float value) { return value != 1.f; });
+
     std::vector<int> skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count);
     float cfg_scale     = guidance.txt_cfg;
     float img_cfg_scale = guidance.img_cfg;
@@ -2287,13 +2351,13 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
             hunyuan_timestep_r_tensor = sd::Tensor<float>::from_vector({sigmas[step + 1]});
         }
         sd::Tensor<float> noised_input = x * c_in;
-        if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) {
+        if (apply_denoise_mask && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) {
             noised_input = noised_input * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
         }
 
         if (cache_runtime.spectrum_enabled && cache_runtime.spectrum.should_predict()) {
             cache_runtime.spectrum.predict(&denoised);
-            if (!denoise_mask.empty()) {
+            if (apply_denoise_mask) {
                 denoised = denoised * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
             }
             if (preview_needed && sd_should_preview_denoised()) {
@@ -2361,6 +2425,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};
@@ -2381,6 +2447,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,
@@ -2516,7 +2585,7 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
         if (cache_runtime.spectrum_enabled) {
             cache_runtime.spectrum.update(denoised);
         }
-        if (!denoise_mask.empty()) {
+        if (apply_denoise_mask) {
             denoised = denoised * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
         }
         if (preview_needed && sd_should_preview_denoised()) {
@@ -2561,7 +2630,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;
@@ -2577,6 +2646,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)) {
@@ -2605,7 +2676,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) {
@@ -2695,7 +2766,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;
@@ -2758,6 +2830,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

@@ -8,6 +8,7 @@
 #include <map>
 #include <memory>
 #include <mutex>
+#include <optional>
 #include <set>
 #include <string>
 #include <vector>
@@ -36,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) ||
@@ -56,8 +57,9 @@ public:
     std::shared_ptr<RNG> rng;
     std::shared_ptr<RNG> sampler_rng = nullptr;
     int n_threads                    = -1;
-    float default_flow_shift         = INFINITY;
-    float active_flow_shift          = INFINITY;
+    std::unique_ptr<sd::ParallelExecutor> tensor_executor;
+    float default_flow_shift = INFINITY;
+    float active_flow_shift  = INFINITY;
 
     std::shared_ptr<Conditioner> cond_stage_model;
     std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision;  // for svd or wan2.1 i2v
@@ -128,6 +130,7 @@ public:
                                 &sd_ctx_params_t::clip_g_path, &sd_ctx_params_t::clip_vision_path,
                                 &sd_ctx_params_t::t5xxl_path, &sd_ctx_params_t::llm_path,
                                 &sd_ctx_params_t::llm_vision_path, &sd_ctx_params_t::diffusion_model_path,
+                                &sd_ctx_params_t::tokenizer,
                                 &sd_ctx_params_t::high_noise_diffusion_model_path, &sd_ctx_params_t::uncond_diffusion_model_path,
                                 &sd_ctx_params_t::embeddings_connectors_path, &sd_ctx_params_t::vae_path,
                                 &sd_ctx_params_t::audio_vae_path, &sd_ctx_params_t::taesd_path,
@@ -205,6 +208,7 @@ public:
         StableDiffusionGGML& sd;
         std::unique_lock<std::recursive_mutex> lock;
         bool acquired = false;
+        std::optional<sd::ParallelScope> tensor_scope;
 
         explicit ContextOperation(StableDiffusionGGML& sd)
             : sd(sd), lock(sd.execution_mutex, std::try_to_lock) {
@@ -214,6 +218,7 @@ public:
             }
             sd.executing_ = true;
             acquired      = true;
+            tensor_scope.emplace(sd.tensor_executor.get());
         }
 
         ~ContextOperation() {
@@ -307,6 +312,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

@@ -285,7 +285,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 +310,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 +475,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 +789,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;
         }
 

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"
@@ -73,12 +75,13 @@ namespace sd::model_builders {
         }
     }
 
-    bool build_core_runners(const Context& ctx, CoreRunners& runners) {
+    bool build_core_runners(const Context& ctx, CoreRunners& runners) try {
         const auto* sd_ctx_params      = &ctx.params;
         const auto& tensor_storage_map = ctx.tensor_storage_map;
         const auto version             = ctx.version;
         const auto& weight_manager     = ctx.weight_manager;
         CoreRunners result;
+        TokenizerConfig tokenizers(sd_ctx_params->tokenizer);
         if (!ensure_backend_pair(ctx.backends, SDBackendModule::TE) ||
             !ensure_backend_pair(ctx.backends, SDBackendModule::DIFFUSION)) {
             return false;
@@ -87,7 +90,8 @@ namespace sd::model_builders {
         if (sd_version_is_sd3(version)) {
             result.conditioner = std::make_shared<SD3CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                    tensor_storage_map,
-                                                                   weight_manager);
+                                                                   weight_manager,
+                                                                   tokenizers);
             result.diffusion   = std::make_shared<MMDiTRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                              tensor_storage_map,
                                                              "model.diffusion_model",
@@ -98,7 +102,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Pid::PiDRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                 tensor_storage_map,
                                                                 "model.diffusion_model.net",
@@ -109,7 +114,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Ideogram4::Ideogram4Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                             tensor_storage_map,
                                                                             "model.diffusion_model",
@@ -120,7 +126,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                true,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Krea2::Krea2Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                     tensor_storage_map,
                                                                     "model.diffusion_model",
@@ -147,11 +154,13 @@ namespace sd::model_builders {
                                                                    version,
                                                                    "",
                                                                    false,
-                                                                   weight_manager);
+                                                                   weight_manager,
+                                                                   tokenizers);
             } else {
                 result.conditioner = std::make_shared<FluxCLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                         tensor_storage_map,
-                                                                        weight_manager);
+                                                                        weight_manager,
+                                                                        tokenizers);
             }
             result.diffusion = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                   tensor_storage_map,
@@ -166,7 +175,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                   tensor_storage_map,
                                                                   "model.diffusion_model",
@@ -178,7 +188,8 @@ namespace sd::model_builders {
                                                                  tensor_storage_map,
                                                                  "text_encoders.llm",
                                                                  "text_embedding_projection",
-                                                                 weight_manager);
+                                                                 weight_manager,
+                                                                 tokenizers);
             result.diffusion   = std::make_shared<LTXV::LTXAVRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                    tensor_storage_map,
                                                                    "model.diffusion_model",
@@ -189,7 +200,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                true,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<MiniMaxH3::MiniMaxH3Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                             tensor_storage_map,
                                                                             "model.diffusion_model",
@@ -200,7 +212,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Hunyuan::HunyuanVideoRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                              tensor_storage_map,
                                                                              "model.diffusion_model",
@@ -258,7 +271,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                enable_vision,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<LingBotVideo::LingBotVideoRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                                   tensor_storage_map,
                                                                                   "model.diffusion_model",
@@ -271,20 +285,29 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                enable_vision,
-                                                               weight_manager);
-            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);
+                                                               weight_manager,
+                                                               tokenizers);
+            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,
                                                                version,
                                                                "",
                                                                true,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<MageFlow::MageFlowRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                           tensor_storage_map,
                                                                           "model.diffusion_model",
@@ -295,7 +318,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                true,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                   tensor_storage_map,
                                                                   "model.diffusion_model",
@@ -305,7 +329,8 @@ namespace sd::model_builders {
         } else if (version == VERSION_HIDREAM_O1) {
             result.conditioner = std::make_shared<HiDreamO1::HiDreamO1Conditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                                    tensor_storage_map,
-                                                                                   weight_manager);
+                                                                                   weight_manager,
+                                                                                   tokenizers);
             result.diffusion   = std::make_shared<HiDreamO1::HiDreamO1Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                             tensor_storage_map,
                                                                             "model",
@@ -327,7 +352,8 @@ namespace sd::model_builders {
         } else if (sd_version_is_anima(version)) {
             result.conditioner = std::make_shared<AnimaConditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                     tensor_storage_map,
-                                                                    weight_manager);
+                                                                    weight_manager,
+                                                                    tokenizers);
             result.diffusion   = std::make_shared<Anima::AnimaRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                     tensor_storage_map,
                                                                     "model.diffusion_model",
@@ -338,19 +364,34 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<ZImage::ZImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                       tensor_storage_map,
                                                                       "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,
                                                                version,
                                                                "",
                                                                true,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Boogu::BooguImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                          tensor_storage_map,
                                                                          "model.diffusion_model",
@@ -362,7 +403,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<ErnieImage::ErnieImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                               tensor_storage_map,
                                                                               "model.diffusion_model",
@@ -373,7 +415,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Lens::LensRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                   tensor_storage_map,
                                                                   "model.diffusion_model",
@@ -387,7 +430,8 @@ namespace sd::model_builders {
                                                                                      tensor_storage_map,
                                                                                      embbeding_map,
                                                                                      version,
-                                                                                     weight_manager);
+                                                                                     weight_manager,
+                                                                                     tokenizers);
             result.diffusion   = std::make_shared<UNetModelRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                  tensor_storage_map,
                                                                  "model.diffusion_model",
@@ -429,8 +473,12 @@ namespace sd::model_builders {
         if (result.ip_adapter) {
             result.ip_adapter->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
         }
+        tokenizers.validate_usage();
         runners = std::move(result);
         return true;
+    } catch (const std::exception& error) {
+        LOG_ERROR("failed to build model runners: %s", error.what());
+        return false;
     }
 
     bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners) {
@@ -558,10 +606,12 @@ namespace sd::model_builders {
         }
 
         if (sd_ctx_params->vae_conv_direct) {
-            LOG_INFO("Using Conv2d direct in the vae model");
+            LOG_INFO("Using Conv2d/Conv3d direct in the vae model");
             result.vae->set_conv2d_direct_enabled(true);
+            result.vae->set_conv3d_direct_enabled(true);
             if (result.preview) {
                 result.preview->set_conv2d_direct_enabled(true);
+                result.preview->set_conv3d_direct_enabled(true);
             }
         }
         if (result.vae) {

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

@@ -1139,10 +1139,10 @@ namespace sd::pipeline {
         return latents;
     }
 
-    static ImageGenerationEmbeds prepare_video_generation_embeds(StableDiffusionGGML* sd,
-                                                                 const sd_vid_gen_params_t* sd_vid_gen_params,
-                                                                 const GenerationRequest& request,
-                                                                 const ImageGenerationLatents& latents) {
+    static std::optional<ImageGenerationEmbeds> prepare_video_generation_embeds(StableDiffusionGGML* sd,
+                                                                                const sd_vid_gen_params_t* sd_vid_gen_params,
+                                                                                const GenerationRequest& request,
+                                                                                const ImageGenerationLatents& latents) {
         ConditionerRunnerEndOnExit conditioner_runner_end{sd->cond_stage_model.get()};
 
         ImageGenerationEmbeds embeds;
@@ -1159,8 +1159,12 @@ namespace sd::pipeline {
         int64_t prepare_start_ms = ggml_time_ms();
         embeds.cond              = sd->cond_stage_model->get_learned_condition(sd->n_threads,
                                                                                condition_params);
-        embeds.cond.c_concat     = latents.concat_latent;
-        embeds.cond.c_vector     = latents.clip_vision_output;
+        if (embeds.cond.empty()) {
+            LOG_ERROR("failed to encode video prompt");
+            return std::nullopt;
+        }
+        embeds.cond.c_concat = latents.concat_latent;
+        embeds.cond.c_vector = latents.clip_vision_output;
         if (sd_version_is_minimax_h3(sd->version)) {
             embeds.cond.c_ref_images       = latents.ref_latents;
             embeds.cond.c_ref_audios       = latents.reference_audio_latents;
@@ -1178,9 +1182,13 @@ namespace sd::pipeline {
             }
         }
         if (request.use_uncond) {
-            condition_params.text  = request.negative_prompt;
-            embeds.uncond          = sd->cond_stage_model->get_learned_condition(sd->n_threads,
-                                                                                 condition_params);
+            condition_params.text = request.negative_prompt;
+            embeds.uncond         = sd->cond_stage_model->get_learned_condition(sd->n_threads,
+                                                                                condition_params);
+            if (embeds.uncond.empty()) {
+                LOG_ERROR("failed to encode negative video prompt");
+                return std::nullopt;
+            }
             embeds.uncond.c_concat = latents.concat_latent;
             embeds.uncond.c_vector = latents.clip_vision_output;
             if (sd_version_is_minimax_h3(sd->version)) {
@@ -1574,10 +1582,14 @@ namespace sd::pipeline {
         }
         ImageGenerationLatents latents = std::move(*latent_inputs_opt);
 
-        ImageGenerationEmbeds embeds = prepare_video_generation_embeds(sd,
-                                                                       sd_vid_gen_params,
-                                                                       request,
-                                                                       latents);
+        auto embeds_opt = prepare_video_generation_embeds(sd,
+                                                          sd_vid_gen_params,
+                                                          request,
+                                                          latents);
+        if (!embeds_opt) {
+            return false;
+        }
+        ImageGenerationEmbeds embeds = std::move(*embeds_opt);
         if (latent_upscale_enabled) {
             LOG_INFO("generate_video %dx%dx%d -> LTX latent spatial upscale",
                      request.width,

src/runtime/denoiser.hpp

@@ -12,6 +12,8 @@
 #include <utility>
 
 #include "core/rng.hpp"
+#include "core/rng_mt19937.hpp"
+#include "core/rng_philox.hpp"
 #include "core/tensor.hpp"
 #include "core/util.h"
 #include "model.h"
@@ -784,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;
 
@@ -1121,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);
@@ -1632,12 +1687,18 @@ static std::tuple<float, float, float> get_ancestral_step(float sigma_from,
     }
 }
 
+class NoiseSampler {
+public:
+    virtual sd::Tensor<float> operator()(double sigma_from, double sigma_to) = 0;
+    virtual ~NoiseSampler()                                                  = default;
+};
+
 static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
                                                 sd::Tensor<float> x,
                                                 const std::vector<float>& sigmas,
-                                                std::shared_ptr<RNG> rng = nullptr,
-                                                bool is_flow_denoiser    = false,
-                                                float eta                = 0.f) {
+                                                NoiseSampler& noise_sampler,
+                                                bool is_flow_denoiser = false,
+                                                float eta             = 0.f) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
         float sigma       = sigmas[i];
@@ -1660,7 +1721,7 @@ static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
                 if (is_flow_denoiser) {
                     x *= alpha_scale;
                 }
-                x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+                x += noise_sampler(sigma, sigma_to) * sigma_up;
             }
         }
     }
@@ -1778,7 +1839,7 @@ static sd::Tensor<float> sample_dpm2(denoise_cb_t model,
 static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
                                                    sd::Tensor<float> x,
                                                    const std::vector<float>& sigmas,
-                                                   std::shared_ptr<RNG> rng,
+                                                   NoiseSampler& noise_sampler,
                                                    float eta) {
     auto t_fn     = [](float sigma) -> float { return -log(sigma); };
     auto sigma_fn = [](float t) -> float { return exp(-t); };
@@ -1810,7 +1871,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
         }
 
         if (sigmas[i + 1] > 0) {
-            x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+            x += noise_sampler(sigmas[i], sigmas[i + 1]) * sigma_up;
         }
     }
     return x;
@@ -1819,7 +1880,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
 static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
                                                         sd::Tensor<float> x,
                                                         const std::vector<float>& sigmas,
-                                                        std::shared_ptr<RNG> rng,
+                                                        NoiseSampler& noise_sampler,
                                                         float eta = 1.0f) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
@@ -1902,7 +1963,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
             x                        = (x * sigma_down_i_ratio) + (D_i * (1.0f - sigma_down_i_ratio));
 
             if (sigma_to > 0.0f && eta > 0.0f) {
-                x = alpha_scale * x + sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+                x = alpha_scale * x + noise_sampler(sigma, sigma_to) * sigma_up;
             }
         }
     }
@@ -1978,169 +2039,17 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
     return x;
 }
 
-// DPM-Solver++(2M) SDE, midpoint variant. Ref: Lu et al. arXiv:2211.01095;
-// k-diffusion sample_dpmpp_2m_sde.
+// DPM-Solver++(2M) SDE, midpoint variant.
+// Ref: Lu et al. arXiv:2211.01095; k-diffusion sample_dpmpp_2m_sde
 static sd::Tensor<float> sample_dpmpp_2m_sde(denoise_cb_t model,
                                              sd::Tensor<float> x,
                                              const std::vector<float>& sigmas,
-                                             std::shared_ptr<RNG> rng,
+                                             NoiseSampler& noise_sampler,
                                              float eta) {
     sd::Tensor<float> old_denoised;
     bool have_old_denoised = false;
     float h_last           = 0.f;
 
-    int steps = static_cast<int>(sigmas.size()) - 1;
-    for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], i + 1);
-        if (denoised_opt.pred.empty()) {
-            return {};
-        }
-        sd::Tensor<float> denoised = std::move(denoised_opt.pred);
-
-        if (sigmas[i + 1] == 0.f) {
-            x = denoised;
-        } else {
-            float t     = -std::log(sigmas[i]);
-            float s     = -std::log(sigmas[i + 1]);
-            float h     = s - t;
-            float eta_h = eta * h;
-            float a     = sigmas[i + 1] / sigmas[i] * std::exp(-eta_h);
-            float b     = -std::expm1(-h - eta_h);
-
-            x = a * x + b * denoised;
-
-            if (have_old_denoised) {
-                float r = h_last / h;
-                x += (0.5f * b / r) * (denoised - old_denoised);
-            }
-            if (eta > 0.f) {
-                x += sd::Tensor<float>::randn_like(x, rng) * (sigmas[i + 1] * std::sqrt(-std::expm1(-2.f * eta_h)));
-            }
-            h_last = h;
-        }
-        old_denoised      = denoised;
-        have_old_denoised = true;
-    }
-    return x;
-}
-
-// Seeded Brownian tree providing deterministic, step-count-stable Gaussian
-// increments for stochastic samplers. Constructed once per generation; each
-// call returns unit-variance noise for interval [sigma_a, sigma_b].
-// Reference: torchsde BrownianTree; k-diffusion BatchedBrownianTree.
-class BrownianTreeNoiseSampler {
-public:
-    BrownianTreeNoiseSampler(const sd::Tensor<float>& x_template,
-                             double sigma_min,
-                             double sigma_max,
-                             uint64_t seed)
-        : t_min_(sigma_min),
-          t_max_(sigma_max),
-          shape_(x_template.shape()),
-          root_seed_(mix64(seed, 0x9E3779B97F4A7C15ULL)) {
-        auto rng = std::make_shared<STDDefaultRNG>();
-        rng->manual_seed(mix64(seed, 0xBF58476D1CE4E5B9ULL));
-        w_at_tmax_ = sd::Tensor<float>::randn(shape_, rng) * std::sqrt(static_cast<float>(t_max_ - t_min_));
-    }
-
-    sd::Tensor<float> operator()(double sigma_a, double sigma_b) {
-        double a   = clamp(std::min(sigma_a, sigma_b));
-        double b   = clamp(std::max(sigma_a, sigma_b));
-        auto dW    = w(b) - w(a);
-        float span = static_cast<float>(std::max(std::abs(sigma_b - sigma_a), 1e-12));
-        return dW * (1.0f / std::sqrt(span));
-    }
-
-private:
-    static constexpr int kMaxDepth = 24;
-
-    static uint64_t mix64(uint64_t v, uint64_t salt) {
-        uint64_t z = v + salt;
-        z          = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
-        z          = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
-        return z ^ (z >> 31);
-    }
-
-    double clamp(double t) const {
-        return std::min(std::max(t, t_min_), t_max_);
-    }
-
-    sd::Tensor<float> w(double t) {
-        auto it = cache_.find(t);
-        if (it != cache_.end()) {
-            return it->second;
-        }
-        sd::Tensor<float> zero = sd::Tensor<float>::zeros(shape_);
-        sd::Tensor<float> out  = bridge(t_min_, t_max_, zero, w_at_tmax_, t, root_seed_, kMaxDepth);
-        cache_.emplace(t, out);
-        return out;
-    }
-
-    sd::Tensor<float> bridge(double a,
-                             double c,
-                             const sd::Tensor<float>& w_a,
-                             const sd::Tensor<float>& w_c,
-                             double t,
-                             uint64_t node_seed,
-                             int depth) {
-        if (depth <= 0 || c - a < 1e-9) {
-            float alpha = (c > a) ? static_cast<float>((t - a) / (c - a)) : 0.5f;
-            return (1.0f - alpha) * w_a + alpha * w_c;
-        }
-        double m       = 0.5 * (a + c);
-        double std_dev = std::sqrt((c - m) * (m - a) / (c - a));
-        auto rng       = std::make_shared<STDDefaultRNG>();
-        rng->manual_seed(node_seed);
-        auto z   = sd::Tensor<float>::randn(shape_, rng);
-        auto w_m = 0.5f * (w_a + w_c) + static_cast<float>(std_dev) * z;
-        if (t == m) {
-            return w_m;
-        }
-        if (t < m) {
-            return bridge(a, m, w_a, w_m, t, mix64(node_seed, 1), depth - 1);
-        }
-        return bridge(m, c, w_m, w_c, t, mix64(node_seed, 2), depth - 1);
-    }
-
-    double t_min_;
-    double t_max_;
-    std::vector<int64_t> shape_;
-    uint64_t root_seed_;
-    sd::Tensor<float> w_at_tmax_;
-    std::map<double, sd::Tensor<float>> cache_;
-};
-
-// DPM-Solver++(2M) SDE, midpoint variant, with step-count-stable Brownian-tree
-// noise. Same trajectory shape at any step count for a given seed. Aliased in
-// k-diffusion / ComfyUI as sample_dpmpp_2m_sde_gpu.
-// Ref: Lu et al. arXiv:2211.01095; torchsde BrownianTree.
-static sd::Tensor<float> sample_dpmpp_2m_sde_bt(denoise_cb_t model,
-                                                sd::Tensor<float> x,
-                                                const std::vector<float>& sigmas,
-                                                std::shared_ptr<RNG> rng,
-                                                float eta) {
-    double sigma_max = 0.0;
-    double sigma_min = std::numeric_limits<double>::infinity();
-    for (float s : sigmas) {
-        if (s > 0.0f) {
-            sigma_max = std::max(sigma_max, static_cast<double>(s));
-            sigma_min = std::min(sigma_min, static_cast<double>(s));
-        }
-    }
-    if (sigma_max <= sigma_min) {
-        return x;
-    }
-    uint64_t tree_seed = 0;
-    {
-        auto draw = rng->randn(2);
-        std::memcpy(&tree_seed, draw.data(), sizeof(tree_seed));
-    }
-    BrownianTreeNoiseSampler noise_sampler(x, sigma_min, sigma_max, tree_seed);
-
-    sd::Tensor<float> old_denoised;
-    bool have_old_denoised = false;
-    float h_last           = 0.f;
-
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
         auto denoised_opt = model(x, sigmas[i], i + 1);
@@ -2181,7 +2090,7 @@ using SamplerExtraArgs = KeyValueArgs;
 static sd::Tensor<float> sample_lcm(denoise_cb_t model,
                                     sd::Tensor<float> x,
                                     const std::vector<float>& sigmas,
-                                    std::shared_ptr<RNG> rng,
+                                    NoiseSampler& noise_sampler,
                                     bool is_flow_denoiser,
                                     const SamplerExtraArgs& extra_sample_args) {
     struct LCMSampleArgs {
@@ -2234,7 +2143,7 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
             if (is_flow_denoiser) {
                 x *= (1 - sigmas[i + 1]);
             }
-            auto noise = sd::Tensor<float>::randn_like(x, rng);
+            auto noise = noise_sampler(sigmas[i], sigmas[i + 1]);
             if (args.noise_clip_std > 0.0f && noise.numel() > 0) {
                 double mean = 0.0;
                 for (int64_t j = 0; j < noise.numel(); ++j) {
@@ -2352,7 +2261,7 @@ static sd::Tensor<float> sample_ipndm_v(denoise_cb_t model,
 static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
                                               sd::Tensor<float> x,
                                               const std::vector<float>& sigmas,
-                                              std::shared_ptr<RNG> rng,
+                                              NoiseSampler& noise_sampler,
                                               bool is_flow_denoiser,
                                               float eta) {
     sd::Tensor<float> old_denoised = x;
@@ -2417,7 +2326,7 @@ static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
             if (is_flow_denoiser) {
                 x *= alpha_scale;
             }
-            x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+            x += noise_sampler(sigma_from, sigma_to) * sigma_up;
         }
 
         old_denoised   = denoised;
@@ -2430,7 +2339,7 @@ static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
 static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
                                        sd::Tensor<float> x,
                                        const std::vector<float>& sigmas,
-                                       std::shared_ptr<RNG> rng,
+                                       NoiseSampler& noise_sampler,
                                        bool is_flow_denoiser,
                                        float eta) {
     const float c2 = 0.5f;
@@ -2493,7 +2402,7 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
             if (is_flow_denoiser) {
                 x *= alpha_scale;
             }
-            x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+            x += noise_sampler(sigma_from, sigma_to) * sigma_up;
         }
     }
     return x;
@@ -2502,7 +2411,7 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
 static sd::Tensor<float> sample_er_sde(denoise_cb_t model,
                                        sd::Tensor<float> x,
                                        std::vector<float> sigmas,
-                                       std::shared_ptr<RNG> rng,
+                                       NoiseSampler& noise_sampler,
                                        bool is_flow_denoiser,
                                        float eta) {
     constexpr int max_stage                  = 3;
@@ -2624,7 +2533,7 @@ static sd::Tensor<float> sample_er_sde(denoise_cb_t model,
             float noise_scale_sq = er_lambda_t * er_lambda_t - er_lambda_s * er_lambda_s * r * r;
             if (s_noise > 0.0f && noise_scale_sq > 0.0f) {
                 float noise_scale = alpha_t * std::sqrt(std::max(noise_scale_sq, 0.0f));
-                x += sd::Tensor<float>::randn_like(x, rng) * noise_scale;
+                x += noise_sampler(sigmas[i], sigmas[i + 1]) * noise_scale;
             }
         }
 
@@ -2637,7 +2546,7 @@ static sd::Tensor<float> sample_er_sde(denoise_cb_t model,
 static sd::Tensor<float> sample_tcd(denoise_cb_t model,
                                     sd::Tensor<float> x,
                                     const std::vector<float>& sigmas,
-                                    std::shared_ptr<RNG> rng,
+                                    NoiseSampler& noise_sampler,
                                     float eta) {
     float beta_start = 0.00085f;
     float beta_end   = 0.0120f;
@@ -2694,7 +2603,7 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
 
         if (eta > 0 && sigma_to > 0.0f) {
             x = std::sqrt(alpha_prod_t_prev / alpha_prod_s) * x +
-                std::sqrt(1.0f / alpha_prod_t_prev - 1.0f / alpha_prod_s) * sd::Tensor<float>::randn_like(x, rng);
+                std::sqrt(1.0f / alpha_prod_t_prev - 1.0f / alpha_prod_s) * noise_sampler(sigma, sigma_to);
         }
     }
     return x;
@@ -2829,7 +2738,7 @@ static sd::Tensor<float> sample_euler_cfg_pp(denoise_cb_t model,
 static sd::Tensor<float> sample_euler_ancestral_cfg_pp(denoise_cb_t model,
                                                        sd::Tensor<float> x,
                                                        const std::vector<float>& sigmas,
-                                                       std::shared_ptr<RNG> rng,
+                                                       NoiseSampler& noise_sampler,
                                                        float eta) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
@@ -2848,7 +2757,7 @@ static sd::Tensor<float> sample_euler_ancestral_cfg_pp(denoise_cb_t model,
         x = denoised + d * sigma_down;
 
         if (sigmas[i + 1] > 0) {
-            x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+            x += noise_sampler(sigmas[i], sigmas[i + 1]) * sigma_up;
         }
     }
     return x;
@@ -2858,7 +2767,7 @@ static sd::Tensor<float> sample_euler_ancestral_cfg_pp(denoise_cb_t model,
 static sd::Tensor<float> sample_gradient_estimation(denoise_cb_t model,
                                                     sd::Tensor<float> x,
                                                     const std::vector<float>& sigmas,
-                                                    std::shared_ptr<RNG> rng,
+                                                    NoiseSampler& noise_sampler,
                                                     bool is_flow_denoiser,
                                                     float eta,
                                                     const SamplerExtraArgs& extra_sample_args) {
@@ -2905,13 +2814,180 @@ static sd::Tensor<float> sample_gradient_estimation(denoise_cb_t model,
                 if (is_flow_denoiser) {
                     x *= alpha_scale;
                 }
-                x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+                x += noise_sampler(sigma, sigma_to) * sigma_up;
             }
         }
     }
     return x;
 }
 
+// independent and identically distributed Gaussian noise (default for most samplers)
+class IIDGaussianNoiseSampler : public NoiseSampler {
+public:
+    IIDGaussianNoiseSampler(const sd::Tensor<float>& x_template, std::shared_ptr<RNG> r)
+        : rng(std::move(r)), shape(x_template.shape()) {}
+    sd::Tensor<float> operator()(double sigma_from, double sigma_to) override {
+        (void)sigma_from;
+        (void)sigma_to;
+        return sd::Tensor<float>::randn(shape, rng);
+    }
+
+private:
+    std::shared_ptr<RNG> rng;
+    std::vector<int64_t> shape;
+};
+
+// A fixed tree seed, shape and sigma range give consistent increments across
+// interval subdivisions. Each query returns normalized Gaussian noise.
+// Reference: torchsde BrownianTree; k-diffusion BatchedBrownianTree.
+class BrownianTreeNoiseSampler : public NoiseSampler {
+public:
+    BrownianTreeNoiseSampler(const sd::Tensor<float>& x_template,
+                             double sigma_min,
+                             double sigma_max,
+                             std::shared_ptr<RNG> seed_rng,
+                             std::shared_ptr<RNG> node_rng)
+        : t_min_(sigma_min),
+          t_max_(sigma_max),
+          shape_(x_template.shape()),
+          seed_rng_(std::move(seed_rng)),
+          node_rng_(std::move(node_rng)) {}
+
+    sd::Tensor<float> operator()(double sigma_a, double sigma_b) override {
+        if (!initialized_) {
+            uint64_t seed = 0;
+            auto draw     = seed_rng_->randn(2);
+            std::memcpy(&seed, draw.data(), sizeof(seed));
+            root_seed_ = mix64(seed, 0x9E3779B97F4A7C15ULL);
+            node_rng_->manual_seed(mix64(seed, 0xBF58476D1CE4E5B9ULL));
+            w_at_tmax_   = sd::Tensor<float>::randn(shape_, node_rng_) * std::sqrt(static_cast<float>(t_max_ - t_min_));
+            initialized_ = true;
+        }
+        double a   = clamp(std::min(sigma_a, sigma_b));
+        double b   = clamp(std::max(sigma_a, sigma_b));
+        auto dW    = w(b) - w(a);
+        float span = static_cast<float>(std::max(std::abs(sigma_b - sigma_a), 1e-12));
+        return dW * (1.0f / std::sqrt(span));
+    }
+
+private:
+    static constexpr int kMaxDepth = 24;
+
+    static uint64_t mix64(uint64_t v, uint64_t salt) {
+        uint64_t z = v + salt;
+        z          = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
+        z          = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
+        return z ^ (z >> 31);
+    }
+
+    double clamp(double t) const {
+        return std::min(std::max(t, t_min_), t_max_);
+    }
+
+    sd::Tensor<float> w(double t) {
+        auto it = cache_.find(t);
+        if (it != cache_.end()) {
+            return it->second;
+        }
+        sd::Tensor<float> zero = sd::Tensor<float>::zeros(shape_);
+        sd::Tensor<float> out  = bridge(t_min_, t_max_, zero, w_at_tmax_, t, root_seed_, kMaxDepth);
+        cache_.emplace(t, out);
+        return out;
+    }
+
+    sd::Tensor<float> bridge(double a,
+                             double c,
+                             const sd::Tensor<float>& w_a,
+                             const sd::Tensor<float>& w_c,
+                             double t,
+                             uint64_t node_seed,
+                             int depth) {
+        if (depth <= 0 || c - a < 1e-9) {
+            float alpha = (c > a) ? static_cast<float>((t - a) / (c - a)) : 0.5f;
+            return (1.0f - alpha) * w_a + alpha * w_c;
+        }
+        double m       = 0.5 * (a + c);
+        double std_dev = std::sqrt((c - m) * (m - a) / (c - a));
+        node_rng_->manual_seed(node_seed);
+        auto z   = sd::Tensor<float>::randn(shape_, node_rng_);
+        auto w_m = 0.5f * (w_a + w_c) + static_cast<float>(std_dev) * z;
+        if (t == m) {
+            return w_m;
+        }
+        if (t < m) {
+            return bridge(a, m, w_a, w_m, t, mix64(node_seed, 1), depth - 1);
+        }
+        return bridge(m, c, w_m, w_c, t, mix64(node_seed, 2), depth - 1);
+    }
+
+    double t_min_;
+    double t_max_;
+    std::vector<int64_t> shape_;
+    std::shared_ptr<RNG> seed_rng_;
+    std::shared_ptr<RNG> node_rng_;
+    uint64_t root_seed_ = 0;
+    bool initialized_   = false;
+    sd::Tensor<float> w_at_tmax_;
+    std::map<double, sd::Tensor<float>> cache_;
+};
+
+static std::unique_ptr<NoiseSampler> make_noise_sampler(const sd::Tensor<float>& x, std::shared_ptr<RNG> rng, sample_method_t method, const std::vector<float>& sigmas, const SamplerExtraArgs& extra_args) {
+    bool brownian_tree            = (method == DPMPP2M_SDE_BT_SAMPLE_METHOD);
+    bool def_brownian_tree        = brownian_tree;
+    std::string brownian_tree_rng = "cpu";
+
+    for (const auto& [key, value] : extra_args) {
+        if (key == "noise_sampler") {
+            if (value == "iid") {
+                brownian_tree = false;
+            } else if (value == "brownian_tree") {
+                brownian_tree = true;
+            } else {
+                LOG_WARN("unknown noise_sampler value '%s'; using default", value.c_str());
+            }
+        } else if (key == "brownian_tree_rng") {
+            if (value == "cpu" || value == "cuda" || value == "std_default" || value == "sampler_rng") {
+                brownian_tree_rng = value;
+            } else {
+                LOG_WARN("ignoring invalid brownian_tree_rng value '%s'; expected cpu, cuda, std_default or sampler_rng", value.c_str());
+            }
+        }
+    }
+
+    if (brownian_tree) {
+        double sigma_max = 0.0;
+        double sigma_min = std::numeric_limits<double>::infinity();
+        for (float s : sigmas) {
+            if (s > 0.0f) {
+                sigma_max = std::max(sigma_max, static_cast<double>(s));
+                sigma_min = std::min(sigma_min, static_cast<double>(s));
+            }
+        }
+
+        if (sigma_max > sigma_min) {
+            std::shared_ptr<RNG> node_rng;
+            if (brownian_tree_rng == "sampler_rng") {
+                node_rng = rng->clone();
+            } else if (brownian_tree_rng == "std_default") {
+                node_rng = std::make_shared<STDDefaultRNG>();
+            } else if (brownian_tree_rng == "cuda") {
+                node_rng = std::make_shared<PhiloxRNG>();
+            } else {
+                node_rng = std::make_shared<MT19937RNG>();
+            }
+            if (!def_brownian_tree) {
+                LOG_INFO("setting noise sampler to Brownian tree (%s RNG)", brownian_tree_rng.c_str());
+            }
+            return std::make_unique<BrownianTreeNoiseSampler>(x, sigma_min, sigma_max, rng, std::move(node_rng));
+        }
+    }
+
+    if (def_brownian_tree) {
+        LOG_INFO("setting noise sampler to independent and identically distributed (iid)");
+    }
+    return std::make_unique<IIDGaussianNoiseSampler>(x, rng);
+}
+
 // k diffusion reverse ODE: dx = (x - D(x;\sigma)) / \sigma dt; \sigma(t) = t
 static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
                                             denoise_cb_t model,
@@ -2928,9 +3004,12 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
         }
     }
     SamplerExtraArgs extra_args = parse_key_value_args(extra_sample_args, "extra sample arg");
+
+    std::unique_ptr<NoiseSampler> noise_sampler = make_noise_sampler(x, rng, method, sigmas, extra_args);
+
     switch (method) {
         case EULER_A_SAMPLE_METHOD:
-            return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
+            return sample_euler_ancestral(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta);
         case EULER_SAMPLE_METHOD:
             return sample_euler(model, std::move(x), sigmas);
         case HEUN_SAMPLE_METHOD:
@@ -2939,42 +3018,41 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
             return sample_dpm2(model, std::move(x), sigmas);
         case DPMPP2S_A_SAMPLE_METHOD:
             if (is_flow_denoiser)
-                return sample_dpmpp_2s_ancestral_flow(model, std::move(x), sigmas, rng, eta);
+                return sample_dpmpp_2s_ancestral_flow(model, std::move(x), sigmas, *noise_sampler, eta);
             else
-                return sample_dpmpp_2s_ancestral(model, std::move(x), sigmas, rng, eta);
+                return sample_dpmpp_2s_ancestral(model, std::move(x), sigmas, *noise_sampler, eta);
         case DPMPP2M_SAMPLE_METHOD:
             return sample_dpmpp_2m(model, std::move(x), sigmas);
         case DPMPP2Mv2_SAMPLE_METHOD:
             return sample_dpmpp_2m_v2(model, std::move(x), sigmas);
         case LCM_SAMPLE_METHOD:
-            return sample_lcm(model, std::move(x), sigmas, rng, is_flow_denoiser, extra_args);
+            return sample_lcm(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, extra_args);
         case IPNDM_SAMPLE_METHOD:
             return sample_ipndm(model, std::move(x), sigmas);
         case IPNDM_V_SAMPLE_METHOD:
             return sample_ipndm_v(model, std::move(x), sigmas);
         case RES_MULTISTEP_SAMPLE_METHOD:
-            return sample_res_multistep(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
+            return sample_res_multistep(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta);
         case RES_2S_SAMPLE_METHOD:
-            return sample_res_2s(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
+            return sample_res_2s(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta);
         case ER_SDE_SAMPLE_METHOD:
-            return sample_er_sde(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
+            return sample_er_sde(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta);
         case DPMPP2M_SDE_SAMPLE_METHOD:
-            return sample_dpmpp_2m_sde(model, std::move(x), sigmas, rng, eta);
         case DPMPP2M_SDE_BT_SAMPLE_METHOD:
-            return sample_dpmpp_2m_sde_bt(model, std::move(x), sigmas, rng, eta);
+            return sample_dpmpp_2m_sde(model, std::move(x), sigmas, *noise_sampler, eta);
         case DDIM_TRAILING_SAMPLE_METHOD:
             // DDIM is equivalent to Euler Ancestral with the Simple scheduler
-            return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
+            return sample_euler_ancestral(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta);
         case TCD_SAMPLE_METHOD:
-            return sample_tcd(model, std::move(x), sigmas, rng, eta);
+            return sample_tcd(model, std::move(x), sigmas, *noise_sampler, eta);
         case LMS_SAMPLE_METHOD:
             return sample_lms(model, std::move(x), sigmas, extra_args);
         case EULER_CFG_PP_SAMPLE_METHOD:
             return sample_euler_cfg_pp(model, std::move(x), sigmas);
         case EULER_A_CFG_PP_SAMPLE_METHOD:
-            return sample_euler_ancestral_cfg_pp(model, std::move(x), sigmas, rng, eta);
+            return sample_euler_ancestral_cfg_pp(model, std::move(x), sigmas, *noise_sampler, eta);
         case EULER_GE_SAMPLE_METHOD:
-            return sample_gradient_estimation(model, std::move(x), sigmas, rng, is_flow_denoiser, eta, extra_args);
+            return sample_gradient_estimation(model, std::move(x), sigmas, *noise_sampler, is_flow_denoiser, eta, extra_args);
         default:
             return {};
     }

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;
@@ -349,6 +364,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "t5xxl_path: %s\n"
              "llm_path: %s\n"
              "llm_vision_path: %s\n"
+             "tokenizer: %s\n"
              "diffusion_model_path: %s\n"
              "high_noise_diffusion_model_path: %s\n"
              "uncond_diffusion_model_path: %s\n"
@@ -377,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",
@@ -387,6 +404,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              SAFE_STR(sd_ctx_params->t5xxl_path),
              SAFE_STR(sd_ctx_params->llm_path),
              SAFE_STR(sd_ctx_params->llm_vision_path),
+             SAFE_STR(sd_ctx_params->tokenizer),
              SAFE_STR(sd_ctx_params->diffusion_model_path),
              SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path),
              SAFE_STR(sd_ctx_params->uncond_diffusion_model_path),
@@ -415,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));
@@ -611,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) {

src/tokenizers/bpe_tokenizer.cpp

@@ -2,6 +2,7 @@
 
 #include <algorithm>
 #include <sstream>
+#include <stdexcept>
 
 #include "core/util.h"
 #include "tokenize_util.h"
@@ -31,8 +32,37 @@ std::vector<std::pair<int, std::u32string>> BPETokenizer::bytes_to_unicode() {
     return byte_unicode_pairs;
 }
 
-std::vector<std::string> BPETokenizer::token_split(const std::string& text) const {
-    return ::token_split(text);
+BPETokenizer::BPETokenizer(const std::string& pattern) {
+    if (!pattern.empty()) {
+        split_regex_ = std::make_unique<sd::Regex>();
+        std::string error;
+        if (!split_regex_->compile(pattern, &error)) {
+            throw std::runtime_error("invalid tokenizer regex: " + error);
+        }
+    }
+}
+
+bool BPETokenizer::token_split(const std::string& text, std::vector<std::string>& tokens, std::string* error) const {
+    tokens.clear();
+    if (error) {
+        error->clear();
+    }
+    if (!split_regex_) {
+        if (!text.empty()) {
+            tokens.push_back(text);
+        }
+        return true;
+    }
+    std::vector<sd::Regex::Match> matches;
+    if (!split_regex_->find_matches(text, matches, error)) {
+        return false;
+    }
+    for (const auto& match : matches) {
+        if (match.first != match.second) {
+            tokens.push_back(text.substr(match.first, match.second - match.first));
+        }
+    }
+    return true;
 }
 
 std::vector<std::u32string> BPETokenizer::split_utf32(const std::string& text, char32_t delimiter) {
@@ -130,7 +160,11 @@ std::vector<std::u32string> BPETokenizer::bpe(const std::u32string& token) const
     return word;
 }
 
-std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t on_new_token_cb) {
+bool BPETokenizer::encode(const std::string& text, std::vector<int>& result, on_new_token_cb_t on_new_token_cb, std::string* error) {
+    result.clear();
+    if (error) {
+        error->clear();
+    }
     std::vector<int32_t> bpe_tokens;
     std::vector<std::string> token_strs;
 
@@ -150,7 +184,10 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
             token_strs.push_back(splited_text);
             continue;
         }
-        auto tokens = token_split(splited_text);
+        std::vector<std::string> tokens;
+        if (!token_split(splited_text, tokens, error)) {
+            return false;
+        }
         for (auto& token : tokens) {
             if (on_new_token_cb != nullptr) {
                 bool skip = on_new_token_cb(token, bpe_tokens);
@@ -206,7 +243,8 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
     }
     ss << "]";
     LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", text.c_str(), bpe_tokens.size(), ss.str().c_str());
-    return bpe_tokens;
+    result = std::move(bpe_tokens);
+    return true;
 }
 
 std::string BPETokenizer::decode_token(int token_id) const {

src/tokenizers/bpe_tokenizer.h

@@ -5,15 +5,19 @@
 #include <cstdint>
 #include <functional>
 #include <map>
-#include <regex>
+#include <memory>
 #include <set>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include "core/regex.h"
 #include "tokenizer.h"
 
 class BPETokenizer : public Tokenizer {
+private:
+    std::unique_ptr<sd::Regex> split_regex_;
+
 protected:
     std::map<int, std::u32string> byte_encoder;
     std::map<std::u32string, int> byte_decoder;
@@ -28,15 +32,15 @@ protected:
 protected:
     static std::vector<std::pair<int, std::u32string>> bytes_to_unicode();
     static std::vector<std::u32string> split_utf32(const std::string& text, char32_t delimiter = U'\n');
-    virtual std::vector<std::string> token_split(const std::string& text) const;
+    bool token_split(const std::string& text, std::vector<std::string>& tokens, std::string* error = nullptr) const;
     std::vector<std::u32string> bpe(const std::u32string& token) const;
     std::string decode_token(int token_id) const override;
 
 public:
-    BPETokenizer()          = default;
+    explicit BPETokenizer(const std::string& pattern);
     virtual ~BPETokenizer() = default;
 
-    std::vector<int> encode(const std::string& text, on_new_token_cb_t on_new_token_cb = nullptr) override;
+    bool encode(const std::string& text, std::vector<int>& tokens, on_new_token_cb_t on_new_token_cb = nullptr, std::string* error = nullptr) override;
 };
 
 #endif  // __SD_TOKENIZERS_BPE_TOKENIZER_H__

src/tokenizers/clip_tokenizer.cpp

@@ -8,10 +8,10 @@
 
 #include "core/util.h"
 #include "ggml.h"
-#include "tokenize_util.h"
 #include "vocab/vocab.h"
 
-CLIPTokenizer::CLIPTokenizer(int pad_token_id, const std::string& merges_utf8_str) {
+CLIPTokenizer::CLIPTokenizer(int pad_token_id, const std::string& merges_utf8_str)
+    : BPETokenizer(R"((?i:'s|'t|'re|'ve|'m|'ll|'d)|\p{L}+|\p{N}|[^\s\p{L}\p{N}]+)") {
     UNK_TOKEN = "<|endoftext|>";
     BOS_TOKEN = "<|startoftext|>";
     EOS_TOKEN = "<|endoftext|>";
@@ -101,17 +101,3 @@ std::string CLIPTokenizer::normalize(const std::string& text) const {
     std::transform(normalized_text.begin(), normalized_text.end(), normalized_text.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
     return normalized_text;
 }
-
-std::vector<std::string> CLIPTokenizer::token_split(const std::string& text) const {
-    std::regex clip_pat(R"('s|'t|'re|'ve|'m|'ll|'d|[[:alpha:]]+|[[:digit:]]|[^[:space:][:alpha:][:digit:]]+)",
-                        std::regex::icase);
-    std::sregex_iterator iter(text.begin(), text.end(), clip_pat);
-    std::sregex_iterator end;
-
-    std::vector<std::string> result;
-    for (; iter != end; ++iter) {
-        result.emplace_back(iter->str());
-    }
-
-    return result;
-}

src/tokenizers/clip_tokenizer.h

@@ -11,7 +11,6 @@ class CLIPTokenizer : public BPETokenizer {
 protected:
     void load_from_merges(const std::string& merges_utf8_str);
     std::string normalize(const std::string& text) const override;
-    std::vector<std::string> token_split(const std::string& text) const override;
 
 public:
     explicit CLIPTokenizer(int pad_token_id = 49407, const std::string& merges_utf8_str = "");

src/tokenizers/gemma_tokenizer.cpp

@@ -46,7 +46,9 @@ void GemmaTokenizer::load_from_merges(const std::string& merges_utf8_str, const
     bpe_len = rank;
 }
 
-GemmaTokenizer::GemmaTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+GemmaTokenizer::GemmaTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str)
+    : BPETokenizer("") {
+    // Gemma replaces spaces with metaspace before its literal-space Split, so no regex boundaries apply.
     byte_level_bpe = false;
     byte_fallback  = true;
     add_bos_token  = true;
@@ -189,164 +191,3 @@ GemmaTokenizer::GemmaTokenizer(const std::string& merges_utf8_str, const std::st
         load_from_merges(load_gemma_merges(), load_gemma_vocab_json());
     }
 }
-
-std::string Gemma2Tokenizer::normalize(const std::string& text) const {
-    std::string normalized = text;
-    size_t pos             = 0;
-    while ((pos = normalized.find(' ', pos)) != std::string::npos) {
-        normalized.replace(pos, 1, "\xE2\x96\x81");
-        pos += 3;
-    }
-    return normalized;
-}
-
-void Gemma2Tokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
-    nlohmann::json vocab;
-    try {
-        vocab = nlohmann::json::parse(vocab_utf8_str);
-    } catch (const nlohmann::json::parse_error&) {
-        GGML_ABORT("invalid vocab json str");
-    }
-    for (const auto& [key, value] : vocab.items()) {
-        std::u32string token = utf8_to_utf32(key);
-        int i                = value;
-        encoder[token]       = i;
-        decoder[i]           = token;
-    }
-    encoder_len = static_cast<int>(vocab.size());
-    LOG_VERBOSE("vocab size: %d", encoder_len);
-
-    std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
-    std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
-    for (const auto& merge : merges) {
-        size_t space_pos = merge.find(' ');
-        merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
-    }
-    LOG_VERBOSE("merges size %zu", merge_pairs.size());
-
-    int rank = 0;
-    for (const auto& merge : merge_pairs) {
-        bpe_ranks[merge] = rank++;
-    }
-    bpe_len = rank;
-}
-
-Gemma2Tokenizer::Gemma2Tokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
-    byte_level_bpe = false;
-    byte_fallback  = true;
-    add_bos_token  = true;
-    PAD_TOKEN      = "<pad>";
-    EOS_TOKEN      = "<eos>";
-    BOS_TOKEN      = "<bos>";
-    UNK_TOKEN      = "<unk>";
-
-    PAD_TOKEN_ID = 0;
-    EOS_TOKEN_ID = 1;
-    BOS_TOKEN_ID = 2;
-    UNK_TOKEN_ID = 3;
-
-    std::vector<std::string> special_tokens_before_merge = {
-        PAD_TOKEN,
-        EOS_TOKEN,
-        BOS_TOKEN,
-        UNK_TOKEN,
-        "<mask>",
-        "<2mass>",
-        "[@BOS@]",
-    };
-    for (int i = 0; i <= 98; i++) {
-        special_tokens_before_merge.push_back("<unused" + std::to_string(i) + ">");
-    }
-    special_tokens_before_merge.push_back("<start_of_turn>");
-    special_tokens_before_merge.push_back("<end_of_turn>");
-    for (int i = 1; i <= 31; i++) {
-        special_tokens_before_merge.push_back(std::string(i, '\n'));
-    }
-    for (int i = 2; i <= 31; i++) {
-        std::string whitespace_token;
-        for (int j = 0; j < i; j++) {
-            whitespace_token += "\xE2\x96\x81";
-        }
-        special_tokens_before_merge.push_back(whitespace_token);
-    }
-    std::vector<std::string> html_tokens = {
-        "<table>",
-        "<caption>",
-        "<thead>",
-        "<tbody>",
-        "<tfoot>",
-        "<tr>",
-        "<th>",
-        "<td>",
-        "</table>",
-        "</caption>",
-        "</thead>",
-        "</tbody>",
-        "</tfoot>",
-        "</tr>",
-        "</th>",
-        "</td>",
-        "<h1>",
-        "<h2>",
-        "<h3>",
-        "<h4>",
-        "<h5>",
-        "<h6>",
-        "<blockquote>",
-        "</h1>",
-        "</h2>",
-        "</h3>",
-        "</h4>",
-        "</h5>",
-        "</h6>",
-        "</blockquote>",
-        "<strong>",
-        "<em>",
-        "<b>",
-        "<i>",
-        "<u>",
-        "<s>",
-        "<sub>",
-        "<sup>",
-        "<code>",
-        "</strong>",
-        "</em>",
-        "</b>",
-        "</i>",
-        "</u>",
-        "</s>",
-        "</sub>",
-        "</sup>",
-        "</code>",
-    };
-    special_tokens_before_merge.insert(special_tokens_before_merge.end(),
-                                       html_tokens.begin(),
-                                       html_tokens.end());
-    for (int i = 0; i <= 0xFF; i++) {
-        char hex_buf[16];
-        snprintf(hex_buf, sizeof(hex_buf), "<0x%02X>", i);
-        special_tokens_before_merge.push_back(hex_buf);
-    }
-
-    std::vector<std::string> special_tokens_after_merge = {
-        "[toxicity=0]",
-    };
-    for (int i = 1; i <= 31; i++) {
-        special_tokens_after_merge.insert(special_tokens_after_merge.begin() + i - 1,
-                                          std::string(i, '\t'));
-    }
-    for (int i = 99; i <= 99; i++) {
-        special_tokens_after_merge.push_back("<unused" + std::to_string(i) + ">");
-    }
-
-    special_tokens = special_tokens_before_merge;
-    special_tokens.insert(special_tokens.end(),
-                          special_tokens_after_merge.begin(),
-                          special_tokens_after_merge.end());
-
-    if (merges_utf8_str.size() > 0 && vocab_utf8_str.size() > 0) {
-        load_from_merges(merges_utf8_str, vocab_utf8_str);
-    } else {
-        load_from_merges(load_gemma2_merges(), load_gemma2_vocab_json());
-    }
-}

src/tokenizers/gemma_tokenizer.h

@@ -14,13 +14,4 @@ public:
     explicit GemmaTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
 };
 
-class Gemma2Tokenizer : public BPETokenizer {
-protected:
-    void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str);
-    std::string normalize(const std::string& text) const override;
-
-public:
-    explicit Gemma2Tokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
-};
-
 #endif  // __SD_TOKENIZERS_GEMMA_TOKENIZER_H__

src/tokenizers/gpt_oss_tokenizer.cpp

@@ -1,91 +0,0 @@
-#include "gpt_oss_tokenizer.h"
-
-#include "core/util.h"
-#include "json.hpp"
-#include "vocab/vocab.h"
-
-void GPTOSSTokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
-    auto byte_unicode_pairs = bytes_to_unicode();
-    byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
-    for (auto& pair : byte_unicode_pairs) {
-        byte_decoder[pair.second] = pair.first;
-    }
-
-    nlohmann::json vocab;
-    try {
-        vocab = nlohmann::json::parse(vocab_utf8_str);
-    } catch (const nlohmann::json::parse_error&) {
-        GGML_ABORT("invalid vocab json str");
-    }
-    for (const auto& [key, value] : vocab.items()) {
-        std::u32string token = utf8_to_utf32(key);
-        int i                = value;
-        encoder[token]       = i;
-        decoder[i]           = token;
-    }
-    encoder_len = static_cast<int>(encoder.size());
-    for (auto& special_token : special_tokens) {
-        auto token           = utf8_to_utf32(special_token);
-        encoder[token]       = encoder_len;
-        decoder[encoder_len] = token;
-        encoder_len++;
-    }
-    encoder_len = static_cast<int>(encoder.size());
-    LOG_VERBOSE("vocab size: %d", encoder_len);
-
-    std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
-    std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
-    for (const auto& merge : merges) {
-        size_t space_pos = merge.find(' ');
-        merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
-    }
-    LOG_VERBOSE("merges size %zu", merge_pairs.size());
-
-    int rank = 0;
-    for (const auto& merge : merge_pairs) {
-        bpe_ranks[merge] = rank++;
-    }
-    bpe_len = rank;
-}
-
-GPTOSSTokenizer::GPTOSSTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
-    BOS_TOKEN = "<|startoftext|>";
-    UNK_TOKEN = "<|endoftext|>";
-    EOS_TOKEN = "<|endoftext|>";
-    PAD_TOKEN = "<|endoftext|>";
-
-    BOS_TOKEN_ID = 199998;
-    EOS_TOKEN_ID = 199999;
-    UNK_TOKEN_ID = 199999;
-    PAD_TOKEN_ID = 199999;
-
-    special_tokens = {
-        "<|startoftext|>",
-        "<|endoftext|>",
-        "<|reserved_200000|>",
-        "<|reserved_200001|>",
-        "<|return|>",
-        "<|constrain|>",
-        "<|reserved_200004|>",
-        "<|channel|>",
-        "<|start|>",
-        "<|end|>",
-        "<|message|>",
-        "<|reserved_200009|>",
-        "<|reserved_200010|>",
-        "<|reserved_200011|>",
-        "<|call|>",
-        "<|reserved_200013|>",
-        "<|reserved_200014|>",
-        "<|reserved_200015|>",
-        "<|reserved_200016|>",
-        "<|reserved_200017|>",
-        "<|endofprompt|>",
-    };
-
-    if (merges_utf8_str.size() > 0) {
-        load_from_merges(merges_utf8_str, vocab_utf8_str);
-    } else {
-        load_from_merges(load_gpt_oss_merges(), load_gpt_oss_vocab_json());
-    }
-}

src/tokenizers/gpt_oss_tokenizer.h

@@ -1,16 +0,0 @@
-#ifndef __SD_TOKENIZERS_GPT_OSS_TOKENIZER_H__
-#define __SD_TOKENIZERS_GPT_OSS_TOKENIZER_H__
-
-#include <string>
-
-#include "bpe_tokenizer.h"
-
-class GPTOSSTokenizer : public BPETokenizer {
-protected:
-    void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str);
-
-public:
-    explicit GPTOSSTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
-};
-
-#endif  // __SD_TOKENIZERS_GPT_OSS_TOKENIZER_H__

src/tokenizers/hf_tokenizer.cpp

@@ -0,0 +1,856 @@
+#include "hf_tokenizer.h"
+
+#include <algorithm>
+#include <array>
+#include <climits>
+#include <cstdlib>
+#include <fstream>
+#include <queue>
+#include <sstream>
+#include <stdexcept>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "core/regex.h"
+#include "core/util.h"
+#include "json.hpp"
+#include "utf8proc.h"
+
+using TokenizerJSON = nlohmann::json;
+
+static void tokenizer_require(bool condition, const std::string& message) {
+    if (!condition) {
+        throw std::runtime_error("tokenizer.json: " + message);
+    }
+}
+
+static std::string tokenizer_utf8(int32_t codepoint) {
+    utf8proc_uint8_t bytes[4];
+    auto count = utf8proc_encode_char(codepoint, bytes);
+    return std::string(reinterpret_cast<const char*>(bytes), count);
+}
+
+static bool tokenizer_error(std::string* error, const std::string& message) {
+    if (error) {
+        *error = "tokenizer.json: " + message;
+    } else {
+        LOG_ERROR("tokenizer.json: %s", message.c_str());
+    }
+    return false;
+}
+
+static bool tokenizer_next(const std::string& text, size_t& offset, int32_t& codepoint, std::string* error) {
+    auto size = utf8proc_iterate(reinterpret_cast<const utf8proc_uint8_t*>(text.data() + offset), text.size() - offset, &codepoint);
+    if (size <= 0) {
+        return tokenizer_error(error, "invalid UTF-8 input");
+    }
+    offset += size;
+    return true;
+}
+
+static int tokenizer_id(const TokenizerJSON& value) {
+    tokenizer_require(value.is_number_integer(), "token ID must be an integer");
+    auto id = value.get<int64_t>();
+    tokenizer_require(id >= 0 && id <= INT_MAX, "token ID outside int32 range");
+    return static_cast<int>(id);
+}
+
+static uint64_t tokenizer_pair(int left, int right) {
+    return (static_cast<uint64_t>(left) << 32) | static_cast<uint32_t>(right);
+}
+
+struct HFTokenizer::Impl {
+    struct Pattern {
+        std::string literal;
+        std::shared_ptr<sd::Regex> regex;
+
+        explicit Pattern(const TokenizerJSON& config) {
+            tokenizer_require(config.is_object() && config.size() == 1, "invalid String/Regex pattern");
+            if (config.contains("String")) {
+                literal = config.at("String").get<std::string>();
+            } else {
+                tokenizer_require(config.contains("Regex"), "unsupported pattern");
+                regex = std::make_shared<sd::Regex>();
+                std::string error;
+                bool ok = regex->compile(config.at("Regex").get<std::string>(), &error);
+                tokenizer_require(ok, "invalid regex: " + error);
+            }
+        }
+
+        bool matches(const std::string& text, std::vector<sd::Regex::Match>& result, std::string* error) const {
+            result.clear();
+            if (regex) {
+                std::string regex_error;
+                if (!regex->find_matches(text, result, &regex_error)) {
+                    return tokenizer_error(error, "regex search failed: " + regex_error);
+                }
+            } else if (literal.empty()) {
+                size_t offset = 0;
+                for (;;) {
+                    result.emplace_back(offset, offset);
+                    if (offset == text.size()) {
+                        break;
+                    }
+                    int32_t cp;
+                    if (!tokenizer_next(text, offset, cp, error)) {
+                        result.clear();
+                        return false;
+                    }
+                }
+            } else {
+                size_t offset = 0;
+                while ((offset = text.find(literal, offset)) != std::string::npos) {
+                    result.emplace_back(offset, offset + literal.size());
+                    offset += literal.size();
+                }
+            }
+            return true;
+        }
+
+        bool replace(const std::string& text, const std::string& replacement, std::string& result, std::string* error) const {
+            result.clear();
+            size_t offset = 0;
+            std::vector<sd::Regex::Match> found;
+            if (!matches(text, found, error)) {
+                return false;
+            }
+            for (const auto& match : found) {
+                result.append(text, offset, match.first - offset);
+                result += replacement;
+                offset = match.second;
+            }
+            result.append(text, offset, std::string::npos);
+            return true;
+        }
+
+        bool split(const std::string& text, const std::string& behavior, bool invert, std::vector<std::string>& result, std::string* error) const {
+            result.clear();
+            struct Part {
+                size_t start, end;
+                bool matched;
+            };
+            std::vector<Part> parts;
+            size_t offset = 0;
+            std::vector<sd::Regex::Match> found;
+            if (!matches(text, found, error)) {
+                return false;
+            }
+            for (const auto& match : found) {
+                if (match.first > offset) {
+                    parts.push_back({offset, match.first, invert});
+                }
+                parts.push_back({match.first, match.second, !invert});
+                offset = match.second;
+            }
+            if (offset < text.size()) {
+                parts.push_back({offset, text.size(), invert});
+            }
+            if (behavior == "MergedWithNext") {
+                std::reverse(parts.begin(), parts.end());
+            }
+            std::vector<Part> merged;
+            bool previous = false;
+            for (const auto& part : parts) {
+                bool join = (behavior == "Contiguous" && part.matched == previous) ||
+                            ((behavior == "MergedWithPrevious" || behavior == "MergedWithNext") && part.matched && !previous);
+                if (join && !merged.empty()) {
+                    merged.back().start = std::min(merged.back().start, part.start);
+                    merged.back().end   = std::max(merged.back().end, part.end);
+                } else if (behavior != "Removed" || !part.matched) {
+                    merged.push_back(part);
+                }
+                previous = part.matched;
+            }
+            if (behavior == "MergedWithNext") {
+                std::reverse(merged.begin(), merged.end());
+            }
+            for (const auto& part : merged) {
+                if (part.start != part.end) {
+                    result.push_back(text.substr(part.start, part.end - part.start));
+                }
+            }
+            return true;
+        }
+    };
+
+    struct Step {
+        std::string type, content, behavior;
+        std::shared_ptr<Pattern> pattern;
+        bool invert = false, prefix_space = false;
+    };
+
+    struct Trie {
+        struct Node {
+            std::unordered_map<unsigned char, size_t> children;
+            int id = -1;
+        };
+        std::vector<Node> nodes{1};
+
+        void add(const std::string& text, int id) {
+            size_t index = 0;
+            for (unsigned char c : text) {
+                auto found = nodes[index].children.find(c);
+                if (found == nodes[index].children.end()) {
+                    size_t next = nodes.size();
+                    nodes[index].children.emplace(c, next);
+                    nodes.emplace_back();
+                    index = next;
+                } else {
+                    index = found->second;
+                }
+            }
+            nodes[index].id = id;
+        }
+
+        std::pair<size_t, int> match(const std::string& text, size_t start) const {
+            size_t index = 0;
+            std::pair<size_t, int> result{start, -1};
+            for (size_t end = start; end < text.size(); ++end) {
+                auto found = nodes[index].children.find(static_cast<unsigned char>(text[end]));
+                if (found == nodes[index].children.end()) {
+                    break;
+                }
+                index = found->second;
+                if (nodes[index].id >= 0) {
+                    result = {end + 1, nodes[index].id};
+                }
+            }
+            return result;
+        }
+    };
+
+    struct Merge {
+        size_t rank;
+        int id;
+    };
+    std::unordered_map<std::string, int> vocab;
+    std::unordered_map<std::string, int> added_vocab;
+    std::unordered_map<int, std::string> tokens;
+    std::unordered_map<uint64_t, Merge> merges;
+    std::unordered_set<int> special_ids;
+    std::vector<std::string> custom_tokens;
+    std::vector<Step> normalizers, pre_tokenizers, decoders;
+    Trie raw_added, normalized_added;
+    std::array<std::string, 256> byte_encoder;
+    std::unordered_map<int32_t, unsigned char> byte_decoder;
+    std::string suffix;
+    int unk       = -1;
+    bool fuse_unk = false, byte_fallback = false, ignore_merges = false, has_decoder = false;
+
+    Impl() {
+        int extra = 256;
+        for (int byte = 0; byte < 256; ++byte) {
+            int cp             = ((byte >= 33 && byte <= 126) || (byte >= 161 && byte <= 172) || byte >= 174) ? byte : extra++;
+            byte_encoder[byte] = tokenizer_utf8(cp);
+            byte_decoder[cp]   = static_cast<unsigned char>(byte);
+        }
+    }
+
+    static void parse_steps(const TokenizerJSON& config, const std::string& stage, std::vector<Step>& out, int depth = 0) {
+        tokenizer_require(depth < 32, stage + " nesting is too deep");
+        if (config.is_null()) {
+            return;
+        }
+        Step step;
+        step.type = config.at("type").get<std::string>();
+        if (step.type == "Sequence") {
+            const char* key = stage == "normalizer" ? "normalizers" : stage == "pre_tokenizer" ? "pretokenizers"
+                                                                                               : "decoders";
+            for (const auto& child : config.at(key)) {
+                parse_steps(child, stage, out, depth + 1);
+            }
+            return;
+        }
+        if ((stage == "normalizer" || stage == "decoder") && step.type == "Replace") {
+            step.pattern = std::make_shared<Pattern>(config.at("pattern"));
+            step.content = config.at("content").get<std::string>();
+        } else if (stage == "normalizer" && (step.type == "NFC" || step.type == "Lowercase")) {
+        } else if (stage == "pre_tokenizer" && step.type == "Split") {
+            step.pattern  = std::make_shared<Pattern>(config.at("pattern"));
+            step.behavior = config.at("behavior").get<std::string>();
+            tokenizer_require(step.behavior == "Removed" || step.behavior == "Isolated" || step.behavior == "Contiguous" || step.behavior == "MergedWithPrevious" || step.behavior == "MergedWithNext", "unsupported Split behavior: " + step.behavior);
+            step.invert = config.value("invert", false);
+        } else if ((stage == "pre_tokenizer" || stage == "decoder") && step.type == "ByteLevel") {
+            step.prefix_space = config.value("add_prefix_space", true);
+            if (stage == "pre_tokenizer" && config.value("use_regex", true)) {
+                step.pattern = std::make_shared<Pattern>(TokenizerJSON{{"Regex", R"('s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+)"}});
+            }
+        } else if (stage == "decoder" && (step.type == "ByteFallback" || step.type == "Fuse")) {
+        } else {
+            tokenizer_require(false, "unsupported " + stage + ": " + step.type);
+        }
+        out.push_back(std::move(step));
+    }
+
+    bool normalize(std::string text, std::string& result, std::string* error) const {
+        result.clear();
+        for (const auto& step : normalizers) {
+            if (step.type == "Replace") {
+                std::string replaced;
+                if (!step.pattern->replace(text, step.content, replaced, error)) {
+                    return false;
+                }
+                text = std::move(replaced);
+            } else if (step.type == "NFC") {
+                utf8proc_uint8_t* output = nullptr;
+                auto size                = utf8proc_map(reinterpret_cast<const utf8proc_uint8_t*>(text.data()), text.size(), &output, static_cast<utf8proc_option_t>(UTF8PROC_STABLE | UTF8PROC_COMPOSE));
+                std::unique_ptr<utf8proc_uint8_t, decltype(&std::free)> buffer(output, &std::free);
+                if (size < 0) {
+                    return tokenizer_error(error, std::string("NFC normalization failed: ") + utf8proc_errmsg(size));
+                }
+                text.assign(reinterpret_cast<const char*>(output), size);
+            } else {
+                std::string lower;
+                for (size_t i = 0; i < text.size();) {
+                    int32_t cp;
+                    if (!tokenizer_next(text, i, cp, error)) {
+                        return false;
+                    }
+                    // Rust char::to_lowercase uses full, context-free lowercase. U+0130 expands.
+                    lower += cp == 0x130 ? "i\xcc\x87" : tokenizer_utf8(utf8proc_tolower(cp));
+                }
+                text = std::move(lower);
+            }
+        }
+        result = std::move(text);
+        return true;
+    }
+
+    bool pre_tokenize(const std::string& text, std::vector<std::string>& result, std::string* error) const {
+        result.clear();
+        std::vector<std::string> pieces{text};
+        for (const auto& step : pre_tokenizers) {
+            std::vector<std::string> next;
+            for (auto piece : pieces) {
+                if (piece.empty()) {
+                    continue;
+                }
+                std::vector<std::string> split;
+                if (step.type == "Split") {
+                    if (!step.pattern->split(piece, step.behavior, step.invert, split, error)) {
+                        return false;
+                    }
+                    next.insert(next.end(), split.begin(), split.end());
+                } else {
+                    if (step.prefix_space && piece.front() != ' ') {
+                        piece.insert(piece.begin(), ' ');
+                    }
+                    if (step.pattern) {
+                        if (!step.pattern->split(piece, "Isolated", false, split, error)) {
+                            return false;
+                        }
+                    } else {
+                        split.push_back(piece);
+                    }
+                    for (const auto& part : split) {
+                        std::string encoded;
+                        for (unsigned char byte : part) {
+                            encoded += byte_encoder[byte];
+                        }
+                        next.push_back(std::move(encoded));
+                    }
+                }
+            }
+            pieces = std::move(next);
+        }
+        result = std::move(pieces);
+        return true;
+    }
+
+    bool bpe(const std::string& text, std::vector<int>& ids, std::string* error) const {
+        ids.clear();
+        if (ignore_merges) {
+            auto found = vocab.find(text);
+            if (found != vocab.end()) {
+                ids.push_back(found->second);
+                return true;
+            }
+        }
+        bool pending_unk = false;
+        for (size_t i = 0; i < text.size();) {
+            int32_t cp;
+            size_t end = i;
+            if (!tokenizer_next(text, end, cp, error)) {
+                ids.clear();
+                return false;
+            }
+            std::string symbol = text.substr(i, end - i);
+            if (end == text.size()) {
+                symbol += suffix;
+            }
+            i          = end;
+            auto found = vocab.find(symbol);
+            if (found != vocab.end()) {
+                if (pending_unk) {
+                    ids.push_back(unk);
+                    pending_unk = false;
+                }
+                ids.push_back(found->second);
+                continue;
+            }
+            if (byte_fallback) {
+                std::vector<int> bytes;
+                for (unsigned char byte : symbol) {
+                    const char* hex   = "0123456789ABCDEF";
+                    std::string token = "<0x00>";
+                    token[3]          = hex[byte >> 4];
+                    token[4]          = hex[byte & 15];
+                    auto fallback     = vocab.find(token);
+                    if (fallback == vocab.end()) {
+                        break;
+                    }
+                    bytes.push_back(fallback->second);
+                }
+                if (bytes.size() == symbol.size()) {
+                    ids.insert(ids.end(), bytes.begin(), bytes.end());
+                    continue;
+                }
+            }
+            if (unk >= 0) {
+                if (pending_unk && !fuse_unk) {
+                    ids.push_back(unk);
+                }
+                pending_unk = true;
+            }
+        }
+        if (pending_unk) {
+            ids.push_back(unk);
+        }
+        struct Symbol {
+            int id;
+            size_t prev, next, generation = 0;
+            bool alive = true;
+        };
+        struct Candidate {
+            size_t rank, left, right, left_generation, right_generation;
+            int id;
+            bool operator<(const Candidate& other) const {
+                return rank != other.rank ? rank > other.rank : left > other.left;
+            }
+        };
+        const size_t none = ids.size();
+        std::vector<Symbol> symbols;
+        for (size_t i = 0; i < ids.size(); ++i) {
+            symbols.push_back({ids[i], i == 0 ? none : i - 1, i + 1});
+        }
+        std::priority_queue<Candidate> queue;
+        auto push = [&](size_t left) {
+            if (left == none || symbols[left].next == none) {
+                return;
+            }
+            size_t right = symbols[left].next;
+            auto found   = merges.find(tokenizer_pair(symbols[left].id, symbols[right].id));
+            if (found != merges.end()) {
+                queue.push({found->second.rank, left, right, symbols[left].generation, symbols[right].generation, found->second.id});
+            }
+        };
+        for (size_t i = 0; i < symbols.size(); ++i) {
+            push(i);
+        }
+        while (!queue.empty()) {
+            Candidate item = queue.top();
+            queue.pop();
+            auto& left  = symbols[item.left];
+            auto& right = symbols[item.right];
+            if (!left.alive || !right.alive || left.next != item.right || left.generation != item.left_generation || right.generation != item.right_generation) {
+                continue;
+            }
+            left.id   = item.id;
+            left.next = right.next;
+            ++left.generation;
+            right.alive = false;
+            if (left.next != none) {
+                symbols[left.next].prev = item.left;
+            }
+            push(left.prev);
+            push(item.left);
+        }
+        ids.clear();
+        for (const auto& symbol : symbols) {
+            if (symbol.alive) {
+                ids.push_back(symbol.id);
+            }
+        }
+        return true;
+    }
+
+    int lookup(const std::string& token) const {
+        auto added = added_vocab.find(token);
+        if (added != added_vocab.end()) {
+            return added->second;
+        }
+        auto found = vocab.find(token);
+        return found == vocab.end() ? -1 : found->second;
+    }
+
+    void add_token(const std::string& token, int id, bool added = false) {
+        auto old = tokens.find(id);
+        tokenizer_require(old == tokens.end() || old->second == token, "conflicting token ID " + std::to_string(id));
+        int old_id = lookup(token);
+        tokenizer_require(old_id < 0 || old_id == id, "conflicting ID for token " + token);
+        tokens[id]                           = token;
+        (added ? added_vocab : vocab)[token] = id;
+    }
+};
+
+HFTokenizer::HFTokenizer(const std::string& path)
+    : impl_(new Impl) {
+    std::ifstream stream(path, std::ios::binary);
+    tokenizer_require(stream.good(), "cannot open " + path);
+    TokenizerJSON config;
+    stream >> config;
+    tokenizer_require(config.value("version", std::string("1.0")) == "1.0", "unsupported version");
+    tokenizer_require(config.value("padding", TokenizerJSON()).is_null(), "JSON padding is unsupported; padding is controlled by the text encoder");
+    tokenizer_require(config.value("truncation", TokenizerJSON()).is_null(), "JSON truncation is unsupported; truncation is controlled by the text encoder");
+    const auto& model = config.at("model");
+    tokenizer_require(model.at("type") == "BPE", "only BPE models are supported");
+    tokenizer_require(model.value("dropout", TokenizerJSON()).is_null() || model.at("dropout") == 0, "BPE dropout is unsupported");
+    const auto& prefix = model.value("continuing_subword_prefix", TokenizerJSON());
+    tokenizer_require(prefix.is_null() || prefix == "", "nonempty continuing_subword_prefix is unsupported");
+    const auto& suffix   = model.value("end_of_word_suffix", TokenizerJSON());
+    impl_->suffix        = suffix.is_null() ? "" : suffix.get<std::string>();
+    impl_->fuse_unk      = model.value("fuse_unk", false);
+    impl_->byte_fallback = model.value("byte_fallback", false);
+    impl_->ignore_merges = model.value("ignore_merges", false);
+    tokenizer_require(model.at("vocab").is_object(), "BPE vocab must be an object");
+    impl_->vocab.reserve(model.at("vocab").size());
+    impl_->tokens.reserve(model.at("vocab").size());
+    for (const auto& entry : model.at("vocab").items()) {
+        impl_->add_token(entry.key(), tokenizer_id(entry.value()));
+    }
+    if (!model.value("unk_token", TokenizerJSON()).is_null()) {
+        UNK_TOKEN  = model.at("unk_token").get<std::string>();
+        impl_->unk = impl_->lookup(UNK_TOKEN);
+        tokenizer_require(impl_->unk >= 0, "unk_token is absent from vocab");
+    }
+    UNK_TOKEN_ID = impl_->unk;
+    tokenizer_require(model.at("merges").is_array(), "BPE merges must be an array");
+    impl_->merges.reserve(model.at("merges").size());
+    size_t rank = 0;
+    for (const auto& merge : model.at("merges")) {
+        std::string left, right;
+        if (merge.is_string()) {
+            auto value = merge.get<std::string>();
+            auto space = value.find(' ');
+            tokenizer_require(space != std::string::npos && value.find(' ', space + 1) == std::string::npos, "invalid legacy BPE merge");
+            left  = value.substr(0, space);
+            right = value.substr(space + 1);
+        } else {
+            tokenizer_require(merge.is_array() && merge.size() == 2, "BPE merge must contain two tokens");
+            left  = merge.at(0).get<std::string>();
+            right = merge.at(1).get<std::string>();
+        }
+        int a = impl_->lookup(left), b = impl_->lookup(right), id = impl_->lookup(left + right);
+        tokenizer_require(a >= 0 && b >= 0 && id >= 0, "BPE merge references a missing vocab token");
+        impl_->merges[tokenizer_pair(a, b)] = {rank++, id};
+    }
+    Impl::parse_steps(config.value("normalizer", TokenizerJSON()), "normalizer", impl_->normalizers);
+    Impl::parse_steps(config.value("pre_tokenizer", TokenizerJSON()), "pre_tokenizer", impl_->pre_tokenizers);
+    impl_->has_decoder = !config.value("decoder", TokenizerJSON()).is_null();
+    Impl::parse_steps(config.value("decoder", TokenizerJSON()), "decoder", impl_->decoders);
+    size_t next_added_id = impl_->vocab.size();
+    for (const auto& token : config.value("added_tokens", TokenizerJSON::array())) {
+        for (const char* flag : {"single_word", "lstrip", "rstrip"}) {
+            tokenizer_require(!token.value(flag, false), std::string("added_tokens.") + flag + "=true is unsupported");
+        }
+        auto content = token.at("content").get<std::string>();
+        tokenizer_require(!content.empty(), "empty added token is unsupported");
+        int id = tokenizer_id(token.at("id"));
+        if (impl_->lookup(content) < 0) {
+            tokenizer_require(static_cast<size_t>(id) == next_added_id++, "nonconsecutive added token IDs would be reassigned by Hugging Face tokenizers");
+        }
+        impl_->add_token(content, id, true);
+        if (token.value("special", false)) {
+            special_tokens.push_back(content);
+            impl_->special_ids.insert(id);
+        }
+        bool normalized     = token.value("normalized", true);
+        std::string pattern = content;
+        if (normalized) {
+            std::string error;
+            bool ok = impl_->normalize(content, pattern, &error);
+            tokenizer_require(ok, error);
+        }
+        tokenizer_require(!pattern.empty(), "added token normalizes to an empty string");
+        (normalized ? impl_->normalized_added : impl_->raw_added).add(pattern, id);
+    }
+    const auto& processor = config.value("post_processor", TokenizerJSON());
+    BOS_TOKEN_ID = EOS_TOKEN_ID = -1;
+    if (!processor.is_null()) {
+        auto type    = processor.at("type").get<std::string>();
+        auto special = [&](const TokenizerJSON& pair) {
+            tokenizer_require(pair.is_array() && pair.size() == 2, "invalid postprocessor special token");
+            int id = tokenizer_id(pair.at(1));
+            tokenizer_require(impl_->lookup(pair.at(0).get<std::string>()) == id, "postprocessor token/ID does not match vocab");
+            return id;
+        };
+        if (type == "RobertaProcessing") {
+            BOS_TOKEN_ID = special(processor.at("cls"));
+            EOS_TOKEN_ID = special(processor.at("sep"));
+        } else if (type == "TemplateProcessing") {
+            bool seen_sequence = false;
+            for (const auto& item : processor.at("single")) {
+                if (item.contains("Sequence")) {
+                    tokenizer_require(!seen_sequence && item.at("Sequence").at("id") == "A", "single template must contain exactly one sequence A");
+                    seen_sequence = true;
+                } else {
+                    auto name         = item.at("SpecialToken").at("id").get<std::string>();
+                    const auto& token = processor.at("special_tokens").at(name);
+                    tokenizer_require(token.at("ids").size() == 1 && token.at("tokens").size() == 1, "multi-ID template special tokens are unsupported");
+                    int id      = special(TokenizerJSON::array({token.at("tokens").at(0), token.at("ids").at(0)}));
+                    int& target = seen_sequence ? EOS_TOKEN_ID : BOS_TOKEN_ID;
+                    tokenizer_require(target < 0, "single template supports at most one prefix and one suffix token");
+                    target = id;
+                }
+            }
+            tokenizer_require(seen_sequence, "single template has no sequence A");
+        } else {
+            tokenizer_require(type == "ByteLevel", "unsupported post_processor: " + type);
+        }
+    }
+    add_bos_token = BOS_TOKEN_ID >= 0;
+    add_eos_token = EOS_TOKEN_ID >= 0;
+    BOS_TOKEN     = decode_token(BOS_TOKEN_ID);
+    EOS_TOKEN     = decode_token(EOS_TOKEN_ID);
+    set_padding(0, false);
+}
+
+HFTokenizer::~HFTokenizer() = default;
+
+void HFTokenizer::set_padding(int token_id, bool left) {
+    PAD_TOKEN_ID = token_id;
+    PAD_TOKEN    = decode_token(token_id);
+    pad_left     = left;
+}
+
+void HFTokenizer::validate_vocab_size(int64_t embedding_rows) const {
+    tokenizer_require(embedding_rows > 0, "text encoder has no token embedding rows");
+    for (const auto& token : impl_->tokens) {
+        tokenizer_require(token.first < embedding_rows, "token ID " + std::to_string(token.first) + " exceeds text encoder vocabulary (" + std::to_string(embedding_rows) + ")");
+    }
+    tokenizer_require(PAD_TOKEN_ID >= 0 && PAD_TOKEN_ID < embedding_rows, "padding ID exceeds text encoder vocabulary");
+}
+
+int HFTokenizer::token_to_id(const std::string& token) const {
+    return impl_->lookup(token);
+}
+
+void HFTokenizer::add_special_token(const std::string& token) {
+    Tokenizer::add_special_token(token);
+    if (!token.empty()) {
+        impl_->custom_tokens.push_back(token);
+    }
+}
+
+bool HFTokenizer::encode(const std::string& text, std::vector<int>& tokens, on_new_token_cb_t callback, std::string* error) {
+    tokens.clear();
+    if (error) {
+        error->clear();
+    }
+    for (size_t i = 0; i < text.size();) {
+        int32_t cp;
+        if (!tokenizer_next(text, i, cp, error)) {
+            return false;
+        }
+    }
+    std::vector<int> result;
+    Impl::Trie raw_custom, normalized_custom;
+    if (callback) {
+        for (size_t index = 0; index < impl_->custom_tokens.size(); ++index) {
+            const auto& token = impl_->custom_tokens[index];
+            raw_custom.add(token, static_cast<int>(index));
+            std::string normalized;
+            if (!impl_->normalize(token, normalized, error)) {
+                return false;
+            }
+            if (!normalized.empty()) {
+                normalized_custom.add(normalized, static_cast<int>(index));
+            }
+        }
+    }
+    auto encode_plain = [&](const std::string& value) {
+        std::vector<std::string> pieces;
+        if (!impl_->pre_tokenize(value, pieces, error)) {
+            return false;
+        }
+        for (auto& piece : pieces) {
+            if (callback && callback(piece, result)) {
+                continue;
+            }
+            std::vector<int> ids;
+            if (!impl_->bpe(piece, ids, error)) {
+                return false;
+            }
+            result.insert(result.end(), ids.begin(), ids.end());
+        }
+        return true;
+    };
+    auto extract = [&](const std::string& value, const Impl::Trie& added, const Impl::Trie& custom_tokens, const auto& encode_gap) {
+        size_t start = 0, i = 0;
+        while (i < value.size()) {
+            auto match      = added.match(value, i);
+            auto custom     = custom_tokens.match(value, i);
+            bool use_custom = custom.second >= 0 && custom.first >= match.first;
+            if (match.second < 0 && !use_custom) {
+                ++i;
+                continue;
+            }
+            if (!encode_gap(value.substr(start, i - start))) {
+                return false;
+            }
+            if (use_custom) {
+                auto token = impl_->custom_tokens[custom.second];
+                if (!callback(token, result)) {
+                    if (match.second >= 0 && match.first == custom.first) {
+                        result.push_back(match.second);
+                    } else if (!encode_gap(value.substr(i, custom.first - i))) {
+                        return false;
+                    }
+                }
+            } else {
+                result.push_back(match.second);
+            }
+            start = i = use_custom ? custom.first : match.first;
+        }
+        return encode_gap(value.substr(start));
+    };
+    auto encode_normalized = [&](const std::string& value) {
+        std::string normalized;
+        if (!impl_->normalize(value, normalized, error)) {
+            return false;
+        }
+        return extract(normalized, impl_->normalized_added, normalized_custom, encode_plain);
+    };
+    if (!extract(text, impl_->raw_added, raw_custom, encode_normalized)) {
+        return false;
+    }
+    std::stringstream ss;
+    ss << "[";
+    for (int id : result) {
+        auto token = impl_->tokens.find(id);
+        if (token != impl_->tokens.end()) {
+            ss << "\"" << token->second << "\", ";
+        } else {
+            ss << "\"<id:" << id << ">\", ";
+        }
+    }
+    ss << "]";
+    LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", text.c_str(), result.size(), ss.str().c_str());
+    tokens = std::move(result);
+    return true;
+}
+
+std::string HFTokenizer::decode_token(int id) const {
+    auto found = impl_->tokens.find(id);
+    return found == impl_->tokens.end() ? "" : found->second;
+}
+
+static std::string tokenizer_lossy_utf8(const std::string& bytes, bool fallback) {
+    std::string result;
+    for (size_t i = 0; i < bytes.size();) {
+        int32_t cp;
+        auto count = utf8proc_iterate(reinterpret_cast<const utf8proc_uint8_t*>(bytes.data() + i), bytes.size() - i, &cp);
+        if (count > 0) {
+            result.append(bytes, i, count);
+            i += count;
+        } else if (fallback) {
+            result.clear();
+            for (size_t j = 0; j < bytes.size(); ++j) {
+                result += "\xef\xbf\xbd";
+            }
+            return result;
+        } else {
+            result += "\xef\xbf\xbd";
+            unsigned char lead = bytes[i++];
+            size_t expected    = lead >= 0xc2 && lead <= 0xdf ? 2 : lead >= 0xe0 && lead <= 0xef ? 3
+                                                                : lead >= 0xf0 && lead <= 0xf4   ? 4
+                                                                                                 : 1;
+            for (size_t j = 1; j < expected && i < bytes.size(); ++j) {
+                unsigned char c = bytes[i];
+                if (c < 0x80 || c > 0xbf || (j == 1 && ((lead == 0xe0 && c < 0xa0) || (lead == 0xed && c > 0x9f) || (lead == 0xf0 && c < 0x90) || (lead == 0xf4 && c > 0x8f)))) {
+                    break;
+                }
+                ++i;
+            }
+        }
+    }
+    return result;
+}
+
+bool HFTokenizer::decode(const std::vector<int>& ids, std::string& text, std::string* error) const {
+    text.clear();
+    if (error) {
+        error->clear();
+    }
+    std::vector<std::string> pieces;
+    for (int id : ids) {
+        if (!impl_->special_ids.count(id) && impl_->tokens.count(id)) {
+            pieces.push_back(decode_token(id));
+        }
+    }
+    for (const auto& step : impl_->decoders) {
+        if (step.type == "Replace") {
+            for (auto& piece : pieces) {
+                std::string replaced;
+                if (!step.pattern->replace(piece, step.content, replaced, error)) {
+                    return false;
+                }
+                piece = std::move(replaced);
+            }
+        } else if (step.type == "ByteLevel" || step.type == "Fuse") {
+            std::string joined;
+            for (const auto& piece : pieces) {
+                std::string bytes;
+                if (step.type == "ByteLevel") {
+                    for (size_t i = 0; i < piece.size();) {
+                        int32_t cp;
+                        if (!tokenizer_next(piece, i, cp, error)) {
+                            return false;
+                        }
+                        auto found = impl_->byte_decoder.find(cp);
+                        if (found == impl_->byte_decoder.end()) {
+                            bytes = piece;
+                            break;
+                        }
+                        bytes += static_cast<char>(found->second);
+                    }
+                } else {
+                    bytes = piece;
+                }
+                joined += bytes;
+            }
+            pieces = {step.type == "ByteLevel" ? tokenizer_lossy_utf8(joined, false) : joined};
+        } else {
+            std::vector<std::string> decoded;
+            std::string bytes;
+            auto flush = [&] {
+                if (!bytes.empty()) {
+                    decoded.push_back(tokenizer_lossy_utf8(bytes, true));
+                    bytes.clear();
+                }
+            };
+            for (const auto& piece : pieces) {
+                auto hex = [](char c) { return c >= '0' && c <= '9' ? c - '0' : c >= 'A' && c <= 'F' ? c - 'A' + 10
+                                                                            : c >= 'a' && c <= 'f'   ? c - 'a' + 10
+                                                                                                     : -1; };
+                if (piece.size() == 6 && piece.compare(0, 3, "<0x") == 0 && piece[5] == '>' && hex(piece[3]) >= 0 && hex(piece[4]) >= 0) {
+                    bytes += static_cast<char>((hex(piece[3]) << 4) | hex(piece[4]));
+                } else {
+                    flush();
+                    decoded.push_back(piece);
+                }
+            }
+            flush();
+            pieces = std::move(decoded);
+        }
+    }
+    std::string result;
+    for (size_t i = 0; i < pieces.size(); ++i) {
+        if (i && !impl_->has_decoder) {
+            result += ' ';
+        }
+        result += pieces[i];
+    }
+    text = std::move(result);
+    return true;
+}

src/tokenizers/hf_tokenizer.h

@@ -0,0 +1,26 @@
+#ifndef __SD_TOKENIZERS_HF_TOKENIZER_H__
+#define __SD_TOKENIZERS_HF_TOKENIZER_H__
+
+#include <memory>
+
+#include "tokenizer.h"
+
+class HFTokenizer : public Tokenizer {
+    struct Impl;
+    std::unique_ptr<Impl> impl_;
+    std::string decode_token(int token_id) const override;
+
+public:
+    explicit HFTokenizer(const std::string& path);
+    ~HFTokenizer() override;
+
+    // Padding belongs to the encoder; tokenizer.json supplies the single-sequence template.
+    void set_padding(int token_id, bool left);
+    void validate_vocab_size(int64_t embedding_rows) const;
+    int token_to_id(const std::string& token) const;
+    void add_special_token(const std::string& token) override;
+    bool encode(const std::string& text, std::vector<int>& tokens, on_new_token_cb_t on_new_token_cb = nullptr, std::string* error = nullptr) override;
+    bool decode(const std::vector<int>& tokens, std::string& text, std::string* error = nullptr) const override;
+};
+
+#endif  // __SD_TOKENIZERS_HF_TOKENIZER_H__

src/tokenizers/mistral_tokenizer.cpp

@@ -42,7 +42,8 @@ void MistralTokenizer::load_from_merges(const std::string& merges_utf8_str, cons
     bpe_len = rank;
 }
 
-MistralTokenizer::MistralTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+MistralTokenizer::MistralTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str)
+    : BPETokenizer(R"([^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+)") {
     add_bos_token = true;
 
     UNK_TOKEN = "<unk>";

src/tokenizers/qwen2_tokenizer.cpp

@@ -87,7 +87,8 @@ Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str)
 }
 
 Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str,
-                               const std::vector<std::string>& special_tokens_override) {
+                               const std::vector<std::string>& special_tokens_override)
+    : BPETokenizer(R"((?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+)") {
     UNK_TOKEN = "<|endoftext|>";
     EOS_TOKEN = "<|endoftext|>";
     PAD_TOKEN = "<|endoftext|>";

src/tokenizers/t5_unigram_tokenizer.cpp

@@ -287,7 +287,11 @@ std::string T5UniGramTokenizer::normalize(const std::string& input) const {
     return normalized;
 }
 
-std::vector<int> T5UniGramTokenizer::encode(const std::string& input, on_new_token_cb_t on_new_token_cb) {
+bool T5UniGramTokenizer::encode(const std::string& input, std::vector<int>& result, on_new_token_cb_t on_new_token_cb, std::string* error) {
+    result.clear();
+    if (error) {
+        error->clear();
+    }
     std::vector<int32_t> tokens;
     std::vector<std::string> token_strs;
     std::string normalized = normalize(input);
@@ -335,5 +339,6 @@ std::vector<int> T5UniGramTokenizer::encode(const std::string& input, on_new_tok
     ss << "]";
     LOG_VERBOSE("split prompt \"%s\" to tokens %s", input.c_str(), ss.str().c_str());
 
-    return tokens;
+    result = std::move(tokens);
+    return true;
 }

src/tokenizers/t5_unigram_tokenizer.h

@@ -64,7 +64,7 @@ public:
     explicit T5UniGramTokenizer(bool is_umt5 = false);
     ~T5UniGramTokenizer();
 
-    std::vector<int> encode(const std::string& input, on_new_token_cb_t on_new_token_cb = nullptr) override;
+    bool encode(const std::string& input, std::vector<int>& tokens, on_new_token_cb_t on_new_token_cb = nullptr, std::string* error = nullptr) override;
 };
 
 #endif  // __SD_TOKENIZERS_T5_UNIGRAM_TOKENIZER_H__

src/tokenizers/tokenize_util.h

@@ -4,7 +4,6 @@
 #include <string>
 #include <vector>
 
-std::vector<std::string> token_split(const std::string& text);
 std::vector<std::string> split_with_special_tokens(const std::string& text, const std::vector<std::string>& special_tokens);
 
-#endif  // __SD_TOKENIZERS_BPE_TOKENIZE_UTIL_H__
+#endif  // __SD_TOKENIZERS_BPE_TOKENIZE_UTIL_H__

src/tokenizers/tokenizer.cpp

@@ -23,17 +23,22 @@ std::string Tokenizer::normalize(const std::string& text) const {
     return text;
 }
 
-std::vector<int> Tokenizer::tokenize(const std::string& text,
-                                     on_new_token_cb_t on_new_token_cb,
-                                     bool padding,
-                                     size_t min_length,
-                                     size_t max_length,
-                                     bool allow_overflow_expand) {
-    std::vector<int> tokens = encode(text, on_new_token_cb);
+bool Tokenizer::tokenize(const std::string& text,
+                         std::vector<int>& tokens,
+                         on_new_token_cb_t on_new_token_cb,
+                         bool padding,
+                         size_t min_length,
+                         size_t max_length,
+                         bool allow_overflow_expand,
+                         std::string* error) {
+    if (!encode(text, tokens, on_new_token_cb, error)) {
+        tokens.clear();
+        return false;
+    }
     if (padding) {
         pad_tokens(tokens, nullptr, nullptr, min_length, max_length, allow_overflow_expand);
     }
-    return tokens;
+    return true;
 }
 
 void Tokenizer::pad_tokens(std::vector<int>& tokens,
@@ -200,8 +205,11 @@ static std::string clean_up_tokenization(std::string& text) {
     return std::regex_replace(text, pattern, ",");
 }
 
-std::string Tokenizer::decode(const std::vector<int>& tokens) const {
-    std::string text;
+bool Tokenizer::decode(const std::vector<int>& tokens, std::string& text, std::string* error) const {
+    text.clear();
+    if (error) {
+        error->clear();
+    }
 
     for (int token_id : tokens) {
         if (token_id == BOS_TOKEN_ID || token_id == EOS_TOKEN_ID || token_id == PAD_TOKEN_ID) {
@@ -218,5 +226,6 @@ std::string Tokenizer::decode(const std::vector<int>& tokens) const {
     }
 
     text = clean_up_tokenization(text);
-    return trim(text);
+    text = trim(text);
+    return true;
 }