feat: support external Hugging Face tokenizer JSON files (#1973)

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

.github/ISSUE_TEMPLATE/bug_report.yml

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

.github/ISSUE_TEMPLATE/config.yml

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

CMakeLists.txt

@@ -229,6 +229,8 @@ file(GLOB SD_LIB_SOURCES CONFIGURE_DEPENDS
     "src/model/*/*.h"
     "src/model/*/*.cpp"
     "src/model/*/*.hpp"
+    "src/pipeline/*.h"
+    "src/pipeline/*.cpp"
     "src/runtime/*.h"
     "src/runtime/*.cpp"
     "src/runtime/*.hpp"
@@ -292,6 +294,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)
@@ -338,6 +342,7 @@ add_subdirectory(thirdparty)
 
 target_sources(${SD_LIB} PRIVATE $<TARGET_OBJECTS:zip>)
 target_link_libraries(${SD_LIB} PUBLIC ggml)
+target_link_libraries(${SD_LIB} PRIVATE onig sd-utf8proc)
 target_include_directories(${SD_LIB} PUBLIC . src include)
 target_include_directories(${SD_LIB} PRIVATE src/core)
 target_include_directories(${SD_LIB} PUBLIC . thirdparty)

README.md

@@ -51,6 +51,7 @@ API and command-line option may change frequently.***
     - [LongCat Image](./docs/longcat_image.md)
     - [Z-Image](./docs/z_image.md)
     - [MiniT2I](./docs/minit2i.md)
+    - [SenseNova U1.5](./docs/sensenova_u1.md)
     - [Ovis-Image](./docs/ovis_image.md)
     - [Anima](./docs/anima.md)
     - [ERNIE-Image](./docs/ernie_image.md)
@@ -147,6 +148,7 @@ For runtime and parameter backend placement, see the [backend selection guide](.
 
 ## More Guides
 
+- [Troubleshooting](./docs/troubleshooting.md)
 - [Backend selection](./docs/backend.md)
 - [RPC](./docs/rpc.md)
 - [LoRA](./docs/lora.md)

docs/backend.md

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

docs/performance.md

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

docs/rpc.md

@@ -57,7 +57,7 @@ The RPC server acts as the worker. You must explicitly enable the **backend** (t
 
 To find the correct flags for your system, refer to the official documentation for the [`llama.cpp`](https://github.com/ggml-org/llama.cpp/blob/master/docs/build.md) repository.
 
-> **Crucial:** You must include the compiler flags required to satisfy the API compatibility with `stable-diffusion.cpp` (`-DGGML_MAX_NAME=128`). Without this flag, `GGML_MAX_NAME` will default to `64` for the server, and data transfers between the client and server will fail. Of course, `-DGGML_RPC` must also be enabled.
+> **Crucial:** You must include the compiler flags required to satisfy the API compatibility with `stable-diffusion.cpp` (`-DGGML_MAX_NAME=160`). Without this flag, `GGML_MAX_NAME` will default to `64` for the server, and data transfers between the client and server will fail. Of course, `-DGGML_RPC` must also be enabled.
 >
 > I recommend disabling the `LLAMA_CURL` flag to avoid unnecessary dependencies, and disabling shared library builds to avoid potential conflicts.
 
@@ -72,8 +72,8 @@ cmake .. -DGGML_RPC=ON \
     -DGGML_VULKAN=ON \        # Ensure backend is enabled
     -DGGML_BUILD_SHARED_LIBS=OFF \
     -DLLAMA_CURL=OFF \
-    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 \
-    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128
+    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=160 \
+    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=160
 cmake --build . --config Release --target rpc-server -j $(nproc)
 ```
 
@@ -86,8 +86,8 @@ cmake .. -DGGML_RPC=ON \
     -DGGML_METAL=ON \
     -DGGML_BUILD_SHARED_LIBS=OFF \
     -DLLAMA_CURL=OFF \
-    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 \
-    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128
+    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=160 \
+    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=160
 cmake --build . --config Release --target rpc-server
 ```
 
@@ -101,8 +101,8 @@ cmake .. -G "Visual Studio 17 2022" -A x64 `
     -DGGML_VULKAN=ON `
     -DGGML_BUILD_SHARED_LIBS=OFF `
     -DLLAMA_CURL=OFF `
-    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 `
-    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128
+    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=160 `
+    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=160
 cmake --build . --config Release --target rpc-server
 ```
 

docs/sensenova_u1.md

@@ -0,0 +1,46 @@
+# How to Use
+
+SenseNova U1.5 is an 8B MoT model that performs diffusion directly in RGB pixel
+space. It does not require a separate text encoder or VAE.
+
+## Download weights
+
+- Download SenseNova U1.5 8B MoT
+    - safetensors: https://huggingface.co/sensenova/SenseNova-U1.5-8B-MoT
+
+Pass the complete downloaded repository directory to `--model`. The directory
+must contain `model.safetensors.index.json`, every referenced Safetensors shard,
+and the tokenizer files.
+
+## Examples
+
+### CUDA
+
+```bash
+./bin/sd-cli \
+  --model /path/to/SenseNova-U1.5-8B-MoT \
+  --prompt "a red cube on a white background" \
+  --width 2048 \
+  --height 2048 \
+  --steps 50 \
+  --cfg-scale 4 \
+  --flow-shift 3 \
+  --seed 42 \
+  --sampling-method euler \
+  --rng cuda \
+  --fa \
+  --output output.png
+```
+
+## Notes
+
+- To match the official non-thinking text-to-image pipeline, use 50 Euler
+  steps, CFG 4, flow shift 3, seed 42, CUDA RNG, and an empty negative prompt.
+- Width and height must be multiples of 32. The trained 1:1 resolution is
+  2048x2048; lower resolutions are useful for smoke tests but are outside the
+  training buckets.
+- The SenseNova prompt template and unconditional prompt are built
+  automatically.
+- This implementation supports non-thinking text-to-image generation. Image
+  editing, visual understanding, interleaved generation, and thinking-mode
+  prompt expansion are not implemented.

docs/tokenizers.md

@@ -0,0 +1,97 @@
+# JSON tokenizers
+
+Use a Hugging Face `tokenizer.json` to supply the tokenizer vocabulary, merges,
+added tokens, and processing stages. Without this option, sd.cpp keeps its
+embedded tokenizer for the selected model.
+
+```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.
+
+| 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 the embedded tokenizers. The CLI passes the string through;
+`TokenizerConfig` parses and validates it when text encoders are initialized.
+
+```c
+sd_ctx_params_t params;
+sd_ctx_params_init(&params);
+params.tokenizer = "clip-l=tokenizer_clip.json,clip-g=tokenizer_clip.json";
+```
+
+Rebuild applications against the updated public header when using the updated
+library.
+
+## Supported components
+
+| Stage | Supported configurations |
+| --- | --- |
+| Normalizer | `Sequence`, `NFC`, `Lowercase`, `Replace` with String/Regex patterns |
+| PreTokenizer | `Sequence`, `Split` with String/Regex patterns, all five delimiter behaviors and `invert`; `ByteLevel` with `add_prefix_space` and `use_regex` |
+| Model | Deterministic `BPE`, string or array-pair merges, `unk_token`, `fuse_unk`, `byte_fallback`, `ignore_merges`, `end_of_word_suffix` |
+| PostProcessor | Single-sequence `TemplateProcessing` with at most one prefix and one suffix token, `RobertaProcessing`, `ByteLevel` |
+| Decoder | `Sequence`, `Replace`, `ByteLevel`, `ByteFallback`, `Fuse` |
+| AddedToken | Special and ordinary added tokens, original IDs, raw or normalized matching, leftmost-longest matching |
+
+`ByteLevel.use_regex` defaults to true when omitted. ByteLevel postprocessing
+changes offsets only and adds no tokens. Added tokens with `single_word`,
+`lstrip`, or `rstrip` enabled, nonzero BPE dropout, nonempty
+`continuing_subword_prefix`, and unsupported component types fail loading.
+New added-token IDs must follow the model vocabulary consecutively; configurations
+whose IDs Hugging Face would reassign are rejected.
+JSON `padding` and `truncation` must be null. This API returns IDs, not offsets,
+type IDs, or paired-input encodings; the pair template is not used.
+
+The pipeline covers the CLIP, Gemma 2, Gemma 3, GPT-OSS, Mistral 3, Qwen 2 and
+Qwen 3 JSON configurations used by the differential test. It does not imply
+support for every tokenizer published under those model names.
+
+## Prompt integration
+
+Prompt attention parsing and model-specific chat/image templates remain in the
+conditioner. Raw `encode()` does not add BOS/EOS. The conditioner concatenates
+weighted prompt fragments, then the existing padding/chunking step applies the
+JSON single-sequence template once per sequence or CLIP chunk. Padding ID,
+direction, length limits and attention masks remain text encoder policies.
+CLIP requires both BOS and EOS because its chunking reserves those positions.
+
+The internal `encode()`, `tokenize()`, and `decode()` interfaces return a success
+flag and write to an output parameter. A successful result may be empty; a failed
+call clears its output. JSON tokenizer input, normalization, and regex failures
+return `false` with diagnostic information instead of throwing. Invalid
+JSON, unsupported stages, conflicting IDs and IDs outside the encoder embedding
+table fail initialization instead of falling back to the embedded tokenizer.

docs/troubleshooting.md

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

docs/wan.md

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

examples/cli/README.md

@@ -22,3 +22,6 @@ Metadata mode inspects PNG/JPEG container metadata without loading any model:
 ./bin/sd-cli -M metadata --image ./output.png --metadata-raw
 ./bin/sd-cli -M metadata --image ./output.png --metadata-all
 ```
+
+For completely black or white images or videos, NaNs, and the `--linear-scale` /
+`--attn-scale` workaround, see [Troubleshooting](../../docs/troubleshooting.md).

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

@@ -302,8 +302,12 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
                         invalid_arg = true;
                         return;
                     }
-                    *option.target = std::stoi(argv[i]);
-                    found_arg      = true;
+                    try {
+                        *option.target = std::stoi(argv[i]);
+                    } catch (const std::invalid_argument&) {
+                        invalid_arg = true;
+                    }
+                    found_arg = true;
                 }))
                 break;
 
@@ -312,8 +316,12 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
                         invalid_arg = true;
                         return;
                     }
-                    *option.target = std::stof(argv[i]);
-                    found_arg      = true;
+                    try {
+                        *option.target = std::stof(argv[i]);
+                    } catch (const std::invalid_argument&) {
+                        invalid_arg = true;
+                    }
+                    found_arg = true;
                 }))
                 break;
 
@@ -337,7 +345,8 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
 
         if (invalid_arg) {
             if (!valid) {
-                LOG_ERROR("error: invalid parameter for argument: %s", arg.c_str());
+                LOG_ERROR("error: invalid parameter for argument \"%s\": \"%s\"",
+                          arg.c_str(), (i >= argc) ? "" : argv[i]);
             }
             return false;
         }
@@ -350,6 +359,25 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
     return true;
 }
 
+static int parse_scale_override(int argc, const char** argv, int index, float& scale) {
+    if (++index >= argc) {
+        return -1;
+    }
+    try {
+        size_t end              = 0;
+        const std::string value = argv[index];
+        float parsed            = std::stof(value, &end);
+        if (end != value.size() || !std::isfinite(parsed) || parsed < 0.f ||
+            (parsed > 0.f && !std::isfinite(1.f / parsed))) {
+            return -1;
+        }
+        scale = parsed;
+    } catch (const std::exception&) {
+        return -1;
+    }
+    return 1;
+}
+
 ArgOptions SDContextParams::get_options() {
     ArgOptions options;
     options.string_options = {
@@ -382,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",
+         (int)',',
+         &tokenizer},
         {"",
          "--llm_vision",
          "path to the llm vit",
@@ -432,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)",
@@ -678,11 +716,23 @@ ArgOptions SDContextParams::get_options() {
     };
 
     options.manual_options = {
+        {"",
+         "--linear-scale",
+         "linear input scale override (float, default: 0 = model default, 1 = no scaling)",
+         [this](int argc, const char** argv, int index) {
+             return parse_scale_override(argc, argv, index, linear_scale);
+         }},
+        {"",
+         "--attn-scale",
+         "flash-attention K/V scale override (float, default: 0 = model default, 1 = no scaling); requires --fa or --diffusion-fa",
+         [this](int argc, const char** argv, int index) {
+             return parse_scale_override(argc, argv, index, attn_scale);
+         }},
         {"",
          "--auto-fit",
-         "on|off (default: on). Use one GPU for diffusion/te/vae computation and place weights on that GPU, "
+         "on|off (default: on). Preserve --backend (otherwise select one GPU) and place weights on the compute GPU, "
          "RAM, another GPU, or disk in that order, according to available memory (--max-vram limits GPU budgets). "
-         "Disabled by explicit --backend or --params-backend; uses automatic graph segmentation when needed",
+         "Disabled by explicit --params-backend; uses automatic graph segmentation when needed",
          on_auto_fit_arg},
         {"",
          "--type",
@@ -851,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"
@@ -858,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"
@@ -886,6 +938,8 @@ std::string SDContextParams::to_string() const {
         << "  vae_on_cpu: " << (vae_on_cpu ? "true" : "false") << ",\n"
         << "  flash_attn: " << (flash_attn ? "true" : "false") << ",\n"
         << "  diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n"
+        << "  linear_scale: " << linear_scale << ",\n"
+        << "  attn_scale: " << attn_scale << ",\n"
         << "  diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n"
         << "  vae_conv_direct: " << (vae_conv_direct ? "true" : "false") << ",\n"
         << "  prediction: " << sd_prediction_name(prediction) << ",\n"
@@ -915,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();
@@ -939,6 +995,8 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.enable_mmap                     = enable_mmap;
     sd_ctx_params.flash_attn                      = flash_attn;
     sd_ctx_params.diffusion_flash_attn            = diffusion_flash_attn;
+    sd_ctx_params.linear_scale                    = linear_scale;
+    sd_ctx_params.attn_scale                      = attn_scale;
     sd_ctx_params.tae_preview_only                = taesd_preview;
     sd_ctx_params.diffusion_conv_direct           = diffusion_conv_direct;
     sd_ctx_params.vae_conv_direct                 = vae_conv_direct;
@@ -1471,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;
@@ -1660,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;
@@ -175,6 +177,8 @@ struct SDContextParams {
     lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO;
 
     bool force_sdxl_vae_conv_scale = false;
+    float linear_scale             = 0.f;
+    float attn_scale               = 0.f;
 
     float flow_shift = INFINITY;
     ArgOptions get_options();

examples/common/media_io.cpp

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

examples/common/media_io.h

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

examples/server/README.md

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

examples/server/async_jobs.cpp

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

include/stable-diffusion.h

@@ -92,6 +92,7 @@ enum prediction_t {
     FLUX_FLOW_PRED,
     SEFI_FLOW_PRED,
     MINIT2I_FLOW_PRED,
+    SENSENOVA_U1_FLOW_PRED,
     PREDICTION_COUNT
 };
 
@@ -207,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;
@@ -240,6 +242,9 @@ typedef struct {
     const char* rpc_servers;
     const char* model_args;
     bool disable_segmented_compute;  // Force monolithic graph execution even when automatic graph cutting would fit memory better
+    float linear_scale;              // Override linear input scaling; 0 keeps the model default
+    float attn_scale;                // Override flash-attention K/V scaling; 0 keeps the model default
+    const char* tokenizer;           // Optional tokenizer.json path or main=FILE,clip-l=FILE,clip-g=FILE assignments
 } sd_ctx_params_t;
 
 typedef struct {
@@ -493,6 +498,9 @@ SD_API void free_sd_audio(sd_audio_t* audio);
 SD_API void sd_sample_params_init(sd_sample_params_t* sample_params);
 SD_API char* sd_sample_params_to_str(const sd_sample_params_t* sample_params);
 
+// Requires a loaded context; returns a static string owned by the library, or "Unknown".
+SD_API const char* sd_get_model_version_name(const sd_ctx_t* sd_ctx);
+
 SD_API enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx);
 SD_API enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_method_t sample_method);
 
@@ -515,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;
 

scripts/format-code.ps1

@@ -11,6 +11,8 @@ $patterns = @(
     "src/extensions/*.cpp"
     "src/extensions/*.h"
     "src/extensions/*.hpp"
+    "src/pipeline/*.cpp"
+    "src/pipeline/*.h"
     "src/runtime/*.cpp"
     "src/runtime/*.h"
     "src/runtime/*.hpp"

scripts/format-code.sh

@@ -9,6 +9,7 @@ for f in src/*.cpp src/*.h src/*.hpp \
          src/conditioning/*.cpp src/conditioning/*.h src/conditioning/*.hpp \
          src/core/*.cpp src/core/*.h src/core/*.hpp \
          src/extensions/*.cpp src/extensions/*.h src/extensions/*.hpp \
+         src/pipeline/*.cpp src/pipeline/*.h \
          src/runtime/*.cpp src/runtime/*.h src/runtime/*.hpp \
          src/model/*/*.cpp src/model/*/*.h src/model/*/*.hpp \
          src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \

src/conditioning/conditioner.hpp

@@ -16,6 +16,8 @@
 #include "model/te/llm.hpp"
 #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;
@@ -149,6 +151,7 @@ public:
     virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {}
     virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
     virtual void set_flash_attention_enabled(bool enabled) = 0;
+    virtual void set_scale_overrides(float linear_scale, float attn_scale) {}
     virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
     virtual void runner_end() {}
 };
@@ -157,7 +160,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;
 
@@ -171,12 +174,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)) {
@@ -231,6 +240,13 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         }
     }
 
+    void set_scale_overrides(float linear_scale, float attn_scale) override {
+        text_model->set_scale_overrides(linear_scale, attn_scale);
+        if (sd_version_is_sdxl(version)) {
+            text_model2->set_scale_overrides(linear_scale, attn_scale);
+        }
+    }
+
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
         text_model->set_weight_adapter(adapter);
         if (sd_version_is_sdxl(version)) {
@@ -356,16 +372,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,
@@ -403,18 +418,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] << ", ";
@@ -451,7 +469,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);
                 }
@@ -552,7 +570,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,
@@ -620,8 +641,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;
@@ -629,8 +650,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;
@@ -648,9 +669,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) {
@@ -736,6 +765,18 @@ struct SD3CLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_scale_overrides(float linear_scale, float attn_scale) override {
+        if (clip_l) {
+            clip_l->set_scale_overrides(linear_scale, attn_scale);
+        }
+        if (clip_g) {
+            clip_g->set_scale_overrides(linear_scale, attn_scale);
+        }
+        if (t5) {
+            t5->set_scale_overrides(linear_scale, attn_scale);
+        }
+    }
+
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
         if (clip_l) {
             clip_l->set_weight_adapter(adapter);
@@ -790,27 +831,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);
@@ -881,7 +931,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,
@@ -924,7 +974,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,
@@ -1002,6 +1052,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,
@@ -1010,7 +1063,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;
@@ -1018,7 +1071,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) {
@@ -1035,6 +1089,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.");
@@ -1106,6 +1165,15 @@ struct FluxCLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_scale_overrides(float linear_scale, float attn_scale) override {
+        if (clip_l) {
+            clip_l->set_scale_overrides(linear_scale, attn_scale);
+        }
+        if (t5) {
+            t5->set_scale_overrides(linear_scale, attn_scale);
+        }
+    }
+
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
         if (clip_l) {
             clip_l->set_weight_adapter(adapter);
@@ -1151,19 +1219,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);
@@ -1213,7 +1287,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,
@@ -1224,7 +1298,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});
                 }
@@ -1243,7 +1320,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);
@@ -1270,6 +1350,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,
@@ -1368,6 +1451,12 @@ struct T5CLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_scale_overrides(float linear_scale, float attn_scale) override {
+        if (t5) {
+            t5->set_scale_overrides(linear_scale, attn_scale);
+        }
+    }
+
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
         if (t5) {
             t5->set_weight_adapter(adapter);
@@ -1407,7 +1496,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);
             }
@@ -1505,6 +1597,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,
@@ -1576,6 +1671,12 @@ struct MiniT2IConditioner : public Conditioner {
         }
     }
 
+    void set_scale_overrides(float linear_scale, float attn_scale) override {
+        if (t5) {
+            t5->set_scale_overrides(linear_scale, attn_scale);
+        }
+    }
+
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
         if (t5) {
             t5->set_weight_adapter(adapter);
@@ -1597,7 +1698,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);
         }
@@ -1623,21 +1727,93 @@ struct MiniT2IConditioner : public Conditioner {
     }
 };
 
+struct SenseNovaU1Conditioner : public Conditioner {
+    static constexpr size_t kMaxPromptTokens = 12288;
+    SenseNovaU1Tokenizer tokenizer;
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        SD_UNUSED(tensors);
+    }
+
+    void set_flash_attention_enabled(bool enabled) override {
+        SD_UNUSED(enabled);
+    }
+
+    static std::string build_query(const std::string& text, bool is_negative) {
+        static const std::string kSystemMessage =
+            "You are an image generation and editing assistant that accurately understands and executes user intent.\n\n"
+            "You support two modes:\n\n1. Think Mode:\nIf the task requires reasoning, you MUST start with a "
+            "<think></think> block. Put all reasoning inside the block using plain text. DO NOT include any image tags. "
+            "Keep it reasonable and directly useful for producing the final image.\n\n2. Non-Think Mode:\nIf no reasoning "
+            "is needed, directly produce the final image.\n\nTask Types:\n\nA. Text-to-Image Generation:\n- Generate a "
+            "high-quality image based on the user's description.\n- Ensure visual clarity, semantic consistency, and "
+            "completeness.\n- DO NOT introduce elements that contradict or override the user's intent.\n\nB. Image Editing:\n"
+            "- Use the provided image(s) as input or reference for modification or transformation.\n- The result can be an "
+            "edited image or a new image based on the reference(s).\n- Preserve all unspecified attributes unless explicitly "
+            "changed.\n\nGeneral Rules:\n- For any visible text in the image, follow the language specified for the rendered "
+            "text in the user's description, not the language of the prompt. If no language is specified, use the user's input "
+            "language.";
+
+        std::string query;
+        if (!is_negative) {
+            query += "<|im_start|>system\n";
+            query += kSystemMessage;
+            query += "<|im_end|>\n";
+        }
+        query += "<|im_start|>user\n";
+        query += text;
+        query += "<|im_end|>\n<|im_start|>assistant\n";
+        query += is_negative ? "<img>" : "<think>\n\n</think>\n\n<img>";
+        return query;
+    }
+
+    SDCondition tokenize_condition(const std::string& text, bool 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(),
+                      kMaxPromptTokens);
+            return {};
+        }
+
+        SDCondition result;
+        result.c_input_ids = sd::Tensor<int32_t>({static_cast<int64_t>(tokens.size())}, tokens);
+        return result;
+    }
+
+    SDCondition get_learned_condition(int n_threads,
+                                      const ConditionerParams& conditioner_params) override {
+        SD_UNUSED(n_threads);
+        return tokenize_condition(conditioner_params.text, false);
+    }
+
+    SDCondition get_unconditional_condition(const std::string& text) {
+        return tokenize_condition(text, true);
+    }
+};
+
 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 {
@@ -1672,6 +1848,10 @@ struct AnimaConditioner : public Conditioner {
         llm->set_flash_attention_enabled(enabled);
     }
 
+    void set_scale_overrides(float linear_scale, float attn_scale) override {
+        llm->set_scale_overrides(linear_scale, attn_scale);
+    }
+
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
         llm->set_weight_adapter(adapter);
     }
@@ -1700,7 +1880,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);
@@ -1713,7 +1896,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);
         }
@@ -1732,6 +1918,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,
@@ -1758,7 +1948,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;
 
@@ -1767,7 +1957,8 @@ 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) {
         LLM::LLMArch arch = LLM::LLMArch::QWEN2_5_VL;
         if (version == VERSION_FLUX2) {
@@ -1789,21 +1980,32 @@ 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 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>();
+            }
+        }
         if (sd_version_is_hunyuan_video(version)) {
             const std::string byt5_prefix = "text_encoders.t5xxl.transformer";
             for (const auto& [name, _] : tensor_storage_map) {
@@ -1876,6 +2078,13 @@ struct LLMEmbedder : public Conditioner {
         }
     }
 
+    void set_scale_overrides(float linear_scale, float attn_scale) override {
+        llm->set_scale_overrides(linear_scale, attn_scale);
+        if (byt5) {
+            byt5->set_scale_overrides(linear_scale, attn_scale);
+        }
+    }
+
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
         if (llm) {
             llm->set_weight_adapter(adapter);
@@ -1935,7 +2144,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);
         }
@@ -1968,6 +2180,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()) {
@@ -2127,7 +2343,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);
@@ -2143,6 +2363,7 @@ struct LLMEmbedder : public Conditioner {
                         prompt += placeholder;
                     }
                     prompt += "<|vision_end|>";
+                    return true;
                 };
 
                 const auto* references = conditioner_params.minimax_h3_references;
@@ -2159,11 +2380,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;
                         }
 
@@ -2187,22 +2410,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 {};
+                        }
                     }
                 }
             }
@@ -2238,7 +2465,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) {
@@ -2291,7 +2521,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|>";
@@ -2433,7 +2667,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|>";
@@ -2501,7 +2739,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|>";
@@ -2706,7 +2948,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;
@@ -2725,7 +2970,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);
@@ -2758,6 +3005,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;
@@ -2808,6 +3058,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));
         }
 
@@ -2906,7 +3159,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;
@@ -2933,17 +3186,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,
@@ -2965,6 +3222,11 @@ struct LTXAVEmbedder : public Conditioner {
         projector->set_flash_attention_enabled(enabled);
     }
 
+    void set_scale_overrides(float linear_scale, float attn_scale) override {
+        llm->set_scale_overrides(linear_scale, attn_scale);
+        projector->set_scale_overrides(linear_scale, attn_scale);
+    }
+
     void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
         llm->set_max_graph_vram_bytes(max_vram_bytes);
         projector->set_max_graph_vram_bytes(max_vram_bytes);
@@ -3017,7 +3279,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);
         }
@@ -3035,6 +3300,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()) {
@@ -3133,7 +3402,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/core/backend_fit.cpp

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

src/core/compute_workspace.cpp

@@ -2,12 +2,14 @@
 
 #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"
 
@@ -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

@@ -325,6 +325,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,
@@ -683,17 +753,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

@@ -9,7 +9,7 @@
 
 #define EPS 1e-05f
 
-static_assert(GGML_MAX_NAME >= 128, "GGML_MAX_NAME must be at least 128");
+static_assert(GGML_MAX_NAME >= 160, "GGML_MAX_NAME must be at least 160");
 
 // n-mode tensor-matrix product
 // example: 2-mode product
@@ -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,]
@@ -219,7 +231,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

@@ -87,6 +87,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;
 }
 
@@ -593,7 +597,7 @@ static ggml_backend_t sd_get_default_backend() {
     return backend;
 }
 
-static bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error) {
+bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error) {
     if (assignment == nullptr) {
         return false;
     }
@@ -660,7 +664,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 +972,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 {
@@ -93,6 +94,7 @@ ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
                                                         sd_graph_eval_callback_t callback_eval,
                                                         void* callback_eval_user_data);
 std::string sd_backend_resolve_name(const std::string& name);
+bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error);
 const char* sd_backend_module_name(SDBackendModule module);
 void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);
 bool add_rpc_devices(const std::string& servers);

src/core/ggml_graph_cut.cpp

@@ -482,6 +482,16 @@ namespace sd::ggml_graph_cut {
         return ggml_nbytes(cache_src);
     }
 
+    static bool can_ignore_op_params(ggml_op op) {
+        // Exempt only parameters that cannot affect graph layout or backend allocation size.
+        switch (op) {
+            case GGML_OP_SCALE:
+                return true;
+            default:
+                return false;
+        }
+    }
+
     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;
@@ -530,8 +540,10 @@ namespace sd::ggml_graph_cut {
             for (auto source : tensor->src) {
                 signature.push_back(source == nullptr ? 0 : indices.at(source));
             }
-            for (int value : tensor->op_params) {
-                signature.push_back(static_cast<uint32_t>(value));
+            if (!can_ignore_op_params(tensor->op)) {
+                for (int value : tensor->op_params) {
+                    signature.push_back(static_cast<uint32_t>(value));
+                }
             }
         }
         return signature;

src/core/ggml_runner.cpp

@@ -1,7 +1,9 @@
 #include <algorithm>
+#include <exception>
 #include <map>
 #include <utility>
 
+#include "core/ggml_extend.h"
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_runner.h"
 #include "core/ggml_tensor_utils.h"
@@ -11,6 +13,21 @@
 
 using namespace sd;
 
+ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
+                                    ggml_tensor* q,
+                                    ggml_tensor* k,
+                                    ggml_tensor* v,
+                                    int64_t n_head,
+                                    ggml_tensor* mask,
+                                    bool skip_reshape,
+                                    bool flash_attn,
+                                    float kv_scale) {
+    if (ctx->attn_scale > 0.f) {
+        kv_scale = ctx->attn_scale;
+    }
+    return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale);
+}
+
 void GGMLRunner::alloc_params_ctx() {
     ggml_init_params params;
     params.mem_size   = static_cast<size_t>(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead());
@@ -66,15 +83,9 @@ void GGMLRunner::rebuild_params_tensor_set() {
 }
 
 ggml_tensor* GGMLRunner::canonical_param_tensor(ggml_tensor* tensor) {
-    if (tensor == nullptr) {
-        return nullptr;
-    }
-    if (params_tensor_set_.find(tensor) != params_tensor_set_.end()) {
-        return tensor;
-    }
-    if (tensor->view_src != nullptr &&
-        params_tensor_set_.find(tensor->view_src) != params_tensor_set_.end()) {
-        return tensor->view_src;
+    for (auto* current = tensor; current != nullptr; current = current->view_src) {
+        if (params_tensor_set_.count(current) != 0)
+            return current;
     }
     return nullptr;
 }
@@ -483,9 +494,10 @@ void GGMLRunner::runner_end() {
     if (auto manager = residency_manager.lock()) {
         manager->clear_prefetched_params(reinterpret_cast<uintptr_t>(this));
         std::vector<ggml_tensor*> tensors;
-        for (auto tensor = ggml_get_first_tensor(params_ctx); tensor != nullptr;
-             tensor      = ggml_get_next_tensor(params_ctx, tensor)) {
-            tensors.push_back(tensor);
+        for (auto tensor : params_tensor_set_) {
+            auto* parameter = manager->resolve_param_tensor(const_cast<ggml_tensor*>(tensor));
+            if (parameter != nullptr)
+                tensors.push_back(parameter);
         }
         manager->evict_compute_backend_params(tensors);
         manager->remove_runtime_owner(reinterpret_cast<uintptr_t>(this));
@@ -515,6 +527,8 @@ GGMLRunnerContext GGMLRunner::get_context() {
     runner_ctx.ggml_ctx              = compute_ctx;
     runner_ctx.backend               = runtime_backend;
     runner_ctx.flash_attn_enabled    = flash_attn_enabled;
+    runner_ctx.linear_scale          = linear_scale;
+    runner_ctx.attn_scale            = attn_scale;
     runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled;
     runner_ctx.circular_x_enabled    = circular_x_enabled;
     runner_ctx.circular_y_enabled    = circular_y_enabled;
@@ -620,9 +634,24 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
     if (!prepare_compute_graph(get_graph, &graph)) {
         return std::nullopt;
     }
+    params_tensor_set_dirty_ = true;
     rebuild_params_tensor_set();
-    auto output = execute_graph(graph, n_threads, no_return, read_outputs);
-    success     = output.has_value();
+    if (auto manager = residency_manager.lock()) {
+        for (int i = 0; i < sd::ggml_graph_cut::leaf_count(graph); ++i) {
+            auto* parameter = manager->resolve_param_tensor(sd::ggml_graph_cut::leaf_tensor(graph, i));
+            if (parameter != nullptr)
+                params_tensor_set_.insert(parameter);
+        }
+    }
+    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);
     }
@@ -935,6 +964,9 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
                 }
             }
         }
+        if (!workspace_.segment_end()) {
+            return fail_segment("workspace synchronization");
+        }
         // Final outputs and their callbacks may still be views of consumed cuts.
         cut_cache_.prune(segment.future_cut_names);
     }

src/core/ggml_runner.h

@@ -68,6 +68,8 @@ struct GGMLRunnerContext {
     ggml_backend_t backend                                           = nullptr;
     ggml_context* ggml_ctx                                           = nullptr;
     bool flash_attn_enabled                                          = false;
+    float linear_scale                                               = 0.f;
+    float attn_scale                                                 = 0.f;
     bool conv2d_direct_enabled                                       = false;
     bool circular_x_enabled                                          = false;
     bool circular_y_enabled                                          = false;
@@ -113,6 +115,16 @@ struct GGMLRunnerContext {
     }
 };
 
+ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
+                                    ggml_tensor* q,
+                                    ggml_tensor* k,
+                                    ggml_tensor* v,
+                                    int64_t n_head,
+                                    ggml_tensor* mask = nullptr,
+                                    bool skip_reshape = false,
+                                    bool flash_attn   = false,
+                                    float kv_scale    = 1.f);
+
 struct GGMLRunner {
 private:
     std::map<ggml_backend_t, size_t> logged_compute_bytes_;
@@ -163,6 +175,8 @@ protected:
     const std::string final_result_name = "ggml_runner_final_result_tensor";
 
     bool flash_attn_enabled    = false;
+    float linear_scale         = 0.f;
+    float attn_scale           = 0.f;
     bool conv2d_direct_enabled = false;
     bool circular_x_enabled    = false;
     bool circular_y_enabled    = false;
@@ -323,6 +337,11 @@ public:
         flash_attn_enabled = enabled;
     }
 
+    void set_scale_overrides(float linear_scale, float attn_scale) {
+        this->linear_scale = linear_scale;
+        this->attn_scale   = attn_scale;
+    }
+
     void set_conv2d_direct_enabled(bool enabled) {
         conv2d_direct_enabled = enabled;
     }

src/core/ordered_map.hpp

@@ -38,9 +38,19 @@ public:
             insert(kv);
     }
 
-    OrderedMap(const OrderedMap&)                = default;
-    OrderedMap(OrderedMap&&) noexcept            = default;
-    OrderedMap& operator=(const OrderedMap&)     = default;
+    OrderedMap(const OrderedMap& other) {
+        for (const auto& value : other) {
+            insert(value);
+        }
+    }
+    OrderedMap(OrderedMap&&) noexcept = default;
+    OrderedMap& operator=(const OrderedMap& other) {
+        if (this != &other) {
+            OrderedMap copy(other);
+            swap(copy);
+        }
+        return *this;
+    }
     OrderedMap& operator=(OrderedMap&&) noexcept = default;
 
     // --- element access ---
@@ -174,4 +184,4 @@ public:
     }
 };
 
-#endif  // __SD_CORE_ORDERED_MAP_HPP__
+#endif  // __SD_CORE_ORDERED_MAP_HPP__

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

@@ -821,7 +821,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 +834,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 +897,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/detailer.cpp

@@ -676,7 +676,7 @@ bool ADetailerGGML::load_from_file(const std::string& detector_path) {
     model_manager = std::make_shared<ModelManager>();
     model_manager->set_n_threads(n_threads);
     model_manager->set_enable_mmap(false);
-    ModelLoader& loader = model_manager->loader();
+    ModelLoader loader;
     if (!loader.init_from_file(detector_path)) {
         LOG_ERROR("failed to load ADetailer detector: '%s'", detector_path.c_str());
         return false;
@@ -696,7 +696,8 @@ bool ADetailerGGML::load_from_file(const std::string& detector_path) {
 
     std::map<std::string, ggml_tensor*> tensors;
     detector->get_param_tensors(tensors);
-    if (!model_manager->register_param_tensors("YOLOv8",
+    if (!model_manager->set_loader(loader) ||
+        !model_manager->register_param_tensors(ModelComponent::Detector,
                                                std::move(tensors),
                                                backend_manager.params_backend_is_disk(SDBackendModule::DETECTOR)
                                                    ? ModelManager::ResidencyMode::Disk

src/device_residency_manager.h

@@ -48,9 +48,10 @@ struct DeviceResidencyManager {
                                                const std::vector<ggml_tensor*>& required_params) const = 0;
     virtual bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
                                         ggml_backend_t compute_backend)                                = 0;
-    virtual bool prepare_params(const std::vector<ggml_tensor*>& tensors)                              = 0;
-    virtual void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors)              = 0;
-    virtual void evict_compute_backend_params(const std::vector<ggml_tensor*>& tensors)                = 0;
+    virtual ggml_tensor* resolve_param_tensor(ggml_tensor* tensor) const { return nullptr; }
+    virtual bool prepare_params(const std::vector<ggml_tensor*>& tensors)                 = 0;
+    virtual void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
+    virtual void evict_compute_backend_params(const std::vector<ggml_tensor*>& tensors)   = 0;
     virtual WeightResidencyInfo inspect_compute_backend_params(
         const std::vector<ggml_tensor*>& tensors) const          = 0;
     virtual void update_runtime_residency(uintptr_t owner_id,

src/extensions/generation_extension.h

@@ -19,7 +19,7 @@ struct GenerationExtensionInitContext {
     const sd_ctx_params_t* params;
     SDVersion version;
     const String2TensorStorage& tensor_storage_map;
-    ModelLoader& model_loader;
+    bool photomaker_source_available;
     std::shared_ptr<ModelManager> model_manager;
     int n_threads;
     std::function<bool(SDBackendModule)> ensure_backend_pair;
@@ -39,7 +39,8 @@ struct GenerationExtensionConditionContext {
 struct GenerationExtension {
     virtual ~GenerationExtension() = default;
 
-    virtual const char* name() const = 0;
+    virtual ModelComponent component() const = 0;
+    const char* name() const { return model_component_name(component()); }
     virtual bool is_enabled() const {
         return false;
     }

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 {
@@ -109,8 +134,8 @@ struct PhotoMakerExtension : public GenerationExtension {
     SDCondition id_condition;
     int start_merge_step = -1;
 
-    const char* name() const override {
-        return "photomaker";
+    ModelComponent component() const override {
+        return ModelComponent::PhotoMaker;
     }
 
     bool is_enabled() const override {
@@ -119,7 +144,7 @@ struct PhotoMakerExtension : public GenerationExtension {
 
     bool init(const GenerationExtensionInitContext& ctx) override {
         model_path = SAFE_STR(ctx.params->photo_maker_path);
-        if (model_path.empty()) {
+        if (model_path.empty() || !ctx.photomaker_source_available) {
             return true;
         }
 
@@ -128,19 +153,14 @@ struct PhotoMakerExtension : public GenerationExtension {
         }
 
         PMVersion pm_version = std::strstr(model_path.c_str(), "v2") != nullptr ? PM_VERSION_2 : PM_VERSION_1;
-        LOG_INFO("loading stacked ID embedding (PHOTOMAKER) model file from '%s'", model_path.c_str());
-        if (!ctx.model_loader.init_from_file_and_convert_name(model_path, "pmid.")) {
-            LOG_WARN("loading stacked ID embedding from '%s' failed", model_path.c_str());
-            return true;
-        }
-
-        pmid_model = std::make_shared<PhotoMakerIDEncoder>(ctx.backend_for(SDBackendModule::PHOTOMAKER),
+        pmid_model           = std::make_shared<PhotoMakerIDEncoder>(ctx.backend_for(SDBackendModule::PHOTOMAKER),
                                                            ctx.tensor_storage_map,
                                                            "pmid",
                                                            ctx.version,
                                                            pm_version,
                                                            20.f,
                                                            ctx.model_manager);
+        pmid_model->set_scale_overrides(ctx.params->linear_scale, ctx.params->attn_scale);
         if (pm_version == PM_VERSION_2) {
             LOG_INFO("using PhotoMaker Version 2");
         }
@@ -228,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");
@@ -268,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/extensions/pulid_extension.cpp

@@ -79,8 +79,8 @@ struct PuLIDExtension : public GenerationExtension {
     sd::Tensor<float> id_embedding;
     float id_weight = 1.0f;
 
-    const char* name() const override {
-        return "pulid";
+    ModelComponent component() const override {
+        return ModelComponent::PuLID;
     }
 
     bool is_enabled() const override {

src/model.h

@@ -35,6 +35,7 @@ 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,
@@ -57,6 +58,7 @@ enum SDVersion {
     VERSION_SEFI_IMAGE,
     VERSION_KREA2,
     VERSION_MAGE_FLOW,
+    VERSION_SENSENOVA_U1_5,
     VERSION_ESRGAN,
     VERSION_COUNT,
 };
@@ -129,7 +131,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;
@@ -237,6 +239,10 @@ static inline bool sd_version_is_mage_flow(SDVersion version) {
     return version == VERSION_MAGE_FLOW;
 }
 
+static inline bool sd_version_is_sensenova_u1(SDVersion version) {
+    return version == VERSION_SENSENOVA_U1_5;
+}
+
 static inline bool sd_version_uses_flux_vae(SDVersion version) {
     if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) {
         return true;
@@ -295,7 +301,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_ideogram4(version) ||
         sd_version_is_sefi_image(version) ||
         sd_version_is_krea2(version) ||
-        sd_version_is_mage_flow(version)) {
+        sd_version_is_mage_flow(version) ||
+        sd_version_is_sensenova_u1(version)) {
         return true;
     }
     return false;

src/model/adapter/ip_adapter.hpp

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

src/model/adapter/lora.hpp

@@ -23,25 +23,31 @@ struct LoraModel : public GGMLRunner {
     std::set<std::string> skipped_incompatible_lora_tensors;
     std::set<std::string> warned_incompatible_model_tensors;
     std::string file_path;
-    std::shared_ptr<ModelManager> model_manager;
-    ggml_backend_t params_backend = nullptr;
-    bool load_failed              = false;
-    bool applied                  = false;
-    bool tensor_preprocessed      = false;
+    ggml_backend_t params_backend              = nullptr;
+    bool load_failed                           = false;
+    bool applied                               = false;
+    bool tensor_preprocessed                   = false;
+    ModelLoader::FileId source_file            = 0;
+    SDVersion source_version                   = VERSION_COUNT;
+    ModelManager::ResidencyMode residency_mode = ModelManager::ResidencyMode::ParamBackend;
+    bool params_follow_compute                 = false;
+    std::vector<ggml_tensor*> registered_params;
+    std::map<ggml_tensor*, float> scalar_values;
 
     typedef std::function<bool(const std::string&)> filter_t;
 
-    LoraModel(const std::string& lora_id,
-              ggml_backend_t backend,
-              ggml_backend_t params_backend_,
-              const std::string& file_path          = "",
-              std::string prefix                    = "",
-              SDVersion version                     = VERSION_COUNT,
-              std::shared_ptr<ModelManager> manager = std::make_shared<ModelManager>())
-        : GGMLRunner(backend, manager), lora_id(lora_id), file_path(file_path), model_manager(std::move(manager)), params_backend(params_backend_) {
-        prefix = "lora." + prefix;
-        if (model_manager == nullptr || !model_manager->loader().init_from_file_and_convert_name(file_path, prefix, version)) {
-            load_failed = true;
+    LoraModel(const std::string& id, ggml_backend_t backend, ggml_backend_t params, std::shared_ptr<ModelManager> manager, ModelLoader::FileId file, SDVersion version, ModelManager::ResidencyMode mode = ModelManager::ResidencyMode::ParamBackend, bool follow_compute = false)
+        : GGMLRunner(backend, manager), lora_id(id), params_backend(params), source_file(file), source_version(version), residency_mode(mode), params_follow_compute(follow_compute) {
+        load_failed = source_file == 0 || manager == nullptr || manager->loader().file_revision(source_file) == 0;
+        if (!load_failed) {
+            file_path = manager->loader().file_path(source_file);
+        }
+    }
+
+    ~LoraModel() override {
+        runner_end();
+        if (auto manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock())) {
+            GGML_ASSERT(manager->unregister_param_tensors(registered_params));
         }
     }
 
@@ -49,95 +55,65 @@ struct LoraModel : public GGMLRunner {
         return "lora";
     }
 
-    bool load_from_file(int n_threads, filter_t filter = nullptr) {
-        LOG_INFO("loading LoRA from '%s'", file_path.c_str());
-
-        if (load_failed) {
-            LOG_ERROR("init lora model loader from file failed: '%s'", file_path.c_str());
+    bool init_params(int n_threads, filter_t filter = nullptr) {
+        auto model_manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock());
+        if (model_manager == nullptr)
             return false;
-        }
-
-        std::unordered_map<std::string, TensorStorage> tensors_to_create;
-        std::mutex lora_mutex;
-        bool dry_run          = true;
-        auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
-            if (dry_run) {
-                const std::string& name = tensor_storage.name;
-
-                if (filter && !filter(name)) {
-                    return true;
-                }
-
-                {
-                    std::lock_guard<std::mutex> lock(lora_mutex);
-                    tensors_to_create[name] = tensor_storage;
-                }
-            } else {
-                const std::string& name = tensor_storage.name;
-                auto iter               = lora_tensors.find(name);
-                if (iter != lora_tensors.end()) {
-                    *dst_tensor = iter->second;
-                }
-            }
-            return true;
-        };
-
-        if (model_manager != nullptr) {
-            model_manager->set_n_threads(n_threads);
-        }
-        ModelLoader& model_loader = model_manager->loader();
-        model_loader.load_tensors(on_new_tensor_cb);
-
-        if (tensors_to_create.empty()) {
-            return true;
-        }
-
-        for (const auto& pair : tensors_to_create) {
-            const auto& name   = pair.first;
-            const auto& ts     = pair.second;
-            ggml_tensor* real  = ggml_new_tensor(params_ctx,
-                                                 ts.type,
-                                                 ts.n_dims,
-                                                 ts.ne);
-            lora_tensors[name] = real;
-        }
-
+        if (load_failed || !registered_params.empty())
+            return false;
+        model_manager->set_n_threads(n_threads);
+        const auto sources = model_manager->loader().file_tensors(source_file, source_version);
         std::map<std::string, ggml_tensor*> tensors;
-        for (const auto& pair : lora_tensors) {
-            tensors[pair.first] = pair.second;
+        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)))
+                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);
+            lora_tensors[name] = tensor;
+            if (scalar) {
+                tensor->data  = &scalar_values[tensor];
+                scalars[name] = tensor;
+                scalar_names.insert(name);
+            } else {
+                tensors[name] = tensor;
+            }
         }
-        if (model_manager == nullptr ||
-            !model_manager->register_param_tensors("LoRA",
-                                                   std::move(tensors),
-                                                   ModelManager::ResidencyMode::ParamBackend,
-                                                   runtime_backend,
-                                                   params_backend) ||
-            !model_manager->validate_registered_tensors()) {
-            LOG_ERROR("lora model manager registration failed");
+        // These values are consumed while constructing the graph, before weight preparation.
+        if (!scalars.empty()) {
+            auto callback = [&](const TensorStorage& source, ggml_tensor** dst) {
+                auto found = scalars.find(source.name);
+                *dst       = found == scalars.end() ? nullptr : found->second;
+                return true;
+            };
+            if (!model_manager->loader().load_file_tensors(source_file, source_version, callback, scalar_names))
+                return false;
+        }
+        if (!model_manager->register_param_tensors(ModelComponent::LoRA, tensors, residency_mode,
+                                                   runtime_backend, params_backend, nullptr, false, params_follow_compute,
+                                                   nullptr, source_file, source_version))
             return false;
-        }
-        std::vector<ggml_tensor*> lora_params;
-        lora_params.reserve(lora_tensors.size());
-        for (const auto& pair : lora_tensors) {
-            lora_params.push_back(pair.second);
-        }
-        if (!model_manager->prepare_params(lora_params)) {
-            LOG_ERROR("lora model manager prepare params failed");
-            return false;
-        }
+        for (const auto& entry : tensors)
+            registered_params.push_back(entry.second);
+        return model_manager->validate_registered_tensors();
+    }
 
-        LOG_VERBOSE("finished loaded lora");
-        return true;
+    float scalar_value(ggml_tensor* tensor) const {
+        auto found = scalar_values.find(tensor);
+        return found != scalar_values.end() ? found->second : ggml_ext_backend_tensor_get_f32(tensor);
     }
 
     void release_loaded_tensors() {
         runner_end();
-        model_manager.reset();
+        if (auto manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock())) {
+            GGML_ASSERT(manager->unregister_param_tensors(registered_params));
+        }
+        registered_params.clear();
         free_params_ctx();
         alloc_params_ctx();
-        model_manager     = std::make_shared<ModelManager>();
-        residency_manager = model_manager;
         lora_tensors.clear();
+        scalar_values.clear();
         original_tensor_to_final_tensor.clear();
         applied_lora_tensors.clear();
         skipped_incompatible_lora_tensors.clear();
@@ -241,12 +217,12 @@ struct LoraModel : public GGMLRunner {
             int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
             iter         = lora_tensors.find(scale_name);
             if (iter != lora_tensors.end()) {
-                scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
+                scale_value = scalar_value(iter->second);
                 applied_lora_tensors.insert(scale_name);
             } else {
                 iter = lora_tensors.find(alpha_name);
                 if (iter != lora_tensors.end()) {
-                    float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
+                    float alpha = scalar_value(iter->second);
                     scale_value = alpha / rank;
                     // LOG_VERBOSE("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
                     applied_lora_tensors.insert(alpha_name);
@@ -395,7 +371,7 @@ struct LoraModel : public GGMLRunner {
             int64_t rank = hada_1_down->ne[ggml_n_dims(hada_1_down) - 1];
             iter         = lora_tensors.find(alpha_name);
             if (iter != lora_tensors.end()) {
-                float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
+                float alpha = scalar_value(iter->second);
                 scale_value = alpha / rank;
                 applied_lora_tensors.insert(alpha_name);
             }
@@ -508,7 +484,7 @@ struct LoraModel : public GGMLRunner {
             float scale_value = 1.0f;
             iter              = lora_tensors.find(alpha_name);
             if (iter != lora_tensors.end()) {
-                float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
+                float alpha = scalar_value(iter->second);
                 scale_value = alpha / rank;
                 applied_lora_tensors.insert(alpha_name);
             }
@@ -669,7 +645,7 @@ struct LoraModel : public GGMLRunner {
                 float scale_value = 1.0f;
                 iter              = lora_tensors.find(alpha_name);
                 if (iter != lora_tensors.end()) {
-                    float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
+                    float alpha = scalar_value(iter->second);
                     scale_value = alpha / rank;
                 }
 
@@ -796,12 +772,12 @@ struct LoraModel : public GGMLRunner {
             int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
             iter         = lora_tensors.find(scale_name);
             if (iter != lora_tensors.end()) {
-                scale_value       = ggml_ext_backend_tensor_get_f32(iter->second);
+                scale_value       = scalar_value(iter->second);
                 scale_tensor_name = scale_name;
             } else {
                 iter = lora_tensors.find(alpha_name);
                 if (iter != lora_tensors.end()) {
-                    float alpha       = ggml_ext_backend_tensor_get_f32(iter->second);
+                    float alpha       = scalar_value(iter->second);
                     scale_value       = alpha / rank;
                     scale_tensor_name = alpha_name;
                     // LOG_VERBOSE("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
@@ -949,7 +925,7 @@ struct LoraModel : public GGMLRunner {
         return gf;
     }
 
-    void apply(std::map<std::string, ggml_tensor*> model_tensors,
+    bool apply(std::map<std::string, ggml_tensor*> model_tensors,
                const std::set<std::string>& model_tensor_names,
                SDVersion version,
                int n_threads,
@@ -970,10 +946,11 @@ struct LoraModel : public GGMLRunner {
         stat(!warn_unused);
         original_tensor_to_final_tensor.clear();
         runner_end();
+        return result.has_value();
     }
 
-    void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads, bool warn_unused = true) {
-        apply(model_tensors, tensor_names(model_tensors), version, n_threads, warn_unused);
+    bool apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads, bool warn_unused = true) {
+        return apply(model_tensors, tensor_names(model_tensors), version, n_threads, warn_unused);
     }
 
     void stat(bool at_runntime = false) {

src/model/adapter/pmid.hpp

@@ -6,10 +6,8 @@
 #include "core/util.h"
 #include "model/common/ggml_block.hpp"
 
-#include "model/adapter/lora.hpp"
 #include "model/common/block.hpp"
 #include "model/te/clip.hpp"
-#include "model_loader.h"
 
 struct FuseBlock : public GGMLBlock {
     // network hparams
@@ -565,94 +563,4 @@ public:
     }
 };
 
-struct PhotoMakerIDEmbed : public GGMLRunner {
-    std::map<std::string, ggml_tensor*> tensors;
-    std::string file_path;
-    std::shared_ptr<ModelManager> model_manager;
-    ggml_backend_t params_backend = nullptr;
-    bool load_failed              = false;
-    bool applied                  = false;
-
-    PhotoMakerIDEmbed(ggml_backend_t backend,
-                      ggml_backend_t params_backend_,
-                      std::shared_ptr<ModelManager> manager = std::make_shared<ModelManager>(),
-                      const std::string& file_path          = "",
-                      const std::string& prefix             = "")
-        : GGMLRunner(backend, manager), file_path(file_path), model_manager(std::move(manager)), params_backend(params_backend_) {
-        if (model_manager == nullptr || !model_manager->loader().init_from_file_and_convert_name(file_path, prefix)) {
-            load_failed = true;
-        }
-    }
-
-    std::string get_desc() {
-        return "id_embeds";
-    }
-
-    bool load_from_file(bool filter_tensor, int n_threads) {
-        LOG_INFO("loading PhotoMaker ID Embeds from '%s'", file_path.c_str());
-
-        if (load_failed) {
-            LOG_ERROR("init photomaker id embed from file failed: '%s'", file_path.c_str());
-            return false;
-        }
-
-        bool dry_run = true;
-        std::mutex tensor_mutex;
-        auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
-            const std::string& name = tensor_storage.name;
-
-            if (filter_tensor && !contains(name, "pmid.id_embeds")) {
-                // LOG_INFO("skipping LoRA tesnor '%s'", name.c_str());
-                return true;
-            }
-            if (dry_run) {
-                std::lock_guard<std::mutex> lock(tensor_mutex);
-                ggml_tensor* real = ggml_new_tensor(params_ctx,
-                                                    tensor_storage.type,
-                                                    tensor_storage.n_dims,
-                                                    tensor_storage.ne);
-                tensors[name]     = real;
-            } else {
-                auto real   = tensors[name];
-                *dst_tensor = real;
-            }
-
-            return true;
-        };
-
-        model_manager->set_n_threads(n_threads);
-        ModelLoader& model_loader = model_manager->loader();
-        model_loader.load_tensors(on_new_tensor_cb);
-        if (!model_manager->register_param_tensors("PhotoMaker ID embeds",
-                                                   tensors,
-                                                   ModelManager::ResidencyMode::ParamBackend,
-                                                   runtime_backend,
-                                                   params_backend) ||
-            !model_manager->validate_registered_tensors()) {
-            LOG_ERROR("PhotoMaker ID embeds model manager registration failed");
-            return false;
-        }
-        std::vector<ggml_tensor*> id_embed_params;
-        id_embed_params.reserve(tensors.size());
-        for (const auto& pair : tensors) {
-            id_embed_params.push_back(pair.second);
-        }
-        if (!model_manager->prepare_params(id_embed_params)) {
-            LOG_ERROR("PhotoMaker ID embeds model manager prepare params failed");
-            return false;
-        }
-
-        LOG_VERBOSE("finished loading PhotoMaker ID Embeds ");
-        return true;
-    }
-
-    ggml_tensor* get() {
-        std::map<std::string, ggml_tensor*>::iterator pos;
-        pos = tensors.find("pmid.id_embeds");
-        if (pos != tensors.end())
-            return pos->second;
-        return nullptr;
-    }
-};
-
 #endif  // __SD_MODEL_ADAPTER_PMID_HPP__

src/model/adapter/pulid.hpp

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

src/model/audio/wav2vec2.hpp

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

src/model/common/block.hpp

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

src/model/common/ggml_block.hpp

@@ -206,6 +206,7 @@ public:
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         ggml_tensor* w            = params["weight"];
+        const float scale         = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale;
         ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr;
         if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) {
             bool supports_fp8_matmul = false;
@@ -309,6 +310,7 @@ public:
 __STATIC_INLINE__ bool support_get_rows(ggml_type wtype) {
     switch (wtype) {
         case GGML_TYPE_F16:
+        case GGML_TYPE_BF16:
         case GGML_TYPE_Q8_0:
         case GGML_TYPE_Q5_1:
         case GGML_TYPE_Q5_0:
@@ -366,6 +368,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;
@@ -764,7 +821,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);
     }
 };
 
@@ -869,7 +926,7 @@ public:
             v = v_proj->forward(ctx, x);
         }
 
-        x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, false);  // [N, n_token, embed_dim]
+        x = ggml_ext_attention_ext(ctx, q, k, v, n_head, mask, false);  // [N, n_token, embed_dim]
 
         x = out_proj->forward(ctx, x);  // [N, n_token, embed_dim]
         return x;

src/model/common/rope.hpp

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

src/model/diffusion/anima.hpp

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

src/model/diffusion/animatediff.hpp

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

src/model/diffusion/control.hpp

@@ -2,8 +2,6 @@
 #define __SD_MODEL_DIFFUSION_CONTROL_HPP__
 
 #include "model/common/block.hpp"
-#include "model_loader.h"
-#include "model_manager.h"
 
 // Match main UNet's MAX_GRAPH_SIZE so SDXL ControlNet (transformer_depth={1,2,10}) fits.
 #define CONTROL_NET_GRAPH_SIZE MAX_GRAPH_SIZE
@@ -317,20 +315,17 @@ struct ControlNet : public GGMLRunner {
     ggml_tensor* guided_hint_output_ggml = nullptr;
     std::vector<sd::Tensor<float>> controls;
     bool guided_hint_cached = false;
-    std::shared_ptr<ModelManager> owned_model_manager;
-    ggml_backend_t params_backend = nullptr;
 
     static const char* guided_hint_cache_name() {
         return "controlnet.guided_hint";
     }
 
     ControlNet(ggml_backend_t backend,
-               ggml_backend_t params_backend_,
                const String2TensorStorage& tensor_storage_map      = {},
                SDVersion version                                   = VERSION_SD1,
                const std::string& prefix                           = "",
                std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
-        : GGMLRunner(backend, weight_manager), version(version), control_net(version), weight_prefix(prefix), params_backend(params_backend_) {
+        : GGMLRunner(backend, weight_manager), version(version), control_net(version), weight_prefix(prefix) {
         control_net.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -445,39 +440,6 @@ struct ControlNet : public GGMLRunner {
         guided_hint_cached = get_cache_tensor_by_name(guided_hint_cache_name()) != nullptr;
         return controls;
     }
-
-    bool load_from_file(const std::string& file_path, int n_threads) {
-        LOG_INFO("loading control net from '%s'", file_path.c_str());
-        std::map<std::string, ggml_tensor*> tensors;
-        control_net.get_param_tensors(tensors);
-
-        auto manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock());
-        if (manager == nullptr) {
-            owned_model_manager = std::make_shared<ModelManager>();
-            residency_manager   = owned_model_manager;
-            manager             = owned_model_manager;
-        }
-
-        ModelLoader& model_loader = manager->loader();
-        if (!model_loader.init_from_file_and_convert_name(file_path)) {
-            LOG_ERROR("init control net model loader from file failed: '%s'", file_path.c_str());
-            return false;
-        }
-
-        manager->set_n_threads(n_threads);
-        if (!manager->register_param_tensors("ControlNet",
-                                             std::move(tensors),
-                                             ModelManager::ResidencyMode::ParamBackend,
-                                             runtime_backend,
-                                             params_backend) ||
-            !manager->validate_registered_tensors()) {
-            LOG_ERROR("register control net tensors with model manager failed");
-            return false;
-        }
-
-        LOG_INFO("control net model loaded");
-        return true;
-    }
 };
 
 #endif  // __SD_MODEL_DIFFUSION_CONTROL_HPP__

src/model/diffusion/ernie_image.hpp

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

src/model/diffusion/flux.hpp

@@ -1714,8 +1714,8 @@ namespace Flux {
             ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_COUNT;
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -1736,7 +1736,8 @@ namespace Flux {
                                                                             VERSION_FLUX2,
                                                                             model_manager);
 
-            if (!model_manager->register_runner_params("Flux test",
+            if (!model_manager->set_loader(model_loader) ||
+                !model_manager->register_runner_params(ModelComponent::Diffusion,
                                                        *flux,
                                                        "model.diffusion_model",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/diffusion/hidream_o1.hpp

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

src/model/diffusion/hunyuan.hpp

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

src/model/diffusion/krea2.hpp

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

src/model/diffusion/ltxv.hpp

@@ -709,8 +709,7 @@ namespace LTXV {
                 k = apply_hidden_rope(ctx->ggml_ctx, k, k_pe, heads, dim_head, rope_interleaved);
             }
 
-            auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
-                                              ctx->backend,
+            auto out = ggml_ext_attention_ext(ctx,
                                               q,
                                               k,
                                               v,
@@ -2087,8 +2086,8 @@ namespace LTXV {
             ggml_backend_t backend = sd_backend_cpu_init();
             LOG_INFO("loading ltxav from '%s'", model_path.c_str());
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(model_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
                 return;
@@ -2107,7 +2106,8 @@ namespace LTXV {
                                                                                "model.diffusion_model",
                                                                                model_manager);
 
-            if (!model_manager->register_runner_params("LTXAV test",
+            if (!model_manager->set_loader(model_loader) ||
+                !model_manager->register_runner_params(ModelComponent::Diffusion,
                                                        *ltxav,
                                                        "model.diffusion_model",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/diffusion/minimax_h3.hpp

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

src/model/diffusion/mmdit.hpp

@@ -365,8 +365,8 @@ public:
     ggml_tensor* forward(GGMLRunnerContext* ctx,
                          ggml_tensor* x) {
         auto qkv = pre_attention(ctx, x);
-        x        = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
-        x        = post_attention(ctx, x);                                                                                                           // [N, n_token, dim]
+        x        = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
+        x        = post_attention(ctx, x);                                                                                   // [N, n_token, dim]
         return x;
     }
 };
@@ -587,8 +587,8 @@ public:
             auto qkv2          = std::get<1>(qkv_intermediates);
             auto intermediates = std::get<2>(qkv_intermediates);
 
-            auto attn_out  = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);     // [N, n_token, dim]
-            auto attn2_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
+            auto attn_out  = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);     // [N, n_token, dim]
+            auto attn2_out = ggml_ext_attention_ext(ctx, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
             x              = post_attention_x(ctx,
                                               attn_out,
                                               attn2_out,
@@ -604,7 +604,7 @@ public:
             auto qkv               = qkv_intermediates.first;
             auto intermediates     = qkv_intermediates.second;
 
-            auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
+            auto attn_out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, dim]
             x             = post_attention(ctx,
                                            attn_out,
                                            intermediates[0],
@@ -648,7 +648,7 @@ block_mixing(GGMLRunnerContext* ctx,
         qkv.push_back(ggml_concat(ctx->ggml_ctx, context_qkv[i], x_qkv[i], 1));
     }
 
-    auto attn = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_context + n_token, hidden_size]
+    auto attn = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_context + n_token, hidden_size]
 
     auto context_attn = ggml_view_3d(ctx->ggml_ctx,
                                      attn,
@@ -680,7 +680,7 @@ block_mixing(GGMLRunnerContext* ctx,
     }
 
     if (x_block->self_attn) {
-        auto attn2 = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, hidden_size]
+        auto attn2 = ggml_ext_attention_ext(ctx, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, hidden_size]
 
         x = x_block->post_attention_x(ctx,
                                       x_attn,
@@ -1064,13 +1064,14 @@ struct MMDiTRunner : public DiffusionModelRunner {
         {
             LOG_INFO("loading from '%s'", file_path.c_str());
 
-            ModelLoader& model_loader = model_manager->loader();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(file_path)) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
             }
 
-            if (!model_manager->register_runner_params("MMDiT test",
+            if (!model_manager->set_loader(std::move(model_loader)) ||
+                !model_manager->register_runner_params(ModelComponent::Diffusion,
                                                        *mmdit,
                                                        "model.diffusion_model",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/diffusion/model.hpp

@@ -69,6 +69,8 @@ struct AnimaDiffusionExtra {
 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 {
@@ -114,6 +116,10 @@ struct MiniT2IDiffusionExtra {
     const sd::Tensor<float>* mask = nullptr;
 };
 
+struct SenseNovaU1DiffusionExtra {
+    const sd::Tensor<int32_t>* input_ids = nullptr;
+};
+
 struct HunyuanVideoDiffusionExtra {
     const sd::Tensor<float>* guidance   = nullptr;
     const sd::Tensor<float>* byt5       = nullptr;
@@ -131,6 +137,7 @@ using DiffusionExtraParams = std::variant<std::monostate,
                                           LTXAVDiffusionExtra,
                                           MiniMaxH3DiffusionExtra,
                                           MiniT2IDiffusionExtra,
+                                          SenseNovaU1DiffusionExtra,
                                           HunyuanVideoDiffusionExtra>;
 
 struct DiffusionParams {

src/model/diffusion/qwen_image.hpp

@@ -773,8 +773,8 @@ namespace Qwen {
             ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_Q8_0;
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -793,7 +793,8 @@ namespace Qwen {
                                                                                             VERSION_QWEN_IMAGE,
                                                                                             model_manager);
 
-            if (!model_manager->register_runner_params("Qwen image test",
+            if (!model_manager->set_loader(model_loader) ||
+                !model_manager->register_runner_params(ModelComponent::Diffusion,
                                                        *qwen_image,
                                                        "model.diffusion_model",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/diffusion/sensenova_u1.h

@@ -0,0 +1,846 @@
+#ifndef __SD_MODEL_DIFFUSION_SENSENOVA_U1_H__
+#define __SD_MODEL_DIFFUSION_SENSENOVA_U1_H__
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <memory>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "core/ggml_extend.h"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/model.hpp"
+#include "model/te/llm.hpp"
+#include "model_loader.h"
+
+namespace SenseNovaU1 {
+    constexpr int SENSENOVA_U1_GRAPH_SIZE = 327680;
+
+    struct SenseNovaU1Config {
+        int64_t hidden_size                  = 4096;
+        int64_t intermediate_size            = 12288;
+        int64_t num_layers                   = 42;
+        int64_t num_heads                    = 32;
+        int64_t num_kv_heads                 = 8;
+        int64_t head_dim                     = 128;
+        int64_t vocab_size                   = 151936;
+        int64_t max_position_embeddings      = 262144;
+        int64_t max_position_embeddings_hw   = 10000;
+        int64_t vision_hidden_size           = 1024;
+        int64_t patch_size                   = 16;
+        int64_t vision_downsample_factor     = 2;
+        int64_t in_channels                  = 3;
+        int64_t timestep_embedding_size      = 256;
+        float rms_norm_eps                   = 1e-6f;
+        float rope_theta                     = 5000000.f;
+        float rope_theta_hw                  = 10000.f;
+        float noise_scale_base_image_seq_len = 64.f;
+        float noise_scale_max_value          = 16.f;
+        float t_eps                          = 0.02f;
+        bool add_noise_scale_embedding       = true;
+
+        int64_t image_token_stride() const {
+            return patch_size * vision_downsample_factor;
+        }
+
+        static SenseNovaU1Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                                     const std::string& prefix) {
+            SenseNovaU1Config config;
+            config.num_layers      = 0;
+            const std::string root = prefix.empty() ? "" : prefix + ".";
+
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, root)) {
+                    continue;
+                }
+                if (ends_with(name, "language_model.model.embed_tokens.weight") && tensor_storage.n_dims == 2) {
+                    config.hidden_size = tensor_storage.ne[0];
+                    config.vocab_size  = tensor_storage.ne[1];
+                } else if (ends_with(name, "language_model.model.layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {
+                    config.intermediate_size = tensor_storage.ne[1];
+                } else if (ends_with(name, "language_model.model.layers.0.self_attn.q_proj.weight") && tensor_storage.n_dims == 2) {
+                    config.num_heads = tensor_storage.ne[1] / config.head_dim;
+                } else if (ends_with(name, "language_model.model.layers.0.self_attn.k_proj.weight") && tensor_storage.n_dims == 2) {
+                    config.num_kv_heads = tensor_storage.ne[1] / config.head_dim;
+                } else if (ends_with(name, "fm_modules.vision_model_mot_gen.embeddings.patch_embedding.weight") && tensor_storage.n_dims == 4) {
+                    config.patch_size         = tensor_storage.ne[0];
+                    config.in_channels        = tensor_storage.ne[2];
+                    config.vision_hidden_size = tensor_storage.ne[3];
+                } else if (ends_with(name, "fm_modules.vision_model_mot_gen.embeddings.dense_embedding.weight") && tensor_storage.n_dims == 4) {
+                    config.vision_downsample_factor = tensor_storage.ne[0];
+                }
+
+                const std::string layer_prefix = root + "language_model.model.layers.";
+                if (starts_with(name, layer_prefix)) {
+                    const char* index_begin = name.c_str() + layer_prefix.size();
+                    config.num_layers       = std::max<int64_t>(config.num_layers, std::strtoll(index_begin, nullptr, 10) + 1);
+                }
+            }
+
+            if (config.num_layers == 0) {
+                config.num_layers = 42;
+            }
+            config.add_noise_scale_embedding = tensor_storage_map.find(root + "fm_modules.noise_scale_embedder.mlp.0.weight") != tensor_storage_map.end();
+
+            LOG_DEBUG("sensenova-u1.5: layers=%" PRId64 ", hidden=%" PRId64 ", intermediate=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", patch=%" PRId64 "x%" PRId64,
+                      config.num_layers,
+                      config.hidden_size,
+                      config.intermediate_size,
+                      config.num_heads,
+                      config.num_kv_heads,
+                      config.patch_size,
+                      config.vision_downsample_factor);
+            return config;
+        }
+    };
+
+    class StorageConv2d : public Conv2d {
+    protected:
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            this->prefix     = prefix;
+            ggml_type wtype  = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F16);
+            params["weight"] = ggml_new_tensor_4d(ctx,
+                                                  wtype,
+                                                  kernel_size.second,
+                                                  kernel_size.first,
+                                                  in_channels,
+                                                  out_channels);
+            if (bias) {
+                params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+            }
+        }
+
+    public:
+        StorageConv2d(int64_t in_channels,
+                      int64_t out_channels,
+                      std::pair<int, int> kernel_size,
+                      std::pair<int, int> stride  = {1, 1},
+                      std::pair<int, int> padding = {0, 0},
+                      bool bias                   = true)
+            : Conv2d(in_channels,
+                     out_channels,
+                     kernel_size,
+                     stride,
+                     padding,
+                     {1, 1},
+                     bias) {}
+    };
+
+    struct TimestepEmbedder : public GGMLBlock {
+        int64_t frequency_embedding_size;
+
+        TimestepEmbedder(int64_t hidden_size, int64_t frequency_embedding_size = 256)
+            : frequency_embedding_size(frequency_embedding_size) {
+            blocks["mlp.0"] = std::make_shared<Linear>(frequency_embedding_size, hidden_size, true);
+            blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timesteps) {
+            auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
+            auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
+            auto x     = ggml_ext_timestep_embedding(ctx->ggml_ctx,
+                                                     timesteps,
+                                                     static_cast<int>(frequency_embedding_size),
+                                                     10000,
+                                                     1.f);
+            x          = mlp_0->forward(ctx, x);
+            x          = ggml_silu_inplace(ctx->ggml_ctx, x);
+            return mlp_2->forward(ctx, x);
+        }
+    };
+
+    inline ggml_tensor* apply_vision_rope(GGMLRunnerContext* ctx,
+                                          ggml_tensor* x,
+                                          ggml_tensor* position_x,
+                                          ggml_tensor* position_y,
+                                          float theta,
+                                          int max_position) {
+        GGML_ASSERT(x->ne[0] % 2 == 0);
+        // ggml_rope_ext addresses positions through ne[2]. The vision
+        // embeddings arrive as [hidden, tokens, batch], so add the singleton
+        // head axis used by the RoPE kernel: [hidden, 1, tokens, batch].
+        x                  = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]);
+        const int64_t half = x->ne[0] / 2;
+        auto x_part        = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, half);
+        auto y_part        = ggml_ext_slice(ctx->ggml_ctx, x, 0, half, x->ne[0]);
+        x_part             = ggml_rope_ext(ctx->ggml_ctx,
+                                           x_part,
+                                           position_x,
+                                           nullptr,
+                                           static_cast<int>(half),
+                                           GGML_ROPE_TYPE_NORMAL,
+                                           max_position,
+                                           theta,
+                                           1.f,
+                                           0.f,
+                                           1.f,
+                                           32.f,
+                                           1.f);
+        y_part             = ggml_rope_ext(ctx->ggml_ctx,
+                                           y_part,
+                                           position_y,
+                                           nullptr,
+                                           static_cast<int>(half),
+                                           GGML_ROPE_TYPE_NORMAL,
+                                           max_position,
+                                           theta,
+                                           1.f,
+                                           0.f,
+                                           1.f,
+                                           32.f,
+                                           1.f);
+        return ggml_concat(ctx->ggml_ctx, x_part, y_part, 0);
+    }
+
+    struct VisionEmbeddings : public GGMLBlock {
+        SenseNovaU1Config config;
+
+        explicit VisionEmbeddings(const SenseNovaU1Config& config)
+            : config(config) {
+            blocks["patch_embedding"] = std::make_shared<StorageConv2d>(config.in_channels,
+                                                                        config.vision_hidden_size,
+                                                                        std::pair<int, int>{static_cast<int>(config.patch_size), static_cast<int>(config.patch_size)},
+                                                                        std::pair<int, int>{static_cast<int>(config.patch_size), static_cast<int>(config.patch_size)},
+                                                                        std::pair<int, int>{0, 0},
+                                                                        true);
+            blocks["dense_embedding"] = std::make_shared<StorageConv2d>(config.vision_hidden_size,
+                                                                        config.hidden_size,
+                                                                        std::pair<int, int>{static_cast<int>(config.vision_downsample_factor), static_cast<int>(config.vision_downsample_factor)},
+                                                                        std::pair<int, int>{static_cast<int>(config.vision_downsample_factor), static_cast<int>(config.vision_downsample_factor)},
+                                                                        std::pair<int, int>{0, 0},
+                                                                        true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* image,
+                             ggml_tensor* position_x,
+                             ggml_tensor* position_y) {
+            auto patch_embedding = std::dynamic_pointer_cast<StorageConv2d>(blocks["patch_embedding"]);
+            auto dense_embedding = std::dynamic_pointer_cast<StorageConv2d>(blocks["dense_embedding"]);
+
+            auto x = patch_embedding->forward(ctx, image);
+            x      = ggml_gelu_erf(ctx->ggml_ctx, x);
+
+            const int64_t grid_w = x->ne[0];
+            const int64_t grid_h = x->ne[1];
+            const int64_t batch  = x->ne[3];
+            x                    = ggml_reshape_3d(ctx->ggml_ctx, x, grid_w * grid_h, x->ne[2], batch);
+            x                    = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x                    = apply_vision_rope(ctx,
+                                                     x,
+                                                     position_x,
+                                                     position_y,
+                                                     config.rope_theta_hw,
+                                                     static_cast<int>(config.max_position_embeddings_hw));
+            x                    = ggml_reshape_4d(ctx->ggml_ctx, x, config.vision_hidden_size, grid_w, grid_h, batch);
+            x                    = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+            x                    = dense_embedding->forward(ctx, x);
+
+            const int64_t token_w = x->ne[0];
+            const int64_t token_h = x->ne[1];
+            x                     = ggml_reshape_3d(ctx->ggml_ctx, x, token_w * token_h, x->ne[2], x->ne[3]);
+            return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+        }
+    };
+
+    inline ggml_tensor* pixel_shuffle(GGMLRunnerContext* ctx,
+                                      ggml_tensor* x,
+                                      int upscale_factor) {
+        GGML_ASSERT(upscale_factor > 0);
+        const int64_t h = x->ne[1];
+        const int64_t w = x->ne[0];
+        GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0);
+        x = ggml_ext_cont(ctx->ggml_ctx,
+                          ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+        x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]);
+        return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true);
+    }
+
+    struct PixelDecoder : public GGMLBlock {
+        explicit PixelDecoder(const SenseNovaU1Config& config) {
+            blocks["conv1"] = std::make_shared<StorageConv2d>(config.hidden_size / 4,
+                                                              1024,
+                                                              std::pair<int, int>{3, 3},
+                                                              std::pair<int, int>{1, 1},
+                                                              std::pair<int, int>{1, 1},
+                                                              true);
+            blocks["conv2"] = std::make_shared<StorageConv2d>(256,
+                                                              192,
+                                                              std::pair<int, int>{3, 3},
+                                                              std::pair<int, int>{1, 1},
+                                                              std::pair<int, int>{1, 1},
+                                                              true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto conv1 = std::dynamic_pointer_cast<StorageConv2d>(blocks["conv1"]);
+            auto conv2 = std::dynamic_pointer_cast<StorageConv2d>(blocks["conv2"]);
+            x          = pixel_shuffle(ctx, x, 2);
+            x          = conv1->forward(ctx, x);
+            x          = ggml_gelu_erf(ctx->ggml_ctx, x);
+            x          = pixel_shuffle(ctx, x, 2);
+            x          = conv2->forward(ctx, x);
+            return pixel_shuffle(ctx, x, 8);
+        }
+    };
+
+    enum class Branch {
+        UNDERSTANDING,
+        GENERATION,
+    };
+
+    struct Attention : public GGMLBlock {
+        SenseNovaU1Config config;
+        int layer_index;
+
+        Attention(const SenseNovaU1Config& config, int layer_index)
+            : config(config), layer_index(layer_index) {
+            blocks["q_proj"]         = std::make_shared<Linear>(config.hidden_size, config.num_heads * config.head_dim, false);
+            blocks["k_proj"]         = std::make_shared<Linear>(config.hidden_size, config.num_kv_heads * config.head_dim, false);
+            blocks["v_proj"]         = std::make_shared<Linear>(config.hidden_size, config.num_kv_heads * config.head_dim, false);
+            blocks["o_proj"]         = std::make_shared<Linear>(config.num_heads * config.head_dim, config.hidden_size, false);
+            blocks["q_proj_mot_gen"] = std::make_shared<Linear>(config.hidden_size, config.num_heads * config.head_dim, false);
+            blocks["k_proj_mot_gen"] = std::make_shared<Linear>(config.hidden_size, config.num_kv_heads * config.head_dim, false);
+            blocks["v_proj_mot_gen"] = std::make_shared<Linear>(config.hidden_size, config.num_kv_heads * config.head_dim, false);
+            blocks["o_proj_mot_gen"] = std::make_shared<Linear>(config.num_heads * config.head_dim, config.hidden_size, false);
+
+            const int64_t axis_dim      = config.head_dim / 2;
+            blocks["q_norm"]            = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
+            blocks["k_norm"]            = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
+            blocks["q_norm_hw"]         = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
+            blocks["k_norm_hw"]         = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
+            blocks["q_norm_mot_gen"]    = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
+            blocks["k_norm_mot_gen"]    = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
+            blocks["q_norm_hw_mot_gen"] = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
+            blocks["k_norm_hw_mot_gen"] = std::make_shared<LLM::LLMRMSNorm>(axis_dim, config.rms_norm_eps);
+        }
+
+        ggml_tensor* apply_axis_rope(GGMLRunnerContext* ctx,
+                                     ggml_tensor* x,
+                                     ggml_tensor* positions,
+                                     int dimensions,
+                                     float theta,
+                                     int max_position) {
+            return ggml_rope_ext(ctx->ggml_ctx,
+                                 x,
+                                 positions,
+                                 nullptr,
+                                 dimensions,
+                                 GGML_ROPE_TYPE_NEOX,
+                                 max_position,
+                                 theta,
+                                 1.f,
+                                 0.f,
+                                 1.f,
+                                 32.f,
+                                 1.f);
+        }
+
+        ggml_tensor* normalize_and_rotate(GGMLRunnerContext* ctx,
+                                          ggml_tensor* x,
+                                          ggml_tensor* position_t,
+                                          ggml_tensor* position_h,
+                                          ggml_tensor* position_w,
+                                          const std::string& norm_name,
+                                          const std::string& norm_hw_name) {
+            const int64_t temporal_dim = config.head_dim / 2;
+            const int64_t spatial_dim  = config.head_dim - temporal_dim;
+            const int64_t axis_dim     = spatial_dim / 2;
+
+            auto temporal = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, temporal_dim);
+            auto spatial  = ggml_ext_slice(ctx->ggml_ctx, x, 0, temporal_dim, config.head_dim);
+            temporal      = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(blocks[norm_name])->forward(ctx, temporal);
+            spatial       = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(blocks[norm_hw_name])->forward(ctx, spatial);
+
+            auto height = ggml_ext_slice(ctx->ggml_ctx, spatial, 0, 0, axis_dim);
+            auto width  = ggml_ext_slice(ctx->ggml_ctx, spatial, 0, axis_dim, spatial_dim);
+            temporal    = apply_axis_rope(ctx,
+                                          temporal,
+                                          position_t,
+                                          static_cast<int>(temporal_dim),
+                                          config.rope_theta,
+                                          static_cast<int>(config.max_position_embeddings));
+            height      = apply_axis_rope(ctx,
+                                          height,
+                                          position_h,
+                                          static_cast<int>(axis_dim),
+                                          config.rope_theta_hw,
+                                          static_cast<int>(config.max_position_embeddings_hw));
+            width       = apply_axis_rope(ctx,
+                                          width,
+                                          position_w,
+                                          static_cast<int>(axis_dim),
+                                          config.rope_theta_hw,
+                                          static_cast<int>(config.max_position_embeddings_hw));
+            return ggml_concat(ctx->ggml_ctx,
+                               ggml_concat(ctx->ggml_ctx, temporal, height, 0),
+                               width,
+                               0);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* position_t,
+                             ggml_tensor* position_h,
+                             ggml_tensor* position_w,
+                             ggml_tensor* attention_mask,
+                             Branch branch,
+                             const std::string& cache_prefix) {
+            const bool generation    = branch == Branch::GENERATION;
+            const std::string suffix = generation ? "_mot_gen" : "";
+            auto q_proj              = std::dynamic_pointer_cast<Linear>(blocks["q_proj" + suffix]);
+            auto k_proj              = std::dynamic_pointer_cast<Linear>(blocks["k_proj" + suffix]);
+            auto v_proj              = std::dynamic_pointer_cast<Linear>(blocks["v_proj" + suffix]);
+            auto o_proj              = std::dynamic_pointer_cast<Linear>(blocks["o_proj" + suffix]);
+
+            const int64_t n_tokens = x->ne[1];
+            const int64_t batch    = x->ne[2];
+            auto q                 = ggml_reshape_4d(ctx->ggml_ctx,
+                                                     q_proj->forward(ctx, x),
+                                                     config.head_dim,
+                                                     config.num_heads,
+                                                     n_tokens,
+                                                     batch);
+            auto k                 = ggml_reshape_4d(ctx->ggml_ctx,
+                                                     k_proj->forward(ctx, x),
+                                                     config.head_dim,
+                                                     config.num_kv_heads,
+                                                     n_tokens,
+                                                     batch);
+            auto v                 = ggml_reshape_4d(ctx->ggml_ctx,
+                                                     v_proj->forward(ctx, x),
+                                                     config.head_dim,
+                                                     config.num_kv_heads,
+                                                     n_tokens,
+                                                     batch);
+
+            q = normalize_and_rotate(ctx,
+                                     q,
+                                     position_t,
+                                     position_h,
+                                     position_w,
+                                     "q_norm" + suffix,
+                                     "q_norm_hw" + suffix);
+            k = normalize_and_rotate(ctx,
+                                     k,
+                                     position_t,
+                                     position_h,
+                                     position_w,
+                                     "k_norm" + suffix,
+                                     "k_norm_hw" + suffix);
+
+            const std::string layer_cache = cache_prefix + "." + std::to_string(layer_index);
+            if (generation) {
+                auto prefix_k = ctx->load_cache_tensor(layer_cache + ".k");
+                auto prefix_v = ctx->load_cache_tensor(layer_cache + ".v");
+                GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr);
+                k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 2);
+                v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
+            } else {
+                // Keep dedicated graph outputs alive until the runner copies them
+                // into its persistent cache buffer after graph execution.
+                auto cache_k = ggml_dup_tensor(ctx->ggml_ctx, k);
+                cache_k      = ggml_cpy(ctx->ggml_ctx, k, cache_k);
+                ggml_set_output(cache_k);
+                auto cache_v = ggml_dup_tensor(ctx->ggml_ctx, v);
+                cache_v      = ggml_cpy(ctx->ggml_ctx, v, cache_v);
+                ggml_set_output(cache_v);
+                ctx->persist_cache_tensor(layer_cache + ".k", cache_k);
+                ctx->persist_cache_tensor(layer_cache + ".v", cache_v);
+            }
+
+            q = ggml_cont(ctx->ggml_ctx,
+                          ggml_ext_torch_permute(ctx->ggml_ctx, q, 0, 2, 1, 3));
+            q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);
+            k = ggml_cont(ctx->ggml_ctx,
+                          ggml_ext_torch_permute(ctx->ggml_ctx, k, 0, 2, 1, 3));
+            k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);
+
+            auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
+                                              ctx->backend,
+                                              q,
+                                              k,
+                                              v,
+                                              config.num_heads,
+                                              attention_mask,
+                                              true,
+                                              ctx->flash_attn_enabled);
+            return o_proj->forward(ctx, out);
+        }
+    };
+
+    struct TransformerBlock : public GGMLBlock {
+        TransformerBlock(const SenseNovaU1Config& config, int layer_index) {
+            blocks["self_attn"]                        = std::make_shared<Attention>(config, layer_index);
+            blocks["mlp"]                              = std::make_shared<LLM::MLP>(config.hidden_size, config.intermediate_size, false);
+            blocks["mlp_mot_gen"]                      = std::make_shared<LLM::MLP>(config.hidden_size, config.intermediate_size, false);
+            blocks["input_layernorm"]                  = std::make_shared<LLM::LLMRMSNorm>(config.hidden_size, config.rms_norm_eps);
+            blocks["input_layernorm_mot_gen"]          = std::make_shared<LLM::LLMRMSNorm>(config.hidden_size, config.rms_norm_eps);
+            blocks["post_attention_layernorm"]         = std::make_shared<LLM::LLMRMSNorm>(config.hidden_size, config.rms_norm_eps);
+            blocks["post_attention_layernorm_mot_gen"] = std::make_shared<LLM::LLMRMSNorm>(config.hidden_size, config.rms_norm_eps);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* position_t,
+                             ggml_tensor* position_h,
+                             ggml_tensor* position_w,
+                             ggml_tensor* attention_mask,
+                             Branch branch,
+                             const std::string& cache_prefix) {
+            const bool generation = branch == Branch::GENERATION;
+            auto input_norm       = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(
+                blocks[generation ? "input_layernorm_mot_gen" : "input_layernorm"]);
+            auto post_norm = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(
+                blocks[generation ? "post_attention_layernorm_mot_gen" : "post_attention_layernorm"]);
+            auto attention = std::dynamic_pointer_cast<Attention>(blocks["self_attn"]);
+            auto mlp       = std::dynamic_pointer_cast<LLM::MLP>(blocks[generation ? "mlp_mot_gen" : "mlp"]);
+
+            auto residual = x;
+            x             = input_norm->forward(ctx, x);
+            x             = attention->forward(ctx,
+                                               x,
+                                               position_t,
+                                               position_h,
+                                               position_w,
+                                               attention_mask,
+                                               branch,
+                                               cache_prefix);
+            x             = ggml_add_inplace(ctx->ggml_ctx, x, residual);
+
+            residual = x;
+            x        = post_norm->forward(ctx, x);
+            x        = mlp->forward(ctx, x);
+            return ggml_add_inplace(ctx->ggml_ctx, x, residual);
+        }
+    };
+
+    struct TextModel : public GGMLBlock {
+        SenseNovaU1Config config;
+
+        explicit TextModel(const SenseNovaU1Config& config)
+            : config(config) {
+            blocks["embed_tokens"] = std::make_shared<Embedding>(config.vocab_size, config.hidden_size);
+            for (int i = 0; i < config.num_layers; ++i) {
+                blocks["layers." + std::to_string(i)] = std::make_shared<TransformerBlock>(config, i);
+            }
+            blocks["norm"]         = std::make_shared<LLM::LLMRMSNorm>(config.hidden_size, config.rms_norm_eps);
+            blocks["norm_mot_gen"] = std::make_shared<LLM::LLMRMSNorm>(config.hidden_size, config.rms_norm_eps);
+        }
+
+        ggml_tensor* embed(GGMLRunnerContext* ctx, ggml_tensor* input_ids) {
+            return std::dynamic_pointer_cast<Embedding>(blocks["embed_tokens"])->forward(ctx, input_ids);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* position_t,
+                             ggml_tensor* position_h,
+                             ggml_tensor* position_w,
+                             ggml_tensor* attention_mask,
+                             Branch branch,
+                             const std::string& cache_prefix) {
+            for (int i = 0; i < config.num_layers; ++i) {
+                auto layer = std::dynamic_pointer_cast<TransformerBlock>(blocks["layers." + std::to_string(i)]);
+                x          = layer->forward(ctx,
+                                            x,
+                                            position_t,
+                                            position_h,
+                                            position_w,
+                                            attention_mask,
+                                            branch,
+                                            cache_prefix);
+            }
+            auto norm = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(
+                blocks[branch == Branch::GENERATION ? "norm_mot_gen" : "norm"]);
+            return norm->forward(ctx, x);
+        }
+    };
+
+    struct SenseNovaU1Model : public GGMLBlock {
+        SenseNovaU1Config config;
+
+        explicit SenseNovaU1Model(const SenseNovaU1Config& config)
+            : config(config) {
+            blocks["language_model.model"]                       = std::make_shared<TextModel>(config);
+            blocks["fm_modules.vision_model_mot_gen.embeddings"] = std::make_shared<VisionEmbeddings>(config);
+            blocks["fm_modules.timestep_embedder"]               = std::make_shared<TimestepEmbedder>(config.hidden_size,
+                                                                                        config.timestep_embedding_size);
+            if (config.add_noise_scale_embedding) {
+                blocks["fm_modules.noise_scale_embedder"] = std::make_shared<TimestepEmbedder>(config.hidden_size,
+                                                                                               config.timestep_embedding_size);
+            }
+            blocks["fm_modules.fm_head"] = std::make_shared<PixelDecoder>(config);
+        }
+
+        std::shared_ptr<TextModel> text_model() {
+            return std::dynamic_pointer_cast<TextModel>(blocks["language_model.model"]);
+        }
+
+        std::shared_ptr<VisionEmbeddings> vision_embeddings() {
+            return std::dynamic_pointer_cast<VisionEmbeddings>(blocks["fm_modules.vision_model_mot_gen.embeddings"]);
+        }
+
+        std::shared_ptr<TimestepEmbedder> timestep_embedder() {
+            return std::dynamic_pointer_cast<TimestepEmbedder>(blocks["fm_modules.timestep_embedder"]);
+        }
+
+        std::shared_ptr<TimestepEmbedder> noise_scale_embedder() {
+            if (!config.add_noise_scale_embedding) {
+                return nullptr;
+            }
+            return std::dynamic_pointer_cast<TimestepEmbedder>(blocks["fm_modules.noise_scale_embedder"]);
+        }
+
+        std::shared_ptr<PixelDecoder> pixel_decoder() {
+            return std::dynamic_pointer_cast<PixelDecoder>(blocks["fm_modules.fm_head"]);
+        }
+    };
+
+    struct SenseNovaU1Runner : public DiffusionModelRunner {
+        SenseNovaU1Config config;
+        SenseNovaU1Model model;
+        std::unordered_set<uint64_t> cached_prefix_hashes;
+        std::vector<int32_t> position_t_vec;
+        std::vector<int32_t> position_h_vec;
+        std::vector<int32_t> position_w_vec;
+        std::vector<float> attention_mask_vec;
+        std::vector<float> noise_scale_vec;
+
+        SenseNovaU1Runner(ggml_backend_t backend,
+                          const String2TensorStorage& tensor_storage_map      = {},
+                          const std::string& prefix                           = "",
+                          std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(SenseNovaU1Config::detect_from_weights(tensor_storage_map, prefix)),
+              model(config) {
+            model.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "SenseNova U1.5";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors,
+                               const std::string& prefix) override {
+            model.get_param_tensors(tensors, prefix);
+        }
+
+        static uint64_t hash_input_ids(const sd::Tensor<int32_t>& input_ids) {
+            uint64_t hash = 1469598103934665603ULL;
+            for (int32_t token : input_ids.values()) {
+                uint32_t value = static_cast<uint32_t>(token);
+                for (int byte = 0; byte < 4; ++byte) {
+                    hash ^= static_cast<uint8_t>(value & 0xffU);
+                    hash *= 1099511628211ULL;
+                    value >>= 8;
+                }
+            }
+            hash ^= static_cast<uint64_t>(input_ids.numel());
+            hash *= 1099511628211ULL;
+            return hash;
+        }
+
+        static std::string cache_prefix(uint64_t hash) {
+            return "snu15." + std::to_string(hash);
+        }
+
+        ggml_tensor* make_position_tensor(const std::vector<int32_t>& values,
+                                          const std::string& name) {
+            auto tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, values.size());
+            ggml_set_name(tensor, name.c_str());
+            set_backend_tensor_data(tensor, values.data());
+            return tensor;
+        }
+
+        ggml_cgraph* build_prefix_graph(const sd::Tensor<int32_t>& input_ids_tensor,
+                                        const std::string& prefix_cache) {
+            ggml_cgraph* graph   = new_graph_custom(SENSENOVA_U1_GRAPH_SIZE);
+            ggml_tensor* ids     = make_input(input_ids_tensor);
+            const int64_t length = input_ids_tensor.numel();
+
+            position_t_vec.resize(length);
+            position_h_vec.assign(length, 0);
+            position_w_vec.assign(length, 0);
+            for (int64_t i = 0; i < length; ++i) {
+                position_t_vec[i] = static_cast<int32_t>(i);
+            }
+            auto position_t = make_position_tensor(position_t_vec, "snu15.prefix.position_t");
+            auto position_h = make_position_tensor(position_h_vec, "snu15.prefix.position_h");
+            auto position_w = make_position_tensor(position_w_vec, "snu15.prefix.position_w");
+
+            attention_mask_vec.assign(static_cast<size_t>(length * length), 0.f);
+            for (int64_t query = 0; query < length; ++query) {
+                for (int64_t key = query + 1; key < length; ++key) {
+                    attention_mask_vec[static_cast<size_t>(query * length + key)] = -INFINITY;
+                }
+            }
+            auto attention_mask = ggml_new_tensor_2d(compute_ctx,
+                                                     GGML_TYPE_F32,
+                                                     length,
+                                                     length);
+            ggml_set_name(attention_mask, "snu15.prefix.attention_mask");
+            set_backend_tensor_data(attention_mask, attention_mask_vec.data());
+
+            auto runner_ctx = get_context();
+            auto text_model = model.text_model();
+            auto hidden     = text_model->embed(&runner_ctx, ids);
+            hidden          = text_model->forward(&runner_ctx,
+                                                  hidden,
+                                                  position_t,
+                                                  position_h,
+                                                  position_w,
+                                                  attention_mask,
+                                                  Branch::UNDERSTANDING,
+                                                  prefix_cache);
+            ggml_build_forward_expand(graph, hidden);
+            return graph;
+        }
+
+        bool ensure_prefix_cache(int n_threads,
+                                 const sd::Tensor<int32_t>& input_ids,
+                                 std::string* prefix_cache) {
+            const uint64_t hash = hash_input_ids(input_ids);
+            *prefix_cache       = cache_prefix(hash);
+            if (cached_prefix_hashes.find(hash) != cached_prefix_hashes.end() &&
+                get_cache_tensor_by_name(*prefix_cache + ".0.k") != nullptr) {
+                return true;
+            }
+
+            if (cached_prefix_hashes.size() >= 2) {
+                free_cache_ctx_and_buffer();
+                cached_prefix_hashes.clear();
+            }
+            auto get_graph = [&]() {
+                return build_prefix_graph(input_ids, *prefix_cache);
+            };
+            auto result = GGMLRunner::compute(get_graph, n_threads, false, true);
+            if (!result.has_value()) {
+                LOG_ERROR("SenseNova U1.5 prefix cache computation failed");
+                return false;
+            }
+            cached_prefix_hashes.insert(hash);
+            return true;
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timestep_tensor,
+                                 const std::string& prefix_cache,
+                                 int64_t prefix_length) {
+            ggml_cgraph* graph = new_graph_custom(SENSENOVA_U1_GRAPH_SIZE);
+            ggml_tensor* x     = make_input(x_tensor);
+            ggml_tensor* t     = make_input(timestep_tensor);
+            GGML_ASSERT(x->ne[3] == 1);
+            GGML_ASSERT(x->ne[0] % config.image_token_stride() == 0);
+            GGML_ASSERT(x->ne[1] % config.image_token_stride() == 0);
+
+            const int64_t grid_w  = x->ne[0] / config.patch_size;
+            const int64_t grid_h  = x->ne[1] / config.patch_size;
+            const int64_t token_w = grid_w / config.vision_downsample_factor;
+            const int64_t token_h = grid_h / config.vision_downsample_factor;
+            const int64_t tokens  = token_w * token_h;
+
+            position_h_vec.resize(grid_w * grid_h);
+            position_w_vec.resize(grid_w * grid_h);
+            for (int64_t index = 0; index < grid_w * grid_h; ++index) {
+                position_h_vec[index] = static_cast<int32_t>(index / grid_w);
+                position_w_vec[index] = static_cast<int32_t>(index % grid_w);
+            }
+            auto vision_position_x = make_position_tensor(position_w_vec, "snu15.vision.position_x");
+            auto vision_position_y = make_position_tensor(position_h_vec, "snu15.vision.position_y");
+
+            auto runner_ctx     = get_context();
+            auto hidden         = model.vision_embeddings()->forward(&runner_ctx,
+                                                                     x,
+                                                                     vision_position_x,
+                                                                     vision_position_y);
+            auto time_embedding = model.timestep_embedder()->forward(&runner_ctx, t);
+            time_embedding      = ggml_reshape_3d(compute_ctx, time_embedding, config.hidden_size, 1, 1);
+            hidden              = ggml_add(compute_ctx, hidden, time_embedding);
+
+            if (config.add_noise_scale_embedding) {
+                const float image_tokens = static_cast<float>(tokens);
+                const float noise_scale  = std::min(config.noise_scale_max_value,
+                                                    std::sqrt(image_tokens / config.noise_scale_base_image_seq_len));
+                noise_scale_vec          = {noise_scale / config.noise_scale_max_value};
+                auto noise_scale_tensor  = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, 1);
+                ggml_set_name(noise_scale_tensor, "snu15.noise_scale");
+                set_backend_tensor_data(noise_scale_tensor, noise_scale_vec.data());
+                auto noise_embedding = model.noise_scale_embedder()->forward(&runner_ctx, noise_scale_tensor);
+                noise_embedding      = ggml_reshape_3d(compute_ctx, noise_embedding, config.hidden_size, 1, 1);
+                hidden               = ggml_add(compute_ctx, hidden, noise_embedding);
+            }
+
+            position_t_vec.assign(tokens, static_cast<int32_t>(prefix_length));
+            position_h_vec.resize(tokens);
+            position_w_vec.resize(tokens);
+            for (int64_t index = 0; index < tokens; ++index) {
+                position_h_vec[index] = static_cast<int32_t>(index / token_w);
+                position_w_vec[index] = static_cast<int32_t>(index % token_w);
+            }
+            auto position_t = make_position_tensor(position_t_vec, "snu15.image.position_t");
+            auto position_h = make_position_tensor(position_h_vec, "snu15.image.position_h");
+            auto position_w = make_position_tensor(position_w_vec, "snu15.image.position_w");
+
+            hidden            = model.text_model()->forward(&runner_ctx,
+                                                            hidden,
+                                                            position_t,
+                                                            position_h,
+                                                            position_w,
+                                                            nullptr,
+                                                            Branch::GENERATION,
+                                                            prefix_cache);
+            hidden            = ggml_reshape_4d(compute_ctx,
+                                                hidden,
+                                                config.hidden_size,
+                                                token_w,
+                                                token_h,
+                                                x->ne[3]);
+            hidden            = ggml_cont(compute_ctx, ggml_permute(compute_ctx, hidden, 2, 0, 1, 3));
+            auto x_prediction = model.pixel_decoder()->forward(&runner_ctx, hidden);
+
+            const float timestep = timestep_tensor.values()[0];
+            const float denom    = std::max(1.f - timestep, config.t_eps);
+            auto velocity        = ggml_scale(compute_ctx,
+                                              ggml_sub(compute_ctx, x_prediction, x),
+                                              1.f / denom);
+            ggml_build_forward_expand(graph, velocity);
+            return graph;
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timestep,
+                                  const sd::Tensor<int32_t>& input_ids) {
+            std::string prefix_cache;
+            if (!ensure_prefix_cache(n_threads, input_ids, &prefix_cache)) {
+                return {};
+            }
+            auto get_graph = [&]() {
+                return build_graph(x, timestep, prefix_cache, input_ids.numel());
+            };
+            return restore_trailing_singleton_dims(
+                GGMLRunner::compute(get_graph, n_threads, false),
+                x.dim());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            const auto* extra = diffusion_extra_as<SenseNovaU1DiffusionExtra>(diffusion_params);
+            GGML_ASSERT(extra->input_ids != nullptr);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           *extra->input_ids);
+        }
+    };
+}  // namespace SenseNovaU1
+
+#endif  // __SD_MODEL_DIFFUSION_SENSENOVA_U1_H__

src/model/diffusion/wan.hpp

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

src/model/diffusion/wan_audio.hpp

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

src/model/diffusion/z_image.hpp

@@ -706,8 +706,8 @@ namespace ZImage {
             ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_Q8_0;
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -728,7 +728,8 @@ namespace ZImage {
                                                                                    VERSION_QWEN_IMAGE,
                                                                                    model_manager);
 
-            if (!model_manager->register_runner_params("ZImage test",
+            if (!model_manager->set_loader(model_loader) ||
+                !model_manager->register_runner_params(ModelComponent::Diffusion,
                                                        *z_image,
                                                        "model.diffusion_model",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/te/clip.hpp

@@ -145,6 +145,10 @@ protected:
         params["position_embedding.weight"] = ggml_new_tensor_2d(ctx, position_wtype, embed_dim, num_positions);
     }
 
+    enum ggml_op param_usage_op(const std::string& name) const override {
+        return name == "token_embedding.weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
+    }
+
 public:
     CLIPEmbeddings(int64_t embed_dim,
                    int64_t vocab_size    = 49408,

src/model/te/llm.hpp

@@ -139,7 +139,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 +231,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 +279,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 +333,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 +1376,7 @@ namespace LLM {
                 x        = ggml_ext_cont(ctx->ggml_ctx, kqv);
                 x        = ggml_reshape_3d(ctx->ggml_ctx, x, head_dim * num_heads, n_token, N);
             } else {
-                x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, false);  // [N, n_token, hidden_size]
+                x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, false);  // [N, n_token, hidden_size]
             }
 
             x = out_proj->forward(ctx, x);  // [N, n_token, hidden_size]
@@ -1886,12 +1903,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) {
@@ -2389,7 +2402,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);
             }
@@ -2576,8 +2592,8 @@ namespace LLM {
             ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_COUNT;
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file_and_convert_name(file_path, "text_encoders.llm.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -2601,7 +2617,8 @@ namespace LLM {
                                                                              true,
                                                                              model_manager);
 
-            if (!model_manager->register_runner_params("LLM test",
+            if (!model_manager->set_loader(model_loader) ||
+                !model_manager->register_runner_params(ModelComponent::Conditioner,
                                                        *llm,
                                                        "text_encoders.llm",
                                                        ModelManager::ResidencyMode::ParamBackend,

src/model/te/t5.hpp

@@ -251,7 +251,7 @@ public:
 
         k = ggml_ext_scale(ctx->ggml_ctx, k, ::sqrtf(static_cast<float>(d_head)), true);
 
-        x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask);  // [N, n_token, d_head * n_head]
+        x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask);  // [N, n_token, d_head * n_head]
 
         x = out_proj->forward(ctx, x);  // [N, n_token, model_dim]
         return {x, past_bias};
@@ -567,7 +567,10 @@ struct T5Embedder {
         for (const auto& item : parsed_attention) {
             const std::string& curr_text = item.first;
             float curr_weight            = item.second;
-            std::vector<int> curr_tokens = tokenizer.encode(curr_text);
+            std::vector<int> curr_tokens;
+            if (!tokenizer.encode(curr_text, curr_tokens)) {
+                return {};
+            }
             tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
             weights.insert(weights.end(), curr_tokens.size(), curr_weight);
         }
@@ -631,8 +634,8 @@ struct T5Embedder {
         ggml_backend_t backend    = sd_backend_cpu_init();
         ggml_type model_data_type = GGML_TYPE_F16;
 
-        auto model_manager        = std::make_shared<ModelManager>();
-        ModelLoader& model_loader = model_manager->loader();
+        auto model_manager = std::make_shared<ModelManager>();
+        ModelLoader model_loader;
         if (!model_loader.init_from_file_and_convert_name(file_path)) {
             LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
             return;
@@ -647,7 +650,8 @@ struct T5Embedder {
 
         std::shared_ptr<T5Embedder> t5 = std::make_shared<T5Embedder>(backend, tensor_storage_map, "", true, model_manager);
 
-        if (!model_manager->register_runner_params("T5 test",
+        if (!model_manager->set_loader(model_loader) ||
+            !model_manager->register_runner_params(ModelComponent::Conditioner,
                                                    *t5,
                                                    "",
                                                    ModelManager::ResidencyMode::ParamBackend,

src/model/upscaler/ltx_latent_upscaler.hpp

@@ -433,12 +433,15 @@ namespace LTXVUpsampler {
     struct LatentUpsamplerRunner : public GGMLRunner {
         LatentUpsamplerConfig config;
         std::unique_ptr<LatentUpsampler> model;
+        std::string weight_prefix;
 
         LatentUpsamplerRunner(ggml_backend_t backend,
                               const String2TensorStorage& tensor_storage_map,
+                              const std::string& prefix                           = "",
                               std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
             : GGMLRunner(backend, weight_manager),
-              config(LatentUpsamplerConfig::detect_from_weights(tensor_storage_map)) {
+              config(LatentUpsamplerConfig::detect_from_weights(tensor_storage_map, prefix)),
+              weight_prefix(prefix) {
             if (config.dims != 3 || (!config.spatial_upsample && !config.temporal_upsample) ||
                 config.spatial_up_num < 1 || config.spatial_down_den < 1 || config.temporal_up_factor < 1) {
                 LOG_ERROR("unsupported LTX latent upsampler config: dims=%d spatial=%d temporal=%d rational=%d scale=%.3f temporal_factor=%d",
@@ -452,7 +455,7 @@ namespace LTXVUpsampler {
             }
 
             model = std::make_unique<LatentUpsampler>(config);
-            model->init(params_ctx, tensor_storage_map, "");
+            model->init(params_ctx, tensor_storage_map, prefix);
         }
 
         std::string get_desc() override {
@@ -461,7 +464,7 @@ namespace LTXVUpsampler {
 
         void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
             if (model) {
-                model->get_param_tensors(tensors);
+                model->get_param_tensors(tensors, weight_prefix);
             }
         }
 

src/model/vae/auto_encoder_kl.hpp

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

src/model/vae/hunyuan_vae.hpp

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

src/model/vae/ltx_audio_vae.hpp

@@ -1078,8 +1078,8 @@ namespace LTXV {
             // ggml_backend_t backend = ggml_backend_cuda_init(0);
             LOG_INFO("loading ltx audio vae from '%s'", model_path.c_str());
 
-            auto model_manager        = std::make_shared<ModelManager>();
-            ModelLoader& model_loader = model_manager->loader();
+            auto model_manager = std::make_shared<ModelManager>();
+            ModelLoader model_loader;
             if (!model_loader.init_from_file(model_path)) {
                 LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
                 return;
@@ -1091,7 +1091,8 @@ namespace LTXV {
                                                                      prefix,
                                                                      model_manager);
 
-            if (!model_manager->register_runner_params("LTX audio VAE test",
+            if (!model_manager->set_loader(std::move(model_loader)) ||
+                !model_manager->register_runner_params(ModelComponent::AudioVAE,
                                                        *ltx_audio_vae,
                                                        ModelManager::ResidencyMode::ParamBackend,
                                                        backend,

src/model/vae/ltx_vae.hpp

@@ -1476,8 +1476,8 @@ struct LTXVideoVAE : public VAE {
         ggml_backend_t backend = sd_backend_cpu_init();
         LOG_INFO("loading ltx vae from '%s'", model_path.c_str());
 
-        auto model_manager        = std::make_shared<ModelManager>();
-        ModelLoader& model_loader = model_manager->loader();
+        auto model_manager = std::make_shared<ModelManager>();
+        ModelLoader model_loader;
         if (!model_loader.init_from_file_and_convert_name(model_path, "vae.")) {
             LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
             return;
@@ -1491,7 +1491,8 @@ struct LTXVideoVAE : public VAE {
                                                                          VERSION_LTXAV,
                                                                          model_manager);
 
-        if (!model_manager->register_runner_params("LTX VAE test",
+        if (!model_manager->set_loader(model_loader) ||
+            !model_manager->register_runner_params(ModelComponent::VAE,
                                                    *vae,
                                                    ModelManager::ResidencyMode::ParamBackend,
                                                    backend,

src/model/vae/mage_vae.hpp

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

src/model/vae/minimax_h3_audio_vae.hpp

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

src/model/vae/minimax_h3_vae.hpp

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

src/model/vae/tae.hpp

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

src/model/vae/vae.hpp

@@ -166,7 +166,7 @@ public:
             scale_factor = 16;
         } else if (sd_version_uses_flux2_vae(version)) {
             scale_factor = 16;
-        } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) {
+        } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version) || sd_version_is_sensenova_u1(version)) {
             scale_factor = 1;
         }
         return scale_factor;

src/model/vae/wan_vae.hpp

@@ -615,8 +615,8 @@ namespace WAN {
             auto v = qkv_vec[2];
             v      = ggml_reshape_3d(ctx->ggml_ctx, v, h * w, c, n);  // [t, c, h * w]
 
-            v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3));                            // [t, h * w, c]
-            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);  // [t, h * w, c]
+            v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3));    // [t, h * w, c]
+            x = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);  // [t, h * w, c]
 
             x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [t, c, h * w]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, c, n);                             // [t, c, h, w]
@@ -1494,13 +1494,14 @@ namespace WAN {
             {
                 LOG_INFO("loading from '%s'", file_path.c_str());
 
-                ModelLoader& model_loader = model_manager->loader();
+                ModelLoader model_loader;
                 if (!model_loader.init_from_file_and_convert_name(file_path, "vae.")) {
                     LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                     return;
                 }
 
-                if (!model_manager->register_runner_params("Wan VAE test",
+                if (!model_manager->set_loader(model_loader) ||
+                    !model_manager->register_runner_params(ModelComponent::VAE,
                                                            *vae,
                                                            ModelManager::ResidencyMode::ParamBackend,
                                                            backend,

src/model_component.h

@@ -0,0 +1,64 @@
+#ifndef __SD_MODEL_COMPONENT_H__
+#define __SD_MODEL_COMPONENT_H__
+
+enum class ModelComponent {
+    Conditioner,
+    Diffusion,
+    HighNoiseDiffusion,
+    CLIPVision,
+    IPAdapter,
+    VAE,
+    PreviewVAE,
+    AudioVAE,
+    AudioEncoder,
+    ControlNet,
+    PhotoMaker,
+    PuLID,
+    LoRA,
+    Upscaler,
+    Detector,
+    LatentUpsampler,
+    Count,
+};
+
+inline const char* model_component_name(ModelComponent component) {
+    switch (component) {
+        case ModelComponent::Conditioner:
+            return "Conditioner model";
+        case ModelComponent::Diffusion:
+            return "Diffusion model";
+        case ModelComponent::HighNoiseDiffusion:
+            return "High noise diffusion model";
+        case ModelComponent::CLIPVision:
+            return "CLIP vision";
+        case ModelComponent::IPAdapter:
+            return "IP-Adapter";
+        case ModelComponent::VAE:
+            return "VAE";
+        case ModelComponent::PreviewVAE:
+            return "preview VAE";
+        case ModelComponent::AudioVAE:
+            return "audio VAE";
+        case ModelComponent::AudioEncoder:
+            return "audio encoder";
+        case ModelComponent::ControlNet:
+            return "ControlNet";
+        case ModelComponent::PhotoMaker:
+            return "photomaker";
+        case ModelComponent::PuLID:
+            return "pulid";
+        case ModelComponent::LoRA:
+            return "LoRA";
+        case ModelComponent::Upscaler:
+            return "ESRGAN";
+        case ModelComponent::Detector:
+            return "YOLOv8";
+        case ModelComponent::LatentUpsampler:
+            return "LTX latent upsampler";
+        case ModelComponent::Count:
+            break;
+    }
+    return "unknown";
+}
+
+#endif  // __SD_MODEL_COMPONENT_H__

src/model_io/safetensors_io.cpp

@@ -144,7 +144,8 @@ static bool read_comfy_quant_config(std::ifstream& file,
 bool read_safetensors_file(const std::string& file_path,
                            std::vector<TensorStorage>& tensor_storages,
                            std::string* error,
-                           std::map<std::string, std::string>* metadata) {
+                           std::map<std::string, std::string>* metadata,
+                           std::set<std::string>* tensor_names) {
     std::ifstream file(file_path, std::ios::binary);
     if (!file.is_open()) {
         set_error(error, "failed to open '" + file_path + "'");
@@ -246,7 +247,8 @@ bool read_safetensors_file(const std::string& file_path,
         std::string dtype    = tensor_info["dtype"];
         nlohmann::json shape = tensor_info["shape"];
 
-        if (dtype == "U8") {
+        // ComfyUI FP8 activation scales cancel when inference uses F16/F32 activations.
+        if (ends_with(name, ".scale_input")) {
             continue;
         }
 
@@ -257,6 +259,26 @@ bool read_safetensors_file(const std::string& file_path,
             return false;
         }
 
+        if (tensor_names != nullptr) {
+            tensor_names->insert(name);
+        }
+        if (dtype == "U8") {
+            uint64_t bytes = 1;
+            for (const auto& dimension : shape) {
+                const int64_t size = dimension.get<int64_t>();
+                if (size < 0 || (bytes != 0 && static_cast<uint64_t>(size) > UINT64_MAX / bytes)) {
+                    set_error(error, "invalid dimensions for tensor '" + name + "'");
+                    return false;
+                }
+                bytes *= size;
+            }
+            if (bytes != end - begin) {
+                set_error(error, "size mismatch for tensor '" + name + "'");
+                return false;
+            }
+            continue;
+        }
+
         ggml_type type = safetensors_dtype_to_ggml_type(dtype);
         if (type == GGML_TYPE_COUNT) {
             set_error(error, "unsupported dtype '" + dtype + "' (tensor '" + name + "')");
@@ -270,8 +292,20 @@ bool read_safetensors_file(const std::string& file_path,
 
         int n_dims              = (int)shape.size();
         int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+        uint64_t elements       = 1;
         for (int i = 0; i < n_dims; i++) {
             ne[i] = shape[i].get<int64_t>();
+            if (ne[i] < 0 || (elements != 0 && static_cast<uint64_t>(ne[i]) > INT64_MAX / elements)) {
+                set_error(error, "invalid dimensions for tensor '" + name + "'");
+                return false;
+            }
+            elements *= ne[i];
+        }
+        const uint64_t storage_size = ggml_type_size(type) * ((dtype == "F64" || dtype == "I64") ? 2 : 1);
+        if (elements % ggml_blck_size(type) != 0 ||
+            elements / ggml_blck_size(type) > INT64_MAX / storage_size) {
+            set_error(error, "invalid storage size for tensor '" + name + "'");
+            return false;
         }
 
         if (n_dims == 5) {

src/model_io/safetensors_io.h

@@ -2,6 +2,7 @@
 #define __SD_MODEL_IO_SAFETENSORS_IO_H__
 
 #include <map>
+#include <set>
 #include <string>
 #include <vector>
 
@@ -12,7 +13,8 @@ bool is_safetensors_file(const std::string& file_path);
 bool read_safetensors_file(const std::string& file_path,
                            std::vector<TensorStorage>& tensor_storages,
                            std::string* error                           = nullptr,
-                           std::map<std::string, std::string>* metadata = nullptr);
+                           std::map<std::string, std::string>* metadata = nullptr,
+                           std::set<std::string>* tensor_names          = nullptr);
 bool read_safetensors_index_file(const std::string& file_path,
                                  std::vector<std::string>& shard_paths,
                                  std::string* error = nullptr);

src/model_io/tensor_storage.h

@@ -28,9 +28,11 @@ struct TensorStorage {
     int n_dims                  = 0;
 
     std::string storage_key;
-    size_t file_index = 0;
-    int index_in_zip  = -1;  // >= means stored in a zip file
-    uint64_t offset   = 0;   // offset in file
+    size_t file_index      = 0;
+    uint64_t file_id       = 0;
+    uint64_t file_revision = 0;
+    int index_in_zip       = -1;  // >= means stored in a zip file
+    uint64_t offset        = 0;   // offset in file
 
     TensorStorage() = default;
 

src/model_loader.cpp

@@ -9,6 +9,7 @@
 #include <mutex>
 #include <regex>
 #include <set>
+#include <stdexcept>
 #include <string>
 #include <thread>
 #include <unordered_map>
@@ -27,6 +28,7 @@
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 #include "ggml.h"
+#include "json.hpp"
 #include "zip.h"
 
 #include "name_conversion.h"
@@ -67,6 +69,8 @@ const char* unused_tensors[] = {
     // "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) {
@@ -151,15 +155,19 @@ ModelLoader::ModelLoader()
 }
 
 size_t ModelLoader::add_file_path(const std::string& file_path) {
-    if (model_files_processed) {
-        file_data.clear();
-        model_files_processed = false;
+    auto it = std::find(file_paths_.begin(), file_paths_.end(), file_path);
+    if (it != file_paths_.end()) {
+        return static_cast<size_t>(it - file_paths_.begin());
     }
+    invalidate_file_data();
     file_paths_.push_back(file_path);
     return file_paths_.size() - 1;
 }
 
 void ModelLoader::add_tensor_storage(const TensorStorage& tensor_storage) {
+    if (tensor_storage_map.count(tensor_storage.name) != 0) {
+        throw std::runtime_error("duplicate tensor in model source: " + tensor_storage.name);
+    }
     tensor_storage_map[tensor_storage.name] = tensor_storage;
 }
 
@@ -169,7 +177,25 @@ void ModelLoader::set_n_threads(int n_threads) {
 }
 
 bool ModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
+    return add_file(file_path, prefix);
+}
+
+bool ModelLoader::parse_file(const std::string& file_path, const std::string& prefix) {
+    FileStamp stamp;
+    if (!read_file_stamp(file_path, stamp)) {
+        return false;
+    }
+    parsed_dependencies_.push_back(stamp);
     if (is_directory(file_path)) {
+        const std::string diffusers_index_path = path_join(file_path, "model_index.json");
+        const std::string diffusers_unet_path  = path_join(file_path, "unet/diffusion_pytorch_model.safetensors");
+        const bool has_diffusers_layout        = file_exists(diffusers_index_path) || file_exists(diffusers_unet_path);
+
+        const std::string safetensors_index_path = path_join(file_path, "model.safetensors.index.json");
+        if (!has_diffusers_layout && file_exists(safetensors_index_path)) {
+            LOG_INFO("load %s using root safetensors index", file_path.c_str());
+            return parse_file(safetensors_index_path, prefix);
+        }
         LOG_INFO("load %s using diffusers format", file_path.c_str());
         return init_from_diffusers_file(file_path, prefix);
     } else if (is_gguf_file(file_path)) {
@@ -198,17 +224,11 @@ bool ModelLoader::init_from_file(const std::string& file_path, const std::string
 }
 
 void ModelLoader::convert_tensors_name() {
-    SDVersion version = (version_ == VERSION_COUNT) ? get_sd_version() : version_;
-    String2TensorStorage new_map;
-
-    for (auto& [_, tensor_storage] : tensor_storage_map) {
-        auto new_name = convert_tensor_name(tensor_storage.name, version);
-        // LOG_VERBOSE("%s -> %s", tensor_storage.name.c_str(), new_name.c_str());
-        tensor_storage.name = new_name;
-        new_map[new_name]   = std::move(tensor_storage);
+    if (names_converted_) {
+        return;
     }
-
-    tensor_storage_map.swap(new_map);
+    names_converted_ = true;
+    rebuild_catalog();
 }
 
 bool ModelLoader::init_from_file_and_convert_name(const std::string& file_path, const std::string& prefix, SDVersion version) {
@@ -257,7 +277,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
 
     std::vector<TensorStorage> tensor_storages;
     std::string error;
-    if (!read_safetensors_file(file_path, tensor_storages, &error, &metadata_)) {
+    if (!read_safetensors_file(file_path, tensor_storages, &error, &metadata_, &parsed_tensor_names_[file_path])) {
         LOG_ERROR("%s", error.c_str());
         return false;
     }
@@ -293,7 +313,26 @@ bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path,
     }
 
     for (const std::string& shard_path : shard_paths) {
-        if (!init_from_file(shard_path, prefix)) {
+        if (!parse_file(shard_path, prefix)) {
+            return false;
+        }
+    }
+
+    std::ifstream index_file(file_path);
+    const auto index = nlohmann::json::parse(index_file);
+    for (const auto& entry : index.at("weight_map").items()) {
+        const auto expected = (std::filesystem::u8path(file_path).parent_path() /
+                               std::filesystem::u8path(entry.value().get<std::string>()))
+                                  .lexically_normal();
+        bool found = false;
+        for (const auto& shard : parsed_tensor_names_) {
+            if (std::filesystem::u8path(shard.first).lexically_normal() == expected) {
+                found = shard.second.count(entry.key()) != 0;
+                break;
+            }
+        }
+        if (!found) {
+            LOG_ERROR("safetensors index tensor '%s' is missing from its declared shard", entry.key().c_str());
             return false;
         }
     }
@@ -369,25 +408,23 @@ bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const s
     std::string clip_path   = path_join(file_path, "text_encoder/model.safetensors");
     std::string clip_g_path = path_join(file_path, "text_encoder_2/model.safetensors");
 
-    if (!init_from_safetensors_file(unet_path, "unet.")) {
+    if (!parse_file(unet_path, prefix + "unet.")) {
         return false;
     }
 
-    if (!init_from_safetensors_file(vae_path, "vae.")) {
-        LOG_WARN("Couldn't find working VAE in %s", file_path.c_str());
-        // return false;
+    if (file_exists(vae_path) && !parse_file(vae_path, prefix + "vae.")) {
+        return false;
     }
-    if (!init_from_safetensors_file(clip_path, "te.")) {
-        LOG_WARN("Couldn't find working text encoder in %s", file_path.c_str());
-        // return false;
+    if (file_exists(clip_path) && !parse_file(clip_path, prefix + "te.")) {
+        return false;
     }
-    if (!init_from_safetensors_file(clip_g_path, "te.1.")) {
-        LOG_VERBOSE("Couldn't find working second text encoder in %s", file_path.c_str());
+    if (file_exists(clip_g_path) && !parse_file(clip_g_path, prefix + "te.1.")) {
+        return false;
     }
     return true;
 }
 
-SDVersion ModelLoader::get_sd_version() {
+SDVersion ModelLoader::get_sd_version() const {
     TensorStorage token_embedding_weight, input_block_weight, context_ebedding_weight;
 
     bool has_multiple_encoders = false;
@@ -398,6 +435,7 @@ SDVersion ModelLoader::get_sd_version() {
     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;
@@ -437,6 +475,9 @@ SDVersion ModelLoader::get_sd_version() {
         if (tensor_storage.name.find("net.img_embedder.proj1.weight") != std::string::npos) {
             return VERSION_MINIT2I;
         }
+        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.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) {
@@ -484,6 +525,11 @@ SDVersion ModelLoader::get_sd_version() {
         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;
         }
@@ -547,6 +593,9 @@ SDVersion ModelLoader::get_sd_version() {
     }
     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;
         }
@@ -623,7 +672,7 @@ SDVersion ModelLoader::get_sd_version() {
     return VERSION_COUNT;
 }
 
-std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() {
+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)) {
@@ -640,7 +689,7 @@ std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() {
     return wtype_stat;
 }
 
-std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() {
+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)) {
@@ -664,7 +713,7 @@ std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() {
     return wtype_stat;
 }
 
-std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() {
+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)) {
@@ -685,7 +734,7 @@ std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() {
     return wtype_stat;
 }
 
-std::map<ggml_type, uint32_t> ModelLoader::get_vae_wtype_stat() {
+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)) {
@@ -743,9 +792,12 @@ TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules) {
 }
 
 void ModelLoader::set_wtype_override(ggml_type wtype, std::string tensor_type_rules) {
-    auto map_rules = parse_tensor_type_rules(tensor_type_rules);
+    wtype_override_    = wtype;
+    tensor_type_rules_ = tensor_type_rules;
+    auto map_rules     = parse_tensor_type_rules(tensor_type_rules);
     for (auto& [name, tensor_storage] : tensor_storage_map) {
-        ggml_type dst_type = wtype;
+        tensor_storage.expected_type = GGML_TYPE_COUNT;
+        ggml_type dst_type           = wtype;
         for (const auto& tensor_type_rule : map_rules) {
             std::regex pattern(tensor_type_rule.first);
             if (std::regex_search(name, pattern)) {
@@ -761,6 +813,8 @@ void ModelLoader::set_wtype_override(ggml_type wtype, std::string tensor_type_ru
         }
         tensor_storage.expected_type = dst_type;
     }
+    invalidate_file_data();
+    ++revision_;
 }
 
 void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
@@ -829,6 +883,13 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
 std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
                                                        std::set<std::string> ignore_tensors,
                                                        bool writable_mmap) {
+    std::set<std::string> names;
+    for (const auto& entry : tensors) {
+        names.insert(entry.first);
+    }
+    if (!validate_sources(&names)) {
+        return {};
+    }
     process_model_files(true, writable_mmap);
 
     std::vector<MmapTensorStore> result;
@@ -919,6 +980,9 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                                bool enable_mmap,
                                const std::set<std::string>* target_tensor_names,
                                bool log_progress) {
+    if (!validate_sources(target_tensor_names)) {
+        return false;
+    }
     process_model_files(enable_mmap, false);
 
     std::atomic<int64_t> read_time_ms(0);
@@ -1242,7 +1306,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                  (convert_time_ms.load() / (float)last_n_threads) / 1000.f,
                  (copy_to_backend_time_ms.load() / (float)last_n_threads) / 1000.f);
     }
-    return success;
+    return success && validate_sources(target_tensor_names);
 }
 
 bool ModelLoader::load_tensor(const TensorStorage& tensor_storage, ggml_tensor* dst_tensor) {
@@ -1259,7 +1323,9 @@ bool ModelLoader::load_tensor(const TensorStorage& tensor_storage, ggml_tensor*
             return true;
         }
 
-        if (current_tensor_storage.file_index != tensor_storage.file_index ||
+        if (current_tensor_storage.file_id != tensor_storage.file_id ||
+            current_tensor_storage.file_revision != tensor_storage.file_revision ||
+            current_tensor_storage.file_index != tensor_storage.file_index ||
             current_tensor_storage.offset != tensor_storage.offset ||
             current_tensor_storage.index_in_zip != tensor_storage.index_in_zip) {
             LOG_ERROR("load tensor failed: storage mismatch for '%s'", tensor_storage.name.c_str());
@@ -1440,7 +1506,7 @@ bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
     return true;
 }
 
-bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) {
+bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) const {
     const std::string& name = tensor_storage.name;
     if (tensor_storage.is_int8_tensorwise) {
         return false;
@@ -1478,7 +1544,7 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage
     return false;
 }
 
-int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type) {
+int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type) const {
     size_t alignment = 128;
     if (backend != nullptr) {
         alignment = ggml_backend_get_alignment(backend);

src/model_loader.h

@@ -2,6 +2,7 @@
 #define __MODEL_LOADER_H__
 
 #include <cstdint>
+#include <filesystem>
 #include <map>
 #include <memory>
 #include <set>
@@ -30,6 +31,46 @@ struct MmapTensorStore {
 bool is_unused_tensor(const std::string& name);
 
 class ModelLoader {
+public:
+    using FileId       = uint64_t;
+    using FileVersions = std::map<FileId, uint64_t>;
+    enum class FileScope { Catalog,
+                           Isolated };
+
+private:
+    struct FileStamp {
+        std::string path;
+        uintmax_t size = 0;
+        std::filesystem::file_time_type modified;
+    };
+
+    struct FileRecord {
+        FileId id         = 0;
+        uint64_t revision = 0;
+        std::string path;
+        std::string prefix;
+        FileScope scope = FileScope::Catalog;
+        std::vector<FileStamp> dependencies;
+        String2TensorStorage tensors;
+        std::map<std::string, std::string> metadata;
+    };
+
+    std::vector<FileRecord> files_;
+    uint64_t revision_        = 0;
+    bool names_converted_     = false;
+    ggml_type wtype_override_ = GGML_TYPE_COUNT;
+    std::string tensor_type_rules_;
+    std::vector<FileStamp> parsed_dependencies_;
+    std::map<std::string, std::set<std::string>> parsed_tensor_names_;
+
+    static bool read_file_stamp(const std::string& path, FileStamp& stamp);
+    static bool file_unchanged(const FileStamp& stamp);
+    bool parse_file(const std::string& path, const std::string& prefix);
+    bool add_file_impl(const std::string& path, const std::string& prefix, FileId* id, bool force, FileScope scope);
+    ModelLoader file_reader(FileId id, SDVersion version) const;
+    void rebuild_catalog();
+    void invalidate_file_data();
+
 protected:
     SDVersion version_ = VERSION_COUNT;
     std::vector<std::string> file_paths_;
@@ -52,16 +93,27 @@ protected:
 public:
     ModelLoader();
 
+    bool add_file(const std::string& path, const std::string& prefix = "", FileId* id = nullptr, bool force = false, FileScope scope = FileScope::Catalog);
+    bool del_file(FileId id);
+    uint64_t file_revision(FileId id) const;
+    std::string file_path(FileId id) const;
+    String2TensorStorage file_tensors(FileId id, SDVersion version) const;
+    bool load_file_tensors(FileId id, SDVersion version, on_new_tensor_cb_t callback, const std::set<std::string>& names, bool use_mmap = false) const;
+    bool refresh_files(bool include_isolated = true);
+    bool files_changed(bool& changed, bool include_isolated = true) const;
+    bool validate_sources(const std::set<std::string>* tensor_names = nullptr) const;
+    uint64_t revision() const { return revision_; }
+    FileVersions file_versions(const std::vector<std::string>& prefixes = {}) const;
     bool init_from_file(const std::string& file_path, const std::string& prefix = "");
     void convert_tensors_name();
     bool init_from_file_and_convert_name(const std::string& file_path,
                                          const std::string& prefix = "",
                                          SDVersion version         = VERSION_COUNT);
-    SDVersion get_sd_version();
-    std::map<ggml_type, uint32_t> get_wtype_stat();
-    std::map<ggml_type, uint32_t> get_conditioner_wtype_stat();
-    std::map<ggml_type, uint32_t> get_diffusion_model_wtype_stat();
-    std::map<ggml_type, uint32_t> get_vae_wtype_stat();
+    SDVersion get_sd_version() const;
+    std::map<ggml_type, uint32_t> get_wtype_stat() const;
+    std::map<ggml_type, uint32_t> get_conditioner_wtype_stat() const;
+    std::map<ggml_type, uint32_t> get_diffusion_model_wtype_stat() const;
+    std::map<ggml_type, uint32_t> get_vae_wtype_stat() const;
     String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
     const String2TensorStorage& get_tensor_storage_map() const { return tensor_storage_map; }
     const std::map<std::string, std::string>& get_metadata() const { return metadata_; }
@@ -92,8 +144,8 @@ public:
         return names;
     }
 
-    bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type);
-    int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
+    bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) const;
+    int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT) const;
     ~ModelLoader() = default;
 };
 

src/model_loader_files.cpp

@@ -0,0 +1,336 @@
+#include "model_loader.h"
+
+#include <algorithm>
+#include <atomic>
+#include <limits>
+
+#include "core/util.h"
+#include "name_conversion.h"
+
+static uint64_t next_source_revision() {
+    static std::atomic<uint64_t> revision{0};
+    return revision.fetch_add(1, std::memory_order_relaxed) + 1;
+}
+
+bool ModelLoader::read_file_stamp(const std::string& path, FileStamp& stamp) {
+    std::error_code error;
+    const auto file_path = std::filesystem::u8path(path);
+    stamp.path           = path;
+    stamp.size           = 0;
+    stamp.modified       = std::filesystem::last_write_time(file_path, error);
+    if (!error && std::filesystem::is_regular_file(file_path, error)) {
+        stamp.size = std::filesystem::file_size(file_path, error);
+    }
+    if (error) {
+        LOG_ERROR("cannot inspect model source '%s': %s", path.c_str(), error.message().c_str());
+        return false;
+    }
+    return true;
+}
+
+bool ModelLoader::file_unchanged(const FileStamp& stamp) {
+    std::error_code error;
+    if (!std::filesystem::exists(std::filesystem::u8path(stamp.path), error)) {
+        return false;
+    }
+    FileStamp current;
+    return read_file_stamp(stamp.path, current) &&
+           current.size == stamp.size && current.modified == stamp.modified;
+}
+
+void ModelLoader::invalidate_file_data() {
+    file_data.clear();
+    model_files_processed = false;
+}
+
+void ModelLoader::rebuild_catalog() {
+    tensor_storage_map.clear();
+    metadata_.clear();
+    for (const auto& file : files_) {
+        if (file.scope == FileScope::Isolated)
+            continue;
+        for (const auto& entry : file.tensors) {
+            tensor_storage_map[entry.first] = entry.second;
+        }
+        for (const auto& entry : file.metadata) {
+            metadata_[entry.first] = entry.second;
+        }
+    }
+    if (names_converted_) {
+        const SDVersion version = version_ == VERSION_COUNT ? get_sd_version() : version_;
+        tensor_storage_map.clear();
+        for (const auto& file : files_) {
+            if (file.scope == FileScope::Isolated)
+                continue;
+            for (const auto& entry : file.tensors) {
+                TensorStorage tensor            = entry.second;
+                tensor.name                     = convert_tensor_name(tensor.name, version);
+                tensor_storage_map[tensor.name] = std::move(tensor);
+            }
+        }
+    }
+    std::set<size_t> used_files;
+    for (const auto& file : files_) {
+        for (const auto& entry : file.tensors) {
+            used_files.insert(entry.second.file_index);
+        }
+    }
+    for (size_t i = 0; i < file_paths_.size(); ++i) {
+        if (used_files.count(i) == 0) {
+            file_paths_[i].clear();
+        }
+    }
+    set_wtype_override(wtype_override_, tensor_type_rules_);
+}
+
+bool ModelLoader::add_file_impl(const std::string& path, const std::string& prefix, FileId* id, bool force, FileScope scope) {
+    FileStamp root;
+    if (!read_file_stamp(path, root)) {
+        return false;
+    }
+    auto existing = std::find_if(files_.begin(), files_.end(), [&](const FileRecord& file) {
+        return file.path == root.path && file.prefix == prefix && file.scope == scope;
+    });
+    if (existing != files_.end() && !force &&
+        std::all_of(existing->dependencies.begin(), existing->dependencies.end(), file_unchanged)) {
+        if (id != nullptr) {
+            *id = existing->id;
+        }
+        return true;
+    }
+
+    ModelLoader parsed;
+    try {
+        if (!parsed.parse_file(root.path, prefix)) {
+            return false;
+        }
+    } catch (const std::exception& error) {
+        LOG_ERROR("invalid model source '%s': %s", path.c_str(), error.what());
+        return false;
+    }
+
+    std::vector<size_t> file_indices;
+    std::vector<FileStamp> physical_files;
+    for (const auto& physical_path : parsed.file_paths_) {
+        FileStamp stamp;
+        if (!read_file_stamp(physical_path, stamp)) {
+            return false;
+        }
+        parsed.parsed_dependencies_.push_back(stamp);
+        file_indices.push_back(add_file_path(stamp.path));
+        physical_files.push_back(std::move(stamp));
+    }
+    for (auto& entry : parsed.tensor_storage_map) {
+        auto& tensor = entry.second;
+        // Pickle preserves rank-zero scalars; GGML uses a one-element dimension.
+        if (tensor.n_dims == 0) {
+            tensor.n_dims = 1;
+        }
+        if (tensor.n_dims < 1 || tensor.n_dims > SD_MAX_DIMS || tensor.type < 0 ||
+            tensor.type >= GGML_TYPE_COUNT || tensor.file_index >= parsed.file_paths_.size()) {
+            LOG_ERROR("invalid tensor metadata for '%s'", tensor.name.c_str());
+            return false;
+        }
+        uint64_t elements = 1;
+        for (int i = 0; i < tensor.n_dims; ++i) {
+            if (tensor.ne[i] < 0 || (elements != 0 && static_cast<uint64_t>(tensor.ne[i]) > INT64_MAX / elements)) {
+                LOG_ERROR("invalid tensor dimensions for '%s'", tensor.name.c_str());
+                return false;
+            }
+            elements *= tensor.ne[i];
+        }
+        const uint64_t block_size = ggml_blck_size(tensor.type);
+        const uint64_t type_size  = ggml_type_size(tensor.type) * ((tensor.is_f64 || tensor.is_i64) ? 2 : 1);
+        if (block_size == 0 || type_size == 0 || elements % block_size != 0 || elements / block_size > INT64_MAX / type_size) {
+            LOG_ERROR("invalid tensor storage size for '%s'", tensor.name.c_str());
+            return false;
+        }
+        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) {
+                LOG_ERROR("tensor '%s' extends beyond its model file", tensor.name.c_str());
+                return false;
+            }
+        }
+    }
+    if (!std::all_of(parsed.parsed_dependencies_.begin(), parsed.parsed_dependencies_.end(), file_unchanged)) {
+        LOG_ERROR("model source changed while reading metadata: '%s'", path.c_str());
+        return false;
+    }
+
+    FileRecord record;
+    // Snapshots and independently created loaders must never alias different versions.
+    record.revision = next_source_revision();
+    record.id       = existing == files_.end() ? record.revision : existing->id;
+    ++revision_;
+    record.path   = root.path;
+    record.prefix = prefix;
+    record.scope  = scope;
+    std::set<std::string> seen_dependencies;
+    for (auto& stamp : parsed.parsed_dependencies_) {
+        if (seen_dependencies.insert(stamp.path).second) {
+            record.dependencies.push_back(std::move(stamp));
+        }
+    }
+    record.metadata = std::move(parsed.metadata_);
+    record.tensors  = std::move(parsed.tensor_storage_map);
+    for (auto& entry : record.tensors) {
+        entry.second.file_index    = file_indices[entry.second.file_index];
+        entry.second.file_id       = record.id;
+        entry.second.file_revision = record.revision;
+    }
+    if (id != nullptr) {
+        *id = record.id;
+    }
+    if (existing == files_.end()) {
+        files_.push_back(std::move(record));
+    } else {
+        *existing = std::move(record);
+    }
+    rebuild_catalog();
+    return true;
+}
+
+bool ModelLoader::add_file(const std::string& path, const std::string& prefix, FileId* id, bool force, FileScope scope) {
+    ModelLoader candidate = *this;
+    FileId added_id       = 0;
+    if (!candidate.add_file_impl(path, prefix, &added_id, force, scope)) {
+        return false;
+    }
+    *this = std::move(candidate);
+    if (id != nullptr) {
+        *id = added_id;
+    }
+    return true;
+}
+
+bool ModelLoader::del_file(FileId id) {
+    auto it = std::find_if(files_.begin(), files_.end(), [id](const FileRecord& file) { return file.id == id; });
+    if (it == files_.end()) {
+        return false;
+    }
+    files_.erase(it);
+    ++revision_;
+    rebuild_catalog();
+    return true;
+}
+
+bool ModelLoader::files_changed(bool& changed, bool include_isolated) const {
+    changed = false;
+    for (const auto& file : files_) {
+        if (!include_isolated && file.scope == FileScope::Isolated)
+            continue;
+        for (const auto& stamp : file.dependencies) {
+            std::error_code error;
+            if (!std::filesystem::exists(std::filesystem::u8path(stamp.path), error) && !error) {
+                // An updated index may no longer reference this dependency.
+                changed = true;
+                continue;
+            }
+            FileStamp current;
+            if (!read_file_stamp(stamp.path, current)) {
+                return false;
+            }
+            changed |= current.size != stamp.size || current.modified != stamp.modified;
+        }
+    }
+    return true;
+}
+
+bool ModelLoader::refresh_files(bool include_isolated) {
+    bool changed;
+    if (!files_changed(changed, include_isolated)) {
+        return false;
+    }
+    if (!changed) {
+        return true;
+    }
+    ModelLoader candidate = *this;
+    for (const auto& file : files_) {
+        if (!include_isolated && file.scope == FileScope::Isolated)
+            continue;
+        if (!candidate.add_file_impl(file.path, file.prefix, nullptr, false, file.scope)) {
+            return false;
+        }
+    }
+    *this = std::move(candidate);
+    return true;
+}
+
+bool ModelLoader::validate_sources(const std::set<std::string>* tensor_names) const {
+    std::set<FileId> required;
+    if (tensor_names != nullptr) {
+        for (const auto& name : *tensor_names) {
+            auto it = tensor_storage_map.find(name);
+            if (it != tensor_storage_map.end()) {
+                required.insert(it->second.file_id);
+            }
+        }
+    }
+    for (const auto& file : files_) {
+        if (tensor_names != nullptr && required.count(file.id) == 0) {
+            continue;
+        }
+        if (!std::all_of(file.dependencies.begin(), file.dependencies.end(), file_unchanged)) {
+            LOG_ERROR("model source changed; refresh it before execution: '%s'", file.path.c_str());
+            return false;
+        }
+    }
+    return true;
+}
+
+ModelLoader::FileVersions ModelLoader::file_versions(const std::vector<std::string>& prefixes) const {
+    FileVersions versions;
+    for (const auto& entry : tensor_storage_map) {
+        if (prefixes.empty() || std::any_of(prefixes.begin(), prefixes.end(), [&](const std::string& prefix) {
+                return starts_with(entry.first, prefix);
+            })) {
+            versions[entry.second.file_id] = entry.second.file_revision;
+        }
+    }
+    return versions;
+}
+
+uint64_t ModelLoader::file_revision(FileId id) const {
+    for (const auto& file : files_) {
+        if (file.id == id)
+            return file.revision;
+    }
+    return 0;
+}
+
+std::string ModelLoader::file_path(FileId id) const {
+    for (const auto& file : files_) {
+        if (file.id == id)
+            return file.path;
+    }
+    return {};
+}
+
+ModelLoader ModelLoader::file_reader(FileId id, SDVersion version) const {
+    ModelLoader reader;
+    reader.file_paths_      = file_paths_;
+    reader.n_threads_       = n_threads_;
+    reader.version_         = version;
+    reader.names_converted_ = true;
+    for (const auto& file : files_) {
+        if (file.id == id) {
+            reader.files_.push_back(file);
+            reader.files_.back().scope = FileScope::Catalog;
+            break;
+        }
+    }
+    reader.rebuild_catalog();
+    return reader;
+}
+
+String2TensorStorage ModelLoader::file_tensors(FileId id, SDVersion version) const {
+    return file_reader(id, version).tensor_storage_map;
+}
+
+bool ModelLoader::load_file_tensors(FileId id, SDVersion version, on_new_tensor_cb_t callback, const std::set<std::string>& names, bool use_mmap) const {
+    if (file_revision(id) == 0)
+        return false;
+    auto reader = file_reader(id, version);
+    return reader.load_tensors(callback, use_mmap, &names, false);
+}

src/model_manager.cpp

@@ -4,6 +4,7 @@
 #include <cstdint>
 #include <iterator>
 #include <mutex>
+#include <tuple>
 #include <unordered_set>
 
 #include "core/ggml_extend_backend.h"
@@ -26,7 +27,8 @@ static bool lora_specs_equal(const std::vector<ModelManager::LoraSpec>& lhs,
             lhs[i].multiplier != rhs[i].multiplier ||
             lhs[i].is_high_noise != rhs[i].is_high_noise ||
             lhs[i].tensor_name_prefix_filter != rhs[i].tensor_name_prefix_filter ||
-            lhs[i].required != rhs[i].required) {
+            lhs[i].required != rhs[i].required ||
+            lhs[i].file_id != rhs[i].file_id || lhs[i].file_revision != rhs[i].file_revision) {
             return false;
         }
     }
@@ -104,25 +106,61 @@ void ModelManager::set_common_ignore_tensors(std::set<std::string> ignore_tensor
     common_ignore_tensors_ = std::move(ignore_tensors);
 }
 
-void ModelManager::set_loras(std::vector<LoraSpec> loras, SDVersion version) {
-    if (loras.empty() && loras_.empty()) {
-        lora_version_ = version;
-        return;
+bool ModelManager::prepare_lora_sources(std::vector<LoraSpec>& loras) {
+    ModelLoader candidate = model_loader_;
+    std::vector<LoraSpec> resolved;
+    std::set<ModelLoader::FileId> sources;
+    for (auto spec : loras) {
+        const std::string prefix = spec.is_high_noise ? "lora.model.high_noise_" : "lora.";
+        if (!candidate.add_file(spec.path, prefix, &spec.file_id, false, ModelLoader::FileScope::Isolated)) {
+            if (spec.required)
+                return false;
+            LOG_WARN("cannot register LoRA source '%s'", spec.path.c_str());
+            continue;
+        }
+        spec.file_revision = candidate.file_revision(spec.file_id);
+        sources.insert(spec.file_id);
+        resolved.push_back(std::move(spec));
     }
-    if (lora_version_ == version && lora_specs_equal(loras_, loras)) {
-        return;
+    for (auto id : lora_sources_) {
+        if (sources.count(id) == 0)
+            candidate.del_file(id);
     }
+    if (!set_loader(std::move(candidate)))
+        return false;
+    lora_sources_ = std::move(sources);
+    loras         = std::move(resolved);
+    return true;
+}
 
+bool ModelManager::set_loras(std::vector<LoraSpec> loras, SDVersion version) {
+    if (std::any_of(loras.begin(), loras.end(), [](const LoraSpec& spec) { return spec.file_id == 0; }) &&
+        !prepare_lora_sources(loras))
+        return false;
+    for (auto& spec : loras) {
+        spec.file_revision = model_loader_.file_revision(spec.file_id);
+        if (spec.file_revision == 0)
+            return false;
+    }
+    if (lora_version_ == version && lora_specs_equal(loras_, loras))
+        return true;
+    if (!workspace_reclaimers_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) {
+            return state->pin_count != 0;
+        })) {
+        LOG_ERROR("cannot change LoRA configuration during execution");
+        return false;
+    }
     loras_        = std::move(loras);
     lora_version_ = version;
     current_lora_epoch_++;
     reset_lora_applied_params();
+    return true;
 }
 
 std::set<std::string> ModelManager::tensor_names() const {
     std::set<std::string> names;
     for (const auto& state : tensor_states_) {
-        if (state != nullptr) {
+        if (state != nullptr && state->component != ModelComponent::LoRA) {
             names.insert(state->name);
         }
     }
@@ -171,7 +209,7 @@ ggml_backend_buffer_type_t ModelManager::split_buffer_type_for(const TensorState
     return state.split_buffer_type;
 }
 
-bool ModelManager::register_param_tensors(const std::string& desc,
+bool ModelManager::register_param_tensors(ModelComponent component,
                                           std::map<std::string, ggml_tensor*> tensors,
                                           ResidencyMode residency_mode,
                                           ggml_backend_t compute_backend,
@@ -179,15 +217,20 @@ bool ModelManager::register_param_tensors(const std::string& desc,
                                           size_t* registered_tensor_size,
                                           bool allow_split_buffer,
                                           bool params_follow_compute_backend,
-                                          const std::map<ggml_tensor*, enum ggml_op>* tensor_ops) {
-    if (desc.empty()) {
-        LOG_ERROR("model manager tensor desc is empty");
+                                          const std::map<ggml_tensor*, enum ggml_op>* tensor_ops,
+                                          ModelLoader::FileId source_file,
+                                          SDVersion source_version) {
+    if (component == ModelComponent::Count) {
+        LOG_ERROR("model manager tensor component is invalid");
         return false;
     }
     if (registered_tensor_size != nullptr) {
         *registered_tensor_size += estimate_tensors_size(tensors);
     }
 
+    const auto scoped_sources = source_file != 0 ? model_loader_.file_tensors(source_file, source_version) : String2TensorStorage{};
+    const auto& sources       = source_file != 0 ? scoped_sources : model_loader_.get_tensor_storage_map();
+    std::unordered_set<ggml_tensor*> new_tensors;
     std::vector<std::unique_ptr<TensorState>> new_states;
     new_states.reserve(tensors.size());
 
@@ -197,16 +240,23 @@ bool ModelManager::register_param_tensors(const std::string& desc,
         if (tensor == nullptr) {
             continue;
         }
-        if (tensor_states_by_name_.find(name) != tensor_states_by_name_.end()) {
+        if (tensor_states_by_tensor_.count(tensor) != 0 || !new_tensors.insert(tensor).second) {
             LOG_ERROR("model manager tensor name '%s' is already registered", name.c_str());
             return false;
         }
         ggml_set_name(tensor, name.c_str());
 
-        auto state             = std::make_unique<TensorState>();
-        state->name            = name;
-        state->tensor          = tensor;
-        state->desc            = desc;
+        auto state            = std::make_unique<TensorState>();
+        state->name           = name;
+        state->tensor         = tensor;
+        state->component      = component;
+        state->source_file    = source_file;
+        state->source_version = source_version;
+        auto source           = sources.find(name);
+        if (source != sources.end()) {
+            state->source     = source->second;
+            state->has_source = true;
+        }
         state->residency_mode  = residency_mode;
         state->compute_backend = compute_backend;
         state->params_backend  = params_backend;
@@ -225,31 +275,45 @@ bool ModelManager::register_param_tensors(const std::string& desc,
     }
 
     for (auto& state : new_states) {
-        TensorState* registered_state                  = state.get();
-        tensor_states_by_name_[registered_state->name] = registered_state;
+        TensorState* registered_state                      = state.get();
+        tensor_states_by_tensor_[registered_state->tensor] = registered_state;
         tensor_states_.push_back(std::move(state));
     }
     return true;
 }
 
-bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* registered_tensor_size) {
-    if (desc.empty()) {
-        return true;
+bool ModelManager::unregister_param_tensors(ModelComponent component, size_t* registered_tensor_size) {
+    std::unordered_set<TensorState*> states;
+    for (auto& state : tensor_states_) {
+        if (state->component == component)
+            states.insert(state.get());
     }
+    return unregister_tensor_states(states, registered_tensor_size);
+}
 
-    std::unordered_set<TensorState*> target_states;
+bool ModelManager::unregister_param_tensors(const std::vector<ggml_tensor*>& tensors) {
+    std::unordered_set<TensorState*> states;
+    for (auto tensor : tensors) {
+        auto found = tensor_states_by_tensor_.find(tensor);
+        if (found != tensor_states_by_tensor_.end())
+            states.insert(found->second);
+    }
+    return unregister_tensor_states(states, nullptr);
+}
+
+bool ModelManager::unregister_tensor_states(const std::unordered_set<TensorState*>& target_states,
+                                            size_t* registered_tensor_size) {
     size_t released_size = 0;
     for (auto& state : tensor_states_) {
-        if (state == nullptr || state->desc != desc) {
+        if (state == nullptr || target_states.count(state.get()) == 0) {
             continue;
         }
         if (state->pin_count > 0) {
             LOG_ERROR("model manager cannot unregister active %s tensor '%s'",
-                      desc.c_str(),
+                      model_component_name(state->component),
                       state->name.c_str());
             return false;
         }
-        target_states.insert(state.get());
         if (state->tensor != nullptr) {
             released_size += ggml_nbytes(state->tensor);
         }
@@ -260,7 +324,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg
     }
 
     clear_all_prefetched_params();
-    release_compute_staging_blocks(false);
+    release_compute_staging_blocks(false, &target_states);
 
     std::vector<ParamsStorageBlock*> storage_blocks_to_release;
     std::unordered_set<TensorState*> affected_storage_states;
@@ -292,7 +356,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg
         }
         if (state->pin_count > 0 || state->staged_to_compute_backend) {
             LOG_ERROR("model manager cannot unregister %s while tensor '%s' is active",
-                      desc.c_str(),
+                      model_component_name(state->component),
                       state->name.c_str());
             return false;
         }
@@ -305,9 +369,9 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg
         }
     }
 
-    for (auto it = tensor_states_by_name_.begin(); it != tensor_states_by_name_.end();) {
+    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_name_.erase(it);
+            it = tensor_states_by_tensor_.erase(it);
         } else {
             ++it;
         }
@@ -559,19 +623,24 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
 }
 
 bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states) {
-    if (loras_.empty()) {
+    if (loras_.empty() || applying_loras_)
         return true;
-    }
+    applying_loras_ = true;
+    struct ApplyGuard {
+        bool& active;
+        ~ApplyGuard() { active = false; }
+    } guard{applying_loras_};
 
     struct LoraApplyGroup {
         std::map<std::string, ggml_tensor*> model_tensors;
         std::vector<TensorState*> states;
     };
 
-    std::map<ggml_backend_t, LoraApplyGroup> groups;
+    using ApplyTarget = std::tuple<ggml_backend_t, ggml_backend_t, ResidencyMode>;
+    std::map<ApplyTarget, LoraApplyGroup> groups;
     for (TensorState* state : states) {
-        if (state == nullptr || state->tensor == nullptr ||
-            should_ignore(*state) || is_optional_missing_tensor(state->name)) {
+        if (state == nullptr || state->tensor == nullptr || state->component == ModelComponent::LoRA ||
+            state->component == ModelComponent::LatentUpsampler || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
             continue;
         }
         if (state->applied_lora_epoch == current_lora_epoch_) {
@@ -596,7 +665,7 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
             LOG_ERROR("model manager lora target tensor '%s' is not prepared", state->name.c_str());
             return false;
         }
-        LoraApplyGroup& group            = groups[state->compute_backend];
+        LoraApplyGroup& group            = groups[{state->compute_backend, state->params_backend, state->residency_mode}];
         group.model_tensors[state->name] = state->tensor;
         group.states.push_back(state);
     }
@@ -607,20 +676,20 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
 
     std::set<std::string> all_tensor_names = tensor_names();
     for (auto& group_pair : groups) {
-        ggml_backend_t compute_backend = group_pair.first;
+        ggml_backend_t compute_backend = std::get<0>(group_pair.first);
         LoraApplyGroup& group          = group_pair.second;
         for (const LoraSpec& lora_spec : loras_) {
             if (group.model_tensors.empty()) {
                 continue;
             }
 
-            std::string id = lora_id(lora_spec);
-            auto lora      = std::make_shared<LoraModel>(id,
-                                                    compute_backend,
-                                                    compute_backend,
-                                                    lora_spec.path,
-                                                    lora_spec.is_high_noise ? "model.high_noise_" : "",
-                                                    lora_version_);
+            std::string id     = lora_id(lora_spec);
+            const auto* target = group.states.front();
+            // The temporary runner is destroyed before this manager call returns.
+            auto borrowed_manager = std::shared_ptr<ModelManager>(this, [](ModelManager*) {});
+            auto lora             = std::make_shared<LoraModel>(id, compute_backend, target->params_backend,
+                                                    borrowed_manager, lora_spec.file_id, lora_version_,
+                                                    target->residency_mode);
 
             LoraModel::filter_t lora_tensor_filter = nullptr;
             if (!lora_spec.tensor_name_prefix_filter.empty()) {
@@ -628,7 +697,7 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
                     return starts_with(tensor_name, lora_spec.tensor_name_prefix_filter);
                 };
             }
-            if (!lora->load_from_file(n_threads_, lora_tensor_filter)) {
+            if (!lora->init_params(n_threads_, lora_tensor_filter)) {
                 LOG_WARN("load lora tensors from %s failed", lora_spec.path.c_str());
                 if (lora_spec.required) {
                     return false;
@@ -643,7 +712,8 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
                 continue;
             }
             lora->multiplier = lora_spec.multiplier;
-            lora->apply(group.model_tensors, all_tensor_names, lora_version_, n_threads_, false);
+            if (!lora->apply(group.model_tensors, all_tensor_names, lora_version_, n_threads_, false))
+                return false;
             lora->release_loaded_tensors();
         }
 
@@ -657,12 +727,13 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
 }
 
 void ModelManager::reset_lora_applied_params() {
-    clear_all_prefetched_params();
-    release_compute_staging_blocks(true);
-    release_params_storage_blocks(true);
+    std::unordered_set<TensorState*> affected;
     for (auto& state : tensor_states_) {
-        state->applied_lora_epoch = UINT64_MAX;
+        if (state->component != ModelComponent::LoRA && state->applied_lora_epoch != UINT64_MAX) {
+            affected.insert(state.get());
+        }
     }
+    invalidate_sources(affected);
 }
 
 bool ModelManager::should_ignore(const TensorState& state) const {
@@ -684,21 +755,19 @@ bool ModelManager::validate_tensor(const TensorState& state) const {
         return true;
     }
 
-    const auto& tensor_storage_map = model_loader_.get_tensor_storage_map();
-    auto ts_it                     = tensor_storage_map.find(state.name);
-    if (ts_it == tensor_storage_map.end()) {
-        LOG_ERROR("%s tensor '%s' not in model metadata", state.desc.c_str(), state.name.c_str());
+    if (!state.has_source) {
+        LOG_ERROR("%s tensor '%s' not in model metadata", model_component_name(state.component), state.name.c_str());
         return false;
     }
 
-    const TensorStorage& tensor_storage = ts_it->second;
+    const TensorStorage& tensor_storage = state.source;
     if (state.tensor->ne[0] != tensor_storage.ne[0] ||
         state.tensor->ne[1] != tensor_storage.ne[1] ||
         state.tensor->ne[2] != tensor_storage.ne[2] ||
         state.tensor->ne[3] != tensor_storage.ne[3]) {
         LOG_ERROR(
             "%s tensor '%s' has wrong shape in model metadata: got [%d, %d, %d, %d], expected [%d, %d, %d, %d]",
-            state.desc.c_str(),
+            model_component_name(state.component),
             state.name.c_str(),
             (int)tensor_storage.ne[0], (int)tensor_storage.ne[1], (int)tensor_storage.ne[2], (int)tensor_storage.ne[3],
             (int)state.tensor->ne[0], (int)state.tensor->ne[1], (int)state.tensor->ne[2], (int)state.tensor->ne[3]);
@@ -746,7 +815,7 @@ bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
 }
 
 bool ModelManager::can_mmap_storage(const TensorState& state) const {
-    if (!enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
+    if (state.source_file != 0 || !enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
         return false;
     }
     if (state.compute_backend == nullptr || state.params_backend == nullptr) {
@@ -857,75 +926,55 @@ bool ModelManager::alloc_params_buffers(const std::vector<TensorState*>& states,
 }
 
 bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
-    std::map<std::string, TensorState*> states_by_name;
-    std::set<std::string> target_tensor_names;
-    for (TensorState* state : states) {
-        if (state == nullptr) {
+    using ReadGroup = std::pair<ModelLoader::FileId, SDVersion>;
+    using ReadBatch = std::map<std::string, std::vector<TensorState*>>;
+    std::map<ReadGroup, std::vector<ReadBatch>> groups;
+    for (auto* state : states) {
+        if (state == nullptr)
             continue;
+        auto& batches = groups[{state->source_file, state->source_version}];
+        // The loader supplies one destination per name; only conflicting types need another batch.
+        auto batch = std::find_if(batches.begin(), batches.end(), [&](const ReadBatch& candidate) {
+            auto found = candidate.find(state->name);
+            return found == candidate.end() || found->second.front()->tensor->type == state->tensor->type;
+        });
+        if (batch == batches.end()) {
+            batches.emplace_back();
+            batch = std::prev(batches.end());
         }
-        states_by_name[state->name] = state;
-        target_tensor_names.insert(state->name);
+        (*batch)[state->name].push_back(state);
     }
-    if (states_by_name.empty()) {
-        return true;
-    }
-
-    std::set<std::string> loaded_names;
-    std::mutex loaded_names_mutex;
-    auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
-        const std::string& name = tensor_storage.name;
-        *dst_tensor             = nullptr;
-
-        auto state_it = states_by_name.find(name);
-        if (state_it == states_by_name.end()) {
-            return true;
+    for (auto& group : groups) {
+        for (auto& batch : group.second) {
+            std::set<std::string> names;
+            std::set<std::string> loaded;
+            std::mutex mutex;
+            for (const auto& entry : batch)
+                names.insert(entry.first);
+            auto callback = [&](const TensorStorage& source, ggml_tensor** dst) {
+                *dst       = nullptr;
+                auto found = batch.find(source.name);
+                if (found == batch.end())
+                    return true;
+                *dst = found->second.front()->tensor;
+                std::lock_guard<std::mutex> lock(mutex);
+                loaded.insert(source.name);
+                return true;
+            };
+            const auto file = group.first.first;
+            bool success    = file == 0 ? model_loader_.load_tensors(callback, enable_mmap_, &names)
+                                        : model_loader_.load_file_tensors(file, group.first.second, callback, names, enable_mmap_);
+            if (!success || loaded != names)
+                return false;
+            for (auto& entry : batch) {
+                auto* first = entry.second.front()->tensor;
+                for (auto* state : entry.second) {
+                    if (state->tensor != first)
+                        ggml_backend_tensor_copy(first, state->tensor);
+                    state->loaded_to_params_backend = true;
+                }
+            }
         }
-
-        TensorState* state = state_it->second;
-        if (state == nullptr || state->tensor == nullptr) {
-            LOG_ERROR("model manager tensor '%s' is null", name.c_str());
-            return false;
-        }
-
-        if (state->tensor->ne[0] != tensor_storage.ne[0] ||
-            state->tensor->ne[1] != tensor_storage.ne[1] ||
-            state->tensor->ne[2] != tensor_storage.ne[2] ||
-            state->tensor->ne[3] != tensor_storage.ne[3]) {
-            LOG_ERROR(
-                "model manager tensor '%s' has wrong shape in model file: got [%d, %d, %d, %d], expected [%d, %d, %d, %d]",
-                name.c_str(),
-                (int)tensor_storage.ne[0], (int)tensor_storage.ne[1], (int)tensor_storage.ne[2], (int)tensor_storage.ne[3],
-                (int)state->tensor->ne[0], (int)state->tensor->ne[1], (int)state->tensor->ne[2], (int)state->tensor->ne[3]);
-            return false;
-        }
-
-        {
-            std::lock_guard<std::mutex> lock(loaded_names_mutex);
-            loaded_names.insert(name);
-        }
-        *dst_tensor = state->tensor;
-        return true;
-    };
-
-    if (!model_loader_.load_tensors(on_new_tensor_cb, enable_mmap_, &target_tensor_names)) {
-        LOG_ERROR("model manager load tensors failed");
-        return false;
-    }
-
-    bool missing = false;
-    for (const auto& pair : states_by_name) {
-        const std::string& name = pair.first;
-        if (loaded_names.find(name) == loaded_names.end()) {
-            LOG_ERROR("model manager tensor '%s' was not loaded", name.c_str());
-            missing = true;
-        }
-    }
-    if (missing) {
-        return false;
-    }
-
-    for (const auto& pair : states_by_name) {
-        pair.second->loaded_to_params_backend = true;
     }
     return true;
 }
@@ -1138,6 +1187,14 @@ void ModelManager::release_all() {
     release_params_storage_blocks(true);
 }
 
+ggml_tensor* ModelManager::resolve_param_tensor(ggml_tensor* tensor) const {
+    for (auto* current = tensor; current != nullptr; current = current->view_src) {
+        if (tensor_states_by_tensor_.count(current) != 0)
+            return current;
+    }
+    return nullptr;
+}
+
 bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*>& tensors,
                                                   std::vector<TensorState*>& required_states,
                                                   ggml_backend_t compute_backend) const {
@@ -1147,21 +1204,13 @@ bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*
         if (tensor == nullptr) {
             continue;
         }
-        const char* raw_name = ggml_get_name(tensor);
-        if (raw_name == nullptr || raw_name[0] == '\0') {
-            LOG_ERROR("model manager unnamed tensor is not registered");
-            return false;
-        }
-        auto state_it = tensor_states_by_name_.find(raw_name);
-        if (state_it == tensor_states_by_name_.end()) {
-            LOG_ERROR("model manager tensor '%s' is not registered", raw_name);
-            return false;
-        }
-        TensorState* state = state_it->second;
-        if (state == nullptr) {
-            LOG_ERROR("model manager tensor '%s' has no tensor state", raw_name);
+        auto param = resolve_param_tensor(tensor);
+        auto found = tensor_states_by_tensor_.find(param);
+        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) {
@@ -1375,8 +1424,8 @@ bool ModelManager::prepare_params(const std::vector<ggml_tensor*>& tensors) {
     }
     if (!apply_loras_to_params(required_states)) {
         finish_compute_backend_usage(required_states);
-        release_compute_staging_blocks(false);
-        release_params_storage_blocks(false);
+        std::unordered_set<TensorState*> failed(required_states.begin(), required_states.end());
+        invalidate_sources(failed);
         return false;
     }
     return true;
@@ -1535,18 +1584,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));
@@ -1570,11 +1636,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));
@@ -1690,12 +1752,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

@@ -10,6 +10,7 @@
 #include <vector>
 
 #include "device_residency_manager.h"
+#include "model_component.h"
 #include "model_loader.h"
 
 class ModelManager : public DeviceResidencyManager {
@@ -24,16 +25,22 @@ public:
         float multiplier   = 1.0f;
         bool is_high_noise = false;
         std::string tensor_name_prefix_filter;
-        bool required = false;
+        bool required               = false;
+        ModelLoader::FileId file_id = 0;
+        uint64_t file_revision      = 0;
     };
 
 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;
-        ggml_tensor* tensor = nullptr;
-        std::string desc;
+        ggml_tensor* tensor      = nullptr;
+        ModelComponent component = ModelComponent::Count;
+        TensorStorage source;
+        bool has_source                 = false;
+        ModelLoader::FileId source_file = 0;
+        SDVersion source_version        = VERSION_COUNT;
 
         ResidencyMode residency_mode                 = ResidencyMode::ParamBackend;
         ggml_backend_t compute_backend               = nullptr;
@@ -79,7 +86,7 @@ private:
 
     ModelLoader model_loader_;
     std::vector<std::unique_ptr<TensorState>> tensor_states_;
-    std::map<std::string, TensorState*> tensor_states_by_name_;
+    std::map<const ggml_tensor*, TensorState*> tensor_states_by_tensor_;
     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_;
@@ -91,6 +98,8 @@ private:
     bool warned_split_lora_skip_ = false;
     std::set<std::string> common_ignore_tensors_;
     std::vector<LoraSpec> loras_;
+    std::set<ModelLoader::FileId> lora_sources_;
+    bool applying_loras_             = false;
     SDVersion lora_version_          = VERSION_COUNT;
     uint64_t current_lora_epoch_     = 0;
     uint64_t residency_epoch_        = 0;
@@ -102,6 +111,7 @@ private:
 
     void finish_compute_backend_usage(const std::vector<TensorState*>& states);
     void release_all();
+    void invalidate_sources(const std::unordered_set<TensorState*>& states);
 
     ggml_backend_t prefetch_backend_for(ggml_backend_t compute_backend);
     bool populate_prefetch_block(PrefetchBlock& block);
@@ -152,15 +162,27 @@ private:
     void free_params_storage_block(ParamsStorageBlock& block);
     void erase_params_storage_block(ParamsStorageBlock* block);
     void reset_lora_applied_params();
+    bool unregister_tensor_states(const std::unordered_set<TensorState*>& states, size_t* size);
     size_t other_runtime_resident_bytes(uintptr_t owner_id,
                                         ggml_backend_t compute_backend) const;
 
 public:
     ~ModelManager() override;
 
-    ModelLoader& loader() { return model_loader_; }
     const ModelLoader& loader() const { return model_loader_; }
 
+    bool set_loader(ModelLoader loader);
+    bool add_file(const std::string& path, const std::string& prefix = "", ModelLoader::FileId* id = nullptr, bool force = false);
+    bool del_file(ModelLoader::FileId id);
+    bool refresh_files();
+    ModelLoader::FileVersions source_versions(const std::set<ModelComponent>& components, const ModelLoader& loader) const;
+    size_t registered_params_size(const std::set<ModelComponent>& components) const;
+
+    void prepare_file_io() { model_loader_.process_model_files(enable_mmap_, writable_mmap_); }
+    bool load_float_tensor(const std::string& name, std::vector<float>& data) {
+        return model_loader_.load_float_tensor(name, data, n_threads_, enable_mmap_);
+    }
+
     void set_n_threads(int n_threads) {
         n_threads_ = n_threads;
         model_loader_.set_n_threads(n_threads);
@@ -172,14 +194,15 @@ public:
     void set_enable_mmap(bool enable_mmap) { enable_mmap_ = enable_mmap; }
     void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; }
     void set_common_ignore_tensors(std::set<std::string> ignore_tensors);
-    void set_loras(std::vector<LoraSpec> loras, SDVersion version);
+    bool prepare_lora_sources(std::vector<LoraSpec>& loras);
+    bool set_loras(std::vector<LoraSpec> loras, SDVersion version);
     void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector<std::pair<ggml_backend_t, size_t>>& device_limits);
 
     static bool tensor_shape_supports_split_buffer(const ggml_tensor* tensor);
 
     std::set<std::string> tensor_names() const;
 
-    bool register_param_tensors(const std::string& desc,
+    bool register_param_tensors(ModelComponent component,
                                 std::map<std::string, ggml_tensor*> tensors,
                                 ResidencyMode residency_mode,
                                 ggml_backend_t compute_backend,
@@ -187,13 +210,18 @@ public:
                                 size_t* registered_tensor_size                         = nullptr,
                                 bool allow_split_buffer                                = false,
                                 bool params_follow_compute_backend                     = false,
-                                const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr);
+                                const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr,
+                                ModelLoader::FileId source_file                        = 0,
+                                SDVersion source_version                               = VERSION_COUNT);
 
-    bool unregister_param_tensors(const std::string& desc,
+    ggml_tensor* resolve_param_tensor(ggml_tensor* tensor) const override;
+    bool unregister_param_tensors(const std::vector<ggml_tensor*>& tensors);
+
+    bool unregister_param_tensors(ModelComponent component,
                                   size_t* registered_tensor_size = nullptr);
 
     template <typename Runner>
-    bool register_runner_params(const std::string& desc,
+    bool register_runner_params(ModelComponent component,
                                 Runner& runner,
                                 ResidencyMode residency_mode,
                                 ggml_backend_t compute_backend,
@@ -201,7 +229,7 @@ public:
                                 size_t* registered_tensor_size = nullptr) {
         std::map<std::string, ggml_tensor*> tensors;
         runner.get_param_tensors(tensors);
-        return register_param_tensors(desc,
+        return register_param_tensors(component,
                                       std::move(tensors),
                                       residency_mode,
                                       compute_backend,
@@ -210,7 +238,7 @@ public:
     }
 
     template <typename Runner>
-    bool register_runner_params(const std::string& desc,
+    bool register_runner_params(ModelComponent component,
                                 Runner& runner,
                                 const std::string& prefix,
                                 ResidencyMode residency_mode,
@@ -219,7 +247,7 @@ public:
                                 size_t* registered_tensor_size = nullptr) {
         std::map<std::string, ggml_tensor*> tensors;
         runner.get_param_tensors(tensors, prefix);
-        return register_param_tensors(desc,
+        return register_param_tensors(component,
                                       std::move(tensors),
                                       residency_mode,
                                       compute_backend,

src/model_manager_files.cpp

@@ -0,0 +1,154 @@
+#include "model_manager.h"
+
+#include <algorithm>
+
+#include "core/util.h"
+
+static bool same_tensor_source(const TensorStorage& a, const TensorStorage& b) {
+    return a.file_id == b.file_id && a.file_revision == b.file_revision &&
+           a.file_index == b.file_index && a.offset == b.offset && a.index_in_zip == b.index_in_zip &&
+           a.storage_key == b.storage_key && a.type == b.type && a.expected_type == b.expected_type &&
+           a.n_dims == b.n_dims && std::equal(a.ne, a.ne + SD_MAX_DIMS, b.ne) &&
+           a.is_f8_e4m3 == b.is_f8_e4m3 && a.is_f8_e5m2 == b.is_f8_e5m2 &&
+           a.is_f64 == b.is_f64 && a.is_i64 == b.is_i64 &&
+           a.is_int8_tensorwise == b.is_int8_tensorwise && a.int8_convrot == b.int8_convrot &&
+           a.int8_convrot_group_size == b.int8_convrot_group_size;
+}
+
+void ModelManager::invalidate_sources(const std::unordered_set<TensorState*>& states) {
+    auto affected = states;
+    for (const auto& block : params_storage_blocks_) {
+        if (std::any_of(block->states.begin(), block->states.end(), [&](TensorState* state) { return states.count(state) != 0; })) {
+            affected.insert(block->states.begin(), block->states.end());
+        }
+    }
+    for (auto it = prefetch_blocks_.begin(); it != prefetch_blocks_.end();) {
+        if (std::any_of(it->second->states.begin(), it->second->states.end(), [&](TensorState* state) { return affected.count(state) != 0; })) {
+            free_prefetch_block(*it->second);
+            it = prefetch_blocks_.erase(it);
+        } else {
+            ++it;
+        }
+    }
+    for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end();) {
+        if (std::any_of((*it)->staged_tensors.begin(), (*it)->staged_tensors.end(), [&](const auto& entry) { return affected.count(entry.first) != 0; })) {
+            ggml_backend_synchronize((*it)->compute_backend);
+            free_compute_staging_block(**it);
+            it = compute_staging_blocks_.erase(it);
+        } else {
+            ++it;
+        }
+    }
+    for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end();) {
+        if (std::any_of((*it)->states.begin(), (*it)->states.end(), [&](TensorState* state) { return affected.count(state) != 0; })) {
+            free_params_storage_block(**it);
+            it = params_storage_blocks_.erase(it);
+        } else {
+            ++it;
+        }
+    }
+    for (auto* state : affected) {
+        state->metadata_validated = false;
+        state->applied_lora_epoch = UINT64_MAX;
+    }
+}
+
+bool ModelManager::set_loader(ModelLoader loader) {
+    if (!workspace_reclaimers_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) {
+            return state->pin_count != 0;
+        })) {
+        LOG_ERROR("cannot update model sources during execution");
+        return false;
+    }
+    std::map<std::pair<ModelLoader::FileId, SDVersion>, String2TensorStorage> scoped;
+    auto sources_for = [&](const TensorState& state) -> const String2TensorStorage& {
+        if (state.source_file == 0)
+            return loader.get_tensor_storage_map();
+        auto key   = std::make_pair(state.source_file, state.source_version);
+        auto found = scoped.find(key);
+        if (found == scoped.end())
+            found = scoped.emplace(key, loader.file_tensors(key.first, key.second)).first;
+        return found->second;
+    };
+    bool lora_changed = false;
+    for (const auto& spec : loras_) {
+        lora_changed |= loader.file_revision(spec.file_id) != spec.file_revision;
+    }
+    std::unordered_set<TensorState*> changed;
+    for (const auto& state : tensor_states_) {
+        const auto& sources = sources_for(*state);
+        auto source         = sources.find(state->name);
+        const bool found    = source != sources.end();
+        if (found != state->has_source || (found && !same_tensor_source(state->source, source->second)) ||
+            (lora_changed && state->component != ModelComponent::LoRA && state->applied_lora_epoch != UINT64_MAX)) {
+            changed.insert(state.get());
+        }
+    }
+    invalidate_sources(changed);
+    for (auto* state : changed) {
+        const auto& sources = sources_for(*state);
+        auto source         = sources.find(state->name);
+        state->has_source   = source != sources.end();
+        state->source       = state->has_source ? source->second : TensorStorage{};
+    }
+    if (lora_changed) {
+        ++current_lora_epoch_;
+        for (auto& spec : loras_)
+            spec.file_revision = loader.file_revision(spec.file_id);
+    }
+    model_loader_ = std::move(loader);
+    model_loader_.set_n_threads(n_threads_);
+    return true;
+}
+
+bool ModelManager::add_file(const std::string& path, const std::string& prefix, ModelLoader::FileId* id, bool force) {
+    ModelLoader candidate = model_loader_;
+    ModelLoader::FileId added_id;
+    if (!candidate.add_file(path, prefix, &added_id, force) || !set_loader(std::move(candidate))) {
+        return false;
+    }
+    if (id != nullptr) {
+        *id = added_id;
+    }
+    return true;
+}
+
+bool ModelManager::del_file(ModelLoader::FileId id) {
+    ModelLoader candidate = model_loader_;
+    return candidate.del_file(id) && set_loader(std::move(candidate));
+}
+
+bool ModelManager::refresh_files() {
+    ModelLoader candidate = model_loader_;
+    return candidate.refresh_files() && set_loader(std::move(candidate));
+}
+
+ModelLoader::FileVersions ModelManager::source_versions(const std::set<ModelComponent>& components, const ModelLoader& loader) const {
+    ModelLoader::FileVersions versions;
+    const auto& sources = loader.get_tensor_storage_map();
+    for (const auto& state : tensor_states_) {
+        if (components.count(state->component) == 0) {
+            continue;
+        }
+        if (state->source_file != 0) {
+            versions[state->source_file] = loader.file_revision(state->source_file);
+            continue;
+        }
+        auto source = sources.find(state->name);
+        if (source != sources.end()) {
+            versions[source->second.file_id] = source->second.file_revision;
+        }
+    }
+    return versions;
+}
+
+size_t ModelManager::registered_params_size(const std::set<ModelComponent>& components) const {
+    size_t bytes = 0;
+    std::unordered_set<const ggml_tensor*> seen;
+    for (const auto& state : tensor_states_) {
+        if (components.count(state->component) != 0 && state->tensor != nullptr && seen.insert(state->tensor).second) {
+            bytes += ggml_nbytes(state->tensor);
+        }
+    }
+    return bytes;
+}

src/pipeline/diffusion_engine.h

@@ -0,0 +1,489 @@
+#ifndef __SD_PIPELINE_DIFFUSION_ENGINE_H__
+#define __SD_PIPELINE_DIFFUSION_ENGINE_H__
+
+#include <atomic>
+#include <cmath>
+#include <functional>
+#include <list>
+#include <map>
+#include <memory>
+#include <mutex>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_graph_cut.h"
+#include "core/tensor.hpp"
+#include "core/util.h"
+#include "model/adapter/lora.hpp"
+#include "model_builders.h"
+#include "model_manager.h"
+#include "stable-diffusion.h"
+
+class RNG;
+struct Denoiser;
+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_mage_flow(version) ||
+           sd_version_is_longcat(version) ||
+           sd_version_is_z_image(version) ||
+           sd_version_is_boogu_image(version);
+}
+
+class StableDiffusionGGML {
+public:
+    SDBackendManager backend_manager;
+
+    SDVersion version;
+    bool external_vae_is_invalid = false;
+
+    bool circular_x = false;
+    bool circular_y = false;
+
+    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::shared_ptr<Conditioner> cond_stage_model;
+    std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision;  // for svd or wan2.1 i2v
+    std::shared_ptr<DiffusionModelRunner> diffusion_model;
+    std::shared_ptr<DiffusionModelRunner> high_noise_diffusion_model;
+    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;
+    sd::Tensor<float> ip_adapter_uncond_tokens;
+    float ip_adapter_strength = 1.0f;
+    std::vector<std::shared_ptr<GenerationExtension>> generation_extensions;
+    struct RuntimeLora {
+        ModelManager::LoraSpec spec;
+        SDBackendModule module;
+        std::shared_ptr<LoraModel> model;
+
+        bool matches(const ModelManager::LoraSpec& other) const {
+            return spec.file_id == other.file_id && spec.file_revision == other.file_revision &&
+                   spec.tensor_name_prefix_filter == other.tensor_name_prefix_filter;
+        }
+    };
+    std::vector<RuntimeLora> runtime_lora_models;
+    bool apply_lora_immediately = false;
+    int animatediff_num_frames  = 0;
+
+    std::string taesd_path;
+    sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr};
+    bool enable_mmap                     = false;
+    sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment;
+    bool disable_prefetch          = false;
+    bool disable_segmented_compute = false;
+    bool eager_load                = false;
+    std::string backend_spec;
+    std::string params_backend_spec;
+    std::string split_mode_spec;
+    bool auto_fit_enabled = false;
+
+    bool diffusion_conv_direct = false;
+
+    bool is_using_v_parameterization     = false;
+    bool is_using_edm_v_parameterization = false;
+
+    std::shared_ptr<ModelManager> model_manager;
+
+    enum class RunnerGroup { Core,
+                             VAE,
+                             ControlNet,
+                             Extensions };
+    using RunnerGroups = std::set<RunnerGroup>;
+
+    struct ModelConfig {
+        sd_ctx_params_t params{};
+        std::list<std::string> strings;
+        std::vector<sd_embedding_t> embeddings;
+        ModelLoader::FileId control_net_file = 0;
+        bool use_tae                         = false;
+        bool use_audio_vae                   = false;
+        bool photomaker_source_available     = false;
+        bool animatediff_loaded              = false;
+
+        explicit ModelConfig(const sd_ctx_params_t& initial)
+            : params(initial) {
+            for (auto member : {&sd_ctx_params_t::model_path, &sd_ctx_params_t::clip_l_path,
+                                &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,
+                                &sd_ctx_params_t::control_net_path, &sd_ctx_params_t::ip_adapter_path,
+                                &sd_ctx_params_t::motion_module_path, &sd_ctx_params_t::photo_maker_path,
+                                &sd_ctx_params_t::pulid_weights_path, &sd_ctx_params_t::tensor_type_rules,
+                                &sd_ctx_params_t::max_vram, &sd_ctx_params_t::backend,
+                                &sd_ctx_params_t::params_backend, &sd_ctx_params_t::split_mode,
+                                &sd_ctx_params_t::rpc_servers, &sd_ctx_params_t::model_args}) {
+                strings.emplace_back(SAFE_STR(initial.*member));
+                params.*member = strings.back().c_str();
+            }
+            for (uint32_t i = 0; i < initial.embedding_count; ++i) {
+                strings.emplace_back(SAFE_STR(initial.embeddings[i].name));
+                const char* name = strings.back().c_str();
+                strings.emplace_back(SAFE_STR(initial.embeddings[i].path));
+                embeddings.push_back({name, strings.back().c_str()});
+            }
+            params.embeddings = embeddings.data();
+        }
+
+        ModelConfig(const ModelConfig& other)
+            : ModelConfig(other.params) {
+            control_net_file            = other.control_net_file;
+            use_tae                     = other.use_tae;
+            use_audio_vae               = other.use_audio_vae;
+            photomaker_source_available = other.photomaker_source_available;
+            animatediff_loaded          = other.animatediff_loaded;
+        }
+        ModelConfig& operator=(const ModelConfig&) = delete;
+
+        void set_control_net(ModelLoader::FileId id, const std::string& path) {
+            control_net_file = id;
+            strings.push_back(path);
+            params.control_net_path = strings.back().c_str();
+        }
+    };
+
+    struct RunnerState {
+        bool ready                = false;
+        uint64_t catalog_revision = 0;
+        std::map<RunnerGroup, ModelLoader::FileVersions> sources;
+    };
+
+    std::recursive_mutex execution_mutex;
+    std::unique_ptr<ModelConfig> config_;
+    RunnerState runner_state_;
+    bool executing_ = false;
+
+    std::shared_ptr<Denoiser> denoiser;
+    std::vector<float> file_alphas_cumprod;
+
+    StableDiffusionGGML();
+    ~StableDiffusionGGML();
+
+    static const std::map<RunnerGroup, std::set<ModelComponent>>& runner_components();
+
+    static RunnerGroups all_runner_groups();
+
+    ModelLoader::FileVersions runner_source_versions(RunnerGroup group, const ModelLoader& loader) const;
+
+    void capture_runner_sources();
+
+    void end_runners();
+
+    bool reset_runners(const RunnerGroups& groups);
+
+    bool refresh_model_sources();
+
+    bool apply_model_update(ModelLoader candidate,
+                            std::unique_ptr<ModelConfig> next_config = nullptr,
+                            RunnerGroups groups                      = {});
+
+    struct ContextOperation {
+        StableDiffusionGGML& sd;
+        std::unique_lock<std::recursive_mutex> lock;
+        bool acquired = false;
+
+        explicit ContextOperation(StableDiffusionGGML& sd)
+            : sd(sd), lock(sd.execution_mutex, std::try_to_lock) {
+            if (!lock.owns_lock() || sd.executing_) {
+                // The caller may be a log callback, so rejecting it must not log.
+                return;
+            }
+            sd.executing_ = true;
+            acquired      = true;
+        }
+
+        ~ContextOperation() {
+            if (acquired) {
+                sd.executing_ = false;
+            }
+        }
+    };
+
+    struct ExecutionScope {
+        ContextOperation operation;
+        bool ready = false;
+
+        explicit ExecutionScope(StableDiffusionGGML& sd)
+            : operation(sd) {
+            ready = operation.acquired && sd.refresh_model_sources();
+        }
+
+        ~ExecutionScope() {
+            if (ready) {
+                operation.sd.end_runners();
+            }
+        }
+    };
+
+    ggml_backend_t backend_for(SDBackendModule module);
+
+    ggml_backend_t params_backend_for(SDBackendModule module);
+
+    std::atomic<sd_cancel_mode_t> cancellation_flag = SD_CANCEL_RESET;
+
+    void set_cancel_flag(enum sd_cancel_mode_t flag);
+
+    void reset_cancel_flag();
+
+    enum sd_cancel_mode_t get_cancel_flag();
+
+    size_t max_graph_vram_bytes_for_module(SDBackendModule module);
+
+    std::vector<size_t> layer_split_vram_limits_for_backends(const std::vector<ggml_backend_t>& backends);
+
+    bool ensure_backend_pair(SDBackendModule module);
+
+    template <typename T>
+    bool register_runner_params(ModelComponent component,
+                                const std::shared_ptr<T>& model,
+                                SDBackendModule module,
+                                size_t* params_mem_size = nullptr);
+
+    template <typename T>
+    bool register_row_split_runner_params(ModelComponent component,
+                                          const std::shared_ptr<T>& model,
+                                          SDBackendModule module,
+                                          const std::vector<ggml_backend_t>& module_backends,
+                                          std::map<std::string, ggml_tensor*> group_tensors,
+                                          const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
+                                          ModelManager::ResidencyMode residency_mode,
+                                          size_t* params_mem_size);
+
+    // Register graph-cut layer-split tensors on the primary backend first.
+    // The first real graph assigns each param tensor to a runtime backend
+    // before weights are loaded or staged.
+    template <typename T>
+    bool register_layer_split_runner_params(ModelComponent component,
+                                            const std::shared_ptr<T>& model,
+                                            SDBackendModule module,
+                                            const std::vector<ggml_backend_t>& module_backends,
+                                            std::map<std::string, ggml_tensor*> group_tensors,
+                                            const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
+                                            ModelManager::ResidencyMode residency_mode,
+                                            size_t* params_mem_size);
+
+    bool unload_control_net();
+
+    bool load_control_net_from_file(const std::string& path);
+
+    void apply_circular_axes(bool circular_x, bool circular_y);
+
+    bool init_backend();
+
+    bool row_split_active();
+
+    bool graph_cut_layer_split_active();
+
+    std::shared_ptr<RNG> get_rng(rng_type_t rng_type);
+
+    void refresh_compvis_denoiser_sigmas();
+
+    void load_alphas_cumprod();
+
+    bool init_model_loader(ModelLoader& model_loader, ModelConfig& configuration);
+
+    bool init(const sd_ctx_params_t* sd_ctx_params);
+
+    bool uses_tae() const;
+
+    bool tae_preview_only() const;
+
+    void configure_weight_loading();
+
+    sd::model_builders::Context model_build_context();
+
+    bool build_core_runners();
+
+    bool build_vae_runners();
+
+    bool build_control_net_runner();
+
+    bool build_extension_runners();
+
+    bool validate_and_load_runners();
+
+    bool build_denoiser();
+
+    bool build_runners(const RunnerGroups& groups);
+
+    bool is_using_v_parameterization_for_sd2(bool is_inpaint = false);
+
+    static std::string lora_log_id(const ModelManager::LoraSpec& lora);
+
+    std::shared_ptr<LoraModel> load_lora_model(const ModelManager::LoraSpec& lora_spec,
+                                               SDBackendModule module,
+                                               LoraModel::filter_t module_filter = nullptr);
+
+    void clear_lora_adapters();
+
+    std::vector<std::shared_ptr<LoraModel>> load_runtime_loras_for_module(const std::vector<ModelManager::LoraSpec>& loras,
+                                                                          const std::set<std::string>& model_tensor_names,
+                                                                          SDBackendModule module,
+                                                                          LoraModel::filter_t module_filter,
+                                                                          bool& success,
+                                                                          std::vector<RuntimeLora>& next_models);
+
+    bool apply_loras_immediately(const std::vector<ModelManager::LoraSpec>& loras);
+
+    bool apply_loras_at_runtime(const std::vector<ModelManager::LoraSpec>& loras);
+
+    void lora_stat();
+
+    bool apply_loras(const sd_lora_t* loras, uint32_t lora_count);
+
+    void reset_generation_extensions();
+
+    void prepare_generation_extensions(const sd_pm_params_t& pm_params,
+                                       const sd_pulid_params_t& pulid_params,
+                                       ConditionerParams& condition_params,
+                                       int total_steps);
+
+    sd::Tensor<float> get_clip_vision_output(const sd::Tensor<float>& image,
+                                             bool return_pooled   = true,
+                                             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,
+                                         const sd::Tensor<float>& init_latent,
+                                         const sd::Tensor<float>& denoise_mask,
+                                         int step);
+
+    std::vector<float> process_ltxav_video_timesteps(const std::vector<float>& timesteps,
+                                                     const sd::Tensor<float>& init_latent,
+                                                     const sd::Tensor<float>& denoise_mask);
+
+    void preview_image(int step,
+                       const sd::Tensor<float>& latents,
+                       enum SDVersion version,
+                       preview_t preview_mode,
+                       std::function<void(int, int, sd_image_t*, bool, void*)> step_callback,
+                       void* step_callback_data,
+                       bool is_noisy);
+
+    std::vector<float> prepare_sample_timesteps(float sigma,
+                                                int shifted_timestep);
+
+    void adjust_sample_step_scalings(int shifted_timestep,
+                                     const std::vector<float>& timesteps_vec,
+                                     float c_in,
+                                     float* c_skip,
+                                     float* c_out);
+
+    struct SamplePreviewContext {
+        sd_preview_cb_t callback = nullptr;
+        void* data               = nullptr;
+        preview_t mode           = PREVIEW_NONE;
+    };
+
+    SamplePreviewContext prepare_sample_preview_context();
+
+    void report_sample_progress(int step,
+                                size_t total_steps,
+                                bool terminal_sigma_is_zero,
+                                int64_t* last_progress_us);
+
+    void compute_sample_controls(const sd::Tensor<float>& control_image,
+                                 const sd::Tensor<float>& noised_input,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const SDCondition& condition,
+                                 std::vector<sd::Tensor<float>>* controls);
+
+    sd::Tensor<float> sample(const std::shared_ptr<DiffusionModelRunner>& work_diffusion_model,
+                             bool inverse_noise_scaling,
+                             const sd::Tensor<float>& init_latent,
+                             sd::Tensor<float> noise,
+                             const SDCondition& cond,
+                             const SDCondition& uncond,
+                             const SDCondition& img_uncond,
+                             const sd::Tensor<float>& control_image,
+                             float control_strength,
+                             const sd_guidance_params_t& guidance,
+                             float eta,
+                             int shifted_timestep,
+                             sample_method_t method,
+                             bool is_flow_denoiser,
+                             const char* extra_sample_args,
+                             const std::vector<float>& sigmas,
+                             const std::vector<sd::Tensor<float>>& ref_latents,
+                             const RefImageParams& ref_image_params,
+                             const sd::Tensor<float>& denoise_mask,
+                             const sd::Tensor<float>& vace_context,
+                             float vace_strength,
+                             int audio_length,
+                             float frame_rate,
+                             const sd_cache_params_t* cache_params,
+                             bool preview_final_step,
+                             const sd::Tensor<float>& video_positions = {});
+
+    int get_vae_scale_factor();
+
+    int get_diffusion_model_down_factor();
+
+    int get_latent_channel();
+
+    int get_image_channels() const;
+
+    int get_image_seq_len(int h, int w);
+
+    sd::Tensor<float> generate_init_latent(int width,
+                                           int height,
+                                           int frames = 1,
+                                           bool video = false);
+
+    int video_frames_to_latent_frames(int frames);
+
+    int latent_frames_to_video_frames(int latent_frames);
+
+    int align_video_frames(int frames);
+
+    sd::Tensor<float> encode_to_vae_latents(const sd::Tensor<float>& x);
+
+    sd::Tensor<float> encode_first_stage(const sd::Tensor<float>& x);
+
+    sd::Tensor<float> decode_first_stage(const sd::Tensor<float>& x, bool decode_video = false);
+
+    sd::Tensor<float> normalize_ltx_video_latents(const sd::Tensor<float>& x);
+
+    sd::Tensor<float> un_normalize_ltx_video_latents(const sd::Tensor<float>& x);
+
+    sd::Tensor<float> decode_ltx_audio_latent(const sd::Tensor<float>& audio_latent);
+
+    void set_flow_shift(float flow_shift = INFINITY);
+
+    bool is_flow_denoiser();
+
+    std::string get_default_ref_image_preset(SDVersion version) const;
+
+    RefImageParams resolve_ref_image_params(const char* ref_image_args) const;
+};
+
+#endif  // __SD_PIPELINE_DIFFUSION_ENGINE_H__

src/pipeline/generation.h

@@ -0,0 +1,75 @@
+#ifndef __SD_PIPELINE_GENERATION_H__
+#define __SD_PIPELINE_GENERATION_H__
+
+#include "conditioning/conditioner.hpp"
+#include "stable-diffusion.h"
+
+class StableDiffusionGGML;
+
+static inline bool sd_version_supports_animatediff(SDVersion version) {
+    return version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX;
+}
+
+namespace sd::pipeline {
+
+    struct ImageGenerationLatents {
+        sd::Tensor<float> init_latent;
+        sd::Tensor<float> concat_latent;
+        sd::Tensor<float> img_uncond_concat_latent;
+        sd::Tensor<float> audio_latent;
+        sd::Tensor<float> video_positions;
+        sd::Tensor<float> control_image;
+        std::vector<sd::Tensor<float>> ref_images;
+        std::vector<sd::Tensor<float>> ref_latents;
+        std::vector<sd::Tensor<float>> reference_audio_latents;
+        std::vector<MiniMaxH3ReferenceBlock> minimax_reference_blocks;
+        std::vector<MiniMaxH3PresentationItem> minimax_presentation_refs;
+        std::vector<int32_t> keyframe_indices;
+        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;
+        int audio_length                       = 0;
+    };
+
+    struct ImageGenerationEmbeds {
+        SDCondition cond;
+        SDCondition uncond;
+        SDCondition img_uncond;
+    };
+
+    struct ConditionerRunnerEndOnExit {
+        Conditioner* conditioner = nullptr;
+        ~ConditionerRunnerEndOnExit() {
+            if (conditioner != nullptr) {
+                conditioner->runner_end();
+            }
+        }
+    };
+
+    // Callers hold ExecutionScope; AnimateDiff reuses the image path within the same scope.
+    bool generate_image(StableDiffusionGGML* sd,
+                        const sd_img_gen_params_t* sd_img_gen_params,
+                        sd_image_t** images_out,
+                        int* num_images_out);
+
+    bool generate_video(StableDiffusionGGML* sd,
+                        const sd_vid_gen_params_t* sd_vid_gen_params,
+                        sd_image_t** frames_out,
+                        int* num_frames_out,
+                        sd_audio_t** audio_out,
+                        int* fps_out);
+
+    sd::Tensor<float> upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd,
+                                                       const char* model_path,
+                                                       const sd::Tensor<float>& packed_latent,
+                                                       int audio_length);
+
+    sd::Tensor<float> ensure_image_tensor_channels(sd::Tensor<float> image, int channels);
+
+}  // namespace sd::pipeline
+
+#endif  // __SD_PIPELINE_GENERATION_H__

src/pipeline/model_builders.cpp

@@ -0,0 +1,643 @@
+#include "model_builders.h"
+
+#include <cstring>
+#include <utility>
+
+#include "conditioning/conditioner.hpp"
+#include "core/ggml_extend_backend.h"
+#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"
+#include "model/diffusion/ernie_image.hpp"
+#include "model/diffusion/flux.hpp"
+#include "model/diffusion/hidream_o1.hpp"
+#include "model/diffusion/hunyuan.hpp"
+#include "model/diffusion/ideogram4.hpp"
+#include "model/diffusion/krea2.hpp"
+#include "model/diffusion/lens.hpp"
+#include "model/diffusion/lingbot_video.hpp"
+#include "model/diffusion/ltxv.hpp"
+#include "model/diffusion/mage_flow.hpp"
+#include "model/diffusion/minimax_h3.hpp"
+#include "model/diffusion/minit2i.hpp"
+#include "model/diffusion/mmdit.hpp"
+#include "model/diffusion/model.hpp"
+#include "model/diffusion/pid.hpp"
+#include "model/diffusion/qwen_image.hpp"
+#include "model/diffusion/sensenova_u1.h"
+#include "model/diffusion/unet.hpp"
+#include "model/diffusion/wan.hpp"
+#include "model/diffusion/z_image.hpp"
+#include "model/vae/auto_encoder_kl.hpp"
+#include "model/vae/hunyuan_vae.hpp"
+#include "model/vae/ltx_audio_vae.hpp"
+#include "model/vae/ltx_vae.hpp"
+#include "model/vae/mage_vae.hpp"
+#include "model/vae/minimax_h3_audio_vae.hpp"
+#include "model/vae/minimax_h3_vae.hpp"
+#include "model/vae/tae.hpp"
+#include "model/vae/vae.hpp"
+#include "model/vae/wan_vae.hpp"
+
+namespace sd::model_builders {
+
+    static bool ensure_backend_pair(SDBackendManager& backends, SDBackendModule module) {
+        if (backends.runtime_backend(module) == nullptr) {
+            LOG_ERROR("failed to initialize %s backend", sd_backend_module_name(module));
+            return false;
+        }
+        if (backends.params_backend(module) == nullptr) {
+            LOG_ERROR("failed to initialize %s params backend", sd_backend_module_name(module));
+            return false;
+        }
+        return true;
+    }
+
+    static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback) {
+        switch (format) {
+            case SD_VAE_FORMAT_FLUX:
+                return VERSION_FLUX;
+            case SD_VAE_FORMAT_SD3:
+                return VERSION_SD3;
+            case SD_VAE_FORMAT_FLUX2:
+                return VERSION_FLUX2;
+            case SD_VAE_FORMAT_WAN:
+                return VERSION_WAN2;
+            case SD_VAE_FORMAT_AUTO:
+            default:
+                return fallback;
+        }
+    }
+
+    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;
+        }
+
+        if (sd_version_is_sd3(version)) {
+            result.conditioner = std::make_shared<SD3CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                   tensor_storage_map,
+                                                                   weight_manager,
+                                                                   tokenizers);
+            result.diffusion   = std::make_shared<MMDiTRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                             tensor_storage_map,
+                                                             "model.diffusion_model",
+                                                             weight_manager);
+        } else if (sd_version_is_pid(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<Pid::PiDRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                tensor_storage_map,
+                                                                "model.diffusion_model.net",
+                                                                weight_manager);
+        } else if (sd_version_is_ideogram4(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<Ideogram4::Ideogram4Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                            tensor_storage_map,
+                                                                            "model.diffusion_model",
+                                                                            weight_manager);
+        } else if (sd_version_is_krea2(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               true,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<Krea2::Krea2Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                    tensor_storage_map,
+                                                                    "model.diffusion_model",
+                                                                    weight_manager);
+        } else if (sd_version_is_flux(version)) {
+            bool is_chroma = false;
+            for (auto pair : tensor_storage_map) {
+                if (pair.first.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
+                    is_chroma = true;
+                    break;
+                }
+            }
+            if (is_chroma) {
+                result.conditioner = std::make_shared<T5CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                      tensor_storage_map,
+                                                                      false,
+                                                                      1,
+                                                                      false,
+                                                                      weight_manager,
+                                                                      sd_ctx_params->model_args);
+            } else if (version == VERSION_OVIS_IMAGE) {
+                result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                   tensor_storage_map,
+                                                                   version,
+                                                                   "",
+                                                                   false,
+                                                                   weight_manager,
+                                                                   tokenizers);
+            } else {
+                result.conditioner = std::make_shared<FluxCLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                        tensor_storage_map,
+                                                                        weight_manager,
+                                                                        tokenizers);
+            }
+            result.diffusion = std::make_shared<Flux::FluxRunner>(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_flux2(version) || sd_version_is_sefi_image(version)) {
+            bool is_chroma     = false;
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<Flux::FluxRunner>(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_ltxav(version)) {
+            result.conditioner = std::make_shared<LTXAVEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                 tensor_storage_map,
+                                                                 "text_encoders.llm",
+                                                                 "text_embedding_projection",
+                                                                 weight_manager,
+                                                                 tokenizers);
+            result.diffusion   = std::make_shared<LTXV::LTXAVRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                   tensor_storage_map,
+                                                                   "model.diffusion_model",
+                                                                   weight_manager);
+        } else if (sd_version_is_minimax_h3(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               true,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<MiniMaxH3::MiniMaxH3Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                            tensor_storage_map,
+                                                                            "model.diffusion_model",
+                                                                            weight_manager);
+        } else if (sd_version_is_hunyuan_video(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<Hunyuan::HunyuanVideoRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                             tensor_storage_map,
+                                                                             "model.diffusion_model",
+                                                                             version,
+                                                                             weight_manager);
+        } else if (sd_version_is_wan(version)) {
+            result.conditioner = std::make_shared<T5CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                  tensor_storage_map,
+                                                                  true,
+                                                                  0,
+                                                                  true,
+                                                                  weight_manager);
+            result.diffusion   = std::make_shared<WAN::WanRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                tensor_storage_map,
+                                                                "model.diffusion_model",
+                                                                version,
+                                                                weight_manager);
+            if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) {
+                result.high_noise_diffusion = std::make_shared<WAN::WanRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                               tensor_storage_map,
+                                                                               "model.high_noise_diffusion_model",
+                                                                               version,
+                                                                               weight_manager);
+            }
+            if (result.diffusion->get_desc() == "Wan2.1-I2V-14B" ||
+                result.diffusion->get_desc() == "Wan2.1-FLF2V-14B" ||
+                result.diffusion->get_desc() == "Wan2.1-I2V-1.3B") {
+                if (!ensure_backend_pair(ctx.backends, SDBackendModule::CLIP_VISION)) {
+                    return false;
+                }
+                result.clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(ctx.backends.runtime_backend(SDBackendModule::CLIP_VISION),
+                                                                                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) {
+                if (starts_with(name, "text_encoders.llm.visual.")) {
+                    enable_vision = true;
+                    break;
+                }
+            }
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               enable_vision,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<LingBotVideo::LingBotVideoRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                                  tensor_storage_map,
+                                                                                  "model.diffusion_model",
+                                                                                  weight_manager,
+                                                                                  sd_ctx_params->model_args);
+        } else if (sd_version_is_qwen_image(version)) {
+            bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED;
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               enable_vision,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                       tensor_storage_map,
+                                                                       "model.diffusion_model",
+                                                                       version,
+                                                                       weight_manager,
+                                                                       sd_ctx_params->model_args);
+        } 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,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<MageFlow::MageFlowRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                          tensor_storage_map,
+                                                                          "model.diffusion_model",
+                                                                          weight_manager);
+        } else if (sd_version_is_longcat(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               true,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                  tensor_storage_map,
+                                                                  "model.diffusion_model",
+                                                                  version,
+                                                                  weight_manager,
+                                                                  sd_ctx_params->model_args);
+        } else if (version == VERSION_HIDREAM_O1) {
+            result.conditioner = std::make_shared<HiDreamO1::HiDreamO1Conditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                                   tensor_storage_map,
+                                                                                   weight_manager,
+                                                                                   tokenizers);
+            result.diffusion   = std::make_shared<HiDreamO1::HiDreamO1Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                            tensor_storage_map,
+                                                                            "model",
+                                                                            weight_manager);
+        } else if (sd_version_is_minit2i(version)) {
+            result.conditioner = std::make_shared<MiniT2IConditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                      tensor_storage_map,
+                                                                      weight_manager);
+            result.diffusion   = std::make_shared<MiniT2I::MiniT2IRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                        tensor_storage_map,
+                                                                        "model.diffusion_model.model.net",
+                                                                        weight_manager);
+        } else if (sd_version_is_sensenova_u1(version)) {
+            result.conditioner = std::make_shared<SenseNovaU1Conditioner>();
+            result.diffusion   = std::make_shared<SenseNovaU1::SenseNovaU1Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                                tensor_storage_map,
+                                                                                "",
+                                                                                weight_manager);
+        } else if (sd_version_is_anima(version)) {
+            result.conditioner = std::make_shared<AnimaConditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                    tensor_storage_map,
+                                                                    weight_manager,
+                                                                    tokenizers);
+            result.diffusion   = std::make_shared<Anima::AnimaRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                    tensor_storage_map,
+                                                                    "model.diffusion_model",
+                                                                    weight_manager);
+        } else if (sd_version_is_z_image(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<ZImage::ZImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                      tensor_storage_map,
+                                                                      "model.diffusion_model",
+                                                                      version,
+                                                                      weight_manager);
+        } else if (sd_version_is_boogu_image(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               true,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<Boogu::BooguImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                         tensor_storage_map,
+                                                                         "model.diffusion_model",
+                                                                         version,
+                                                                         weight_manager);
+        } else if (sd_version_is_ernie_image(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<ErnieImage::ErnieImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                              tensor_storage_map,
+                                                                              "model.diffusion_model",
+                                                                              weight_manager);
+        } else if (sd_version_is_lens(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map,
+                                                               version,
+                                                               "",
+                                                               false,
+                                                               weight_manager,
+                                                               tokenizers);
+            result.diffusion   = std::make_shared<Lens::LensRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                  tensor_storage_map,
+                                                                  "model.diffusion_model",
+                                                                  weight_manager);
+        } else {  // SD1.x SD2.x SDXL
+            std::map<std::string, std::string> embbeding_map;
+            for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) {
+                embbeding_map.emplace(SAFE_STR(sd_ctx_params->embeddings[i].name), SAFE_STR(sd_ctx_params->embeddings[i].path));
+            }
+            result.conditioner = std::make_shared<FrozenCLIPEmbedderWithCustomWords>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                                     tensor_storage_map,
+                                                                                     embbeding_map,
+                                                                                     version,
+                                                                                     weight_manager,
+                                                                                     tokenizers);
+            result.diffusion   = std::make_shared<UNetModelRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                 tensor_storage_map,
+                                                                 "model.diffusion_model",
+                                                                 version,
+                                                                 weight_manager);
+            if (sd_ctx_params->diffusion_conv_direct) {
+                LOG_INFO("Using Conv2d direct in the diffusion model");
+                result.diffusion->set_conv2d_direct_enabled(true);
+            }
+        }
+
+        if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0 && result.clip_vision == nullptr) {
+            if (!ensure_backend_pair(ctx.backends, SDBackendModule::CLIP_VISION)) {
+                return false;
+            }
+            result.clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(ctx.backends.runtime_backend(SDBackendModule::CLIP_VISION),
+                                                                            tensor_storage_map,
+                                                                            weight_manager);
+        }
+
+        if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) {
+            result.ip_adapter = std::make_shared<IPAdapter::IPAdapterRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                             tensor_storage_map,
+                                                                             "ip_adapter",
+                                                                             weight_manager);
+        }
+        if (result.conditioner) {
+            result.conditioner->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
+        }
+        if (result.diffusion) {
+            result.diffusion->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
+        }
+        if (result.high_noise_diffusion) {
+            result.high_noise_diffusion->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
+        }
+        if (result.clip_vision) {
+            result.clip_vision->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
+        }
+        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) {
+        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;
+        VAERunners result;
+        if (!ensure_backend_pair(ctx.backends, SDBackendModule::VAE)) {
+            return false;
+        }
+
+        auto create_tae = [&](bool decode_only) -> std::shared_ptr<VAE> {
+            if (sd_version_uses_wan_vae(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) {
+                return std::make_shared<TinyVideoAutoEncoder>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                              tensor_storage_map,
+                                                              "decoder",
+                                                              decode_only,
+                                                              version,
+                                                              weight_manager);
+
+            } else {
+                auto model = std::make_shared<TinyImageAutoEncoder>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                                    tensor_storage_map,
+                                                                    "decoder.layers",
+                                                                    decode_only,
+                                                                    version,
+                                                                    weight_manager);
+                return model;
+            }
+        };
+
+        sd_vae_format_t vae_format = sd_ctx_params->vae_format;
+        if (vae_format < SD_VAE_FORMAT_AUTO || vae_format >= SD_VAE_FORMAT_COUNT) {
+            LOG_WARN("invalid VAE format override, using auto");
+            vae_format = SD_VAE_FORMAT_AUTO;
+        }
+        SDVersion vae_version = version;
+        if (sd_version_is_pid(version) && vae_format != SD_VAE_FORMAT_AUTO) {
+            vae_version = sd_vae_format_to_version(vae_format, vae_version);
+        }
+
+        auto create_vae = [&]() -> std::shared_ptr<VAE> {
+            if (sd_version_is_ltxav(version)) {
+                return std::make_shared<LTXVideoVAE>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                     tensor_storage_map,
+                                                     "first_stage_model",
+                                                     false,
+                                                     version,
+                                                     weight_manager);
+            } else if (sd_version_is_minimax_h3(version)) {
+                return std::make_shared<MiniMaxH3VAE::MiniMaxH3VideoVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                                               tensor_storage_map,
+                                                                               "first_stage_model",
+                                                                               weight_manager);
+            } else if (sd_version_is_mage_flow(vae_version)) {
+                return std::make_shared<MageVAE::MageVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                                tensor_storage_map,
+                                                                "first_stage_model",
+                                                                weight_manager);
+            } else if (sd_version_uses_hunyuan_video_vae(vae_version)) {
+                return std::make_shared<Hunyuan::HunyuanVideoVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                                        tensor_storage_map,
+                                                                        "first_stage_model",
+                                                                        false,
+                                                                        vae_version,
+                                                                        weight_manager);
+            } else if (sd_version_uses_wan_vae(vae_version)) {
+                return std::make_shared<WAN::WanVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                           tensor_storage_map,
+                                                           "first_stage_model",
+                                                           false,
+                                                           vae_version,
+                                                           weight_manager);
+            } else {
+                auto model = std::make_shared<AutoEncoderKL>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                             tensor_storage_map,
+                                                             "first_stage_model",
+                                                             false,
+                                                             false,
+                                                             vae_version,
+                                                             weight_manager);
+                if (sd_version_is_sdxl(version) &&
+                    (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || options.external_vae_is_invalid)) {
+                    float vae_conv_2d_scale = 1.f / 32.f;
+                    LOG_WARN(
+                        "No valid VAE specified with --vae or --force-sdxl-vae-conv-scale flag set, "
+                        "using Conv2D scale %.3f",
+                        vae_conv_2d_scale);
+                    model->set_conv2d_scale(vae_conv_2d_scale);
+                }
+                return model;
+            }
+        };
+
+        if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version) || sd_version_is_sensenova_u1(version)) {
+            LOG_INFO("using FakeVAE");
+            result.vae = std::make_shared<FakeVAE>(version,
+                                                   ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                   weight_manager);
+        } else if (options.use_tae && !options.tae_preview_only) {
+            LOG_INFO("using TAE for encoding / decoding");
+            result.vae = create_tae(false);
+        } else {
+            LOG_INFO("using VAE for encoding / decoding");
+            result.vae = create_vae();
+            if (options.use_tae && options.tae_preview_only) {
+                LOG_INFO("using TAE for preview");
+                result.preview = create_tae(true);
+            }
+        }
+
+        if (options.use_audio_vae) {
+            if (sd_version_is_minimax_h3(version)) {
+                result.audio = std::make_shared<MiniMaxH3::AudioVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                                           tensor_storage_map,
+                                                                           "",
+                                                                           weight_manager);
+            } else {
+                result.audio = std::make_shared<LTXV::LTXAudioVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
+                                                                         tensor_storage_map,
+                                                                         "",
+                                                                         weight_manager);
+            }
+        }
+
+        if (sd_ctx_params->vae_conv_direct) {
+            LOG_INFO("Using Conv2d direct in the vae model");
+            result.vae->set_conv2d_direct_enabled(true);
+            if (result.preview) {
+                result.preview->set_conv2d_direct_enabled(true);
+            }
+        }
+        if (result.vae) {
+            result.vae->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
+        }
+        if (result.preview) {
+            result.preview->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
+        }
+        if (result.audio) {
+            result.audio->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
+        }
+        runners = std::move(result);
+        return true;
+    }
+
+    bool build_control_net_runner(const Context& ctx, std::shared_ptr<ControlNet>& runner) {
+        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;
+        if (!ensure_backend_pair(ctx.backends, SDBackendModule::CONTROL_NET)) {
+            return false;
+        }
+        auto control_net = std::make_shared<ControlNet>(ctx.backends.runtime_backend(SDBackendModule::CONTROL_NET),
+                                                        tensor_storage_map,
+                                                        version,
+                                                        "",
+                                                        weight_manager);
+        if (sd_ctx_params->diffusion_conv_direct) {
+            LOG_INFO("Using Conv2d direct in the control net");
+            control_net->set_conv2d_direct_enabled(true);
+        }
+        control_net->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
+        runner = std::move(control_net);
+        return true;
+    }
+
+    bool build_extension_runners(const GenerationExtensionInitContext& ctx,
+                                 std::vector<std::shared_ptr<GenerationExtension>>& extensions) {
+        std::vector<std::shared_ptr<GenerationExtension>> result;
+        for (auto extension : {create_photomaker_extension(), create_pulid_extension()}) {
+            if (!extension->init(ctx)) {
+                return false;
+            }
+            if (extension->is_enabled()) {
+                result.push_back(std::move(extension));
+            }
+        }
+        extensions = std::move(result);
+        return true;
+    }
+
+}  // namespace sd::model_builders

src/pipeline/model_builders.h

@@ -0,0 +1,67 @@
+#ifndef __SD_PIPELINE_MODEL_BUILDERS_H__
+#define __SD_PIPELINE_MODEL_BUILDERS_H__
+
+#include <memory>
+#include <vector>
+
+#include "model.h"
+#include "stable-diffusion.h"
+
+class SDBackendManager;
+struct DeviceResidencyManager;
+struct Conditioner;
+struct FrozenCLIPVisionEmbedder;
+struct DiffusionModelRunner;
+struct VAE;
+struct AudioVAERunner;
+struct ControlNet;
+namespace Wav2Vec2 {
+    class Wav2Vec2ModelRunner;
+}
+struct GenerationExtension;
+struct GenerationExtensionInitContext;
+namespace IPAdapter {
+    struct IPAdapterRunner;
+}
+
+namespace sd::model_builders {
+
+    struct Context {
+        const sd_ctx_params_t& params;
+        SDVersion version;
+        const String2TensorStorage& tensor_storage_map;
+        SDBackendManager& backends;
+        std::shared_ptr<DeviceResidencyManager> weight_manager;
+    };
+
+    struct CoreRunners {
+        std::shared_ptr<Conditioner> conditioner;
+        std::shared_ptr<DiffusionModelRunner> diffusion;
+        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 {
+        bool use_tae                 = false;
+        bool tae_preview_only        = false;
+        bool use_audio_vae           = false;
+        bool external_vae_is_invalid = false;
+    };
+
+    struct VAERunners {
+        std::shared_ptr<VAE> vae;
+        std::shared_ptr<VAE> preview;
+        std::shared_ptr<AudioVAERunner> audio;
+    };
+
+    bool build_core_runners(const Context& ctx, CoreRunners& runners);
+    bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners);
+    bool build_control_net_runner(const Context& ctx, std::shared_ptr<ControlNet>& runner);
+    bool build_extension_runners(const GenerationExtensionInitContext& ctx,
+                                 std::vector<std::shared_ptr<GenerationExtension>>& extensions);
+
+}  // namespace sd::model_builders
+
+#endif  // __SD_PIPELINE_MODEL_BUILDERS_H__