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
 
@@ -126,32 +127,87 @@ Direct ("immediately") LoRA application cannot patch row-split tensors; with
 explicit `--lora-apply-mode immediately` skips the split tensors with a
 warning.
 
-## Automatic placement (`--auto-fit`)
+## Automatic placement (`--auto-fit on|off`)
 
-`--auto-fit` derives the `diffusion` / `te` / `vae` placements from the model
-metadata and the per-device memory budgets, then feeds them into the same
-backend assignment mechanism described above (the chosen specs are printed).
-`--backend` and `--params-backend` are ignored while auto-fit is enabled.
+`--auto-fit` requires `on` or `off` and defaults to `on` when omitted.
+Explicit `--params-backend` assignments disable auto-fit,
+regardless of argument order, even with `--auto-fit on`.
+
+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
-sd-cli -m model.safetensors -p "a cat" --auto-fit --max-vram cuda0=8,cuda1=14
-sd-cli -m model.safetensors -p "a cat" --auto-fit --split-mode row
+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
 ```
 
 Budgets reuse `--max-vram`: a positive per-device value caps what auto-fit
 plans with on that device, a negative value means "free memory minus that many
 GiB", and with no budget set each device's free memory minus a 512 MiB margin
-is used. (The same values still drive graph-cut segmented execution for
-modules that end up on a single device.)
+is used. These resolved GPU budgets, including the safety margin, also drive
+the runner's graph-cut capacity checks.
 
-When everything fits resident, components are simply spread across the
-available GPUs. When it does not, auto-fit switches to time-share mode: the
-heavy components get `disk` params residency (loaded for their phase, freed
-after), and a component too large for any single device is split across all
-GPUs with the layer/row split mechanism (`--split-mode` selects which, layer
-by default). Components that fit nowhere fall back to the CPU. If a VAE decode
-still runs out of memory, tiling is enabled and the decode retried once.
+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 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.
+
+GPU cache space follows the same component priority. Before a lower-priority
+component can become permanently resident, the planner leaves room for the full
+weights and estimated compute space of higher-priority offloaded components.
+If offloaded diffusion already needs the entire main GPU budget, TE and VAE also
+use offloaded parameters. Their GPU copies can then be released after their
+phases, leaving more room to reuse diffusion weights across sampling steps.
+CPU parameter residency allows GPU weight caching; it does not force every
+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.
+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
+sufficient. Automatic graph segmentation can then load the required weights
+for each segment and reclaim idle GPU copies. `--disable-segmented-compute`
+still disables segmentation.
+
+Initial compute reserves are estimates (2 GiB for diffusion and text encoders,
+1 GiB for VAE); higher-priority placements also leave staging space for the
+largest weight tensor of each lower-priority offloaded component. Actual segment
+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, 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
 
@@ -203,7 +259,7 @@ sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend disk
 
 This runs all modules on `cuda0`, reloads parameters from the model file as needed, and releases those parameter buffers after use.
 
-`disk` is never selected implicitly. If `--params-backend` is not set, parameters use the runtime backend.
+Outside `--auto-fit`, `disk` is never selected implicitly. If `--params-backend` is not set, parameters use the runtime backend.
 
 Per-module assignments can be mixed:
 
@@ -252,4 +308,8 @@ 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. The old CPU/offload fields are no longer part of the C API. Explicit `--backend` and `--params-backend` assignments are preferred for new commands.
+Library callers should set `backend` and `params_backend` directly. `sd_ctx_params_init()`
+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/model_config.md

@@ -67,21 +67,21 @@ Detection should respect `prefix`. For nested weights, construct full names from
 
 Do not add persistent config fields such as `inferred_from_weights` only to
 record whether detection happened. If the function needs to decide whether to
-print a debug line, keep that as local control flow inside `detect_from_weights`.
+print a verbose line, keep that as local control flow inside `detect_from_weights`.
 
 ## Logging
 
-When config values are inferred from weights, print one `LOG_DEBUG` line at the
+When config values are inferred from weights, print one `LOG_VERBOSE` line at the
 end of `detect_from_weights`.
 
 Example:
 
 ```cpp
-LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
-          config.num_layers,
-          config.vocab_size,
-          config.hidden_size,
-          config.intermediate_size);
+LOG_VERBOSE("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
+            config.num_layers,
+            config.vocab_size,
+            config.hidden_size,
+            config.intermediate_size);
 ```
 
 Only print the config detection log when the function actually inferred values

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

@@ -6,6 +6,11 @@ For detailed command-line arguments, run:
 ./bin/sd-cli -h
 ```
 
+Logging defaults to `info`. Use `--log-level <level>` to select `debug`, `verbose`,
+`info`, `warn`, or `error` (from most to least detailed). Each level includes
+messages at that level and all less detailed levels. `-v` and `--verbose` are
+equivalent to `--log-level verbose`. If repeated, the last logging option wins.
+
 For direct image repair or automatic post-generation YOLOv8 detection followed by cropped inpainting, see
 [ADetailer](../../docs/adetailer.md).
 
@@ -17,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

@@ -40,9 +40,9 @@ struct SDCliParams {
     std::string image_path;
     std::string metadata_format = "text";
 
-    bool verbose          = false;
-    bool canny_preprocess = false;
-    bool convert_name     = false;
+    sd_log_level_t log_level = SD_LOG_INFO;
+    bool canny_preprocess    = false;
+    bool convert_name        = false;
 
     preview_t preview_method = PREVIEW_NONE;
     int preview_interval     = 1;
@@ -115,10 +115,6 @@ struct SDCliParams {
              "--convert-name",
              "convert tensor name (for convert mode)",
              true, &convert_name},
-            {"-v",
-             "--verbose",
-             "print extra info",
-             true, &verbose},
             {"",
              "--color",
              "colors the logging tags according to level",
@@ -220,6 +216,7 @@ struct SDCliParams {
              on_imatrix_in_arg},
         };
 
+        add_log_options(options, log_level);
         return options;
     };
 
@@ -269,7 +266,7 @@ struct SDCliParams {
             << "  output_path: \"" << output_path << "\",\n"
             << "  image_path: \"" << image_path << "\",\n"
             << "  metadata_format: \"" << metadata_format << "\",\n"
-            << "  verbose: " << (verbose ? "true" : "false") << ",\n"
+            << "  log_level: " << log_level_name(log_level) << ",\n"
             << "  color: " << (color ? "true" : "false") << ",\n"
             << "  canny_preprocess: " << (canny_preprocess ? "true" : "false") << ",\n"
             << "  convert_name: " << (convert_name ? "true" : "false") << ",\n"
@@ -307,6 +304,9 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP
         exit(cli_params.normal_exit ? 0 : 1);
     }
 
+    log_level = cli_params.log_level;
+    log_color = cli_params.color;
+
     bool valid = cli_params.resolve_and_validate();
     if (valid && cli_params.mode != METADATA) {
         valid = ctx_params.resolve_and_validate(cli_params.mode) &&
@@ -323,15 +323,14 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP
 
 void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
     SDCliParams* cli_params = (SDCliParams*)data;
-    log_print(level, log, cli_params->verbose, cli_params->color);
+    log_print(level, log, cli_params->log_level, cli_params->color);
 }
 
 bool load_images_from_dir(const std::string dir,
                           std::vector<SDImageOwner>& images,
                           int expected_width  = 0,
                           int expected_height = 0,
-                          int max_image_num   = 0,
-                          bool verbose        = false) {
+                          int max_image_num   = 0) {
     if (!fs::exists(dir) || !fs::is_directory(dir)) {
         LOG_ERROR("'%s' is not a valid directory\n", dir.c_str());
         return false;
@@ -355,7 +354,7 @@ bool load_images_from_dir(const std::string dir,
         std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
 
         if (ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".bmp" || ext == ".webp") {
-            LOG_DEBUG("load image %zu from '%s'", images.size(), path.c_str());
+            LOG_VERBOSE("load image %zu from '%s'", images.size(), path.c_str());
             int width             = 0;
             int height            = 0;
             uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, expected_width, expected_height);
@@ -420,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());
         }
     }
@@ -541,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;
@@ -651,8 +655,6 @@ int main(int argc, const char* argv[]) {
 
     parse_args(argc, argv, cli_params, ctx_params, gen_params);
     sd_set_log_callback(sd_log_cb, (void*)&cli_params);
-    log_verbose = cli_params.verbose;
-    log_color   = cli_params.color;
 
     if (cli_params.mode == METADATA) {
         MetadataReadOptions options;
@@ -690,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,
@@ -700,11 +700,11 @@ int main(int argc, const char* argv[]) {
                             cli_params.preview_noisy,
                             (void*)&cli_params);
 
-    LOG_DEBUG("version: %s", version_string().c_str());
-    LOG_DEBUG("%s", sd_get_system_info());
-    LOG_DEBUG("%s", cli_params.to_string().c_str());
-    LOG_DEBUG("%s", ctx_params.to_string().c_str());
-    LOG_DEBUG("%s", gen_params.to_string().c_str());
+    LOG_VERBOSE("version: %s", version_string().c_str());
+    LOG_VERBOSE("%s", sd_get_system_info());
+    LOG_VERBOSE("%s", cli_params.to_string().c_str());
+    LOG_VERBOSE("%s", ctx_params.to_string().c_str());
+    LOG_VERBOSE("%s", gen_params.to_string().c_str());
 
     if (!cli_params.imatrix_out.empty()) {
         if (fs::exists(cli_params.imatrix_out) &&
@@ -808,7 +808,7 @@ int main(int argc, const char* argv[]) {
         gen_params.ref_videos.reserve(gen_params.ref_video_paths.size());
         for (const auto& path : gen_params.ref_video_paths) {
             std::vector<SDImageOwner> frames;
-            if (!load_images_from_dir(path, frames, 0, 0, 0, cli_params.verbose) || frames.empty()) {
+            if (!load_images_from_dir(path, frames) || frames.empty()) {
                 LOG_ERROR("load reference video frames from '%s' failed", path.c_str());
                 return 1;
             }
@@ -890,8 +890,7 @@ int main(int argc, const char* argv[]) {
                                   gen_params.control_frames,
                                   gen_params.get_resolved_width(),
                                   gen_params.get_resolved_height(),
-                                  gen_params.video_frames,
-                                  cli_params.verbose)) {
+                                  gen_params.video_frames)) {
             return 1;
         }
     }
@@ -902,8 +901,7 @@ int main(int argc, const char* argv[]) {
                                   gen_params.pm_id_images,
                                   0,
                                   0,
-                                  0,
-                                  cli_params.verbose)) {
+                                  0)) {
             return 1;
         }
     }
@@ -956,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

@@ -239,6 +239,26 @@ void ArgOptions::print() const {
     }
 }
 
+void add_log_options(ArgOptions& options, sd_log_level_t& level) {
+    options.manual_options.push_back({"", "--log-level",
+                                      "minimum log level, one of [debug, verbose, info, warn, error] (default: info)",
+                                      [&level](int argc, const char** argv, int index) {
+                                          if (++index >= argc) {
+                                              return -1;
+                                          }
+                                          if (!parse_log_level(argv[index], level)) {
+                                              LOG_ERROR("invalid log level %s, must be one of [debug, verbose, info, warn, error]", argv[index]);
+                                              return -1;
+                                          }
+                                          return 1;
+                                      }});
+    options.manual_options.push_back({"-v", "--verbose", "equivalent to --log-level verbose",
+                                      [&level](int, const char**, int) {
+                                          level = SD_LOG_VERBOSE;
+                                          return 0;
+                                      }});
+}
+
 bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& options_list) {
     bool invalid_arg = false;
     std::string arg;
@@ -282,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;
 
@@ -292,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;
 
@@ -317,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;
         }
@@ -330,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 = {
@@ -362,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",
@@ -412,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)",
@@ -528,12 +586,6 @@ ArgOptions SDContextParams::get_options() {
          "--eager-load",
          "load all params into the params backend at model-load time instead of lazily on first use (defaults to false)",
          true, &eager_load},
-        {"",
-         "--auto-fit",
-         "pick the diffusion/te/vae device placements automatically from the model size and the per-device "
-         "memory budgets (--max-vram; defaults to free memory minus a small margin). Overrides --backend and "
-         "--params-backend; may split modules across GPUs (--split-mode still selects layer or row)",
-         true, &auto_fit},
         {"",
          "--force-sdxl-vae-conv-scale",
          "force use of conv scale on sdxl vae",
@@ -576,6 +628,23 @@ ArgOptions SDContextParams::get_options() {
          true, &vae_conv_direct},
     };
 
+    auto on_auto_fit_arg = [&](int argc, const char** argv, int index) {
+        if (++index >= argc) {
+            LOG_ERROR("--auto-fit requires 'on' or 'off'");
+            return -1;
+        }
+        const std::string arg = argv[index];
+        if (arg == "on") {
+            auto_fit = true;
+        } else if (arg == "off") {
+            auto_fit = false;
+        } else {
+            LOG_ERROR("invalid --auto-fit value '%s'; expected 'on' or 'off'", argv[index]);
+            return -1;
+        }
+        return 1;
+    };
+
     auto on_type_arg = [&](int argc, const char** argv, int index) {
         if (++index >= argc) {
             return -1;
@@ -647,6 +716,24 @@ 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). 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 --params-backend; uses automatic graph segmentation when needed",
+         on_auto_fit_arg},
         {"",
          "--type",
          "weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K). "
@@ -814,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"
@@ -821,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"
@@ -849,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"
@@ -878,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();
@@ -902,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;
@@ -1434,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;
@@ -1623,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

@@ -107,6 +107,7 @@ struct ArgOptions {
     void print() const;
 };
 
+void add_log_options(ArgOptions& options, sd_log_level_t& level);
 bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& options_list);
 bool decode_base64_image(const std::string& encoded_input,
                          int target_channels,
@@ -123,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;
@@ -130,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;
@@ -157,7 +160,7 @@ struct SDContextParams {
     std::string params_backend;
     std::string split_mode;
     std::string model_args;
-    bool auto_fit = false;
+    bool auto_fit = true;
     std::string rpc_servers;
     std::string effective_backend;
     std::string effective_params_backend;
@@ -174,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/log.cpp

@@ -2,8 +2,8 @@
 
 #include <vector>
 
-bool log_verbose = false;
-bool log_color   = false;
+sd_log_level_t log_level = SD_LOG_INFO;
+bool log_color           = false;
 
 std::string sd_basename(const std::string& path) {
     size_t pos = path.find_last_of('/');
@@ -51,12 +51,40 @@ void print_utf8(FILE* stream, const char* utf8) {
 #endif
 }
 
-void log_print(enum sd_log_level_t level, const char* log, bool verbose, bool color) {
+const char* log_level_name(sd_log_level_t level) {
+    switch (level) {
+        case SD_LOG_DEBUG:
+            return "debug";
+        case SD_LOG_VERBOSE:
+            return "verbose";
+        case SD_LOG_INFO:
+            return "info";
+        case SD_LOG_WARN:
+            return "warn";
+        case SD_LOG_ERROR:
+            return "error";
+        default:
+            return "unknown";
+    }
+}
+
+bool parse_log_level(const std::string& name, sd_log_level_t& level) {
+    const sd_log_level_t levels[] = {SD_LOG_DEBUG, SD_LOG_VERBOSE, SD_LOG_INFO, SD_LOG_WARN, SD_LOG_ERROR};
+    for (sd_log_level_t candidate : levels) {
+        if (name == log_level_name(candidate)) {
+            level = candidate;
+            return true;
+        }
+    }
+    return false;
+}
+
+void log_print(enum sd_log_level_t level, const char* log, sd_log_level_t min_level, bool color) {
     int tag_color;
     const char* level_str;
     FILE* out_stream = (level == SD_LOG_ERROR) ? stderr : stdout;
 
-    if (!log || (!verbose && level <= SD_LOG_DEBUG)) {
+    if (!log || level < min_level) {
         return;
     }
 
@@ -65,6 +93,10 @@ void log_print(enum sd_log_level_t level, const char* log, bool verbose, bool co
             tag_color = 37;
             level_str = "DEBUG";
             break;
+        case SD_LOG_VERBOSE:
+            tag_color = 37;
+            level_str = "VERBOSE";
+            break;
         case SD_LOG_INFO:
             tag_color = 34;
             level_str = "INFO";
@@ -84,9 +116,9 @@ void log_print(enum sd_log_level_t level, const char* log, bool verbose, bool co
     }
 
     if (color) {
-        fprintf(out_stream, "\033[%d;1m[%-5s]\033[0m ", tag_color, level_str);
+        fprintf(out_stream, "\033[%d;1m[%-7s]\033[0m ", tag_color, level_str);
     } else {
-        fprintf(out_stream, "[%-5s] ", level_str);
+        fprintf(out_stream, "[%-7s] ", level_str);
     }
     fflush(out_stream);
     print_utf8(out_stream, log);
@@ -110,7 +142,7 @@ void example_log_printf(sd_log_level_t level, const char* file, int line, const
         strncat(log_buffer, "\n", LOG_BUFFER_SIZE - len);
     }
 
-    log_print(level, log_buffer, log_verbose, log_color);
+    log_print(level, log_buffer, log_level, log_color);
 
     va_end(args);
 }

examples/common/log.h

@@ -16,15 +16,18 @@
 
 #include "stable-diffusion.h"
 
-extern bool log_verbose;
+extern sd_log_level_t log_level;
 extern bool log_color;
 
 std::string sd_basename(const std::string& path);
 void print_utf8(FILE* stream, const char* utf8);
-void log_print(sd_log_level_t level, const char* log, bool verbose, bool color);
+const char* log_level_name(sd_log_level_t level);
+bool parse_log_level(const std::string& name, sd_log_level_t& level);
+void log_print(sd_log_level_t level, const char* log, sd_log_level_t min_level, bool color);
 void example_log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...);
 
 #define LOG_DEBUG(format, ...) example_log_printf(SD_LOG_DEBUG, __FILE__, __LINE__, format, ##__VA_ARGS__)
+#define LOG_VERBOSE(format, ...) example_log_printf(SD_LOG_VERBOSE, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_INFO(format, ...) example_log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_WARN(format, ...) example_log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_ERROR(format, ...) example_log_printf(SD_LOG_ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__)

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 {};
@@ -836,7 +860,7 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
     const uint32_t audio_data_size       = has_audio ? static_cast<uint32_t>(audio_pcm.size()) : 0;
 
     if (mjpg_quality != quality)
-        LOG_DEBUG("create_mjpg_avi...(): compression quality was limited from %i to %i", quality, mjpg_quality);
+        LOG_VERBOSE("create_mjpg_avi...(): compression quality was limited from %i to %i", quality, mjpg_quality);
 
     std::vector<uint8_t> avi_data;
     avi_data.reserve(static_cast<size_t>(num_images) * 1024);
@@ -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

@@ -13,9 +13,14 @@ What this example does:
 * `--llm` selects the text encoder / language model used by this pipeline
 * `--diffusion-fa` enables flash attention in the diffusion model
 * `--offload-to-cpu` reduces VRAM pressure by keeping weights in RAM when possible
-* `-v` enables verbose logging
+* `-v` enables verbose logging (equivalent to `--log-level verbose`)
 * `--cfg-scale 1.0` sets the default CFG scale for generation
 
+Logging defaults to `info`. Use `--log-level <level>` to select `debug`, `verbose`,
+`info`, `warn`, or `error` (from most to least detailed). Each level includes
+messages at that level and all less detailed levels. `-v` and `--verbose` are
+equivalent to `--log-level verbose`. If repeated, the last logging option wins.
+
 After the server starts successfully:
 
 * the web UI is available at `http://127.0.0.1:1234/`
@@ -124,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;
 }
 

examples/server/main.cpp

@@ -44,6 +44,9 @@ static void parse_args(int argc,
         exit(svr_params.normal_exit ? 0 : 1);
     }
 
+    log_level = svr_params.log_level;
+    log_color = svr_params.color;
+
     const bool random_seed_requested = default_gen_params.seed < 0;
 
     if (!svr_params.resolve_and_validate() ||
@@ -62,7 +65,7 @@ static void parse_args(int argc,
 
 void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
     SDSvrParams* svr_params = (SDSvrParams*)data;
-    log_print(level, log, svr_params->verbose, svr_params->color);
+    log_print(level, log, svr_params->log_level, svr_params->color);
 }
 
 int main(int argc, const char** argv) {
@@ -76,14 +79,12 @@ int main(int argc, const char** argv) {
     parse_args(argc, argv, svr_params, ctx_params, default_gen_params);
 
     sd_set_log_callback(sd_log_cb, (void*)&svr_params);
-    log_verbose = svr_params.verbose;
-    log_color   = svr_params.color;
 
-    LOG_DEBUG("version: %s", version_string().c_str());
-    LOG_DEBUG("%s", sd_get_system_info());
-    LOG_DEBUG("%s", svr_params.to_string().c_str());
-    LOG_DEBUG("%s", ctx_params.to_string().c_str());
-    LOG_DEBUG("%s", default_gen_params.to_string().c_str());
+    LOG_VERBOSE("version: %s", version_string().c_str());
+    LOG_VERBOSE("%s", sd_get_system_info());
+    LOG_VERBOSE("%s", svr_params.to_string().c_str());
+    LOG_VERBOSE("%s", ctx_params.to_string().c_str());
+    LOG_VERBOSE("%s", default_gen_params.to_string().c_str());
 
     sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(false);
     SDCtxPtr sd_ctx(new_sd_ctx(&sd_ctx_params));

examples/server/routes_openai.cpp

@@ -270,7 +270,7 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
                 return;
             }
 
-            LOG_DEBUG("%s\n", request.gen_params.to_string().c_str());
+            LOG_VERBOSE("%s\n", request.gen_params.to_string().c_str());
 
             SDImageVec results;
             if (!execute_sync_img_gen_request(*runtime, request, results, error_message)) {
@@ -344,7 +344,7 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
                 return;
             }
 
-            LOG_DEBUG("%s\n", request.gen_params.to_string().c_str());
+            LOG_VERBOSE("%s\n", request.gen_params.to_string().c_str());
 
             SDImageVec results;
             if (!execute_sync_img_gen_request(*runtime, request, results, error_message)) {

examples/server/routes_sdapi.cpp

@@ -330,7 +330,7 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
                 return;
             }
 
-            LOG_DEBUG("%s\n", request.gen_params.to_string().c_str());
+            LOG_VERBOSE("%s\n", request.gen_params.to_string().c_str());
 
             sd_img_gen_params_t img_gen_params = request.to_sd_img_gen_params_t();
             SDImageVec results;

examples/server/runtime.cpp

@@ -199,7 +199,6 @@ ArgOptions SDSvrParams::get_options() {
     };
 
     options.bool_options = {
-        {"-v", "--verbose", "print extra info", true, &verbose},
         {"", "--color", "colors the logging tags according to level", true, &color},
     };
 
@@ -212,6 +211,7 @@ ArgOptions SDSvrParams::get_options() {
     options.manual_options = {
         {"-h", "--help", "show this help message and exit", on_help_arg},
     };
+    add_log_options(options, log_level);
     return options;
 }
 
@@ -243,6 +243,7 @@ bool SDSvrParams::resolve_and_validate() {
 std::string SDSvrParams::to_string() const {
     std::ostringstream oss;
     oss << "SDSvrParams {\n"
+        << "  log_level: " << log_level_name(log_level) << ",\n"
         << "  listen_ip: " << listen_ip << ",\n"
         << "  listen_port: \"" << listen_port << "\",\n"
         << "  serve_html_path: \"" << serve_html_path << "\",\n"

examples/server/runtime.h

@@ -22,7 +22,7 @@ struct SDSvrParams {
     int listen_port       = 1234;
     std::string serve_html_path;
     bool normal_exit = false;
-    bool verbose     = false;
+    sd_log_level_t log_level = SD_LOG_INFO;
     bool color       = false;
 
     ArgOptions get_options();

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
 };
 
@@ -147,6 +148,7 @@ enum sd_type_t {
 
 enum sd_log_level_t {
     SD_LOG_DEBUG,
+    SD_LOG_VERBOSE,
     SD_LOG_INFO,
     SD_LOG_WARN,
     SD_LOG_ERROR
@@ -206,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;
@@ -239,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 {
@@ -492,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);
 
@@ -514,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/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

@@ -2,364 +2,479 @@
 
 #include <algorithm>
 #include <cctype>
+#include <cstddef>
 #include <cstdint>
+#include <fstream>
 #include <utility>
 #include <vector>
 
+#if defined(_WIN32)
+#ifndef NOMINMAX
+#define NOMINMAX
+#endif
+#include <windows.h>
+#elif defined(__APPLE__)
+#include <mach/mach.h>
+#endif
+
 #include "core/ggml_extend_backend.h"
 #include "core/util.h"
 #include "ggml-backend.h"
 
 namespace sd::backend_fit {
-    namespace {
 
-        constexpr int64_t MiB = 1024ll * 1024;
+    static constexpr int64_t MiB = 1024ll * 1024;
 
-        enum class ComponentKind {
-            DIT         = 0,
-            VAE         = 1,
-            CONDITIONER = 2,
-        };
+    enum class ComponentKind {
+        DIT,
+        CONDITIONER,
+        VAE,
+    };
 
-        struct Component {
+    struct Component {
+        ComponentKind kind;
+        const char* name;
+        int64_t params_bytes  = 0;
+        int64_t reserve_bytes = 0;
+        int64_t staging_bytes = 0;
+    };
+
+    struct Device {
+        std::string name;
+        std::string description;
+        int64_t free_bytes   = 0;
+        int64_t budget_bytes = 0;
+    };
+
+    enum class ParamsLocation {
+        MAIN_GPU,
+        CPU,
+        OTHER_GPU,
+        DISK,
+    };
+
+    struct Decision {
+        ParamsLocation params_location = ParamsLocation::DISK;
+        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;
+    };
+
+    static bool classify_tensor(const std::string& name, ComponentKind& out) {
+        auto contains = [&](const char* s) { return name.find(s) != std::string::npos; };
+
+        if (contains("model.diffusion_model.") || contains("unet.")) {
+            out = ComponentKind::DIT;
+            return true;
+        }
+        if (contains("first_stage_model.") ||
+            name.rfind("vae.", 0) == 0 ||
+            name.rfind("tae.", 0) == 0) {
+            out = ComponentKind::VAE;
+            return true;
+        }
+        if (contains("text_encoders") ||
+            contains("cond_stage_model") ||
+            contains("te.text_model.") ||
+            contains("conditioner") ||
+            name.rfind("text_encoder.", 0) == 0 ||
+            name.rfind("text_embedding_projection.", 0) == 0 ||
+            contains(".aggregate_embed.")) {
+            out = ComponentKind::CONDITIONER;
+            return true;
+        }
+        return false;
+    }
+
+    static std::vector<Component> estimate_components(ModelLoader& loader, ggml_type override_wtype) {
+        int64_t bytes[3]          = {0, 0, 0};
+        int64_t largest_tensor[3] = {0, 0, 0};
+        for (const auto& [name, stored_tensor] : loader.get_tensor_storage_map()) {
+            TensorStorage ts = stored_tensor;
             ComponentKind kind;
-            const char* name;
-            int64_t params_bytes  = 0;
-            int64_t reserve_bytes = 0;
-            bool splittable       = false;
-        };
-
-        struct Device {
-            ggml_backend_dev_t dev = nullptr;
-            std::string name;
-            std::string description;
-            int64_t free_bytes   = 0;
-            int64_t total_bytes  = 0;
-            int64_t budget_bytes = 0;
-        };
-
-        struct Decision {
-            ComponentKind kind;
-            bool on_cpu = false;
-            std::vector<size_t> device_idxs;
-        };
-
-        struct Plan {
-            bool valid      = false;
-            bool time_share = false;
-            std::vector<Decision> decisions;
-        };
-
-        bool classify_tensor(const std::string& name, ComponentKind& out) {
-            auto contains = [&](const char* s) { return name.find(s) != std::string::npos; };
-
-            if (contains("model.diffusion_model.") || contains("unet.")) {
-                out = ComponentKind::DIT;
-                return true;
+            if (is_unused_tensor(ts.name) || !classify_tensor(ts.name, kind)) {
+                continue;
             }
-            if (contains("first_stage_model.") ||
-                name.rfind("vae.", 0) == 0 ||
-                name.rfind("tae.", 0) == 0) {
-                out = ComponentKind::VAE;
-                return true;
+            if (ts.expected_type != GGML_TYPE_COUNT) {
+                ts.type = ts.expected_type;
+            } else if (override_wtype != GGML_TYPE_COUNT && loader.tensor_should_be_converted(ts, override_wtype)) {
+                ts.type = override_wtype;
             }
-            if (contains("text_encoders") ||
-                contains("cond_stage_model") ||
-                contains("te.text_model.") ||
-                contains("conditioner") ||
-                name.rfind("text_encoder.", 0) == 0 ||
-                name.rfind("text_embedding_projection.", 0) == 0 ||
-                contains(".aggregate_embed.")) {
-                out = ComponentKind::CONDITIONER;
-                return true;
-            }
-            return false;
+            const int64_t tensor_bytes = (int64_t)ts.nbytes() + 64;
+            bytes[int(kind)] += tensor_bytes;
+            largest_tensor[int(kind)] = std::max(largest_tensor[int(kind)], tensor_bytes);
         }
 
-        std::vector<Component> estimate_components(ModelLoader& loader, ggml_type override_wtype) {
-            const auto& storage = loader.get_tensor_storage_map();
+        return {
+            {ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, largest_tensor[int(ComponentKind::DIT)]},
+            {ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, largest_tensor[int(ComponentKind::CONDITIONER)]},
+            {ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, largest_tensor[int(ComponentKind::VAE)]},
+        };
+    }
 
-            int64_t bytes[3] = {0, 0, 0};
-            for (const auto& [name, ts_const] : storage) {
-                TensorStorage ts = ts_const;
-                if (is_unused_tensor(ts.name)) {
-                    continue;
-                }
-                ComponentKind kind;
-                if (!classify_tensor(ts.name, kind)) {
-                    continue;
-                }
-                if (override_wtype != GGML_TYPE_COUNT &&
-                    loader.tensor_should_be_converted(ts, override_wtype)) {
-                    ts.type = override_wtype;
-                } else if (ts.expected_type != GGML_TYPE_COUNT && ts.expected_type != ts.type) {
-                    ts.type = ts.expected_type;
-                }
-                bytes[int(kind)] += (int64_t)ts.nbytes() + 64;
+    static std::string budget_key(std::string name) {
+        std::transform(name.begin(), name.end(), name.begin(), [](unsigned char c) { return (char)std::tolower(c); });
+        return name;
+    }
+
+    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);
+            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;
+            device.name        = ggml_backend_dev_name(dev);
+            device.description = ggml_backend_dev_description(dev);
+            size_t free_bytes = 0, total_bytes = 0;
+            ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
+            device.free_bytes = (int64_t)free_bytes;
 
-            std::vector<Component> out;
-            out.push_back({ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, true});
-            out.push_back({ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, false});
-            out.push_back({ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, true});
-            return out;
-        }
-
-        std::vector<Device> enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) {
-            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) {
-                    continue;
-                }
-                Device d;
-                d.dev             = dev;
-                d.name            = ggml_backend_dev_name(dev);
-                d.description     = ggml_backend_dev_description(dev);
-                size_t free_bytes = 0, total_bytes = 0;
-                ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
-                d.free_bytes  = (int64_t)free_bytes;
-                d.total_bytes = (int64_t)total_bytes;
-
-                std::string budget_key = d.name;
-                std::transform(budget_key.begin(), budget_key.end(), budget_key.begin(),
-                               [](unsigned char c) { return (char)std::tolower(c); });
-                float gib = budgets.default_gib;
-                auto it   = budgets.backend_gib.find(budget_key);
-                if (it != budgets.backend_gib.end()) {
-                    gib = it->second;
-                }
-                if (gib > 0.f) {
-                    d.budget_bytes = std::min<int64_t>((int64_t)(gib * 1024.0 * 1024.0 * 1024.0), d.free_bytes);
-                } else if (gib < 0.f) {
-                    d.budget_bytes = d.free_bytes + (int64_t)(gib * 1024.0 * 1024.0 * 1024.0);
-                } else {
-                    d.budget_bytes = d.free_bytes - 512 * MiB;
-                }
-                d.budget_bytes = std::max<int64_t>(d.budget_bytes, 0);
-                out.push_back(d);
+            float gib = budgets.default_gib;
+            auto it   = budgets.backend_gib.find(budget_key(device.name));
+            if (it != budgets.backend_gib.end()) {
+                gib = it->second;
             }
-            return out;
+            if (gib > 0.f) {
+                device.budget_bytes = (int64_t)std::min(gib * 1024.0 * MiB, (double)device.free_bytes);
+            } else if (gib < 0.f) {
+                device.budget_bytes = (int64_t)std::max<double>(device.free_bytes + gib * 1024.0 * MiB, 0);
+            } else {
+                device.budget_bytes = std::max<int64_t>(device.free_bytes - 512 * MiB, 0);
+            }
+            out.push_back(std::move(device));
         }
+        return out;
+    }
 
-        Plan compute_plan(const std::vector<Component>& components, const std::vector<Device>& devices) {
-            Plan plan;
-            if (devices.empty()) {
+    static int64_t available_ram_bytes() {
+#if defined(_WIN32)
+        MEMORYSTATUSEX status{};
+        status.dwLength = sizeof(status);
+        if (GlobalMemoryStatusEx(&status)) {
+            return (int64_t)status.ullAvailPhys;
+        }
+#elif defined(__linux__)
+        std::ifstream meminfo("/proc/meminfo");
+        std::string key, unit;
+        int64_t kib = 0;
+        while (meminfo >> key >> kib >> unit) {
+            if (key == "MemAvailable:" && unit == "kB" && kib >= 0) {
+                return kib * 1024;
+            }
+        }
+#elif defined(__APPLE__)
+        const mach_port_t host = mach_host_self();
+        vm_size_t page_size    = 0;
+        vm_statistics64_data_t stats{};
+        mach_msg_type_number_t count = HOST_VM_INFO64_COUNT;
+        const bool ok                = host_page_size(host, &page_size) == KERN_SUCCESS &&
+                        host_statistics64(host, HOST_VM_INFO64, (host_info64_t)&stats, &count) == KERN_SUCCESS;
+        mach_port_deallocate(mach_task_self(), host);
+        if (ok) {
+            return ((int64_t)stats.free_count + stats.inactive_count) * page_size;
+        }
+#endif
+        return -1;
+    }
+
+    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 &&
+                (main_device == SIZE_MAX || devices[di].budget_bytes > devices[main_device].budget_bytes)) {
+                main_device = di;
+            }
+        }
+        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());
-            for (size_t i = 0; i < order.size(); i++) {
-                order[i] = i;
-            }
-            std::sort(order.begin(), order.end(), [&](size_t a, size_t b) {
-                return components[a].params_bytes > components[b].params_bytes;
-            });
+        std::vector<size_t> order(components.size());
+        for (size_t ci = 0; ci < components.size(); ++ci) {
+            order[ci] = ci;
+        }
+        std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) {
+            return components[a].kind < components[b].kind;
+        });
 
-            {
-                std::vector<int64_t> params_sum(devices.size(), 0);
-                std::vector<int64_t> max_reserve(devices.size(), 0);
-                std::vector<Decision> decisions(components.size());
-                bool ok = true;
-                for (size_t ci : order) {
-                    const Component& comp = components[ci];
-                    decisions[ci].kind    = comp.kind;
-                    if (comp.params_bytes == 0) {
-                        continue;
-                    }
-                    int best = -1;
-                    for (size_t di = 0; di < devices.size(); di++) {
-                        int64_t need = params_sum[di] + comp.params_bytes + std::max(max_reserve[di], comp.reserve_bytes);
-                        if (need <= devices[di].budget_bytes &&
-                            (best < 0 || devices[di].budget_bytes - params_sum[di] > devices[best].budget_bytes - params_sum[best])) {
-                            best = (int)di;
-                        }
-                    }
-                    if (best < 0) {
-                        ok = false;
-                        break;
-                    }
-                    params_sum[best] += comp.params_bytes;
-                    max_reserve[best] = std::max(max_reserve[best], comp.reserve_bytes);
-                    decisions[ci].device_idxs.push_back((size_t)best);
-                }
-                if (ok) {
-                    plan.valid      = true;
-                    plan.time_share = false;
-                    plan.decisions  = std::move(decisions);
-                    return plan;
-                }
-            }
+        std::vector<int64_t> remaining;
+        for (const Device& device : devices) {
+            remaining.push_back(std::max<int64_t>(device.budget_bytes, 0));
+        }
+        ram_budget_bytes = std::max<int64_t>(ram_budget_bytes, 0);
+        plan.decisions.resize(components.size());
 
-            plan.decisions.assign(components.size(), {});
-            for (size_t ci : order) {
-                const Component& comp = components[ci];
-                Decision& decision    = plan.decisions[ci];
-                decision.kind         = comp.kind;
-                if (comp.params_bytes == 0) {
+        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;
                 }
-                int best = -1;
-                for (size_t di = 0; di < devices.size(); di++) {
-                    if (comp.params_bytes + comp.reserve_bytes <= devices[di].budget_bytes &&
-                        (best < 0 || devices[di].budget_bytes > devices[best].budget_bytes)) {
-                        best = (int)di;
-                    }
+                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];
+            if (comp.params_bytes == 0) {
+                continue;
+            }
+
+            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   = 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;
                 }
-                if (best >= 0) {
-                    decision.device_idxs.push_back((size_t)best);
+                continue;
+            }
+            if (comp.params_bytes <= ram_budget_bytes) {
+                decision.params_location = ParamsLocation::CPU;
+                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) {
+                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;
+                }
+            }
+            if (best != SIZE_MAX) {
+                decision.params_location = ParamsLocation::OTHER_GPU;
+                decision.params_device   = best;
+                remaining[best] -= comp.params_bytes;
+            }
+        }
+        plan.valid = true;
+        return plan;
+    }
+
+    static std::string params_backend_name(const Decision& decision, const std::vector<Device>& devices) {
+        switch (decision.params_location) {
+            case ParamsLocation::MAIN_GPU:
+            case ParamsLocation::OTHER_GPU:
+                return devices[decision.params_device].name;
+            case ParamsLocation::CPU:
+                return "cpu";
+            case ParamsLocation::DISK:
+                return "disk";
+        }
+        return "disk";
+    }
+
+    static void print_plan(const Plan& plan,
+                           const std::vector<Component>& components,
+                           const std::vector<Device>& devices,
+                           int64_t free_ram,
+                           int64_t ram_budget) {
+        LOG_INFO("auto-fit plan:");
+        LOG_INFO("  devices:");
+        for (const Device& device : devices) {
+            LOG_INFO("    %-12s %-32s free %6lld MiB, budget %6lld MiB",
+                     device.name.c_str(), device.description.c_str(),
+                     (long long)(device.free_bytes / MiB), (long long)(device.budget_bytes / MiB));
+        }
+        if (free_ram < 0) {
+            LOG_WARN("auto-fit: available RAM is unknown; skipping CPU parameter residency");
+        } else {
+            LOG_INFO("    RAM          free %6lld MiB, params budget %6lld MiB",
+                     (long long)(free_ram / MiB), (long long)(ram_budget / MiB));
+        }
+        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 = 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),
+                     plan.runtimes[ci].name.c_str(), params.c_str());
+        }
+    }
+
+    static void append_assignment(std::string& spec, const char* key, const std::string& value) {
+        if (!spec.empty()) {
+            spec += ",";
+        }
+        spec += key;
+        spec += "=";
+        spec += value;
+    }
+
+    static const char* module_key(ComponentKind kind) {
+        switch (kind) {
+            case ComponentKind::DIT:
+                return "diffusion";
+            case ComponentKind::CONDITIONER:
+                return "te";
+            case ComponentKind::VAE:
+                return "vae";
+        }
+        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;
                 }
-                if (comp.splittable && devices.size() > 1) {
-                    int64_t capacity = 0;
-                    for (const Device& d : devices) {
-                        capacity += std::max<int64_t>(d.budget_bytes - comp.reserve_bytes, 0);
-                    }
-                    if (comp.params_bytes <= capacity) {
-                        std::vector<size_t> idxs(devices.size());
-                        for (size_t i = 0; i < idxs.size(); i++) {
-                            idxs[i] = i;
-                        }
-                        std::sort(idxs.begin(), idxs.end(), [&](size_t a, size_t b) {
-                            return devices[a].budget_bytes > devices[b].budget_bytes;
-                        });
-                        decision.device_idxs = std::move(idxs);
-                        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);
                     }
                 }
-                decision.on_cpu = true;
             }
-            plan.valid      = true;
-            plan.time_share = true;
-            return plan;
+            runtimes.push_back(std::move(runtime));
         }
