feat: add native CUDA Sol-Attn support

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

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -6,7 +6,9 @@ body:
   - type: markdown
     attributes:
       value: |
-        Please use this template and include as many details as possible to help us reproduce and fix the issue.
+        Before submitting a bug report, please read the [Troubleshooting guide](https://github.com/leejet/stable-diffusion.cpp/blob/master/docs/troubleshooting.md) and try the steps relevant to your problem.
+
+        If the problem persists, complete this form and include what you tried and the results, along with enough details to help us reproduce and fix the issue.
   - type: textarea
     id: commit
     attributes:

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,4 @@
+contact_links:
+  - name: Troubleshooting
+    url: https://github.com/leejet/stable-diffusion.cpp/blob/master/docs/troubleshooting.md
+    about: Read the troubleshooting guide first. If the problem persists, submit a bug report.

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)
@@ -294,6 +296,8 @@ endif()
 
 if(MSVC)
     target_compile_options(${SD_LIB} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/bigobj>)
+    # ggml backends can throw C++ exceptions through their C API.
+    target_compile_options(${SD_LIB} PRIVATE $<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/EHsc->)
 endif()
 
 if(APPLE)
@@ -323,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)
@@ -148,6 +152,7 @@ For runtime and parameter backend placement, see the [backend selection guide](.
 
 ## More Guides
 
+- [Troubleshooting](./docs/troubleshooting.md)
 - [Backend selection](./docs/backend.md)
 - [RPC](./docs/rpc.md)
 - [LoRA](./docs/lora.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

@@ -5,7 +5,8 @@
 - `--backend` selects the runtime backend used to execute model graphs.
 - `--params-backend` selects where model parameters are kept.
 
-If `--params-backend` is not set, parameters use the same backend as their module runtime backend.
+If `--params-backend` is not set, auto-fit chooses parameter placement. With
+`--auto-fit off`, parameters use the same backend as their module runtime backend.
 
 ## Syntax
 
@@ -129,17 +130,21 @@ warning.
 ## Automatic placement (`--auto-fit on|off`)
 
 `--auto-fit` requires `on` or `off` and defaults to `on` when omitted.
-Explicit `--backend` or `--params-backend` assignments disable auto-fit,
+Explicit `--params-backend` assignments disable auto-fit,
 regardless of argument order, even with `--auto-fit on`.
 
-When enabled, auto-fit uses one GPU for `diffusion` / `te` / `vae` computation. It chooses
-the GPU with the largest available memory budget (the first device on a tie),
-then derives parameter placements from the model metadata and the remaining
-memory budgets. The chosen backend specifications are printed.
+Auto-fit preserves explicit `--backend` assignments, including per-module
+assignments and device lists. For modules without a runtime assignment, it chooses
+the GPU with the largest available memory budget (the first device on a tie).
+It then derives parameter placements from the model metadata, each module's
+compute devices, and the remaining memory budgets. The chosen backend
+specifications are printed.
 
 ```shell
 sd-cli -m model.safetensors -p "a cat" --auto-fit on
 sd-cli -m model.safetensors -p "a cat" --auto-fit on --max-vram cuda0=8,cuda1=14
+sd-cli -m model.safetensors -p "a cat" --backend cuda0
+sd-cli -m model.safetensors -p "a cat" --backend diffusion=cuda0,te=cpu,vae=cuda1
 sd-cli -m model.safetensors -p "a cat" --auto-fit off
 ```
 
@@ -149,11 +154,16 @@ GiB", and with no budget set each device's free memory minus a 512 MiB margin
 is used. These resolved GPU budgets, including the safety margin, also drive
 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.
+
 Components are considered in `diffusion`, `te`, `vae` order so that repeatedly
 used diffusion weights have priority. Each component's weights use the first
 storage location with enough remaining budget:
 
-1. The main GPU, leaving estimated space for computation and weight staging.
+1. The component's compute GPU, leaving estimated space for computation and weight staging.
 2. CPU RAM, reserving the larger of 2 GiB or 10% of available RAM for other work.
 3. Another GPU, choosing the one with the largest remaining budget that fits.
 4. Disk, reloading weights on demand.
@@ -170,10 +180,17 @@ weight to be copied again at every step.
 RAM and GPU budgets are shared across components. Each component uses a single
 parameter backend; several other GPUs' capacities are not combined to store
 one component. If available RAM cannot be queried, RAM residency is skipped.
-Other GPUs store weights only: weights are copied to the main GPU for execution.
-Auto-fit does not select multi-GPU layer/row computation, so `--split-mode` does
-not change its placements. Use explicit backend assignments for multi-GPU
-computation.
+Weights stored on another GPU are copied to the component's compute devices for
+execution. CPU modules use RAM or disk. Compute reserves and cache priority are
+accounted for separately on each device, so a CPU module does not reserve GPU
+space. Storage on another module's GPU also leaves room for that module's work.
+
+Auto-fit does not select multi-GPU layer/row computation itself. Explicit device
+lists and `--split-mode` still control that computation. Before the runners have
+built their split plans, auto-fit conservatively counts the full component size
+on each listed GPU when checking residency and cache space. This can offload
+parameters even when a split layout would fit; use `--auto-fit off` to keep the
+default split-device parameter placement.
 
 For example, a diffusion model whose full weights exceed the main GPU's budget
 can use `--backend diffusion=cuda0 --params-backend diffusion=cpu` when RAM is
@@ -292,6 +309,7 @@ The example CLI/server still accepts these older CPU placement flags as compatib
 Because this default is inserted first, later explicit `--params-backend` entries can still override it, for example `--offload-to-cpu --params-backend te=disk` keeps non-TE parameters on CPU and reloads TE parameters from disk.
 
 Library callers should set `backend` and `params_backend` directly. `sd_ctx_params_init()`
-enables `auto_fit` by default; nonempty `backend` or `params_backend` assignments disable it.
+enables `auto_fit` by default; a nonempty `params_backend` assignment disables it.
+The `backend` assignment constrains auto-fit's compute placement.
 The old CPU/offload fields are no longer part of the C API. Explicit `--backend` and
 `--params-backend` assignments are preferred for new commands.

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

@@ -27,7 +27,7 @@ Using `--offload-to-cpu` allows you to offload weights to the CPU, saving VRAM w
 
 ## Use params backend to reduce VRAM or RAM usage.
 
-`--params-backend` controls where model parameters are kept. If it is not set, parameters use the same backend as `--backend`, so a GPU runtime backend also keeps parameters in VRAM.
+`--params-backend` controls where model parameters are kept. If it is not set, auto-fit chooses parameter placement while preserving `--backend`. With `--auto-fit off`, parameters use the same backend as `--backend`, so a GPU runtime backend also keeps parameters in VRAM.
 
 Use CPU params to reduce VRAM usage:
 

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

@@ -0,0 +1,106 @@
+# Sol-Attn
+
+`--sol-attn` enables native CUDA Sol-Attn in the diffusion model, including the
+high-noise diffusion model when present. It uses the shared attention dispatcher
+without classifying tokens as text, images, or video. Python, PyTorch, Triton,
+and CuTe DSL are not needed to build or run it.
+
+This implementation follows the diagonal-threshold algorithm in
+[NVlabs/Sana's Sol-Attn](https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn).
+It summarizes 64-token KV blocks, selects exact blocks using proxy scores and
+an online threshold, and approximates the remaining blocks using their K means
+and V sums. Adjacent blocks remain exact. Both contributions share an online
+softmax normalizer. Q/K/V and probability tiles use BF16 Tensor Cores with FP32
+accumulation; the BF16 result is returned through the existing FP32 interface.
+
+## Build
+
+Use patched GGML, CUDA Toolkit 12.0 or newer, and an NVIDIA GPU with compute
+capability 8.0 or newer. Compile kernels for the target GPU:
+
+```sh
+cmake -S . -B build -DSD_CUDA=ON -DSD_USE_UPSTREAM_GGML=OFF
+cmake --build build --config Release
+```
+
+The feature is compiled with the CUDA backend; no separate build option is
+required. Upstream GGML and non-CUDA backends do not support it. A system GGML
+must provide the matching patched API and CUDA implementation. Tensor-parallel
+row splitting is not supported; layer splitting requires supported devices.
+
+## Use
+
+Add `--sol-attn` to an existing generation command:
+
+```sh
+sd-cli ... --sol-attn
+sd-cli ... --sol-attn --sol-attn-tau 1.0
+```
+
+The default threshold coefficient is `1.0`. Larger coefficients select fewer
+blocks for exact attention. The coefficient must be finite; zero does not mean
+dense attention. Omit `--sol-attn` to disable the feature.
+
+The native kernel supports unmasked, noncausal attention with head dimension
+128, equal Q/K/V sequence lengths and head counts, and multiple batches. Other
+attention operations fall back to FlashAttention when available, then ordinary
+attention. Existing attention scaling overrides remain effective. `--fa` and
+`--diffusion-fa` may be used together with Sol-Attn; `--sage-attn` is mutually
+exclusive. Text encoders and VAEs retain their existing attention selection.
+
+Initialization reports an error if the requested diffusion backend cannot run
+Sol-Attn. Graph logs report the number of Sol-Attn and FlashAttention nodes and
+warn when no Sol-Attn nodes are selected. CUDA execution errors are not silently
+converted into dense attention.
+
+This is approximate attention. Validate quality and end-to-end speed with the
+same prompt, seed, dimensions, frame count, and sampling settings. Include
+packing, preprocessing, offload, and decode time in comparisons. Short sequences
+may not benefit. Upstream combined pipeline speedups are not measurements of
+this native kernel. Exact-covariance thresholds, text sinks, Morton ordering,
+and step/layer schedules are not implemented.
+
+## Validation
+
+On an RTX 4090 with CUDA 12.4, Wan 2.1 T2V 1.3B was tested at 832x480,
+33 frames, 20 Euler steps, seed 42, CFG 6, and flow shift 3, using the prompt
+`a lovely cat` and the same negative prompt for every run:
+
+| Attention | Sampling time | Total process time |
+| --- | ---: | ---: |
+| FlashAttention | 45.73 s | 74.63 s |
+| Sol-Attn, tau 1 | 37.17 s | 66.20 s |
+| Sol-Attn, tau 0 | 40.66 s | 68.50 s |
+
+These are single-run measurements. The graph selected 30 Sol-Attn nodes and
+30 FlashAttention nodes. At tau 1, sampled video frames showed washed-out
+colors and reduced detail. Tau 0 improved clarity in this example, but still
+changed the composition. Neither setting guarantees the baseline's quality.
+For this Wan command, `--sol-attn --sol-attn-tau 0` is a more conservative
+starting point. In the one-frame case, tau 1 increased warm sampling time from
+0.140 to 0.148 seconds per step.
+
+Validation also covered 15 numerical reference cases, 11 layout/scaling/fallback
+cases, CUDA memory checking, and 36 existing SageAttention regression cases.
+CLI and server CUDA builds and the upstream GGML CPU library build passed.
+Other GPU architectures, multi-GPU execution, and other models have not been
+tested.
+
+## Library API
+
+Configure Sol-Attn in `sd_ctx_params_t` before creating the context:
+
+```cpp
+sd_ctx_params_t params;
+sd_ctx_params_init(&params);
+// Set model paths and other context options here.
+params.sol_attn = true;
+params.sol_attn_tau = 1.0f;
+sd_ctx_t* ctx = new_sd_ctx(&params);
+```
+
+`sd_ctx_params_init` defaults `sol_attn` to false and `sol_attn_tau` to 1.0.
+`new_sd_ctx` returns null for a nonfinite threshold, unavailable requested
+backends, or a conflict with SageAttention. The context owns a copy of these
+settings; changing the input structure after creation does not reconfigure it.
+Applications must be rebuilt against the updated `sd_ctx_params_t` definition.

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

@@ -0,0 +1,45 @@
+# Troubleshooting
+
+## Completely black or white images or videos / NaNs
+
+Some ggml backends can encounter numerical overflow during inference, producing
+NaN (not-a-number) values. This can result in completely black or white images or videos.
+Whether it happens can depend on the backend, device, model, and weight format.
+
+Known overflow issues have been addressed as far as possible, but the maintainer
+has limited hardware and cannot test every combination. Some cases may therefore
+still need a manual workaround.
+
+These options are supported by both `sd-cli` and `sd-server`. If you encounter
+this problem, add them to your CLI generation command or server startup command:
+
+```sh
+--linear-scale 0.0078125 --attn-scale 0.0078125
+```
+
+For `sd-server`, restart the server after changing these startup options. Run the
+same prompt and seed again to see whether the output recovers. If the problem
+persists, try smaller positive values, for example:
+
+```sh
+--linear-scale 0.00390625 --attn-scale 0.00390625
+```
+
+These options reduce intermediate values and compensate afterwards to preserve
+the intended output scale:
+
+- `--linear-scale` scales Linear inputs before matrix multiplication and rescales
+  the result.
+- `--attn-scale` scales attention keys and values (K/V). It takes effect only in
+  the Flash Attention path, where `--fa` or `--diffusion-fa` is enabled and the
+  backend supports it.
+
+The two values can be set independently and apply across model components. The
+default `0` preserves each model's built-in settings; `1` explicitly disables the
+corresponding scaling. Overrides must be finite positive values. C API users can
+set `linear_scale` and `attn_scale` in `sd_ctx_params_t`.
+
+If the problem persists after trying the relevant steps above,
+[submit a bug report](https://github.com/leejet/stable-diffusion.cpp/issues/new?template=bug_report.yml).
+Include your full command, backend and hardware, model and weight format, logs,
+and the scale values you tried with their results.

docs/wan.md

@@ -34,6 +34,10 @@
         - Wan2.2 I2V A14B
             - safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/tree/main/split_files/diffusion_models
             - gguf: https://huggingface.co/QuantStack/Wan2.2-I2V-A14B-GGUF/tree/main
+        - Wan2.2 S2V 14B
+            - safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/tree/main/split_files/diffusion_models
+            - gguf: https://huggingface.co/QuantStack/Wan2.2-S2V-14B-GGUF/tree/main
+            - int8_convrot safetensors: https://huggingface.co/noctrex/Wan2.2-S2V-14B-int8_convrot
 - Download vae
     - wan_2.1_vae (for all the wan model except Wan2.2 TI2V 5B)
         - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors
@@ -49,6 +53,9 @@
 - Download clip_vison_h (for Wan2.1 I2V/FLF2V only)
     - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/clip_vision/clip_vision_h.safetensors
 
+- Download audio_encoder (for Wan2.2 S2V only)
+    - safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/blob/main/split_files/audio_encoders/wav2vec2_large_english_fp16.safetensors
+
 
 ## Examples
 
@@ -94,6 +101,48 @@
 
 <video src=../assets/wan/Wan2.2_14B_i2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>
 
+### Wan2.2 S2V 14B
+
+Audio-driven video (speech-to-video). The reference image (`-i`) is the speaker
+portrait, `--audio` is the driving audio track and `--audio-encoder` is the
+wav2vec2 audio encoder. Wan2.2 S2V requires the wan_2.1 vae (16 channel), not
+the wan2.2 vae.
+
+```
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model  ..\models\diffusion_models\wan2.2_s2v-14B-Q8_0.gguf --audio-encoder ..\models\audio_encoders\wav2vec2_large_english_fp16.safetensors --vae ..\models\vae\wan_2.1_vae.safetensors --t5xxl ..\models\text_encoders\umt5-xxl-encoder-Q8_0.gguf  -p "a person is talking" --cfg-scale 6.0 --steps 20 --sampling-method euler -v -n "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" -W 832 -H 480 --diffusion-fa --offload-to-cpu --vae-tiling --video-frames 81 -i ..\assets\cat_with_sd_cpp_42.png --audio .\input\speech.wav --flow-shift 3.0
+```
+
+Notes:
+
+- Recommended settings: `--sampling-method euler --steps 20 --cfg-scale 6.0`.
+  `dpm++2m` produces heavy artifacts on S2V. 4 steps with the lightning LoRA
+  (below) is the fast option.
+- Resolutions: width and height must be multiples of 16; the examples use
+  multiples of 64. 832x480 is a fast starting point; generation cost scales
+  with pixel area.
+- `--audio` accepts a WAV file; it is downmixed to mono and resampled to 16 kHz
+  internally. Audio longer than the video is truncated, video longer than the
+  audio is padded with silence. Pick `--video-frames` to match the audio:
+  roughly `audio_seconds * 16` frames, capped at one chunk (77-81 frames,
+  ~5 s at the model's 16 fps). 33, 77 and 81 map to clean latent frame counts.
+- S2V always uses 16 fps. Other requested frame rates are automatically
+  changed to 16 with a warning, including the CLI and server video output.
+  `generate_video()` returns the actual frame rate through `fps_out`; C API
+  callers should use that value when encoding the output video.
+- One generation covers the first S2V chunk window (`--video-frames` frames).
+  Long-video chunked extend mode is not implemented yet.
+- Speed: the lightx2v lightning LoRA works with S2V at 4 steps and
+  `--cfg-scale 1.0`. Use the **low_noise** variant;
+  the high_noise variant produces artifacts on S2V:
+
+  ```
+  --lora-model-dir ..\models\loras
+  -p "...<lora:lightx2v-Wan2.2-T2V-A14B-4steps-lora-rank64-Seko-V2.0-low_noise:1.0>"
+  --cfg-scale 1.0 --steps 4
+  ```
+
+  Expect some quality/dynamics loss compared to the full 20-step run.
+
 ### Wan2.2 T2V A14B T2I
 
 ```

examples/cli/README.md

@@ -22,3 +22,10 @@ Metadata mode inspects PNG/JPEG container metadata without loading any model:
 ./bin/sd-cli -M metadata --image ./output.png --metadata-raw
 ./bin/sd-cli -M metadata --image ./output.png --metadata-all
 ```
+
+For completely black or white images or videos, NaNs, and the `--linear-scale` /
+`--attn-scale` workaround, see [Troubleshooting](../../docs/troubleshooting.md).
+
+For native CUDA sparse attention in the diffusion model, use `--sol-attn`.
+See [Sol-Attn](../../docs/sol_attention.md) for requirements, supported shapes,
+and the `--sol-attn-tau` threshold coefficient.

examples/cli/main.cpp

@@ -419,7 +419,8 @@ void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy,
             LOG_ERROR("save preview image to '%s' failed", path.string().c_str());
         }
     } else {
-        if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, cli_params->preview_fps, cli_params->compression_quality) != 0) {
+        int fps = cli_params->preview_method == PREVIEW_PROJ ? cli_params->preview_fps / 4 : cli_params->preview_fps;
+        if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, fps, cli_params->compression_quality) != 0) {
             LOG_ERROR("save preview video to '%s' failed", cli_params->preview_path.c_str());
         }
     }
