fix: avoid passing ggml logs as format strings (#2002)

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

AGENTS.md

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

CMakeLists.txt

@@ -95,6 +95,8 @@ option(SD_MUSA                       "sd: musa backend" OFF)
 option(SD_BUILD_SHARED_LIBS          "sd: build shared libs" OFF)
 option(SD_BUILD_SHARED_GGML_LIB      "sd: build ggml as a separate shared lib" OFF)
 option(SD_USE_SYSTEM_GGML            "sd: use system-installed GGML library" OFF)
+option(SD_USE_UPSTREAM_GGML          "sd: build with upstream GGML instead of the patched GGML extensions" OFF)
+set(SD_GGML_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/ggml" CACHE PATH "sd: ggml source directory (also supplies private headers for system ggml)")
 #option(SD_BUILD_SERVER               "sd: build server example"                           ON)
 
 set(CMAKE_C_STANDARD 11)
@@ -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

@@ -46,6 +46,21 @@ Some older code in the project may not fully follow the current conventions. Ple
 
 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)

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

@@ -154,6 +154,11 @@ 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:

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`.

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/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/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/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/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

@@ -410,6 +410,11 @@ ArgOptions SDContextParams::get_options() {
          "path to the llm text encoder. For example: (qwenvl2.5 for qwen-image, mistral-small3.2 for flux2, ...)",
          0,
          &llm_path},
+        {"",
+         "--tokenizer",
+         "tokenizer.json path, or comma-separated main=FILE,clip-l=FILE,clip-g=FILE assignments; required for PiD and Lens",
+         (int)',',
+         &tokenizer},
         {"",
          "--llm_vision",
          "path to the llm vit",
@@ -460,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)",
@@ -614,7 +624,7 @@ ArgOptions SDContextParams::get_options() {
          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},
     };
 
@@ -891,6 +901,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"
@@ -898,6 +909,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"
@@ -957,12 +969,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();
@@ -1095,7 +1109,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; 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},
         {"",
@@ -1515,6 +1529,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;
@@ -1704,6 +1726,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;

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/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

@@ -208,6 +208,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;
@@ -243,6 +244,7 @@ typedef struct {
     bool disable_segmented_compute;  // Force monolithic graph execution even when automatic graph cutting would fit memory better
     float linear_scale;              // Override linear input scaling; 0 keeps the model default
     float attn_scale;                // Override flash-attention K/V scaling; 0 keeps the model default
+    const char* tokenizer;           // tokenizer.json path or main=FILE,clip-l=FILE,clip-g=FILE assignments; required for PiD and Lens
 } sd_ctx_params_t;
 
 typedef struct {
@@ -521,11 +523,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/conditioner.hpp

@@ -7,6 +7,7 @@
 #include <limits>
 #include <optional>
 #include <sstream>
+#include <stdexcept>
 #include "core/ggml_tensor_utils.h"
 
 #include "core/tensor_ggml.hpp"
@@ -17,6 +18,7 @@
 #include "model/te/t5.hpp"
 #include "model_loader.h"
 #include "tokenizers/sensenova_u1_tokenizer.h"
+#include "tokenizers/tokenizer_config.h"
 
 struct SDCondition {
     sd::Tensor<float> c_crossattn;
@@ -159,7 +161,7 @@ public:
 // Ref: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/cad87bf4e3e0b0a759afa94e933527c3123d59bc/modules/sd_hijack_clip.py#L283
 struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
     SDVersion version = VERSION_SD1;
-    CLIPTokenizer tokenizer;
+    std::shared_ptr<Tokenizer> tokenizer;
     std::shared_ptr<CLIPTextModelRunner> text_model;
     std::shared_ptr<CLIPTextModelRunner> text_model2;
 
@@ -173,12 +175,18 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                                       const String2TensorStorage& tensor_storage_map,
                                       const std::map<std::string, std::string>& orig_embedding_map,
                                       SDVersion version                                   = VERSION_SD1,
-                                      std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
-        : version(version), tokenizer(sd_version_is_sd2(version) ? 0 : 49407) {
+                                      std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                                      const TokenizerConfig& tokenizers                   = {})
+        : version(version) {
+        const int pad_id = sd_version_is_sd2(version) ? 0 : 49407;
+        tokenizer        = tokenizers.create(TokenizerConfig::MAIN, 49408, pad_id, false, true);
+        if (!tokenizer) {
+            tokenizer = std::make_shared<CLIPTokenizer>(pad_id);
+        }
         for (const auto& kv : orig_embedding_map) {
             std::string name    = normalize_embedding_name(kv.first);
             embedding_map[name] = kv.second;
-            tokenizer.add_special_token(name);
+            tokenizer->add_special_token(name);
         }
         bool force_clip_f32 = !embedding_map.empty();
         if (sd_version_is_sd1(version)) {
@@ -365,16 +373,15 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         return load_embedding(name, iter->second, bpe_tokens);
     }
 
-    std::vector<int> convert_token_to_id(std::string text) {
+    bool convert_token_to_id(const std::string& text, std::vector<int>& tokens) {
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
             return append_embedding_tokens(str, bpe_tokens);
         };
-        std::vector<int> curr_tokens = tokenizer.encode(text, on_new_token_cb);
-        return curr_tokens;
+        return tokenizer->encode(text, tokens, on_new_token_cb);
     }
 
-    std::string decode(const std::vector<int>& tokens) {
-        return tokenizer.decode(tokens);
+    bool decode(const std::vector<int>& tokens, std::string& text) {
+        return tokenizer->decode(tokens, text);
     }
 
     std::pair<std::vector<int>, std::vector<float>> tokenize(std::string text,
@@ -412,18 +419,21 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
 
                 if (padding_size > 0) {
                     LOG_VERBOSE("BREAK token encountered, padding current chunk by %zu tokens.", padding_size);
-                    tokens.insert(tokens.end(), padding_size, tokenizer.EOS_TOKEN_ID);
+                    tokens.insert(tokens.end(), padding_size, tokenizer->EOS_TOKEN_ID);
                     weights.insert(weights.end(), padding_size, 1.0f);
                 }
                 continue;  // Skip to the next item after handling BREAK
             }
 
-            std::vector<int> curr_tokens = tokenizer.encode(curr_text, on_new_token_cb);
+            std::vector<int> curr_tokens;
+            if (!tokenizer->encode(curr_text, curr_tokens, on_new_token_cb)) {
+                return {};
+            }
             tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
             weights.insert(weights.end(), curr_tokens.size(), curr_weight);
         }
 
-        tokenizer.pad_tokens(tokens, &weights, nullptr, min_length, max_length, allow_overflow_expand);
+        tokenizer->pad_tokens(tokens, &weights, nullptr, min_length, max_length, allow_overflow_expand);
 
         // for (int i = 0; i < tokens.size(); i++) {
         //     std::cout << tokens[i] << ":" << weights[i] << ", ";
@@ -460,7 +470,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
             sd::Tensor<int32_t> input_ids2;
             size_t max_token_idx = 0;
             if (sd_version_is_sdxl(version)) {
-                auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), tokenizer.EOS_TOKEN_ID);
+                auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), tokenizer->EOS_TOKEN_ID);
                 if (it != chunk_tokens.end()) {
                     std::fill(std::next(it), chunk_tokens.end(), 0);
                 }
@@ -561,7 +571,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
 
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
-        auto tokens_and_weights     = tokenize(conditioner_params.text, text_model->model.n_token, text_model->model.n_token, true);
+        auto tokens_and_weights = tokenize(conditioner_params.text, text_model->model.n_token, text_model->model.n_token, true);
+        if (tokens_and_weights.first.empty()) {
+            return {};
+        }
         std::vector<int>& tokens    = tokens_and_weights.first;
         std::vector<float>& weights = tokens_and_weights.second;
         return get_learned_condition_common(n_threads,
@@ -629,8 +642,8 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
 };
 
 struct SD3CLIPEmbedder : public Conditioner {
-    CLIPTokenizer clip_l_tokenizer;
-    CLIPTokenizer clip_g_tokenizer;
+    std::shared_ptr<Tokenizer> clip_l_tokenizer;
+    std::shared_ptr<Tokenizer> clip_g_tokenizer;
     T5UniGramTokenizer t5_tokenizer;
     std::shared_ptr<CLIPTextModelRunner> clip_l;
     std::shared_ptr<CLIPTextModelRunner> clip_g;
@@ -638,8 +651,8 @@ struct SD3CLIPEmbedder : public Conditioner {
 
     SD3CLIPEmbedder(ggml_backend_t backend,
                     const String2TensorStorage& tensor_storage_map      = {},
-                    std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
-        : clip_g_tokenizer(0) {
+                    std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                    const TokenizerConfig& tokenizers                   = {}) {
         bool use_clip_l = false;
         bool use_clip_g = false;
         bool use_t5     = false;
@@ -657,9 +670,17 @@ struct SD3CLIPEmbedder : public Conditioner {
             return;
         }
         if (use_clip_l) {
+            clip_l_tokenizer = tokenizers.create(TokenizerConfig::CLIP_L, 49408, 49407, false, true);
+            if (!clip_l_tokenizer) {
+                clip_l_tokenizer = std::make_shared<CLIPTokenizer>();
+            }
             clip_l = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, false, false, weight_manager);
         }
         if (use_clip_g) {
+            clip_g_tokenizer = tokenizers.create(TokenizerConfig::CLIP_G, 49408, 0, false, true);
+            if (!clip_g_tokenizer) {
+                clip_g_tokenizer = std::make_shared<CLIPTokenizer>(0);
+            }
             clip_g = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "text_encoders.clip_g.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false, false, weight_manager);
         }
         if (use_t5) {
@@ -811,27 +832,36 @@ struct SD3CLIPEmbedder : public Conditioner {
             const std::string& curr_text = item.first;
             float curr_weight            = item.second;
             if (clip_l) {
-                std::vector<int> curr_tokens = clip_l_tokenizer.encode(curr_text, on_new_token_cb);
+                std::vector<int> curr_tokens;
+                if (!clip_l_tokenizer->encode(curr_text, curr_tokens, on_new_token_cb)) {
+                    return {};
+                }
                 clip_l_tokens.insert(clip_l_tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight);
             }
             if (clip_g) {
-                std::vector<int> curr_tokens = clip_g_tokenizer.encode(curr_text, on_new_token_cb);
+                std::vector<int> curr_tokens;
+                if (!clip_g_tokenizer->encode(curr_text, curr_tokens, on_new_token_cb)) {
+                    return {};
+                }
                 clip_g_tokens.insert(clip_g_tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 clip_g_weights.insert(clip_g_weights.end(), curr_tokens.size(), curr_weight);
             }
             if (t5) {
-                std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
+                std::vector<int> curr_tokens;
+                if (!t5_tokenizer.encode(curr_text, curr_tokens)) {
+                    return {};
+                }
                 t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
             }
         }
 
         if (clip_l) {
-            clip_l_tokenizer.pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, min_length, max_length, allow_overflow_expand);
+            clip_l_tokenizer->pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, min_length, max_length, allow_overflow_expand);
         }
         if (clip_g) {
-            clip_g_tokenizer.pad_tokens(clip_g_tokens, &clip_g_weights, nullptr, min_length, max_length, allow_overflow_expand);
+            clip_g_tokenizer->pad_tokens(clip_g_tokens, &clip_g_weights, nullptr, min_length, max_length, allow_overflow_expand);
         }
         if (t5) {
             t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr, min_length, max_length, true);
@@ -902,7 +932,7 @@ struct SD3CLIPEmbedder : public Conditioner {
                 chunk_hidden_states_l = ::apply_token_weights(std::move(chunk_hidden_states_l), chunk_weights);
 
                 if (chunk_idx == 0) {
-                    auto it       = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_l_tokenizer.EOS_TOKEN_ID);
+                    auto it       = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_l_tokenizer->EOS_TOKEN_ID);
                     max_token_idx = std::min<size_t>(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1);
                     pooled_l      = clip_l->compute(n_threads,
                                                     input_ids,
@@ -945,7 +975,7 @@ struct SD3CLIPEmbedder : public Conditioner {
                 chunk_hidden_states_g = ::apply_token_weights(std::move(chunk_hidden_states_g), chunk_weights);
 
                 if (chunk_idx == 0) {
-                    auto it       = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_g_tokenizer.EOS_TOKEN_ID);
+                    auto it       = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_g_tokenizer->EOS_TOKEN_ID);
                     max_token_idx = std::min<size_t>(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1);
                     pooled_g      = clip_g->compute(n_threads,
                                                     input_ids,
@@ -1023,6 +1053,9 @@ struct SD3CLIPEmbedder : public Conditioner {
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
         auto tokens_and_weights = tokenize(conditioner_params.text, 77, 77, true);
+        if (tokens_and_weights.empty()) {
+            return {};
+        }
         return get_learned_condition_common(n_threads,
                                             tokens_and_weights,
                                             conditioner_params.clip_skip,
@@ -1031,7 +1064,7 @@ struct SD3CLIPEmbedder : public Conditioner {
 };
 
 struct FluxCLIPEmbedder : public Conditioner {
-    CLIPTokenizer clip_l_tokenizer;
+    std::shared_ptr<Tokenizer> clip_l_tokenizer;
     T5UniGramTokenizer t5_tokenizer;
     std::shared_ptr<CLIPTextModelRunner> clip_l;
     std::shared_ptr<T5Runner> t5;
@@ -1039,7 +1072,8 @@ struct FluxCLIPEmbedder : public Conditioner {
 
     FluxCLIPEmbedder(ggml_backend_t backend,
                      const String2TensorStorage& tensor_storage_map      = {},
-                     std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) {
+                     std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                     const TokenizerConfig& tokenizers                   = {}) {
         bool use_clip_l = false;
         bool use_t5     = false;
         for (auto pair : tensor_storage_map) {
@@ -1056,6 +1090,11 @@ struct FluxCLIPEmbedder : public Conditioner {
         }
 
         if (use_clip_l) {
+            auto slot        = tokenizers.has(TokenizerConfig::CLIP_L) ? TokenizerConfig::CLIP_L : TokenizerConfig::MAIN;
+            clip_l_tokenizer = tokenizers.create(slot, 49408, 49407, false, true);
+            if (!clip_l_tokenizer) {
+                clip_l_tokenizer = std::make_shared<CLIPTokenizer>();
+            }
             clip_l = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, true, false, weight_manager);
         } else {
             LOG_WARN("clip_l text encoder not found! Prompt adherence might be degraded.");
@@ -1181,19 +1220,25 @@ struct FluxCLIPEmbedder : public Conditioner {
             const std::string& curr_text = item.first;
             float curr_weight            = item.second;
             if (clip_l) {
-                std::vector<int> curr_tokens = clip_l_tokenizer.encode(curr_text, on_new_token_cb);
+                std::vector<int> curr_tokens;
+                if (!clip_l_tokenizer->encode(curr_text, curr_tokens, on_new_token_cb)) {
+                    return {};
+                }
                 clip_l_tokens.insert(clip_l_tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight);
             }
             if (t5) {
-                std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
+                std::vector<int> curr_tokens;
+                if (!t5_tokenizer.encode(curr_text, curr_tokens)) {
+                    return {};
+                }
                 t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
             }
         }
 
         if (clip_l) {
-            clip_l_tokenizer.pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, 77, 77, true);
+            clip_l_tokenizer->pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, 77, 77, true);
         }
         if (t5) {
             t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr, min_length, max_length, true);
@@ -1243,7 +1288,7 @@ struct FluxCLIPEmbedder : public Conditioner {
                     sd::Tensor<int32_t> input_ids({static_cast<int64_t>(chunk_tokens.size())}, chunk_tokens);
                     size_t max_token_idx = 0;
 
-                    auto it       = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_l_tokenizer.EOS_TOKEN_ID);
+                    auto it       = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_l_tokenizer->EOS_TOKEN_ID);
                     max_token_idx = std::min<size_t>(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1);
 
                     pooled = clip_l->compute(n_threads,
@@ -1254,7 +1299,10 @@ struct FluxCLIPEmbedder : public Conditioner {
                                              true,
                                              clip_skip,
                                              false);
-                    GGML_ASSERT(!pooled.empty());
+                    if (pooled.empty()) {
+                        LOG_ERROR("Flux CLIP-L encoding failed");
+                        return {};
+                    }
                 } else {
                     pooled = sd::Tensor<float>::zeros({768});
                 }
@@ -1273,7 +1321,10 @@ struct FluxCLIPEmbedder : public Conditioner {
                                                   input_ids,
                                                   sd::Tensor<float>(),
                                                   false);
-                GGML_ASSERT(!chunk_hidden_states.empty());
+                if (chunk_hidden_states.empty()) {
+                    LOG_ERROR("Flux T5 encoding failed at chunk %d/%zu", chunk_idx + 1, chunk_count);
+                    return {};
+                }
                 chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
                 if (zero_out_masked) {
                     chunk_hidden_states.fill_(0.0f);
@@ -1300,6 +1351,9 @@ struct FluxCLIPEmbedder : public Conditioner {
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
         auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, chunk_len);
+        if (tokens_and_weights.empty()) {
+            return {};
+        }
         return get_learned_condition_common(n_threads,
                                             tokens_and_weights,
                                             conditioner_params.clip_skip,
@@ -1443,7 +1497,10 @@ struct T5CLIPEmbedder : public Conditioner {
                 const std::string& curr_text = item.first;
                 float curr_weight            = item.second;
 
-                std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
+                std::vector<int> curr_tokens;
+                if (!t5_tokenizer.encode(curr_text, curr_tokens)) {
+                    return {};
+                }
                 t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
             }
@@ -1541,6 +1598,9 @@ struct T5CLIPEmbedder : public Conditioner {
     SDCondition get_learned_condition(int n_threads,
                                       const ConditionerParams& conditioner_params) override {
         auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, chunk_len);
+        if (std::get<0>(tokens_and_weights).empty()) {
+            return {};
+        }
         return get_learned_condition_common(n_threads,
                                             tokens_and_weights,
                                             conditioner_params.clip_skip,
@@ -1639,7 +1699,10 @@ struct MiniT2IConditioner : public Conditioner {
             return result;
         }
 
-        std::vector<int> tokens = tokenizer.encode(conditioner_params.text);
+        std::vector<int> tokens;
+        if (!tokenizer.encode(conditioner_params.text, tokens)) {
+            return {};
+        }
         if (tokens.size() > prompt_length) {
             tokens.resize(prompt_length);
         }
@@ -1706,7 +1769,10 @@ struct SenseNovaU1Conditioner : public Conditioner {
     }
 
     SDCondition tokenize_condition(const std::string& text, bool is_negative) {
-        auto tokens = tokenizer.encode(build_query(text, is_negative));
+        std::vector<int> tokens;
+        if (!tokenizer.encode(build_query(text, is_negative), tokens)) {
+            return {};
+        }
         if (tokens.empty() || tokens.size() > kMaxPromptTokens) {
             LOG_ERROR("SenseNova U1.5 prompt token count %zu is outside [1, %zu]",
                       tokens.size(),
@@ -1731,20 +1797,24 @@ struct SenseNovaU1Conditioner : public Conditioner {
 };
 
 struct AnimaConditioner : public Conditioner {
-    std::shared_ptr<BPETokenizer> qwen_tokenizer;
+    std::shared_ptr<Tokenizer> qwen_tokenizer;
     T5UniGramTokenizer t5_tokenizer;
     std::shared_ptr<LLM::LLMRunner> llm;
 
     AnimaConditioner(ggml_backend_t backend,
                      const String2TensorStorage& tensor_storage_map      = {},
-                     std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) {
-        qwen_tokenizer = std::make_shared<Qwen2Tokenizer>();
+                     std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                     const TokenizerConfig& tokenizers                   = {}) {
         llm            = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::QWEN3,
                                                backend,
                                                tensor_storage_map,
                                                "text_encoders.llm",
                                                false,
                                                weight_manager);
+        qwen_tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 151643);
+        if (!qwen_tokenizer) {
+            qwen_tokenizer = std::make_shared<Qwen2Tokenizer>();
+        }
     }
 
     void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
@@ -1811,7 +1881,10 @@ struct AnimaConditioner : public Conditioner {
 
         for (const auto& item : parsed_attention) {
             const std::string& curr_text = item.first;
-            std::vector<int> curr_tokens = qwen_tokenizer->tokenize(curr_text, nullptr);
+            std::vector<int> curr_tokens;
+            if (!qwen_tokenizer->tokenize(curr_text, curr_tokens, nullptr)) {
+                return {};
+            }
             qwen_tokens.insert(qwen_tokens.end(), curr_tokens.begin(), curr_tokens.end());
             // Anima uses uniform Qwen token weights.
             qwen_weights.insert(qwen_weights.end(), curr_tokens.size(), 1.f);
@@ -1824,7 +1897,10 @@ struct AnimaConditioner : public Conditioner {
         for (const auto& item : parsed_attention) {
             const std::string& curr_text = item.first;
             float curr_weight            = item.second;
-            std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
+            std::vector<int> curr_tokens;
+            if (!t5_tokenizer.encode(curr_text, curr_tokens)) {
+                return {};
+            }
             t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
             t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
         }
@@ -1843,6 +1919,10 @@ struct AnimaConditioner : public Conditioner {
         auto& t5_tokens    = std::get<2>(tokenized);
         auto& t5_weights   = std::get<3>(tokenized);
 
+        if (qwen_tokens.empty()) {
+            return {};
+        }
+
         sd::Tensor<int32_t> input_ids({static_cast<int64_t>(qwen_tokens.size()), 1}, qwen_tokens);
         auto hidden_states = llm->compute(n_threads,
                                           input_ids,
@@ -1869,7 +1949,7 @@ struct AnimaConditioner : public Conditioner {
 
 struct LLMEmbedder : public Conditioner {
     SDVersion version;
-    std::shared_ptr<BPETokenizer> tokenizer;
+    std::shared_ptr<Tokenizer> tokenizer;
     std::shared_ptr<LLM::LLMRunner> llm;
     std::shared_ptr<T5Runner> byt5;
 
@@ -1878,8 +1958,17 @@ struct LLMEmbedder : public Conditioner {
                 SDVersion version                                   = VERSION_QWEN_IMAGE,
                 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                   = {})
         : version(version) {
+        if (!tokenizers.has(TokenizerConfig::MAIN)) {
+            if (sd_version_is_lens(version)) {
+                throw std::runtime_error("Lens requires an external GPT-OSS tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer");
+            }
+            if (sd_version_is_pid(version)) {
+                throw std::runtime_error("PiD requires an external Gemma 2 tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer");
+            }
+        }
         LLM::LLMArch arch = LLM::LLMArch::QWEN2_5_VL;
         if (version == VERSION_FLUX2) {
             arch = LLM::LLMArch::MISTRAL_SMALL_3_2;
@@ -1889,7 +1978,8 @@ struct LLMEmbedder : public Conditioner {
             arch = LLM::LLMArch::GPT_OSS_20B;
         } else if (sd_version_is_pid(version)) {
             arch = LLM::LLMArch::GEMMA2_2B;
-        } else if (sd_version_is_lingbot_video(version) ||
+        } else if (version == VERSION_QWEN_IMAGE_2_1 ||
+                   sd_version_is_lingbot_video(version) ||
                    sd_version_is_ideogram4(version) ||
                    sd_version_is_boogu_image(version) ||
                    sd_version_is_sefi_image(version) ||
@@ -1900,21 +1990,28 @@ struct LLMEmbedder : public Conditioner {
         } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
             arch = LLM::LLMArch::QWEN3;
         }
-        if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
-            tokenizer = std::make_shared<MistralTokenizer>();
-        } else if (arch == LLM::LLMArch::GPT_OSS_20B) {
-            tokenizer = std::make_shared<GPTOSSTokenizer>();
-        } else if (arch == LLM::LLMArch::GEMMA2_2B) {
-            tokenizer = std::make_shared<Gemma2Tokenizer>();
-        } else {
-            tokenizer = std::make_shared<Qwen2Tokenizer>();
-        }
-        llm = std::make_shared<LLM::LLMRunner>(arch,
+        llm        = std::make_shared<LLM::LLMRunner>(arch,
                                                backend,
                                                tensor_storage_map,
                                                "text_encoders.llm",
                                                enable_vision,
                                                weight_manager);
+        int pad_id = 151643;
+        if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
+            pad_id = 11;
+        } else if (arch == LLM::LLMArch::GPT_OSS_20B) {
+            pad_id = 199999;
+        } else if (arch == LLM::LLMArch::GEMMA2_2B) {
+            pad_id = 0;
+        }
+        tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, pad_id);
+        if (!tokenizer) {
+            if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
+                tokenizer = std::make_shared<MistralTokenizer>();
+            } else {
+                tokenizer = std::make_shared<Qwen2Tokenizer>();
+            }
+        }
         if (sd_version_is_hunyuan_video(version)) {
             const std::string byt5_prefix = "text_encoders.t5xxl.transformer";
             for (const auto& [name, _] : tensor_storage_map) {
@@ -2053,7 +2150,10 @@ struct LLMEmbedder : public Conditioner {
         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, nullptr);
+            std::vector<int> curr_tokens;
+            if (!tokenizer->encode(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);
         }
@@ -2086,6 +2186,10 @@ struct LLMEmbedder : public Conditioner {
         auto& weights            = std::get<1>(tokens_weights_mask);
         auto& mask               = std::get<2>(tokens_weights_mask);
 
+        if (tokens.empty()) {
+            return {};
+        }
+
         sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
         sd::Tensor<float> attention_mask;
         if (!mask.empty()) {
@@ -2245,7 +2349,11 @@ struct LLMEmbedder : public Conditioner {
                     GGML_ASSERT(image_outputs.size() == 4);
                     auto image_embed = std::move(image_outputs[0]);
                     prompt += "<|vision_start|>";
-                    int image_embed_idx = static_cast<int>(tokenizer->encode(prompt, nullptr).size());
+                    std::vector<int> prefix_tokens;
+                    if (!tokenizer->encode(prompt, prefix_tokens, nullptr)) {
+                        return false;
+                    }
+                    int image_embed_idx = static_cast<int>(prefix_tokens.size());
                     image_embeds.emplace_back(image_embed_idx, image_embed);
                     if (deepstack_image_embeds.empty()) {
                         deepstack_image_embeds.resize(image_outputs.size() - 1);
@@ -2261,6 +2369,7 @@ struct LLMEmbedder : public Conditioner {
                         prompt += placeholder;
                     }
                     prompt += "<|vision_end|>";
+                    return true;
                 };
 
                 const auto* references = conditioner_params.minimax_h3_references;
@@ -2277,11 +2386,13 @@ struct LLMEmbedder : public Conditioner {
                             GGML_ASSERT(item.frames.size() == 1);
                             auto resized = resize_for_vision(item.frames[0]);
                             prompt += "<Picture " + std::to_string(++picture_index) + ">: ";
-                            add_vision_outputs(llm->encode_image_outputs(n_threads,
-                                                                         resized,
-                                                                         false),
-                                               static_cast<int>(resized.shape()[1]) / patch_size,
-                                               static_cast<int>(resized.shape()[0]) / patch_size);
+                            if (!add_vision_outputs(llm->encode_image_outputs(n_threads,
+                                                                              resized,
+                                                                              false),
+                                                    static_cast<int>(resized.shape()[1]) / patch_size,
+                                                    static_cast<int>(resized.shape()[0]) / patch_size)) {
+                                return {};
+                            }
                             continue;
                         }
 
@@ -2305,22 +2416,26 @@ struct LLMEmbedder : public Conditioner {
                                                                                    second.shape()[3]});
                             }
                             auto pair = sd::ops::concat(first.unsqueeze(2), second.unsqueeze(2), 2);
-                            add_vision_outputs(llm->encode_video_block_outputs(n_threads,
-                                                                               pair,
-                                                                               false),
-                                               static_cast<int>(first.shape()[1]) / patch_size,
-                                               static_cast<int>(first.shape()[0]) / patch_size);
+                            if (!add_vision_outputs(llm->encode_video_block_outputs(n_threads,
+                                                                                    pair,
+                                                                                    false),
+                                                    static_cast<int>(first.shape()[1]) / patch_size,
+                                                    static_cast<int>(first.shape()[0]) / patch_size)) {
+                                return {};
+                            }
                         }
                     }
                 } else if (conditioner_params.ref_images != nullptr) {
                     for (size_t i = 0; i < conditioner_params.ref_images->size(); ++i) {
                         auto resized = resize_for_vision((*conditioner_params.ref_images)[i]);
                         prompt += "<Picture " + std::to_string(i + 1) + ">: ";
-                        add_vision_outputs(llm->encode_image_outputs(n_threads,
-                                                                     resized,
-                                                                     false),
-                                           static_cast<int>(resized.shape()[1]) / patch_size,
-                                           static_cast<int>(resized.shape()[0]) / patch_size);
+                        if (!add_vision_outputs(llm->encode_image_outputs(n_threads,
+                                                                          resized,
+                                                                          false),
+                                                static_cast<int>(resized.shape()[1]) / patch_size,
+                                                static_cast<int>(resized.shape()[0]) / patch_size)) {
+                            return {};
+                        }
                     }
                 }
             }
@@ -2356,7 +2471,10 @@ struct LLMEmbedder : public Conditioner {
                 "enhanced description for the prompt below and avoid including any additional "
                 "commentary or evaluations:<|im_end|>\n<|im_start|>user\n";
 
-            auto prefix_tokens               = tokenizer->encode(prompt_prefix, nullptr);
+            std::vector<int> prefix_tokens;
+            if (!tokenizer->encode(prompt_prefix, prefix_tokens, nullptr)) {
+                return {};
+            }
             prompt_template_encode_start_idx = 0;
             for (int token : prefix_tokens) {
                 if (token != pad_token) {
@@ -2409,7 +2527,11 @@ struct LLMEmbedder : public Conditioner {
                     GGML_ASSERT(!image_embed.empty());
 
                     std::string image_prefix = prompt + img_prompt + "<|vision_start|>";
-                    int image_embed_idx      = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
+                    std::vector<int> prefix_tokens;
+                    if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) {
+                        return {};
+                    }
+                    int image_embed_idx = static_cast<int>(prefix_tokens.size());
                     image_embeds.emplace_back(image_embed_idx, image_embed);
 
                     img_prompt += "<|vision_start|>";
@@ -2426,6 +2548,67 @@ struct LLMEmbedder : public Conditioner {
             prompt += conditioner_params.text;
             prompt_attn_range = {0, 0};
             prompt += "<|im_end|>\n<|im_start|>assistant\n";
+        } else if (version == VERSION_QWEN_IMAGE_2_1) {
+            if (!llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
+                LOG_ERROR("Qwen Image 2.1 editing requires Qwen3-VL vision weights; provide --llm_vision or a combined encoder");
+                return {};
+            }
+            prompt = "<|im_start|>system\nComprehend and analyze the provided prompt.<|im_end|>\n";
+            std::vector<int> system_tokens;
+            if (!tokenizer->encode(prompt, system_tokens, nullptr)) {
+                return {};
+            }
+            prompt_template_encode_start_idx = static_cast<int>(system_tokens.size());
+            out_layers                       = {static_cast<int>(llm->config.num_layers)};
+            prompt += "<|im_start|>user\n";
+            if (llm->enable_vision && conditioner_params.ref_images != nullptr) {
+                for (size_t i = 0; i < conditioner_params.ref_images->size(); ++i) {
+                    const auto& image = (*conditioner_params.ref_images)[i];
+                    int64_t width     = image.shape()[0];
+                    int64_t height    = image.shape()[1];
+                    int64_t pixels    = width * height;
+                    if (width % 32 != 0 || height % 32 != 0) {
+                        LOG_ERROR("Qwen Image 2.1 reference dimensions must be multiples of 32");
+                        return {};
+                    }
+                    auto rgb = sd::Tensor<float>({width, height, 3, 1});
+                    for (int64_t p = 0; p < pixels; ++p) {
+                        float alpha = image.shape()[2] == 4 ? image[p + 3 * pixels] : 1.f;
+                        for (int c = 0; c < 3; ++c) {
+                            rgb[p + c * pixels] = 2.f * (image[p + c * pixels] * alpha + 1.f - alpha) - 1.f;
+                        }
+                    }
+                    auto outputs = llm->encode_image_outputs(n_threads, rgb, false);
+                    if (outputs.empty()) {
+                        return {};
+                    }
+                    prompt += (i == 0 ? "" : " ") + std::string("<image") + std::to_string(i + 1) + "><|vision_start|>";
+                    std::vector<int> prefix_tokens;
+                    if (!tokenizer->encode(prompt, prefix_tokens, nullptr)) {
+                        return {};
+                    }
+                    int index = static_cast<int>(prefix_tokens.size());
+                    int count = static_cast<int>(outputs[0].shape()[1]);
+                    image_embeds.emplace_back(index, std::move(outputs[0]));
+                    if (deepstack_image_embeds.empty()) {
+                        deepstack_image_embeds.resize(outputs.size() - 1);
+                    }
+                    for (size_t layer = 1; layer < outputs.size(); ++layer) {
+                        deepstack_image_embeds[layer - 1].emplace_back(index, std::move(outputs[layer]));
+                    }
+                    image_grids.push_back({index, count,
+                                           static_cast<int>(height) / llm->config.vision.patch_size,
+                                           static_cast<int>(width) / llm->config.vision.patch_size});
+                    for (int j = 0; j < count; ++j) {
+                        prompt += "<|image_pad|>";
+                    }
+                    prompt += "<|vision_end|>";
+                }
+            }
+            prompt_attn_range.first = static_cast<int>(prompt.size());
+            prompt += conditioner_params.text.empty() ? " " : conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+            prompt += "<|im_end|>\n<|im_start|>assistant\n";
         } else if (sd_version_is_qwen_image(version) || sd_version_is_mage_flow(version)) {
             if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
                 LOG_INFO("%s", sd_version_is_mage_flow(version) ? "MageFlowEditPipeline" : "QwenImageEditPlusPipeline");
@@ -2551,7 +2734,11 @@ struct LLMEmbedder : public Conditioner {
                     GGML_ASSERT(!image_embed.empty());
 
                     std::string image_prefix = prompt_prefix + img_prompt + "<|vision_start|>";
-                    int image_embed_idx      = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
+                    std::vector<int> prefix_tokens;
+                    if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) {
+                        return {};
+                    }
+                    int image_embed_idx = static_cast<int>(prefix_tokens.size());
                     image_embeds.emplace_back(image_embed_idx, image_embed);
 
                     img_prompt += "<|vision_start|>";
@@ -2619,7 +2806,11 @@ struct LLMEmbedder : public Conditioner {
                     GGML_ASSERT(!image_embed.empty());
 
                     std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
-                    int image_embed_idx      = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
+                    std::vector<int> prefix_tokens;
+                    if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) {
+                        return {};
+                    }
+                    int image_embed_idx = static_cast<int>(prefix_tokens.size());
                     image_embeds.emplace_back(image_embed_idx, image_embed);
 
                     img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
@@ -2824,7 +3015,10 @@ struct LLMEmbedder : public Conditioner {
                 "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.\n"
                 "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:\n"
                 "User Prompt: ";
-            auto chi_tokens       = std::get<0>(tokenize(chi_prompt, {0, 0}));
+            auto chi_tokens = std::get<0>(tokenize(chi_prompt, {0, 0}));
+            if (chi_tokens.empty()) {
+                return {};
+            }
             size_t num_chi_tokens = chi_tokens.size();
             max_length            = (int)num_chi_tokens + pixeldit_max_length - 2;
             min_length            = max_length;
@@ -2843,7 +3037,9 @@ struct LLMEmbedder : public Conditioner {
                                                0,
                                                false,
                                                max_length);
-            GGML_ASSERT(!hidden_states.empty());
+            if (hidden_states.empty()) {
+                return {};
+            }
 
             if (hidden_states.shape()[1] > pixeldit_max_length) {
                 auto bos      = sd::ops::slice(hidden_states, 1, 0, 1);
@@ -2876,6 +3072,9 @@ struct LLMEmbedder : public Conditioner {
                                            max_length,
                                            deepstack_image_embeds,
                                            image_grids);
+        if (hidden_states.empty()) {
+            return {};
+        }
         std::vector<sd::Tensor<float>> extra_hidden_states_vec;
         if (sd_version_is_hunyuan_video(version) && byt5) {
             std::vector<std::string> quoted_texts;
@@ -2926,6 +3125,9 @@ struct LLMEmbedder : public Conditioner {
                                                      prompt_template_encode_start_idx,
                                                      spell_quotes,
                                                      max_length);
+            if (extra_hidden_states.empty()) {
+                return {};
+            }
             extra_hidden_states_vec.push_back(std::move(extra_hidden_states));
         }
 
@@ -2934,6 +3136,21 @@ struct LLMEmbedder : public Conditioner {
         SDCondition result;
         result.c_crossattn        = std::move(hidden_states);
         result.extra_c_crossattns = std::move(extra_hidden_states_vec);
+        if (version == VERSION_QWEN_IMAGE_2_1) {
+            auto slots = sd::Tensor<int32_t>::zeros({result.c_crossattn.shape()[1]});
+            for (size_t i = 0; i < image_embeds.size(); ++i) {
+                int64_t begin = image_embeds[i].first - prompt_template_encode_start_idx;
+                int64_t end   = begin + image_embeds[i].second.shape()[1];
+                if (begin < 0 || end > slots.numel()) {
+                    LOG_ERROR("Qwen Image 2.1 image slots exceed the encoded prompt");
+                    return {};
+                }
+                for (int64_t j = begin; j < end; ++j) {
+                    slots[j] = static_cast<int32_t>(i + 1);
+                }
+            }
+            result.c_token_types = std::move(slots);
+        }
         if (sd_version_is_minimax_h3(version)) {
             std::vector<int32_t> tags(static_cast<size_t>(result.c_crossattn.shape()[1]), 1);
             for (const auto& [index, image_embed] : image_embeds) {
@@ -3024,7 +3241,7 @@ struct LTXAVEmbedder : public Conditioner {
     static constexpr int64_t kNumStates  = 49;
     static constexpr int64_t kMinLength  = 1024;
 
-    std::shared_ptr<GemmaTokenizer> tokenizer;
+    std::shared_ptr<Tokenizer> tokenizer;
     std::shared_ptr<LLM::LLMRunner> llm;
     std::shared_ptr<LTXAVTextProjectionRunner> projector;
     std::string projector_prefix;
@@ -3051,17 +3268,21 @@ struct LTXAVEmbedder : public Conditioner {
                   const String2TensorStorage& tensor_storage_map      = {},
                   const std::string& llm_prefix                       = "text_encoders.llm",
                   const std::string& projector_prefix                 = "text_embedding_projection",
-                  std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+                  std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                  const TokenizerConfig& tokenizers                   = {})
         : projector_prefix(projector_prefix) {
         LLM::LLMArch arch = detect_gemma_arch(tensor_storage_map, llm_prefix);
         LOG_INFO("ltxav text encoder: %s", arch == LLM::LLMArch::GEMMA4_12B ? "gemma 4" : "gemma 3");
-        tokenizer       = std::make_shared<GemmaTokenizer>();
-        llm             = std::make_shared<LLM::LLMRunner>(arch,
+        llm       = std::make_shared<LLM::LLMRunner>(arch,
                                                backend,
                                                tensor_storage_map,
                                                llm_prefix,
                                                false,
                                                weight_manager);
+        tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 0, true);
+        if (!tokenizer) {
+            tokenizer = std::make_shared<GemmaTokenizer>();
+        }
         dual_projection = tensor_storage_map.find(projector_prefix + ".video_aggregate_embed.weight") != tensor_storage_map.end();
         projector       = std::make_shared<LTXAVTextProjectionRunner>(backend,
                                                                 tensor_storage_map,
@@ -3140,7 +3361,10 @@ struct LTXAVEmbedder : public Conditioner {
         std::vector<int> tokens;
         std::vector<float> weights;
         for (const auto& item : parsed_attention) {
-            auto curr_tokens = tokenizer->encode(item.first, nullptr);
+            std::vector<int> curr_tokens;
+            if (!tokenizer->encode(item.first, curr_tokens, nullptr)) {
+                return {};
+            }
             tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
             weights.insert(weights.end(), curr_tokens.size(), item.second);
         }
@@ -3158,6 +3382,10 @@ struct LTXAVEmbedder : public Conditioner {
         auto& weights            = std::get<1>(tokens_weights_mask);
         auto& mask               = std::get<2>(tokens_weights_mask);
 
+        if (tokens.empty()) {
+            return {};
+        }
+
         sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, std::vector<int32_t>(tokens.begin(), tokens.end()));
         sd::Tensor<float> attention_mask;
         if (!mask.empty()) {
@@ -3256,7 +3484,9 @@ struct LTXAVEmbedder : public Conditioner {
         prompt_attn_range.second = static_cast<int>(prompt.size());
 
         auto hidden_states = encode_prompt(n_threads, prompt, prompt_attn_range);
-        GGML_ASSERT(!hidden_states.empty());
+        if (hidden_states.empty()) {
+            return {};
+        }
 
         int64_t t1 = ggml_time_ms();
         LOG_VERBOSE("computing LTXAV condition graph completed, taking %" PRId64 " ms", t1 - t0);

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

@@ -102,7 +102,7 @@ namespace sd::backend_fit {
         for (const auto& [name, stored_tensor] : loader.get_tensor_storage_map()) {
             TensorStorage ts = stored_tensor;
             ComponentKind kind;
-            if (is_unused_tensor(ts.name) || !classify_tensor(ts.name, kind)) {
+            if (!classify_tensor(ts.name, kind)) {
                 continue;
             }
             if (ts.expected_type != GGML_TYPE_COUNT) {

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;
@@ -573,6 +660,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 +675,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,6 +707,9 @@ 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;
     };
 
@@ -683,17 +784,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]
@@ -219,7 +232,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);
@@ -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,

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,4 +1,5 @@
 #include <algorithm>
+#include <exception>
 #include <map>
 #include <utility>
 
@@ -341,21 +342,17 @@ void GGMLRunner::copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_c
     }
 }
 
-bool GGMLRunner::resolve_graph_cut_plan(ggml_cgraph* gf,
-                                        GraphCutPlan* plan_out) {
-    GGML_ASSERT(plan_out != nullptr);
+const GGMLRunner::GraphCutPlan& GGMLRunner::resolve_graph_cut_plan(ggml_cgraph* gf) {
     GGML_ASSERT(gf != nullptr);
-    *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
-                                                 gf,
-                                                 &graph_cut_plan_cache_,
-                                                 params_tensor_set_,
-                                                 get_desc().c_str());
-    return true;
+    return sd::ggml_graph_cut::resolve_plan(runtime_backend,
+                                            gf,
+                                            &graph_cut_plan_cache_,
+                                            params_tensor_set_,
+                                            get_desc().c_str());
 }
 
-bool GGMLRunner::resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
-                                                    GraphCutPlan* plan_out) {
-    return resolve_graph_cut_plan(gf, plan_out);
+const GGMLRunner::GraphCutPlan& GGMLRunner::resolve_graph_cut_layer_split_plan(ggml_cgraph* gf) {
+    return resolve_graph_cut_plan(gf);
 }
 
 bool GGMLRunner::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
@@ -368,10 +365,7 @@ bool GGMLRunner::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
         return false;
     }
 
-    GraphCutPlan plan;
-    if (!resolve_graph_cut_layer_split_plan(gf, &plan)) {
-        return false;
-    }
+    const auto& plan = resolve_graph_cut_layer_split_plan(gf);
     if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
         auto manager = residency_manager.lock();
         if (manager == nullptr) {
@@ -529,6 +523,7 @@ GGMLRunnerContext GGMLRunner::get_context() {
     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;
@@ -642,8 +637,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);
     }
@@ -811,24 +813,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);
         }
@@ -841,8 +844,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(),
@@ -901,7 +904,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);
             }
@@ -956,6 +962,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

@@ -71,6 +71,7 @@ struct GGMLRunnerContext {
     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;
@@ -178,6 +179,7 @@ protected:
     float linear_scale         = 0.f;
     float attn_scale           = 0.f;
     bool conv2d_direct_enabled = false;
+    bool conv3d_direct_enabled = false;
     bool circular_x_enabled    = false;
     bool circular_y_enabled    = false;
 
@@ -263,11 +265,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);
 
@@ -346,6 +346,10 @@ public:
         conv2d_direct_enabled = enabled;
     }
 
+    void set_conv3d_direct_enabled(bool enabled) {
+        conv3d_direct_enabled = enabled;
+    }
+
     void set_circular_axes(bool circular_x, bool circular_y) {
         circular_x_enabled = circular_x;
         circular_y_enabled = circular_y;

src/core/parallel.cpp

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

src/core/parallel.h

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

src/core/regex.cpp

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

src/core/regex.h

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

src/core/rng.hpp

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

src/core/rng_mt19937.hpp

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

src/core/rng_philox.hpp

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

src/core/tensor.hpp

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

src/core/util.cpp

@@ -62,17 +62,34 @@ void replace_all_chars(std::string& str, char target, char replacement) {
     }
 }
 
+static std::string sd_vformat(const char* fmt, va_list ap) {
+    char small[128];
+    va_list ap2;
+    va_copy(ap2, ap);
+    int size = vsnprintf(small, sizeof small, fmt, ap);
+    if (size < 0) {
+        va_end(ap2);
+        return {};
+    }
+    size_t needed = (size_t)size;
+    if (needed < sizeof small) {
+        va_end(ap2);
+        return std::string(small, needed);
+    }
+    std::string out(needed, '\0');
+    int size2 = vsnprintf(out.data(), needed + 1, fmt, ap2);
+    va_end(ap2);
+    if (size2 < 0)
+        out.clear();
+    return out;
+}
+
 std::string sd_format(const char* fmt, ...) {
     va_list ap;
-    va_list ap2;
     va_start(ap, fmt);
-    va_copy(ap2, ap);
-    int size = vsnprintf(nullptr, 0, fmt, ap);
-    std::vector<char> buf(size + 1);
-    int size2 = vsnprintf(buf.data(), size + 1, fmt, ap2);
-    va_end(ap2);
+    std::string result = sd_vformat(fmt, ap);
     va_end(ap);
-    return std::string(buf.data(), size);
+    return result;
 }
 
 int round_up_to(int value, int base) {
@@ -414,14 +431,43 @@ std::vector<std::string> split_string(const std::string& str, char delimiter) {
 }
 
 ggml_type sd_type_to_ggml_type(sd_type_t sdtype) {
+    if (sdtype == SD_TYPE_F8_E4M3 || sdtype == SD_TYPE_F8_E5M2) {
+#ifndef SD_USE_UPSTREAM_GGML
+        return sdtype == SD_TYPE_F8_E4M3 ? GGML_TYPE_F8_E4M3 : GGML_TYPE_F8_E5M2;
+#else
+        return GGML_TYPE_COUNT;
+#endif
+    }
     const int type_value = static_cast<int>(sdtype);
-    if (type_value < std::min<int>(SD_TYPE_COUNT, GGML_TYPE_COUNT)) {
+    if (type_value >= 0 && type_value < std::min<int>(SD_TYPE_COUNT, GGML_TYPE_COUNT)) {
         return static_cast<ggml_type>(type_value);
     } else {
         return GGML_TYPE_COUNT;
     }
 }
 
+bool validate_tensor_types(sd_type_t type, const char* tensor_type_rules) {
+    if (type != SD_TYPE_COUNT && sd_type_to_ggml_type(type) == GGML_TYPE_COUNT) {
+        LOG_ERROR("weight type %s is not supported by this ggml build", sd_type_name(type));
+        return false;
+    }
+#ifdef SD_USE_UPSTREAM_GGML
+    for (const auto& rule : split_string(SAFE_STR(tensor_type_rules), ',')) {
+        const auto pos = rule.find('=');
+        if (pos != std::string::npos) {
+            const auto name = rule.substr(pos + 1);
+            if (name == "f8_e4m3" || name == "f8_e5m2") {
+                LOG_ERROR("FP8 is not supported by this ggml build (tensor type rule '%s')", rule.c_str());
+                return false;
+            }
+        }
+    }
+#else
+    GGML_UNUSED(tensor_type_rules);
+#endif
+    return true;
+}
+
 KeyValueArgs parse_key_value_args(const char* args, const char* context) {
     KeyValueArgs pairs;
 
@@ -595,47 +641,45 @@ std::string trim(const std::string& s) {
 static sd_log_cb_t sd_log_cb = nullptr;
 void* sd_log_cb_data         = nullptr;
 
-#define LOG_BUFFER_SIZE 4096
+static void sd_log_dispatch(sd_log_level_t level, const std::string& origin, const std::string& text) {
+    if (sd_log_cb == nullptr)
+        return;
+    std::string message = origin + " - " + text;
+    if (message.back() != '\n') {
+        message += '\n';
+    }
+    sd_log_cb(level, message.c_str(), sd_log_cb_data);
+}
 
 void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...) {
     va_list args;
     va_start(args, format);
-
-    static char log_buffer[LOG_BUFFER_SIZE + 1];
-    int written = snprintf(log_buffer, LOG_BUFFER_SIZE, "%s:%-4d - ", sd_basename(file).c_str(), line);
-
-    if (written >= 0 && written < LOG_BUFFER_SIZE) {
-        vsnprintf(log_buffer + written, LOG_BUFFER_SIZE - written, format, args);
-    }
-    size_t len = strlen(log_buffer);
-    if (log_buffer[len - 1] != '\n') {
-        strncat(log_buffer, "\n", LOG_BUFFER_SIZE - len);
-    }
-
-    if (sd_log_cb) {
-        sd_log_cb(level, log_buffer, sd_log_cb_data);
-    }
-
+    std::string message = sd_vformat(format, args);
     va_end(args);
+    std::string origin = sd_format("%s:%-4d", sd_basename(file).c_str(), line);
+    sd_log_dispatch(level, origin, message);
 }
 
 void sd_ggml_log_callback(ggml_log_level level, const char* text, void*) {
+    sd_log_level_t sd_level = SD_LOG_VERBOSE;
     switch (level) {
         case GGML_LOG_LEVEL_DEBUG:
-            LOG_VERBOSE(text);
+            sd_level = SD_LOG_VERBOSE;
             break;
         case GGML_LOG_LEVEL_INFO:
-            LOG_INFO(text);
+            sd_level = SD_LOG_INFO;
             break;
         case GGML_LOG_LEVEL_WARN:
-            LOG_WARN(text);
+            sd_level = SD_LOG_WARN;
             break;
         case GGML_LOG_LEVEL_ERROR:
-            LOG_ERROR(text);
+            sd_level = SD_LOG_ERROR;
             break;
         default:
-            LOG_VERBOSE(text);
+            sd_level = SD_LOG_VERBOSE;
+            break;
     }
+    sd_log_dispatch(sd_level, "ggml", text);
 }
 
 void sd_set_log_callback(sd_log_cb_t cb, void* data) {
@@ -755,6 +799,13 @@ sd::Tensor<float> clip_preprocess(const sd::Tensor<float>& image, int target_wid
     int64_t resized_width  = static_cast<int64_t>(scale * static_cast<float>(image.shape()[0]));
     int64_t resized_height = static_cast<int64_t>(scale * static_cast<float>(image.shape()[1]));
 
+    // The resized image must cover the crop window. Floating-point rounding can
+    // leave a side one pixel short of the crop target (e.g. 730 -> 735.999...
+    // -> 735 after truncation), so clamp to keep the center crop in bounds.
+    // Truncation is otherwise preserved to avoid changing existing results.
+    resized_width  = std::max<int64_t>(resized_width, target_width);
+    resized_height = std::max<int64_t>(resized_height, target_height);
+
     sd::Tensor<float> resized = sd::ops::interpolate(
         image,
         {resized_width, resized_height, image.shape()[2], image.shape()[3]});
@@ -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 {
@@ -223,6 +248,10 @@ struct PhotoMakerExtension : public GenerationExtension {
                                                                                   trigger_token_count);
         SDCondition prepared_id_condition = std::get<0>(cond_tup);
         auto class_tokens_mask            = std::get<1>(cond_tup);
+        if (prepared_id_condition.empty()) {
+            LOG_ERROR("failed to encode PhotoMaker prompt");
+            return false;
+        }
         if (std::find(class_tokens_mask.begin(), class_tokens_mask.end(), true) == class_tokens_mask.end()) {
             LOG_WARN("PhotoMaker trigger word '%s' was not found in prompt", trigger_word.c_str());
             LOG_WARN("Turn off PhotoMaker for this request");
@@ -263,11 +292,16 @@ struct PhotoMakerExtension : public GenerationExtension {
 
         prepared_id_condition.c_crossattn = std::move(res);
         int64_t t1                        = ggml_time_ms();
-        id_condition                      = std::move(prepared_id_condition);
-        start_merge_step                  = int(ctx.pm_params.style_strength / 100.f * ctx.total_steps);
-        ctx.condition_params.text         = remove_photomaker_trigger_from_prompt(*clip_conditioner,
-                                                                                  ctx.condition_params.text,
-                                                                                  trigger_word);
+        std::string prompt;
+        if (!remove_photomaker_trigger_from_prompt(*clip_conditioner,
+                                                   ctx.condition_params.text,
+                                                   trigger_word,
+                                                   prompt)) {
+            return false;
+        }
+        id_condition              = std::move(prepared_id_condition);
+        start_merge_step          = int(ctx.pm_params.style_strength / 100.f * ctx.total_steps);
+        ctx.condition_params.text = std::move(prompt);
         LOG_INFO("Photomaker ID Stacking, taking %" PRId64 " ms", t1 - t0);
         LOG_INFO("PHOTOMAKER: start_merge_step: %d", start_merge_step);
 

src/model.h

@@ -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,
@@ -130,7 +132,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 +146,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;

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

@@ -208,6 +208,7 @@ public:
         ggml_tensor* w            = params["weight"];
         const float scale         = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale;
         ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr;
+#ifndef SD_USE_UPSTREAM_GGML
         if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) {
             bool supports_fp8_matmul = false;
             if (ctx->backend != nullptr) {
@@ -221,6 +222,7 @@ public:
                 w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16);
             }
         }
+#endif
         ggml_tensor* b = nullptr;
         if (bias) {
             b = params["bias"];
@@ -238,6 +240,7 @@ public:
             if (ctx->weight_adapter && b != nullptr) {
                 b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
             }
+#ifndef SD_USE_UPSTREAM_GGML
             if (int8_convrot && scale == 1.f) {
                 const auto cache_key = std::make_pair(x, int8_convrot_group_size);
                 auto cached          = ctx->int8_convrot_cache.find(cache_key);
@@ -248,6 +251,7 @@ public:
                     x = cached->second;
                 }
             }
+#endif
             out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
                                                 x,
                                                 w,
@@ -368,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;
@@ -673,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);
     }
 };
 
@@ -766,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);
     }
 };
 

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

src/model/diffusion/hidream_o1.hpp

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

src/model/diffusion/model.hpp

@@ -66,9 +66,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 {
@@ -130,6 +136,7 @@ using DiffusionExtraParams = std::variant<std::monostate,
                                           SkipLayerDiffusionExtra,
                                           FluxDiffusionExtra,
                                           AnimaDiffusionExtra,
+                                          QwenImage21DiffusionExtra,
                                           WanDiffusionExtra,
                                           HiDreamO1DiffusionExtra,
                                           LTXAVDiffusionExtra,

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;
                 }
@@ -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/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;
@@ -139,7 +140,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) {
@@ -230,8 +232,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 +280,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 +334,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;
         }
     };
@@ -1359,7 +1377,7 @@ namespace LLM {
                 x        = ggml_ext_cont(ctx->ggml_ctx, kqv);
                 x        = ggml_reshape_3d(ctx->ggml_ctx, x, head_dim * num_heads, n_token, N);
             } else {
-                x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, false);  // [N, n_token, hidden_size]
+                x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled);  // [N, n_token, hidden_size]
             }
 
             x = out_proj->forward(ctx, x);  // [N, n_token, hidden_size]
@@ -1886,12 +1904,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) {
@@ -2338,7 +2352,7 @@ namespace LLM {
     };
 
     struct LLMEmbedder {
-        std::shared_ptr<BPETokenizer> tokenizer;
+        std::shared_ptr<Tokenizer> tokenizer;
         LLMRunner model;
 
         LLMEmbedder(LLMArch arch,
@@ -2346,14 +2360,27 @@ 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;
+            }
+            tokenizer = tokenizers.create(TokenizerConfig::MAIN, model.config.vocab_size, pad_id);
+            if (!tokenizer) {
+                if (arch == LLMArch::GPT_OSS_20B || arch == LLMArch::GEMMA2_2B) {
+                    throw std::runtime_error("GPT-OSS and Gemma 2 require an external tokenizer.json in the main tokenizer slot");
+                }
+                if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
+                    tokenizer = std::make_shared<MistralTokenizer>();
+                } else {
+                    tokenizer = std::make_shared<Qwen2Tokenizer>();
+                }
             }
         }
 
@@ -2389,7 +2416,10 @@ namespace LLM {
             for (const auto& item : parsed_attention) {
                 const std::string& curr_text = item.first;
                 float curr_weight            = item.second;
-                std::vector<int> curr_tokens = tokenizer->tokenize(curr_text, nullptr);
+                std::vector<int> curr_tokens;
+                if (!tokenizer->tokenize(curr_text, curr_tokens, nullptr)) {
+                    return {};
+                }
                 tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
                 weights.insert(weights.end(), curr_tokens.size(), curr_weight);
             }

src/model/te/t5.hpp

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

src/model/vae/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++;
             }
         }
@@ -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,

src/model_io/safetensors_io.cpp

@@ -86,10 +86,17 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) {
         ttype = GGML_TYPE_F32;
     } else if (dtype == "F64") {
         ttype = GGML_TYPE_F32;
+#ifdef SD_USE_UPSTREAM_GGML
+    } else if (dtype == "F8_E4M3") {
+        ttype = GGML_TYPE_F16;
+    } else if (dtype == "F8_E5M2") {
+        ttype = GGML_TYPE_F16;
+#else
     } else if (dtype == "F8_E4M3") {
         ttype = GGML_TYPE_F8_E4M3;
     } else if (dtype == "F8_E5M2") {
         ttype = GGML_TYPE_F8_E5M2;
+#endif
     } else if (dtype == "I32") {
         ttype = GGML_TYPE_I32;
     } else if (dtype == "I64") {
@@ -230,6 +237,12 @@ bool read_safetensors_file(const std::string& file_path,
         if (!read_comfy_quant_config(file, file_path, name, data_start + begin, end - begin, config, error)) {
             return false;
         }
+#ifdef SD_USE_UPSTREAM_GGML
+        if (config.format == "int8_tensorwise") {
+            set_error(error, "INT8 tensorwise/convrot is not supported by this ggml build (tensor '" + name + "')");
+            return false;
+        }
+#endif
         const std::string module_name = name.substr(0, name.size() - std::string(".comfy_quant").size());
         comfy_quant_configs.emplace(module_name, std::move(config));
     }
@@ -247,6 +260,11 @@ bool read_safetensors_file(const std::string& file_path,
         std::string dtype    = tensor_info["dtype"];
         nlohmann::json shape = tensor_info["shape"];
 
+        // ComfyUI FP8 activation scales cancel when inference uses F16/F32 activations.
+        if (ends_with(name, ".scale_input")) {
+            continue;
+        }
+
         size_t begin = tensor_info["data_offsets"][0].get<size_t>();
         size_t end   = tensor_info["data_offsets"][1].get<size_t>();
         if (begin > end || end > file_size_ - data_start) {
@@ -357,10 +375,20 @@ bool read_safetensors_file(const std::string& file_path,
         bool tensor_size_ok;
         if (dtype == "F8_E4M3") {
             tensor_storage.is_f8_e4m3 = true;
-            tensor_size_ok            = (tensor_storage.nbytes() == tensor_data_size);
+#ifdef SD_USE_UPSTREAM_GGML
+            // f8 -> f16
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
+#else
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
+#endif
         } else if (dtype == "F8_E5M2") {
             tensor_storage.is_f8_e5m2 = true;
-            tensor_size_ok            = (tensor_storage.nbytes() == tensor_data_size);
+#ifdef SD_USE_UPSTREAM_GGML
+            // f8 -> f16
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
+#else
+            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
+#endif
         } else if (dtype == "F64") {
             tensor_storage.is_f64 = true;
             // f64 -> f32

src/model_io/tensor_storage.h

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

src/model_loader.cpp

@@ -35,53 +35,66 @@
 
 /*================================================= Preprocess ==================================================*/
 
-const char* unused_tensors[] = {
-    "betas",
-    "alphas_cumprod_prev",
-    "sqrt_alphas_cumprod",
-    "sqrt_one_minus_alphas_cumprod",
-    "log_one_minus_alphas_cumprod",
-    "sqrt_recip_alphas_cumprod",
-    "sqrt_recipm1_alphas_cumprod",
-    "posterior_variance",
-    "posterior_log_variance_clipped",
-    "posterior_mean_coef1",
-    "posterior_mean_coef2",
-    "cond_stage_model.transformer.text_model.embeddings.position_ids",
-    "cond_stage_model.1.model.text_model.embeddings.position_ids",
-    "cond_stage_model.transformer.vision_model.embeddings.position_ids",
-    "cond_stage_model.model.logit_scale",
-    "conditioner.embedders.0.transformer.text_model.embeddings.position_ids",
-    "conditioner.embedders.0.model.logit_scale",
-    "conditioner.embedders.1.model.logit_scale",
-    "model.diffusion_model.time_embedding.cond_proj.weight",
-    "unet.time_embedding.cond_proj.weight",
-    "model_ema.decay",
-    "model_ema.num_updates",
-    "model_ema.diffusion_model",
-    "embedding_manager",
-    "denoiser.sigmas",
-    "text_encoders.t5xxl.transformer.encoder.embed_tokens.weight",  // only used during training
-    "ztsnr",                                                        // Found in some SDXL vpred models
-    "edm_vpred.sigma_min",                                          // Found in CosXL
-    // TODO: find another way to avoid the "unknown tensor" for these two
-    // "edm_vpred.sigma_max", // Used to detect CosXL
-    // "v_pred", // Used to detect SDXL vpred models
-    "text_encoders.llm.output.weight",
-    "text_encoders.llm.lm_head.",
-    "language_model.lm_head.",
-    "vision_model.",
-};
-
-bool is_unused_tensor(const std::string& name) {
-    for (size_t i = 0; i < sizeof(unused_tensors) / sizeof(const char*); i++) {
-        if (starts_with(name, unused_tensors[i])) {
-            return true;
-        }
+#ifdef SD_USE_UPSTREAM_GGML
+uint16_t f8_e4m3_to_f16(uint8_t f8) {
+    const uint32_t exponent_bias = 7;
+    if (f8 == 0xff) {
+        return ggml_fp32_to_fp16(-NAN);
+    } else if (f8 == 0x7f) {
+        return ggml_fp32_to_fp16(NAN);
     }
-    return false;
+
+    uint32_t sign     = f8 & 0x80;
+    uint32_t exponent = (f8 & 0x78) >> 3;
+    uint32_t mantissa = f8 & 0x07;
+    uint32_t result   = sign << 24;
+    if (exponent == 0) {
+        if (mantissa > 0) {
+            exponent = 0x7f - exponent_bias;
+
+            // yes, 2 times
+            if ((mantissa & 0x04) == 0) {
+                mantissa &= 0x03;
+                mantissa <<= 1;
+                exponent -= 1;
+            }
+            if ((mantissa & 0x04) == 0) {
+                mantissa &= 0x03;
+                mantissa <<= 1;
+                exponent -= 1;
+            }
+
+            result |= (mantissa & 0x03) << 21;
+            result |= exponent << 23;
+        }
+    } else {
+        result |= mantissa << 20;
+        exponent += 0x7f - exponent_bias;
+        result |= exponent << 23;
+    }
+
+    return ggml_fp32_to_fp16(*reinterpret_cast<const float*>(&result));
 }
 
+uint16_t f8_e5m2_to_f16(uint8_t fp8) {
+    return static_cast<uint16_t>(fp8) << 8;
+}
+
+void f8_e4m3_to_f16_vec(uint8_t* src, uint16_t* dst, int64_t n) {
+    // support inplace op
+    for (int64_t i = n - 1; i >= 0; i--) {
+        dst[i] = f8_e4m3_to_f16(src[i]);
+    }
+}
+
+void f8_e5m2_to_f16_vec(uint8_t* src, uint16_t* dst, int64_t n) {
+    // support inplace op
+    for (int64_t i = n - 1; i >= 0; i--) {
+        dst[i] = f8_e5m2_to_f16(src[i]);
+    }
+}
+#endif
+
 void f64_to_f32_vec(double* src, float* dst, int64_t n) {
     // support inplace op
     for (int64_t i = 0; i < n; i++) {
@@ -285,10 +298,6 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
     size_t file_index = add_file_path(file_path);
 
     for (auto& tensor_storage : tensor_storages) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
         if (!starts_with(tensor_storage.name, prefix)) {
             tensor_storage.name = prefix + tensor_storage.name;
         }
@@ -357,10 +366,6 @@ bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, cons
     size_t file_index = add_file_path(file_path);
 
     for (auto& tensor_storage : tensor_storages) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
         if (!starts_with(tensor_storage.name, prefix)) {
             tensor_storage.name = prefix + tensor_storage.name;
         }
@@ -435,6 +440,7 @@ SDVersion ModelLoader::get_sd_version() const {
     bool is_flux2                    = false;
     bool has_single_block_47         = false;
     bool is_wan                      = false;
+    bool is_s2v                      = false;
     int64_t patch_embedding_channels = 0;
     bool has_img_emb                 = false;
     bool has_middle_block_1          = false;
@@ -477,6 +483,9 @@ SDVersion ModelLoader::get_sd_version() const {
         if (tensor_storage.name.find("language_model.model.layers.0.self_attn.q_proj_mot_gen.weight") != std::string::npos) {
             return VERSION_SENSENOVA_U1_5;
         }
+        if (tensor_storage.name == "model.diffusion_model.txt_in.text_norm.weight") {
+            return VERSION_QWEN_IMAGE_2_1;
+        }
         if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
             auto img_in = tensor_storage_map.find("model.diffusion_model.img_in.weight");
             if (img_in != tensor_storage_map.end() && img_in->second.ne[0] == 128) {
@@ -524,6 +533,11 @@ SDVersion ModelLoader::get_sd_version() const {
         if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
             is_wan = true;
         }
+        if (tensor_storage.name.find("casual_audio_encoder.weights") != std::string::npos ||
+            tensor_storage.name.find("audio_injector.injector.0.q.weight") != std::string::npos) {
+            // S2V and T2V-14B share patch_embedding shapes.
+            is_s2v = true;
+        }
         if (tensor_storage.name.find("model.diffusion_model.patch_embedder.weight") != std::string::npos) {
             return VERSION_LINGBOT_VIDEO;
         }
@@ -587,6 +601,9 @@ SDVersion ModelLoader::get_sd_version() const {
     }
     if (is_wan) {
         LOG_VERBOSE("patch_embedding_channels %d", patch_embedding_channels);
+        if (is_s2v) {
+            return VERSION_WAN2_2_S2V;
+        }
         if (patch_embedding_channels == 184320 && !has_img_emb) {
             return VERSION_WAN2_2_I2V;
         }
@@ -666,10 +683,6 @@ SDVersion ModelLoader::get_sd_version() const {
 std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() const {
     std::map<ggml_type, uint32_t> wtype_stat;
     for (auto& [name, tensor_storage] : tensor_storage_map) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
         auto iter = wtype_stat.find(tensor_storage.type);
         if (iter != wtype_stat.end()) {
             iter->second++;
@@ -683,10 +696,6 @@ std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() const {
 std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() const {
     std::map<ggml_type, uint32_t> wtype_stat;
     for (auto& [name, tensor_storage] : tensor_storage_map) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
         if ((tensor_storage.name.find("text_encoders") == std::string::npos &&
              tensor_storage.name.find("cond_stage_model") == std::string::npos &&
              tensor_storage.name.find("te.text_model.") == std::string::npos &&
@@ -707,10 +716,6 @@ std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() const {
 std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() const {
     std::map<ggml_type, uint32_t> wtype_stat;
     for (auto& [name, tensor_storage] : tensor_storage_map) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
         if (tensor_storage.name.find("model.diffusion_model.") == std::string::npos && tensor_storage.name.find("unet.") == std::string::npos) {
             continue;
         }
@@ -728,10 +733,6 @@ std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() cons
 std::map<ggml_type, uint32_t> ModelLoader::get_vae_wtype_stat() const {
     std::map<ggml_type, uint32_t> wtype_stat;
     for (auto& [name, tensor_storage] : tensor_storage_map) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
-
         if (tensor_storage.name.find("vae.") == std::string::npos &&
             tensor_storage.name.find("first_stage_model") == std::string::npos) {
             continue;
@@ -815,9 +816,6 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
 
     std::vector<TensorStorage> processed_tensor_storages;
     for (const auto& [name, tensor_storage] : tensor_storage_map) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
         processed_tensor_storages.push_back(tensor_storage);
     }
 
@@ -923,6 +921,10 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
 
             if (tensor_storage.is_f64 ||
                 tensor_storage.is_i64 ||
+#ifdef SD_USE_UPSTREAM_GGML
+                tensor_storage.is_f8_e4m3 ||
+                tensor_storage.is_f8_e5m2 ||
+#endif
                 tensor_storage.type != dst_tensor->type) {
                 continue;
             }
@@ -1212,6 +1214,12 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                         f64_to_f32_vec((double*)read_buf, (float*)target_buf, tensor_storage.nelements());
                     } else if (tensor_storage.is_i64) {
                         i64_to_i32_vec((int64_t*)read_buf, (int32_t*)target_buf, tensor_storage.nelements());
+#ifdef SD_USE_UPSTREAM_GGML
+                    } else if (tensor_storage.is_f8_e4m3) {
+                        f8_e4m3_to_f16_vec((uint8_t*)read_buf, (uint16_t*)target_buf, tensor_storage.nelements());
+                    } else if (tensor_storage.is_f8_e5m2) {
+                        f8_e5m2_to_f16_vec((uint8_t*)read_buf, (uint16_t*)target_buf, tensor_storage.nelements());
+#endif
                     }
                     if (tensor_storage.type != dst_tensor->type) {
                         if (convert_buf == nullptr) {
@@ -1543,9 +1551,6 @@ int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type)
     int64_t mem_size = 0;
     std::vector<TensorStorage> processed_tensor_storages;
     for (auto [name, tensor_storage] : tensor_storage_map) {
-        if (is_unused_tensor(tensor_storage.name)) {
-            continue;
-        }
         if (tensor_should_be_converted(tensor_storage, type)) {
             tensor_storage.type = type;
         }

src/model_loader.h

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

src/model_loader_files.cpp

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

src/model_manager.cpp

@@ -274,6 +274,7 @@ bool ModelManager::register_param_tensors(ModelComponent component,
         new_states.push_back(std::move(state));
     }
 
+    resolved_tensor_states_.clear();
     for (auto& state : new_states) {
         TensorState* registered_state                      = state.get();
         tensor_states_by_tensor_[registered_state->tensor] = registered_state;
@@ -369,6 +370,7 @@ bool ModelManager::unregister_tensor_states(const std::unordered_set<TensorState
         }
     }
 
+    resolved_tensor_states_.clear();
     for (auto it = tensor_states_by_tensor_.begin(); it != tensor_states_by_tensor_.end();) {
         if (target_states.count(it->second) > 0) {
             it = tensor_states_by_tensor_.erase(it);
@@ -1199,22 +1201,52 @@ bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*
                                                   std::vector<TensorState*>& required_states,
                                                   ggml_backend_t compute_backend) const {
     required_states.clear();
+    required_states.reserve(tensors.size());
+    auto append_states = [&](const std::vector<TensorState*>& states) {
+        for (TensorState* state : states) {
+            if (compute_backend == nullptr || state->compute_backend == nullptr ||
+                state->compute_backend == compute_backend) {
+                required_states.push_back(state);
+            }
+        }
+    };
+    for (auto it = resolved_tensor_states_.begin(); it != resolved_tensor_states_.end(); ++it) {
+        if (it->tensors == tensors) {
+            append_states(it->states);
+            resolved_tensor_states_.splice(resolved_tensor_states_.begin(), resolved_tensor_states_, it);
+            return true;
+        }
+    }
+    std::vector<TensorState*> states;
+    states.reserve(tensors.size());
     std::unordered_set<TensorState*> seen;
+    seen.reserve(tensors.size());
+    bool cacheable = true;
     for (ggml_tensor* tensor : tensors) {
         if (tensor == nullptr) {
             continue;
         }
-        auto param = resolve_param_tensor(tensor);
-        auto found = tensor_states_by_tensor_.find(param);
+        auto found = tensor_states_by_tensor_.find(tensor);
+        // Unregistered views can be rebound without changing the parameter list.
+        cacheable &= found != tensor_states_by_tensor_.end();
+        for (auto view = tensor->view_src; found == tensor_states_by_tensor_.end() && view != nullptr; view = view->view_src) {
+            found = tensor_states_by_tensor_.find(view);
+        }
         if (found == tensor_states_by_tensor_.end()) {
             LOG_ERROR("model manager tensor '%s' is not registered", ggml_get_name(tensor));
             return false;
         }
         TensorState* state = found->second;
-        if ((compute_backend == nullptr || state->compute_backend == nullptr ||
-             state->compute_backend == compute_backend) &&
-            seen.insert(state).second) {
-            required_states.push_back(state);
+        if (seen.insert(state).second) {
+            states.push_back(state);
+        }
+    }
+    append_states(states);
+    if (cacheable && !tensors.empty()) {
+        static constexpr size_t MAX_RESOLVED_LISTS = 4;
+        resolved_tensor_states_.push_front({tensors, std::move(states)});
+        if (resolved_tensor_states_.size() > MAX_RESOLVED_LISTS) {
+            resolved_tensor_states_.pop_back();
         }
     }
     return true;
@@ -1274,8 +1306,7 @@ size_t ModelManager::compute_backend_alloc_size(const std::vector<TensorState*>&
     size_t total_size = 0;
     std::unordered_set<TensorState*> seen;
     for (TensorState* state : states) {
-        if (state == nullptr || state->tensor == nullptr || !seen.insert(state).second ||
-            should_ignore(*state) || is_optional_missing_tensor(state->name)) {
+        if (state == nullptr || state->tensor == nullptr) {
             continue;
         }
         const bool compute_resident =
@@ -1285,6 +1316,9 @@ size_t ModelManager::compute_backend_alloc_size(const std::vector<TensorState*>&
         if (missing_only && compute_resident) {
             continue;
         }
+        if (!seen.insert(state).second || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
+            continue;
+        }
 
         ggml_backend_buffer_type_t buffer_type = nullptr;
         if (state->compute_backend == state->params_backend) {
@@ -1584,18 +1618,35 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
     if (request.compute_backend == nullptr || sd_backend_is_cpu(request.compute_backend)) {
         return result;
     }
-    auto add                     = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; };
-    const size_t missing         = compute_backend_alloc_size(states, true);
-    result.required_device_bytes = add(request.pending_allocation_bytes, missing);
-    result.required_budget_bytes = add(request.runtime_peak_bytes(), missing);
-    auto device                  = ggml_backend_get_device(request.compute_backend);
-    if (device != nullptr) {
+    auto add             = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; };
+    const size_t missing = compute_backend_alloc_size(states, true);
+    // Backend scratch buffers and pipelines are not included in graph measurements.
+    constexpr size_t safety_margin = 512ULL * 1024ULL * 1024ULL;
+    result.required_device_bytes   = add(add(request.pending_allocation_bytes, missing), safety_margin);
+    result.required_budget_bytes   = add(request.runtime_peak_bytes(), missing);
+    auto available_device_bytes    = [&](ggml_backend_t backend) {
+        auto device = ggml_backend_get_device(backend);
+        if (device == nullptr) {
+            return SIZE_MAX;
+        }
         size_t free_bytes = 0, total_bytes = 0;
         ggml_backend_dev_memory(device, &free_bytes, &total_bytes);
-        if (free_bytes != 0 || total_bytes != 0) {
-            result.available_device_bytes = free_bytes;
+        if (free_bytes == 0 && total_bytes == 0) {
+            return SIZE_MAX;
         }
-    }
+        // Vulkan's heap budget subtraction can underflow when usage exceeds the budget.
+        if (total_bytes > 0 && free_bytes > total_bytes) {
+            return size_t{0};
+        }
+        const size_t resident = add(compute_backend_resident_bytes(backend),
+                                       add(other_runtime_resident_bytes(request.owner_id, backend),
+                                           request.runtime_resident_bytes));
+        if (total_bytes > 0) {
+            free_bytes = std::min(free_bytes, resident < total_bytes ? total_bytes - resident : 0);
+        }
+        return free_bytes;
+    };
+    result.available_device_bytes = available_device_bytes(request.compute_backend);
     if (request.max_backend_bytes > 0) {
         const size_t resident         = add(compute_backend_resident_bytes(request.compute_backend),
                                             other_runtime_resident_bytes(request.owner_id, request.compute_backend));
@@ -1619,11 +1670,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
     // GGML exposes only a split buffer's total size, not per-device allocations.
     // Charge that upper bound on every participant instead of undercounting a shard.
     for (const auto& entry : split_devices) {
-        size_t free_bytes = 0, total_bytes = 0;
-        ggml_backend_dev_memory(ggml_backend_get_device(entry.first), &free_bytes, &total_bytes);
-        if (free_bytes != 0 || total_bytes != 0) {
-            result.available_device_bytes = std::min(result.available_device_bytes, free_bytes);
-        }
+        result.available_device_bytes = std::min(result.available_device_bytes, available_device_bytes(entry.first));
         if (entry.second > 0) {
             const size_t resident         = add(compute_backend_resident_bytes(entry.first),
                                                 other_runtime_resident_bytes(request.owner_id, entry.first));
@@ -1739,12 +1786,18 @@ bool ModelManager::ensure_compute_backend_capacity(
         }
     }
 
-    const auto capacity = check_capacity(request, required_states);
-    LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %.2f MB device / %.2f MB budget",
+    const auto capacity                = check_capacity(request, required_states);
+    const std::string available_device = capacity.available_device_bytes == SIZE_MAX
+                                             ? "unknown"
+                                             : sd_format("%.2f MB", capacity.available_device_bytes / (1024.0 * 1024.0));
+    const std::string available_budget = capacity.available_budget_bytes == SIZE_MAX
+                                             ? "unlimited"
+                                             : sd_format("%.2f MB", capacity.available_budget_bytes / (1024.0 * 1024.0));
+    LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget",
              ggml_backend_name(compute_backend),
              capacity.required_device_bytes / (1024.0 * 1024.0),
              capacity.required_budget_bytes / (1024.0 * 1024.0),
-             capacity.available_device_bytes / (1024.0 * 1024.0),
-             capacity.available_budget_bytes / (1024.0 * 1024.0));
+             available_device.c_str(),
+             available_budget.c_str());
     return false;
 }

src/model_manager.h

@@ -2,6 +2,7 @@
 #define __MODEL_MANAGER_H__
 
 #include <cstdint>
+#include <list>
 #include <map>
 #include <memory>
 #include <set>
@@ -31,7 +32,7 @@ public:
     };
 
 private:
-    static constexpr size_t MAX_RESIDENCY_BLOCK_BYTES = 64ULL * 1024ULL * 1024ULL;
+    static constexpr size_t MAX_RESIDENCY_BLOCK_BYTES = 1024ULL * 1024ULL * 1024ULL;
 
     struct TensorState {
         std::string name;
@@ -84,9 +85,15 @@ private:
         size_t resident_bytes          = 0;
     };
 
+    struct ResolvedTensorStates {
+        std::vector<ggml_tensor*> tensors;
+        std::vector<TensorState*> states;
+    };
+
     ModelLoader model_loader_;
     std::vector<std::unique_ptr<TensorState>> tensor_states_;
     std::map<const ggml_tensor*, TensorState*> tensor_states_by_tensor_;
+    mutable std::list<ResolvedTensorStates> resolved_tensor_states_;
     std::vector<std::unique_ptr<ParamsStorageBlock>> params_storage_blocks_;
     std::vector<std::unique_ptr<ComputeStagingBlock>> compute_staging_blocks_;
     std::map<ggml_backend_t, ggml_backend_buffer_type_t> split_buffer_types_;

src/name_conversion.cpp

@@ -999,7 +999,30 @@ std::string convert_diffusers_vae_to_original_sd1(std::string name) {
     return result;
 }
 
-std::string convert_diffusers_to_original_wan_vae(std::string name) {
+std::string convert_diffusers_to_original_wan_vae(std::string name, bool qwen_image_2_1 = false) {
+    if (qwen_image_2_1) {
+        for (int i = 0; i < 5; ++i) {
+            const auto index = std::to_string(i);
+            for (const auto& side : {std::string("encoder"), std::string("decoder")}) {
+                const bool encoder           = side == "encoder";
+                const std::string old_prefix = side + (encoder ? ".down_blocks." : ".up_blocks.") + index + ".";
+                const std::string new_prefix = side + (encoder ? ".downsamples." : ".upsamples.") + index + ".";
+                if (!starts_with(name, old_prefix)) {
+                    continue;
+                }
+                name.replace(0, old_prefix.size(), new_prefix);
+                const std::string layers = encoder ? "downsamples." : "upsamples.";
+                for (int j = 0; j < (encoder ? 2 : 3); ++j) {
+                    const auto old_resnet = new_prefix + "resnets." + std::to_string(j) + ".";
+                    const auto new_resnet = new_prefix + layers + std::to_string(j) + ".";
+                    replace_with_prefix_map(name, std::vector<std::pair<std::string, std::string>>{{old_resnet + "conv_shortcut.", new_resnet + "shortcut."},
+                                                                                                   {old_resnet, new_resnet + "residual."}});
+                }
+                replace_with_prefix_map(name, std::vector<std::pair<std::string, std::string>>{{new_prefix + (encoder ? "downsampler." : "upsampler."),
+                                                                                                new_prefix + layers + (encoder ? "2." : "3.")}});
+            }
+        }
+    }
     static const std::vector<std::pair<std::string, std::string>> prefix_map = {
         {"quant_conv.", "conv1."},
         {"post_quant_conv.", "conv2."},
@@ -1055,7 +1078,11 @@ std::string convert_diffusers_to_original_wan_vae(std::string name) {
     };
 
     replace_with_name_map(name, shared_name_map);
-    replace_with_prefix_map(name, prefix_map);
+    if (qwen_image_2_1) {
+        replace_with_prefix_map(name, std::vector<std::pair<std::string, std::string>>{{"quant_conv.", "conv1."}, {"post_quant_conv.", "conv2."}});
+    } else {
+        replace_with_prefix_map(name, prefix_map);
+    }
 
     // Only apply the ResNet-specific renaming if the tensor belongs to a ResNet block.
     // This prevents generic ".conv1." or ".conv2." matching on top-level encoder/decoder convolutions.
@@ -1071,7 +1098,7 @@ std::string convert_first_stage_model_name(std::string name, std::string prefix,
         return name;
     }
     if (sd_version_uses_wan_vae(version)) {
-        return convert_diffusers_to_original_wan_vae(name);
+        return convert_diffusers_to_original_wan_vae(name, version == VERSION_QWEN_IMAGE_2_1);
     }
     static std::unordered_map<std::string, std::string> vae_name_map = {
         {"decoder.post_quant_conv.", "post_quant_conv."},
@@ -1489,7 +1516,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
 
     replace_with_prefix_map(name, prefix_map);
 
-    if (sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
+    if (version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
         const std::string hf_vision_prefix = "text_encoders.llm.model.visual.";
         if (starts_with(name, hf_vision_prefix)) {
             name = "text_encoders.llm.visual." + name.substr(hf_vision_prefix.size());

src/pipeline/diffusion_engine.cpp

@@ -33,12 +33,14 @@
 #include "extensions/generation_extension.h"
 #include "model/adapter/ip_adapter.hpp"
 #include "model/adapter/lora.hpp"
+#include "model/audio/wav2vec2.hpp"
 #include "model/diffusion/animatediff.hpp"
 #include "model/diffusion/control.hpp"
 #include "model/diffusion/model.hpp"
 #include "model/vae/audio_vae.hpp"
 #include "model/vae/ltx_vae.hpp"
 #include "model/vae/vae.hpp"
+#include "runtime/audio_processing.h"
 #include "runtime/denoiser.hpp"
 #include "runtime/guidance.h"
 #include "runtime/preview_interval.h"
@@ -74,9 +76,11 @@ const char* model_version_to_str[] = {
     "Wan 2.x",
     "Wan 2.2 I2V",
     "Wan 2.2 TI2V",
+    "Wan 2.2 S2V",
     "LingBot Video",
     "Qwen Image",
     "Qwen Image Layered",
+    "Qwen Image 2.1",
     "Hunyuan Video",
     "Anima",
     "Flux.2",
@@ -136,7 +140,7 @@ StableDiffusionGGML::~StableDiffusionGGML() = default;
 
 const std::map<StableDiffusionGGML::RunnerGroup, std::set<ModelComponent>>& StableDiffusionGGML::runner_components() {
     static const std::map<RunnerGroup, std::set<ModelComponent>> components{
-        {RunnerGroup::Core, {ModelComponent::Conditioner, ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion, ModelComponent::CLIPVision, ModelComponent::IPAdapter}},
+        {RunnerGroup::Core, {ModelComponent::Conditioner, ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion, ModelComponent::CLIPVision, ModelComponent::IPAdapter, ModelComponent::AudioEncoder}},
         {RunnerGroup::VAE, {ModelComponent::VAE, ModelComponent::PreviewVAE, ModelComponent::AudioVAE}},
         {RunnerGroup::ControlNet, {ModelComponent::ControlNet}},
         {RunnerGroup::Extensions, {ModelComponent::PhotoMaker, ModelComponent::PuLID}},
@@ -804,6 +808,13 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf
         }
     }
 
+    if (strlen(SAFE_STR(sd_ctx_params->audio_encoder_path)) > 0) {
+        LOG_INFO("loading audio encoder (wav2vec2) from '%s'", sd_ctx_params->audio_encoder_path);
+        if (!model_loader.init_from_file(sd_ctx_params->audio_encoder_path, "wav2vec2.")) {
+            LOG_WARN("loading audio encoder weights from '%s' failed", sd_ctx_params->audio_encoder_path);
+        }
+    }
+
     if (strlen(SAFE_STR(sd_ctx_params->motion_module_path)) > 0) {
         LOG_INFO("loading motion module (AnimateDiff) from '%s'", sd_ctx_params->motion_module_path);
         if (!model_loader.init_from_file(sd_ctx_params->motion_module_path,
@@ -847,14 +858,25 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf
 }
 
 bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) {
+#ifdef SD_USE_UPSTREAM_GGML
+    LOG_WARN(
+        "Using upstream GGML: INT8 tensorwise/convrot is disabled and FP8 weights are "
+        "converted to F16 at load time. Some operators may be unsupported and performance "
+        "may be lower than with patched GGML.");
+#endif
+    if (!validate_tensor_types(sd_ctx_params->wtype, sd_ctx_params->tensor_type_rules)) {
+        return false;
+    }
     for (float scale : {sd_ctx_params->linear_scale, sd_ctx_params->attn_scale}) {
         if (!std::isfinite(scale) || scale < 0.f || (scale > 0.f && !std::isfinite(1.f / scale))) {
             LOG_ERROR("scale overrides must be finite positive values, or 0 to keep model defaults");
             return false;
         }
     }
-    auto configuration        = std::make_unique<ModelConfig>(*sd_ctx_params);
-    n_threads                 = sd_ctx_params->n_threads;
+    auto configuration = std::make_unique<ModelConfig>(*sd_ctx_params);
+    n_threads          = sd_ctx_params->n_threads;
+    tensor_executor    = std::make_unique<sd::ParallelExecutor>(n_threads > 0 ? n_threads : sd_get_num_physical_cores());
+    sd::ParallelScope tensor_scope(tensor_executor.get());
     enable_mmap               = sd_ctx_params->enable_mmap;
     disable_prefetch          = sd_ctx_params->disable_prefetch;
     disable_segmented_compute = sd_ctx_params->disable_segmented_compute;
@@ -1010,6 +1032,7 @@ bool StableDiffusionGGML::build_core_runners() {
     high_noise_diffusion_model = std::move(runners.high_noise_diffusion);
     clip_vision                = std::move(runners.clip_vision);
     ip_adapter                 = std::move(runners.ip_adapter);
+    audio_encoder              = std::move(runners.audio_encoder);
 
     cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::TE));
     diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION));
@@ -1019,11 +1042,15 @@ bool StableDiffusionGGML::build_core_runners() {
     if (clip_vision) {
         clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CLIP_VISION));
     }
+    if (audio_encoder) {
+        audio_encoder->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::AUDIO_ENCODER));
+    }
     return register_runner_params(ModelComponent::Conditioner, cond_stage_model, SDBackendModule::TE) &&
            register_runner_params(ModelComponent::Diffusion, diffusion_model, SDBackendModule::DIFFUSION) &&
            register_runner_params(ModelComponent::HighNoiseDiffusion, high_noise_diffusion_model, SDBackendModule::DIFFUSION) &&
            register_runner_params(ModelComponent::CLIPVision, clip_vision, SDBackendModule::CLIP_VISION) &&
-           register_runner_params(ModelComponent::IPAdapter, ip_adapter, SDBackendModule::DIFFUSION);
+           register_runner_params(ModelComponent::IPAdapter, ip_adapter, SDBackendModule::DIFFUSION) &&
+           register_runner_params(ModelComponent::AudioEncoder, audio_encoder, SDBackendModule::AUDIO_ENCODER);
 }
 
 bool StableDiffusionGGML::build_vae_runners() {
@@ -1121,6 +1148,12 @@ bool StableDiffusionGGML::validate_and_load_runners() {
     ignore_tensors.insert("model.diffusion_model.__32x32__");
     ignore_tensors.insert("model.diffusion_model.__index_timestep_zero__");
 
+    if (audio_encoder != nullptr) {
+        // These wav2vec2 tensors are unused during feature extraction.
+        ignore_tensors.insert("wav2vec2.lm_head.");
+        ignore_tensors.insert("wav2vec2.masked_spec_embed");
+    }
+
     if (audio_vae_model) {
         if (!sd_version_is_minimax_h3(version)) {
             ignore_tensors.insert("audio_vae.encoder");
@@ -1755,6 +1788,29 @@ sd::Tensor<float> StableDiffusionGGML::get_clip_vision_output(const sd::Tensor<f
     return output;
 }
 
+// Returns 50 Hz wav2vec2 states in sd::Tensor layout: [dim, frames, layers].
+sd::Tensor<float> StableDiffusionGGML::get_audio_embedding(const sd_audio_t& audio) {
+    if (audio_encoder == nullptr) {
+        LOG_ERROR("audio encoder model is not loaded");
+        return {};
+    }
+    if (audio.data == nullptr || audio.sample_count == 0 || audio.channels == 0 || audio.sample_rate == 0) {
+        LOG_ERROR("invalid driving audio");
+        return {};
+    }
+    auto mono = sd::audio::downmix_to_mono(audio.data, audio.sample_count, audio.channels);
+    if (mono.empty()) {
+        LOG_ERROR("audio mono downmix failed");
+        return {};
+    }
+    mono = sd::audio::resample_audio(mono.data(), mono.size(), audio.sample_rate, 16000);
+    if (mono.empty()) {
+        LOG_ERROR("audio resample to 16 kHz failed");
+        return {};
+    }
+    return audio_encoder->compute(n_threads, mono);
+}
+
 void StableDiffusionGGML::compute_ip_adapter_tokens(const sd_image_t& image, float strength) {
     ip_adapter_tokens        = {};
     ip_adapter_uncond_tokens = {};
@@ -1810,6 +1866,10 @@ std::vector<float> StableDiffusionGGML::process_timesteps(const std::vector<floa
             }
         }
         return new_timesteps;
+    }
+    if (diffusion_model->get_desc() == "Wan2.2-S2V-14B") {
+        int64_t frame_count = init_latent.shape()[2];
+        return std::vector<float>(static_cast<size_t>(frame_count), timesteps[0]);
     } else {
         return timesteps;
     }
@@ -2110,6 +2170,10 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
 
     RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr};
 
+    const bool apply_denoise_mask = !denoise_mask.empty() &&
+                                    std::any_of(denoise_mask.values().begin(), denoise_mask.values().end(),
+                                                [](float value) { return value != 1.f; });
+
     std::vector<int> skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count);
     float cfg_scale     = guidance.txt_cfg;
     float img_cfg_scale = guidance.img_cfg;
@@ -2239,13 +2303,13 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
             hunyuan_timestep_r_tensor = sd::Tensor<float>::from_vector({sigmas[step + 1]});
         }
         sd::Tensor<float> noised_input = x * c_in;
-        if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) {
+        if (apply_denoise_mask && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) {
             noised_input = noised_input * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
         }
 
         if (cache_runtime.spectrum_enabled && cache_runtime.spectrum.should_predict()) {
             cache_runtime.spectrum.predict(&denoised);
-            if (!denoise_mask.empty()) {
+            if (apply_denoise_mask) {
                 denoised = denoised * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
             }
             if (preview_needed && sd_should_preview_denoised()) {
@@ -2313,12 +2377,15 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
             } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
                 diffusion_params.extra = FluxDiffusionExtra{&guidance_tensor,
                                                             local_skip_layers};
+            } else if (version == VERSION_QWEN_IMAGE_2_1) {
+                diffusion_params.extra = QwenImage21DiffusionExtra{&condition.c_token_types};
             } else if (sd_version_is_anima(version)) {
                 diffusion_params.extra = AnimaDiffusionExtra{condition.c_t5_ids.empty() ? nullptr : &condition.c_t5_ids,
                                                              condition.c_t5_weights.empty() ? nullptr : &condition.c_t5_weights};
             } else if (sd_version_is_wan(version)) {
                 diffusion_params.extra = WanDiffusionExtra{vace_context.empty() ? nullptr : &vace_context,
-                                                           vace_strength};
+                                                           vace_strength,
+                                                           condition.c_ref_audios.empty() ? nullptr : &condition.c_ref_audios[0]};
             } else if (sd_version_is_hunyuan_video(version)) {
                 diffusion_params.extra = HunyuanVideoDiffusionExtra{
                     &guidance_tensor,
@@ -2467,7 +2534,7 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
         if (cache_runtime.spectrum_enabled) {
             cache_runtime.spectrum.update(denoised);
         }
-        if (!denoise_mask.empty()) {
+        if (apply_denoise_mask) {
             denoised = denoised * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
         }
         if (preview_needed && sd_should_preview_denoised()) {
@@ -2512,7 +2579,7 @@ int StableDiffusionGGML::get_diffusion_model_down_factor() {
     if (sd_version_is_dit(version)) {
         if (sd_version_is_sensenova_u1(version)) {
             down_factor = 32;
-        } else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version)) {
+        } else if (version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version)) {
             down_factor = 2;
         } else {
             down_factor = 1;
@@ -2528,6 +2595,8 @@ int StableDiffusionGGML::get_latent_channel() {
             latent_channel = 128;
         } else if (sd_version_is_minimax_h3(version)) {
             latent_channel = 24;
+        } else if (version == VERSION_QWEN_IMAGE_2_1) {
+            latent_channel = 64;
         } else if (version == VERSION_WAN2_2_TI2V) {
             latent_channel = 48;
         } else if (sd_version_is_hunyuan_video(version)) {
@@ -2556,7 +2625,7 @@ int StableDiffusionGGML::get_latent_channel() {
 }
 
 int StableDiffusionGGML::get_image_channels() const {
-    return version == VERSION_QWEN_IMAGE_LAYERED ? 4 : 3;
+    return version == VERSION_QWEN_IMAGE_LAYERED || version == VERSION_QWEN_IMAGE_2_1 ? 4 : 3;
 }
 
 int StableDiffusionGGML::get_image_seq_len(int h, int w) {

src/pipeline/diffusion_engine.h

@@ -8,6 +8,7 @@
 #include <map>
 #include <memory>
 #include <mutex>
+#include <optional>
 #include <set>
 #include <string>
 #include <vector>
@@ -27,13 +28,16 @@ struct LoraModel;
 struct ConditionerParams;
 struct SDCondition;
 struct RefImageParams;
+namespace Wav2Vec2 {
+    class Wav2Vec2ModelRunner;
+}
 
 extern const char* model_version_to_str[];
 
 static inline bool sd_version_supports_ref_latent_img_cfg(SDVersion version) {
     return version == VERSION_FLUX ||
            sd_version_is_flux2(version) ||
-           sd_version_is_qwen_image(version) ||
+           (sd_version_is_qwen_image(version) && version != VERSION_QWEN_IMAGE_2_1) ||
            sd_version_is_mage_flow(version) ||
            sd_version_is_longcat(version) ||
            sd_version_is_z_image(version) ||
@@ -53,8 +57,9 @@ public:
     std::shared_ptr<RNG> rng;
     std::shared_ptr<RNG> sampler_rng = nullptr;
     int n_threads                    = -1;
-    float default_flow_shift         = INFINITY;
-    float active_flow_shift          = INFINITY;
+    std::unique_ptr<sd::ParallelExecutor> tensor_executor;
+    float default_flow_shift = INFINITY;
+    float active_flow_shift  = INFINITY;
 
     std::shared_ptr<Conditioner> cond_stage_model;
     std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision;  // for svd or wan2.1 i2v
@@ -63,6 +68,7 @@ public:
     std::shared_ptr<VAE> first_stage_model;
     std::shared_ptr<VAE> preview_vae;
     std::shared_ptr<AudioVAERunner> audio_vae_model;
+    std::shared_ptr<Wav2Vec2::Wav2Vec2ModelRunner> audio_encoder;
     std::shared_ptr<ControlNet> control_net;
     std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
     sd::Tensor<float> ip_adapter_tokens;
@@ -124,6 +130,7 @@ public:
                                 &sd_ctx_params_t::clip_g_path, &sd_ctx_params_t::clip_vision_path,
                                 &sd_ctx_params_t::t5xxl_path, &sd_ctx_params_t::llm_path,
                                 &sd_ctx_params_t::llm_vision_path, &sd_ctx_params_t::diffusion_model_path,
+                                &sd_ctx_params_t::tokenizer,
                                 &sd_ctx_params_t::high_noise_diffusion_model_path, &sd_ctx_params_t::uncond_diffusion_model_path,
                                 &sd_ctx_params_t::embeddings_connectors_path, &sd_ctx_params_t::vae_path,
                                 &sd_ctx_params_t::audio_vae_path, &sd_ctx_params_t::taesd_path,
@@ -201,6 +208,7 @@ public:
         StableDiffusionGGML& sd;
         std::unique_lock<std::recursive_mutex> lock;
         bool acquired = false;
+        std::optional<sd::ParallelScope> tensor_scope;
 
         explicit ContextOperation(StableDiffusionGGML& sd)
             : sd(sd), lock(sd.execution_mutex, std::try_to_lock) {
@@ -210,6 +218,7 @@ public:
             }
             sd.executing_ = true;
             acquired      = true;
+            tensor_scope.emplace(sd.tensor_executor.get());
         }
 
         ~ContextOperation() {
@@ -363,6 +372,8 @@ public:
                                              int clip_skip        = -1,
                                              bool zero_out_masked = false);
 
+    sd::Tensor<float> get_audio_embedding(const sd_audio_t& audio);
+
     void compute_ip_adapter_tokens(const sd_image_t& image, float strength);
 
     std::vector<float> process_timesteps(const std::vector<float>& timesteps,

src/pipeline/generation.h

@@ -28,6 +28,7 @@ namespace sd::pipeline {
         sd::Tensor<float> denoise_mask;
         sd::Tensor<float> clip_vision_output;
         sd::Tensor<float> vace_context;
+        sd::Tensor<float> s2v_audio_embed;
         int64_t ref_image_num                  = 0;
         int64_t video_conditioning_frame_count = 0;
         int64_t video_target_frame_count       = 0;
@@ -59,7 +60,8 @@ namespace sd::pipeline {
                         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);
 
     sd::Tensor<float> upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd,
                                                        const char* model_path,

src/pipeline/image.cpp

@@ -285,7 +285,8 @@ namespace sd::pipeline {
                     vae_width  = request->width;
                     vae_height = request->height;
                 } else {
-                    int target_pixels  = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : 1024 * 1024;
+                    int default_pixels = sd->version == VERSION_QWEN_IMAGE_2_1 ? request->width * request->height : 1024 * 1024;
+                    int target_pixels  = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : default_pixels;
                     int vae_image_size = std::min(target_pixels, request->width * request->height);
                     vae_width          = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]);
                     vae_height         = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0];
@@ -309,6 +310,9 @@ namespace sd::pipeline {
                             resized_ref_img.shape()[0]);
 
                 ref_latent = sd->encode_first_stage(resized_ref_img);
+                if (sd->version == VERSION_QWEN_IMAGE_2_1) {
+                    ref_images[i] = std::move(resized_ref_img);
+                }
             } else {
                 ref_latent = sd->encode_first_stage(ref_images[i]);
             }
@@ -438,6 +442,10 @@ namespace sd::pipeline {
         condition_params.zero_out_masked = false;
         auto cond                        = sd->cond_stage_model->get_learned_condition(sd->n_threads,
                                                                                        condition_params);
+        if (cond.empty()) {
+            LOG_ERROR("failed to encode prompt");
+            return std::nullopt;
+        }
         if (cond.c_concat.empty() && ref_image_params.pass_to_dit) {
             cond.c_concat = latents->concat_latent;  // TODO: optimize
         }
@@ -469,6 +477,10 @@ namespace sd::pipeline {
                 condition_params.zero_out_masked = zero_out_masked;
                 uncond                           = sd->cond_stage_model->get_learned_condition(sd->n_threads,
                                                                                                condition_params);
+                if (uncond.empty()) {
+                    LOG_ERROR("failed to encode negative prompt");
+                    return std::nullopt;
+                }
             }
             if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
                 uncond.c_concat = latents->concat_latent;  // TODO: optimize
@@ -494,6 +506,10 @@ namespace sd::pipeline {
                 }
                 img_uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
                                                                          condition_params);
+                if (img_uncond.empty()) {
+                    LOG_ERROR("failed to encode image guidance prompt");
+                    return std::nullopt;
+                }
                 if (img_uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
                     img_uncond.c_concat = latents->img_uncond_concat_latent;  // TODO: optimize
                 }

src/pipeline/model_builders.cpp

@@ -8,6 +8,7 @@
 #include "core/util.h"
 #include "extensions/generation_extension.h"
 #include "model/adapter/ip_adapter.hpp"
+#include "model/audio/wav2vec2.hpp"
 #include "model/diffusion/anima.hpp"
 #include "model/diffusion/boogu.hpp"
 #include "model/diffusion/control.hpp"
@@ -27,6 +28,7 @@
 #include "model/diffusion/model.hpp"
 #include "model/diffusion/pid.hpp"
 #include "model/diffusion/qwen_image.hpp"
+#include "model/diffusion/qwen_image_2_1.hpp"
 #include "model/diffusion/sensenova_u1.h"
 #include "model/diffusion/unet.hpp"
 #include "model/diffusion/wan.hpp"
@@ -72,12 +74,13 @@ namespace sd::model_builders {
         }
     }
 
-    bool build_core_runners(const Context& ctx, CoreRunners& runners) {
+    bool build_core_runners(const Context& ctx, CoreRunners& runners) try {
         const auto* sd_ctx_params      = &ctx.params;
         const auto& tensor_storage_map = ctx.tensor_storage_map;
         const auto version             = ctx.version;
         const auto& weight_manager     = ctx.weight_manager;
         CoreRunners result;
+        TokenizerConfig tokenizers(sd_ctx_params->tokenizer);
         if (!ensure_backend_pair(ctx.backends, SDBackendModule::TE) ||
             !ensure_backend_pair(ctx.backends, SDBackendModule::DIFFUSION)) {
             return false;
@@ -86,7 +89,8 @@ namespace sd::model_builders {
         if (sd_version_is_sd3(version)) {
             result.conditioner = std::make_shared<SD3CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                    tensor_storage_map,
-                                                                   weight_manager);
+                                                                   weight_manager,
+                                                                   tokenizers);
             result.diffusion   = std::make_shared<MMDiTRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                              tensor_storage_map,
                                                              "model.diffusion_model",
@@ -97,7 +101,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Pid::PiDRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                 tensor_storage_map,
                                                                 "model.diffusion_model.net",
@@ -108,7 +113,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Ideogram4::Ideogram4Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                             tensor_storage_map,
                                                                             "model.diffusion_model",
@@ -119,7 +125,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                true,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Krea2::Krea2Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                     tensor_storage_map,
                                                                     "model.diffusion_model",
@@ -146,11 +153,13 @@ namespace sd::model_builders {
                                                                    version,
                                                                    "",
                                                                    false,
-                                                                   weight_manager);
+                                                                   weight_manager,
+                                                                   tokenizers);
             } else {
                 result.conditioner = std::make_shared<FluxCLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                         tensor_storage_map,
-                                                                        weight_manager);
+                                                                        weight_manager,
+                                                                        tokenizers);
             }
             result.diffusion = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                   tensor_storage_map,
@@ -165,7 +174,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                   tensor_storage_map,
                                                                   "model.diffusion_model",
@@ -177,7 +187,8 @@ namespace sd::model_builders {
                                                                  tensor_storage_map,
                                                                  "text_encoders.llm",
                                                                  "text_embedding_projection",
-                                                                 weight_manager);
+                                                                 weight_manager,
+                                                                 tokenizers);
             result.diffusion   = std::make_shared<LTXV::LTXAVRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                    tensor_storage_map,
                                                                    "model.diffusion_model",
@@ -188,7 +199,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                true,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<MiniMaxH3::MiniMaxH3Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                             tensor_storage_map,
                                                                             "model.diffusion_model",
@@ -199,7 +211,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Hunyuan::HunyuanVideoRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                              tensor_storage_map,
                                                                              "model.diffusion_model",
@@ -234,6 +247,16 @@ namespace sd::model_builders {
                                                                                 tensor_storage_map,
                                                                                 weight_manager);
             }
+            if (version == VERSION_WAN2_2_S2V &&
+                tensor_storage_map.count("wav2vec2.encoder.layer_norm.bias") > 0) {
+                if (!ensure_backend_pair(ctx.backends, SDBackendModule::AUDIO_ENCODER)) {
+                    return false;
+                }
+                result.audio_encoder = std::make_shared<Wav2Vec2::Wav2Vec2ModelRunner>(ctx.backends.runtime_backend(SDBackendModule::AUDIO_ENCODER),
+                                                                                       tensor_storage_map,
+                                                                                       "wav2vec2.",
+                                                                                       weight_manager);
+            }
         } else if (sd_version_is_lingbot_video(version)) {
             bool enable_vision = false;
             for (const auto& [name, _] : tensor_storage_map) {
@@ -247,7 +270,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                enable_vision,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<LingBotVideo::LingBotVideoRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                                   tensor_storage_map,
                                                                                   "model.diffusion_model",
@@ -260,20 +284,29 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                enable_vision,
-                                                               weight_manager);
-            result.diffusion   = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
-                                                                       tensor_storage_map,
-                                                                       "model.diffusion_model",
-                                                                       version,
-                                                                       weight_manager,
-                                                                       sd_ctx_params->model_args);
+                                                               weight_manager,
+                                                               tokenizers);
+            if (version == VERSION_QWEN_IMAGE_2_1) {
+                result.diffusion = std::make_shared<Qwen::QwenImage21Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                             tensor_storage_map,
+                                                                             "model.diffusion_model",
+                                                                             weight_manager);
+            } else {
+                result.diffusion = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                           tensor_storage_map,
+                                                                           "model.diffusion_model",
+                                                                           version,
+                                                                           weight_manager,
+                                                                           sd_ctx_params->model_args);
+            }
         } else if (sd_version_is_mage_flow(version)) {
             result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                tensor_storage_map,
                                                                version,
                                                                "",
                                                                true,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<MageFlow::MageFlowRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                           tensor_storage_map,
                                                                           "model.diffusion_model",
@@ -284,7 +317,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                true,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                   tensor_storage_map,
                                                                   "model.diffusion_model",
@@ -294,7 +328,8 @@ namespace sd::model_builders {
         } else if (version == VERSION_HIDREAM_O1) {
             result.conditioner = std::make_shared<HiDreamO1::HiDreamO1Conditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                                    tensor_storage_map,
-                                                                                   weight_manager);
+                                                                                   weight_manager,
+                                                                                   tokenizers);
             result.diffusion   = std::make_shared<HiDreamO1::HiDreamO1Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                             tensor_storage_map,
                                                                             "model",
@@ -316,7 +351,8 @@ namespace sd::model_builders {
         } else if (sd_version_is_anima(version)) {
             result.conditioner = std::make_shared<AnimaConditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                     tensor_storage_map,
-                                                                    weight_manager);
+                                                                    weight_manager,
+                                                                    tokenizers);
             result.diffusion   = std::make_shared<Anima::AnimaRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                     tensor_storage_map,
                                                                     "model.diffusion_model",
@@ -327,7 +363,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<ZImage::ZImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                       tensor_storage_map,
                                                                       "model.diffusion_model",
@@ -339,7 +376,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                true,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Boogu::BooguImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                          tensor_storage_map,
                                                                          "model.diffusion_model",
@@ -351,7 +389,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<ErnieImage::ErnieImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                               tensor_storage_map,
                                                                               "model.diffusion_model",
@@ -362,7 +401,8 @@ namespace sd::model_builders {
                                                                version,
                                                                "",
                                                                false,
-                                                               weight_manager);
+                                                               weight_manager,
+                                                               tokenizers);
             result.diffusion   = std::make_shared<Lens::LensRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                   tensor_storage_map,
                                                                   "model.diffusion_model",
@@ -376,7 +416,8 @@ namespace sd::model_builders {
                                                                                      tensor_storage_map,
                                                                                      embbeding_map,
                                                                                      version,
-                                                                                     weight_manager);
+                                                                                     weight_manager,
+                                                                                     tokenizers);
             result.diffusion   = std::make_shared<UNetModelRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                  tensor_storage_map,
                                                                  "model.diffusion_model",
@@ -418,8 +459,12 @@ namespace sd::model_builders {
         if (result.ip_adapter) {
             result.ip_adapter->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
         }
+        tokenizers.validate_usage();
         runners = std::move(result);
         return true;
+    } catch (const std::exception& error) {
+        LOG_ERROR("failed to build model runners: %s", error.what());
+        return false;
     }
 
     bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners) {
@@ -547,10 +592,12 @@ namespace sd::model_builders {
         }
 
         if (sd_ctx_params->vae_conv_direct) {
-            LOG_INFO("Using Conv2d direct in the vae model");
+            LOG_INFO("Using Conv2d/Conv3d direct in the vae model");
             result.vae->set_conv2d_direct_enabled(true);
+            result.vae->set_conv3d_direct_enabled(true);
             if (result.preview) {
                 result.preview->set_conv2d_direct_enabled(true);
+                result.preview->set_conv3d_direct_enabled(true);
             }
         }
         if (result.vae) {

src/pipeline/model_builders.h

@@ -15,6 +15,9 @@ struct DiffusionModelRunner;
 struct VAE;
 struct AudioVAERunner;
 struct ControlNet;
+namespace Wav2Vec2 {
+    class Wav2Vec2ModelRunner;
+}
 struct GenerationExtension;
 struct GenerationExtensionInitContext;
 namespace IPAdapter {
@@ -37,6 +40,7 @@ namespace sd::model_builders {
         std::shared_ptr<DiffusionModelRunner> high_noise_diffusion;
         std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision;
         std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
+        std::shared_ptr<Wav2Vec2::Wav2Vec2ModelRunner> audio_encoder;
     };
 
     struct VAEOptions {

src/pipeline/request.cpp

@@ -65,7 +65,7 @@ namespace sd::pipeline {
             return LCM_SCHEDULER;
         } else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) {
             return SIMPLE_SCHEDULER;
-        } else if (sd != nullptr && sd_version_is_flux(sd->version)) {
+        } else if (sd != nullptr && (sd_version_is_flux(sd->version) || sd->version == VERSION_QWEN_IMAGE_2_1)) {
             return FLUX_SCHEDULER;
         } else if (sd != nullptr && sd_version_is_flux2(sd->version)) {
             return FLUX2_SCHEDULER;
@@ -161,7 +161,10 @@ namespace sd::pipeline {
         frames           = sd->align_video_frames(requested_frames);
         clip_skip        = sd_vid_gen_params->clip_skip;
         fps              = std::max(1, sd_vid_gen_params->fps);
-        if (sd_version_is_minimax_h3(sd->version) && fps != 24) {
+        if (sd->version == VERSION_WAN2_2_S2V && sd_vid_gen_params->fps != 16) {
+            LOG_WARN("Wan2.2 S2V uses 16 fps; overriding requested fps %d", sd_vid_gen_params->fps);
+            fps = 16;
+        } else if (sd_version_is_minimax_h3(sd->version) && fps != 24) {
             LOG_WARN("MiniMax-H3 uses 24 fps; overriding requested fps %d", fps);
             fps = 24;
         }

src/pipeline/video.cpp

@@ -5,6 +5,7 @@
 #include <cstdlib>
 #include <optional>
 
+#include "conditioning/wan_audio.h"
 #include "core/rng.hpp"
 #include "core/rng_philox.hpp"
 #include "diffusion_engine.h"
@@ -420,6 +421,45 @@ namespace sd::pipeline {
         return audio;
     }
 
+    // Build the first 16 fps audio window, zero-padding past the track end.
+    static sd::Tensor<float> build_s2v_audio_window(const sd::Tensor<float>& stacked, int64_t batch_frames) {
+        const int64_t embed_dim  = stacked.shape()[0];
+        const int64_t in_frames  = stacked.shape()[1];
+        const int64_t num_layers = stacked.shape()[2];
+        if (embed_dim <= 0 || in_frames <= 0 || num_layers <= 0 || batch_frames <= 0) {
+            return {};
+        }
+        std::vector<float> layer_first(static_cast<size_t>(num_layers) * in_frames * embed_dim);
+        for (int64_t l = 0; l < num_layers; ++l) {
+            for (int64_t f = 0; f < in_frames; ++f) {
+                const float* src = stacked.data() + l * embed_dim * in_frames + f * embed_dim;
+                std::copy_n(src,
+                            static_cast<size_t>(embed_dim),
+                            layer_first.data() + (static_cast<size_t>(l) * in_frames + f) * embed_dim);
+            }
+        }
+        sd::wan_audio::BucketPlan plan;
+        std::vector<float> buckets = sd::wan_audio::build_audio_buckets(layer_first.data(),
+                                                                        static_cast<int>(num_layers),
+                                                                        static_cast<int>(in_frames),
+                                                                        static_cast<int>(embed_dim),
+                                                                        static_cast<int>(batch_frames),
+                                                                        &plan);
+        if (buckets.empty() || plan.bucket_frames < batch_frames) {
+            return {};
+        }
+        // Reorder frame-major buckets into sd::Tensor's [dim, frame, layer] layout.
+        sd::Tensor<float> window({embed_dim, batch_frames, num_layers});
+        for (int64_t f = 0; f < batch_frames; ++f) {
+            for (int64_t l = 0; l < num_layers; ++l) {
+                const float* src = buckets.data() + (static_cast<size_t>(f) * num_layers + l) * embed_dim;
+                float* dst       = window.data() + l * embed_dim * batch_frames + f * embed_dim;
+                std::copy_n(src, static_cast<size_t>(embed_dim), dst);
+            }
+        }
+        return window;
+    }
+
     static std::optional<ImageGenerationLatents> prepare_video_generation_latents(StableDiffusionGGML* sd,
                                                                                   const sd_vid_gen_params_t* sd_vid_gen_params,
                                                                                   GenerationRequest* request) {
@@ -1033,6 +1073,53 @@ namespace sd::pipeline {
             latents.vace_context = sd::ops::concat(vace_context, mask_context, 3);  // [b, 2*c + vae_scale_factor*vae_scale_factor, t + 1 or t, h/vae_scale_factor, w/vae_scale_factor]
             int64_t t2           = ggml_time_ms();
             LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
+        } else if (sd->diffusion_model->get_desc() == "Wan2.2-S2V-14B") {
+            LOG_INFO("S2V");
+            if (!end_image.empty()) {
+                LOG_WARN("Wan2.2 S2V ignores end_image");
+            }
+            if (sd_vid_gen_params->ref_audios_count > 1) {
+                LOG_ERROR("Wan2.2 S2V supports a single driving audio track");
+                return std::nullopt;
+            }
+            int64_t t1 = ggml_time_ms();
+            if (!start_image.empty()) {
+                auto ref_img     = start_image.reshape({start_image.shape()[0],
+                                                        start_image.shape()[1],
+                                                        1,
+                                                        start_image.shape()[2],
+                                                        1});
+                auto encoded_ref = sd->encode_first_stage(ref_img);
+                if (encoded_ref.empty()) {
+                    LOG_ERROR("failed to encode S2V reference image");
+                    return std::nullopt;
+                }
+                // Wan consumes reference latents in 4D.
+                latents.ref_latents.push_back(encoded_ref.reshape({encoded_ref.shape()[0],
+                                                                   encoded_ref.shape()[1],
+                                                                   encoded_ref.shape()[2],
+                                                                   encoded_ref.shape()[3]}));
+            }
+            if (sd_vid_gen_params->ref_audios_count == 1) {
+                if (sd->audio_encoder == nullptr) {
+                    LOG_ERROR("S2V audio conditioning requires --audio-encoder (wav2vec2)");
+                    return std::nullopt;
+                }
+                auto stacked = sd->get_audio_embedding(sd_vid_gen_params->ref_audios[0]);
+                if (stacked.empty()) {
+                    LOG_ERROR("failed to compute wav2vec2 embedding for driving audio");
+                    return std::nullopt;
+                }
+                int64_t latent_t        = sd->video_frames_to_latent_frames(request->frames);
+                int64_t batch_frames    = latent_t * 4;
+                latents.s2v_audio_embed = build_s2v_audio_window(stacked, batch_frames);
+                if (latents.s2v_audio_embed.empty()) {
+                    LOG_ERROR("failed to build S2V audio window");
+                    return std::nullopt;
+                }
+            }
+            int64_t t2 = ggml_time_ms();
+            LOG_INFO("s2v conditioning prepared, taking %" PRId64 " ms", t2 - t1);
         }
 
         if (latents.init_latent.empty()) {
@@ -1052,10 +1139,10 @@ namespace sd::pipeline {
         return latents;
     }
 
-    static ImageGenerationEmbeds prepare_video_generation_embeds(StableDiffusionGGML* sd,
-                                                                 const sd_vid_gen_params_t* sd_vid_gen_params,
-                                                                 const GenerationRequest& request,
-                                                                 const ImageGenerationLatents& latents) {
+    static std::optional<ImageGenerationEmbeds> prepare_video_generation_embeds(StableDiffusionGGML* sd,
+                                                                                const sd_vid_gen_params_t* sd_vid_gen_params,
+                                                                                const GenerationRequest& request,
+                                                                                const ImageGenerationLatents& latents) {
         ConditionerRunnerEndOnExit conditioner_runner_end{sd->cond_stage_model.get()};
 
         ImageGenerationEmbeds embeds;
@@ -1072,8 +1159,12 @@ namespace sd::pipeline {
         int64_t prepare_start_ms = ggml_time_ms();
         embeds.cond              = sd->cond_stage_model->get_learned_condition(sd->n_threads,
                                                                                condition_params);
-        embeds.cond.c_concat     = latents.concat_latent;
-        embeds.cond.c_vector     = latents.clip_vision_output;
+        if (embeds.cond.empty()) {
+            LOG_ERROR("failed to encode video prompt");
+            return std::nullopt;
+        }
+        embeds.cond.c_concat = latents.concat_latent;
+        embeds.cond.c_vector = latents.clip_vision_output;
         if (sd_version_is_minimax_h3(sd->version)) {
             embeds.cond.c_ref_images       = latents.ref_latents;
             embeds.cond.c_ref_audios       = latents.reference_audio_latents;
@@ -1084,10 +1175,20 @@ namespace sd::pipeline {
                     latents.keyframe_indices);
             }
         }
+        if (sd->version == VERSION_WAN2_2_S2V) {
+            embeds.cond.c_ref_images = latents.ref_latents;
+            if (!latents.s2v_audio_embed.empty()) {
+                embeds.cond.c_ref_audios = {latents.s2v_audio_embed};
+            }
+        }
         if (request.use_uncond) {
-            condition_params.text  = request.negative_prompt;
-            embeds.uncond          = sd->cond_stage_model->get_learned_condition(sd->n_threads,
-                                                                                 condition_params);
+            condition_params.text = request.negative_prompt;
+            embeds.uncond         = sd->cond_stage_model->get_learned_condition(sd->n_threads,
+                                                                                condition_params);
+            if (embeds.uncond.empty()) {
+                LOG_ERROR("failed to encode negative video prompt");
+                return std::nullopt;
+            }
             embeds.uncond.c_concat = latents.concat_latent;
             embeds.uncond.c_vector = latents.clip_vision_output;
             if (sd_version_is_minimax_h3(sd->version)) {
@@ -1096,6 +1197,12 @@ namespace sd::pipeline {
                 embeds.uncond.c_reference_blocks = latents.minimax_reference_blocks;
                 embeds.uncond.c_position_ids     = embeds.cond.c_position_ids;
             }
+            if (sd->version == VERSION_WAN2_2_S2V) {
+                embeds.uncond.c_ref_images = latents.ref_latents;
+                if (!latents.s2v_audio_embed.empty()) {
+                    embeds.uncond.c_ref_audios = {sd::Tensor<float>::zeros_like(latents.s2v_audio_embed)};
+                }
+            }
         }
 
         int64_t t1 = ggml_time_ms();
@@ -1422,10 +1529,14 @@ namespace sd::pipeline {
                         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) {
         if (sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd->version)) {
             LOG_INFO("AnimateDiff dispatch: %d frames, %dx%d",
                      sd_vid_gen_params->video_frames, sd_vid_gen_params->width, sd_vid_gen_params->height);
+            if (fps_out != nullptr) {
+                *fps_out = std::max(1, sd_vid_gen_params->fps);
+            }
             return generate_animatediff_video(sd, sd_vid_gen_params, frames_out, num_frames_out);
         }
 
@@ -1437,6 +1548,9 @@ namespace sd::pipeline {
         sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params;
         sd->apply_circular_axes(sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y);
         GenerationRequest request(sd, sd_vid_gen_params);
+        if (fps_out != nullptr) {
+            *fps_out = request.fps;
+        }
         bool latent_upscale_enabled     = request.hires.enabled;
         GenerationRequest hires_request = request;
         if (latent_upscale_enabled) {
@@ -1468,10 +1582,14 @@ namespace sd::pipeline {
         }
         ImageGenerationLatents latents = std::move(*latent_inputs_opt);
 
-        ImageGenerationEmbeds embeds = prepare_video_generation_embeds(sd,
-                                                                       sd_vid_gen_params,
-                                                                       request,
-                                                                       latents);
+        auto embeds_opt = prepare_video_generation_embeds(sd,
+                                                          sd_vid_gen_params,
+                                                          request,
+                                                          latents);
+        if (!embeds_opt) {
+            return false;
+        }
+        ImageGenerationEmbeds embeds = std::move(*embeds_opt);
         if (latent_upscale_enabled) {
             LOG_INFO("generate_video %dx%dx%d -> LTX latent spatial upscale",
                      request.width,
@@ -1725,6 +1843,33 @@ namespace sd::pipeline {
         LOG_INFO("generating latent video completed, taking %.2fs", (latent_end - latent_start) * 1.0f / 1000);
 
         sd_audio_t* generated_audio = nullptr;
+        if (sd->version == VERSION_WAN2_2_S2V && sd_vid_gen_params->ref_audios_count > 0) {
+            // Return the driving track for muxing with the generated video.
+            const sd_audio_t& driving = sd_vid_gen_params->ref_audios[0];
+            generated_audio           = (sd_audio_t*)malloc(sizeof(sd_audio_t));
+            if (generated_audio != nullptr) {
+                generated_audio->sample_rate  = driving.sample_rate;
+                generated_audio->channels     = driving.channels;
+                generated_audio->sample_count = driving.sample_count;
+                generated_audio->data         = (float*)malloc(sizeof(float) * driving.sample_count * driving.channels);
+                if (generated_audio->data == nullptr) {
+                    free(generated_audio);
+                    generated_audio = nullptr;
+                } else {
+                    memcpy(generated_audio->data,
+                           driving.data,
+                           sizeof(float) * driving.sample_count * driving.channels);
+                }
+            }
+            if (generated_audio != nullptr) {
+                LOG_DEBUG("s2v output audio: %u Hz, %u channels, %llu samples",
+                          generated_audio->sample_rate,
+                          generated_audio->channels,
+                          (unsigned long long)generated_audio->sample_count);
+            } else {
+                LOG_DEBUG("s2v output audio copy failed (out of memory)");
+            }
+        }
         if ((sd_version_is_ltxav(sd->version) || sd_version_is_minimax_h3(sd->version)) &&
             latents.audio_length > 0 &&
             sd->audio_vae_model != nullptr) {
@@ -1774,6 +1919,7 @@ namespace sd::pipeline {
             return false;
         }
         auto result = decode_video_outputs(sd, latent_upscale_enabled ? hires_request : request, final_latent, num_frames_out);
+        LOG_DEBUG("decode_video_outputs returned %s", result == nullptr ? "nullptr (failed)" : "frames");
         if (result == nullptr) {
             free_sd_audio(generated_audio);
             return false;
@@ -1786,6 +1932,20 @@ namespace sd::pipeline {
         if (frames_out != nullptr) {
             *frames_out = result;
         }
+        if (sd->version == VERSION_WAN2_2_S2V && generated_audio != nullptr) {
+            // Limit the driving track to the generated video's duration.
+            int fps               = request.fps;
+            uint64_t video_frames = num_frames_out != nullptr ? (uint64_t)*num_frames_out : 0;
+            uint64_t want_samples = (uint64_t)((double)video_frames / fps * generated_audio->sample_rate);
+            LOG_DEBUG("s2v audio truncate: %llu samples -> %llu (video %llu frames @ %d fps)",
+                      (unsigned long long)generated_audio->sample_count,
+                      (unsigned long long)want_samples,
+                      (unsigned long long)video_frames,
+                      fps);
+            if (want_samples > 0 && want_samples < generated_audio->sample_count) {
+                generated_audio->sample_count = want_samples;
+            }
+        }
         if (audio_out != nullptr) {
             *audio_out = generated_audio;
         } else {

src/runtime/audio_processing.cpp

@@ -0,0 +1,97 @@
+#include "audio_processing.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+#include <numeric>
+
+namespace sd::audio {
+
+    // Match torchaudio's Hann-windowed sinc resampler.
+    std::vector<float> resample_audio(const float* samples,
+                                      uint64_t sample_count,
+                                      uint32_t orig_sample_rate,
+                                      uint32_t target_sample_rate) {
+        if (samples == nullptr || sample_count == 0 || orig_sample_rate == 0 || target_sample_rate == 0) {
+            return {};
+        }
+        if (orig_sample_rate == target_sample_rate) {
+            return std::vector<float>(samples, samples + sample_count);
+        }
+
+        constexpr int kLowpassFilterWidth = 6;
+        constexpr double kRolloff         = 0.99;
+        constexpr double kPi              = 3.14159265358979323846;
+
+        const uint64_t gcd        = std::gcd(static_cast<uint64_t>(orig_sample_rate),
+                                             static_cast<uint64_t>(target_sample_rate));
+        const int64_t orig_freq   = static_cast<int64_t>(orig_sample_rate / gcd);
+        const int64_t new_freq    = static_cast<int64_t>(target_sample_rate / gcd);
+        const double base_freq    = static_cast<double>(std::min(orig_freq, new_freq)) * kRolloff;
+        const int64_t width       = static_cast<int64_t>(std::ceil(kLowpassFilterWidth * orig_freq / base_freq));
+        const int64_t kernel_size = 2 * width + orig_freq;
+
+        std::vector<double> kernel(static_cast<size_t>(new_freq) * kernel_size);
+        for (int64_t j = 0; j < new_freq; ++j) {
+            for (int64_t i = 0; i < kernel_size; ++i) {
+                double t = -static_cast<double>(j) / new_freq + static_cast<double>(i - width) / orig_freq;
+                t *= base_freq;
+                t                           = std::clamp(t, -static_cast<double>(kLowpassFilterWidth), static_cast<double>(kLowpassFilterWidth));
+                const double cos_arg        = std::cos(t * kPi / kLowpassFilterWidth / 2);
+                const double window         = cos_arg * cos_arg;
+                double s                    = t * kPi;
+                const double sinc           = (s == 0.0) ? 1.0 : std::sin(s) / s;
+                kernel[j * kernel_size + i] = sinc * window * (base_freq / orig_freq);
+            }
+        }
+
+        const uint64_t num_phases    = static_cast<uint64_t>(sample_count / orig_freq) + 1;
+        const uint64_t target_length = (static_cast<uint64_t>(new_freq) * sample_count +
+                                        static_cast<uint64_t>(orig_freq) - 1) /
+                                       static_cast<uint64_t>(orig_freq);
+        std::vector<float> out(target_length);
+        for (uint64_t phase = 0; phase < num_phases; ++phase) {
+            const int64_t src_base = static_cast<int64_t>(phase * orig_freq) - width;
+            for (int64_t j = 0; j < new_freq; ++j) {
+                const uint64_t out_index = phase * new_freq + j;
+                if (out_index >= target_length) {
+                    break;
+                }
+                const double* k = &kernel[j * kernel_size];
+                double acc      = 0.0;
+                for (int64_t i = 0; i < kernel_size; ++i) {
+                    const int64_t src = src_base + i;
+                    if (src >= 0 && src < static_cast<int64_t>(sample_count)) {
+                        acc += samples[src] * k[i];
+                    }
+                }
+                out[out_index] = static_cast<float>(acc);
+            }
+        }
+        return out;
+    }
+
+    std::vector<float> downmix_to_mono(const float* interleaved_samples,
+                                       uint64_t sample_count,
+                                       uint32_t channels) {
+        std::vector<float> mono;
+        if (interleaved_samples == nullptr || sample_count == 0 || channels == 0) {
+            return mono;
+        }
+        mono.resize(static_cast<size_t>(sample_count));
+        if (channels == 1) {
+            std::memcpy(mono.data(), interleaved_samples, static_cast<size_t>(sample_count) * sizeof(float));
+            return mono;
+        }
+        const float scale = 1.0f / static_cast<float>(channels);
+        for (uint64_t i = 0; i < sample_count; ++i) {
+            float sum = 0.0f;
+            for (uint32_t c = 0; c < channels; ++c) {
+                sum += interleaved_samples[i * channels + c];
+            }
+            mono[static_cast<size_t>(i)] = sum * scale;
+        }
+        return mono;
+    }
+
+}  // namespace sd::audio

src/runtime/audio_processing.h

@@ -0,0 +1,22 @@
+#ifndef __SD_RUNTIME_AUDIO_PROCESSING_H__
+#define __SD_RUNTIME_AUDIO_PROCESSING_H__
+
+#include <cstdint>
+#include <vector>
+
+namespace sd::audio {
+
+    // Returns the input unchanged when sample rates are equal, and an empty vector on invalid input.
+    std::vector<float> resample_audio(const float* samples,
+                                      uint64_t sample_count,
+                                      uint32_t orig_sample_rate,
+                                      uint32_t target_sample_rate);
+
+    // Average interleaved channels; return an empty vector on invalid input.
+    std::vector<float> downmix_to_mono(const float* interleaved_samples,
+                                       uint64_t sample_count,
+                                       uint32_t channels);
+
+}  // namespace sd::audio
+
+#endif  // __SD_RUNTIME_AUDIO_PROCESSING_H__

src/runtime/denoiser.hpp

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

src/stable-diffusion.cpp

@@ -26,13 +26,26 @@ static float get_cache_reuse_threshold(const sd_cache_params_t& params) {
 }
 
 const char* sd_type_name(enum sd_type_t type) {
-    if ((int)type < std::min<int>(SD_TYPE_COUNT, GGML_TYPE_COUNT)) {
-        return ggml_type_name((ggml_type)type);
+    if (type == SD_TYPE_F8_E4M3) {
+        return "f8_e4m3";
+    }
+    if (type == SD_TYPE_F8_E5M2) {
+        return "f8_e5m2";
+    }
+    const auto ggml_type = sd_type_to_ggml_type(type);
+    if (ggml_type != GGML_TYPE_COUNT) {
+        return ggml_type_name(ggml_type);
     }
     return NONE_STR;
 }
 
 enum sd_type_t str_to_sd_type(const char* str) {
+    if (!strcmp(str, "f8_e4m3")) {
+        return SD_TYPE_F8_E4M3;
+    }
+    if (!strcmp(str, "f8_e5m2")) {
+        return SD_TYPE_F8_E5M2;
+    }
     for (int i = 0; i < std::min<int>(SD_TYPE_COUNT, GGML_TYPE_COUNT); i++) {
         auto trait = ggml_get_type_traits((ggml_type)i);
         if (!strcmp(str, trait->type_name)) {
@@ -349,12 +362,14 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "t5xxl_path: %s\n"
              "llm_path: %s\n"
              "llm_vision_path: %s\n"
+             "tokenizer: %s\n"
              "diffusion_model_path: %s\n"
              "high_noise_diffusion_model_path: %s\n"
              "uncond_diffusion_model_path: %s\n"
              "embeddings_connectors_path: %s\n"
              "vae_path: %s\n"
              "audio_vae_path: %s\n"
+             "audio_encoder_path: %s\n"
              "taesd_path: %s\n"
              "control_net_path: %s\n"
              "photo_maker_path: %s\n"
@@ -386,12 +401,14 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              SAFE_STR(sd_ctx_params->t5xxl_path),
              SAFE_STR(sd_ctx_params->llm_path),
              SAFE_STR(sd_ctx_params->llm_vision_path),
+             SAFE_STR(sd_ctx_params->tokenizer),
              SAFE_STR(sd_ctx_params->diffusion_model_path),
              SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path),
              SAFE_STR(sd_ctx_params->uncond_diffusion_model_path),
              SAFE_STR(sd_ctx_params->embeddings_connectors_path),
              SAFE_STR(sd_ctx_params->vae_path),
              SAFE_STR(sd_ctx_params->audio_vae_path),
+             SAFE_STR(sd_ctx_params->audio_encoder_path),
              SAFE_STR(sd_ctx_params->taesd_path),
              SAFE_STR(sd_ctx_params->control_net_path),
              SAFE_STR(sd_ctx_params->photo_maker_path),
@@ -736,8 +753,12 @@ 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) {
     if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) {
+        if (fps_out != nullptr) {
+            *fps_out = 0;
+        }
         return false;
     }
 
@@ -753,10 +774,13 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
 
     StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd);
     if (!execution.ready) {
+        if (fps_out != nullptr) {
+            *fps_out = 0;
+        }
         return false;
     }
 
-    return sd::pipeline::generate_video(sd_ctx->sd, sd_vid_gen_params, frames_out, num_frames_out, audio_out);
+    return sd::pipeline::generate_video(sd_ctx->sd, sd_vid_gen_params, frames_out, num_frames_out, audio_out, fps_out);
 }
 
 SD_API void free_sd_images(sd_image_t* result_images, int num_images) {

src/tokenizers/bpe_tokenizer.cpp

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

src/tokenizers/bpe_tokenizer.h

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

src/tokenizers/clip_tokenizer.cpp

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

src/tokenizers/clip_tokenizer.h

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

src/tokenizers/gemma_tokenizer.cpp

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

src/tokenizers/gemma_tokenizer.h

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

src/tokenizers/gpt_oss_tokenizer.cpp

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

src/tokenizers/gpt_oss_tokenizer.h

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

src/tokenizers/hf_tokenizer.cpp

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

src/tokenizers/hf_tokenizer.h

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

src/tokenizers/mistral_tokenizer.cpp

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