-
-        void print_plan(const Plan& plan,
-                        const std::vector<Component>& components,
-                        const std::vector<Device>& devices) {
-            LOG_INFO("auto-fit plan%s:", plan.time_share ? " (time-share: params load per phase and free after)" : "");
-            LOG_INFO("  devices:");
-            for (const Device& d : devices) {
-                LOG_INFO("    %-12s %-32s free %6lld MiB, budget %6lld MiB",
-                         d.name.c_str(), d.description.c_str(),
-                         (long long)(d.free_bytes / MiB), (long long)(d.budget_bytes / MiB));
-            }
-            LOG_INFO("  components:");
-            for (size_t ci = 0; ci < components.size(); ci++) {
-                const Component& comp    = components[ci];
-                const Decision& decision = plan.decisions[ci];
-                std::string target;
-                if (comp.params_bytes == 0) {
-                    target = "(not present)";
-                } else if (decision.on_cpu) {
-                    target = "CPU";
-                } else {
-                    for (size_t k = 0; k < decision.device_idxs.size(); k++) {
-                        if (k > 0) {
-                            target += " & ";
-                        }
-                        target += devices[decision.device_idxs[k]].name;
-                    }
-                    if (decision.device_idxs.size() > 1) {
-                        target += " (split)";
-                    }
-                }
-                LOG_INFO("    %-12s params %6lld MiB, compute reserve %5lld MiB -> %s",
-                         comp.name,
-                         (long long)(comp.params_bytes / MiB),
-                         (long long)(comp.reserve_bytes / MiB),
-                         target.c_str());
-            }
-        }
-
-        void append_assignment(std::string& spec, const char* key, const std::string& value) {
-            if (!spec.empty()) {
-                spec += ",";
-            }
-            spec += key;
-            spec += "=";
-            spec += value;
-        }
-
-        void append_component_decision(const std::vector<Component>& components,
-                                       const std::vector<Device>& devices,
-                                       const Plan& plan,
-                                       ComponentKind kind,
-                                       const char* module_key,
-                                       std::string& runtime_spec,
-                                       std::string& params_spec) {
-            for (size_t ci = 0; ci < components.size(); ci++) {
-                if (components[ci].kind != kind || components[ci].params_bytes == 0) {
-                    continue;
-                }
-                const Decision& decision = plan.decisions[ci];
-                if (decision.on_cpu) {
-                    append_assignment(runtime_spec, module_key, "cpu");
-                    return;
-                }
-                if (decision.device_idxs.empty()) {
-                    return;
-                }
-                std::string device_list;
-                for (size_t k = 0; k < decision.device_idxs.size(); k++) {
-                    if (k > 0) {
-                        device_list += "&";
-                    }
-                    device_list += devices[decision.device_idxs[k]].name;
-                }
-                append_assignment(runtime_spec, module_key, device_list);
-                if (plan.time_share) {
-                    append_assignment(params_spec, module_key, "disk");
-                }
-                return;
-            }
-        }
-
-    }  // namespace
+        return true;
+    }
 
     bool derive_backend_specs(ModelLoader& loader,
                               ggml_type override_wtype,
                               sd::ggml_graph_cut::MaxVramAssignment& budgets,
                               std::string& runtime_spec,
                               std::string& params_spec) {
-        if (!runtime_spec.empty() || !params_spec.empty()) {
-            LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend");
+        std::string error;
+        if (!budgets.canonicalize_backend_keys(&error)) {
+            LOG_ERROR("%s", error.c_str());
+            return false;
         }
 
-        {
-            std::string error;
-            if (!budgets.canonicalize_backend_keys(&error)) {
-                LOG_ERROR("%s", error.c_str());
-                return false;
-            }
+        // 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;
         }
-
-        auto components = estimate_components(loader, override_wtype);
-        auto devices    = enumerate_gpu_devices(budgets);
-        auto plan       = compute_plan(components, devices);
+        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, runtimes);
+        params_spec.clear();
         if (!plan.valid) {
-            LOG_WARN("auto-fit: no usable GPU devices; using the default backend");
-            runtime_spec.clear();
-            params_spec.clear();
+            if (devices.empty()) {
+                LOG_WARN("auto-fit: no GPU devices; using the default backend");
+            } else {
+                LOG_WARN("auto-fit: no GPU memory budget available; using CPU");
+                runtime_spec = "cpu";
+            }
             return true;
         }
 
-        print_plan(plan, components, devices);
+        print_plan(plan, components, devices, free_ram, ram_budget);
+        for (size_t ci = 0; ci < components.size(); ++ci) {
+            if (components[ci].params_bytes == 0) {
+                continue;
+            }
+            const char* key = module_key(components[ci].kind);
+            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));
+            }
+        }
 
-        std::string derived_runtime_spec;
-        std::string derived_params_spec;
-        append_component_decision(components, devices, plan, ComponentKind::DIT, "diffusion", derived_runtime_spec, derived_params_spec);
-        append_component_decision(components, devices, plan, ComponentKind::CONDITIONER, "te", derived_runtime_spec, derived_params_spec);
-        append_component_decision(components, devices, plan, ComponentKind::VAE, "vae", derived_runtime_spec, derived_params_spec);
-
-        runtime_spec = std::move(derived_runtime_spec);
-        params_spec  = std::move(derived_params_spec);
+        // Keep the planner's safety margin when the runner resolves its device limits.
+        for (const Device& device : devices) {
+            if (device.budget_bytes > 0) {
+                budgets.backend_gib[budget_key(device.name)] = (float)(device.budget_bytes / (1024.0 * MiB));
+            }
+        }
+        budgets.resolved_backend_bytes.clear();
 
         LOG_INFO("auto-fit: --backend \"%s\"%s%s%s",
                  runtime_spec.empty() ? "(default)" : runtime_spec.c_str(),
                  params_spec.empty() ? "" : " --params-backend \"",
-                 params_spec.c_str(),
-                 params_spec.empty() ? "" : "\"");
+                 params_spec.c_str(), params_spec.empty() ? "" : "\"");
         return true;
     }
 
@@ -369,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;
             }
@@ -381,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