@@ -540,12 +541,16 @@ bool save_results(const SDCliParams& cli_params,
     if (cli_params.mode == VID_GEN && num_results > 1) {
         if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm")
             ext = ".avi";
+        std::string params = gen_params.embed_image_metadata
+                                 ? get_image_params(ctx_params, gen_params, gen_params.seed, cli_params.mode)
+                                 : "";
+
         fs::path video_path = base_path;
         video_path += ext;
         std::string final_ext_lower = ext.string();
         std::transform(final_ext_lower.begin(), final_ext_lower.end(), final_ext_lower.begin(), ::tolower);
         const bool mux_audio = generated_audio != nullptr && (final_ext_lower == ".avi" || final_ext_lower == ".webm");
-        if (create_video_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps, cli_params.compression_quality, mux_audio ? generated_audio : nullptr) == 0) {
+        if (create_video_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps, cli_params.compression_quality, mux_audio ? generated_audio : nullptr, params) == 0) {
             LOG_INFO("save result video to '%s'", video_path.string().c_str());
             if (generated_audio != nullptr && !mux_audio) {
                 fs::path wav_path = video_path;
@@ -687,8 +692,6 @@ int main(int argc, const char* argv[]) {
         }
     }
     cli_params.preview_fps = gen_params.fps;
-    if (cli_params.preview_method == PREVIEW_PROJ)
-        cli_params.preview_fps /= 4;
 
     sd_set_preview_callback(step_callback,
                             cli_params.preview_method,
@@ -951,9 +954,10 @@ int main(int argc, const char* argv[]) {
         } else if (cli_params.mode == VID_GEN) {
             sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
             sd_image_t* generated_video        = nullptr;
-            if (!generate_video(sd_ctx.get(), &vid_gen_params, &generated_video, &num_results, &generated_audio)) {
+            if (!generate_video(sd_ctx.get(), &vid_gen_params, &generated_video, &num_results, &generated_audio, &cli_params.preview_fps)) {
                 generated_video = nullptr;
             }
+            gen_params.fps = cli_params.preview_fps;
             results.adopt(generated_video, num_results);
         }
 

examples/common/common.cpp

@@ -359,6 +359,25 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
     return true;
 }
 
+static int parse_scale_override(int argc, const char** argv, int index, float& scale) {
+    if (++index >= argc) {
+        return -1;
+    }
+    try {
+        size_t end              = 0;
+        const std::string value = argv[index];
+        float parsed            = std::stof(value, &end);
+        if (end != value.size() || !std::isfinite(parsed) || parsed < 0.f ||
+            (parsed > 0.f && !std::isfinite(1.f / parsed))) {
+            return -1;
+        }
+        scale = parsed;
+    } catch (const std::exception&) {
+        return -1;
+    }
+    return 1;
+}
+
 ArgOptions SDContextParams::get_options() {
     ArgOptions options;
     options.string_options = {
@@ -391,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",
@@ -441,6 +465,11 @@ ArgOptions SDContextParams::get_options() {
          "path to standalone LTX audio vae model",
          0,
          &audio_vae_path},
+        {"",
+         "--audio-encoder",
+         "path to wav2vec2 audio encoder model (Wan2.2 S2V)",
+         0,
+         &audio_encoder_path},
         {"",
          "--taesd",
          "path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)",
@@ -589,13 +618,21 @@ 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},
+        {"",
+         "--sol-attn",
+         "use native CUDA Sol-Attn in the diffusion model, with flash/default attention fallback",
+         true, &sol_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},
     };
 
@@ -686,12 +723,26 @@ ArgOptions SDContextParams::get_options() {
         return 1;
     };
 
+    options.float_options.push_back({"", "--sol-attn-tau", "Sol-Attn routing threshold coefficient (default: 1; higher selects fewer exact blocks)", &sol_attn_tau});
+
     options.manual_options = {
+        {"",
+         "--linear-scale",
+         "linear input scale override (float, default: 0 = model default, 1 = no scaling)",
+         [this](int argc, const char** argv, int index) {
+             return parse_scale_override(argc, argv, index, linear_scale);
+         }},
+        {"",
+         "--attn-scale",
+         "flash-attention K/V scale override (float, default: 0 = model default, 1 = no scaling); requires --fa or --diffusion-fa",
+         [this](int argc, const char** argv, int index) {
+             return parse_scale_override(argc, argv, index, attn_scale);
+         }},
         {"",
          "--auto-fit",
-         "on|off (default: on). Use one GPU for diffusion/te/vae computation and place weights on that GPU, "
+         "on|off (default: on). Preserve --backend (otherwise select one GPU) and place weights on the compute GPU, "
          "RAM, another GPU, or disk in that order, according to available memory (--max-vram limits GPU budgets). "
-         "Disabled by explicit --backend or --params-backend; uses automatic graph segmentation when needed",
+         "Disabled by explicit --params-backend; uses automatic graph segmentation when needed",
          on_auto_fit_arg},
         {"",
          "--type",
@@ -777,6 +828,14 @@ bool SDContextParams::resolve(SDMode mode) {
 }
 
 bool SDContextParams::validate(SDMode mode) {
+    if (sol_attn && sage_attn) {
+        LOG_ERROR("--sol-attn and --sage-attn cannot be enabled together");
+        return false;
+    }
+    if (!std::isfinite(sol_attn_tau)) {
+        LOG_ERROR("--sol-attn-tau must be finite");
+        return false;
+    }
     if (mode == CONVERT) {
         const bool has_convert_input = model_path.length() != 0 ||
                                        clip_l_path.length() != 0 ||
@@ -860,6 +919,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"
@@ -867,6 +927,7 @@ std::string SDContextParams::to_string() const {
         << "  vae_path: \"" << vae_path << "\",\n"
         << "  vae_format: \"" << vae_format << "\",\n"
         << "  audio_vae_path: \"" << audio_vae_path << "\",\n"
+        << "  audio_encoder_path: \"" << audio_encoder_path << "\",\n"
         << "  taesd_path: \"" << taesd_path << "\",\n"
         << "  esrgan_path: \"" << esrgan_path << "\",\n"
         << "  control_net_path: \"" << control_net_path << "\",\n"
@@ -895,6 +956,11 @@ 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"
+        << "  sol_attn: " << (sol_attn ? "true" : "false") << ",\n"
+        << "  sol_attn_tau: " << sol_attn_tau << ",\n"
+        << "  linear_scale: " << linear_scale << ",\n"
+        << "  attn_scale: " << attn_scale << ",\n"
         << "  diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n"
         << "  vae_conv_direct: " << (vae_conv_direct ? "true" : "false") << ",\n"
         << "  prediction: " << sd_prediction_name(prediction) << ",\n"
@@ -924,12 +990,14 @@ 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();
     sd_ctx_params.embeddings_connectors_path      = embeddings_connectors_path.c_str();
     sd_ctx_params.vae_path                        = vae_path.c_str();
     sd_ctx_params.audio_vae_path                  = audio_vae_path.c_str();
+    sd_ctx_params.audio_encoder_path              = audio_encoder_path.c_str();
     sd_ctx_params.taesd_path                      = taesd_path.c_str();
     sd_ctx_params.control_net_path                = control_net_path.c_str();
     sd_ctx_params.ip_adapter_path                 = ip_adapter_path.c_str();
@@ -948,6 +1016,11 @@ 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.sol_attn                        = sol_attn;
+    sd_ctx_params.sol_attn_tau                    = sol_attn_tau;
+    sd_ctx_params.linear_scale                    = linear_scale;
+    sd_ctx_params.attn_scale                      = attn_scale;
     sd_ctx_params.tae_preview_only                = taesd_preview;
     sd_ctx_params.diffusion_conv_direct           = diffusion_conv_direct;
     sd_ctx_params.vae_conv_direct                 = vae_conv_direct;
@@ -1060,7 +1133,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},
         {"",
@@ -1480,6 +1553,14 @@ ArgOptions SDGenerationParams::get_options() {
         return 1;
     };
 
+    auto on_audio_arg = [&](int argc, const char** argv, int index) {
+        if (++index >= argc) {
+            return -1;
+        }
+        ref_audio_paths.push_back(argv[index]);
+        return 1;
+    };
+
     auto on_cache_mode_arg = [&](int argc, const char** argv, int index) {
         if (++index >= argc) {
             return -1;
@@ -1669,6 +1750,10 @@ ArgOptions SDGenerationParams::get_options() {
          "--ref-audio",
          "standalone WAV reference for MiniMax-H3 Ref2VA (can be used multiple times)",
          on_ref_audio_arg},
+        {"",
+         "--audio",
+         "driving audio track (Wan2.2 S2V; can be used once)",
+         on_audio_arg},
         {"",
          "--cache-mode",
          "caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT Chebyshev+Taylor forecasting)",

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;
@@ -131,6 +132,7 @@ struct SDContextParams {
     std::string vae_path;
     std::string vae_format = "auto";
     std::string audio_vae_path;
+    std::string audio_encoder_path;
     std::string taesd_path;
     std::string esrgan_path;
     std::string control_net_path;
@@ -168,6 +170,9 @@ struct SDContextParams {
     bool vae_on_cpu            = false;
     bool flash_attn            = false;
     bool diffusion_flash_attn  = false;
+    bool sage_attn             = false;
+    bool sol_attn              = false;
+    float sol_attn_tau         = 1.f;
     bool diffusion_conv_direct = false;
     bool vae_conv_direct       = false;
 
@@ -175,6 +180,8 @@ struct SDContextParams {
     lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO;
 
     bool force_sdxl_vae_conv_scale = false;
+    float linear_scale             = 0.f;
+    float attn_scale               = 0.f;
 
     float flow_shift = INFINITY;
     ArgOptions get_options();

examples/common/media_io.cpp

@@ -810,7 +810,31 @@ uint8_t* load_image_from_memory(const char* image_bytes,
     return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel);
 }
 
-std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
+static void append_avi_metadata(std::vector<uint8_t>& data, const std::string& parameters) {
+    if (parameters.empty()) {
+        return;
+    }
+
+    std::vector<uint8_t> info_content;
+
+    write_fourcc(info_content, "INFO");
+
+    const size_t comment_size = parameters.size() + 1;
+    write_fourcc(info_content, "ICMT");
+    write_u32_le(info_content, static_cast<uint32_t>(comment_size));
+    info_content.insert(info_content.end(), parameters.begin(), parameters.end());
+    info_content.push_back(0);
+    if (comment_size & 1u) {
+        info_content.push_back(0);
+    }
+
+    write_fourcc(data, "LIST");
+    write_u32_le(data, static_cast<uint32_t>(info_content.size()));
+    data.insert(data.end(), info_content.begin(), info_content.end());
+    size_t start_pos = data.size();
+}
+
+std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio, const std::string& parameters) {
     if (num_images == 0) {
         fprintf(stderr, "Error: Image array is empty.\n");
         return {};
@@ -1000,6 +1024,8 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
     const size_t movi_size = avi_data.size() - movi_size_pos - 4;
     patch_u32_le(avi_data, movi_size_pos, static_cast<uint32_t>(movi_size));
 
+    append_avi_metadata(avi_data, parameters);
+
     write_fourcc(avi_data, "idx1");
     write_u32_le(avi_data, static_cast<uint32_t>(index.size() * 16));
     for (const auto& entry : index) {
@@ -1015,8 +1041,8 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
     return avi_data;
 }
 
-int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
-    std::vector<uint8_t> avi_data = create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality, audio);
+int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio, const std::string& parameters) {
+    std::vector<uint8_t> avi_data = create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality, audio, parameters);
     if (avi_data.empty()) {
         return -1;
     }
@@ -1146,7 +1172,7 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
 #endif
 
 #ifdef SD_USE_WEBM
-std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
+std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio, const std::string& parameters) {
     if (num_images == 0) {
         fprintf(stderr, "Error: Image array is empty.\n");
         return {};
@@ -1213,6 +1239,21 @@ std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, in
         segment.GetSegmentInfo()->set_writing_app("stable-diffusion.cpp");
         segment.GetSegmentInfo()->set_muxing_app("stable-diffusion.cpp");
 
+        LOG_DEBUG("Embedding parameters to metadata: %s", parameters.c_str());
+        if (!parameters.empty()) {
+            mkvmuxer::Tag* tag = segment.AddTag();
+
+            if (tag) {
+                if (!tag->add_simple_tag("COMMENT", parameters.c_str())) {
+                    LOG_WARN("Failed to add COMMENT simple tag.");
+                }
+            } else {
+                LOG_WARN("Failed to add tag to segment.");
+            }
+        } else {
+            LOG_INFO("Paramaters is empty, COMMENT tag not embedded.\n");
+        }
+
         const uint64_t frame_duration_ns = std::max<uint64_t>(
             1, static_cast<uint64_t>(std::llround(1000000000.0 / static_cast<double>(fps))));
         uint64_t timestamp_ns = 0;
@@ -1271,8 +1312,8 @@ std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, in
     return writer.data();
 }
 
-int create_webm_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
-    std::vector<uint8_t> webm_data = create_webm_from_sd_images_to_vector(images, num_images, fps, quality, audio);
+int create_webm_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio, const std::string& parameters) {
+    std::vector<uint8_t> webm_data = create_webm_from_sd_images_to_vector(images, num_images, fps, quality, audio, parameters);
     if (webm_data.empty()) {
         return -1;
     }
@@ -1289,7 +1330,8 @@ std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& ou
                                                            int num_images,
                                                            int fps,
                                                            int quality,
-                                                           const sd_audio_t* audio) {
+                                                           const sd_audio_t* audio,
+                                                           const std::string& parameters) {
     std::string format = output_format;
     std::transform(format.begin(), format.end(), format.begin(),
                    [](unsigned char c) { return static_cast<char>(tolower(c)); });
@@ -1299,7 +1341,7 @@ std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& ou
 
 #ifdef SD_USE_WEBM
     if (format == "webm") {
-        return create_webm_from_sd_images_to_vector(images, num_images, fps, quality, audio);
+        return create_webm_from_sd_images_to_vector(images, num_images, fps, quality, audio, parameters);
     }
 #endif
 
@@ -1309,14 +1351,14 @@ std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& ou
     }
 #endif
 
-    return create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality, audio);
+    return create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality, audio, parameters);
 }
 
-int create_video_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
+int create_video_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio, const std::string& parameters) {
     std::string path                = filename ? filename : "";
     auto pos                        = path.find_last_of('.');
     std::string ext                 = pos == std::string::npos ? "" : path.substr(pos);
-    std::vector<uint8_t> video_data = create_video_from_sd_images_to_vector(ext, images, num_images, fps, quality, audio);
+    std::vector<uint8_t> video_data = create_video_from_sd_images_to_vector(ext, images, num_images, fps, quality, audio, parameters);
     if (video_data.empty()) {
         return -1;
     }

examples/common/media_io.h

@@ -57,13 +57,15 @@ int create_mjpg_avi_from_sd_images(const char* filename,
                                    sd_image_t* images,
                                    int num_images,
                                    int fps,
-                                   int quality             = 90,
-                                   const sd_audio_t* audio = nullptr);
+                                   int quality                   = 90,
+                                   const sd_audio_t* audio       = nullptr,
+                                   const std::string& parameters = "");
 std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images,
                                                               int num_images,
                                                               int fps,
-                                                              int quality             = 90,
-                                                              const sd_audio_t* audio = nullptr);
+                                                              int quality                   = 90,
+                                                              const sd_audio_t* audio       = nullptr,
+                                                              const std::string& parameters = "");
 
 #ifdef SD_USE_WEBP
 int create_animated_webp_from_sd_images(const char* filename,
@@ -82,27 +84,31 @@ int create_webm_from_sd_images(const char* filename,
                                sd_image_t* images,
                                int num_images,
                                int fps,
-                               int quality             = 90,
-                               const sd_audio_t* audio = nullptr);
+                               int quality                   = 90,
+                               const sd_audio_t* audio       = nullptr,
+                               const std::string& parameters = "");
 std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images,
                                                           int num_images,
                                                           int fps,
-                                                          int quality             = 90,
-                                                          const sd_audio_t* audio = nullptr);
+                                                          int quality                   = 90,
+                                                          const sd_audio_t* audio       = nullptr,
+                                                          const std::string& parameters = "");
 #endif
 
 int create_video_from_sd_images(const char* filename,
                                 sd_image_t* images,
                                 int num_images,
                                 int fps,
-                                int quality             = 90,
-                                const sd_audio_t* audio = nullptr);
+                                int quality                   = 90,
+                                const sd_audio_t* audio       = nullptr,
+                                const std::string& parameters = "");
 std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& output_format,
                                                            sd_image_t* images,
                                                            int num_images,
                                                            int fps,
-                                                           int quality             = 90,
-                                                           const sd_audio_t* audio = nullptr);
+                                                           int quality                   = 90,
+                                                           const sd_audio_t* audio       = nullptr,
+                                                           const std::string& parameters = "");
 
 bool write_wav_to_file(const std::string& path,
                        const float* interleaved_samples,

examples/server/README.md

@@ -129,3 +129,6 @@ For detailed command-line arguments, run:
 ```bash
 ./bin/sd-server -h
 ```
+
+For completely black or white images or videos, NaNs, and the `--linear-scale` /
+`--attn-scale` startup options, see [Troubleshooting](../../docs/troubleshooting.md).

examples/server/async_jobs.cpp

@@ -237,6 +237,9 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
                          int& output_fps,
                          std::string& error_message) {
     sd_vid_gen_params_t params = job.vid_gen.to_sd_vid_gen_params_t();
+    std::string str_params     = job.vid_gen.gen_params.embed_image_metadata
+                                     ? get_image_params(*runtime.ctx_params, job.vid_gen.gen_params, job.vid_gen.gen_params.seed, VID_GEN)
+                                     : "";
 
     SDImageVec results;
     int num_results             = 0;
@@ -245,7 +248,7 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
     {
         std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
         sd_image_t* raw_results = nullptr;
-        if (!generate_video(runtime.sd_ctx, &params, &raw_results, &num_results, &generated_audio)) {
+        if (!generate_video(runtime.sd_ctx, &params, &raw_results, &num_results, &generated_audio, &output_fps)) {
             raw_results = nullptr;
         }
         results.adopt(raw_results, num_results);
@@ -261,9 +264,10 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
     std::vector<uint8_t> video_bytes = create_video_from_sd_images_to_vector(job.vid_gen.output_format,
                                                                              results.data(),
                                                                              num_results,
-                                                                             job.vid_gen.gen_params.fps,
+                                                                             output_fps,
                                                                              job.vid_gen.output_compression,
-                                                                             generated_audio);
+                                                                             generated_audio,
+                                                                             str_params);
     free_sd_audio(generated_audio);
     if (video_bytes.empty()) {
         error_message = "failed to encode generated video container";
@@ -273,7 +277,6 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
     output_media_b64       = base64_encode(video_bytes);
     output_media_mime_type = video_mime_type(job.vid_gen.output_format);
     output_frame_count     = num_results;
-    output_fps             = job.vid_gen.gen_params.fps;
     return true;
 }
 

include/stable-diffusion.h

@@ -79,6 +79,7 @@ enum scheduler_t {
     FLUX2_SCHEDULER,
     FLUX_SCHEDULER,
     BETA_SCHEDULER,
+    LLADA_IMAGE_SCHEDULER,
     SCHEDULER_COUNT
 };
 
@@ -208,6 +209,7 @@ typedef struct {
     const char* embeddings_connectors_path;
     const char* vae_path;
     const char* audio_vae_path;
+    const char* audio_encoder_path;
     const char* taesd_path;
     const char* control_net_path;
     const char* ip_adapter_path;
@@ -241,6 +243,12 @@ typedef struct {
     const char* rpc_servers;
     const char* model_args;
     bool disable_segmented_compute;  // Force monolithic graph execution even when automatic graph cutting would fit memory better
+    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;
+    bool sol_attn;
+    float sol_attn_tau;
 } sd_ctx_params_t;
 
 typedef struct {
@@ -519,11 +527,13 @@ enum sd_cancel_mode_t {
 SD_API void sd_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode);
 
 SD_API void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params);
+// If non-NULL, fps_out receives the effective encoding frame rate before preview callbacks.
 SD_API bool generate_video(sd_ctx_t* sd_ctx,
                            const sd_vid_gen_params_t* sd_vid_gen_params,
                            sd_image_t** frames_out,
                            int* num_frames_out,
-                           sd_audio_t** audio_out);
+                           sd_audio_t** audio_out,
+                           int* fps_out);
 
 typedef struct upscaler_ctx_t upscaler_ctx_t;
 

src/conditioning/wan_audio.cpp

@@ -0,0 +1,103 @@
+#include "wan_audio.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstddef>
+
+namespace sd::wan_audio {
+
+    static BucketPlan plan_buckets(int audio_frames, int batch_frames, int video_rate, int fps) {
+        BucketPlan plan;
+        plan.audio_frames  = audio_frames;
+        plan.batch_frames  = batch_frames;
+        plan.video_rate    = video_rate;
+        plan.fps           = fps;
+        const double scale = static_cast<double>(video_rate) / fps;
+        // Keep a trailing chunk even when audio ends on a chunk boundary.
+        plan.num_chunks          = static_cast<int>(audio_frames / (batch_frames * scale)) + 1;
+        plan.bucket_frames       = plan.num_chunks * batch_frames;
+        plan.padded_audio_frames = static_cast<int>(
+            std::ceil(plan.bucket_frames / static_cast<double>(fps) * video_rate));
+        return plan;
+    }
+
+    // Match NumPy's round-half-even sampling.
+    static int bucket_source_frame(int bucket_frame, int video_rate, int fps) {
+        return static_cast<int>(std::nearbyint(static_cast<double>(bucket_frame) * video_rate / fps));
+    }
+
+    static int interpolated_frame_count(int in_frames, int input_fps, int output_fps) {
+        return static_cast<int>(in_frames / static_cast<double>(input_fps) * output_fps);
+    }
+
+    // Match PyTorch linear interpolation with align_corners=True.
+    static std::vector<float> linear_interpolate_frames(const std::vector<float>& in,
+                                                        int num_layers,
+                                                        int in_frames,
+                                                        int dim,
+                                                        int out_frames) {
+        std::vector<float> out(static_cast<size_t>(num_layers) * out_frames * dim, 0.0f);
+        if (in.empty() || in_frames <= 0 || out_frames <= 0 || num_layers <= 0 || dim <= 0) {
+            return out;
+        }
+        const double scale = out_frames > 1 ? static_cast<double>(in_frames - 1) / (out_frames - 1) : 0.0;
+        for (int layer = 0; layer < num_layers; ++layer) {
+            for (int out_i = 0; out_i < out_frames; ++out_i) {
+                const double pos     = out_i * scale;
+                const int src0       = static_cast<int>(pos);
+                const int src1       = std::min(src0 + 1, in_frames - 1);
+                const float frac     = static_cast<float>(pos - src0);
+                const float* in_row  = &in[(static_cast<size_t>(layer) * in_frames + src0) * dim];
+                const float* in_next = &in[(static_cast<size_t>(layer) * in_frames + src1) * dim];
+                float* out_row       = &out[(static_cast<size_t>(layer) * out_frames + out_i) * dim];
+                for (int d = 0; d < dim; ++d) {
+                    out_row[d] = in_row[d] * (1.0f - frac) + in_next[d] * frac;
+                }
+            }
+        }
+        return out;
+    }
+
+    std::vector<float> build_audio_buckets(const float* stacked_states,
+                                           int num_layers,
+                                           int in_frames,
+                                           int dim,
+                                           int batch_frames,
+                                           BucketPlan* plan_out,
+                                           int input_fps,
+                                           int video_rate,
+                                           int fps) {
+        if (stacked_states == nullptr || num_layers <= 0 || in_frames <= 0 || dim <= 0 || batch_frames <= 0) {
+            return {};
+        }
+        const int audio_frames = interpolated_frame_count(in_frames, input_fps, video_rate);
+        if (audio_frames <= 0) {
+            return {};
+        }
+        const std::vector<float> interpolated =
+            linear_interpolate_frames(std::vector<float>(stacked_states,
+                                                         stacked_states + static_cast<size_t>(num_layers) * in_frames * dim),
+                                      num_layers,
+                                      in_frames,
+                                      dim,
+                                      audio_frames);
+        const BucketPlan plan = plan_buckets(audio_frames, batch_frames, video_rate, fps);
+        if (plan_out != nullptr) {
+            *plan_out = plan;
+        }
+        std::vector<float> buckets(static_cast<size_t>(plan.bucket_frames) * num_layers * dim, 0.0f);
+        for (int frame = 0; frame < plan.bucket_frames; ++frame) {
+            const int src = bucket_source_frame(frame, video_rate, fps);
+            if (src >= plan.audio_frames) {
+                continue;
+            }
+            for (int layer = 0; layer < num_layers; ++layer) {
+                std::copy_n(interpolated.data() + (static_cast<size_t>(layer) * audio_frames + src) * dim,
+                            static_cast<size_t>(dim),
+                            buckets.data() + (static_cast<size_t>(frame) * num_layers + layer) * dim);
+            }
+        }
+        return buckets;
+    }
+
+}  // namespace sd::wan_audio

src/conditioning/wan_audio.h

@@ -0,0 +1,32 @@
+#ifndef __SD_CONDITIONING_WAN_AUDIO_H__
+#define __SD_CONDITIONING_WAN_AUDIO_H__
+
+#include <vector>
+
+namespace sd::wan_audio {
+
+    struct BucketPlan {
+        int audio_frames;  // frames at video_rate
+        int batch_frames;  // latent_t * 4
+        int video_rate;
+        int fps;         // bucket frame rate
+        int num_chunks;  // includes trailing padding
+        int bucket_frames;
+        int padded_audio_frames;
+    };
+
+    // [layers, frames, dim] at input_fps -> [bucket_frames, layers, dim] at fps.
+    // Pads past the audio end; returns an empty vector on invalid input.
+    std::vector<float> build_audio_buckets(const float* stacked_states,
+                                           int num_layers,
+                                           int in_frames,
+                                           int dim,
+                                           int batch_frames,
+                                           BucketPlan* plan_out = nullptr,
+                                           int input_fps        = 50,
+                                           int video_rate       = 30,
+                                           int fps              = 16);
+
+}  // namespace sd::wan_audio
+
+#endif  // __SD_CONDITIONING_WAN_AUDIO_H__

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

@@ -58,9 +58,15 @@ namespace sd::backend_fit {
         size_t params_device           = SIZE_MAX;
     };
 
+    struct Runtime {
+        std::string name;
+        std::vector<size_t> devices;
+    };
+
     struct Plan {
         bool valid         = false;
         size_t main_device = SIZE_MAX;
+        std::vector<Runtime> runtimes;
         std::vector<Decision> decisions;
     };
 
@@ -96,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) {
@@ -121,11 +127,14 @@ namespace sd::backend_fit {
         return name;
     }
 
-    static std::vector<Device> enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) {
+    static std::vector<Device> enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets,
+                                                     bool include_other_devices) {
         std::vector<Device> out;
         for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
             ggml_backend_dev_t dev = ggml_backend_dev_get(i);
-            if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) {
+            const auto type        = ggml_backend_dev_type(dev);
+            if (type != GGML_BACKEND_DEVICE_TYPE_GPU &&
+                (!include_other_devices || type == GGML_BACKEND_DEVICE_TYPE_CPU)) {
                 continue;
             }
             Device device;
@@ -183,18 +192,29 @@ namespace sd::backend_fit {
         return -1;
     }
 
-    static Plan compute_plan(const std::vector<Component>& components,
-                             const std::vector<Device>& devices,
-                             int64_t ram_budget_bytes) {
-        Plan plan;
+    static size_t select_main_device(const std::vector<Device>& devices) {
+        size_t main_device = SIZE_MAX;
         for (size_t di = 0; di < devices.size(); ++di) {
             if (devices[di].budget_bytes > 0 &&
-                (plan.main_device == SIZE_MAX || devices[di].budget_bytes > devices[plan.main_device].budget_bytes)) {
-                plan.main_device = di;
+                (main_device == SIZE_MAX || devices[di].budget_bytes > devices[main_device].budget_bytes)) {
+                main_device = di;
             }
         }
-        if (plan.main_device == SIZE_MAX) {
-            return plan;
+        return main_device;
+    }
+
+    static Plan compute_plan(const std::vector<Component>& components,
+                             const std::vector<Device>& devices,
+                             int64_t ram_budget_bytes,
+                             const std::vector<Runtime>& runtimes = {}) {
+        Plan plan;
+        plan.main_device = select_main_device(devices);
+        plan.runtimes    = runtimes;
+        if (plan.runtimes.empty()) {
+            if (plan.main_device == SIZE_MAX) {
+                return plan;
+            }
+            plan.runtimes.resize(components.size(), {devices[plan.main_device].name, {plan.main_device}});
         }
 
         std::vector<size_t> order(components.size());
@@ -212,6 +232,27 @@ namespace sd::backend_fit {
         ram_budget_bytes = std::max<int64_t>(ram_budget_bytes, 0);
         plan.decisions.resize(components.size());
 
+        auto uses_device = [&](size_t ci, size_t di) {
+            const auto& runtime_devices = plan.runtimes[ci].devices;
+            return std::find(runtime_devices.begin(), runtime_devices.end(), di) != runtime_devices.end();
+        };
+        auto headroom_for = [&](size_t ci, size_t di) {
+            // Higher-priority offloaded weights need cache space on their compute devices.
+            int64_t headroom = 0;
+            for (size_t other = 0; other < components.size(); ++other) {
+                if (components[other].params_bytes == 0 || !uses_device(other, di)) {
+                    continue;
+                }
+                const bool resident          = other == ci || plan.decisions[other].params_location == ParamsLocation::MAIN_GPU;
+                const int64_t cached_weights = components[other].kind < components[ci].kind
+                                                   ? components[other].params_bytes
+                                                   : components[other].staging_bytes;
+                headroom                     = std::max(headroom, components[other].reserve_bytes +
+                                                                      (resident ? 0 : cached_weights));
+            }
+            return headroom;
+        };
+
         for (size_t ci : order) {
             const Component& comp = components[ci];
             Decision& decision    = plan.decisions[ci];
@@ -219,24 +260,19 @@ namespace sd::backend_fit {
                 continue;
             }
 
-            // Higher-priority offloaded weights need GPU cache space across graph runs.
-            int64_t headroom = 0;
-            for (size_t other = 0; other < components.size(); ++other) {
-                if (components[other].params_bytes == 0) {
-                    continue;
-                }
-                const bool resident          = other == ci || plan.decisions[other].params_location == ParamsLocation::MAIN_GPU;
-                const int64_t cached_weights = components[other].kind < comp.kind
-                                                   ? components[other].params_bytes
-                                                   : components[other].staging_bytes;
-                headroom                     = std::max(headroom, components[other].reserve_bytes +
-                                                                      (resident ? 0 : cached_weights));
-            }
-            int64_t& main_remaining = remaining[plan.main_device];
-            if (headroom <= main_remaining && comp.params_bytes <= main_remaining - headroom) {
+            const auto& runtime_devices = plan.runtimes[ci].devices;
+            const bool fits_runtime     = !runtime_devices.empty() &&
+                                      std::all_of(runtime_devices.begin(), runtime_devices.end(), [&](size_t di) {
+                                          const int64_t headroom = headroom_for(ci, di);
+                                          return headroom <= remaining[di] && comp.params_bytes <= remaining[di] - headroom;
+                                      });
+            if (fits_runtime) {
                 decision.params_location = ParamsLocation::MAIN_GPU;
-                decision.params_device   = plan.main_device;
-                main_remaining -= comp.params_bytes;
+                decision.params_device   = runtime_devices.front();
+                // Exact split allocations are unavailable until the runners build their plans.
+                for (size_t di : runtime_devices) {
+                    remaining[di] -= comp.params_bytes;
+                }
                 continue;
             }
             if (comp.params_bytes <= ram_budget_bytes) {
@@ -244,10 +280,14 @@ namespace sd::backend_fit {
                 ram_budget_bytes -= comp.params_bytes;
                 continue;
             }
+            if (runtime_devices.empty()) {
+                continue;
+            }
 
             size_t best = SIZE_MAX;
             for (size_t di = 0; di < devices.size(); ++di) {
-                if (di != plan.main_device && comp.params_bytes <= remaining[di] &&
+                const int64_t headroom = headroom_for(ci, di);
+                if (!uses_device(ci, di) && headroom <= remaining[di] && comp.params_bytes <= remaining[di] - headroom &&
                     (best == SIZE_MAX || remaining[di] > remaining[best])) {
                     best = di;
                 }
@@ -280,7 +320,7 @@ namespace sd::backend_fit {
                            const std::vector<Device>& devices,
                            int64_t free_ram,
                            int64_t ram_budget) {
-        LOG_INFO("auto-fit plan (single-GPU compute on %s):", devices[plan.main_device].name.c_str());
+        LOG_INFO("auto-fit plan:");
         LOG_INFO("  devices:");
         for (const Device& device : devices) {
             LOG_INFO("    %-12s %-32s free %6lld MiB, budget %6lld MiB",
@@ -293,17 +333,19 @@ namespace sd::backend_fit {
             LOG_INFO("    RAM          free %6lld MiB, params budget %6lld MiB",
                      (long long)(free_ram / MiB), (long long)(ram_budget / MiB));
         }
-        LOG_INFO("  main-GPU weight cache priority: diffusion > te > vae");
-        LOG_INFO("  components (params: main GPU -> RAM -> other GPU -> disk):");
+        LOG_INFO("  compute-device weight cache priority: diffusion > te > vae");
+        LOG_INFO("  components (params: compute device -> RAM -> other GPU -> disk):");
         for (size_t ci = 0; ci < components.size(); ++ci) {
             const Component& comp = components[ci];
             if (comp.params_bytes == 0) {
                 continue;
             }
-            const std::string params = params_backend_name(plan.decisions[ci], devices);
+            const std::string params = plan.decisions[ci].params_location == ParamsLocation::MAIN_GPU
+                                           ? plan.runtimes[ci].name
+                                           : params_backend_name(plan.decisions[ci], devices);
             LOG_INFO("    %-12s params %6lld MiB, compute reserve %5lld MiB -> compute %s, params %s",
                      comp.name, (long long)(comp.params_bytes / MiB), (long long)(comp.reserve_bytes / MiB),
-                     devices[plan.main_device].name.c_str(), params.c_str());
+                     plan.runtimes[ci].name.c_str(), params.c_str());
         }
     }
 
@@ -328,6 +370,51 @@ namespace sd::backend_fit {
         return "";
     }
 
+    static bool resolve_runtimes(const std::vector<Component>& components,
+                                 const std::vector<Device>& devices,
+                                 std::string& runtime_spec,
+                                 std::vector<Runtime>& runtimes,
+                                 std::string& error) {
+        SDBackendAssignment assignment;
+        if (!sd_parse_backend_assignment(runtime_spec, &assignment, &error)) {
+            return false;
+        }
+        const size_t main_device        = select_main_device(devices);
+        const SDBackendModule modules[] = {SDBackendModule::DIFFUSION, SDBackendModule::TE, SDBackendModule::VAE};
+        for (const Component& comp : components) {
+            std::string name = assignment.get(modules[int(comp.kind)]);
+            if (name.empty()) {
+                name = main_device == SIZE_MAX ? "cpu" : devices[main_device].name;
+                if (comp.params_bytes > 0) {
+                    append_assignment(runtime_spec, module_key(comp.kind), name);
+                }
+            }
+            Runtime runtime;
+            for (const std::string& part : split_string(name, '&')) {
+                if (trim(part).empty()) {
+                    continue;
+                }
+                const std::string resolved = sd_backend_resolve_name(part);
+                if (resolved.empty()) {
+                    error = "backend '" + part + "' was not found";
+                    return false;
+                }
+                if (!runtime.name.empty()) {
+                    runtime.name += "&";
+                }
+                runtime.name += resolved;
+                for (size_t di = 0; di < devices.size(); ++di) {
+                    if (devices[di].name == resolved &&
+                        std::find(runtime.devices.begin(), runtime.devices.end(), di) == runtime.devices.end()) {
+                        runtime.devices.push_back(di);
+                    }
+                }
+            }
+            runtimes.push_back(std::move(runtime));
+        }
+        return true;
+    }
+
     bool derive_backend_specs(ModelLoader& loader,
                               ggml_type override_wtype,
                               sd::ggml_graph_cut::MaxVramAssignment& budgets,
@@ -339,12 +426,18 @@ namespace sd::backend_fit {
             return false;
         }
 
-        const auto components    = estimate_components(loader, override_wtype);
-        const auto devices       = enumerate_gpu_devices(budgets);
+        // Resolve once to ensure dynamic backends are loaded before enumerating devices.
+        sd_backend_resolve_name("");
+        const auto components = estimate_components(loader, override_wtype);
+        const auto devices    = enumerate_gpu_devices(budgets, !runtime_spec.empty());
+        std::vector<Runtime> runtimes;
+        if (!runtime_spec.empty() && !resolve_runtimes(components, devices, runtime_spec, runtimes, error)) {
+            LOG_ERROR("%s", error.c_str());
+            return false;
+        }
         const int64_t free_ram   = available_ram_bytes();
         const int64_t ram_budget = std::max<int64_t>(free_ram - std::max<int64_t>(2048 * MiB, free_ram / 10), 0);
-        const auto plan          = compute_plan(components, devices, ram_budget);
-        runtime_spec.clear();
+        const auto plan          = compute_plan(components, devices, ram_budget, runtimes);
         params_spec.clear();
         if (!plan.valid) {
             if (devices.empty()) {
@@ -362,7 +455,9 @@ namespace sd::backend_fit {
                 continue;
             }
             const char* key = module_key(components[ci].kind);
-            append_assignment(runtime_spec, key, devices[plan.main_device].name);
+            if (runtimes.empty()) {
+                append_assignment(runtime_spec, key, plan.runtimes[ci].name);
+            }
             if (plan.decisions[ci].params_location != ParamsLocation::MAIN_GPU) {
                 append_assignment(params_spec, key, params_backend_name(plan.decisions[ci], devices));
             }

src/core/compute_workspace.cpp

@@ -2,14 +2,16 @@
 
 #include <algorithm>
 #include <cstring>
+#include <exception>
 #include <map>
 #include <unordered_map>
 #include <unordered_set>
 
 #include "core/ggml_extend_backend.h"
 #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() {
@@ -228,11 +230,23 @@ namespace sd {
         }
     }
 
-    void ComputeWorkspace::segment_end() {
-        if (active_) {
-            synchronize();
-            active_ = false;
+    bool ComputeWorkspace::segment_end() noexcept {
+        if (!active_) {
+            return true;
         }
+        // Outer cleanup guards must not retry a failed backend submission.
+        active_ = false;
+        try {
+            synchronize();
+            return true;
+        } catch (const std::exception& error) {
+            LOG_ERROR("%s workspace synchronization failed during segment cleanup: %s",
+                      ggml_backend_name(backend_), error.what());
+        } catch (...) {
+            LOG_ERROR("%s workspace synchronization failed during segment cleanup: unknown exception",
+                      ggml_backend_name(backend_));
+        }
+        return false;
     }
 
     bool ComputeWorkspace::release() {

src/core/compute_workspace.h

@@ -51,7 +51,7 @@ namespace sd {
             const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
             const AssignNodes& assign_nodes);
         void synchronize() const;
-        void segment_end();
+        bool segment_end() noexcept;
         bool release();
         bool active() const { return active_; }
         ggml_backend_sched_t scheduler() const { return scheduler_; }

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,
@@ -325,6 +337,76 @@ ggml_tensor* ggml_ext_pad(ggml_context* ctx,
     return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y);
 }
 
+static ggml_tensor* conv_1d(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, int s0, int p0, int d0, bool force_prec_f32) {
+    ggml_tensor* result;
+    if (force_prec_f32) {
+        ggml_tensor* patches = ggml_im2col(ctx, w, x, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F32);
+        result               = ggml_mul_mat(ctx,
+                                            ggml_reshape_2d(ctx, patches, patches->ne[0], patches->ne[2] * patches->ne[1]),
+                                            ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1], w->ne[2]));
+        result               = ggml_reshape_3d(ctx, result, patches->ne[1], w->ne[2], patches->ne[2]);
+    } else {
+        result = ggml_conv_1d(ctx, w, x, s0, p0, d0);
+    }
+    if (x->ne[2] > 1) {
+        // mul_mat packs positions and batches before output channels: [OL, N, OC].
+        result = ggml_reshape_3d(ctx, result, result->ne[0], x->ne[2], w->ne[2]);
+        result = ggml_cont(ctx, ggml_permute(ctx, result, 0, 2, 1, 3));
+    }
+    return result;
+}
+
+ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int s0,
+                              int p0,
+                              int d0,
+                              int64_t groups,
+                              bool force_prec_f32) {
+    GGML_ASSERT(s0 > 0 && p0 >= 0 && d0 > 0 && groups > 0);
+    GGML_ASSERT(x->type == GGML_TYPE_F32 && x->ne[3] == 1 && w->ne[3] == 1);
+    GGML_ASSERT(x->ne[1] % groups == 0 && w->ne[2] % groups == 0);
+    GGML_ASSERT(w->ne[1] == x->ne[1] / groups);
+    GGML_ASSERT(b == nullptr || (b->type == GGML_TYPE_F32 && ggml_is_vector(b) && b->ne[0] == w->ne[2]));
+
+    // im2col requires contiguous time rows; group views must retain the real channel and batch strides.
+    if (!ggml_is_contiguous(x)) {
+        x = ggml_cont(ctx, x);
+    }
+    if (force_prec_f32 && w->type != GGML_TYPE_F32) {
+        w = ggml_cast(ctx, w, GGML_TYPE_F32);
+    }
+    if (!ggml_is_contiguous(w)) {
+        w = ggml_cont(ctx, w);
+    }
+
+    ggml_tensor* result = nullptr;
+    if (groups == 1) {
+        result = conv_1d(ctx, x, w, s0, p0, d0, force_prec_f32);
+    } else {
+        const int64_t ic_g = x->ne[1] / groups;
+        const int64_t oc_g = w->ne[2] / groups;
+        std::vector<ggml_tensor*> outputs;
+        outputs.reserve(groups);
+        for (int64_t group = 0; group < groups; ++group) {
+            ggml_tensor* x_i = ggml_view_3d(ctx, x, x->ne[0], ic_g, x->ne[2], x->nb[1], x->nb[2], group * ic_g * x->nb[1]);
+            ggml_tensor* w_i = ggml_view_3d(ctx, w, w->ne[0], ic_g, oc_g, w->nb[1], w->nb[2], group * oc_g * w->nb[2]);
+            outputs.push_back(conv_1d(ctx, x_i, w_i, s0, p0, d0, force_prec_f32));
+        }
+        result = ggml_ext_vec_concat(ctx, outputs, 1);
+    }
+    if (b != nullptr) {
+        if (!ggml_is_contiguous(b)) {
+            b = ggml_cont(ctx, b);
+        }
+        b      = ggml_reshape_3d(ctx, b, 1, w->ne[2], 1);
+        result = ggml_add_inplace(ctx, result, b);
+    }
+    return result;
+}
+
 ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
                               ggml_tensor* x,
                               ggml_tensor* w,
@@ -382,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;
@@ -535,7 +622,10 @@ 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,
+                                    bool sol_attn,
+                                    float sol_attn_tau) {  // avoid overflow
     int64_t L_q;
     int64_t L_k;
     int64_t C;
@@ -573,6 +663,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);
         }
@@ -580,6 +678,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);
         }
@@ -609,10 +710,59 @@ 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 (sol_attn && mask == nullptr && d_head == 128 && L_q == L_k && n_head == n_kv_head) {
+        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));
+        if (v_in->type != GGML_TYPE_F32) {
+            v_in = ggml_cast(ctx, v_in, GGML_TYPE_F32);
+        }
+        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);
+        }
+        auto out = ggml_sol_attn(ctx, q_in, k_in, v_in, scale / kv_scale, sol_attn_tau);
+        if (ggml_backend_supports_op(backend, out)) {
+            kqv = kv_scale != 1.0f ? ggml_ext_scale(ctx, out, 1.0f / kv_scale) : out;
+        }
+    }
+    if (kqv == nullptr && 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 || sol_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) {
@@ -683,17 +833,16 @@ ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
                                  ggml_tensor* x,
                                  ggml_tensor* w,
                                  ggml_tensor* b,
-                                 int num_groups) {
+                                 int num_groups,
+                                 float eps) {
     if (ggml_n_dims(x) >= 3 && w != nullptr && b != nullptr) {
         w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1);
         b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
     }
 
-    const float eps = 1e-6f;  // default eps parameter
-    x               = ggml_group_norm(ctx, x, num_groups, eps);
+    x = ggml_group_norm(ctx, x, num_groups, eps);
     if (w != nullptr && b != nullptr) {
         x = ggml_mul_inplace(ctx, x, w);
-        // b = ggml_repeat(ctx, b, x);
         x = ggml_add_inplace(ctx, x, b);
     }
     return x;

src/core/ggml_extend.h

@@ -103,6 +103,18 @@ ggml_tensor* ggml_ext_pad(ggml_context* ctx,
                           bool circular_x = false,
                           bool circular_y = false);
 
+// ggml layout: x [L, IC, N], w [K, IC/groups, OC], b [OC], result [OL, OC, N].
+// force_prec_f32 keeps both input patches and weights in F32.
+ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int s0              = 1,
+                              int p0              = 0,
+                              int d0              = 1,
+                              int64_t groups      = 1,
+                              bool force_prec_f32 = false);
+
 // w: [OC，IC, KH, KW]
 // x: [N, IC, IH, IW]
 // b: [OC,]
@@ -141,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]
@@ -204,10 +217,13 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
                                     ggml_tensor* k,
                                     ggml_tensor* v,
                                     int64_t n_head,
-                                    ggml_tensor* mask = nullptr,
-                                    bool skip_reshape = false,
-                                    bool flash_attn   = false,
-                                    float kv_scale    = 1.0f);
+                                    ggml_tensor* mask  = nullptr,
+                                    bool skip_reshape  = false,
+                                    bool flash_attn    = false,
+                                    float kv_scale     = 1.0f,
+                                    bool sage_attn     = false,
+                                    bool sol_attn      = false,
+                                    float sol_attn_tau = 1.0f);
 
 ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
                                  ggml_tensor* x,
@@ -219,7 +235,8 @@ ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
                                  ggml_tensor* x,
                                  ggml_tensor* w,
                                  ggml_tensor* b,
-                                 int num_groups = 32);
+                                 int num_groups = 32,
+                                 float eps      = 1e-6f);
 
 ggml_tensor* ggml_ext_timestep_embedding(
     ggml_context* ctx,

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) {
@@ -87,6 +91,10 @@ static bool parse_backend_module(const std::string& raw_name, SDBackendModule* m
         *module = SDBackendModule::DETECTOR;
         return true;
     }
+    if (name == "audioencoder" || name == "audio") {
+        *module = SDBackendModule::AUDIO_ENCODER;
+        return true;
+    }
     return false;
 }
 