@@ -0,0 +1,796 @@
+#include "core/ggml_extend.h"
+
+#include <cmath>
+#include <utility>
+
+#include "core/ggml_extend_backend.h"
+
+ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode) {
+    // reshape A
+    // swap 0th and nth axis
+    a           = ggml_cont(ctx, ggml_permute(ctx, a, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0));
+    int64_t ne1 = a->ne[1];
+    int64_t ne2 = a->ne[2];
+    int64_t ne3 = a->ne[3];
+    // make 2D
+    a = ggml_cont(ctx, ggml_reshape_2d(ctx, a, a->ne[0], (ne3 * ne2 * ne1)));
+
+    ggml_tensor* result = ggml_cont(ctx, ggml_transpose(ctx, ggml_mul_mat(ctx, a, b)));
+
+    // reshape output (same shape as a after permutation except first dim)
+    result = ggml_reshape_4d(ctx, result, result->ne[0], ne1, ne2, ne3);
+    // swap back 0th and nth axis
+    result = ggml_permute(ctx, result, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0);
+    return result;
+}
+
+ggml_tensor* ggml_ext_kronecker(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b) {
+    return ggml_mul(ctx,
+                    ggml_interpolate(ctx,
+                                     a,
+                                     a->ne[0] * b->ne[0],
+                                     a->ne[1] * b->ne[1],
+                                     a->ne[2] * b->ne[2],
+                                     a->ne[3] * b->ne[3],
+                                     GGML_SCALE_MODE_NEAREST),
+                    b);
+}
+
+ggml_tensor* ggml_ext_cont(ggml_context* ctx,
+                           ggml_tensor* x) {
+    if (ggml_is_contiguous(x)) {
+        return x;
+    }
+    return ggml_cont(ctx, x);
+}
+
+ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    int axis0,
+                                    int axis1,
+                                    int axis2,
+                                    int axis3) {
+    int torch_axes[4] = {axis0, axis1, axis2, axis3};
+
+    int ggml_axes[4] = {0};
+    for (int i = 0; i < 4; ++i) {
+        int found = 0;
+        for (int j = 0; j < 4; ++j) {
+            if (torch_axes[j] == i) {
+                ggml_axes[i] = j;
+                found        = 1;
+                break;
+            }
+        }
+        GGML_ASSERT(found && "Invalid permute input: must be a permutation of 0-3");
+    }
+
+    return ggml_permute(ctx, x, ggml_axes[0], ggml_axes[1], ggml_axes[2], ggml_axes[3]);
+}
+
+ggml_tensor* ggml_ext_slice(ggml_context* ctx,
+                            ggml_tensor* x,
+                            int dim,
+                            int64_t start,
+                            int64_t end,
+                            bool cont) {
+    GGML_ASSERT(dim >= 0 && dim < 4);
+    if (x->ne[dim] == 1) {
+        return x;
+    }
+    while (start < 0) {
+        start = x->ne[dim] + start;
+    }
+    while (end < 0) {
+        end = x->ne[dim] + end;
+    }
+    GGML_ASSERT(end > start);
+    GGML_ASSERT(start >= 0 && start < x->ne[dim]);
+    GGML_ASSERT(end > start && end <= x->ne[dim]);
+
+    int64_t slice_size  = end - start;
+    int64_t slice_ne[4] = {x->ne[0], x->ne[1], x->ne[2], x->ne[3]};
+    slice_ne[dim]       = slice_size;
+
+    x = ggml_view_4d(ctx, x,
+                     slice_ne[0], slice_ne[1], slice_ne[2], slice_ne[3],
+                     x->nb[1], x->nb[2], x->nb[3], start * x->nb[dim]);
+
+    if (cont) {
+        x = ggml_cont(ctx, x);
+    }
+
+    return x;
+}
+
+std::vector<ggml_tensor*> ggml_ext_chunk(ggml_context* ctx,
+                                         ggml_tensor* x,
+                                         int num,
+                                         int64_t dim,
+                                         bool cont) {
+    GGML_ASSERT(dim >= 0 && dim < 4);
+    GGML_ASSERT(x->ne[dim] % num == 0);
+
+    std::vector<ggml_tensor*> chunks;
+    int64_t chunk_size  = x->ne[dim] / num;
+    int64_t stride      = chunk_size * x->nb[dim];
+    int64_t chunk_ne[4] = {x->ne[0], x->ne[1], x->ne[2], x->ne[3]};
+    chunk_ne[dim]       = chunk_size;
+    for (int i = 0; i < num; i++) {
+        auto chunk = ggml_view_4d(
+            ctx, x,
+            chunk_ne[0], chunk_ne[1], chunk_ne[2], chunk_ne[3],
+            x->nb[1], x->nb[2], x->nb[3], stride * i);
+        if (cont) {
+            chunk = ggml_cont(ctx, chunk);
+        }
+        chunks.push_back(chunk);
+    }
+
+    return chunks;
+}
+
+ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first) {
+    // x: [ne3, ne2, ne1, ne0]
+    // return: [ne3, ne2, ne1, ne0/2]
+
+    auto x_vec = ggml_ext_chunk(ctx, x, 2, 0, false);
+    ggml_tensor* gate;
+    if (gate_first) {
+        gate = x_vec[0];
+        x    = x_vec[1];
+    } else {
+        x    = x_vec[0];
+        gate = x_vec[1];
+    }
+    gate = ggml_cont(ctx, gate);
+    gate = ggml_silu_inplace(ctx, gate);
+
+    x = ggml_mul(ctx, x, gate);  // [ne3, ne2, ne1, ne0/2]
+
+    return x;
+}
+
+ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx,
+                                    ggml_tensor* a) {
+    const float eps = 1e-6f;  // default eps parameter
+    return ggml_group_norm(ctx, a, 32, eps);
+}
+
+static bool ggml_ext_is_padded_1d(const ggml_tensor* x) {
+    return x->nb[0] == ggml_type_size(x->type) &&
+           x->nb[2] == x->nb[1] * x->ne[1] &&
+           x->nb[3] == x->nb[2] * x->ne[2];
+}
+
+ggml_tensor* ggml_ext_scale(ggml_context* ctx,
+                            ggml_tensor* x,
+                            float factor,
+                            bool inplace) {
+    if (!ggml_ext_is_padded_1d(x)) {
+        x = ggml_cont(ctx, x);
+    }
+    if (inplace) {
+        x = ggml_scale_inplace(ctx, x, factor);
+    } else {
+        x = ggml_scale(ctx, x, factor);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_gelu(ggml_context* ctx,
+                           ggml_tensor* x,
+                           bool inplace) {
+    if (!ggml_is_contiguous(x)) {
+        x = ggml_cont(ctx, x);
+    }
+    if (inplace) {
+        x = ggml_gelu_inplace(ctx, x);
+    } else {
+        x = ggml_gelu(ctx, x);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 bool inplace) {
+    if (!ggml_is_contiguous(x)) {
+        x = ggml_cont(ctx, x);
+    }
+    if (inplace) {
+        x = ggml_gelu_quick_inplace(ctx, x);
+    } else {
+        x = ggml_gelu_quick(ctx, x);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_linear(ggml_context* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* w,
+                             ggml_tensor* b,
+                             bool force_prec_f32,
+                             float scale) {
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, scale);
+    }
+    if (x->ne[2] * x->ne[3] > 1024) {
+        // workaround: avoid ggml cuda error
+        int64_t ne2 = x->ne[2];
+        int64_t ne3 = x->ne[3];
+        x           = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]);
+        x           = ggml_mul_mat(ctx, w, x);
+        if (force_prec_f32) {
+            ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+        }
+        x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3);
+    } else {
+        x = ggml_mul_mat(ctx, w, x);
+        if (force_prec_f32) {
+            ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+        }
+    }
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, 1.f / scale);
+    }
+    if (b != nullptr) {
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
+                                           ggml_tensor* x,
+                                           ggml_tensor* w,
+                                           ggml_tensor* weight_scale,
+                                           ggml_tensor* b,
+                                           int convrot_group_size,
+                                           float scale) {
+    GGML_ASSERT(x->type == GGML_TYPE_F32 || (x->type == GGML_TYPE_I8 && scale == 1.f));
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, scale);
+    }
+
+    ggml_tensor* fused_bias = scale == 1.f ? b : nullptr;
+    if (x->ne[2] * x->ne[3] > 1024) {
+        int64_t ne2 = x->ne[2];
+        int64_t ne3 = x->ne[3];
+        x           = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]);
+        x           = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size);
+        x           = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3);
+    } else {
+        x = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size);
+    }
+
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, 1.f / scale);
+        if (b != nullptr) {
+            x = ggml_add_inplace(ctx, x, b);
+        }
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
+                              ggml_backend_t backend,
+                              ggml_tensor* x,
+                              int lp0,
+                              int rp0,
+                              int lp1,
+                              int rp1,
+                              int lp2,
+                              int rp2,
+                              int lp3,
+                              int rp3,
+                              bool circular_x,
+                              bool circular_y) {
+    if (circular_x && circular_y) {
+        return ggml_pad_ext_circular(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
+    }
+
+    if (circular_x && (lp0 != 0 || rp0 != 0)) {
+        x   = ggml_pad_ext_circular(ctx, x, lp0, rp0, 0, 0, 0, 0, 0, 0);
+        lp0 = rp0 = 0;
+    }
+    if (circular_y && (lp1 != 0 || rp1 != 0)) {
+        x   = ggml_pad_ext_circular(ctx, x, 0, 0, lp1, rp1, 0, 0, 0, 0);
+        lp1 = rp1 = 0;
+    }
+
+    if (lp0 != 0 || rp0 != 0 || lp1 != 0 || rp1 != 0 || lp2 != 0 || rp2 != 0 || lp3 != 0 || rp3 != 0) {
+        ggml_tensor* padded = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
+        if (backend == nullptr || ggml_backend_supports_op(backend, padded)) {
+            x = padded;
+        } else {
+            // Some backends (e.g. Metal) only implement right-padding for
+            // GGML_OP_PAD (see #850): pad right by lp+rp instead, then roll
+            // the padding around to the left. shift < ne always holds because
+            // ne grew by lp+rp.
+            x = ggml_pad_ext(ctx, x, 0, lp0 + rp0, 0, lp1 + rp1, 0, lp2 + rp2, 0, lp3 + rp3);
+            x = ggml_roll(ctx, x, lp0, lp1, lp2, lp3);
+        }
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_pad(ggml_context* ctx,
+                          ggml_tensor* x,
+                          int p0,
+                          int p1,
+                          int p2,
+                          int p3,
+                          bool circular_x,
+                          bool circular_y) {
+    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,
+                              ggml_tensor* b,
+                              int s0,
+                              int s1,
+                              int p0,
+                              int p1,
+                              int d0,
+                              int d1,
+                              bool direct,
+                              bool circular_x,
+                              bool circular_y,
+                              float scale) {
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, scale);
+    }
+    if (w->ne[2] != x->ne[2] && ggml_n_dims(w) == 2) {
+        w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], w->ne[1]);
+    }
+
+    if ((p0 != 0 || p1 != 0) && (circular_x || circular_y)) {
+        x  = ggml_ext_pad_ext(ctx, nullptr, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y);
+        p0 = 0;
+        p1 = 0;
+    }
+
+    if (direct) {
+        x = ggml_conv_2d_direct(ctx, w, x, s0, s1, p0, p1, d0, d1);
+    } else {
+        x = ggml_conv_2d(ctx, w, x, s0, s1, p0, p1, d0, d1);
+    }
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, 1.f / scale);
+    }
+    if (b != nullptr) {
+        b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
+                              ggml_backend_t backend,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int64_t IC,
+                              int s0,
+                              int s1,
+                              int s2,
+                              int p0,
+                              int p1,
+                              int p2,
+                              int d0,
+                              int d1,
+                              int d2,
+                              bool force_prec_f32) {
+    if (force_prec_f32) {
+        ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
+
+        int64_t OC = w->ne[3] / IC;
+        int64_t N  = x->ne[3] / IC;
+        x          = ggml_mul_mat(ctx,
+                                  ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]),
+                                  ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1] * w->ne[2] * IC, OC));
+        ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+
+        int64_t OD = im2col->ne[3] / N;
+        x          = ggml_reshape_4d(ctx, x, im2col->ne[1] * im2col->ne[2], OD, N, OC);
+        x          = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2));
+        x          = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N);
+    } else {
+        // ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends
+        // (e.g. Metal, see #850) do not implement. Fall back to
+        // GGML_OP_CONV_3D on those backends.
+        bool im2col_3d_supported = true;
+        if (backend != nullptr) {
+            ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
+            im2col_3d_supported = ggml_backend_supports_op(backend, im2col);
+        }
+        if (im2col_3d_supported) {
+            x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
+        } else {
+            int64_t OC = w->ne[3] / IC;
+            int64_t N  = x->ne[3] / IC;
+            x          = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC);
+        }
+    }
+
+    if (b != nullptr) {
+        b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]);  // [OC, 1, 1, 1]
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_conv_3d_nx1x1(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    ggml_tensor* w,
+                                    ggml_tensor* b,
+                                    int s2,
+                                    int p2,
+                                    int d2) {
+    x = ggml_conv_2d(ctx, w, x, 1, s2, 0, p2, 1, d2);  // [N, OC, T, OH * OW]
+    if (b != nullptr) {
+        b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+        x = ggml_add(ctx, x, b);
+    }
+    return x;  // [N, OC, T, OH * OW]
+}
+
+std::vector<ggml_tensor*> split_qkv(ggml_context* ctx,
+                                    ggml_tensor* qkv) {
+    qkv = ggml_reshape_4d(ctx, qkv, qkv->ne[0] / 3, 3, qkv->ne[1], qkv->ne[2]);  // [N, L, 3, C]
+    qkv = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 3, 1, 2));                    // [3, N, L, C]
+
+    int64_t offset = qkv->nb[2] * qkv->ne[2];
+    auto q         = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 0);  // [N, L, C]
+    auto k         = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 1);  // [N, L, C]
+    auto v         = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 2);  // [N, L, C]
+    return {q, k, v};
+}
+
+std::vector<ggml_tensor*> split_image_qkv(ggml_context* ctx,
+                                          ggml_tensor* qkv) {
+    int64_t W   = qkv->ne[0];
+    int64_t H   = qkv->ne[1];
+    int64_t C   = qkv->ne[2] / 3;
+    int64_t N   = qkv->ne[3];
+    int64_t nb1 = qkv->nb[1];
+    int64_t nb2 = qkv->nb[2];
+    qkv         = ggml_reshape_4d(ctx, qkv, W * H, C, 3, N);                     // [N, 3, C, H*W]
+    qkv         = ggml_cont(ctx, ggml_ext_torch_permute(ctx, qkv, 0, 1, 3, 2));  // [3, N, C, H*W]
+
+    int64_t offset = qkv->nb[2] * qkv->ne[2];
+    auto q         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 0);  // [N, C, H, W]
+    auto k         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 1);  // [N, C, H, W]
+    auto v         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 2);  // [N, C, H, W]
+    return {q, k, v};
+}
+
+ggml_tensor* ggml_ext_full(ggml_context* ctx,
+                           float value,
+                           int64_t ne0,
+                           int64_t ne1,
+                           int64_t ne2,
+                           int64_t ne3) {
+    auto one = ggml_get_tensor(ctx, "ggml_runner_build_in_tensor:one");
+    auto t   = ggml_ext_scale(ctx, one, value);             // [1,]
+    t        = ggml_repeat_4d(ctx, t, ne0, ne1, ne2, ne3);  // [ne0, ne1, ne2, ne3]
+    return t;
+}
+
+ggml_tensor* ggml_ext_zeros(ggml_context* ctx,
+                            int64_t ne0,
+                            int64_t ne1,
+                            int64_t ne2,
+                            int64_t ne3) {
+    return ggml_ext_full(ctx, 0.f, ne0, ne1, ne2, ne3);
+}
+
+ggml_tensor* ggml_ext_zeros_like(ggml_context* ctx,
+                                 ggml_tensor* x) {
+    return ggml_ext_zeros(ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]);
+}
+
+ggml_tensor* ggml_ext_ones(ggml_context* ctx,
+                           int64_t ne0,
+                           int64_t ne1,
+                           int64_t ne2,
+                           int64_t ne3) {
+    return ggml_ext_full(ctx, 1.f, ne0, ne1, ne2, ne3);
+}
+
+ggml_tensor* ggml_ext_ones_like(ggml_context* ctx,
+                                ggml_tensor* x) {
+    return ggml_ext_ones(ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]);
+}
+
+ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a) {
+    if (sd_backend_is(backend, "Vulkan")) {
+        auto zero_index = ggml_get_tensor(ctx, "ggml_runner_build_in_tensor:zero_int");
+        auto out        = ggml_reshape_1d(ctx, a, ggml_nelements(a));
+        out             = ggml_get_rows(ctx, out, zero_index);
+        out             = ggml_reshape(ctx, out, a);
+        // auto out = ggml_cast(ctx, a, GGML_TYPE_F32);
+        return out;
+    } else {
+        auto out         = ggml_reshape_2d(ctx, a, 1, ggml_nelements(a));
+        ggml_tensor* one = ggml_ext_ones(ctx, 1, 1, 1, 1);  // [1,]
+        if (ggml_is_transposed(out)) {
+            out = ggml_mul_mat(ctx, one, out);
+        } else {
+            out = ggml_mul_mat(ctx, out, one);
+        }
+        out = ggml_reshape(ctx, out, a);
+        return out;
+    }
+}
+
+ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
+                                    ggml_backend_t backend,
+                                    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) {  // avoid overflow
+    int64_t L_q;
+    int64_t L_k;
+    int64_t C;
+    int64_t N;
+    int64_t d_head;
+    int64_t n_kv_head;
+    if (!skip_reshape) {
+        L_q       = q->ne[1];
+        L_k       = k->ne[1];
+        C         = q->ne[0];
+        N         = q->ne[2];
+        d_head    = C / n_head;
+        n_kv_head = k->ne[0] / d_head;
+
+        q = ggml_reshape_4d(ctx, q, d_head, n_head, L_q, N);       // [N, L_q, n_head, d_head]
+        q = ggml_ext_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));  // [N, n_head, L_q, d_head]
+        q = ggml_reshape_3d(ctx, q, d_head, L_q, n_head * N);      // [N * n_head, L_q, d_head]
+
+        k = ggml_reshape_4d(ctx, k, d_head, n_kv_head, L_k, N);    // [N, L_k, n_kv_head, d_head]
+        k = ggml_ext_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));  // [N, n_kv_head, L_k, d_head]
+        k = ggml_reshape_3d(ctx, k, d_head, L_k, n_kv_head * N);   // [N * n_kv_head, L_k, d_head]
+
+        v = ggml_reshape_4d(ctx, v, d_head, n_kv_head, L_k, N);  // [N, L_k, n_kv_head, d_head]
+    } else {
+        L_q       = q->ne[1];
+        L_k       = k->ne[1];
+        d_head    = v->ne[0];
+        N         = v->ne[3];
+        n_kv_head = k->ne[2] / N;
+        C         = d_head * n_head;
+    }
+
+    float scale = (1.0f / sqrt((float)d_head));
+
+    ggml_tensor* kqv = nullptr;
+
+    auto build_kqv = [&](ggml_tensor* q_in, ggml_tensor* k_in, ggml_tensor* v_in, ggml_tensor* mask_in) -> ggml_tensor* {
+        if (kv_scale != 1.0f) {
+            k_in = ggml_ext_scale(ctx, k_in, kv_scale);
+        }
+        k_in = ggml_cast(ctx, k_in, GGML_TYPE_F16);
+
+        v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3));
+        v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N);
+        if (kv_scale != 1.0f) {
+            v_in = ggml_ext_scale(ctx, v_in, kv_scale);
+        }
+        v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16);
+
+        if (mask_in != nullptr) {
+            // ggml_flash_attn_ext expects the mask as a contiguous F16 tensor shaped
+            // [n_kv, n_q, (heads), (batch)] (ne0 = key length, ne1 = query length) and,
+            // unlike the manual-attention path, does not broadcast the query dimension.
+            // Some callers (e.g. Chroma/T5) pass a per-key padding mask broadcast over
+            // queries ([n_kv, 1, ...]); materialize the query dimension to L_q so the
+            // kernel indexes it correctly. (A bare ggml_transpose here produced a
+            // [1, n_kv, ...] mask that the kernel silently misreads, yielding NaN/blank
+            // output for masked flash attention.)
+            if (mask_in->ne[1] != L_q) {
+                mask_in = ggml_repeat(ctx, mask_in,
+                                      ggml_new_tensor_4d(ctx, mask_in->type, mask_in->ne[0], L_q, mask_in->ne[2], mask_in->ne[3]));
+            }
+            mask_in = ggml_cast(ctx, mask_in, GGML_TYPE_F16);
+        }
+
+        auto out = ggml_flash_attn_ext(ctx, q_in, k_in, v_in, mask_in, scale / kv_scale, 0, 0);
+        if (!ggml_backend_supports_op(backend, out)) {
+            return nullptr;
+        }
+        ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32);
+        if (kv_scale != 1.0f) {
+            out = ggml_ext_scale(ctx, out, 1.0f / kv_scale);
+        }
+        return out;
+    };
+
+    if (flash_attn) {
+        // LOG_VERBOSE("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
+        bool can_use_flash_attn = true;
+        if (mask != nullptr) {
+            // TODO: figure out if we can bend t5 to work too
+            can_use_flash_attn = can_use_flash_attn && mask->ne[3] == 1;
+        }
+
+        if (can_use_flash_attn) {
+            kqv = build_kqv(q, k, v, mask);
+            if (kqv != nullptr) {
+                kqv = ggml_view_4d(ctx,
+                                   kqv,
+                                   d_head,
+                                   n_head,
+                                   L_q,
+                                   N,
+                                   kqv->nb[1],
+                                   kqv->nb[2],
+                                   kqv->nb[1] * n_head,
+                                   0);
+            }
+        }
+    }
+
+    if (kqv == nullptr) {
+        // if (flash_attn) {
+        //     LOG_VERBOSE("fallback to default attention, L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
+        // }
+        v = ggml_ext_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3));  // [N, n_kv_head, d_head, L_k]
+        v = ggml_reshape_3d(ctx, v, L_k, d_head, n_kv_head * N);   // [N * n_kv_head, d_head, L_k]
+
+        auto kq = ggml_mul_mat(ctx, k, q);  // [N * n_head, L_q, L_k]
+        ggml_mul_mat_set_prec(kq, GGML_PREC_F32);
+        kq = ggml_scale_inplace(ctx, kq, scale);
+        if (mask) {
+            kq = ggml_add_inplace(ctx, kq, mask);
+        }
+        kq = ggml_soft_max_inplace(ctx, kq);
+
+        kqv = ggml_mul_mat(ctx, v, kq);  // [N * n_head, L_q, d_head]
+
+        kqv = ggml_reshape_4d(ctx, kqv, d_head, L_q, n_head, N);  // [N, n_head, L_q, d_head]
+        kqv = ggml_permute(ctx, kqv, 0, 2, 1, 3);                 // [N, L_q, n_head, d_head]
+    }
+
+    kqv = ggml_ext_cont(ctx, kqv);
+    kqv = ggml_reshape_3d(ctx, kqv, d_head * n_head, L_q, N);  // [N, L_q, C]
+
+    return kqv;
+}
+
+ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* w,
+                                 ggml_tensor* b,
+                                 float eps) {
+    x = ggml_norm(ctx, x, eps);
+    if (w != nullptr) {
+        x = ggml_mul_inplace(ctx, x, w);
+        if (b != nullptr) {
+            x = ggml_add_inplace(ctx, x, b);
+        }
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* w,
+                                 ggml_tensor* b,
+                                 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);
+    }
+
+    x = ggml_group_norm(ctx, x, num_groups, eps);
+    if (w != nullptr && b != nullptr) {
+        x = ggml_mul_inplace(ctx, x, w);
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_timestep_embedding(
+    ggml_context* ctx,
+    ggml_tensor* timesteps,
+    int dim,
+    int max_period,
+    float time_factor) {
+    timesteps = ggml_ext_scale(ctx, timesteps, time_factor);
+    return ggml_timestep_embedding(ctx, timesteps, dim, max_period);
+}
+
+ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx,
+                                 std::vector<ggml_tensor*>& tensors,
+                                 int dim) {
+    while (tensors.size() > 1) {
+        std::vector<ggml_tensor*> next_level;
+        for (size_t i = 0; i < tensors.size(); i += 2) {
+            if (i + 1 < tensors.size()) {
+                next_level.push_back(ggml_concat(ctx, tensors[i], tensors[i + 1], dim));
+            } else {
+                next_level.push_back(tensors[i]);
+            }
+        }
+        tensors = std::move(next_level);
+    }
+    return tensors[0];
+}

src/core/ggml_extend.h

@@ -0,0 +1,248 @@
+#ifndef __SD_CORE_GGML_EXTEND_H__
+#define __SD_CORE_GGML_EXTEND_H__
+
+#include <cstdint>
+#include <vector>
+
+#include "ggml-backend.h"
+#include "ggml.h"
+
+#define EPS 1e-05f
+
+static_assert(GGML_MAX_NAME >= 160, "GGML_MAX_NAME must be at least 160");
+
+// n-mode tensor-matrix product
+// example: 2-mode product
+// A: [ne03, k, ne01, ne00]
+// B: k rows, m columns => [k, m]
+// result is [ne03, m, ne01, ne00]
+ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode = 0);
+
+// Kronecker product
+// [ne03,ne02,ne01,ne00] x [ne13,ne12,ne11,ne10] => [ne03*ne13,ne02*ne12,ne01*ne11,ne00*ne10]
+ggml_tensor* ggml_ext_kronecker(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b);
+
+ggml_tensor* ggml_ext_cont(ggml_context* ctx,
+                           ggml_tensor* x);
+
+// torch like permute
+ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    int axis0,
+                                    int axis1,
+                                    int axis2,
+                                    int axis3);
+
+ggml_tensor* ggml_ext_slice(ggml_context* ctx,
+                            ggml_tensor* x,
+                            int dim,
+                            int64_t start,
+                            int64_t end,
+                            bool cont = true);
+
+// example: [N, 3*C, H, W] => ([N, C, H, W], [N, C, H, W], [N, C, H, W])
+std::vector<ggml_tensor*> ggml_ext_chunk(ggml_context* ctx,
+                                         ggml_tensor* x,
+                                         int num,
+                                         int64_t dim,
+                                         bool cont = true);
+
+ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first = true);
+
+ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx,
+                                    ggml_tensor* a);
+
+ggml_tensor* ggml_ext_scale(ggml_context* ctx,
+                            ggml_tensor* x,
+                            float factor,
+                            bool inplace = false);
+
+ggml_tensor* ggml_ext_gelu(ggml_context* ctx,
+                           ggml_tensor* x,
+                           bool inplace = false);
+
+ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 bool inplace = false);
+
+ggml_tensor* ggml_ext_linear(ggml_context* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* w,
+                             ggml_tensor* b,
+                             bool force_prec_f32 = false,
+                             float scale         = 1.f);
+
+ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
+                                           ggml_tensor* x,
+                                           ggml_tensor* w,
+                                           ggml_tensor* weight_scale,
+                                           ggml_tensor* b,
+                                           int convrot_group_size,
+                                           float scale = 1.f);
+
+ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
+                              ggml_backend_t backend,
+                              ggml_tensor* x,
+                              int lp0,
+                              int rp0,
+                              int lp1,
+                              int rp1,
+                              int lp2,
+                              int rp2,
+                              int lp3,
+                              int rp3,
+                              bool circular_x = false,
+                              bool circular_y = false);
+
+ggml_tensor* ggml_ext_pad(ggml_context* ctx,
+                          ggml_tensor* x,
+                          int p0,
+                          int p1,
+                          int p2          = 0,
+                          int p3          = 0,
+                          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,]
+// result: [N, OC, OH, OW]
+ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int s0          = 1,
+                              int s1          = 1,
+                              int p0          = 0,
+                              int p1          = 0,
+                              int d0          = 1,
+                              int d1          = 1,
+                              bool direct     = false,
+                              bool circular_x = false,
+                              bool circular_y = false,
+                              float scale     = 1.f);
+
+// w: [OC，IC, KD, 1 * 1]
+// x: [N, IC, IH, IW]
+// b: [OC,]
+// result: [N*OC, OD, OH, OW]
+ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
+                              ggml_backend_t backend,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int64_t IC,
+                              int s0              = 1,
+                              int s1              = 1,
+                              int s2              = 1,
+                              int p0              = 0,
+                              int p1              = 0,
+                              int p2              = 0,
+                              int d0              = 1,
+                              int d1              = 1,
+                              int d2              = 1,
+                              bool force_prec_f32 = false);
+
+// w: [OC，IC, KD, 1 * 1]
+// x: [N, IC, ID, IH*IW]
+// b: [OC,]
+// result: [N, OC, OD, OH*OW]
+ggml_tensor* ggml_ext_conv_3d_nx1x1(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    ggml_tensor* w,
+                                    ggml_tensor* b,
+                                    int s2 = 1,
+                                    int p2 = 1,
+                                    int d2 = 1);
+
+// qkv: [N, L, 3*C]
+// return: ([N, L, C], [N, L, C], [N, L, C])
+std::vector<ggml_tensor*> split_qkv(ggml_context* ctx,
+                                    ggml_tensor* qkv);
+
+// qkv: [N, 3*C, H, W]
+// return: ([N, C, H, W], [N, C, H, W], [N, C, H, W])
+std::vector<ggml_tensor*> split_image_qkv(ggml_context* ctx,
+                                          ggml_tensor* qkv);
+
+// Constant and cast helpers require the built-in tensors initialized by GGMLRunner.
+ggml_tensor* ggml_ext_full(ggml_context* ctx,
+                           float value,
+                           int64_t ne0,
+                           int64_t ne1,
+                           int64_t ne2,
+                           int64_t ne3);
+
+ggml_tensor* ggml_ext_zeros(ggml_context* ctx,
+                            int64_t ne0,
+                            int64_t ne1,
+                            int64_t ne2,
+                            int64_t ne3);
+
+ggml_tensor* ggml_ext_zeros_like(ggml_context* ctx,
+                                 ggml_tensor* x);
+
+ggml_tensor* ggml_ext_ones(ggml_context* ctx,
+                           int64_t ne0,
+                           int64_t ne1,
+                           int64_t ne2,
+                           int64_t ne3);
+
+ggml_tensor* ggml_ext_ones_like(ggml_context* ctx,
+                                ggml_tensor* x);
+
+ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a);
+
+// q: [N, L_q, C(n_head*d_head)] or [N*n_head, L_q, d_head]
+// k: [N, L_k, n_kv_head*d_head] or [N*n_kv_head, L_k, d_head]
+// v: [N, L_k, n_kv_head*d_head] or [N, L_k, n_kv_head, d_head]
+// mask: [N, L_q, L_k]
+// return: [N, L_q, C]
+ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
+                                    ggml_backend_t backend,
+                                    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.0f);
+
+ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* w,
+                                 ggml_tensor* b,
+                                 float eps = EPS);
+
+ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* w,
+                                 ggml_tensor* b,
+                                 int num_groups = 32,
+                                 float eps      = 1e-6f);
+
+ggml_tensor* ggml_ext_timestep_embedding(
+    ggml_context* ctx,
+    ggml_tensor* timesteps,
+    int dim,
+    int max_period    = 10000,
+    float time_factor = 1.0f);
+
+ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx,
+                                 std::vector<ggml_tensor*>& tensors,
+                                 int dim);
+
+#endif  // __SD_CORE_GGML_EXTEND_H__

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;
 }
 
@@ -392,7 +396,7 @@ static bool backend_name_exists(const std::string& name) {
 
 static ggml_backend_t init_named_backend(const std::string& name) {
     ggml_backend_load_all_once();
-    LOG_DEBUG("Initializing backend: %s", name.c_str());
+    LOG_VERBOSE("Initializing backend: %s", name.c_str());
     if (trim_copy(name).empty()) {
         return ggml_backend_init_best();
     }
@@ -542,10 +546,10 @@ static ggml_backend_t sd_get_default_backend() {
         if (dev_count == 0) {
             LOG_ERROR("No devices found!");
         } else {
-            LOG_DEBUG("Found %zu backend devices:", dev_count);
+            LOG_VERBOSE("Found %zu backend devices:", dev_count);
             for (size_t i = 0; i < dev_count; ++i) {
                 auto dev = ggml_backend_dev_get(i);
-                LOG_DEBUG("#%zu: %s", i, ggml_backend_dev_name(dev));
+                LOG_VERBOSE("#%zu: %s", i, ggml_backend_dev_name(dev));
             }
         }
     });
@@ -587,13 +591,13 @@ static ggml_backend_t sd_get_default_backend() {
     }
 
     if (sd_backend_is_cpu(backend)) {
-        LOG_DEBUG("Using CPU backend");
+        LOG_VERBOSE("Using CPU 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,40 @@ const char* sd_backend_module_name(SDBackendModule module) {
             return "upscaler";
         case SDBackendModule::DETECTOR:
             return "detector";
+        case SDBackendModule::AUDIO_ENCODER:
+            return "audio_encoder";
     }
     return "unknown";
 }
+
+void ggml_ext_backend_tensor_get_and_sync(ggml_backend_t backend, const ggml_tensor* tensor, void* data, size_t offset, size_t size) {
+    if ((sd_backend_is(backend, "ROCm") || sd_backend_is(backend, "CUDA") || sd_backend_is(backend, "SYCL")) &&
+        !sd_backend_is_cpu(backend)) {
+        ggml_backend_tensor_get_async(backend, tensor, data, offset, size);
+        ggml_backend_synchronize(backend);
+        return;
+    }
+
+    ggml_backend_tensor_get(tensor, data, offset, size);
+}
+
+float ggml_ext_backend_tensor_get_f32(ggml_tensor* tensor) {
+    GGML_ASSERT(tensor->type == GGML_TYPE_F32 || tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_I32 || tensor->type == GGML_TYPE_BF16);
+    float value;
+    if (tensor->type == GGML_TYPE_F32) {
+        ggml_backend_tensor_get(tensor, &value, 0, sizeof(value));
+    } else if (tensor->type == GGML_TYPE_BF16) {
+        ggml_bf16_t bf16_value;
+        ggml_backend_tensor_get(tensor, &bf16_value, 0, sizeof(bf16_value));
+        value = ggml_bf16_to_fp32(bf16_value);
+    } else if (tensor->type == GGML_TYPE_F16) {
+        ggml_fp16_t f16_value;
+        ggml_backend_tensor_get(tensor, &f16_value, 0, sizeof(f16_value));
+        value = ggml_fp16_to_fp32(f16_value);
+    } else {  // GGML_TYPE_I32
+        int int32_value;
+        ggml_backend_tensor_get(tensor, &int32_value, 0, sizeof(int32_value));
+        value = (float)int32_value;
+    }
+    return value;
+}

src/core/ggml_extend_backend.h

@@ -21,6 +21,7 @@ enum class SDBackendModule {
     PHOTOMAKER,
     UPSCALER,
     DETECTOR,
+    AUDIO_ENCODER,
 };
 
 struct SDBackendAssignment {
@@ -93,7 +94,10 @@ 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);
+void ggml_ext_backend_tensor_get_and_sync(ggml_backend_t backend, const ggml_tensor* tensor, void* data, size_t offset, size_t size);
+float ggml_ext_backend_tensor_get_f32(ggml_tensor* tensor);
 #endif  // __SD_CORE_GGML_EXTEND_BACKEND_H__

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,17 +540,21 @@ 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;
     }
 