@@ -425,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);
@@ -593,7 +665,7 @@ static ggml_backend_t sd_get_default_backend() {
     return backend;
 }
 
-static bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error) {
+bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error) {
     if (assignment == nullptr) {
         return false;
     }
@@ -660,7 +732,13 @@ void SDBackendAssignment::set_module(SDBackendModule module, const std::string&
 }
 
 void SDBackendHandleDeleter::operator()(ggml_backend_t backend) const {
-    ggml_backend_free(backend);
+    try {
+        ggml_backend_free(backend);
+    } catch (const std::exception& error) {
+        LOG_ERROR("backend cleanup failed: %s", error.what());
+    } catch (...) {
+        LOG_ERROR("backend cleanup failed: unknown exception");
+    }
 }
 
 SDBackendManager::~SDBackendManager() {
@@ -962,6 +1040,8 @@ const char* sd_backend_module_name(SDBackendModule module) {
             return "upscaler";
         case SDBackendModule::DETECTOR:
             return "detector";
+        case SDBackendModule::AUDIO_ENCODER:
+            return "audio_encoder";
     }
     return "unknown";
 }

src/core/ggml_extend_backend.h

@@ -21,6 +21,7 @@ enum class SDBackendModule {
     PHOTOMAKER,
     UPSCALER,
     DETECTOR,
+    AUDIO_ENCODER,
 };
 
 struct SDBackendAssignment {
@@ -86,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,
@@ -93,6 +95,7 @@ ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
                                                         sd_graph_eval_callback_t callback_eval,
                                                         void* callback_eval_user_data);
 std::string sd_backend_resolve_name(const std::string& name);
+bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error);
 const char* sd_backend_module_name(SDBackendModule module);
 void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);
 bool add_rpc_devices(const std::string& servers);

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