-    bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) {
+    static bool plan_matches_graph(ggml_cgraph* gf,
+                                   const Plan& plan,
+                                   const std::vector<uint64_t>& layout) {
         GGML_ASSERT(gf != nullptr);
         if (plan.leaf_names.size() != static_cast<size_t>(gf->n_leafs) ||
-            plan.layout != graph_layout(gf, false)) {
+            plan.layout != layout) {
             return false;
         }
         for (int i = 0; i < gf->n_leafs; ++i) {
@@ -558,6 +572,11 @@ namespace sd::ggml_graph_cut {
         return cut_markers == plan.cut_markers;
     }
 
+    bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) {
+        GGML_ASSERT(gf != nullptr);
+        return plan_matches_graph(gf, plan, graph_layout(gf, false));
+    }
+
     ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index) {
         GGML_ASSERT(gf != nullptr);
         if (output_index >= segment.output_node_indices.size()) {
@@ -938,20 +957,26 @@ namespace sd::ggml_graph_cut {
         GGML_ASSERT(gf != nullptr);
         GGML_ASSERT(cache != nullptr);
 
-        if (cache->graph_cut_plan.available &&
-            plan_matches_graph(gf, cache->graph_cut_plan)) {
-            return cache->graph_cut_plan;
+        const auto layout = graph_layout(gf, false);
+        auto& plans       = cache->graph_cut_plans;
+        for (auto it = plans.begin(); it != plans.end(); ++it) {
+            if (it->available && plan_matches_graph(gf, *it, layout)) {
+                plans.splice(plans.begin(), plans, it);
+                return plans.front();
+            }
         }
 
-        int64_t t_plan_begin  = ggml_time_ms();
-        Plan plan             = build_plan(backend, gf, params_tensor_set, log_desc);
-        cache->graph_cut_plan = plan;
+        int64_t t_plan_begin = ggml_time_ms();
+        plans.push_front(build_plan(backend, gf, params_tensor_set, log_desc));
+        if (plans.size() > PlanCache::MAX_PLANS) {
+            plans.pop_back();
+        }
         if (log_desc != nullptr) {
             LOG_INFO("%s build cached graph cut plan done (taking %lld ms)",
                      log_desc,
                      ggml_time_ms() - t_plan_begin);
         }
-        return plan;
+        return plans.front();
     }
 
 }  // namespace sd::ggml_graph_cut

src/core/ggml_graph_cut.h

@@ -3,6 +3,7 @@
 
 #include <array>
 #include <cstdint>
+#include <list>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
@@ -48,7 +49,8 @@ namespace sd::ggml_graph_cut {
     };
 
     struct PlanCache {
-        Plan graph_cut_plan;
+        static constexpr size_t MAX_PLANS = 4;
+        std::list<Plan> graph_cut_plans;
     };
 
     static constexpr const char* GGML_RUNNER_CUT_PREFIX = "ggml_runner_cut:";

src/core/ggml_runner.cpp

@@ -1,13 +1,696 @@
 #include <algorithm>
+#include <exception>
 #include <map>
 #include <utility>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/layer_split_partition.h"
 #include "core/segment_graph_bindings.h"
 #include "core/segment_weight_pipeline.h"
 
 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());
+    params.mem_buffer = nullptr;
+    params.no_alloc   = true;
+
+    params_ctx = ggml_init(params);
+    GGML_ASSERT(params_ctx != nullptr);
+    params_tensor_set_.clear();
+    params_tensor_set_dirty_ = true;
+}
+
+void GGMLRunner::free_params_ctx() {
+    if (params_ctx != nullptr) {
+        ggml_free(params_ctx);
+        params_ctx = nullptr;
+    }
+    params_tensor_set_.clear();
+    params_tensor_set_dirty_ = true;
+}
+
+void GGMLRunner::alloc_compute_ctx() {
+    ggml_init_params params;
+    params.mem_size   = static_cast<size_t>(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead());
+    params.mem_buffer = nullptr;
+    params.no_alloc   = true;
+
+    compute_ctx = ggml_init(params);
+    GGML_ASSERT(compute_ctx != nullptr);
+}
+
+void GGMLRunner::free_compute_ctx() {
+    debug_tensors.clear();
+    if (compute_ctx != nullptr) {
+        ggml_free(compute_ctx);
+        compute_ctx = nullptr;
+    }
+    backend_tensor_data_map.clear();
+}
+
+void GGMLRunner::rebuild_params_tensor_set() {
+    if (!params_tensor_set_dirty_) {
+        return;
+    }
+    params_tensor_set_.clear();
+    if (params_ctx == nullptr) {
+        return;
+    }
+    for (ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != nullptr; t = ggml_get_next_tensor(params_ctx, t)) {
+        params_tensor_set_.insert(t);
+    }
+    params_tensor_set_dirty_ = false;
+}
+
+ggml_tensor* GGMLRunner::canonical_param_tensor(ggml_tensor* tensor) {
+    for (auto* current = tensor; current != nullptr; current = current->view_src) {
+        if (params_tensor_set_.count(current) != 0)
+            return current;
+    }
+    return nullptr;
+}
+
+std::vector<ggml_tensor*> GGMLRunner::collect_used_param_tensors(ggml_cgraph* gf) {
+    std::vector<ggml_tensor*> used_params;
+    rebuild_params_tensor_set();
+    if (gf == nullptr || params_tensor_set_.empty()) {
+        return used_params;
+    }
+
+    std::unordered_set<const ggml_tensor*> seen_params;
+    const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
+    seen_params.reserve(static_cast<size_t>(n_leafs));
+    for (int i = 0; i < n_leafs; ++i) {
+        ggml_tensor* leaf       = sd::ggml_graph_cut::leaf_tensor(gf, i);
+        ggml_tensor* param_leaf = canonical_param_tensor(leaf);
+        if (param_leaf != nullptr &&
+            seen_params.insert(param_leaf).second) {
+            used_params.push_back(param_leaf);
+        }
+    }
+    return used_params;
+}
+
+void GGMLRunner::evict_compute_backend_param_tensors(const std::vector<ggml_tensor*>& tensors) {
+    if (tensors.empty()) {
+        return;
+    }
+    auto manager = residency_manager.lock();
+    if (manager != nullptr) {
+        manager->evict_compute_backend_params(tensors);
+    }
+}
+
+void GGMLRunner::prepare_build_in_tensor_before() {
+    one_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, 1);
+    ggml_set_name(one_tensor, "ggml_runner_build_in_tensor:one");
+    set_backend_tensor_data(one_tensor, one_vec.data());
+
+    zero_int_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, 1);
+    ggml_set_name(zero_int_tensor, "ggml_runner_build_in_tensor:zero_int");
+    set_backend_tensor_data(zero_int_tensor, zero_int_vec.data());
+}
+
+void GGMLRunner::prepare_build_in_tensor_after(ggml_cgraph* gf) {
+    ggml_build_forward_expand(gf, one_tensor);
+    ggml_build_forward_expand(gf, zero_int_tensor);
+}
+
+ggml_cgraph* GGMLRunner::new_graph_custom(size_t graph_size) {
+    if (weight_adapter) {
+        graph_size += weight_adapter->get_extra_graph_size();
+    }
+    return ggml_new_graph_custom(compute_ctx, graph_size, false);
+}
+
+ggml_cgraph* GGMLRunner::get_compute_graph(get_graph_cb_t get_graph) {
+    prepare_build_in_tensor_before();
+    ggml_cgraph* gf = get_graph();
+    if (gf == nullptr) {
+        return nullptr;
+    }
+    if (ggml_graph_n_nodes(gf) > 0) {
+        auto result = ggml_graph_node(gf, -1);
+        ggml_set_name(result, final_result_name.c_str());
+    }
+    for (const auto& entry : debug_tensors) {
+        if (entry.first != nullptr) {
+            ggml_build_forward_expand(gf, entry.first);
+        }
+    }
+    for (const auto& entry : cache_.outputs()) {
+        if (entry.second != nullptr) {
+            ggml_build_forward_expand(gf, entry.second);
+        }
+    }
+    prepare_build_in_tensor_after(gf);
+    return gf;
+}
+
+bool GGMLRunner::prepare_compute_graph(get_graph_cb_t get_graph,
+                                       ggml_cgraph** gf_out) {
+    GGML_ASSERT(gf_out != nullptr);
+
+    reset_compute_ctx();
+    ggml_cgraph* gf = get_compute_graph(get_graph);
+    if (gf == nullptr) {
+        free_compute_ctx();
+        return false;
+    }
+
+    *gf_out = gf;
+    return true;
+}
+
+ggml_backend_t GGMLRunner::backend_for_weight(const ggml_tensor* tensor) const {
+    if (tensor == nullptr || tensor->buffer == nullptr) {
+        return nullptr;
+    }
+    if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS ||
+        ggml_backend_buffer_is_host(tensor->buffer)) {
+        return nullptr;
+    }
+    ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
+    if (dev == nullptr) {
+        return nullptr;
+    }
+    if (ggml_backend_get_device(runtime_backend) == dev) {
+        return runtime_backend;
+    }
+    for (ggml_backend_t backend : extra_runtime_backends) {
+        if (ggml_backend_get_device(backend) == dev) {
+            return backend;
+        }
+    }
+    return nullptr;
+}
+
+void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) {
+    if (sched == nullptr || gf == nullptr) {
+        return;
+    }
+    ggml_backend_t current = runtime_backend;
+    const int n_nodes      = ggml_graph_n_nodes(gf);
+    for (int i = 0; i < n_nodes; i++) {
+        ggml_tensor* node    = ggml_graph_node(gf, i);
+        auto node_assignment = graph_cut_layer_split_node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i));
+        if (node_assignment != graph_cut_layer_split_node_assignments_.end()) {
+            current = node_assignment->second;
+        }
+        for (int s = 0; s < GGML_MAX_SRC; s++) {
+            ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
+            if (weight_backend != nullptr) {
+                if (node_assignment == graph_cut_layer_split_node_assignments_.end()) {
+                    current = weight_backend;
+                }
+            }
+        }
+        if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
+            node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
+            continue;
+        }
+        if (ggml_backend_supports_op(current, node)) {
+            ggml_backend_sched_set_tensor_backend(sched, node, current);
+        }
+    }
+}
+
+size_t GGMLRunner::retained_runtime_buffer_bytes(ggml_backend_t backend) const {
+    backend      = backend == nullptr ? runtime_backend : backend;
+    size_t bytes = workspace_.bytes(backend);
+    if (backend == runtime_backend) {
+        const size_t cache_bytes = cache_.resident_bytes(ggml_backend_get_device(backend));
+        bytes                    = cache_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cache_bytes;
+        const size_t cut_bytes   = cut_cache_.resident_bytes(ggml_backend_get_device(backend));
+        bytes                    = cut_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cut_bytes;
+    }
+    return bytes;
+}
+
+void GGMLRunner::sync_runtime_residency() {
+    if (auto manager = residency_manager.lock()) {
+        manager->update_runtime_residency(reinterpret_cast<uintptr_t>(this),
+                                          runtime_backend, retained_runtime_buffer_bytes());
+        for (auto backend : extra_runtime_backends) {
+            manager->update_runtime_residency(reinterpret_cast<uintptr_t>(this),
+                                              backend, retained_runtime_buffer_bytes(backend));
+        }
+    }
+}
+
+std::optional<sd::Tensor<float>> GGMLRunner::read_graph_tensor(ggml_tensor* tensor, const char* label) {
+    if (tensor == nullptr) {
+        LOG_ERROR("%s %s tensor is null", get_desc().c_str(), label);
+        return std::nullopt;
+    }
+    if (tensor->type != GGML_TYPE_F32) {
+        LOG_ERROR("%s %s tensor type mismatch: got %s",
+                  get_desc().c_str(),
+                  label,
+                  ggml_type_name(tensor->type));
+        return std::nullopt;
+    }
+    ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor);
+    if (buf == nullptr) {
+        LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p",
+                  get_desc().c_str(),
+                  label,
+                  tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
+                  ggml_op_name(tensor->op),
+                  tensor->buffer,
+                  tensor->view_src,
+                  tensor->view_src ? tensor->view_src->buffer : nullptr,
+                  tensor->data);
+        return std::nullopt;
+    }
+
+    return sd::make_sd_tensor_from_ggml<float>(tensor);
+}
+
+void GGMLRunner::copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy) {
+    GGML_ASSERT(gf != nullptr);
+    std::unordered_set<const ggml_tensor*> graph_tensor_set;
+    const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
+    const int n_nodes = ggml_graph_n_nodes(gf);
+    graph_tensor_set.reserve(static_cast<size_t>(n_leafs + n_nodes));
+    for (int i = 0; i < n_leafs; ++i) {
+        graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i));
+    }
+    for (int i = 0; i < n_nodes; ++i) {
+        graph_tensor_set.insert(ggml_graph_node(gf, i));
+    }
+
+    for (auto& kv : backend_tensor_data_map) {
+        auto tensor = kv.first;
+        auto data   = kv.second;
+        if (tensor == nullptr || data == nullptr) {
+            continue;
+        }
+        const char* name = ggml_get_name(tensor);
+        if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) {
+            continue;
+        }
+        if (tensor->buffer == nullptr) {
+            LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s",
+                     get_desc().c_str(),
+                     name != nullptr ? name : "",
+                     (long long)tensor->ne[0],
+                     (long long)tensor->ne[1],
+                     (long long)tensor->ne[2],
+                     (long long)tensor->ne[3],
+                     ggml_type_name(tensor->type));
+            continue;
+        }
+
+        ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+        if (buf == nullptr) {
+            LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p",
+                     get_desc().c_str(),
+                     tensor && tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
+                     tensor ? ggml_op_name(tensor->op) : "<null>",
+                     data,
+                     tensor ? tensor->view_src : nullptr,
+                     (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr);
+            continue;
+        }
+
+        ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
+    }
+
+    if (clear_after_copy) {
+        backend_tensor_data_map.clear();
+    }
+}
+
+bool GGMLRunner::resolve_graph_cut_plan(ggml_cgraph* gf,
+                                        GraphCutPlan* plan_out) {
+    GGML_ASSERT(plan_out != nullptr);
+    GGML_ASSERT(gf != nullptr);
+    *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
+                                                 gf,
+                                                 &graph_cut_plan_cache_,
+                                                 params_tensor_set_,
+                                                 get_desc().c_str());
+    return true;
+}
+
+bool GGMLRunner::resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
+                                                    GraphCutPlan* plan_out) {
+    return resolve_graph_cut_plan(gf, plan_out);
+}
+
+bool GGMLRunner::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
+    graph_cut_layer_split_node_assignments_.clear();
+    if (!graph_cut_layer_split_enabled) {
+        return true;
+    }
+    if (!is_multi_device()) {
+        LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str());
+        return false;
+    }
+
+    GraphCutPlan plan;
+    if (!resolve_graph_cut_layer_split_plan(gf, &plan)) {
+        return false;
+    }
+    if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
+        auto manager = residency_manager.lock();
+        if (manager == nullptr) {
+            LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
+            return false;
+        }
+        std::vector<ggml_tensor*> graph_params = collect_used_param_tensors(gf);
+        if (!graph_params.empty() &&
+            !manager->assign_compute_backend(graph_params, runtime_backend)) {
+            LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
+                      get_desc().c_str(),
+                      sd::layer_split_backend_device_display_name(runtime_backend).c_str());
+            return false;
+        }
+        for (ggml_tensor* param : graph_params) {
+            if (param != nullptr) {
+                graph_cut_layer_split_assignments_[param] = runtime_backend;
+            }
+        }
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        for (int i = 0; i < n_nodes; i++) {
+            ggml_tensor* node = ggml_graph_node(gf, i);
+            if (node != nullptr) {
+                graph_cut_layer_split_node_assignments_[node] = runtime_backend;
+            }
+        }
+        if (!graph_cut_layer_split_primary_notice_logged_) {
+            LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                     get_desc().c_str(),
+                     sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
+                     graph_params.size());
+            graph_cut_layer_split_primary_notice_logged_ = true;
+        } else {
+            LOG_VERBOSE("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                        get_desc().c_str(),
+                        sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
+                        graph_params.size());
+        }
+        return true;
+    }
+
+    std::vector<ggml_backend_t> split_backends;
+    split_backends.reserve(extra_runtime_backends.size() + 1);
+    split_backends.push_back(runtime_backend);
+    for (ggml_backend_t backend : extra_runtime_backends) {
+        if (backend != nullptr) {
+            split_backends.push_back(backend);
+        }
+    }
+
+    auto manager = residency_manager.lock();
+    if (manager == nullptr) {
+        LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
+        return false;
+    }
+
+    sd::GraphCutLayerSplitAssignment assignment;
+    auto canonicalize_param = [this](ggml_tensor* tensor) {
+        return canonical_param_tensor(tensor);
+    };
+    if (!sd::partition_graph_cut_layer_split(get_desc().c_str(),
+                                             gf,
+                                             plan,
+                                             split_backends,
+                                             graph_cut_layer_split_backend_vram_limits_,
+                                             max_graph_vram_bytes,
+                                             graph_cut_layer_split_assignments_,
+                                             canonicalize_param,
+                                             &assignment)) {
+        return false;
+    }
+
+    for (size_t i = 0; i < split_backends.size(); i++) {
+        if (assignment.tensors_by_backend[i].empty()) {
+            continue;
+        }
+        if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) {
+            LOG_ERROR("%s graph-cut layer split failed to assign params to %s",
+                      get_desc().c_str(),
+                      sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
+            return false;
+        }
+    }
+
+    graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments);
+    sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment);
+
+    return true;
+}
+
+bool GGMLRunner::runner_start() {
+    if (runner_started_) {
+        return true;
+    }
+    cache_.clear();
+    workspace_.set_extra_backends(extra_runtime_backends);
+    if (auto manager = residency_manager.lock()) {
+        manager->set_workspace_reclaimer(reinterpret_cast<uintptr_t>(this), [this]() {
+            if (!workspace_.release()) {
+                return false;
+            }
+            sync_runtime_residency();
+            return true;
+        });
+    }
+    runner_started_ = true;
+    return true;
+}
+
+void GGMLRunner::runner_end() {
+    GGML_ASSERT(!graph_active_);
+    if (!runner_started_) {
+        return;
+    }
+    workspace_.release();
+    cache_.clear();
+    logged_compute_bytes_.clear();
+    logged_segment_count_ = 0;
+    if (auto manager = residency_manager.lock()) {
+        manager->clear_prefetched_params(reinterpret_cast<uintptr_t>(this));
+        std::vector<ggml_tensor*> tensors;
+        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));
+    }
+    runner_started_ = false;
+}
+
+GGMLRunner::GGMLRunner(ggml_backend_t backend,
+                       std::shared_ptr<DeviceResidencyManager> manager)
+    : runtime_backend(backend),
+      cache_(backend),
+      cut_cache_(backend),
+      workspace_(backend),
+      residency_manager(manager) {
+    GGML_ASSERT(runtime_backend != nullptr);
+    alloc_params_ctx();
+}
+
+GGMLRunner::~GGMLRunner() {
+    runner_end();
+    free_compute_ctx();
+    free_params_ctx();
+}
+
+GGMLRunnerContext GGMLRunner::get_context() {
+    GGMLRunnerContext runner_ctx;
+    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;
+    runner_ctx.weight_adapter        = weight_adapter;
+    runner_ctx.debug_tensors         = &debug_tensors;
+    runner_ctx.get_cache_tensor      = [this](const std::string& name) {
+        return this->get_cache_tensor_by_name(name);
+    };
+    runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) {
+        this->cache(name, tensor);
+    };
+    runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) {
+        this->set_backend_tensor_data(tensor, data);
+    };
+    return runner_ctx;
+}
+
+void GGMLRunner::reset_compute_ctx() {
+    free_compute_ctx();
+    alloc_compute_ctx();
+}
+
+void GGMLRunner::free_cache_ctx_and_buffer() {
+    cache_.clear();
+    sync_runtime_residency();
+}
+
+void GGMLRunner::set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
+    // The scheduler only allocates standalone data tensors when they are
+    // marked as graph inputs. The flag is harmless for single-backend graphs.
+    ggml_set_input(tensor);
+    backend_tensor_data_map[tensor] = data;
+}
+
+ggml_tensor* GGMLRunner::to_backend(ggml_tensor* tensor) {
+    GGML_ASSERT(compute_ctx != nullptr);
+    if (tensor == nullptr) {
+        return nullptr;
+    }
+    // it's performing a compute, check if backend isn't cpu
+    if (!sd_backend_is_cpu(runtime_backend) && (tensor->buffer == nullptr || ggml_backend_buffer_is_host(tensor->buffer))) {
+        // pass input tensors to gpu memory
+        auto backend_tensor = ggml_dup_tensor(compute_ctx, tensor);
+
+        set_backend_tensor_data(backend_tensor, tensor->data);
+        return backend_tensor;
+    } else {
+        return tensor;
+    }
+}
+
+void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
+    if (tensor != nullptr && tensor->view_src != nullptr) {
+        tensor = ggml_cont(compute_ctx, tensor);
+    }
+    if (tensor != nullptr) {
+        ggml_set_output(tensor);
+    }
+    cache_.stage(name, tensor);
+}
+
+std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
+                                                     int n_threads,
+                                                     bool auto_runner_end,
+                                                     bool no_return,
+                                                     const std::function<bool()>& read_outputs) {
+    if (graph_active_) {
+        LOG_ERROR("%s does not support reentrant graph execution", get_desc().c_str());
+        return std::nullopt;
+    }
+    if (!runner_start()) {
+        runner_end();
+        return std::nullopt;
+    }
+    struct RunnerEndGuard {
+        GGMLRunner& runner;
+        bool enabled;
+        ~RunnerEndGuard() {
+            if (enabled) {
+                runner.runner_end();
+            }
+        }
+    } runner_guard{*this, auto_runner_end};
+    graph_active_ = true;
+    bool success  = false;
+    struct GraphEndGuard {
+        GGMLRunner& runner;
+        const bool& success;
+        ~GraphEndGuard() {
+            runner.workspace_.segment_end();
+            runner.cache_.graph_end(false);
+            runner.cut_cache_.clear();
+            runner.free_compute_ctx();
+            runner.graph_active_ = false;
+            if (!success) {
+                runner.workspace_.release();
+            }
+            runner.sync_runtime_residency();
+        }
+    } graph_guard{*this, success};
+
+    ggml_cgraph* graph = nullptr;
+    if (!prepare_compute_graph(get_graph, &graph)) {
+        return std::nullopt;
+    }
+    params_tensor_set_dirty_ = true;
+    rebuild_params_tensor_set();
+    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);
+    }
+    return output;
+}
+
+void GGMLRunner::set_graph_cut_layer_split_enabled(bool enabled) {
+    graph_cut_layer_split_enabled = enabled;
+    if (!enabled) {
+        graph_cut_layer_split_assignments_.clear();
+        graph_cut_layer_split_node_assignments_.clear();
+        graph_cut_layer_split_primary_notice_logged_ = false;
+    }
+}
+
+void GGMLRunner::set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {
+    graph_cut_layer_split_backend_vram_limits_ = limits;
+    graph_cut_layer_split_assignments_.clear();
+    graph_cut_layer_split_node_assignments_.clear();
+    graph_cut_layer_split_primary_notice_logged_ = false;
+}
+
+void GGMLRunner::set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
+    extra_runtime_backends.clear();
+    for (ggml_backend_t backend : backends) {
+        if (backend == nullptr || backend == runtime_backend) {
+            continue;
+        }
+        if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) ==
+            extra_runtime_backends.end()) {
+            extra_runtime_backends.push_back(backend);
+        }
+    }
+    workspace_.set_extra_backends(extra_runtime_backends);
+    graph_cut_layer_split_assignments_.clear();
+    graph_cut_layer_split_node_assignments_.clear();
+    graph_cut_layer_split_primary_notice_logged_ = false;
+}
+
 static size_t add_bytes(size_t a, size_t b) {
     return b > SIZE_MAX - a ? SIZE_MAX : a + b;
 }
@@ -170,8 +853,8 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
     }
     const bool segments_changed = plan.segments.size() != logged_segment_count_;
     if (segments_changed && (segmented || logged_segment_count_ > 1)) {
-        LOG_DEBUG("%s using %zu segment%s", get_desc().c_str(),
-                  plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
+        LOG_VERBOSE("%s using %zu segment%s", get_desc().c_str(),
+                    plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
     }
     SegmentGraphBindings bindings(cut_cache_, plan, graph);
     SegmentWeightPipeline weights(manager, runtime_backend, reinterpret_cast<uintptr_t>(this),
@@ -261,6 +944,8 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
         if (!prefetch_requests.empty()) {
             weights.enqueue_next(index, prefetch_requests.front());
         }
+        LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(),
+                  index + 1, plan.segments.size(), segment.group_name.c_str());
         if (!execute_segment(segment_graph, n_threads) ||
             !cache_.capture(segment_graph) ||
             !cut_cache_.capture(graph, segment, get_desc().c_str())) {
@@ -273,21 +958,24 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
             }
             if (!no_return) {
                 auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str());
-                output      = read_graph_tensor<float>(result, "output");
+                output      = read_graph_tensor(result, "output");
                 if (!output.has_value()) {
                     return fail_segment("output readback");
                 }
             }
         }
+        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);
     }
     if (segments_changed || peak_compute_bytes != logged_compute_bytes_) {
         for (const auto& entry : peak_compute_bytes) {
-            LOG_DEBUG("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)",
-                      get_desc().c_str(), entry.second / (1024.0 * 1024.0),
-                      sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first),
-                      plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
+            LOG_VERBOSE("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)",
+                        get_desc().c_str(), entry.second / (1024.0 * 1024.0),
+                        sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first),
+                        plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
         }
         logged_compute_bytes_ = std::move(peak_compute_bytes);
         logged_segment_count_ = plan.segments.size();

src/core/ggml_runner.h

@@ -0,0 +1,369 @@
+#ifndef __SD_CORE_GGML_RUNNER_H__
+#define __SD_CORE_GGML_RUNNER_H__
+
+#include <cstddef>
+#include <functional>
+#include <map>
+#include <memory>
+#include <optional>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "core/compute_workspace.h"
+#include "core/ggml_graph_cut.h"
+#include "core/runner_cache.h"
+#include "core/tensor_ggml.hpp"
+#include "core/util.h"
+#include "device_residency_manager.h"
+
+/* SDXL with LoRA requires more space */
+#define MAX_PARAMS_TENSOR_NUM 32768
+#define MAX_GRAPH_SIZE 327680
+
+struct WeightAdapter {
+    struct ForwardParams {
+        enum class op_type_t {
+            OP_LINEAR,
+            OP_CONV2D,
+        } op_type;
+        struct {
+            bool force_prec_f32 = false;
+            float scale         = 1.f;
+        } linear;
+        struct conv2d_params_t {
+            int s0          = 1;
+            int s1          = 1;
+            int p0          = 0;
+            int p1          = 0;
+            int d0          = 1;
+            int d1          = 1;
+            bool direct     = false;
+            bool circular_x = false;
+            bool circular_y = false;
+            float scale     = 1.f;
+        } conv2d;
+    };
+    virtual ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name) = 0;
+    virtual ggml_tensor* forward_with_lora(ggml_context* ctx,
+                                           ggml_backend_t backend,
+                                           ggml_tensor* x,
+                                           ggml_tensor* w,
+                                           ggml_tensor* b,
+                                           const std::string& prefix,
+                                           ForwardParams forward_params)                                                              = 0;
+    virtual ggml_tensor* add_lora_to_output(ggml_context* ctx,
+                                            ggml_backend_t backend,
+                                            ggml_tensor* x,
+                                            ggml_tensor* w,
+                                            ggml_tensor* output,
+                                            const std::string& prefix,
+                                            ForwardParams forward_params)                                                             = 0;
+    virtual size_t get_extra_graph_size()                                                                                             = 0;
+};
+
+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;
+    ggml_tensor* ip_context                                          = nullptr;
+    float ip_scale                                                   = 1.0f;
+    std::shared_ptr<WeightAdapter> weight_adapter                    = nullptr;
+    std::vector<std::pair<ggml_tensor*, std::string>>* debug_tensors = nullptr;
+    std::function<ggml_tensor*(const std::string&)> get_cache_tensor;
+    std::function<void(const std::string&, ggml_tensor*)> cache_tensor;
+    std::function<void(ggml_tensor*, const void*)> set_backend_tensor_data;
+    std::map<std::pair<ggml_tensor*, int>, ggml_tensor*> int8_convrot_cache;
+
+    void capture_tensor(const std::string& name, ggml_tensor* tensor) {
+        if (debug_tensors == nullptr || tensor == nullptr) {
+            return;
+        }
+        ggml_tensor* snapshot = ggml_cont(ggml_ctx, tensor);
+        ggml_tensor* dst      = ggml_dup_tensor(ggml_ctx, snapshot);
+        snapshot              = ggml_cpy(ggml_ctx, snapshot, dst);
+        ggml_set_output(snapshot);
+        debug_tensors->push_back({snapshot, name});
+    }
+
+    ggml_tensor* load_cache_tensor(const std::string& name) const {
+        if (!get_cache_tensor) {
+            return nullptr;
+        }
+        return get_cache_tensor(name);
+    }
+
+    void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const {
+        if (!cache_tensor || tensor == nullptr) {
+            return;
+        }
+        cache_tensor(name, tensor);
+    }
+
+    void bind_backend_tensor_data(ggml_tensor* tensor, const void* data) const {
+        if (!set_backend_tensor_data || tensor == nullptr || data == nullptr) {
+            return;
+        }
+        set_backend_tensor_data(tensor, data);
+    }
+};
+
+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_;
+    size_t logged_segment_count_ = 0;
+
+    sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes);
+    std::vector<DeviceMemoryRequest> memory_requests(const std::vector<sd::BackendBufferSize>& sizes,
+                                                     size_t pending_cache_bytes) const;
+    bool fits(const std::vector<DeviceMemoryRequest>& requests,
+              const std::vector<ggml_tensor*>& params) const;
+    bool execute_segment(ggml_cgraph* graph, int n_threads);
+    std::optional<sd::Tensor<float>> execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs);
+
+protected:
+    typedef std::function<ggml_cgraph*()> get_graph_cb_t;
+    using GraphCutPlan = sd::ggml_graph_cut::Plan;
+
+    ggml_backend_t runtime_backend = nullptr;
+
+    ggml_context* params_ctx = nullptr;
+
+    sd::RunnerCache cache_;
+    sd::GraphCutTensorCache cut_cache_;
+    sd::ComputeWorkspace workspace_;
+    ggml_context* compute_ctx = nullptr;
+    bool runner_started_      = false;
+    bool graph_active_        = false;
+
+    size_t max_graph_vram_bytes        = 0;
+    bool graph_cut_layer_split_enabled = false;
+    std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
+
+    std::vector<ggml_backend_t> extra_runtime_backends;  // borrowed (SDBackendManager-owned)
+    bool multi_device_eval_callback_warned = false;
+
+    std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
+    std::weak_ptr<DeviceResidencyManager> residency_manager;
+    bool params_tensor_set_dirty_ = true;
+
+    std::vector<float> one_vec = {1.f};
+    ggml_tensor* one_tensor    = nullptr;
+
+    std::vector<int> zero_int_vec = {0};
+    ggml_tensor* zero_int_tensor  = nullptr;
+
+    std::map<ggml_tensor*, const void*> backend_tensor_data_map;
+    std::vector<std::pair<ggml_tensor*, std::string>> debug_tensors;
+    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;
+
+    sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_;
+    std::unordered_set<const ggml_tensor*> params_tensor_set_;
+    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_assignments_;
+    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_node_assignments_;
+    bool graph_cut_layer_split_primary_notice_logged_ = false;
+
+    template <typename T>
+    static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
+        if (!tensor.has_value()) {
+            return {};
+        }
+        return std::move(*tensor);
+    }
+
+    template <typename T>
+    static sd::Tensor<T> restore_trailing_singleton_dims(std::optional<sd::Tensor<T>> tensor,
+                                                         size_t expected_dim) {
+        return restore_trailing_singleton_dims(take_or_empty(std::move(tensor)), expected_dim);
+    }
+
+    template <typename T>
+    static sd::Tensor<T> restore_trailing_singleton_dims(sd::Tensor<T> tensor,
+                                                         size_t expected_dim) {
+        if (tensor.empty()) {
+            return tensor;
+        }
+        while (static_cast<size_t>(tensor.dim()) < expected_dim) {
+            tensor.unsqueeze_(tensor.dim());
+        }
+        return tensor;
+    }
+
+    void alloc_params_ctx();
+
+    void free_params_ctx();
+
+    void alloc_compute_ctx();
+
+    void free_compute_ctx();
+
+    void rebuild_params_tensor_set();
+
+    ggml_tensor* canonical_param_tensor(ggml_tensor* tensor);
+
+    std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* gf);
+
+    void evict_compute_backend_param_tensors(const std::vector<ggml_tensor*>& tensors);
+
+    void prepare_build_in_tensor_before();
+
+    void prepare_build_in_tensor_after(ggml_cgraph* gf);
+
+    ggml_cgraph* new_graph_custom(size_t graph_size);
+
+    ggml_cgraph* get_compute_graph(get_graph_cb_t get_graph);
+
+    bool prepare_compute_graph(get_graph_cb_t get_graph,
+                               ggml_cgraph** gf_out);
+
+    ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const;
+
+    // Weightless ops have no scheduler anchor, so pin them to the most recent
+    // weight device. Views must stay unpinned or cross-device copies can be
+    // skipped for their consumers.
+    void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph = nullptr);
+
+    bool is_multi_device() const {
+        return !extra_runtime_backends.empty();
+    }
+
+    size_t reusable_compute_buffer_bytes() const {
+        return workspace_.bytes(runtime_backend);
+    }
+
+    size_t retained_runtime_buffer_bytes(ggml_backend_t backend = nullptr) const;
+
+    void sync_runtime_residency();
+
+    std::optional<sd::Tensor<float>> read_graph_tensor(ggml_tensor* tensor, const char* label);
+
+    void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true);
+
+    bool resolve_graph_cut_plan(ggml_cgraph* gf,
+                                GraphCutPlan* plan_out);
+
+    bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
+                                            GraphCutPlan* plan_out);
+
+    bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf);
+
+public:
+    bool runner_start();
+
+    bool runner_started() const { return runner_started_; }
+
+    void runner_end();
+
+public:
+    virtual std::string get_desc() = 0;
+
+    GGMLRunner(ggml_backend_t backend,
+               std::shared_ptr<DeviceResidencyManager> manager = nullptr);
+
+    virtual ~GGMLRunner();
+
+    virtual GGMLRunnerContext get_context();
+
+    void reset_compute_ctx();
+
+public:
+    void free_cache_ctx_and_buffer();
+
+    // do copy after alloc graph
+    void set_backend_tensor_data(ggml_tensor* tensor, const void* data);
+
+    template <typename T>
+    ggml_tensor* make_input(const sd::Tensor<T>& tensor) {
+        ggml_tensor* input = sd::make_ggml_tensor(compute_ctx, tensor, false);
+        set_backend_tensor_data(input, tensor.data());
+        return input;
+    }
+
+    template <typename T>
+    ggml_tensor* make_optional_input(const sd::Tensor<T>& tensor) {
+        if (tensor.empty()) {
+            return nullptr;
+        }
+        return make_input(tensor);
+    }
+
+    template <typename T>
+    ggml_tensor* make_optional_input(const sd::Tensor<T>* tensor) {
+        if (tensor == nullptr) {
+            return nullptr;
+        }
+        return make_input(*tensor);
+    }
+
+    ggml_tensor* to_backend(ggml_tensor* tensor);
+
+    void cache(const std::string name, ggml_tensor* tensor);
+
+    ggml_tensor* get_cache_tensor_by_name(const std::string& name) {
+        return cache_.get(name);
+    }
+
+    std::optional<sd::Tensor<float>> compute(get_graph_cb_t get_graph,
+                                             int n_threads,
+                                             bool auto_runner_end                      = true,
+                                             bool no_return                            = false,
+                                             const std::function<bool()>& read_outputs = {});
+
+    void set_flash_attention_enabled(bool enabled) {
+        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;
+    }
+
+    void set_circular_axes(bool circular_x, bool circular_y) {
+        circular_x_enabled = circular_x;
+        circular_y_enabled = circular_y;
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {
+        weight_adapter = adapter;
+    }
+
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) {
+        max_graph_vram_bytes = max_vram_bytes;
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled);
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits);
+
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends);
+};
+
+#endif  // __SD_CORE_GGML_RUNNER_H__