@@ -1,7 +1,9 @@
 #include <algorithm>
+#include <exception>
 #include <map>
 #include <utility>
 
+#include "core/ggml_extend.h"
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_runner.h"
 #include "core/ggml_tensor_utils.h"
@@ -11,6 +13,21 @@
 
 using namespace sd;
 
+ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
+                                    ggml_tensor* q,
+                                    ggml_tensor* k,
+                                    ggml_tensor* v,
+                                    int64_t n_head,
+                                    ggml_tensor* mask,
+                                    bool skip_reshape,
+                                    bool flash_attn,
+                                    float kv_scale) {
+    if (ctx->attn_scale > 0.f) {
+        kv_scale = ctx->attn_scale;
+    }
+    return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled, ctx->sol_attn_enabled, ctx->sol_attn_tau);
+}
+
 void GGMLRunner::alloc_params_ctx() {
     ggml_init_params params;
     params.mem_size   = static_cast<size_t>(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead());
@@ -147,6 +164,22 @@ ggml_cgraph* GGMLRunner::get_compute_graph(get_graph_cb_t get_graph) {
         }
     }
     prepare_build_in_tensor_after(gf);
+#ifndef SD_USE_UPSTREAM_GGML
+    if (sol_attn_enabled && !sol_attn_graph_logged) {
+        int sol_nodes   = 0;
+        int flash_nodes = 0;
+        for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) {
+            const auto op = ggml_graph_node(gf, i)->op;
+            sol_nodes += op == GGML_OP_SOL_ATTN;
+            flash_nodes += op == GGML_OP_FLASH_ATTN_EXT;
+        }
+        LOG_INFO("Sol-Attn graph: %d Sol-Attn nodes, %d FlashAttention nodes", sol_nodes, flash_nodes);
+        if (sol_nodes == 0) {
+            LOG_WARN("This graph has no attention operations supported by Sol-Attn");
+        }
+        sol_attn_graph_logged = true;
+    }
+#endif
     return gf;
 }
 
@@ -325,21 +358,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) {
@@ -352,10 +381,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) {
@@ -510,7 +536,13 @@ 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.sol_attn_enabled      = sol_attn_enabled;
+    runner_ctx.sol_attn_tau          = sol_attn_tau;
+    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;
@@ -624,8 +656,15 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
                 params_tensor_set_.insert(parameter);
         }
     }
-    auto output = execute_graph(graph, n_threads, no_return, read_outputs);
-    success     = output.has_value();
+    std::optional<sd::Tensor<float>> output;
+    try {
+        output = execute_graph(graph, n_threads, no_return, read_outputs);
+    } catch (const std::exception& error) {
+        LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(),
+                  ggml_backend_name(runtime_backend), error.what());
+        return std::nullopt;
+    }
+    success = output.has_value();
     if (success) {
         cache_.graph_end(true);
     }
@@ -793,24 +832,25 @@ 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()) {
         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);
         }
@@ -823,8 +863,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(),
@@ -883,7 +923,10 @@ 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);
             }
@@ -938,6 +981,9 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
                 }
             }
         }
+        if (!workspace_.segment_end()) {
+            return fail_segment("workspace synchronization");
+        }
         // Final outputs and their callbacks may still be views of consumed cuts.
         cut_cache_.prune(segment.future_cut_names);
     }

src/core/ggml_runner.h

@@ -68,7 +68,13 @@ struct GGMLRunnerContext {
     ggml_backend_t backend                                           = nullptr;
     ggml_context* ggml_ctx                                           = nullptr;
     bool flash_attn_enabled                                          = false;
+    bool sage_attn_enabled                                           = false;
+    bool sol_attn_enabled                                            = false;
+    float sol_attn_tau                                               = 1.f;
+    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;
@@ -113,6 +119,16 @@ struct GGMLRunnerContext {
     }
 };
 
+ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
+                                    ggml_tensor* q,
+                                    ggml_tensor* k,
+                                    ggml_tensor* v,
+                                    int64_t n_head,
+                                    ggml_tensor* mask = nullptr,
+                                    bool skip_reshape = false,
+                                    bool flash_attn   = false,
+                                    float kv_scale    = 1.f);
+
 struct GGMLRunner {
 private:
     std::map<ggml_backend_t, size_t> logged_compute_bytes_;
@@ -163,7 +179,14 @@ protected:
     const std::string final_result_name = "ggml_runner_final_result_tensor";
 
     bool flash_attn_enabled    = false;
+    bool sage_attn_enabled     = false;
+    bool sol_attn_enabled      = false;
+    float sol_attn_tau         = 1.f;
+    bool sol_attn_graph_logged = 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;
 
@@ -249,11 +272,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);
 
@@ -323,10 +344,37 @@ public:
         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_sol_attention_enabled(bool enabled, float tau) {
+        if (sol_attn_enabled != enabled || sol_attn_tau != tau) {
+            free_cache_ctx_and_buffer();
+            graph_cut_plan_cache_.graph_cut_plans.clear();
+            sol_attn_enabled      = enabled;
+            sol_attn_tau          = tau;
+            sol_attn_graph_logged = false;
+        }
+    }
+
+    void set_scale_overrides(float linear_scale, float attn_scale) {
+        this->linear_scale = linear_scale;
+        this->attn_scale   = attn_scale;
+    }
+
     void set_conv2d_direct_enabled(bool enabled) {
         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]});
@@ -821,7 +872,11 @@ std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::str
     float round_bracket_multiplier  = 1.1f;
     float square_bracket_multiplier = 1 / 1.1f;
 
-    std::regex re_attention(R"(\\\(|\\\)|\\\[|\\\]|\\\\|\\|\(|\[|:([+-]?[.\d]+)\)|\)|\]|\bBREAK\b|[^\\()\[\]:B]+|:|\bB)");
+    // libstdc++ std::regex recurses per matched character, so unbounded runs
+    // overflow the stack. Split runs are merged back below.
+    const int max_plain_text_run = 1024;
+    std::regex re_attention(R"(\\\(|\\\)|\\\[|\\\]|\\\\|\\|\(|\[|\)|\]|\bBREAK\b|[^\\()\[\]:B]{1,)" +
+                            std::to_string(max_plain_text_run) + R"(}|:|\bB)");
     std::regex re_break(R"(\s*\bBREAK\b\s*)");
 
     auto multiply_range = [&](int start_position, float multiplier) {
@@ -830,22 +885,55 @@ std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::str
         }
     };
 
+    // Kept out of the regex: bounding the repetition rejects valid long weights,
+    // leaving it unbounded overflows the stack.
+    auto lex_weight = [](const std::string& s, float& value) -> size_t {
+        size_t end = 0;
+        if (end < s.size() && (s[end] == '+' || s[end] == '-')) {
+            ++end;
+        }
+        while (end < s.size() && (std::isdigit((unsigned char)s[end]) || s[end] == '.')) {
+            ++end;
+        }
+        if (end >= s.size() || s[end] != ')') {
+            return 0;
+        }
+        std::string number   = s.substr(0, end);
+        char* number_end     = nullptr;
+        float parsed         = std::strtof(number.c_str(), &number_end);
+        const char* expected = number.c_str() + number.size();
+        // Without this ".", "+." and "1.2.3" would silently become weights.
+        if (number.empty() || number_end != expected || !std::isfinite(parsed)) {
+            return 0;
+        }
+        value = parsed;
+        return end + 1;
+    };
+
     std::smatch m, m2;
     std::string remaining_text = text;
 
     while (std::regex_search(remaining_text, m, re_attention)) {
         std::string text   = m[0];
-        std::string weight = m[1];
+        std::string suffix = m.suffix();
+
+        if (text == ":") {
+            float weight_value   = 1.0f;
+            size_t weight_length = lex_weight(suffix, weight_value);
+            if (weight_length > 0) {
+                if (!round_brackets.empty()) {
+                    multiply_range(round_brackets.back(), weight_value);
+                    round_brackets.pop_back();
+                }
+                remaining_text = suffix.substr(weight_length);
+                continue;
+            }
+        }
 
         if (text == "(") {
             round_brackets.push_back((int)res.size());
         } else if (text == "[") {
             square_brackets.push_back((int)res.size());
-        } else if (!weight.empty()) {
-            if (!round_brackets.empty()) {
-                multiply_range(round_brackets.back(), std::stof(weight));
-                round_brackets.pop_back();
-            }
         } else if (text == ")" && !round_brackets.empty()) {
             multiply_range(round_brackets.back(), round_bracket_multiplier);
             round_brackets.pop_back();
@@ -860,7 +948,7 @@ std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::str
             res.push_back({text, 1.0f});
         }
 
-        remaining_text = m.suffix();
+        remaining_text = suffix;
     }
 
     for (int pos : round_brackets) {

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 {
@@ -135,6 +160,7 @@ struct PhotoMakerExtension : public GenerationExtension {
                                                            pm_version,
                                                            20.f,
                                                            ctx.model_manager);
+        pmid_model->set_scale_overrides(ctx.params->linear_scale, ctx.params->attn_scale);
         if (pm_version == PM_VERSION_2) {
             LOG_INFO("using PhotoMaker Version 2");
         }
@@ -222,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");
@@ -262,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

@@ -35,9 +35,11 @@ enum SDVersion {
     VERSION_WAN2,
     VERSION_WAN2_2_I2V,
     VERSION_WAN2_2_TI2V,
+    VERSION_WAN2_2_S2V,
     VERSION_LINGBOT_VIDEO,
     VERSION_QWEN_IMAGE,
     VERSION_QWEN_IMAGE_LAYERED,
+    VERSION_QWEN_IMAGE_2_1,
     VERSION_HUNYUAN_VIDEO,
     VERSION_ANIMA,
     VERSION_FLUX2,
@@ -58,6 +60,7 @@ enum SDVersion {
     VERSION_KREA2,
     VERSION_MAGE_FLOW,
     VERSION_SENSENOVA_U1_5,
+    VERSION_LLADA_IMAGE,
     VERSION_ESRGAN,
     VERSION_COUNT,
 };
@@ -130,7 +133,7 @@ static inline bool sd_version_is_minimax_h3(SDVersion version) {
 }
 
 static inline bool sd_version_is_wan(SDVersion version) {
-    if (version == VERSION_WAN2 || version == VERSION_WAN2_2_I2V || version == VERSION_WAN2_2_TI2V) {
+    if (version == VERSION_WAN2 || version == VERSION_WAN2_2_I2V || version == VERSION_WAN2_2_TI2V || version == VERSION_WAN2_2_S2V) {
         return true;
     }
     return false;
@@ -144,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;
@@ -171,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;
@@ -250,7 +260,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;
@@ -291,6 +301,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/ip_adapter.hpp

@@ -95,7 +95,7 @@ namespace IPAdapter {
                 int64_t L          = kv->ne[1];
                 ggml_tensor* k     = ggml_cont(ctx->ggml_ctx, ggml_view_3d(ctx->ggml_ctx, kv, dim, L, N, kv->nb[1], kv->nb[2], 0));
                 ggml_tensor* v     = ggml_cont(ctx->ggml_ctx, ggml_view_3d(ctx->ggml_ctx, kv, dim, L, N, kv->nb[1], kv->nb[2], dim * kv->nb[0]));
-                ggml_tensor* attn  = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, heads, nullptr, false, false);
+                ggml_tensor* attn  = ggml_ext_attention_ext(ctx, q, k, v, heads, nullptr, false, false);
                 attn               = to_out->forward(ctx, attn);
                 latents            = ggml_add(ctx->ggml_ctx, latents, attn);
 

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

@@ -63,12 +63,11 @@ public:
         k              = ggml_cont(ctx->ggml_ctx, k);
         v              = ggml_cont(ctx->ggml_ctx, v);
 
-        ggml_tensor* attn_out = ggml_ext_attention_ext(
-            ctx->ggml_ctx, ctx->backend,
-            q, k, v,
-            heads,
-            /*mask=*/nullptr,
-            /*diag_mask_inf=*/false);
+        ggml_tensor* attn_out = ggml_ext_attention_ext(ctx,
+                                                       q, k, v,
+                                                       heads,
+                                                       /*mask=*/nullptr,
+                                                       /*diag_mask_inf=*/false);
 
         ggml_tensor* out = to_out->forward(ctx, attn_out);
         return out;

src/model/audio/wav2vec2.hpp

@@ -0,0 +1,413 @@
+#ifndef __SD_MODEL_AUDIO_WAV2VEC2_HPP__
+#define __SD_MODEL_AUDIO_WAV2VEC2_HPP__
+
+#include <algorithm>
+#include <cinttypes>
+#include <cmath>
+#include <cstdio>
+#include <map>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "model.h"
+#include "model/common/ggml_block.hpp"
+
+namespace Wav2Vec2 {
+
+    struct Wav2Vec2Config {
+        int64_t embed_dim             = 1024;
+        int64_t conv_dim              = 512;
+        int num_heads                 = 16;
+        int num_layers                = 24;
+        std::string feat_extract_norm = "layer";
+        bool conv_bias                = true;
+        bool do_normalize             = true;
+        bool do_stable_layer_norm     = true;
+
+        static Wav2Vec2Config detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            Wav2Vec2Config config;
+            auto it = tensor_storage_map.find(prefix + "encoder.layer_norm.bias");
+            if (it == tensor_storage_map.end()) {
+                LOG_WARN("wav2vec2: %sencoder.layer_norm.bias not found, using large defaults", prefix.c_str());
+                return config;
+            }
+            config.embed_dim = it->second.ne[0];
+            if (config.embed_dim == 1024) {
+                config.embed_dim            = 1024;
+                config.num_heads            = 16;
+                config.num_layers           = 24;
+                config.feat_extract_norm    = "layer";
+                config.conv_bias            = true;
+                config.do_normalize         = true;
+                config.do_stable_layer_norm = true;
+            } else if (config.embed_dim == 768) {
+                config.embed_dim            = 768;
+                config.num_heads            = 12;
+                config.num_layers           = 12;
+                config.feat_extract_norm    = "group";
+                config.conv_bias            = false;
+                config.do_normalize         = false;
+                config.do_stable_layer_norm = false;
+            } else {
+                LOG_WARN("wav2vec2: unsupported embed_dim %" PRId64 ", using large defaults", config.embed_dim);
+                config.embed_dim = 1024;
+            }
+            return config;
+        }
+    };
+
+    struct Wav2Vec2NoLayerNormConvLayer : public UnaryBlock {
+        Wav2Vec2NoLayerNormConvLayer(int64_t in_channels, int64_t out_channels, int kernel_size, int stride, bool bias) {
+            blocks["conv"] = std::make_shared<Conv1d>(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto conv = std::dynamic_pointer_cast<Conv1d>(blocks["conv"]);
+            x         = conv->forward(ctx, x);
+            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
+        }
+    };
+
+    struct Wav2Vec2LayerNormConvLayer : public UnaryBlock {
+        Wav2Vec2LayerNormConvLayer(int64_t in_channels, int64_t out_channels, int kernel_size, int stride, bool bias) {
+            blocks["conv"]       = std::make_shared<Conv1d>(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias, true);
+            blocks["layer_norm"] = std::make_shared<LayerNorm>(out_channels);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto conv       = std::dynamic_pointer_cast<Conv1d>(blocks["conv"]);
+            auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
+            x               = conv->forward(ctx, x);
+            // LayerNorm normalizes channels: [N, C, L] -> [N, L, C].
+            x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
+            x = layer_norm->forward(ctx, x);
+            x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
+            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
+        }
+    };
+
+    struct Wav2Vec2GroupNormConvLayer : public UnaryBlock {
+        Wav2Vec2GroupNormConvLayer(int64_t in_channels, int64_t out_channels, int kernel_size, int stride, bool bias) {
+            blocks["conv"]       = std::make_shared<Conv1d>(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias, true);
+            blocks["layer_norm"] = std::make_shared<GroupNorm>((int)out_channels, out_channels, 1e-05f);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto conv       = std::dynamic_pointer_cast<Conv1d>(blocks["conv"]);
+            auto layer_norm = std::dynamic_pointer_cast<GroupNorm>(blocks["layer_norm"]);
+            x               = conv->forward(ctx, x);
+            // ggml GroupNorm needs [N, C, H, W], with H=1 for audio.
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]);
+            x = layer_norm->forward(ctx, x);
+            x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[2], x->ne[3]);
+            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
+        }
+    };
+
+    struct Wav2Vec2FeatureEncoder : public UnaryBlock {
+        Wav2Vec2FeatureEncoder(const Wav2Vec2Config& config) {
+            GGML_ASSERT(config.feat_extract_norm == "layer" || config.feat_extract_norm == "group");
+            const int kernels[7] = {10, 3, 3, 3, 3, 2, 2};
+            const int strides[7] = {5, 2, 2, 2, 2, 2, 2};
+            int64_t in_channels  = 1;
+            for (int i = 0; i < 7; ++i) {
+                const std::string name = "conv_layers." + std::to_string(i);
+                if (config.feat_extract_norm == "layer") {
+                    blocks[name] = std::make_shared<Wav2Vec2LayerNormConvLayer>(in_channels, config.conv_dim, kernels[i], strides[i], config.conv_bias);
+                } else if (i == 0) {
+                    blocks[name] = std::make_shared<Wav2Vec2GroupNormConvLayer>(in_channels, config.conv_dim, kernels[i], strides[i], config.conv_bias);
+                } else {
+                    blocks[name] = std::make_shared<Wav2Vec2NoLayerNormConvLayer>(in_channels, config.conv_dim, kernels[i], strides[i], config.conv_bias);
+                }
+                in_channels = config.conv_dim;
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            for (int i = 0; i < 7; ++i) {
+                auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv_layers." + std::to_string(i)]);
+                x         = conv->forward(ctx, x);
+            }
+            return ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
+        }
+    };
+
+    struct Wav2Vec2FeatureProjection : public UnaryBlock {
+        Wav2Vec2FeatureProjection(const Wav2Vec2Config& config) {
+            blocks["layer_norm"] = std::make_shared<LayerNorm>(config.conv_dim);
+            blocks["projection"] = std::make_shared<Linear>(config.conv_dim, config.embed_dim);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto ln         = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
+            auto projection = std::dynamic_pointer_cast<Linear>(blocks["projection"]);
+            x               = ln->forward(ctx, x);
+            x               = projection->forward(ctx, x);
+            return x;
+        }
+    };
+
+    class Wav2Vec2PositionalConvEmbedding : public UnaryBlock {
+    private:
+        int64_t embed_dim_;
+        static constexpr int groups_      = 16;
+        static constexpr int kernel_size_ = 128;
+        std::string weight_g_name_;
+        std::string weight_v_name_;
+
+        ggml_tensor* weight(GGMLRunnerContext* ctx) {
+            auto g       = params[weight_g_name_];
+            auto v       = ggml_cast(ctx->ggml_ctx, params[weight_v_name_], GGML_TYPE_F32);
+            auto squared = ggml_mul(ctx->ggml_ctx, v, v);
+            // PyTorch weight_norm(dim=2) reduces both channel axes, retaining each kernel tap.
+            squared   = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, squared, 2, 0, 1, 3));
+            squared   = ggml_reshape_2d(ctx->ggml_ctx, squared, embed_dim_ / groups_ * embed_dim_, kernel_size_);
+            auto norm = ggml_sqrt(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, squared));
+            norm      = ggml_reshape_3d(ctx->ggml_ctx, norm, kernel_size_, 1, 1);
+            return ggml_mul(ctx->ggml_ctx, v, ggml_div(ctx->ggml_ctx, g, norm));
+        }
+
+    public:
+        Wav2Vec2PositionalConvEmbedding(const Wav2Vec2Config& config)
+            : embed_dim_(config.embed_dim) {
+            GGML_ASSERT(embed_dim_ > 0 && embed_dim_ % groups_ == 0);
+        }
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            bool legacy    = tensor_storage_map.count(prefix + "conv.weight_g") > 0;
+            weight_g_name_ = legacy ? "conv.weight_g" : "conv.parametrizations.weight.original0";
+            weight_v_name_ = legacy ? "conv.weight_v" : "conv.parametrizations.weight.original1";
+            auto g         = tensor_storage_map.find(prefix + weight_g_name_);
+            auto v         = tensor_storage_map.find(prefix + weight_v_name_);
+            GGML_ASSERT(g != tensor_storage_map.end() && v != tensor_storage_map.end());
+            GGML_ASSERT(g->second.ne[0] == kernel_size_ && g->second.ne[1] == 1 && g->second.ne[2] == 1 && g->second.ne[3] == 1);
+            GGML_ASSERT(v->second.ne[0] == kernel_size_ && v->second.ne[1] == embed_dim_ / groups_ && v->second.ne[2] == embed_dim_ && v->second.ne[3] == 1);
+
+            params[weight_g_name_] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kernel_size_, 1, 1);
+            params[weight_v_name_] = ggml_new_tensor_3d(ctx, get_type(prefix + weight_v_name_, tensor_storage_map, GGML_TYPE_F16),
+                                                        kernel_size_, embed_dim_ / groups_, embed_dim_);
+            if (tensor_storage_map.count(prefix + "conv.bias") > 0) {
+                params["conv.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, embed_dim_);
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto w = weight(ctx);
+            auto b = params.count("conv.bias") > 0 ? params["conv.bias"] : nullptr;
+            x      = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x      = ggml_ext_conv_1d(ctx->ggml_ctx, x, w, b, 1, kernel_size_ / 2, 1, groups_, true);
+            // Apply GELU out of place before cropping to keep graph buffer reuse safe.
+            x = ggml_gelu_erf(ctx->ggml_ctx, x);
+            x = ggml_view_3d(ctx->ggml_ctx, x, x->ne[0] - 1, x->ne[1], x->ne[2], x->nb[1], x->nb[2], 0);
+            return ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
+        }
+    };
+
+    struct Wav2Vec2FeedForward : public UnaryBlock {
+        Wav2Vec2FeedForward(const Wav2Vec2Config& config) {
+            blocks["intermediate_dense"] = std::make_shared<Linear>(config.embed_dim, config.embed_dim * 4);
+            blocks["output_dense"]       = std::make_shared<Linear>(config.embed_dim * 4, config.embed_dim);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto intermediate_dense = std::dynamic_pointer_cast<Linear>(blocks["intermediate_dense"]);
+            auto output_dense       = std::dynamic_pointer_cast<Linear>(blocks["output_dense"]);
+            x                       = intermediate_dense->forward(ctx, x);
+            x                       = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x                       = output_dense->forward(ctx, x);
+            return x;
+        }
+    };
+
+    struct Wav2Vec2EncoderLayer : public UnaryBlock {
+        bool do_stable_layer_norm;
+
+        Wav2Vec2EncoderLayer(const Wav2Vec2Config& config)
+            : do_stable_layer_norm(config.do_stable_layer_norm) {
+            blocks["attention"]        = std::make_shared<MultiheadAttention>(config.embed_dim, config.num_heads, true, true);
+            blocks["layer_norm"]       = std::make_shared<LayerNorm>(config.embed_dim);
+            blocks["feed_forward"]     = std::make_shared<Wav2Vec2FeedForward>(config);
+            blocks["final_layer_norm"] = std::make_shared<LayerNorm>(config.embed_dim);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto attention        = std::dynamic_pointer_cast<MultiheadAttention>(blocks["attention"]);
+            auto layer_norm       = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
+            auto feed_forward     = std::dynamic_pointer_cast<Wav2Vec2FeedForward>(blocks["feed_forward"]);
+            auto final_layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer_norm"]);
+
+            ggml_tensor* residual = x;
+            if (do_stable_layer_norm) {
+                x = layer_norm->forward(ctx, x);
+                x = attention->forward(ctx, x);
+                x = ggml_add(ctx->ggml_ctx, residual, x);
+                x = ggml_add(ctx->ggml_ctx, x, feed_forward->forward(ctx, final_layer_norm->forward(ctx, x)));
+            } else {
+                x = attention->forward(ctx, x);
+                x = ggml_add(ctx->ggml_ctx, residual, x);
+                x = layer_norm->forward(ctx, x);
+                x = final_layer_norm->forward(ctx, ggml_add(ctx->ggml_ctx, x, feed_forward->forward(ctx, x)));
+            }
+            return x;
+        }
+    };
+
+    struct Wav2Vec2Encoder : public GGMLBlock {
+        int num_layers;
+        bool do_stable_layer_norm;
+
+        Wav2Vec2Encoder(const Wav2Vec2Config& config)
+            : num_layers(config.num_layers), do_stable_layer_norm(config.do_stable_layer_norm) {
+            blocks["pos_conv_embed"] = std::make_shared<Wav2Vec2PositionalConvEmbedding>(config);
+            for (int i = 0; i < config.num_layers; ++i) {
+                blocks["layers." + std::to_string(i)] = std::make_shared<Wav2Vec2EncoderLayer>(config);
+            }
+            blocks["layer_norm"] = std::make_shared<LayerNorm>(config.embed_dim);
+        }
+
+        // For N == 1, all_layers stacks pre-layer states and the final state as [embed_dim, L, num_layers + 1].
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor** all_layers = nullptr) {
+            auto pos_conv_embed = std::dynamic_pointer_cast<Wav2Vec2PositionalConvEmbedding>(blocks["pos_conv_embed"]);
+            auto layer_norm     = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
+
+            std::vector<ggml_tensor*> collected;
+            if (all_layers != nullptr) {
+                collected.reserve(num_layers + 1);
+            }
+
+            x = ggml_add(ctx->ggml_ctx, x, pos_conv_embed->forward(ctx, x));
+            if (!do_stable_layer_norm) {
+                x = layer_norm->forward(ctx, x);
+            }
+            for (int i = 0; i < num_layers; ++i) {
+                if (all_layers != nullptr) {
+                    collected.push_back(x);
+                }
+                auto layer = std::dynamic_pointer_cast<Wav2Vec2EncoderLayer>(blocks["layers." + std::to_string(i)]);
+                x          = layer->forward(ctx, x);
+            }
+            if (do_stable_layer_norm) {
+                x = layer_norm->forward(ctx, x);
+            }
+            if (all_layers != nullptr) {
+                collected.push_back(x);
+                ggml_tensor* stack = collected[0];
+                for (size_t i = 1; i < collected.size(); ++i) {
+                    stack = ggml_concat(ctx->ggml_ctx, stack, collected[i], 2);
+                }
+                *all_layers = stack;
+            }
+            return x;
+        }
+    };
+
+    struct Wav2Vec2Model : public GGMLBlock {
+        Wav2Vec2Config config;
+
+        Wav2Vec2Model() = default;
+        Wav2Vec2Model(const Wav2Vec2Config& config_)
+            : config(config_) {
+            blocks["feature_extractor"]  = std::make_shared<Wav2Vec2FeatureEncoder>(config);
+            blocks["feature_projection"] = std::make_shared<Wav2Vec2FeatureProjection>(config);
+            blocks["encoder"]            = std::make_shared<Wav2Vec2Encoder>(config);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor** all_layers = nullptr) {
+            auto feature_extractor  = std::dynamic_pointer_cast<Wav2Vec2FeatureEncoder>(blocks["feature_extractor"]);
+            auto feature_projection = std::dynamic_pointer_cast<Wav2Vec2FeatureProjection>(blocks["feature_projection"]);
+            auto encoder            = std::dynamic_pointer_cast<Wav2Vec2Encoder>(blocks["encoder"]);
+
+            x = feature_extractor->forward(ctx, x);
+            x = feature_projection->forward(ctx, x);
+            x = encoder->forward(ctx, x, all_layers);
+            return x;
+        }
+    };
+
+    class Wav2Vec2ModelRunner : public GGMLRunner {
+    private:
+        Wav2Vec2Config config;
+
+    public:
+        Wav2Vec2Model model;
+        std::string weight_prefix;
+
+        Wav2Vec2ModelRunner(ggml_backend_t backend,
+                            const String2TensorStorage& tensor_storage_map      = {},
+                            const std::string prefix                            = "wav2vec2.",
+                            std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : GGMLRunner(backend, weight_manager),
+              config(Wav2Vec2Config::detect_from_weights(tensor_storage_map, prefix)),
+              model(config),
+              weight_prefix(prefix) {
+            // GGMLBlock appends its own separator; loader prefixes already include one.
+            std::string block_prefix = weight_prefix;
+            if (!block_prefix.empty() && block_prefix.back() == '.') {
+                block_prefix.pop_back();
+            }
+            model.init(params_ctx, tensor_storage_map, block_prefix);
+            LOG_INFO("%s", get_desc().c_str());
+        }
+
+        std::string get_desc() override {
+            return "wav2vec2";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
+            std::string block_prefix = weight_prefix;
+            if (!block_prefix.empty() && block_prefix.back() == '.') {
+                block_prefix.pop_back();
+            }
+            model.get_param_tensors(tensors, block_prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& waveform_tensor) {
+            ggml_cgraph* gf         = ggml_new_graph(compute_ctx);
+            ggml_tensor* waveform   = make_input(waveform_tensor);
+            auto runner_ctx         = get_context();
+            ggml_tensor* all_layers = nullptr;
+            model.forward(&runner_ctx, waveform, &all_layers);
+            GGML_ASSERT(all_layers != nullptr);
+            ggml_build_forward_expand(gf, all_layers);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(const int n_threads, const std::vector<float>& mono_waveform) {
+            GGML_ASSERT(!mono_waveform.empty());
+            const int64_t num_samples = (int64_t)mono_waveform.size();
+            sd::Tensor<float> waveform({num_samples, 1, 1});
+            std::copy(mono_waveform.begin(), mono_waveform.end(), waveform.data());
+            normalize(waveform.data(), num_samples);
+
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(waveform);
+            };
+            return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
+        }
+
+    private:
+        static void normalize(float* x, int64_t n) {
+            double mean = 0.0;
+            for (int64_t i = 0; i < n; ++i) {
+                mean += x[i];
+            }
+            mean /= n;
+            double var = 0.0;
+            for (int64_t i = 0; i < n; ++i) {
+                const double d = x[i] - mean;
+                var += d * d;
+            }
+            var /= n;
+            const float scale = (float)(1.0 / std::sqrt(var + 1e-7));
+            for (int64_t i = 0; i < n; ++i) {
+                x[i] = (float)((x[i] - mean) * scale);
+            }
+        }
+    };
+
+}  // namespace Wav2Vec2
+
+#endif  // __SD_MODEL_AUDIO_WAV2VEC2_HPP__

src/model/common/block.hpp

@@ -380,14 +380,14 @@ public:
         if (xtra_dim) {
             context->ne[0] = 320;  // reset dim to orig
         }
-        x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, inner_dim]
+        x = ggml_ext_attention_ext(ctx, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, inner_dim]
 
         if (has_ip && ctx->ip_context != nullptr && ctx->ip_scale != 0.0f) {
             auto to_k_ip = std::dynamic_pointer_cast<Linear>(blocks["to_k_ip"]);
             auto to_v_ip = std::dynamic_pointer_cast<Linear>(blocks["to_v_ip"]);
             auto k_ip    = to_k_ip->forward(ctx, ctx->ip_context);
             auto v_ip    = to_v_ip->forward(ctx, ctx->ip_context);
-            auto x_ip    = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_ip, v_ip, n_head, nullptr, false, ctx->flash_attn_enabled);
+            auto x_ip    = ggml_ext_attention_ext(ctx, q, k_ip, v_ip, n_head, nullptr, false, ctx->flash_attn_enabled);
             x            = ggml_add(ctx->ggml_ctx, x, ggml_scale(ctx->ggml_ctx, x_ip, ctx->ip_scale));
         }
 