src/core/ggml_tensor_utils.cpp

@@ -0,0 +1,428 @@
+#include "core/ggml_tensor_utils.h"
+
+#include <algorithm>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+
+#include "core/ggml_extend_backend.h"
+#include "core/rng.hpp"
+
+void ggml_ext_im_set_randn_f32(ggml_tensor* tensor, std::shared_ptr<RNG> rng) {
+    uint32_t n                        = (uint32_t)ggml_nelements(tensor);
+    std::vector<float> random_numbers = rng->randn(n);
+    for (uint32_t i = 0; i < n; i++) {
+        ggml_ext_im_set_f32_1d(tensor, i, random_numbers[i]);
+    }
+}
+
+void print_ggml_tensor(ggml_tensor* tensor, bool shape_only, const char* mark) {
+    printf("%s (%s): shape(%zu, %zu, %zu, %zu)\n", mark, ggml_type_name(tensor->type), tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
+    fflush(stdout);
+    if (shape_only) {
+        return;
+    }
+    int range = 3;
+    for (int i3 = 0; i3 < tensor->ne[3]; i3++) {
+        if (i3 >= range && i3 + range < tensor->ne[3]) {
+            continue;
+        }
+        for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
+            if (i2 >= range && i2 + range < tensor->ne[2]) {
+                continue;
+            }
+            for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
+                if (i1 >= range && i1 + range < tensor->ne[1]) {
+                    continue;
+                }
+                for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
+                    if (i0 >= range && i0 + range < tensor->ne[0]) {
+                        continue;
+                    }
+                    if (tensor->type == GGML_TYPE_F32) {
+                        printf("  [%d, %d, %d, %d] = %f\n", i3, i2, i1, i0, ggml_ext_tensor_get_f32(tensor, i0, i1, i2, i3));
+                    } else if (tensor->type == GGML_TYPE_F16) {
+                        printf("  [%d, %d, %d, %d] = %f\n", i3, i2, i1, i0, ggml_fp16_to_fp32(ggml_ext_tensor_get_f16(tensor, i0, i1, i2, i3)));
+                    } else if (tensor->type == GGML_TYPE_I32) {
+                        printf("  [%d, %d, %d, %d] = %i3\n", i3, i2, i1, i0, ggml_ext_tensor_get_i32(tensor, i0, i1, i2, i3));
+                    }
+                    fflush(stdout);
+                }
+            }
+        }
+    }
+}
+
+void ggml_ext_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t, int64_t, int64_t, int64_t)>& fn) {
+    int64_t n0 = tensor->ne[0];
+    int64_t n1 = tensor->ne[1];
+    int64_t n2 = tensor->ne[2];
+    int64_t n3 = tensor->ne[3];
+
+    for (int64_t i3 = 0; i3 < n3; i3++) {
+        for (int64_t i2 = 0; i2 < n2; i2++) {
+            for (int64_t i1 = 0; i1 < n1; i1++) {
+                for (int64_t i0 = 0; i0 < n0; i0++) {
+                    fn(tensor, i0, i1, i2, i3);
+                }
+            }
+        }
+    }
+}
+
+void ggml_ext_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t)>& fn) {
+    int64_t n0 = tensor->ne[0];
+    int64_t n1 = tensor->ne[1];
+    int64_t n2 = tensor->ne[2];
+    int64_t n3 = tensor->ne[3];
+
+    for (int64_t i = 0; i < ggml_nelements(tensor); i++) {
+        fn(tensor, i);
+    }
+}
+
+void ggml_ext_tensor_diff(
+    ggml_tensor* a,
+    ggml_tensor* b,
+    float gap) {
+    GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));
+    ggml_ext_tensor_iter(a, [&](ggml_tensor* a, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+        float a_value = ggml_ext_tensor_get_f32(a, i0, i1, i2, i3);
+        float b_value = ggml_ext_tensor_get_f32(b, i0, i1, i2, i3);
+        if (abs(a_value - b_value) > gap) {
+            LOG_WARN("[%ld, %ld, %ld, %ld] %f %f", i3, i2, i1, i0, a_value, b_value);
+        }
+    });
+}
+
+ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        LOG_ERROR("failed to open '%s'", file_path.c_str());
+        return nullptr;
+    }
+    int32_t n_dims;
+    int32_t length;
+    int32_t ttype;
+
+    file.read(reinterpret_cast<char*>(&n_dims), sizeof(n_dims));
+    file.read(reinterpret_cast<char*>(&length), sizeof(length));
+    file.read(reinterpret_cast<char*>(&ttype), sizeof(ttype));
+
+    LOG_VERBOSE("load_tensor_from_file %d %d %d", n_dims, length, ttype);
+
+    if (file.eof()) {
+        LOG_ERROR("incomplete file '%s'", file_path.c_str());
+        return nullptr;
+    }
+
+    int32_t nelements = 1;
+    int32_t ne[4]     = {1, 1, 1, 1};
+    for (int i = 0; i < n_dims; ++i) {
+        file.read(reinterpret_cast<char*>(&ne[i]), sizeof(ne[i]));
+        nelements *= ne[i];
+    }
+    std::string name(length, 0);
+    file.read(&name[0], length);
+    ggml_tensor* tensor = ggml_new_tensor_4d(ctx, (ggml_type)ttype, ne[0], ne[1], ne[2], ne[3]);
+    const size_t bpe    = ggml_type_size(ggml_type(ttype));
+    file.read(reinterpret_cast<char*>(tensor->data), ggml_nbytes(tensor));
+    return tensor;
+}
+
+// __STATIC_INLINE__ void save_tensor_to_file(const std::string& file_name, ggml_tensor* tensor, const std::string & name) {
+//     std::string file_name_ = file_name + ".tensor";
+//     std::string name_ = name;
+//     std::ofstream file("./" + file_name_, std::ios::binary);
+//     file.write(reinterpret_cast<char*>(&tensor->n_dims), sizeof(tensor->n_dims));
+//     int len = (int)name_.size();
+//     file.write(reinterpret_cast<char*>(&len), sizeof(len));
+//     int ttype = (int)tensor->type;
+//     file.write(reinterpret_cast<char*>(&ttype), sizeof(ttype));
+//     for (int i = 0; i < tensor->n_dims; ++i) {
+//         int ne_ = (int) tensor->ne[i];
+//         file.write(reinterpret_cast<char*>(&ne_), sizeof(ne_));
+//     }
+//     file.write(&name_[0], len);
+//     char* data = nullptr;
+//     file.write((char*)tensor->data, ggml_nbytes(tensor));
+//     file.close();
+// }
+
+uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t* image_data) {
+    int64_t width    = input->ne[0];
+    int64_t height   = input->ne[1];
+    int64_t channels = input->ne[2];
+    GGML_ASSERT(input->type == GGML_TYPE_F32);
+    if (image_data == nullptr) {
+        image_data = (uint8_t*)malloc(width * height * channels);
+    }
+    for (int iy = 0; iy < height; iy++) {
+        for (int ix = 0; ix < width; ix++) {
+            for (int k = 0; k < channels; k++) {
+                float value                                               = ggml_ext_tensor_get_f32(input, ix, iy, k);
+                *(image_data + iy * width * channels + ix * channels + k) = (uint8_t)(value * 255.0f);
+            }
+        }
+    }
+    return image_data;
+}
+
+uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, int idx, bool video) {
+    int64_t width  = input->ne[0];
+    int64_t height = input->ne[1];
+    int64_t channels;
+    if (video) {
+        channels = input->ne[3];
+    } else {
+        channels = input->ne[2];
+    }
+    GGML_ASSERT(channels == 3 && input->type == GGML_TYPE_F32);
+    uint8_t* image_data = (uint8_t*)malloc(width * height * channels);
+    for (int ih = 0; ih < height; ih++) {
+        for (int iw = 0; iw < width; iw++) {
+            for (int ic = 0; ic < channels; ic++) {
+                float value;
+                if (video) {
+                    value = ggml_ext_tensor_get_f32(input, iw, ih, idx, ic);
+                } else {
+                    value = ggml_ext_tensor_get_f32(input, iw, ih, ic, idx);
+                }
+                *(image_data + ih * width * channels + iw * channels + ic) = (uint8_t)(value * 255.0f);
+            }
+        }
+    }
+    return image_data;
+}
+
+void sd_image_to_ggml_tensor(sd_image_t image,
+                             ggml_tensor* tensor,
+                             bool scale) {
+    GGML_ASSERT(image.width == tensor->ne[0]);
+    GGML_ASSERT(image.height == tensor->ne[1]);
+    GGML_ASSERT(image.channel == tensor->ne[2]);
+    GGML_ASSERT(1 == tensor->ne[3]);
+    GGML_ASSERT(tensor->type == GGML_TYPE_F32);
+    ggml_ext_tensor_iter(tensor, [&](ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+        float value = sd_image_get_f32(image, i0, i1, i2, scale);
+        ggml_ext_tensor_set_f32(tensor, value, i0, i1, i2, i3);
+    });
+}
+
+void ggml_ext_tensor_apply_mask(ggml_tensor* image_data,
+                                ggml_tensor* mask,
+                                ggml_tensor* output,
+                                float masked_value) {
+    int64_t width    = output->ne[0];
+    int64_t height   = output->ne[1];
+    int64_t channels = output->ne[2];
+    float rescale_mx = 1.f * mask->ne[0] / output->ne[0];
+    float rescale_my = 1.f * mask->ne[1] / output->ne[1];
+    GGML_ASSERT(output->type == GGML_TYPE_F32);
+    for (int ix = 0; ix < width; ix++) {
+        for (int iy = 0; iy < height; iy++) {
+            int mx  = (int)(ix * rescale_mx);
+            int my  = (int)(iy * rescale_my);
+            float m = ggml_ext_tensor_get_f32(mask, mx, my);
+            m       = round(m);  // inpaint models need binary masks
+            ggml_ext_tensor_set_f32(mask, m, mx, my);
+            for (int k = 0; k < channels; k++) {
+                float value = ggml_ext_tensor_get_f32(image_data, ix, iy, k);
+                value       = (1 - m) * (value - masked_value) + masked_value;
+                ggml_ext_tensor_set_f32(output, value, ix, iy, k);
+            }
+        }
+    }
+}
+
+float ggml_ext_tensor_mean(ggml_tensor* src) {
+    float mean        = 0.0f;
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        mean += data[i] / nelements * 1.0f;
+    }
+    return mean;
+}
+
+void ggml_ext_tensor_add_inplace(ggml_tensor* a, ggml_tensor* b) {
+    GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));
+    int64_t nelements = ggml_nelements(a);
+    float* vec_a      = (float*)a->data;
+    float* vec_b      = (float*)b->data;
+    for (int i = 0; i < nelements; i++) {
+        vec_a[i] = vec_a[i] + vec_b[i];
+    }
+}
+
+void ggml_ext_tensor_scale_inplace(ggml_tensor* src, float scale) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        data[i] = data[i] * scale;
+    }
+}
+
+void ggml_ext_tensor_clamp_inplace(ggml_tensor* src, float min, float max) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        float val = data[i];
+        data[i]   = val < min ? min : (val > max ? max : val);
+    }
+}
+
+ggml_tensor* ggml_ext_tensor_concat(ggml_context* ctx,
+                                    ggml_tensor* a,
+                                    ggml_tensor* b,
+                                    int dim) {
+    int64_t ne[GGML_MAX_DIMS];
+    for (int d = 0; d < GGML_MAX_DIMS; ++d) {
+        if (d == dim) {
+            ne[d] = a->ne[d] + b->ne[d];
+            continue;
+        }
+        GGML_ASSERT(a->ne[d] == b->ne[d]);
+        ne[d] = a->ne[d];
+    }
+    ggml_tensor* result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne);
+    int64_t o[4]        = {0, 0, 0, 0};
+    o[dim]              = a->ne[dim];
+
+    float v;
+    for (int i3 = 0; i3 < result->ne[3]; i3++) {
+        for (int i2 = 0; i2 < result->ne[2]; i2++) {
+            for (int i1 = 0; i1 < result->ne[1]; i1++) {
+                for (int i0 = 0; i0 < result->ne[0]; i0++) {
+                    if (i0 < a->ne[0] && i1 < a->ne[1] && i2 < a->ne[2] && i3 < a->ne[3]) {
+                        v = ggml_ext_tensor_get_f32(a, i0, i1, i2, i3);
+                    } else {
+                        v = ggml_ext_tensor_get_f32(b, i0 - o[0], i1 - o[1], i2 - o[2], i3 - o[3]);
+                    }
+
+                    ggml_ext_tensor_set_f32(result, v, i0, i1, i2, i3);
+                }
+            }
+        }
+    }
+    return result;
+}
+
+void scale_to_minus1_1(ggml_tensor* src) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        float val = data[i];
+        data[i]   = val * 2.0f - 1.0f;
+    }
+}
+
+void scale_to_0_1(ggml_tensor* src) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        float val = data[i];
+        data[i]   = (val + 1.0f) * 0.5f;
+    }
+}
+
+ggml_tensor* vector_to_ggml_tensor(ggml_context* ctx,
+                                   const std::vector<float>& vec) {
+    ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, vec.size());
+    memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t));
+    return t;
+}
+
+ggml_tensor* vector_to_ggml_tensor_i32(ggml_context* ctx,
+                                       const std::vector<int>& vec) {
+    ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, vec.size());
+    memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t));
+    return t;
+}
+
+std::vector<float> arange(float start, float end, float step) {
+    std::vector<float> result;
+
+    for (float value = start; value < end; value += step) {
+        result.push_back(value);
+    }
+
+    return result;
+}
+
+std::vector<float> timestep_embedding(std::vector<float> timesteps,
+                                      int dim,
+                                      int max_period,
+                                      bool flip_sin_to_cos,
+                                      float scale) {
+    // timesteps: [N,]
+    // embedding: [N, dim]
+    size_t N = timesteps.size();
+    std::vector<float> embedding(N * dim, 0.f);
+    int half = dim / 2;
+    std::vector<float> freqs(half);
+    for (int i = 0; i < half; ++i) {
+        freqs[i] = (float)std::exp(-std::log(max_period) * i / half);
+    }
+    for (int i = 0; i < N; ++i) {
+        for (int j = 0; j < half; ++j) {
+            float arg = timesteps[i] * freqs[j] * scale;
+            if (flip_sin_to_cos) {
+                embedding[i * dim + j]        = std::cos(arg);
+                embedding[i * dim + j + half] = std::sin(arg);
+            } else {
+                embedding[i * dim + j]        = std::sin(arg);
+                embedding[i * dim + j + half] = std::cos(arg);
+            }
+        }
+    }
+    return embedding;
+}
+
+void set_timestep_embedding(std::vector<float> timesteps,
+                            ggml_tensor* embedding,
+                            int dim,
+                            int max_period) {
+    std::vector<float> embedding_vec = timestep_embedding(timesteps, dim, max_period);
+    memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding));
+}
+
+void set_timestep_embedding(std::vector<float> timesteps,
+                            sd::Tensor<float>* embedding,
+                            int dim,
+                            int max_period) {
+    GGML_ASSERT(embedding != nullptr);
+    std::vector<float> embedding_vec = timestep_embedding(timesteps, dim, max_period);
+    if (embedding->numel() != static_cast<int64_t>(embedding_vec.size())) {
+        embedding->resize({dim, static_cast<int64_t>(timesteps.size())});
+    }
+    std::copy(embedding_vec.begin(), embedding_vec.end(), embedding->values().begin());
+}
+
+ggml_tensor* new_timestep_embedding(ggml_context* ctx,
+                                    std::vector<float> timesteps,
+                                    int dim,
+                                    int max_period) {
+    // timesteps: [N,]
+    // embedding: [N, dim]
+    std::vector<float> embedding_vec = timestep_embedding(timesteps, dim, max_period);
+    ggml_tensor* embedding           = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, timesteps.size());
+    if (embedding->data != nullptr) {
+        memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding));
+    } else {
+        ggml_backend_tensor_set(embedding, embedding_vec.data(), 0, ggml_nbytes(embedding));
+    }
+    return embedding;
+}
+
+size_t ggml_tensor_num(ggml_context* ctx) {
+    size_t num = 0;
+    for (ggml_tensor* t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+        num++;
+    }
+    return num;
+}

src/core/ggml_tensor_utils.h

@@ -0,0 +1,210 @@
+#ifndef __SD_CORE_GGML_TENSOR_UTILS_H__
+#define __SD_CORE_GGML_TENSOR_UTILS_H__
+
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <functional>
+#include <memory>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+#include "core/tensor.hpp"
+#include "core/util.h"
+#include "ggml-backend.h"
+#include "ggml.h"
+#include "stable-diffusion.h"
+
+class RNG;
+
+__STATIC_INLINE__ int align_up_offset(int n, int multiple) {
+    return (multiple - n % multiple) % multiple;
+}
+
+__STATIC_INLINE__ int align_up(int n, int multiple) {
+    return n + align_up_offset(n, multiple);
+}
+
+void ggml_ext_im_set_randn_f32(ggml_tensor* tensor, std::shared_ptr<RNG> rng);
+
+__STATIC_INLINE__ void ggml_ext_tensor_set_f32(ggml_tensor* tensor, float value, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    GGML_ASSERT(tensor->nb[0] == sizeof(float));
+    *(float*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]) = value;
+}
+
+__STATIC_INLINE__ float ggml_ext_tensor_get_f32(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    if (tensor->buffer != nullptr) {
+        float value;
+        ggml_backend_tensor_get(tensor, &value, i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0], sizeof(float));
+        return value;
+    }
+    GGML_ASSERT(tensor->nb[0] == sizeof(float));
+    return *(float*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]);
+}
+
+__STATIC_INLINE__ int ggml_ext_tensor_get_i32(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    if (tensor->buffer != nullptr) {
+        int value;
+        ggml_backend_tensor_get(tensor, &value, i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0], sizeof(int));
+        return value;
+    }
+    GGML_ASSERT(tensor->nb[0] == sizeof(int));
+    return *(int*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]);
+}
+
+__STATIC_INLINE__ ggml_fp16_t ggml_ext_tensor_get_f16(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    GGML_ASSERT(tensor->nb[0] == sizeof(ggml_fp16_t));
+    return *(ggml_fp16_t*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]);
+}
+
+__STATIC_INLINE__ float sd_image_get_f32(sd_image_t image, int64_t iw, int64_t ih, int64_t ic, bool scale = true) {
+    float value = *(image.data + ih * image.width * image.channel + iw * image.channel + ic);
+    if (scale) {
+        value /= 255.f;
+    }
+    return value;
+}
+
+void print_ggml_tensor(ggml_tensor* tensor, bool shape_only = false, const char* mark = "");
+
+template <typename T>
+__STATIC_INLINE__ void print_sd_tensor(const sd::Tensor<T>& tensor, bool shape_only = false, const char* mark = "") {
+    printf("%s: shape(", mark);
+    for (size_t i = 0; i < static_cast<size_t>(tensor.dim()); ++i) {
+        printf("%s%lld", i == 0 ? "" : ", ", static_cast<long long>(tensor.shape()[i]));
+    }
+    printf(")\n");
+    fflush(stdout);
+    if (shape_only) {
+        return;
+    }
+    if (tensor.empty()) {
+        return;
+    }
+    int range                  = 3;
+    std::vector<int64_t> shape = tensor.shape();
+    while (shape.size() < 4) {
+        shape.push_back(1);
+    }
+    for (int64_t i3 = 0; i3 < shape[3]; i3++) {
+        if (i3 >= range && i3 + range < shape[3]) {
+            continue;
+        }
+        for (int64_t i2 = 0; i2 < shape[2]; i2++) {
+            if (i2 >= range && i2 + range < shape[2]) {
+                continue;
+            }
+            for (int64_t i1 = 0; i1 < shape[1]; i1++) {
+                if (i1 >= range && i1 + range < shape[1]) {
+                    continue;
+                }
+                for (int64_t i0 = 0; i0 < shape[0]; i0++) {
+                    if (i0 >= range && i0 + range < shape[0]) {
+                        continue;
+                    }
+                    size_t offset = static_cast<size_t>(i0 + shape[0] * (i1 + shape[1] * (i2 + shape[2] * i3)));
+                    printf("  [%lld, %lld, %lld, %lld] = ", static_cast<long long>(i3), static_cast<long long>(i2), static_cast<long long>(i1), static_cast<long long>(i0));
+                    if constexpr (std::is_same_v<T, float>) {
+                        printf("%f\n", tensor[static_cast<int64_t>(offset)]);
+                    } else if constexpr (std::is_same_v<T, ggml_fp16_t>) {
+                        printf("%f\n", ggml_fp16_to_fp32(tensor[static_cast<int64_t>(offset)]));
+                    } else if constexpr (std::is_same_v<T, int32_t>) {
+                        printf("%d\n", tensor[static_cast<int64_t>(offset)]);
+                    } else if constexpr (std::is_same_v<T, int64_t>) {
+                        printf("%lld\n", static_cast<long long>(tensor[static_cast<int64_t>(offset)]));
+                    }
+                    fflush(stdout);
+                }
+            }
+        }
+    }
+}
+
+void ggml_ext_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t, int64_t, int64_t, int64_t)>& fn);
+
+void ggml_ext_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t)>& fn);
+
+void ggml_ext_tensor_diff(
+    ggml_tensor* a,
+    ggml_tensor* b,
+    float gap = 0.1f);
+
+ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path);
+
+__STATIC_INLINE__ float sigmoid(float x) {
+    return 1 / (1.0f + expf(-x));
+}
+
+// SPECIAL OPERATIONS WITH TENSORS
+
+uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t* image_data = nullptr);
+
+uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, int idx, bool video = false);
+
+void sd_image_to_ggml_tensor(sd_image_t image,
+                             ggml_tensor* tensor,
+                             bool scale = true);
+
+void ggml_ext_tensor_apply_mask(ggml_tensor* image_data,
+                                ggml_tensor* mask,
+                                ggml_tensor* output,
+                                float masked_value = 0.5f);
+
+float ggml_ext_tensor_mean(ggml_tensor* src);
+
+// a = a+b
+void ggml_ext_tensor_add_inplace(ggml_tensor* a, ggml_tensor* b);
+
+void ggml_ext_tensor_scale_inplace(ggml_tensor* src, float scale);
+
+void ggml_ext_tensor_clamp_inplace(ggml_tensor* src, float min, float max);
+
+ggml_tensor* ggml_ext_tensor_concat(ggml_context* ctx,
+                                    ggml_tensor* a,
+                                    ggml_tensor* b,
+                                    int dim);
+
+// convert values from [0, 1] to [-1, 1]
+void scale_to_minus1_1(ggml_tensor* src);
+
+// convert values from [-1, 1] to [0, 1]
+void scale_to_0_1(ggml_tensor* src);
+
+ggml_tensor* vector_to_ggml_tensor(ggml_context* ctx,
+                                   const std::vector<float>& vec);
+
+ggml_tensor* vector_to_ggml_tensor_i32(ggml_context* ctx,
+                                       const std::vector<int>& vec);
+
+std::vector<float> arange(float start, float end, float step = 1.f);
+
+// Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151
+std::vector<float> timestep_embedding(std::vector<float> timesteps,
+                                      int dim,
+                                      int max_period       = 10000,
+                                      bool flip_sin_to_cos = true,
+                                      float scale          = 1.f);
+
+void set_timestep_embedding(std::vector<float> timesteps,
+                            ggml_tensor* embedding,
+                            int dim,
+                            int max_period = 10000);
+
+void set_timestep_embedding(std::vector<float> timesteps,
+                            sd::Tensor<float>* embedding,
+                            int dim,
+                            int max_period = 10000);
+
+ggml_tensor* new_timestep_embedding(ggml_context* ctx,
+                                    std::vector<float> timesteps,
+                                    int dim,
+                                    int max_period = 10000);
+
+size_t ggml_tensor_num(ggml_context* ctx);
+
+#endif  // __SD_CORE_GGML_TENSOR_UTILS_H__