src/model/common/ggml_block.hpp

@@ -206,7 +206,9 @@ public:
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         ggml_tensor* w            = params["weight"];
+        const float scale         = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale;
         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) {
@@ -220,6 +222,7 @@ public:
                 w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16);
             }
         }
+#endif
         ggml_tensor* b = nullptr;
         if (bias) {
             b = params["bias"];
@@ -237,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);
@@ -247,6 +251,7 @@ public:
                     x = cached->second;
                 }
             }
+#endif
             out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
                                                 x,
                                                 w,
@@ -367,6 +372,61 @@ public:
     }
 };
 
+class Conv1d : public UnaryBlock {
+protected:
+    int64_t in_channels;
+    int64_t out_channels;
+    int64_t groups;
+    int kernel_size;
+    int stride;
+    int padding;
+    int dilation;
+    bool bias;
+    bool force_prec_f32;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        ggml_type wtype  = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F16);
+        params["weight"] = ggml_new_tensor_3d(ctx, wtype, kernel_size, in_channels / groups, out_channels);
+        if (bias) {
+            params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+        }
+    }
+
+public:
+    Conv1d(int64_t in_channels,
+           int64_t out_channels,
+           int kernel_size,
+           int stride          = 1,
+           int padding         = 0,
+           int dilation        = 1,
+           int64_t groups      = 1,
+           bool bias           = true,
+           bool force_prec_f32 = false)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          groups(groups),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias),
+          force_prec_f32(force_prec_f32) {
+        GGML_ASSERT(in_channels > 0 && out_channels > 0 && groups > 0);
+        GGML_ASSERT(in_channels % groups == 0 && out_channels % groups == 0);
+        GGML_ASSERT(kernel_size > 0 && stride > 0 && padding >= 0 && dilation > 0);
+    }
+
+    std::string get_desc() override {
+        return "Conv1d";
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        GGML_ASSERT(x->ne[1] == in_channels);
+        return ggml_ext_conv_1d(ctx->ggml_ctx, x, params["weight"], bias ? params["bias"] : nullptr,
+                                stride, padding, dilation, groups, force_prec_f32);
+    }
+};
+
 class Conv2d : public UnaryBlock {
 protected:
     int64_t in_channels;
@@ -672,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);
     }
 };
 
@@ -765,7 +825,7 @@ public:
                 b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
             }
         }
-        return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups);
+        return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups, eps);
     }
 };
 
@@ -779,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");
@@ -870,7 +939,7 @@ public:
             v = v_proj->forward(ctx, x);
         }
 
-        x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, false);  // [N, n_token, embed_dim]
+        x = ggml_ext_attention_ext(ctx, q, k, v, n_head, mask, false);  // [N, n_token, embed_dim]
 
         x = out_proj->forward(ctx, x);  // [N, n_token, embed_dim]
         return x;

src/model/common/rope.hpp

@@ -818,8 +818,9 @@ namespace Rope {
                                                     int pw,
                                                     int bs,
                                                     int theta,
-                                                    const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids = gen_vid_ids(t, h, w, pt, ph, pw, bs);
+                                                    const std::vector<int>& axes_dim,
+                                                    int t_offset = 0) {
+        std::vector<std::vector<float>> ids = gen_vid_ids(t, h, w, pt, ph, pw, bs, t_offset);
         return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
     }
 
@@ -928,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,
@@ -1024,7 +1164,7 @@ namespace Rope {
         q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved);  // [N*n_head, L, d_head]
         k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved);  // [N*n_head, L, d_head]
 
-        auto x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, true, ctx->flash_attn_enabled, kv_scale);  // [N, L, n_head*d_head]
+        auto x = ggml_ext_attention_ext(ctx, q, k, v, n_head, mask, true, ctx->flash_attn_enabled, kv_scale);  // [N, L, n_head*d_head]
         return x;
     }
 };  // namespace Rope

src/model/diffusion/anima.hpp

@@ -237,8 +237,7 @@ namespace Anima {
                 }
                 auto q_rope = Rope::apply_rope(ctx->ggml_ctx, q4, pe_q, false);
                 auto k_rope = Rope::apply_rope(ctx->ggml_ctx, k4, pe_k, false);
-                attn_out    = ggml_ext_attention_ext(ctx->ggml_ctx,
-                                                     ctx->backend,
+                attn_out    = ggml_ext_attention_ext(ctx,
                                                      q_rope,
                                                      k_rope,
                                                      v4,
@@ -249,8 +248,7 @@ namespace Anima {
             } else {
                 auto q_flat = ggml_reshape_3d(ctx->ggml_ctx, q4, head_dim * num_heads, L_q, N);
                 auto k_flat = ggml_reshape_3d(ctx->ggml_ctx, k4, head_dim * num_heads, L_k, N);
-                attn_out    = ggml_ext_attention_ext(ctx->ggml_ctx,
-                                                     ctx->backend,
+                attn_out    = ggml_ext_attention_ext(ctx,
                                                      q_flat,
                                                      k_flat,
                                                      v,

src/model/diffusion/animatediff.hpp

@@ -61,7 +61,7 @@ namespace AnimateDiff {
             auto k = to_k->forward(ctx, x_pe);
             auto v = to_v->forward(ctx, x_pe);
 
-            auto a = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, (int)num_heads, nullptr, false);
+            auto a = ggml_ext_attention_ext(ctx, q, k, v, (int)num_heads, nullptr, false);
             return to_out->forward(ctx, a);
         }
     };

src/model/diffusion/ernie_image.hpp

@@ -183,7 +183,7 @@ namespace ErnieImage {
             k = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, k, 0, 2, 1, 3));  // [N, heads, S, head_dim]
             k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);
 
-            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled);  // [N, S, hidden_size]
+            x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled);  // [N, S, hidden_size]
             x = to_out_0->forward(ctx, x);
             return x;
         }

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);
@@ -504,6 +510,10 @@ namespace HiDreamO1 {
             vision_runner->set_flash_attention_enabled(enabled);
         }
 
+        void set_scale_overrides(float linear_scale, float attn_scale) override {
+            vision_runner->set_scale_overrides(linear_scale, attn_scale);
+        }
+
         void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
             vision_runner->set_weight_adapter(adapter);
         }
@@ -534,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);
@@ -608,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|>";
@@ -619,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/hunyuan.hpp

@@ -54,7 +54,7 @@ namespace Hunyuan {
             auto k       = qkv_vec[1];
             auto v       = qkv_vec[2];
 
-            auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled);
+            auto attn_out = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled);
             attn_out      = self_attn_proj->forward(ctx, attn_out);
 
             // adaLN_modulation

src/model/diffusion/krea2.hpp

@@ -232,8 +232,7 @@ namespace Krea2 {
             q          = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, q), head_dim_ * heads, Lq, N);
             k          = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, k), head_dim_ * kv_heads, Lk, N);
             v          = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, v), head_dim_ * kv_heads, Lk, N);