src/core/layer_split_partition.cpp

@@ -251,13 +251,13 @@ namespace sd {
                          assignment.tensors_by_backend[i].size(),
                          assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
             } else {
-                LOG_DEBUG("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
-                          desc,
-                          layer_split_backend_device_display_name(split_backends[i]).c_str(),
-                          first_segment,
-                          last_segment,
-                          assignment.tensors_by_backend[i].size(),
-                          assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
+                LOG_VERBOSE("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
+                            desc,
+                            layer_split_backend_device_display_name(split_backends[i]).c_str(),
+                            first_segment,
+                            last_segment,
+                            assignment.tensors_by_backend[i].size(),
+                            assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
             }
         }
     }

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

@@ -15,6 +15,7 @@
 #include <thread>
 #include <unordered_set>
 #include <vector>
+#include "core/ggml_tensor_utils.h"
 #include "runtime/preprocessing.hpp"
 
 #if defined(__APPLE__) && defined(__MACH__)
@@ -618,6 +619,25 @@ void log_printf(sd_log_level_t level, const char* file, int line, const char* fo
     va_end(args);
 }
 
+void sd_ggml_log_callback(ggml_log_level level, const char* text, void*) {
+    switch (level) {
+        case GGML_LOG_LEVEL_DEBUG:
+            LOG_VERBOSE(text);
+            break;
+        case GGML_LOG_LEVEL_INFO:
+            LOG_INFO(text);
+            break;
+        case GGML_LOG_LEVEL_WARN:
+            LOG_WARN(text);
+            break;
+        case GGML_LOG_LEVEL_ERROR:
+            LOG_ERROR(text);
+            break;
+        default:
+            LOG_VERBOSE(text);
+    }
+}
+
 void sd_set_log_callback(sd_log_cb_t cb, void* data) {
     sd_log_cb      = cb;
     sd_log_cb_data = data;
@@ -801,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) {
@@ -810,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();
@@ -840,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/core/util.h

@@ -11,6 +11,14 @@
 #include "ggml-backend.h"
 #include "stable-diffusion.h"
 
+#ifndef __STATIC_INLINE__
+#define __STATIC_INLINE__ static inline
+#endif
+
+#ifndef SD_UNUSED
+#define SD_UNUSED(x) (void)(x)
+#endif
+
 #define SAFE_STR(s) ((s) ? (s) : "")
 #define BOOL_STR(b) ((b) ? "true" : "false")
 
@@ -79,6 +87,7 @@ void pretty_progress(int step, int steps, float time);
 void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds);
 
 void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...);
+void sd_ggml_log_callback(ggml_log_level level, const char* text, void*);
 
 ggml_type sd_type_to_ggml_type(sd_type_t sdtype);
 
@@ -105,6 +114,7 @@ void* sd_get_backend_eval_callback_data();
 bool sd_backend_is(ggml_backend_t backend, const std::string& name);
 
 #define LOG_DEBUG(format, ...) log_printf(SD_LOG_DEBUG, __FILE__, __LINE__, format, ##__VA_ARGS__)
+#define LOG_VERBOSE(format, ...) log_printf(SD_LOG_VERBOSE, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_INFO(format, ...) log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_WARN(format, ...) log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_ERROR(format, ...) log_printf(SD_LOG_ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__)

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

@@ -1,3 +1,4 @@
+#include <cinttypes>
 #include "extensions/generation_extension.h"
 
 #include <algorithm>
@@ -17,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();
     }
@@ -48,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);
@@ -68,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,
@@ -77,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 {
@@ -108,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 {
@@ -118,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;
         }
 
@@ -127,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");
         }
@@ -227,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");
@@ -267,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

@@ -1,8 +1,10 @@
 #ifndef __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
 #define __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "model/common/block.hpp"
+#include "model/common/ggml_block.hpp"
 #include "model_loader.h"
 
 namespace IPAdapter {
@@ -93,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);
 
@@ -200,7 +202,7 @@ namespace IPAdapter {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(image_embeds);
             };
-            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
+            return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
         }
     };
 

src/model/adapter/lora.hpp

@@ -2,7 +2,13 @@
 #define __SD_MODEL_ADAPTER_LORA_HPP__
 
 #include <mutex>
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model.h"
+#include "model/adapter/lora_ops.h"
 #include "model_loader.h"
 #include "model_manager.h"
 
@@ -17,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));
         }
     }
 
@@ -43,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_DEBUG("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();
@@ -235,14 +217,14 @@ 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_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
+                    // LOG_VERBOSE("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
                     applied_lora_tensors.insert(alpha_name);
                 }
             }
@@ -389,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);
             }
@@ -502,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);
             }
@@ -663,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;
                 }
 
@@ -790,15 +772,15 @@ 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_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
+                    // LOG_VERBOSE("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
                 }
             }
             scale_value *= multiplier;
@@ -943,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,
@@ -957,17 +939,18 @@ struct LoraModel : public GGMLRunner {
             }
             return true;
         };
-        auto result = GGMLRunner::compute<float>(get_graph, n_threads, false, true, read_outputs);
+        auto result = GGMLRunner::compute(get_graph, n_threads, false, true, read_outputs);
         if (!result.has_value()) {
             LOG_ERROR("LoRA graph execution failed");
         }
         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/lora_ops.cpp

@@ -0,0 +1,207 @@
+#include "model/adapter/lora_ops.h"
+
+#include <cmath>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+
+ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
+                                 ggml_tensor* lora_down,
+                                 ggml_tensor* lora_up,
+                                 ggml_tensor* lora_mid) {
+    ggml_tensor* updown;
+    // flat lora tensors to multiply it
+    int64_t lora_up_rows  = lora_up->ne[ggml_n_dims(lora_up) - 1];
+    lora_up               = ggml_reshape_2d(ctx, lora_up, ggml_nelements(lora_up) / lora_up_rows, lora_up_rows);
+    auto lora_down_n_dims = ggml_n_dims(lora_down);
+    // assume n_dims should always be a multiple of 2 (otherwise rank 1 doesn't work)
+    lora_down_n_dims       = (lora_down_n_dims + lora_down_n_dims % 2);
+    int64_t lora_down_rows = lora_down->ne[lora_down_n_dims - 1];
+    lora_down              = ggml_reshape_2d(ctx, lora_down, ggml_nelements(lora_down) / lora_down_rows, lora_down_rows);
+
+    // ggml_mul_mat requires tensor b transposed
+    lora_down = ggml_cont(ctx, ggml_transpose(ctx, lora_down));
+    if (lora_mid == nullptr) {
+        updown = ggml_mul_mat(ctx, lora_up, lora_down);
+        updown = ggml_cont(ctx, ggml_transpose(ctx, updown));
+    } else {
+        // undoing tucker decomposition for conv layers.
+        // lora_mid  has shape (3,    3,   Rank, Rank)
+        // lora_down has shape (Rank, In,  1,    1)
+        // lora_up   has shape (Rank, Out, 1,    1)
+        // conv layer shape is (3,    3,   Out,  In)
+        updown = ggml_ext_mul_n_mode(ctx, ggml_ext_mul_n_mode(ctx, lora_mid, lora_down, 3), lora_up, 2);
+        updown = ggml_cont(ctx, updown);
+    }
+    return updown;
+}
+
+ggml_tensor* ggml_ext_lokr_forward(
+    ggml_context* ctx,
+    ggml_backend_t backend,
+    ggml_tensor* h,    // Input: [q, batch] or [W, H, q, batch]
+    ggml_tensor* w1,   // Outer C (Full rank)
+    ggml_tensor* w1a,  // Outer A (Low rank part 1)
+    ggml_tensor* w1b,  // Outer B (Low rank part 2)
+    ggml_tensor* w2,   // Inner BA (Full rank)
+    ggml_tensor* w2a,  // Inner A (Low rank part 1)
+    ggml_tensor* w2b,  // Inner B (Low rank part 2)
+    bool is_conv,
+    WeightAdapter::ForwardParams::conv2d_params_t conv_params,
+    float scale) {
+    GGML_ASSERT((w1 != nullptr || (w1a != nullptr && w1b != nullptr)));
+    GGML_ASSERT((w2 != nullptr || (w2a != nullptr && w2b != nullptr)));
+
+    int uq = (w1 != nullptr) ? (int)w1->ne[0] : (int)w1a->ne[0];
+    int up = (w1 != nullptr) ? (int)w1->ne[1] : (int)w1b->ne[1];
+
+    int q_actual = is_conv ? (int)h->ne[2] : (int)h->ne[0];
+    int vq       = q_actual / uq;
+
+    int vp = (w2 != nullptr) ? (is_conv ? (int)w2->ne[3] : (int)w2->ne[1])
+                             : (int)w2a->ne[1];
+    GGML_ASSERT(q_actual == (uq * vq) && "Input dimension mismatch for LoKR split");
+
+    ggml_tensor* hb;
+
+    if (!is_conv) {
+        int batch          = (int)h->ne[1];
+        int merge_batch_uq = batch;
+        int merge_batch_vp = batch;
+
+        if (sd_backend_is(backend, "Vulkan")) {
+            if (batch > 1) {
+                // no access to backend here, worst case is slightly worse perfs for other backends when built alongside Vulkan backend
+                int max_batch    = 65535;
+                int max_batch_uq = max_batch / uq;
+                merge_batch_uq   = 1;
+                for (int i = max_batch_uq; i > 0; i--) {
+                    if (batch % i == 0) {
+                        merge_batch_uq = i;
+                        break;
+                    }
+                }
+
+                int max_batch_vp = max_batch / vp;
+                merge_batch_vp   = 1;
+                for (int i = max_batch_vp; i > 0; i--) {
+                    if (batch % i == 0) {
+                        merge_batch_vp = i;
+                        break;
+                    }
+                }
+            }
+        }
+
+        ggml_tensor* h_split = ggml_reshape_3d(ctx, h, vq, uq * merge_batch_uq, batch / merge_batch_uq);
+        if (w2 != nullptr) {
+            hb = ggml_mul_mat(ctx, w2, h_split);
+        } else {
+            hb = ggml_mul_mat(ctx, w2b, ggml_mul_mat(ctx, w2a, h_split));
+        }
+
+        if (batch > 1) {
+            hb = ggml_reshape_3d(ctx, hb, vp, uq, batch);
+        }
+        ggml_tensor* hb_t = ggml_cont(ctx, ggml_transpose(ctx, hb));
+        hb_t              = ggml_reshape_3d(ctx, hb_t, uq, vp * merge_batch_vp, batch / merge_batch_vp);
+
+        ggml_tensor* hc_t;
+        if (w1 != nullptr) {
+            hc_t = ggml_mul_mat(ctx, w1, hb_t);
+        } else {
+            hc_t = ggml_mul_mat(ctx, w1b, ggml_mul_mat(ctx, w1a, hb_t));
+        }
+
+        if (batch > 1) {
+            hc_t = ggml_reshape_3d(ctx, hc_t, up, vp, batch);
+        }
+
+        ggml_tensor* hc  = ggml_transpose(ctx, hc_t);
+        ggml_tensor* out = ggml_reshape_2d(ctx, ggml_cont(ctx, hc), up * vp, batch);
+        return ggml_ext_scale(ctx, out, scale);
+    } else {
+        int batch = (int)h->ne[3];
+        // 1. Reshape input: [W, H, vq*uq, batch] -> [W, H, vq, uq * batch]
+        ggml_tensor* h_split = ggml_reshape_4d(ctx, h, h->ne[0], h->ne[1], vq, uq * batch);
+
+        if (w2 != nullptr) {
+            hb = ggml_ext_conv_2d(ctx, h_split, w2, nullptr,
+                                  conv_params.s0,
+                                  conv_params.s1,
+                                  conv_params.p0,
+                                  conv_params.p1,
+                                  conv_params.d0,
+                                  conv_params.d1,
+                                  conv_params.direct,
+                                  conv_params.circular_x,
+                                  conv_params.circular_y,
+                                  conv_params.scale);
+        } else {
+            // swap a and b order for conv lora
+            ggml_tensor* a = w2b;
+            ggml_tensor* b = w2a;
+
+            // unpack conv2d weights if needed
+            if (ggml_n_dims(a) < 4) {
+                int k = (int)sqrt(a->ne[0] / h_split->ne[2]);
+                GGML_ASSERT(k * k * h_split->ne[2] == a->ne[0]);
+                a = ggml_reshape_4d(ctx, a, k, k, a->ne[0] / (k * k), a->ne[1]);
+            } else if (a->ne[2] != h_split->ne[2]) {
+                int k = (int)sqrt(a->ne[2] / h_split->ne[2]);
+                GGML_ASSERT(k * k * h_split->ne[2] == a->ne[2]);
+                a = ggml_reshape_4d(ctx, a, a->ne[0] * k, a->ne[1] * k, a->ne[2] / (k * k), a->ne[3]);
+            }
+            ggml_tensor* ha = ggml_ext_conv_2d(ctx, h_split, a, nullptr,
+                                               conv_params.s0,
+                                               conv_params.s1,
+                                               conv_params.p0,
+                                               conv_params.p1,
+                                               conv_params.d0,
+                                               conv_params.d1,
+                                               conv_params.direct,
+                                               conv_params.circular_x,
+                                               conv_params.circular_y,
+                                               conv_params.scale);
+
+            // not supporting lora_mid here
+            hb = ggml_ext_conv_2d(ctx,
+                                  ha,
+                                  b,
+                                  nullptr,
+                                  1,
+                                  1,
+                                  0,
+                                  0,
+                                  1,
+                                  1,
+                                  conv_params.direct,
+                                  conv_params.circular_x,
+                                  conv_params.circular_y,
+                                  conv_params.scale);
+        }
+
+        // Current hb shape: [W_out, H_out, vp, uq * batch]
+        int w_out = (int)hb->ne[0];
+        int h_out = (int)hb->ne[1];
+
+        // ggml_tensor* hb_cat = ggml_reshape_4d(ctx, hb, w_out , h_out , vp * uq, batch);
+        // [W_out, H_out, vp * uq,  batch]
+        // Now left to compute (W1 kr Id) * hb_cat == (W1 kr W2) cv h
+
+        // merge the uq groups of size vp*w_out*h_out
+        ggml_tensor* hb_merged = ggml_reshape_2d(ctx, hb, w_out * h_out * vp, uq * batch);
+        ggml_tensor* hc_t;
+        ggml_tensor* hb_merged_t = ggml_cont(ctx, ggml_transpose(ctx, hb_merged));
+        if (w1 != nullptr) {
+            // Would be great to be able to transpose w1 instead to avoid transposing both hb and hc
+            hc_t = ggml_mul_mat(ctx, w1, hb_merged_t);
+        } else {
+            hc_t = ggml_mul_mat(ctx, w1b, ggml_mul_mat(ctx, w1a, hb_merged_t));
+        }
+        ggml_tensor* hc = ggml_transpose(ctx, hc_t);
+        // ungroup
+        ggml_tensor* out = ggml_reshape_4d(ctx, ggml_cont(ctx, hc), w_out, h_out, up * vp, batch);
+        return ggml_ext_scale(ctx, out, scale);
+    }
+}

src/model/adapter/lora_ops.h

@@ -0,0 +1,25 @@
+#ifndef __SD_MODEL_ADAPTER_LORA_OPS_H__
+#define __SD_MODEL_ADAPTER_LORA_OPS_H__
+
+#include "core/ggml_runner.h"
+
+ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
+                                 ggml_tensor* lora_down,
+                                 ggml_tensor* lora_up,
+                                 ggml_tensor* lora_mid = nullptr);
+
+ggml_tensor* ggml_ext_lokr_forward(
+    ggml_context* ctx,
+    ggml_backend_t backend,
+    ggml_tensor* h,    // Input: [q, batch] or [W, H, q, batch]
+    ggml_tensor* w1,   // Outer C (Full rank)
+    ggml_tensor* w1a,  // Outer A (Low rank part 1)
+    ggml_tensor* w1b,  // Outer B (Low rank part 2)
+    ggml_tensor* w2,   // Inner BA (Full rank)
+    ggml_tensor* w2a,  // Inner A (Low rank part 1)
+    ggml_tensor* w2b,  // Inner B (Low rank part 2)
+    bool is_conv,
+    WeightAdapter::ForwardParams::conv2d_params_t conv_params,
+    float scale);
+
+#endif  // __SD_MODEL_ADAPTER_LORA_OPS_H__

src/model/adapter/pmid.hpp

@@ -1,12 +1,13 @@
 #ifndef __SD_MODEL_ADAPTER_PMID_HPP__
 #define __SD_MODEL_ADAPTER_PMID_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#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
@@ -558,97 +559,7 @@ public:
             return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds);
         };
 
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
-    }
-};
-
-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_DEBUG("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;
+        return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
     }
 };
 

src/model/adapter/pulid.hpp

@@ -1,8 +1,10 @@
 #ifndef __PULID_HPP__
 #define __PULID_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "model/common/block.hpp"