-            return ggml_ext_attention_ext(ctx->ggml_ctx,
-                                          ctx->backend,
+            return ggml_ext_attention_ext(ctx,
                                           q,
                                           k,
                                           v,

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

@@ -709,8 +709,7 @@ namespace LTXV {
                 k = apply_hidden_rope(ctx->ggml_ctx, k, k_pe, heads, dim_head, rope_interleaved);
             }
 
-            auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
-                                              ctx->backend,
+            auto out = ggml_ext_attention_ext(ctx,
                                               q,
                                               k,
                                               v,

src/model/diffusion/minimax_h3.hpp

@@ -215,8 +215,7 @@ namespace MiniMaxH3 {
                 q = attention_layout(ctx->ggml_ctx, q);
                 k = attention_layout(ctx->ggml_ctx, k);
             }
-            auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
-                                              ctx->backend,
+            auto out = ggml_ext_attention_ext(ctx,
                                               q,
                                               k,
                                               v,

src/model/diffusion/mmdit.hpp

@@ -365,8 +365,8 @@ public:
     ggml_tensor* forward(GGMLRunnerContext* ctx,
                          ggml_tensor* x) {
         auto qkv = pre_attention(ctx, x);
-        x        = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
-        x        = post_attention(ctx, x);                                                                                                           // [N, n_token, dim]
+        x        = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
+        x        = post_attention(ctx, x);                                                                                   // [N, n_token, dim]
         return x;
     }
 };
@@ -587,8 +587,8 @@ public:
             auto qkv2          = std::get<1>(qkv_intermediates);
             auto intermediates = std::get<2>(qkv_intermediates);
 
-            auto attn_out  = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);     // [N, n_token, dim]
-            auto attn2_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
+            auto attn_out  = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);     // [N, n_token, dim]
+            auto attn2_out = ggml_ext_attention_ext(ctx, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
             x              = post_attention_x(ctx,
                                               attn_out,
                                               attn2_out,
@@ -604,7 +604,7 @@ public:
             auto qkv               = qkv_intermediates.first;
             auto intermediates     = qkv_intermediates.second;
 
-            auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
+            auto attn_out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
             x             = post_attention(ctx,
                                            attn_out,
                                            intermediates[0],
@@ -648,7 +648,7 @@ block_mixing(GGMLRunnerContext* ctx,
         qkv.push_back(ggml_concat(ctx->ggml_ctx, context_qkv[i], x_qkv[i], 1));
     }
 
-    auto attn = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_context + n_token, hidden_size]
+    auto attn = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_context + n_token, hidden_size]
 
     auto context_attn = ggml_view_3d(ctx->ggml_ctx,
                                      attn,
@@ -680,7 +680,7 @@ block_mixing(GGMLRunnerContext* ctx,
     }
 
     if (x_block->self_attn) {
-        auto attn2 = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, hidden_size]
+        auto attn2 = ggml_ext_attention_ext(ctx, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, hidden_size]
 
         x = x_block->post_attention_x(ctx,
                                       x_attn,

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,9 +69,15 @@ 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;
+    // S2V audio, sd::Tensor layout: [dim, T_latent*4, layers].
+    const sd::Tensor<float>* audio_embed = nullptr;
 };
 
 struct HiDreamO1DiffusionExtra {
@@ -125,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/sensenova_u1.h

@@ -146,7 +146,7 @@ namespace SenseNovaU1 {
             auto x     = ggml_ext_timestep_embedding(ctx->ggml_ctx,
                                                      timesteps,
                                                      static_cast<int>(frequency_embedding_size),
-                                                     10000.f,
+                                                     10000,
                                                      1.f);
             x          = mlp_0->forward(ctx, x);
             x          = ggml_silu_inplace(ctx->ggml_ctx, x);

src/model/diffusion/wan.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_WAN_HPP__
 #define __SD_MODEL_DIFFUSION_WAN_HPP__
 
+#include <algorithm>
 #include <cinttypes>
 #include <map>
 #include <memory>
@@ -19,25 +20,30 @@ namespace WAN {
     constexpr int WAN_GRAPH_SIZE = 10240;
 
     struct WanConfig {
-        std::string model_type                 = "t2v";
-        std::tuple<int, int, int> patch_size   = {1, 2, 2};
-        int64_t text_len                       = 512;
-        int64_t in_dim                         = 16;
-        int64_t dim                            = 2048;
-        int64_t ffn_dim                        = 8192;
-        int freq_dim                           = 256;
-        int64_t text_dim                       = 4096;
-        int64_t out_dim                        = 16;
-        int64_t num_heads                      = 16;
-        int num_layers                         = 32;
-        int vace_layers                        = 0;
-        int64_t vace_in_dim                    = 96;
-        std::map<int, int> vace_layers_mapping = {};
-        bool qk_norm                           = true;
-        bool cross_attn_norm                   = true;
-        float eps                              = 1e-6f;
-        int64_t flf_pos_embed_token_number     = 0;
-        int theta                              = 10000;
+        std::string model_type                  = "t2v";
+        std::tuple<int, int, int> patch_size    = {1, 2, 2};
+        int64_t text_len                        = 512;
+        int64_t in_dim                          = 16;
+        int64_t dim                             = 2048;
+        int64_t ffn_dim                         = 8192;
+        int freq_dim                            = 256;
+        int64_t text_dim                        = 4096;
+        int64_t out_dim                         = 16;
+        int64_t num_heads                       = 16;
+        int num_layers                          = 32;
+        int vace_layers                         = 0;
+        int64_t vace_in_dim                     = 96;
+        std::map<int, int> vace_layers_mapping  = {};
+        int64_t audio_dim                       = 1024;
+        int num_audio_token                     = 4;  // excludes the learned padding token
+        std::vector<int> audio_inject_layers    = {};
+        std::map<int, int> audio_inject_mapping = {};  // block index -> injector index
+        std::string adain_mode                  = "attn_norm";
+        bool qk_norm                            = true;
+        bool cross_attn_norm                    = true;
+        float eps                               = 1e-6f;
+        int64_t flf_pos_embed_token_number      = 0;
+        int theta                               = 10000;
         // wan2.1 1.3B: 1536/12, wan2.1/2.2 14B: 5120/40, wan2.2 5B: 3074/24
         std::vector<int> axes_dim = {44, 42, 42};
         int64_t axes_dim_sum      = 128;
@@ -74,6 +80,10 @@ namespace WAN {
                 if (name.find("img_emb") != std::string::npos) {
                     config.model_type = "i2v";
                 }
+                if (name.find("audio_injector") != std::string::npos || name.find("casual_audio_encoder") != std::string::npos) {
+                    config.model_type          = "s2v";
+                    config.audio_inject_layers = {0, 4, 8, 12, 16, 20, 24, 27, 30, 33, 36, 39};
+                }
                 if (name.find("img_emb.emb_pos") != std::string::npos) {
                     config.flf_pos_embed_token_number = 514;
                 }
@@ -193,7 +203,7 @@ namespace WAN {
             k      = norm_k->forward(ctx, k);
             auto v = v_proj->forward(ctx, context);  // [N, n_context, dim]
 
-            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
+            x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
 
             x = o_proj->forward(ctx, x);  // [N, n_token, dim]
             return x;
@@ -255,8 +265,8 @@ namespace WAN {
             k_img      = norm_k_img->forward(ctx, k_img);
             auto v_img = v_img_proj->forward(ctx, context_img);  // [N, context_img_len, dim]
 
-            auto img_x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_img, v_img, num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
-            x          = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled);          // [N, n_token, dim]
+            auto img_x = ggml_ext_attention_ext(ctx, q, k_img, v_img, num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
+            x          = ggml_ext_attention_ext(ctx, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled);          // [N, n_token, dim]
 
             x = ggml_add(ctx->ggml_ctx, x, img_x);
 
@@ -265,6 +275,13 @@ namespace WAN {
         }
     };
 
+}  // namespace WAN
+
+// Audio injection reuses WanT2VCrossAttention defined above.
+#include "model/diffusion/wan_audio.hpp"
+
+namespace WAN {
+
     static ggml_tensor* modulate_add(ggml_context* ctx, ggml_tensor* x, ggml_tensor* e) {
         // x: [N, n_token, dim]
         // e: [N, 1, dim] or [N, T, 1, dim]
@@ -532,6 +549,13 @@ namespace WAN {
     protected:
         WanConfig config;
 
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            if (config.model_type == "s2v") {
+                enum ggml_type wtype                 = GGML_TYPE_F32;  // elementwise add vs F32 activations
+                params["trainable_cond_mask.weight"] = ggml_new_tensor_2d(ctx, wtype, config.dim, 3);
+            }
+        }
+
     public:
         Wan() {}
         Wan(WanConfig config)
@@ -554,7 +578,7 @@ namespace WAN {
 
             // blocks
             for (int i = 0; i < config.num_layers; i++) {
-                auto block                            = std::shared_ptr<GGMLBlock>(new WanAttentionBlock(config.model_type == "t2v",
+                auto block                            = std::shared_ptr<GGMLBlock>(new WanAttentionBlock(config.model_type != "i2v",
                                                                                                          config.dim,
                                                                                                          config.ffn_dim,
                                                                                                          config.num_heads,
@@ -595,6 +619,14 @@ namespace WAN {
 
                 blocks["vace_patch_embedding"] = std::shared_ptr<GGMLBlock>(new Conv3d(config.vace_in_dim, config.dim, config.patch_size, config.patch_size));
             }
+
+            if (config.model_type == "s2v") {
+                blocks["casual_audio_encoder"] = std::make_shared<WanCausalAudioEncoder>(config.audio_dim, config.dim, config.num_audio_token);
+                blocks["audio_injector"]       = std::make_shared<WanAudioInjector>(config.dim, config.num_heads, (int)config.audio_inject_layers.size(), config.qk_norm, config.eps);
+                for (size_t i = 0; i < config.audio_inject_layers.size(); i++) {
+                    config.audio_inject_mapping[config.audio_inject_layers[i]] = (int)i;
+                }
+            }
         }
 
         ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
@@ -642,18 +674,24 @@ namespace WAN {
                                   ggml_tensor* timestep,
                                   ggml_tensor* context,
                                   ggml_tensor* pe,
-                                  ggml_tensor* clip_fea     = nullptr,
-                                  ggml_tensor* vace_context = nullptr,
-                                  float vace_strength       = 1.f,
-                                  int64_t N                 = 1) {
+                                  ggml_tensor* clip_fea         = nullptr,
+                                  ggml_tensor* vace_context     = nullptr,
+                                  float vace_strength           = 1.f,
+                                  int64_t N                     = 1,
+                                  ggml_tensor* audio_embed      = nullptr,
+                                  ggml_tensor* reference_latent = nullptr) {
             // x: [N*C, T, H, W], C => in_dim
             // vace_context: [N*vace_in_dim, T, H, W]
             // timestep: [N,] or [T]
             // context: [N, L, text_dim]
-            // return: [N, t_len*h_len*w_len, out_dim*pt*ph*pw]
+            // audio_embed: [layers, T*4, audio_dim]
+            // reference_latent: [N*C, T_ref, H, W]
+            // return: [N, (t_len [+ t_ref_len]) * h_len*w_len, out_dim*pt*ph*pw]
 
             GGML_ASSERT(N == 1);
 
+            int64_t T = x->ne[2];
+
             auto patch_embedding = std::dynamic_pointer_cast<Conv3d>(blocks["patch_embedding"]);
 
             auto text_embedding_0 = std::dynamic_pointer_cast<Linear>(blocks["text_embedding.0"]);
@@ -670,6 +708,40 @@ namespace WAN {
             x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3] / N, N);  // [N, dim, t_len*h_len*w_len]
             x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [N, t_len*h_len*w_len, dim]
 
+            ggml_tensor* audio_local  = nullptr;
+            ggml_tensor* audio_global = nullptr;
+            int64_t seq_len           = x->ne[1];
+            int64_t t_ref_len         = 0;
+            if (config.model_type == "s2v") {
+                if (audio_embed != nullptr) {
+                    GGML_ASSERT(audio_embed->ne[1] == T * 4);
+                    auto audio_encoder = std::dynamic_pointer_cast<WanCausalAudioEncoder>(blocks["casual_audio_encoder"]);
+                    auto audio_emb     = audio_encoder->forward(ctx, audio_embed);
+                    audio_local        = audio_emb.first;
+                    audio_global       = audio_emb.second;
+                    GGML_ASSERT(audio_local->ne[2] == T);
+                }
+
+                // video tokens get cond_mask[0], reference tokens cond_mask[1]
+                auto cond_mask = params["trainable_cond_mask.weight"];
+                auto cm0       = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 0, 1), config.dim, 1, 1);
+                x              = ggml_add(ctx->ggml_ctx, x, cm0);
+
+                if (reference_latent != nullptr) {
+                    t_ref_len = reference_latent->ne[2];
+                    auto ref  = patch_embedding->forward(ctx, reference_latent);
+                    ref       = ggml_reshape_3d(ctx->ggml_ctx, ref, ref->ne[0] * ref->ne[1] * ref->ne[2], ref->ne[3] / N, N);
+                    ref       = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ref, 1, 0, 2, 3));  // [N, t_ref*h_len*w_len, dim]
+                    auto cm1  = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 1, 2), config.dim, 1, 1);
+                    ref       = ggml_add(ctx->ggml_ctx, ref, cm1);
+                    x         = ggml_concat(ctx->ggml_ctx, x, ref, 1);
+
+                    // Reference tokens use timestep 0.
+                    GGML_ASSERT(timestep->ne[0] == T);
+                    timestep = ggml_ext_pad(ctx->ggml_ctx, timestep, (int)t_ref_len, 0, 0, 0);
+                }
+            }
+
             // time_embedding
             auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, config.freq_dim);
             e      = time_embedding_0->forward(ctx, e);
@@ -714,6 +786,11 @@ namespace WAN {
 
             auto x_orig = x;
 
+            std::shared_ptr<WanAudioInjector> audio_injector;
+            if (audio_local != nullptr) {
+                audio_injector = std::dynamic_pointer_cast<WanAudioInjector>(blocks["audio_injector"]);
+            }
+
             for (int i = 0; i < config.num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<WanAttentionBlock>(blocks["blocks." + std::to_string(i)]);
 
@@ -731,6 +808,13 @@ namespace WAN {
                     c_skip      = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength);
                     x           = ggml_add(ctx->ggml_ctx, x, c_skip);
                 }
+
+                if (audio_injector != nullptr) {
+                    auto inject_iter = config.audio_inject_mapping.find(i);
+                    if (inject_iter != config.audio_inject_mapping.end()) {
+                        x = audio_injector->forward(ctx, x, seq_len, T, inject_iter->second, audio_local, audio_global);
+                    }
+                }
                 sd::ggml_graph_cut::mark_graph_cut(x, "wan.blocks." + std::to_string(i), "x");
                 if (c != nullptr) {
                     sd::ggml_graph_cut::mark_graph_cut(c, "wan.blocks." + std::to_string(i), "c");
@@ -747,11 +831,13 @@ namespace WAN {
                              ggml_tensor* timestep,
                              ggml_tensor* context,
                              ggml_tensor* pe,
-                             ggml_tensor* clip_fea        = nullptr,
-                             ggml_tensor* time_dim_concat = nullptr,
-                             ggml_tensor* vace_context    = nullptr,
-                             float vace_strength          = 1.f,
-                             int64_t N                    = 1) {
+                             ggml_tensor* clip_fea         = nullptr,
+                             ggml_tensor* time_dim_concat  = nullptr,
+                             ggml_tensor* vace_context     = nullptr,
+                             float vace_strength           = 1.f,
+                             int64_t N                     = 1,
+                             ggml_tensor* audio_embed      = nullptr,
+                             ggml_tensor* reference_latent = nullptr) {
             // Forward pass of DiT.
             // x: [N*C, T, H, W]
             // timestep: [N,]
@@ -779,7 +865,12 @@ namespace WAN {
                 t_len           = ((x->ne[2] + (std::get<0>(config.patch_size) / 2)) / std::get<0>(config.patch_size));
             }
 
-            auto out = forward_orig(ctx, x, timestep, context, pe, clip_fea, vace_context, vace_strength, N);  // [N, t_len*h_len*w_len, pt*ph*pw*C]
+            auto out = forward_orig(ctx, x, timestep, context, pe, clip_fea, vace_context, vace_strength, N, audio_embed, reference_latent);  // [N, (t_len [+t_ref]) *h_len*w_len, pt*ph*pw*C]
+
+            if (reference_latent != nullptr) {
+                // Exclude reference tokens from the generated video.
+                out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, t_len * h_len * w_len);
+            }
 
             out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len);  // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
 
@@ -839,7 +930,10 @@ namespace WAN {
                     config.text_len  = 512;
                 }
             } else if (config.num_layers == 40) {
-                if (config.model_type == "t2v") {
+                if (version == VERSION_WAN2_2_S2V) {
+                    desc          = "Wan2.2-S2V-14B";
+                    config.in_dim = 16;
+                } else if (config.model_type == "t2v") {
                     if (version == VERSION_WAN2_2_I2V) {
                         desc          = "Wan2.2-I2V-14B";
                         config.in_dim = 36;
@@ -891,7 +985,9 @@ namespace WAN {
                                  const sd::Tensor<float>& c_concat_tensor        = {},
                                  const sd::Tensor<float>& time_dim_concat_tensor = {},
                                  const sd::Tensor<float>& vace_context_tensor    = {},
-                                 float vace_strength                             = 1.f) {
+                                 float vace_strength                             = 1.f,
+                                 const sd::Tensor<float>& audio_embed_tensor     = {},
+                                 const sd::Tensor<float>& ref_latent_tensor      = {}) {
             ggml_cgraph* gf = new_graph_custom(WAN_GRAPH_SIZE);
 
             ggml_tensor* x               = make_input(x_tensor);
@@ -901,16 +997,33 @@ namespace WAN {
             ggml_tensor* c_concat        = make_optional_input(c_concat_tensor);
             ggml_tensor* time_dim_concat = make_optional_input(time_dim_concat_tensor);
             ggml_tensor* vace_context    = make_optional_input(vace_context_tensor);
+            ggml_tensor* audio_embed     = make_optional_input(audio_embed_tensor);
+            ggml_tensor* ref_latent      = make_optional_input(ref_latent_tensor);
 
-            pe_vec      = Rope::gen_wan_pe(static_cast<int>(x->ne[2]),
-                                           static_cast<int>(x->ne[1]),
-                                           static_cast<int>(x->ne[0]),
-                                           std::get<0>(config.patch_size),
-                                           std::get<1>(config.patch_size),
-                                           std::get<2>(config.patch_size),
-                                           1,
-                                           config.theta,
-                                           config.axes_dim);
+            pe_vec = Rope::gen_wan_pe(static_cast<int>(x->ne[2]),
+                                      static_cast<int>(x->ne[1]),
+                                      static_cast<int>(x->ne[0]),
+                                      std::get<0>(config.patch_size),
+                                      std::get<1>(config.patch_size),
+                                      std::get<2>(config.patch_size),
+                                      1,
+                                      config.theta,
+                                      config.axes_dim);
+            if (ref_latent != nullptr) {
+                // Match S2V's reference-frame temporal offset.
+                int t_start = std::max(30, static_cast<int>(x->ne[2]) + 9);
+                auto ref_pe = Rope::gen_wan_pe(static_cast<int>(ref_latent->ne[2]),
+                                               static_cast<int>(ref_latent->ne[1]),
+                                               static_cast<int>(ref_latent->ne[0]),
+                                               std::get<0>(config.patch_size),
+                                               std::get<1>(config.patch_size),
+                                               std::get<2>(config.patch_size),
+                                               1,
+                                               config.theta,
+                                               config.axes_dim,
+                                               t_start);
+                pe_vec.insert(pe_vec.end(), ref_pe.begin(), ref_pe.end());
+            }
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             // LOG_VERBOSE("pos_len %d", pos_len);
             auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
@@ -933,7 +1046,10 @@ namespace WAN {
                                            clip_fea,
                                            time_dim_concat,
                                            vace_context,
-                                           vace_strength);
+                                           vace_strength,
+                                           1,
+                                           audio_embed,
+                                           ref_latent);
 
             ggml_build_forward_expand(gf, out);
 
@@ -948,9 +1064,11 @@ namespace WAN {
                                   const sd::Tensor<float>& c_concat        = {},
                                   const sd::Tensor<float>& time_dim_concat = {},
                                   const sd::Tensor<float>& vace_context    = {},
-                                  float vace_strength                      = 1.f) {
+                                  float vace_strength                      = 1.f,
+                                  const sd::Tensor<float>& audio_embed     = {},
+                                  const sd::Tensor<float>& ref_latent      = {}) {
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
+                return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength, audio_embed, ref_latent);
             };
 
             return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
@@ -961,6 +1079,12 @@ namespace WAN {
             GGML_ASSERT(diffusion_params.x != nullptr);
             GGML_ASSERT(diffusion_params.timesteps != nullptr);
             const auto* extra = diffusion_extra_as<WanDiffusionExtra>(diffusion_params);
+            static const std::vector<sd::Tensor<float>> no_ref_latents;
+            const auto& ref_latents = config.model_type == "s2v" && diffusion_params.ref_latents != nullptr
+                                          ? *diffusion_params.ref_latents
+                                          : no_ref_latents;
+            const sd::Tensor<float> empty_tensor;
+            const sd::Tensor<float>& ref_latent = ref_latents.empty() ? empty_tensor : ref_latents[0];
             return compute(n_threads,
                            *diffusion_params.x,
                            *diffusion_params.timesteps,
@@ -969,7 +1093,9 @@ namespace WAN {
                            tensor_or_empty(diffusion_params.c_concat),
                            sd::Tensor<float>(),
                            tensor_or_empty(extra->vace_context),
-                           extra->vace_strength);
+                           extra->vace_strength,
+                           tensor_or_empty(extra->audio_embed),
+                           ref_latent);
         }
 
         void test() {

src/model/diffusion/wan_audio.hpp

@@ -0,0 +1,215 @@
+#ifndef __SD_MODEL_DIFFUSION_WAN_AUDIO_HPP__
+#define __SD_MODEL_DIFFUSION_WAN_AUDIO_HPP__
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "model/common/ggml_block.hpp"
+
+namespace WAN {
+
+    class WanCausalConv1d : public UnaryBlock {
+    private:
+        int kernel_size_;
+
+    public:
+        WanCausalConv1d(int64_t in_dim,
+                        int64_t out_dim,
+                        int kernel_size = 3,
+                        int stride      = 1)
+            : kernel_size_(kernel_size) {
+            blocks["conv"] = std::make_shared<Conv1d>(in_dim, out_dim, kernel_size, stride, 0, 1, 1, true, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            // Replicate the first sample for causal left padding.
+            if (kernel_size_ > 1) {
+                auto first = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, 1);
+                for (int i = 0; i < kernel_size_ - 1; i++) {
+                    x = ggml_concat(ctx->ggml_ctx, first, x, 0);
+                }
+            }
+            return std::dynamic_pointer_cast<Conv1d>(blocks["conv"])->forward(ctx, x);
+        }
+    };
+
+    class WanMotionEncoder : public GGMLBlock {
+    private:
+        int64_t hidden_dim_;
+        int num_token_;
+        bool need_global_;
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            // The padding token is combined with F32 activations.
+            params["padding_tokens"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_dim_);
+        }
+
+        ggml_tensor* conv_norm_silu(GGMLRunnerContext* ctx,
+                                    ggml_tensor* x,
+                                    const std::string& conv_key,
+                                    const std::string& norm_key,
+                                    bool to_conv_layout) {
+            x = std::dynamic_pointer_cast<WanCausalConv1d>(blocks[conv_key])->forward(ctx, x);
+            x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
+            x = std::dynamic_pointer_cast<LayerNorm>(blocks[norm_key])->forward(ctx, x);
+            x = ggml_silu(ctx->ggml_ctx, x);
+            if (to_conv_layout) {
+                x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            }
+            return x;
+        }
+
+    public:
+        WanMotionEncoder(int64_t in_dim,
+                         int64_t hidden_dim,
+                         int num_token,
+                         bool need_global = true)
+            : hidden_dim_(hidden_dim), num_token_(num_token), need_global_(need_global) {
+            blocks["conv1_local"] = std::make_shared<WanCausalConv1d>(in_dim, hidden_dim / 4 * num_token);
+            if (need_global) {
+                blocks["conv1_global"] = std::make_shared<WanCausalConv1d>(in_dim, hidden_dim / 4);
+            }
+            blocks["norm1"] = std::make_shared<LayerNorm>(hidden_dim / 4, 1e-6f, false);
+            blocks["conv2"] = std::make_shared<WanCausalConv1d>(hidden_dim / 4, hidden_dim / 2, 3, 2);
+            blocks["norm2"] = std::make_shared<LayerNorm>(hidden_dim / 2, 1e-6f, false);
+            blocks["conv3"] = std::make_shared<WanCausalConv1d>(hidden_dim / 2, hidden_dim, 3, 2);
+            blocks["norm3"] = std::make_shared<LayerNorm>(hidden_dim, 1e-6f, false);
+            if (need_global) {
+                blocks["final_linear"] = std::make_shared<Linear>(hidden_dim, hidden_dim);
+            }
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto local = std::dynamic_pointer_cast<WanCausalConv1d>(blocks["conv1_local"])->forward(ctx, x);
+            auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
+            std::vector<ggml_tensor*> tokens;
+            // Each token group is normalized independently over channels.
+            for (auto& group : ggml_ext_chunk(ctx->ggml_ctx, local, num_token_, 1)) {
+                ggml_tensor* s = ggml_permute(ctx->ggml_ctx, group, 1, 0, 2, 3);
+                s              = norm1->forward(ctx, s);
+                s              = ggml_silu(ctx->ggml_ctx, s);
+                s              = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 0, 2, 3));
+                s              = conv_norm_silu(ctx, s, "conv2", "norm2", true);
+                s              = conv_norm_silu(ctx, s, "conv3", "norm3", false);
+                tokens.push_back(ggml_reshape_3d(ctx->ggml_ctx, s, s->ne[0], 1, s->ne[1]));
+            }
+            auto padding = ggml_reshape_3d(ctx->ggml_ctx, params["padding_tokens"], hidden_dim_, 1, 1);
+            padding      = ggml_repeat(ctx->ggml_ctx, padding, tokens[0]);
+            tokens.push_back(padding);
+            ggml_tensor* local_out = ggml_ext_vec_concat(ctx->ggml_ctx, tokens, 1);
+
+            if (!need_global_) {
+                return {local_out, nullptr};
+            }
+            ggml_tensor* g = conv_norm_silu(ctx, x, "conv1_global", "norm1", true);
+            g              = conv_norm_silu(ctx, g, "conv2", "norm2", true);
+            g              = conv_norm_silu(ctx, g, "conv3", "norm3", false);
+            g              = std::dynamic_pointer_cast<Linear>(blocks["final_linear"])->forward(ctx, g);
+            return {local_out, g};
+        }
+    };
+
+    class WanCausalAudioEncoder : public GGMLBlock {
+    private:
+        int num_layers_;
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+            // Preserve the checkpoint shape for loading; layer mixing requires F32.
+            auto it = tensor_storage_map.find(prefix + "weights");
+            if (it != tensor_storage_map.end()) {
+                params["weights"] = ggml_new_tensor(ctx, GGML_TYPE_F32, it->second.n_dims, it->second.ne);
+            } else {
+                params["weights"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_layers_);
+            }
+        }
+
+    public:
+        WanCausalAudioEncoder(int64_t audio_dim,
+                              int64_t dim,
+                              int num_token,
+                              int num_layers = 25)
+            : num_layers_(num_layers) {
+            blocks["encoder"] = std::make_shared<WanMotionEncoder>(audio_dim, dim, num_token, true);
+        }
+
+        // features: [layers, frames, audio_dim]; outputs: [T, tokens+1, dim] and [T, dim].
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* features) {
+            auto weights = ggml_silu(ctx->ggml_ctx, params["weights"]);
+            auto x       = ggml_mul(ctx->ggml_ctx, features, ggml_reshape_3d(ctx->ggml_ctx, weights, 1, 1, num_layers_));
+            x            = ggml_div(ctx->ggml_ctx, x, ggml_sum(ctx->ggml_ctx, weights));
+            // Move the layer axis to ggml dimension 0 for reduction.
+            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+            x = ggml_sum_rows(ctx->ggml_ctx, x);
+            x = ggml_reshape_2d(ctx->ggml_ctx, x, x->ne[1], x->ne[2]);
+            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            return std::dynamic_pointer_cast<WanMotionEncoder>(blocks["encoder"])->forward(ctx, x);
+        }
+    };
+
+    class WanAudioInjector : public GGMLBlock {
+    private:
+        int64_t dim_;
+
+    public:
+        WanAudioInjector(int64_t dim,
+                         int64_t num_heads,
+                         int count,
+                         bool qk_norm = true,
+                         float eps    = 1e-6f)
+            : dim_(dim) {
+            for (int i = 0; i < count; i++) {
+                blocks["injector." + std::to_string(i)] =
+                    std::make_shared<WanT2VCrossAttention>(dim, num_heads, qk_norm, eps);
+                blocks["injector_adain_layers." + std::to_string(i) + ".linear"] =
+                    std::make_shared<Linear>(dim, dim * 2);
+            }
+            // S2V AdaLayerNorm uses its own epsilon, independent of attention norms.
+            blocks["adain_norm"] = std::make_shared<LayerNorm>(dim, 1e-5f, false);
+        }
+
+        // Inject into the video prefix; trailing reference tokens pass through unchanged.
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             int64_t seq_len,
+                             int64_t T,
+                             int injector_id,
+                             ggml_tensor* audio_local,
+                             ggml_tensor* audio_global) {
+            int64_t n_tok   = seq_len / T;
+            int64_t n_token = x->ne[1];
+
+            auto adain_linear = std::dynamic_pointer_cast<Linear>(blocks["injector_adain_layers." + std::to_string(injector_id) + ".linear"]);
+            auto injector     = std::dynamic_pointer_cast<WanT2VCrossAttention>(blocks["injector." + std::to_string(injector_id)]);
+            auto adain_norm   = std::dynamic_pointer_cast<LayerNorm>(blocks["adain_norm"]);
+
+            auto temb  = ggml_silu(ctx->ggml_ctx, audio_global);
+            temb       = adain_linear->forward(ctx, temb);
+            auto shift = ggml_ext_slice(ctx->ggml_ctx, temb, 0, 0, dim_);
+            auto scale = ggml_ext_slice(ctx->ggml_ctx, temb, 0, dim_, dim_ * 2);
+            shift      = ggml_reshape_3d(ctx->ggml_ctx, shift, dim_, 1, T);
+            scale      = ggml_reshape_3d(ctx->ggml_ctx, scale, dim_, 1, T);
+
+            auto x_vid = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, seq_len);
+            auto h     = ggml_reshape_3d(ctx->ggml_ctx, x_vid, dim_, n_tok, T);
+            h          = adain_norm->forward(ctx, h);
+            h          = ggml_add(ctx->ggml_ctx, h, ggml_mul(ctx->ggml_ctx, h, scale));
+            h          = ggml_add(ctx->ggml_ctx, h, shift);
+
+            auto res = injector->forward(ctx, h, audio_local, 0);
+            res      = ggml_reshape_2d(ctx->ggml_ctx, res, dim_, seq_len);
+
+            auto x_head = ggml_add(ctx->ggml_ctx, x_vid, res);
+            if (seq_len < n_token) {
+                auto x_tail = ggml_ext_slice(ctx->ggml_ctx, x, 1, seq_len, n_token);
+                return ggml_concat(ctx->ggml_ctx, x_head, x_tail, 1);
+            }
+            return x_head;
+        }
+    };
+
+}  // namespace WAN
+
+#endif  // __SD_MODEL_DIFFUSION_WAN_AUDIO_HPP__

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;
@@ -139,7 +153,8 @@ namespace LLM {
 
         static LLMConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
                                              const std::string& prefix,
-                                             LLMArch arch) {
+                                             LLMArch arch,
+                                             bool& enable_vision) {
             LLMConfig config;
             config.arch = arch;
             if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
@@ -210,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;
@@ -230,8 +270,9 @@ namespace LLM {
                 config.num_experts_per_tok     = 4;
             }
 
-            config.num_layers          = 0;
-            int detected_vision_layers = 0;
+            config.num_layers             = 0;
+            int detected_vision_layers    = 0;
+            bool out_hidden_size_detected = false;
             for (const auto& [name, tensor_storage] : tensor_storage_map) {
                 if (!starts_with(name, prefix)) {
                     continue;
@@ -277,6 +318,7 @@ namespace LLM {
                     if (ends_with(name, "visual.merger.linear_fc2.weight") ||
                         ends_with(name, "visual.merger.mlp.2.weight")) {
                         config.vision.out_hidden_size = tensor_storage.ne[1];
+                        out_hidden_size_detected      = true;
                     }
                     continue;
                 }
@@ -330,6 +372,20 @@ namespace LLM {
                         config.vocab_size,
                         config.hidden_size,
                         config.intermediate_size);
+            if (enable_vision && !config.have_vision_weight) {
+                LOG_WARN("no vision weights detected, vision disabled");
+                enable_vision = false;
+            }
+            // The default would reject valid models, so only compare a detected dim.
+            if (enable_vision && out_hidden_size_detected &&
+                config.vision.out_hidden_size != config.hidden_size) {
+                LOG_ERROR("vision projector output size (%" PRId64 ") does not match LLM hidden size (%" PRId64
+                          "), "
+                          "the vision weights (mmproj) likely belong to a different LLM variant, vision disabled",
+                          config.vision.out_hidden_size,
+                          config.hidden_size);
+                enable_vision = false;
+            }
             return config;
         }
     };
@@ -401,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;
@@ -587,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;
     }
@@ -1104,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;
 
@@ -1129,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) {
@@ -1143,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);
         }
 
@@ -1168,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]
@@ -1318,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);
@@ -1359,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->ggml_ctx, ctx->backend, 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]
@@ -1414,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,
@@ -1467,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);
@@ -1632,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"]);
@@ -1886,12 +2207,8 @@ namespace LLM {
                   bool enable_vision_                                 = false,
                   std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
             : GGMLRunner(backend, weight_manager),
-              config(LLMConfig::detect_from_weights(tensor_storage_map, prefix, arch)),
+              config(LLMConfig::detect_from_weights(tensor_storage_map, prefix, arch, enable_vision_)),
               enable_vision(enable_vision_) {
-            if (enable_vision && !config.have_vision_weight) {
-                LOG_WARN("no vision weights detected, vision disabled");
-                enable_vision = false;
-            }
             if (enable_vision) {
                 LOG_VERBOSE("enable llm vision");
                 if (config.llama_cpp_style) {
@@ -1976,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;
@@ -2028,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;
@@ -2101,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;
@@ -2338,7 +2673,7 @@ namespace LLM {
     };
 
     struct LLMEmbedder {
-        std::shared_ptr<BPETokenizer> tokenizer;
+        std::shared_ptr<Tokenizer> tokenizer;
         LLMRunner model;
 
         LLMEmbedder(LLMArch arch,
@@ -2346,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>();
+                }
             }
         }
 
@@ -2389,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

@@ -251,7 +251,7 @@ public:
 
         k = ggml_ext_scale(ctx->ggml_ctx, k, ::sqrtf(static_cast<float>(d_head)), true);
 
-        x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask);  // [N, n_token, d_head * n_head]
+        x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask);  // [N, n_token, d_head * n_head]
 
         x = out_proj->forward(ctx, x);  // [N, n_token, model_dim]
         return {x, past_bias};
@@ -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/auto_encoder_kl.hpp

@@ -142,7 +142,7 @@ public:
             v = ggml_reshape_3d(ctx->ggml_ctx, v, c, h * w, n);                        // [N, h * w, in_channels]
         }
 
-        h_ = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
+        h_ = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
 
         if (use_linear) {
             h_ = proj_out->forward(ctx, h_);  // [N, h * w, in_channels]

src/model/vae/hunyuan_vae.hpp

@@ -193,7 +193,7 @@ namespace Hunyuan {
             v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b);                                  // [b, c, t*h*w]
             v = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3));  // [b, t*h*w, c]
 
-            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);  // [b, t*h*w, c]
+            x = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);  // [b, t*h*w, c]
 
             x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [b, c, t*h*w]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b);                         // [b*c, t, h, w]

src/model/vae/mage_vae.hpp

@@ -253,7 +253,7 @@ namespace MageVAE {
             q      = to_patches(ctx->ggml_ctx, q);
             k      = to_patches(ctx->ggml_ctx, k);
             v      = to_patches(ctx->ggml_ctx, v);
-            h      = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
+            h      = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
             h      = from_patches(ctx->ggml_ctx, h, np, batch, hp, wp);
             if (pad_h > 0) {
                 h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, height);

src/model/vae/minimax_h3_audio_vae.hpp

@@ -174,8 +174,7 @@ namespace MiniMaxH3 {
             auto mask        = ggml_diag_mask_inf(ctx->ggml_ctx,
                                                   ggml_ext_zeros(ctx->ggml_ctx, sequence, sequence, 1, 1),
                                                   0);
-            auto attn_out    = ggml_ext_attention_ext(ctx->ggml_ctx,
-                                                      ctx->backend,
+            auto attn_out    = ggml_ext_attention_ext(ctx,
                                                       q,
                                                       k,
                                                       v,

src/model/vae/minimax_h3_vae.hpp

@@ -291,8 +291,7 @@ namespace MiniMaxH3VAE {
             k                   = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f);
             q                   = apply_partial_rope(ctx->ggml_ctx, q, pe);
             k                   = apply_partial_rope(ctx->ggml_ctx, k, pe);
-            auto out            = ggml_ext_attention_ext(ctx->ggml_ctx,
-                                                         ctx->backend,
+            auto out            = ggml_ext_attention_ext(ctx,
                                                          q,
                                                          k,
                                                          v,

src/model/vae/tae.hpp

@@ -564,7 +564,7 @@ public:
             int64_t chunk_frames = 5 * decoder->t_upscale;
             int64_t pad          = (chunk_frames - (num_frames % chunk_frames)) % chunk_frames;
 
-            result = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, result, 0, 0, 0, 0, 0, 0, 0, pad, false, false);
+            result = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, result, 0, 0, 0, 0, 0, 0, 0, static_cast<int>(pad), false, false);
 
             int64_t num_chunks                  = (num_frames + pad) / chunk_frames;
             auto to_trim                        = decoder->t_upscale - 1;

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++;
             }
         }
@@ -615,8 +631,8 @@ namespace WAN {
             auto v = qkv_vec[2];
             v      = ggml_reshape_3d(ctx->ggml_ctx, v, h * w, c, n);  // [t, c, h * w]
 
-            v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3));                            // [t, h * w, c]
-            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);  // [t, h * w, c]
+            v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3));    // [t, h * w, c]
+            x = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);  // [t, h * w, c]
 
             x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [t, c, h * w]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, c, n);                             // [t, c, h, w]
@@ -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_component.h

@@ -10,6 +10,7 @@ enum class ModelComponent {
     VAE,
     PreviewVAE,
     AudioVAE,
+    AudioEncoder,
     ControlNet,
     PhotoMaker,
     PuLID,
@@ -38,6 +39,8 @@ inline const char* model_component_name(ModelComponent component) {
             return "preview VAE";
         case ModelComponent::AudioVAE:
             return "audio VAE";
+        case ModelComponent::AudioEncoder:
+            return "audio encoder";
         case ModelComponent::ControlNet:
             return "ControlNet";
         case ModelComponent::PhotoMaker:

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,