+#include "model/common/ggml_block.hpp"
 
 class PuLIDPerceiverAttentionCA : public GGMLBlock {
 public:
@@ -61,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

@@ -1,9 +1,11 @@
 #ifndef __SD_MODEL_COMMON_BLOCK_HPP__
 #define __SD_MODEL_COMMON_BLOCK_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "core/util.h"
 #include "ggml-backend.h"
+#include "model/common/ggml_block.hpp"
 
 class DownSampleBlock : public GGMLBlock {
 protected:
@@ -340,7 +342,7 @@ public:
           enable_ip(enable_ip) {
         int64_t inner_dim = d_head * n_head;
         if (context_dim == 320 && d_head == 320) {
-            // LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09");
+            // LOG_VERBOSE("CrossAttention: temp set dim to 1024 for sdxs_09");
             xtra_dim    = true;
             context_dim = 1024;
         }
@@ -370,7 +372,7 @@ public:
 
         auto q = to_q->forward(ctx, x);  // [N, n_token, inner_dim]
         if (xtra_dim) {
-            // LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09");
+            // LOG_VERBOSE("CrossAttention: temp set dim to 1024 for sdxs_09");
             context->ne[0] = 1024;  // patch dim
         }
         auto k = to_k->forward(ctx, context);  // [N, n_context, inner_dim]
@@ -378,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

@@ -0,0 +1,936 @@
+#ifndef __SD_MODEL_COMMON_GGML_BLOCK_HPP__
+#define __SD_MODEL_COMMON_GGML_BLOCK_HPP__
+
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <tuple>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "model.h"
+
+class GGMLBlock {
+protected:
+    typedef std::unordered_map<std::string, ggml_tensor*> ParameterMap;
+    typedef std::unordered_map<std::string, std::shared_ptr<GGMLBlock>> GGMLBlockMap;
+    GGMLBlockMap blocks;
+    ParameterMap params;
+
+    ggml_type get_type(const std::string& name, const String2TensorStorage& tensor_storage_map, ggml_type default_type) {
+        ggml_type wtype = default_type;
+        auto iter       = tensor_storage_map.find(name);
+        if (iter != tensor_storage_map.end()) {
+            const TensorStorage& tensor_storage = iter->second;
+            if (tensor_storage.expected_type != GGML_TYPE_COUNT) {
+                wtype = tensor_storage.expected_type;
+            } else {
+                wtype = tensor_storage.type;
+            }
+        }
+        return wtype;
+    }
+
+    void init_blocks(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+            block->init(ctx, tensor_storage_map, prefix + pair.first);
+        }
+    }
+
+    virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {}
+
+    virtual enum ggml_op param_usage_op(const std::string& name) const {
+        (void)name;
+        return GGML_OP_NONE;
+    }
+
+public:
+    void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") {
+        if (prefix.size() > 0) {
+            prefix = prefix + ".";
+        }
+        init_params(ctx, tensor_storage_map, prefix);
+        init_blocks(ctx, tensor_storage_map, prefix);
+    }
+
+    size_t get_params_num() {
+        size_t num_tensors = params.size();
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+
+            num_tensors += block->get_params_num();
+        }
+        return num_tensors;
+    };
+
+    size_t get_params_mem_size() {
+        size_t mem_size = 0;
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+
+            mem_size += block->get_params_mem_size();
+        }
+
+        for (auto& pair : params) {
+            mem_size += ggml_nbytes(pair.second);
+        }
+
+        return mem_size;
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, std::string prefix = "") {
+        if (prefix.size() > 0) {
+            prefix = prefix + ".";
+        }
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+            block->get_param_tensors(tensors, prefix + pair.first);
+        }
+
+        for (auto& pair : params) {
+            ggml_tensor* param           = pair.second;
+            tensors[prefix + pair.first] = pair.second;
+            ggml_set_name(param, (prefix + pair.first).c_str());
+        }
+    }
+
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
+        for (auto& pair : blocks) {
+            pair.second->get_param_tensor_ops(tensor_ops);
+        }
+        for (auto& pair : params) {
+            enum ggml_op op = param_usage_op(pair.first);
+            if (op != GGML_OP_NONE) {
+                tensor_ops[pair.second] = op;
+            }
+        }
+    }
+
+    virtual std::string get_desc() {
+        return "GGMLBlock";
+    }
+
+    void get_all_blocks(std::vector<GGMLBlock*>& result) {
+        result.push_back(this);
+        for (auto& block_iter : blocks) {
+            if (block_iter.second) {
+                block_iter.second->get_all_blocks(result);
+            }
+        }
+    }
+};
+
+class UnaryBlock : public GGMLBlock {
+public:
+    virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) = 0;
+};
+
+class Identity : public UnaryBlock {
+public:
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        return x;
+    }
+};
+
+class Linear : public UnaryBlock {
+protected:
+    int64_t in_features;
+    int64_t out_features;
+    bool bias;
+    bool force_f32;
+    bool force_prec_f32;
+    bool has_weight_scale       = false;
+    bool int8_convrot           = false;
+    int int8_convrot_group_size = 0;
+    float scale;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        this->prefix            = prefix;
+        has_weight_scale        = false;
+        int8_convrot            = false;
+        int8_convrot_group_size = 0;
+        enum ggml_type wtype    = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
+        if (in_features % ggml_blck_size(wtype) != 0 || force_f32) {
+            wtype = GGML_TYPE_F32;
+        }
+        params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features);
+        if (bias) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_features);
+        }
+        auto weight_storage           = tensor_storage_map.find(prefix + "weight");
+        const bool is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise;
+        auto weight_scale_storage     = tensor_storage_map.find(prefix + "weight_scale");
+        if (weight_scale_storage != tensor_storage_map.end()) {
+            const int64_t scale_nelements = weight_scale_storage->second.nelements();
+            GGML_ASSERT(scale_nelements == 1 || scale_nelements == out_features);
+            params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements);
+            has_weight_scale       = true;
+        }
+        if (is_int8_tensorwise) {
+            GGML_ASSERT(wtype == GGML_TYPE_I8);
+            GGML_ASSERT(has_weight_scale);
+            int8_convrot            = weight_storage->second.int8_convrot;
+            int8_convrot_group_size = weight_storage->second.int8_convrot_group_size;
+        }
+    }
+
+public:
+    Linear(int64_t in_features,
+           int64_t out_features,
+           bool bias           = true,
+           bool force_f32      = false,
+           bool force_prec_f32 = false,
+           float scale         = 1.f)
+        : in_features(in_features),
+          out_features(out_features),
+          bias(bias),
+          force_f32(force_f32),
+          force_prec_f32(force_prec_f32),
+          scale(scale) {}
+
+    void set_scale(float scale_) {
+        scale = scale_;
+    }
+
+    void set_force_prec_f32(bool force_prec_f32_) {
+        force_prec_f32 = force_prec_f32_;
+    }
+
+    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;
+            if (ctx->backend != nullptr) {
+                ggml_tensor* fp8_matmul = ggml_mul_mat(ctx->ggml_ctx, w, x);
+                if (force_prec_f32) {
+                    ggml_mul_mat_set_prec(fp8_matmul, GGML_PREC_F32);
+                }
+                supports_fp8_matmul = ggml_backend_supports_op(ctx->backend, fp8_matmul);
+            }
+            if (!supports_fp8_matmul) {
+                w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16);
+            }
+        }
+        ggml_tensor* b = nullptr;
+        if (bias) {
+            b = params["bias"];
+        }
+        ggml_tensor* linear_bias = has_weight_scale ? nullptr : b;
+        ggml_tensor* out         = nullptr;
+        if (w->type == GGML_TYPE_I8) {
+            if (x->type != GGML_TYPE_F32) {
+                x = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, x);
+            }
+            if (!ggml_is_contiguous(x)) {
+                x = ggml_cont(ctx->ggml_ctx, x);
+            }
+            ggml_tensor* lora_input = x;
+            if (ctx->weight_adapter && b != nullptr) {
+                b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+            }
+            if (int8_convrot && scale == 1.f) {
+                const auto cache_key = std::make_pair(x, int8_convrot_group_size);
+                auto cached          = ctx->int8_convrot_cache.find(cache_key);
+                if (cached == ctx->int8_convrot_cache.end()) {
+                    x = ggml_quantize_i8_convrot(ctx->ggml_ctx, x, int8_convrot_group_size);
+                    ctx->int8_convrot_cache.emplace(cache_key, x);
+                } else {
+                    x = cached->second;
+                }
+            }
+            out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
+                                                x,
+                                                w,
+                                                weight_scale,
+                                                b,
+                                                int8_convrot ? int8_convrot_group_size : 0,
+                                                scale);
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
+                forward_params.linear.force_prec_f32 = force_prec_f32;
+                forward_params.linear.scale          = scale;
+                out                                  = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx,
+                                                                                               ctx->backend,
+                                                                                               lora_input,
+                                                                                               w,
+                                                                                               out,
+                                                                                               prefix,
+                                                                                               forward_params);
+            }
+            return out;
+        }
+        if (has_weight_scale) {
+            out = ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr, force_prec_f32, scale);
+            out = ggml_mul(ctx->ggml_ctx, out, weight_scale);
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
+                forward_params.linear.force_prec_f32 = force_prec_f32;
+                forward_params.linear.scale          = scale;
+                out                                  = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx,
+                                                                                               ctx->backend,
+                                                                                               x,
+                                                                                               w,
+                                                                                               out,
+                                                                                               prefix,
+                                                                                               forward_params);
+                if (b != nullptr) {
+                    b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                }
+            }
+            if (b != nullptr) {
+                out = ggml_add_inplace(ctx->ggml_ctx, out, b);
+            }
+            return out;
+        }
+        if (ctx->weight_adapter) {
+            WeightAdapter::ForwardParams forward_params;
+            forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
+            forward_params.linear.force_prec_f32 = force_prec_f32;
+            forward_params.linear.scale          = scale;
+            out                                  = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, linear_bias, prefix, forward_params);
+        } else {
+            out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale);
+        }
+        return out;
+    }
+};
+
+__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:
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
+            return true;
+        default:
+            return false;
+    }
+}
+
+class Embedding : public UnaryBlock {
+protected:
+    int64_t embedding_dim;
+    int64_t num_embeddings;
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
+        enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
+        if (!support_get_rows(wtype)) {
+            wtype = GGML_TYPE_F32;
+        }
+        params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
+    }
+
+    enum ggml_op param_usage_op(const std::string& name) const override {
+        return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
+    }
+
+public:
+    Embedding(int64_t num_embeddings, int64_t embedding_dim)
+        : embedding_dim(embedding_dim),
+          num_embeddings(num_embeddings) {
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* input_ids) override {
+        // input_ids: [N, n_token]
+        auto weight = params["weight"];
+
+        // There are issues with ggml batch inference, so we are expanding it here first.
+        // TODO: fix ggml batch inference
+        int64_t n = input_ids->ne[1];
+        input_ids = ggml_reshape_1d(ctx->ggml_ctx, input_ids, input_ids->ne[0] * input_ids->ne[1]);
+
+        input_ids      = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]);
+        auto embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids);
+        embedding      = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n);
+
+        // [N, n_token, embedding_dim]
+        return embedding;
+    }
+};
+
+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;
+    int64_t out_channels;
+    std::pair<int, int> kernel_size;
+    std::pair<int, int> stride;
+    std::pair<int, int> padding;
+    std::pair<int, int> dilation;
+    bool bias;
+    float scale = 1.f;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F16;
+        params["weight"]     = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels, out_channels);
+        if (bias) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_channels);
+        }
+    }
+
+public:
+    Conv2d(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},
+           std::pair<int, int> dilation = {1, 1},
+           bool bias                    = true)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias) {}
+
+    void set_scale(float scale_value) {
+        scale = scale_value;
+    }
+
+    std::string get_desc() override {
+        return "Conv2d";
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
+        ggml_tensor* b = nullptr;
+        if (bias) {
+            b = params["bias"];
+        }
+        if (ctx->weight_adapter) {
+            WeightAdapter::ForwardParams forward_params;
+            forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+            forward_params.conv2d.s0         = stride.second;
+            forward_params.conv2d.s1         = stride.first;
+            forward_params.conv2d.p0         = padding.second;
+            forward_params.conv2d.p1         = padding.first;
+            forward_params.conv2d.d0         = dilation.second;
+            forward_params.conv2d.d1         = dilation.first;
+            forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
+            forward_params.conv2d.circular_x = ctx->circular_x_enabled;
+            forward_params.conv2d.circular_y = ctx->circular_y_enabled;
+            forward_params.conv2d.scale      = scale;
+            return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params);
+        }
+        return ggml_ext_conv_2d(ctx->ggml_ctx,
+                                x,
+                                w,
+                                b,
+                                stride.second,
+                                stride.first,
+                                padding.second,
+                                padding.first,
+                                dilation.second,
+                                dilation.first,
+                                ctx->conv2d_direct_enabled,
+                                ctx->circular_x_enabled,
+                                ctx->circular_y_enabled,
+                                scale);
+    }
+};
+
+class Conv2d_grouped : public UnaryBlock {
+protected:
+    int64_t in_channels;
+    int64_t out_channels;
+    int groups;
+    std::pair<int, int> kernel_size;
+    std::pair<int, int> stride;
+    std::pair<int, int> padding;
+    std::pair<int, int> dilation;
+    bool bias;
+    float scale = 1.f;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F16;
+        params["weight"]     = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels / groups, out_channels);
+        if (bias) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_channels);
+        }
+    }
+
+public:
+    Conv2d_grouped(int64_t in_channels,
+                   int64_t out_channels,
+                   int groups,
+                   std::pair<int, int> kernel_size,
+                   std::pair<int, int> stride   = {1, 1},
+                   std::pair<int, int> padding  = {0, 0},
+                   std::pair<int, int> dilation = {1, 1},
+                   bool bias                    = true)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          groups(groups),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias) {}
+
+    void set_scale(float scale_value) {
+        scale = scale_value;
+    }
+
+    std::string get_desc() override {
+        return "Conv2d_grouped";
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
+        ggml_tensor* b = nullptr;
+        if (bias) {
+            b = params["bias"];
+        }
+
+        if (groups == 1) {
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+                forward_params.conv2d.s0         = stride.second;
+                forward_params.conv2d.s1         = stride.first;
+                forward_params.conv2d.p0         = padding.second;
+                forward_params.conv2d.p1         = padding.first;
+                forward_params.conv2d.d0         = dilation.second;
+                forward_params.conv2d.d1         = dilation.first;
+                forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
+                forward_params.conv2d.circular_x = ctx->circular_x_enabled;
+                forward_params.conv2d.circular_y = ctx->circular_y_enabled;
+                forward_params.conv2d.scale      = scale;
+                return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params);
+            }
+            return ggml_ext_conv_2d(ctx->ggml_ctx, x, w, b,
+                                    stride.second, stride.first,
+                                    padding.second, padding.first,
+                                    dilation.second, dilation.first,
+                                    ctx->conv2d_direct_enabled,
+                                    ctx->circular_x_enabled,
+                                    ctx->circular_y_enabled,
+                                    scale);
+        }
+
+        if (groups == in_channels && groups == out_channels) {
+            ggml_tensor* res;
+            if (ctx->conv2d_direct_enabled) {
+                res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, w, x,
+                                             stride.second, stride.first,
+                                             padding.second, padding.first,
+                                             dilation.second, dilation.first);
+            } else {
+                res = ggml_conv_2d_dw(ctx->ggml_ctx, w, x,
+                                      stride.second, stride.first,
+                                      padding.second, padding.first,
+                                      dilation.second, dilation.first);
+            }
+            if (b) {
+                b   = ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, b->ne[0], 1);
+                res = ggml_add_inplace(ctx->ggml_ctx, res, b);
+            }
+            return res;
+        }
+
+        int64_t ic_g = in_channels / groups;
+        int64_t oc_g = out_channels / groups;
+
+        std::vector<ggml_tensor*> out_slices(groups);
+
+        for (int i = 0; i < groups; ++i) {
+            size_t x_offset  = i * ic_g * x->nb[2];
+            ggml_tensor* x_i = ggml_view_4d(ctx->ggml_ctx, x,
+                                            x->ne[0], x->ne[1], ic_g, x->ne[3],
+                                            x->nb[1], x->nb[2], x->nb[3],
+                                            x_offset);
+
+            size_t w_offset  = i * oc_g * w->nb[3];
+            ggml_tensor* w_i = ggml_view_4d(ctx->ggml_ctx, w,
+                                            w->ne[0], w->ne[1], w->ne[2], oc_g,
+                                            w->nb[1], w->nb[2], w->nb[3],
+                                            w_offset);
+
+            ggml_tensor* b_i = nullptr;
+            if (b) {
+                size_t b_offset = i * oc_g * b->nb[0];
+                b_i             = ggml_view_1d(ctx->ggml_ctx, b, oc_g, b_offset);
+            }
+
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+                forward_params.conv2d.s0         = stride.second;
+                forward_params.conv2d.s1         = stride.first;
+                forward_params.conv2d.p0         = padding.second;
+                forward_params.conv2d.p1         = padding.first;
+                forward_params.conv2d.d0         = dilation.second;
+                forward_params.conv2d.d1         = dilation.first;
+                forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
+                forward_params.conv2d.circular_x = ctx->circular_x_enabled;
+                forward_params.conv2d.circular_y = ctx->circular_y_enabled;
+                forward_params.conv2d.scale      = scale;
+                out_slices[i]                    = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x_i, w_i, b_i, prefix, forward_params);
+            } else {
+                out_slices[i] = ggml_ext_conv_2d(ctx->ggml_ctx, x_i, w_i, b_i,
+                                                 stride.second, stride.first,
+                                                 padding.second, padding.first,
+                                                 dilation.second, dilation.first,
+                                                 ctx->conv2d_direct_enabled,
+                                                 ctx->circular_x_enabled,
+                                                 ctx->circular_y_enabled,
+                                                 scale);
+            }
+        }
+
+        ggml_tensor* out = ggml_ext_vec_concat(ctx->ggml_ctx, out_slices, 2);
+
+        return out;
+    }
+};
+
+class Conv3d : public UnaryBlock {
+protected:
+    int64_t in_channels;
+    int64_t out_channels;
+    std::tuple<int, int, int> kernel_size;
+    std::tuple<int, int, int> stride;
+    std::tuple<int, int, int> padding;
+    std::tuple<int, int, int> dilation;
+    bool bias;
+    bool force_prec_f32;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F16;
+        params["weight"]     = ggml_new_tensor_4d(ctx,
+                                                  wtype,
+                                                  std::get<2>(kernel_size),
+                                                  std::get<1>(kernel_size),
+                                                  std::get<0>(kernel_size),
+                                                  in_channels * out_channels);
+        if (bias) {
+            params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+        }
+    }
+
+public:
+    Conv3d(int64_t in_channels,
+           int64_t out_channels,
+           std::tuple<int, int, int> kernel_size,
+           std::tuple<int, int, int> stride   = {1, 1, 1},
+           std::tuple<int, int, int> padding  = {0, 0, 0},
+           std::tuple<int, int, int> dilation = {1, 1, 1},
+           bool bias                          = true,
+           bool force_prec_f32                = false)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias),
+          force_prec_f32(force_prec_f32) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
+        ggml_tensor* b = nullptr;
+        if (ctx->weight_adapter) {
+            w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+            if (w->type != GGML_TYPE_F16) {
+                w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F16);
+            }
+        }
+        if (bias) {
+            b = params["bias"];
+            if (ctx->weight_adapter) {
+                b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+            }
+        }
+        return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
+                                std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
+                                std::get<2>(padding), std::get<1>(padding), std::get<0>(padding),
+                                std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
+                                force_prec_f32);
+    }
+};
+
+class LayerNorm : public UnaryBlock {
+protected:
+    int64_t normalized_shape;
+    float eps;
+    bool elementwise_affine;
+    bool bias;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        this->prefix = prefix;
+        if (elementwise_affine) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["weight"]     = ggml_new_tensor_1d(ctx, wtype, normalized_shape);
+            if (bias) {
+                enum ggml_type wtype = GGML_TYPE_F32;
+                params["bias"]       = ggml_new_tensor_1d(ctx, wtype, normalized_shape);
+            }
+        }
+    }
+
+public:
+    LayerNorm(int64_t normalized_shape,
+              float eps               = 1e-05f,
+              bool elementwise_affine = true,
+              bool bias               = true)
+        : normalized_shape(normalized_shape),
+          eps(eps),
+          elementwise_affine(elementwise_affine),
+          bias(bias) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = nullptr;
+        ggml_tensor* b = nullptr;
+
+        if (elementwise_affine) {
+            w = params["weight"];
+            if (ctx->weight_adapter) {
+                w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+            }
+            if (bias) {
+                b = params["bias"];
+                if (ctx->weight_adapter) {
+                    b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                }
+            }
+        }
+        return ggml_ext_layer_norm(ctx->ggml_ctx, x, w, b, eps);
+    }
+};
+
+class GroupNorm : public GGMLBlock {
+protected:
+    int num_groups;
+    int64_t num_channels;
+    float eps;
+    bool affine;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        this->prefix = prefix;
+        if (affine) {
+            enum ggml_type wtype      = GGML_TYPE_F32;
+            enum ggml_type bias_wtype = GGML_TYPE_F32;
+            params["weight"]          = ggml_new_tensor_1d(ctx, wtype, num_channels);
+            params["bias"]            = ggml_new_tensor_1d(ctx, bias_wtype, num_channels);
+        }
+    }
+
+public:
+    GroupNorm(int num_groups,
+              int64_t num_channels,
+              float eps   = 1e-05f,
+              bool affine = true)
+        : num_groups(num_groups),
+          num_channels(num_channels),
+          eps(eps),
+          affine(affine) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        ggml_tensor* w = nullptr;
+        ggml_tensor* b = nullptr;
+        if (affine) {
+            w = params["weight"];
+            b = params["bias"];
+            if (ctx->weight_adapter) {
+                w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+                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, eps);
+    }
+};
+
+class GroupNorm32 : public GroupNorm {
+public:
+    GroupNorm32(int64_t num_channels)
+        : GroupNorm(32, num_channels, 1e-06f) {}
+};
+
+class RMSNorm : public UnaryBlock {
+protected:
+    int64_t hidden_size;
+    float eps;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F32;
+        params["weight"]     = ggml_new_tensor_1d(ctx, wtype, hidden_size);
+    }
+
+public:
+    RMSNorm(int64_t hidden_size,
+            float eps = 1e-06f)
+        : hidden_size(hidden_size),
+          eps(eps) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
+        if (ctx->weight_adapter) {
+            w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+        }
+        x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
+        x = ggml_mul_inplace(ctx->ggml_ctx, x, w);
+        return x;
+    }
+};
+
+class MultiheadAttention : public GGMLBlock {
+protected:
+    int64_t embed_dim;
+    int64_t n_head;
+    bool proj_in;
+    std::string q_proj_name;
+    std::string k_proj_name;
+    std::string v_proj_name;
+    std::string in_proj_name;
+    std::string out_proj_name;
+
+public:
+    MultiheadAttention(int64_t embed_dim,
+                       int64_t n_head,
+                       bool qkv_proj_bias        = true,
+                       bool out_proj_bias        = true,
+                       bool proj_in              = false,
+                       std::string q_proj_name   = "q_proj",
+                       std::string k_proj_name   = "k_proj",
+                       std::string v_proj_name   = "v_proj",
+                       std::string in_proj_name  = "in_proj",
+                       std::string out_proj_name = "out_proj")
+        : embed_dim(embed_dim),
+          n_head(n_head),
+          proj_in(proj_in),
+          q_proj_name(q_proj_name),
+          k_proj_name(k_proj_name),
+          v_proj_name(v_proj_name),
+          in_proj_name(in_proj_name),
+          out_proj_name(out_proj_name) {
+        if (proj_in) {
+            blocks[in_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim * 3, qkv_proj_bias));
+        } else {
+            blocks[q_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, qkv_proj_bias));
+            blocks[k_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, qkv_proj_bias));
+            blocks[v_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, qkv_proj_bias));
+        }
+        blocks[out_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, out_proj_bias));
+    }
+
+    // x: [N, n_token, embed_dim]
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         ggml_tensor* mask = nullptr) {
+        auto out_proj = std::dynamic_pointer_cast<Linear>(blocks[out_proj_name]);
+
+        ggml_tensor* q;
+        ggml_tensor* k;
+        ggml_tensor* v;
+        if (proj_in) {
+            auto in_proj = std::dynamic_pointer_cast<Linear>(blocks[in_proj_name]);
+            auto qkv     = in_proj->forward(ctx, x);
+            auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
+            q            = qkv_vec[0];
+            k            = qkv_vec[1];
+            v            = qkv_vec[2];
+        } else {
+            auto q_proj = std::dynamic_pointer_cast<Linear>(blocks[q_proj_name]);
+            auto k_proj = std::dynamic_pointer_cast<Linear>(blocks[k_proj_name]);
+            auto v_proj = std::dynamic_pointer_cast<Linear>(blocks[v_proj_name]);
+
+            q = q_proj->forward(ctx, x);
+            k = k_proj->forward(ctx, x);
+            v = v_proj->forward(ctx, x);
+        }
+
+        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;
+    }
+};
+
+#endif  // __SD_MODEL_COMMON_GGML_BLOCK_HPP__

src/model/common/rope.hpp

@@ -2,9 +2,13 @@
 #define __SD_MODEL_COMMON_ROPE_HPP__
 
 #include <algorithm>
+#include <cassert>
 #include <cmath>
+#include <set>
 #include <vector>
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
 
 namespace Rope {
     enum class EmbedNDLayout {
@@ -814,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);
     }
 
@@ -1020,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/detector/yolov8.h

@@ -8,8 +8,9 @@
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_runner.h"
 #include "core/util.h"
+#include "model/common/ggml_block.hpp"
 
 struct YOLOv8Config {
     std::array<int, 23> out_channels{};
@@ -68,12 +69,12 @@ struct YOLOv8Config {
         }
 
         if (config.valid) {
-            LOG_DEBUG("yolov8: classes=%d, reg_max=%d, p3=%d, p4=%d, p5=%d",
-                      config.num_classes,
-                      config.reg_max,
-                      config.out_channels[15],
-                      config.out_channels[18],
-                      config.out_channels[21]);
+            LOG_VERBOSE("yolov8: classes=%d, reg_max=%d, p3=%d, p4=%d, p5=%d",
+                        config.num_classes,
+                        config.reg_max,
+                        config.out_channels[15],
+                        config.out_channels[18],
+                        config.out_channels[21]);
         }
         return config;
     }
@@ -355,7 +356,7 @@ struct YOLOv8Runner : public GGMLRunner {
 
     sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& input) {
         auto get_graph = [&]() { return build_graph(input); };
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false));
+        return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false));
     }
 };
 

src/model/diffusion/anima.hpp

@@ -2,6 +2,7 @@
 #define __SD_MODEL_DIFFUSION_ANIMA_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <memory>
 #include <utility>
@@ -46,11 +47,11 @@ namespace Anima {
             }
             if (detected_layers > 0) {
                 config.num_layers = detected_layers;
-                LOG_DEBUG("anima: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", head_dim = %" PRId64,
-                          config.num_layers,
-                          config.hidden_size,
-                          config.num_heads,
-                          config.head_dim);
+                LOG_VERBOSE("anima: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", head_dim = %" PRId64,
+                            config.num_layers,
+                            config.hidden_size,
+                            config.num_heads,
+                            config.head_dim);
             }
             return config;
         }
@@ -236,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,
@@ -248,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,
@@ -717,7 +716,7 @@ namespace Anima {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, t5_ids, t5_weights, ref_latents);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/animatediff.hpp

@@ -1,8 +1,10 @@
 #ifndef __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__
 #define __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "model/common/block.hpp"
+#include "model/common/ggml_block.hpp"
 
 // AnimateDiff (https://arxiv.org/abs/2307.04725) SD 1.5 motion modules.
 namespace AnimateDiff {
@@ -59,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/boogu.hpp

@@ -2,11 +2,15 @@
 #define __SD_MODEL_DIFFUSION_BOOGU_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <tuple>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/model.hpp"
@@ -109,16 +113,16 @@ namespace Boogu {
             }
             config.timestep_embed_dim = std::min<int64_t>(config.hidden_size, 1024);
 
-            LOG_DEBUG("boogu_image: layers=%" PRId64 ", double_stream_layers=%" PRId64 ", refiner_layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", head_dim=%" PRId64 ", in_channels=%" PRId64 ", out_channels=%" PRId64,
-                      config.num_layers,
-                      config.num_double_stream_layers,
-                      config.num_refiner_layers,
-                      config.hidden_size,
-                      config.num_attention_heads,
-                      config.num_kv_heads,
-                      config.head_dim,
-                      config.in_channels,
-                      config.out_channels);
+            LOG_VERBOSE("boogu_image: layers=%" PRId64 ", double_stream_layers=%" PRId64 ", refiner_layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", head_dim=%" PRId64 ", in_channels=%" PRId64 ", out_channels=%" PRId64,
+                        config.num_layers,
+                        config.num_double_stream_layers,
+                        config.num_refiner_layers,
+                        config.hidden_size,
+                        config.num_attention_heads,
+                        config.num_kv_heads,
+                        config.head_dim,
+                        config.in_channels,
+                        config.out_channels);
             return config;
         }
     };
@@ -815,7 +819,7 @@ namespace Boogu {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

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);
     }
 
@@ -435,7 +430,7 @@ struct ControlNet : public GGMLRunner {
             }
             return true;
         };
-        auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false, true, read_outputs);
+        auto compute_result = GGMLRunner::compute(get_graph, n_threads, false, true, read_outputs);
         control_outputs_ggml.clear();
         guided_hint_output_ggml = nullptr;
         if (!compute_result.has_value()) {
@@ -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/dit.hpp

@@ -1,7 +1,8 @@
 #ifndef __SD_MODEL_DIFFUSION_DIT_HPP__
 #define __SD_MODEL_DIFFUSION_DIT_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 
 namespace DiT {
     inline ggml_tensor* patchify(ggml_context* ctx,

src/model/diffusion/ernie_image.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__
 #define __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__
 
+#include <cinttypes>
 #include <memory>
 #include <vector>
 
@@ -72,13 +73,13 @@ namespace ErnieImage {
             for (int axis_dim : config.axes_dim) {
                 config.axes_dim_sum += axis_dim;
             }
-            LOG_DEBUG("ernie_image: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
-                      config.num_layers,
-                      config.hidden_size,
-                      config.num_heads,
-                      config.ffn_hidden_size,
-                      config.in_channels,
-                      config.out_channels);
+            LOG_VERBOSE("ernie_image: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
+                        config.num_layers,
+                        config.hidden_size,
+                        config.num_heads,
+                        config.ffn_hidden_size,
+                        config.in_channels,
+                        config.out_channels);
             return config;
         }
     };
@@ -182,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;
         }
@@ -440,7 +441,7 @@ namespace ErnieImage {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/flux.hpp

@@ -1,8 +1,11 @@
 #ifndef __SD_MODEL_DIFFUSION_FLUX_HPP__
 #define __SD_MODEL_DIFFUSION_FLUX_HPP__
 
+#include <cinttypes>
 #include <memory>
 #include <vector>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "core/util.h"
 #include "model/adapter/pulid.hpp"
@@ -123,16 +126,16 @@ namespace Flux {
                     config.guidance_embed = true;
                 }
                 if (name.find("__x0__") != std::string::npos) {
-                    LOG_DEBUG("using x0 prediction");
+                    LOG_VERBOSE("using x0 prediction");
                     config.chroma_radiance_params.use_x0 = true;
                 }
                 if (name.find("__32x32__") != std::string::npos) {
-                    LOG_DEBUG("using patch size 32");
+                    LOG_VERBOSE("using patch size 32");
                     config.patch_size = 32;
                 }
                 if (name.find("img_in_patch.weight") != std::string::npos) {
                     actual_radiance_patch_size = tensor_storage.ne[0];
-                    LOG_DEBUG("actual radiance patch size: %" PRId64, actual_radiance_patch_size);
+                    LOG_VERBOSE("actual radiance patch size: %" PRId64, actual_radiance_patch_size);
                 }
                 if (name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
                     config.is_chroma = true;
@@ -169,7 +172,7 @@ namespace Flux {
             }
             if (actual_radiance_patch_size > 0 && actual_radiance_patch_size != config.patch_size) {
                 GGML_ASSERT(config.patch_size == 2 * actual_radiance_patch_size);
-                LOG_DEBUG("using fake x2 patch size");
+                LOG_VERBOSE("using fake x2 patch size");
                 config.chroma_radiance_params.fake_patch_size_x2 = true;
             }
             if (head_dim > 0) {
@@ -179,13 +182,13 @@ namespace Flux {
             for (int axis_dim : config.axes_dim) {
                 config.axes_dim_sum += axis_dim;
             }
-            LOG_DEBUG("flux: depth = %d, depth_single_blocks = %d, guidance_embed = %s, context_in_dim = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %d",
-                      config.depth,
-                      config.depth_single_blocks,
-                      config.guidance_embed ? "true" : "false",
-                      config.context_in_dim,
-                      config.hidden_size,
-                      config.num_heads);
+            LOG_VERBOSE("flux: depth = %d, depth_single_blocks = %d, guidance_embed = %s, context_in_dim = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %d",
+                        config.depth,
+                        config.depth_single_blocks,
+                        config.guidance_embed ? "true" : "false",
+                        config.context_in_dim,
+                        config.hidden_size,
+                        config.num_heads);
             return config;
         }
     };
@@ -1560,7 +1563,7 @@ namespace Flux {
                                             config.axes_dim,
                                             sd_version_is_longcat(version));
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
-            // LOG_DEBUG("pos_len %d", pos_len);
+            // 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);
             // pe->data = pe_vec.data();
             // print_ggml_tensor(pe);
@@ -1626,7 +1629,7 @@ namespace Flux {
                 return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers, pulid_id, pulid_id_weight);
             };
 
-            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
             return result;
         }
 
@@ -1702,7 +1705,7 @@ namespace Flux {
                 GGML_ASSERT(!out_opt.empty());
                 out = std::move(out_opt);
                 print_sd_tensor(out);
-                LOG_DEBUG("flux test done in %lldms", t1 - t0);
+                LOG_VERBOSE("flux test done in %lldms", t1 - t0);
             }
         }
 
@@ -1711,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;
@@ -1733,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

@@ -329,7 +329,7 @@ namespace HiDreamO1 {
             auto get_graph = [&]() {
                 return build_graph(image);
             };
-            auto output = GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end);
+            auto output = GGMLRunner::compute(get_graph, n_threads, auto_runner_end);
             return output.has_value() ? std::move(output.value()) : sd::Tensor<float>();
         }
     };
@@ -457,7 +457,7 @@ namespace HiDreamO1 {
             auto get_graph = [&]() {
                 return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -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

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_HUNYUAN_HPP__
 #define __SD_MODEL_DIFFUSION_HUNYUAN_HPP__
 
+#include <cinttypes>
 #include <memory>
 
 #include "model/common/block.hpp"
@@ -53,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
@@ -266,16 +267,16 @@ namespace Hunyuan {
             GGML_ASSERT(config.hidden_size / config.num_heads == config.axes_dim_sum);
 
             if (inferred) {
-                LOG_DEBUG("hunyuan video: depth = %d, single depth = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", hidden_size = %" PRId64 ", context_in_dim = %" PRId64 ", patch_size = %dx%dx%d",
-                          config.depth,
-                          config.depth_single_blocks,
-                          config.in_channels,
-                          config.out_channels,
-                          config.hidden_size,
-                          config.context_in_dim,
-                          std::get<0>(config.patch_size),
-                          std::get<1>(config.patch_size),
-                          std::get<2>(config.patch_size));
+                LOG_VERBOSE("hunyuan video: depth = %d, single depth = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", hidden_size = %" PRId64 ", context_in_dim = %" PRId64 ", patch_size = %dx%dx%d",
+                            config.depth,
+                            config.depth_single_blocks,
+                            config.in_channels,
+                            config.out_channels,
+                            config.hidden_size,
+                            config.context_in_dim,
+                            std::get<0>(config.patch_size),
+                            std::get<1>(config.patch_size),
+                            std::get<2>(config.patch_size));
             }
             return config;
         }
@@ -615,7 +616,7 @@ namespace Hunyuan {
                                                          config.theta,
                                                          config.axes_dim);
             int64_t pos_len = static_cast<int64_t>(pe_vec.size() / config.axes_dim_sum / 2);
-            // LOG_DEBUG("pos_len %d", pos_len);
+            // 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);
             // pe->data = pe_vec.data();
             // print_ggml_tensor(pe, true, "pe");
@@ -654,7 +655,7 @@ namespace Hunyuan {
                 return build_graph(x, timesteps, context, c_concat, y, guidance, byt5, vision, timestep_r);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/ideogram4.hpp

@@ -2,14 +2,18 @@
 #define __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <cstdlib>
 #include <memory>
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/model.hpp"
 
@@ -58,11 +62,11 @@ namespace Ideogram4 {
             }
             if (detected_layers > 0) {
                 config.num_layers = detected_layers;
-                LOG_DEBUG("ideogram4: num_layers = %" PRId64 ", emb_dim = %" PRId64 ", num_heads = %" PRId64 ", intermediate_size = %" PRId64,
-                          config.num_layers,
-                          config.emb_dim,
-                          config.num_heads,
-                          config.intermediate_size);
+                LOG_VERBOSE("ideogram4: num_layers = %" PRId64 ", emb_dim = %" PRId64 ", num_heads = %" PRId64 ", intermediate_size = %" PRId64,
+                            config.num_layers,
+                            config.emb_dim,
+                            config.num_heads,
+                            config.intermediate_size);
             }
             return config;
         }
@@ -465,7 +469,7 @@ namespace Ideogram4 {
                 }
             }
             if (has_uncond_model) {
-                LOG_DEBUG("using uncond model");
+                LOG_VERBOSE("using uncond model");
                 uncond_model = Ideogram4Transformer(config);
                 uncond_model.init(params_ctx, tensor_storage_map, uncond_prefix);
             }
@@ -537,7 +541,7 @@ namespace Ideogram4 {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, use_uncond_model);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/krea2.hpp

@@ -12,8 +12,11 @@
 #include <utility>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/flux.hpp"
@@ -143,16 +146,16 @@ namespace Krea2 {
             }
             config.update_axes_dim();
 
-            LOG_DEBUG("krea2: layers=%" PRId64 ", features=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", text_dim=%" PRId64 ", text_layers=%" PRId64 ", text_heads=%" PRId64 ", text_kv_heads=%" PRId64 ", channels=%" PRId64,
-                      config.layers,
-                      config.features,
-                      config.heads,
-                      config.kv_heads,
-                      config.text_dim,
-                      config.text_layers,
-                      config.text_heads,
-                      config.text_kv_heads,
-                      config.in_channels);
+            LOG_VERBOSE("krea2: layers=%" PRId64 ", features=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", text_dim=%" PRId64 ", text_layers=%" PRId64 ", text_heads=%" PRId64 ", text_kv_heads=%" PRId64 ", channels=%" PRId64,
+                        config.layers,
+                        config.features,
+                        config.heads,
+                        config.kv_heads,
+                        config.text_dim,
+                        config.text_layers,
+                        config.text_heads,
+                        config.text_kv_heads,
+                        config.in_channels);
             return config;
         }
     };
@@ -229,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,
@@ -775,7 +777,7 @@ namespace Krea2 {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents, ref_image_params);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/lens.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_LENS_HPP__
 #define __SD_MODEL_DIFFUSION_LENS_HPP__
 
+#include <cinttypes>
 #include <memory>
 #include <vector>
 
@@ -66,14 +67,14 @@ namespace Lens {
             for (int axis_dim : config.axes_dim) {
                 config.axes_dim_sum += axis_dim;
             }
-            LOG_DEBUG("lens: num_layers = %d, selected_layer_count = %d, hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", attention_head_dim = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
-                      config.num_layers,
-                      config.selected_layer_count,
-                      config.num_attention_heads * config.attention_head_dim,
-                      config.num_attention_heads,
-                      config.attention_head_dim,
-                      config.in_channels,
-                      config.out_channels);
+            LOG_VERBOSE("lens: num_layers = %d, selected_layer_count = %d, hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", attention_head_dim = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
+                        config.num_layers,
+                        config.selected_layer_count,
+                        config.num_attention_heads * config.attention_head_dim,
+                        config.num_attention_heads,
+                        config.attention_head_dim,
+                        config.in_channels,
+                        config.out_channels);
             return config;
         }
     };
@@ -408,7 +409,7 @@ namespace Lens {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/lingbot_video.hpp

@@ -127,17 +127,17 @@ namespace LingBotVideo {
                 config.topk_group            = 2;
                 config.routed_scaling_factor = 2.5f;
             }
-            LOG_DEBUG("lingbot_video: depth = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64 ", text_dim = %" PRId64 ", experts = %" PRId64 ", experts_per_tok = %" PRId64 ", n_group = %" PRId64 ", topk_group = %" PRId64 ", route_scale = %.2f, sparse_layers = %zu",
-                      config.depth,
-                      config.hidden_size,
-                      config.num_attention_heads,
-                      config.text_dim,
-                      config.num_experts,
-                      config.num_experts_per_tok,
-                      config.n_group,
-                      config.topk_group,
-                      config.routed_scaling_factor,
-                      config.sparse_layers.size());
+            LOG_VERBOSE("lingbot_video: depth = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64 ", text_dim = %" PRId64 ", experts = %" PRId64 ", experts_per_tok = %" PRId64 ", n_group = %" PRId64 ", topk_group = %" PRId64 ", route_scale = %.2f, sparse_layers = %zu",
+                        config.depth,
+                        config.hidden_size,
+                        config.num_attention_heads,
+                        config.text_dim,
+                        config.num_experts,
+                        config.num_experts_per_tok,
+                        config.n_group,
+                        config.topk_group,
+                        config.routed_scaling_factor,
+                        config.sparse_layers.size());
             return config;
         }
     };
@@ -674,7 +674,7 @@ namespace LingBotVideo {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/ltxv.hpp

@@ -2,12 +2,15 @@
 #define __SD_MODEL_DIFFUSION_LTXV_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <memory>
 #include <string>
 #include <tuple>
 #include <utility>
 #include <vector>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "model/common/block.hpp"
 #include "model/common/rope.hpp"
@@ -274,12 +277,12 @@ namespace LTXV {
                     config.audio_connector_apply_gated_attention = true;
                 }
             }
-            LOG_DEBUG("ltxav: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", audio_hidden_size = %" PRId64 ", audio_num_attention_heads = %" PRId64,
-                      config.num_layers,
-                      config.hidden_size,
-                      config.num_attention_heads,
-                      config.audio_hidden_size,
-                      config.audio_num_attention_heads);
+            LOG_VERBOSE("ltxav: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", audio_hidden_size = %" PRId64 ", audio_num_attention_heads = %" PRId64,
+                        config.num_layers,
+                        config.hidden_size,
+                        config.num_attention_heads,
+                        config.audio_hidden_size,
+                        config.audio_num_attention_heads);
             return config;
         }
     };
@@ -706,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,
@@ -1998,7 +2000,7 @@ namespace LTXV {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, audio_x, audio_timesteps, audio_length, frame_rate, video_positions);
             };
-            auto out = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            auto out = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
             return out;
         }
 
@@ -2070,7 +2072,7 @@ namespace LTXV {
 
             GGML_ASSERT(!out_opt.empty());
             print_sd_tensor(out_opt, false, "ltxav_out");
-            LOG_DEBUG("ltxav test done in %lldms", t1 - t0);
+            LOG_VERBOSE("ltxav test done in %lldms", t1 - t0);
         }
 
         static void load_from_file_and_test(const std::string& model_path,
@@ -2084,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;
@@ -2104,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/mage_flow.hpp

@@ -142,7 +142,7 @@ namespace MageFlow {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/minimax_h3.hpp

@@ -2,12 +2,14 @@
 #define __SD_MODEL_DIFFUSION_MINIMAX_H3_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <set>
 #include <string>
 #include <tuple>
 #include <utility>
 #include <vector>
+#include "core/ggml_tensor_utils.h"
 
 #include "core/ggml_graph_cut.h"
 #include "model/diffusion/dit.hpp"
@@ -106,14 +108,14 @@ namespace MiniMaxH3 {
                 config.rope_inv_freq_len = inv_freq->ne[0];
             }
 
-            LOG_DEBUG("minimax_h3: layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64
-                      ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64,
-                      config.num_layers,
-                      config.hidden_size,
-                      config.num_attention_heads,
-                      config.attention_head_dim,
-                      config.ffn_hidden_size,
-                      config.adaln_curve_grid);
+            LOG_VERBOSE("minimax_h3: layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64
+                        ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64,
+                        config.num_layers,
+                        config.hidden_size,
+                        config.num_attention_heads,
+                        config.attention_head_dim,
+                        config.ffn_hidden_size,
+                        config.adaln_curve_grid);
             return config;
         }
     };
@@ -213,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,
@@ -1166,9 +1167,9 @@ namespace MiniMaxH3 {
                                    extra->video_sigma_shift,
                                    extra->audio_sigma_shift);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
-                                                                              n_threads,
-                                                                              false),
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph,
+                                                                       n_threads,
+                                                                       false),
                                                    params.x->dim());
         }
     };

src/model/diffusion/minit2i.hpp

@@ -2,6 +2,7 @@
 #define __SD_MODEL_DIFFUSION_MINIT2I_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <cstdint>
 #include <cstdlib>
@@ -9,7 +10,10 @@
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/model.hpp"
@@ -108,15 +112,15 @@ namespace MiniT2I {
                 config.head_dim  = config.hidden_size == 1248 ? 52 : 64;
                 config.num_heads = config.hidden_size / config.head_dim;
             }
-            LOG_DEBUG("minit2i: hidden_size=%" PRId64 ", txt_hidden_size=%" PRId64 ", heads=%" PRId64 ", head_dim=%" PRId64 ", double_blocks=%" PRId64 ", txt_blocks=%" PRId64 ", patch=%" PRId64 ", in_channels=%" PRId64,
-                      config.hidden_size,
-                      config.txt_hidden_size,
-                      config.num_heads,
-                      config.head_dim,
-                      config.depth_double,
-                      config.txt_preamble_depth,
-                      config.patch_size,
-                      config.in_channels);
+            LOG_VERBOSE("minit2i: hidden_size=%" PRId64 ", txt_hidden_size=%" PRId64 ", heads=%" PRId64 ", head_dim=%" PRId64 ", double_blocks=%" PRId64 ", txt_blocks=%" PRId64 ", patch=%" PRId64 ", in_channels=%" PRId64,
+                        config.hidden_size,
+                        config.txt_hidden_size,
+                        config.num_heads,
+                        config.head_dim,
+                        config.depth_double,
+                        config.txt_preamble_depth,
+                        config.patch_size,
+                        config.in_channels);
             return config;
         }
     };
@@ -589,7 +593,7 @@ namespace MiniT2I {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, mask);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/mmdit.hpp

@@ -2,12 +2,18 @@
 #define __SD_MODEL_DIFFUSION_MMDIT_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <memory>
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
 #include "model/common/block.hpp"
+#include "model/common/ggml_block.hpp"
 #include "model/diffusion/model.hpp"
 #include "model_loader.h"
 
@@ -120,16 +126,16 @@ struct MMDiTConfig {
         }
 
         if (has_weight_config) {
-            LOG_DEBUG("mmdit: num_layers = %" PRId64 ", num_mmdit_x_layers = %" PRId64 ", hidden_size = %" PRId64 ", patch_size = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", context_size = %" PRId64 ", adm_in_channels = %" PRId64 ", qk_norm = %s",
-                      config.depth,
-                      config.d_self + 1,
-                      config.hidden_size,
-                      config.patch_size,
-                      config.in_channels,
-                      config.out_channels,
-                      config.context_size,
-                      config.adm_in_channels,
-                      config.qk_norm.empty() ? "none" : config.qk_norm.c_str());
+            LOG_VERBOSE("mmdit: num_layers = %" PRId64 ", num_mmdit_x_layers = %" PRId64 ", hidden_size = %" PRId64 ", patch_size = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", context_size = %" PRId64 ", adm_in_channels = %" PRId64 ", qk_norm = %s",
+                        config.depth,
+                        config.d_self + 1,
+                        config.hidden_size,
+                        config.patch_size,
+                        config.in_channels,
+                        config.out_channels,
+                        config.context_size,
+                        config.adm_in_channels,
+                        config.qk_norm.empty() ? "none" : config.qk_norm.c_str());
         }
         return config;
     }
@@ -359,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;
     }
 };
@@ -581,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,
@@ -598,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],
@@ -642,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,
@@ -674,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,
@@ -987,7 +993,7 @@ struct MMDiTRunner : public DiffusionModelRunner {
             return build_graph(x, timesteps, context, y, skip_layers);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
     }
 
     sd::Tensor<float> compute(int n_threads,
@@ -1045,7 +1051,7 @@ struct MMDiTRunner : public DiffusionModelRunner {
             GGML_ASSERT(!out_opt.empty());
             out = std::move(out_opt);
             print_sd_tensor(out);
-            LOG_DEBUG("mmdit test done in %lldms", t1 - t0);
+            LOG_VERBOSE("mmdit test done in %lldms", t1 - t0);
         }
     }
 
@@ -1058,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

@@ -5,7 +5,7 @@
 #include <utility>
 #include <variant>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_runner.h"
 #include "core/tensor_ggml.hpp"
 #include "model/common/rope.hpp"
 #include "model_manager.h"
@@ -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/pid.hpp

@@ -1,13 +1,18 @@
 #ifndef __SD_MODEL_DIFFUSION_PID_HPP__
 #define __SD_MODEL_DIFFUSION_PID_HPP__
 
+#include <cinttypes>
 #include <cmath>
 #include <cstdlib>
 #include <memory>
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/mmdit.hpp"
@@ -109,16 +114,16 @@ namespace Pid {
                 config.lq_latent_channels    = latent_proj_in_channels;
                 config.lq_latent_down_factor = latent_proj_in_channels >= 64 ? 16 : 8;
             }
-            LOG_DEBUG("pid: version = %s, patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_hidden_dim = %" PRId64 ", lq_latent_down_factor = %" PRId64 ", lq_latent_unpatchify_factor = %" PRId64 ", lq_interval = %" PRId64,
-                      config.pit_lq_inject ? "1.5" : "1",
-                      config.patch_depth,
-                      config.pixel_depth,
-                      config.patch_mlp_hidden_dim,
-                      config.lq_latent_channels,
-                      config.lq_hidden_dim,
-                      config.lq_latent_down_factor,
-                      config.lq_latent_unpatchify_factor,
-                      config.lq_interval);
+            LOG_VERBOSE("pid: version = %s, patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_hidden_dim = %" PRId64 ", lq_latent_down_factor = %" PRId64 ", lq_latent_unpatchify_factor = %" PRId64 ", lq_interval = %" PRId64,
+                        config.pit_lq_inject ? "1.5" : "1",
+                        config.patch_depth,
+                        config.pixel_depth,
+                        config.patch_mlp_hidden_dim,
+                        config.lq_latent_channels,
+                        config.lq_hidden_dim,
+                        config.lq_latent_down_factor,
+                        config.lq_latent_unpatchify_factor,
+                        config.lq_interval);
             return config;
         }
     };
@@ -938,7 +943,7 @@ namespace Pid {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, lq_latent, degrade_sigma);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/qwen_image.hpp

@@ -2,6 +2,8 @@
 #define __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__
 
 #include <memory>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "core/util.h"
 #include "model/common/block.hpp"
@@ -49,9 +51,9 @@ namespace Qwen {
                     }
                 }
             }
-            LOG_DEBUG("qwen_image: num_layers = %d, zero_cond_t = %s",
-                      config.num_layers,
-                      config.zero_cond_t ? "true" : "false");
+            LOG_VERBOSE("qwen_image: num_layers = %d, zero_cond_t = %s",
+                        config.num_layers,
+                        config.zero_cond_t ? "true" : "false");
             return config;
         }
     };
@@ -646,7 +648,7 @@ namespace Qwen {
                                                   circular_x_enabled,
                                                   config.axes_dim);
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
-            // LOG_DEBUG("pos_len %d", pos_len);
+            // 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);
             // pe->data = pe_vec.data();
             // print_ggml_tensor(pe, true, "pe");
@@ -707,7 +709,7 @@ namespace Qwen {
                 return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -760,7 +762,7 @@ namespace Qwen {
                 GGML_ASSERT(!out_opt.empty());
                 out = std::move(out_opt);
                 print_sd_tensor(out);
-                LOG_DEBUG("qwen_image test done in %lldms", t1 - t0);
+                LOG_VERBOSE("qwen_image test done in %lldms", t1 - t0);
             }
         }
 
@@ -771,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;
@@ -791,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/sefi_image.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__
 #define __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__
 
+#include <cinttypes>
 #include <memory>
 
 #include "model/common/block.hpp"
@@ -34,10 +35,10 @@ namespace SefiImage {
                     config.hidden_size              = tensor_storage.ne[1] * 2;
                 }
             }
-            LOG_DEBUG("sefi_image: semantic_channels = %" PRId64 ", texture_latent_channels = %" PRId64 ", hidden_size = %" PRId64,
-                      config.semantic_channels,
-                      config.texture_latent_channels,
-                      config.hidden_size);
+            LOG_VERBOSE("sefi_image: semantic_channels = %" PRId64 ", texture_latent_channels = %" PRId64 ", hidden_size = %" PRId64,
+                        config.semantic_channels,
+                        config.texture_latent_channels,
+                        config.hidden_size);
             return config;
         }
     };

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

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <vector>
+#include "core/ggml_tensor_utils.h"
 
 #include "model.h"
 #include "model/common/block.hpp"
@@ -128,15 +129,15 @@ struct UNetConfig {
             }
         }
 
-        LOG_DEBUG("unet: in_channels = %d, out_channels = %d, model_channels = %d, time_embed_dim = %d, context_dim = %d, adm_in_channels = %d, num_res_blocks = %d, tiny_unet = %s",
-                  config.in_channels,
-                  config.out_channels,
-                  config.model_channels,
-                  config.time_embed_dim,
-                  config.context_dim,
-                  config.adm_in_channels,
-                  config.num_res_blocks,
-                  config.tiny_unet ? "true" : "false");
+        LOG_VERBOSE("unet: in_channels = %d, out_channels = %d, model_channels = %d, time_embed_dim = %d, context_dim = %d, adm_in_channels = %d, num_res_blocks = %d, tiny_unet = %s",
+                    config.in_channels,
+                    config.out_channels,
+                    config.model_channels,
+                    config.time_embed_dim,
+                    config.context_dim,
+                    config.adm_in_channels,
+                    config.num_res_blocks,
+                    config.tiny_unet ? "true" : "false");
         return config;
     }
 };
@@ -835,7 +836,7 @@ struct UNetModelRunner : public DiffusionModelRunner {
             return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, ip_context, ip_scale);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
     }
 
     sd::Tensor<float> compute(int n_threads,
@@ -904,7 +905,7 @@ struct UNetModelRunner : public DiffusionModelRunner {
             GGML_ASSERT(!out_opt.empty());
             out = std::move(out_opt);
             print_sd_tensor(out);
-            LOG_DEBUG("unet test done in %lldms", t1 - t0);
+            LOG_VERBOSE("unet test done in %lldms", t1 - t0);
         }
     }
 };

src/model/diffusion/wan.hpp

@@ -1,9 +1,13 @@
 #ifndef __SD_MODEL_DIFFUSION_WAN_HPP__
 #define __SD_MODEL_DIFFUSION_WAN_HPP__
 
+#include <algorithm>
+#include <cinttypes>
 #include <map>
 #include <memory>
 #include <utility>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "model/common/block.hpp"
 #include "model/common/rope.hpp"
@@ -16,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;
@@ -71,17 +80,21 @@ 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;
                 }
             }
-            LOG_DEBUG("wan: model_type = %s, num_layers = %d, vace_layers = %d, dim = %" PRId64 ", ffn_dim = %" PRId64 ", num_heads = %" PRId64,
-                      config.model_type.c_str(),
-                      config.num_layers,
-                      config.vace_layers,
-                      config.dim,
-                      config.ffn_dim,
-                      config.num_heads);
+            LOG_VERBOSE("wan: model_type = %s, num_layers = %d, vace_layers = %d, dim = %" PRId64 ", ffn_dim = %" PRId64 ", num_heads = %" PRId64,
+                        config.model_type.c_str(),
+                        config.num_layers,
+                        config.vace_layers,
+                        config.dim,
+                        config.ffn_dim,
+                        config.num_heads);
             return config;
         }
     };
@@ -190,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;
@@ -252,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);
 
@@ -262,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]
@@ -529,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)
@@ -551,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,
@@ -592,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,
@@ -639,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"]);
@@ -667,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);
@@ -711,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)]);
 
@@ -728,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");
@@ -744,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,]
@@ -776,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]
 
@@ -836,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;
@@ -888,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);
@@ -898,18 +997,35 @@ 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_DEBUG("pos_len %d", pos_len);
+            // 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);
             // pe->data = pe_vec.data();
             // print_ggml_tensor(pe);
@@ -930,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);
 
@@ -945,12 +1064,14 @@ 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<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -958,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,
@@ -966,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() {
@@ -1007,7 +1136,7 @@ namespace WAN {
                 GGML_ASSERT(!out_opt.empty());
                 out = std::move(out_opt);
                 print_sd_tensor(out);
-                LOG_DEBUG("wan test done in %lldms", t1 - t0);
+                LOG_VERBOSE("wan test done in %lldms", t1 - t0);
             }
         }
 
@@ -1017,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;
@@ -1037,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

@@ -2,8 +2,14 @@
 #define __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/diffusion/flux.hpp"
 #include "model/diffusion/mmdit.hpp"
 #include "model/diffusion/model.hpp"
@@ -107,14 +113,14 @@ namespace ZImage {
                     config.num_kv_heads = std::max<int64_t>(1, (qkv_heads - config.num_heads) / 2);
                 }
             }
-            LOG_DEBUG("z_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
-                      config.num_layers,
-                      config.num_refiner_layers,
-                      config.hidden_size,
-                      config.num_heads,
-                      config.num_kv_heads,
-                      config.in_channels,
-                      config.out_channels);
+            LOG_VERBOSE("z_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
+                        config.num_layers,
+                        config.num_refiner_layers,
+                        config.hidden_size,
+                        config.num_heads,
+                        config.num_kv_heads,
+                        config.in_channels,
+                        config.out_channels);
             return config;
         }
     };
@@ -603,7 +609,7 @@ namespace ZImage {
                                                circular_x_enabled,
                                                config.axes_dim);
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
-            // LOG_DEBUG("pos_len %d", pos_len);
+            // 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);
             // pe->data = pe_vec.data();
             // print_ggml_tensor(pe, true, "pe");
@@ -636,7 +642,7 @@ namespace ZImage {
                 return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -689,7 +695,7 @@ namespace ZImage {
                 GGML_ASSERT(!out_opt.empty());
                 out = std::move(out_opt);
                 print_sd_tensor(out);
-                LOG_DEBUG("z_image test done in %lldms", t1 - t0);
+                LOG_VERBOSE("z_image test done in %lldms", t1 - t0);
             }
         }
 
@@ -700,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;
@@ -722,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

@@ -1,8 +1,11 @@
 #ifndef __SD_MODEL_TE_CLIP_HPP__
 #define __SD_MODEL_TE_CLIP_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
 #include "model.h"
+#include "model/common/ggml_block.hpp"
 #include "tokenizers/clip_tokenizer.h"
 
 /*================================================ FrozenCLIPEmbedder ================================================*/
@@ -100,13 +103,13 @@ public:
                          const std::string& graph_cut_prefix = "") {
         // x: [N, n_token, d_model]
         int layer_idx = n_layer - 1;
-        // LOG_DEBUG("clip_skip %d", clip_skip);
+        // LOG_VERBOSE("clip_skip %d", clip_skip);
         if (clip_skip > 0) {
             layer_idx = n_layer - clip_skip;
         }
 
         for (int i = 0; i < n_layer; i++) {
-            // LOG_DEBUG("layer %d", i);
+            // LOG_VERBOSE("layer %d", i);
             if (i == layer_idx + 1) {
                 break;
             }
@@ -116,7 +119,7 @@ public:
             if (!graph_cut_prefix.empty()) {
                 sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".layers." + std::to_string(i), "x");
             }
-            // LOG_DEBUG("layer %d", i);
+            // LOG_VERBOSE("layer %d", i);
         }
         return x;
     }
@@ -142,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,
@@ -320,7 +327,7 @@ public:
             if (text_projection != nullptr) {
                 pooled = ggml_ext_linear(ctx->ggml_ctx, pooled, text_projection, nullptr);
             } else {
-                LOG_DEBUG("identity projection");
+                LOG_VERBOSE("identity projection");
             }
             return pooled;  // [hidden_size, 1, 1]
         }
@@ -572,7 +579,7 @@ struct CLIPTextModelRunner : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip);
         };
-        auto result = GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end);
+        auto result = GGMLRunner::compute(get_graph, n_threads, auto_runner_end);
         if (return_pooled) {
             return take_or_empty(std::move(result));
         }

src/model/te/llm.hpp

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <array>
+#include <cinttypes>
 #include <cmath>
 #include <fstream>
 #include <functional>
@@ -18,8 +19,13 @@
 #include <utility>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
 #include "json.hpp"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model_loader.h"
 #include "model_manager.h"
@@ -133,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) {
@@ -224,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;
@@ -271,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;
                 }
@@ -319,11 +328,25 @@ namespace LLM {
                     config.vision.deepstack_visual_indexes = {8, 16, 24};
                 }
             }
-            LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
-                      config.num_layers,
-                      config.vocab_size,
-                      config.hidden_size,
-                      config.intermediate_size);
+            LOG_VERBOSE("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
+                        config.num_layers,
+                        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;
         }
     };
@@ -1353,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]
@@ -1880,16 +1903,12 @@ 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_DEBUG("enable llm vision");
+                LOG_VERBOSE("enable llm vision");
                 if (config.llama_cpp_style) {
-                    LOG_DEBUG("llama.cpp style vision weight");
+                    LOG_VERBOSE("llama.cpp style vision weight");
                 }
             }
             model = LLM(config, enable_vision, config.llama_cpp_style);
@@ -2091,7 +2110,7 @@ namespace LLM {
                                    out_layers,
                                    return_all_hidden_states);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end),
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_runner_end),
                                                    input_ids.dim() + 1);
         }
 
@@ -2175,7 +2194,7 @@ namespace LLM {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_encode_image_graph(image);
             };
-            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
+            return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
         }
 
         ggml_cgraph* build_encode_image_outputs_graph(const sd::Tensor<float>& image_tensor) {
@@ -2287,7 +2306,7 @@ namespace LLM {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_encode_image_outputs_graph(image);
             };
-            auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
+            auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
             if (combined.empty()) {
                 return {};
             }
@@ -2313,7 +2332,7 @@ namespace LLM {
             auto get_graph    = [&]() -> ggml_cgraph* {
                 return build_encode_video_block_outputs_graph(pixel_values, grid_h, grid_w);
             };
-            auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
+            auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
             if (combined.empty()) {
                 return {};
             }
@@ -2375,7 +2394,7 @@ namespace LLM {
                     ss << "['" << item.first << "', " << item.second << "], ";
                 }
                 ss << "]";
-                LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str());
+                LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
             }
 
             std::vector<int> tokens;
@@ -2383,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);
             }
@@ -2426,7 +2448,7 @@ namespace LLM {
                     out = std::move(out_opt);
                     print_sd_tensor(out, false, "image_embed");
                     image_embed = out;
-                    LOG_DEBUG("llm encode_image test done in %lldms", t1 - t0);
+                    LOG_VERBOSE("llm encode_image test done in %lldms", t1 - t0);
                 }
 
                 std::string placeholder  = "<|image_pad|>";
@@ -2466,7 +2488,7 @@ namespace LLM {
                 GGML_ASSERT(!out_opt.empty());
                 out = std::move(out_opt);
                 print_sd_tensor(out);
-                LOG_DEBUG("llm test done in %lldms", t1 - t0);
+                LOG_VERBOSE("llm test done in %lldms", t1 - t0);
             } else if (test_vit) {
                 // auto image = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 280, 280, 3);
                 // ggml_set_f32(image, 0.f);
@@ -2485,7 +2507,7 @@ namespace LLM {
                 // auto ref_out = load_tensor_from_file(ctx, "qwen2vl.bin");
                 // ggml_ext_tensor_diff(ref_out, out, 0.01f);
 
-                LOG_DEBUG("llm test done in %lldms", t1 - t0);
+                LOG_VERBOSE("llm test done in %lldms", t1 - t0);
             } else if (test_mistral) {
                 std::pair<int, int> prompt_attn_range;
                 std::string text        = "[SYSTEM_PROMPT]You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object\nattribution and actions without speculation.[/SYSTEM_PROMPT][INST]";
@@ -2510,7 +2532,7 @@ namespace LLM {
                 GGML_ASSERT(!out_opt.empty());
                 out = std::move(out_opt);
                 print_sd_tensor(out);
-                LOG_DEBUG("llm test done in %lldms", t1 - t0);
+                LOG_VERBOSE("llm test done in %lldms", t1 - t0);
             } else if (test_qwen3) {
                 std::pair<int, int> prompt_attn_range;
                 std::string text        = "<|im_start|>user\n";
@@ -2535,7 +2557,7 @@ namespace LLM {
                 GGML_ASSERT(!out_opt.empty());
                 out = std::move(out_opt);
                 print_sd_tensor(out);
-                LOG_DEBUG("llm test done in %lldms", t1 - t0);
+                LOG_VERBOSE("llm test done in %lldms", t1 - t0);
             } else {
                 std::pair<int, int> prompt_attn_range;
                 std::string text        = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n";
@@ -2560,7 +2582,7 @@ namespace LLM {
                 GGML_ASSERT(!out_opt.empty());
                 out = std::move(out_opt);
                 print_sd_tensor(out);
-                LOG_DEBUG("llm test done in %lldms", t1 - t0);
+                LOG_VERBOSE("llm test done in %lldms", t1 - t0);
             }
         }
 
@@ -2570,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;
@@ -2595,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

@@ -10,7 +10,12 @@
 #include <string>
 #include <unordered_map>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model_loader.h"
 #include "model_manager.h"
 #include "tokenizers/t5_unigram_tokenizer.h"
@@ -246,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};
@@ -455,7 +460,7 @@ struct T5Runner : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input_ids, attention_mask);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end), 3);
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_runner_end), 3);
     }
 
     static std::vector<int> _relative_position_bucket(const std::vector<int>& relative_position,
@@ -554,7 +559,7 @@ struct T5Embedder {
                 ss << "['" << item.first << "', " << item.second << "], ";
             }
             ss << "]";
-            LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str());
+            LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
         }
 
         std::vector<int> tokens;
@@ -562,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);
         }
@@ -612,7 +620,7 @@ struct T5Embedder {
             GGML_ASSERT(!out_opt.empty());
             out = std::move(out_opt);
             print_sd_tensor(out);
-            LOG_DEBUG("t5 test done in %lldms", t1 - t0);
+            LOG_VERBOSE("t5 test done in %lldms", t1 - t0);
         }
     }
 
@@ -626,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;
@@ -642,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,