fix: handle invalid option numbers (#1961)

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

.github/workflows/build.yml

@@ -23,6 +23,7 @@ on:
         "**/*.c",
         "**/*.cpp",
         "**/*.cu",
+        "ggml",
         "examples/server/frontend",
         "examples/server/frontend/**",
       ]
@@ -40,6 +41,7 @@ on:
         "**/*.c",
         "**/*.cpp",
         "**/*.cu",
+        "ggml",
         "examples/server/frontend",
         "examples/server/frontend/**",
       ]

.gitignore

@@ -1,6 +1,7 @@
 build*/
 cmake-build-*/
 test/
+tests/
 .vscode/
 .idea/
 .cache/

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"
@@ -290,6 +292,10 @@ else()
     add_library(${SD_LIB} STATIC ${SD_LIB_SOURCES})
 endif()
 
+if(MSVC)
+    target_compile_options(${SD_LIB} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/bigobj>)
+endif()
+
 if(APPLE)
     sd_set_macos_rpaths(${SD_LIB})
 endif()

README.md

@@ -15,6 +15,7 @@ API and command-line option may change frequently.***
 
 ## 🔥Important News
 
+* **2026/08/20** 🚀 stable-diffusion.cpp now supports **LTX-2.5**
 * **2026/08/04** 🚀 stable-diffusion.cpp adds **Day-1 support for MiniMax-H3**
 * **2026/06/25** 🚀 stable-diffusion.cpp now supports **Krea2**
 * **2026/06/04** 🚀 stable-diffusion.cpp now supports **Ideogram4**
@@ -68,7 +69,7 @@ API and command-line option may change frequently.***
   - Video Models
     - [Wan2.1/Wan2.2](./docs/wan.md)
     - [MiniMax-H3](./docs/minimax_h3.md)
-    - [LTX-2.3](./docs/ltx2.md)
+    - [LTX-2.3/LTX-2.5](./docs/ltx2.md)
     - [HunyuanVideo 1.5](./docs/hunyuan_video.md)
     - [LingBot-Video](./docs/lingbot_video.md)
   - [PhotoMaker](./docs/photo_maker.md) support.
@@ -152,6 +153,7 @@ For runtime and parameter backend placement, see the [backend selection guide](.
 - [LCM/LCM-LoRA](./docs/lcm.md)
 - [Docker](./docs/docker.md)
 - [Quantization and GGUF](./docs/quantization_and_gguf.md)
+- [INT8 convrot safetensors](./docs/int8_convrot.md)
 - [Inference acceleration via caching](./docs/caching.md)
 
 ## Bindings

docs/animatediff.md

@@ -79,7 +79,7 @@ Low-VRAM streaming (verified with a 2 GiB cap on RTX 3060):
 .\bin\Release\sd-cli.exe -M vid_gen \
     --model ..\models\checkpoints\realisticVisionV60B1.safetensors \
     --motion-module ..\models\animatediff\mm_sd15_v3.safetensors \
-    --max-vram 2.0 --stream-layers --diffusion-fa \
+    --max-vram 2.0 --diffusion-fa \
     -p "photo of coastline, rocks, storm weather, wind, waves, lightning" \
     --cfg-scale 8.0 --sampling-method euler --scheduler discrete \
     -H 384 -W 384 --video-frames 8 --fps 8 --steps 20 -s 42 \

docs/backend.md

@@ -41,7 +41,11 @@ sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend disk
 sd-cli -m model.safetensors -p "a cat" --backend diffusion=cuda0,vae=vulkan0 --max-vram cuda0=6,vulkan0=2
 ```
 
-The budget applies to every module running on that backend.
+The value is a shared per-device budget for managed weights and registered
+runner compute/cache buffers. Live free memory can lower the effective limit
+for each graph run. Driver contexts and allocations made outside the managed
+model runners are not part of this accounting, so it is not a hard physical
+VRAM cap.
 
 Module names are case-insensitive. Hyphens and underscores in module names are ignored, so `clip_vision`, `clip-vision`, and `clipvision` are equivalent.
 
@@ -79,9 +83,10 @@ with `--params-backend diffusion=disk`, released directly from) its own device;
 an explicit assignment such as `te=cpu` keeps the parameters on that backend
 and stages each range to its device on demand.
 
-Layer split cannot be combined with `--max-vram` graph-cut segmentation or
-`--stream-layers` for the split module; those are single-device mechanisms and
-are disabled for it.
+Layer split uses the fixed graph-cut plan to assign blocks across devices, but
+single-device segmented execution and next-segment prefetch are disabled for
+the split module. `--max-vram` can still provide the per-device limits used by
+layer split and auto-fit.
 
 Use `--list-devices` to see the device names available on the system.
 
@@ -104,43 +109,88 @@ Compared to a layer split this uses all GPUs within every layer (instead of
 sequentially device by device) at the cost of a cross-device reduction per
 matmul - usually the faster option when the devices have fast interconnect.
 
-Row split requires backend support for split buffers and is currently
-available on CUDA only; on other backends (or when the listed devices belong
-to different backend registries) the module falls back to a layer split.
+Row split requires a compatible split-buffer export from the linked GGML
+backend. If it is unavailable (or the listed devices belong to different backend
+registries), the module falls back to a layer split.
 Embeddings, normalization weights, biases and other non-block tensors stay in
 regular buffers on the main device.
 
+Row-split execution can use graph segments, but split weights are loaded
+synchronously instead of using the normal single-device prefetch path. Because
+GGML does not expose exact shard allocation sizes, the managed budget currently
+counts a split buffer's full size on each participating device. This is a
+conservative bound and can reject otherwise feasible layouts.
+
 Direct ("immediately") LoRA application cannot patch row-split tensors; with
 `--split-mode row` the automatic LoRA mode selects runtime application, and an
 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 `--backend` or `--params-backend` assignments disable auto-fit,
+regardless of argument order, even with `--auto-fit on`.
+
+When enabled, auto-fit uses one GPU for `diffusion` / `te` / `vae` computation. It chooses
+the GPU with the largest available memory budget (the first device on a tie),
+then derives parameter placements from the model metadata and the remaining
+memory budgets. The chosen backend specifications are printed.
 
 ```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" --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.
+Components are considered in `diffusion`, `te`, `vae` order so that repeatedly
+used diffusion weights have priority. Each component's weights use the first
+storage location with enough remaining budget:
+
+1. The main GPU, leaving estimated space for computation and weight staging.
+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.
+Other GPUs store weights only: weights are copied to the main GPU for execution.
+Auto-fit does not select multi-GPU layer/row computation, so `--split-mode` does
+not change its placements. Use explicit backend assignments for multi-GPU
+computation.
+
+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
 
@@ -192,7 +242,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:
 
@@ -241,4 +291,7 @@ The example CLI/server still accepts these older CPU placement flags as compatib
 
 Because this default is inserted first, later explicit `--params-backend` entries can still override it, for example `--offload-to-cpu --params-backend te=disk` keeps non-TE parameters on CPU and reloads TE parameters from disk.
 
-Library callers should set `backend` and `params_backend` directly. 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; nonempty `backend` or `params_backend` assignments disable it.
+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/ideogram4.md

@@ -11,7 +11,19 @@
 - Download Qwen3-VL-8B-Instruct
     - gguf: https://huggingface.co/unsloth/Qwen3-VL-8B-Instruct-GGUF/tree/main
 
-## Convert weights
+## Use original FP8 weights
+
+The original Ideogram4 FP8 safetensors can be loaded directly. FP8 tensors stay
+at one byte per element in RAM and VRAM. Backends that cannot multiply FP8
+weights directly cast only the active layer to a temporary BF16 tensor during
+execution; the loader does not expand the entire checkpoint to BF16.
+
+Use `ideogram4_fp8.safetensors` and `ideogram4_uncond_fp8.safetensors` directly
+with `--diffusion-model` and `--uncond-diffusion-model`, respectively.
+
+## Optional conversion for quantization
+
+The following conversion is only needed when creating a quantized GGUF model.
 
 fp8 scale -> bf16
 

docs/int8_convrot.md

@@ -0,0 +1,102 @@
+# INT8 Convrot Safetensors
+
+sd.cpp can load and execute ComfyUI `int8_tensorwise` safetensors with `convrot` metadata directly. The stored INT8 weights are not converted to another weight type at load time.
+
+## Checkpoint format
+
+Each quantized linear module contains the following tensors:
+
+- `<module>.weight`: an I8 weight matrix.
+- `<module>.weight_scale`: one floating-point scale for each output row. ComfyUI's two-dimensional `[out_features, 1]` representation is normalized to a one-dimensional tensor while loading.
+- `<module>.comfy_quant`: a U8 tensor containing the JSON quantization configuration.
+
+A supported configuration has this form:
+
+```json
+{
+  "format": "int8_tensorwise",
+  "convrot": true,
+  "convrot_groupsize": 256
+}
+```
+
+The convrot group size must be a power of four and must divide the input feature dimension. The commonly used configuration is H256, with `convrot_groupsize` set to `256`.
+
+## How INT8 convrot works
+
+Convrot combines an offline rotation of the weights with the same rotation of the activations at runtime. The rotation uses a normalized regular Hadamard matrix constructed recursively from
+
+```text
+     [ 1  1  1 -1 ]
+     [ 1  1 -1  1 ]
+H4 = [ 1 -1  1  1 ] / 2
+     [-1  1  1  1 ]
+```
+
+For a group size `G = 4^n`, the transform is the normalized Kronecker power of `H4`. It is applied independently to every contiguous group of `G` input features. The resulting block-diagonal rotation matrix `R` is orthogonal and symmetric, so `R R^T = I`.
+
+For an original floating-point linear layer
+
+```text
+Y = X W^T + b
+```
+
+the checkpoint stores a rotated weight matrix `W_rot = W R`, quantized per output row. At runtime sd.cpp computes `X_rot = X R`. Ignoring quantization error,
+
+```text
+X_rot W_rot^T = X R (W R)^T = X R R^T W^T = X W^T
+```
+
+The rotation therefore preserves the linear operation. Its purpose is to spread isolated large values across each feature group, reducing the effect of outliers on tensorwise INT8 quantization.
+
+### Weight quantization
+
+The rotated weights are quantized offline with one scale per output row:
+
+```text
+s_w[o]      = max_i(abs(W_rot[o, i])) / 127
+Q_w[o, i]   = clamp(round(W_rot[o, i] / s_w[o]), -127, 127)
+```
+
+`Q_w` is stored in `<module>.weight`, and `s_w` is stored in `<module>.weight_scale`.
+
+### Runtime activation quantization
+
+For every activation row, sd.cpp applies the group-wise Hadamard rotation and then calculates one dynamic scale across the entire rotated row:
+
+```text
+s_x[r]      = max_i(abs(X_rot[r, i])) / 127
+Q_x[r, i]   = clamp(round(X_rot[r, i] / s_x[r]), -127, 127)
+```
+
+The matrix multiplication accumulates into signed 32-bit integers:
+
+```text
+A[r, o] = sum_i(Q_x[r, i] * Q_w[o, i])
+```
+
+The floating-point output is reconstructed as
+
+```text
+Y[r, o] ~= A[r, o] * s_x[r] * s_w[o] + b[o]
+```
+
+The packed runtime activation tensor contains the I8 activation rows and their floating-point row scales. Linear layers that share the same input and convrot group size reuse this packed tensor, avoiding repeated rotation and activation quantization within the graph.
+
+## Backend support
+
+- CPU provides the portable regular Hadamard, activation quantization, INT8 matrix multiplication, and scale restoration implementations.
+- NVIDIA CUDA devices with compute capability 7.5 or newer use the native accelerated path. For H256, CUDA fuses the rotation, row-wise maximum reduction, and activation quantization. It uses cuBLAS for I8 x I8 to I32 GEMM and a CUDA kernel for scale restoration and bias addition.
+- Vulkan and other GPU backends do not currently have dedicated INT8 convrot kernels. They use the backend scheduler to fall back to CPU, which is expected to be substantially slower than the CUDA path.
+
+LoRA adapters are applied at runtime without modifying the INT8 weights. The INT8 convrot path computes the base linear output, while LoRA, LoHa, LoKr, and raw weight-difference adapters compute their output corrections from the original, unrotated activation and add them to the base output. `--lora-apply-mode auto` selects this path for models containing INT8 tensorwise weights. If `immediately` is requested, sd.cpp falls back to runtime application because merging an adapter would require dequantizing and rotating its weight update, then recalculating the per-row scales and requantizing the result.
+
+The dedicated CUDA convrot activation path currently requires a group size of `256`; other supported group sizes use CPU execution.
+
+## Example
+
+ComfyUI INT8 convrot safetensors can be passed to `--diffusion-model` without conversion:
+
+```powershell
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\krea2_turbo_int8_convrot.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\models\vae\wan_2.1_vae.safetensors -p "a lovely cat holding a sign says 'krea2.cpp'" --steps 8 --cfg-scale 1 --diffusion-fa -v --offload-to-cpu
+```

docs/ltx2.md

@@ -1,7 +1,17 @@
 # How to Use
 
+Both LTX-2.3 and LTX-2.5 are supported. The two share a transformer, video VAE and audio
+VAE architecture; LTX-2.5 drops the video FFN biases, adds a learned keyframe
+absolute-position embedding, and pairs with a Gemma 4 text encoder instead of Gemma 3.
+Everything is detected from the weights, so the command lines differ only in which files
+you pass.
+
+# LTX-2.3
+
 ## Download weights
 
+### LTX-2.3
+
 - Download LTX-2.3
     - safetensors: https://huggingface.co/Kijai/LTX2.3_comfy/tree/main/diffusion_models
     - gguf: https://huggingface.co/unsloth/LTX-2.3-GGUF/tree/main
@@ -16,6 +26,30 @@
 - Download LTX spatial latent upscaler
     - safetensors: https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-spatial-upscaler-x2-1.1.safetensors
 
+### LTX-2.5
+
+- Download LTX-2.5
+    - safetensors: https://huggingface.co/Lightricks/LTX-2.5/tree/main/diffusion_models
+    - gguf: https://huggingface.co/vantagewithai/LTX-2.5-GGUF/tree/main
+- Download the text encoder. This is a Gemma 4 12B fine-tuned for LTX with the text
+  projection bundled in, so no separate `--embeddings-connectors` file is needed. Google's
+  stock Gemma 4 is not a substitute.
+    - safetensors: https://huggingface.co/Lightricks/LTX-2.5/blob/main/text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors
+- Download the video vae. Use the **conv** variant: `ltx-2.5-video-vae-conv-bf16.safetensors`.
+  The default `ltx-2.5-video-vae-bf16.safetensors` is a diffusion decoder, which is not
+  implemented here.
+    - safetensors: https://huggingface.co/Lightricks/LTX-2.5/blob/main/vae/ltx-2.5-video-vae-conv-bf16.safetensors
+- Download the audio vae
+    - safetensors: https://huggingface.co/Lightricks/LTX-2.5/blob/main/vae/ltx-2.5-audio-vae-bf16.safetensors
+- Download the LTX spatial latent upscaler
+    - safetensors: https://huggingface.co/Lightricks/LTX-2.5/blob/main/latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors
+
+To run the text encoder quantized, convert it once with sd-cli:
+
+```
+.\bin\Release\sd-cli.exe -M convert -m ..\models\text_encoders\gemma4-12b-with-proj-ltx-2.5-bf16.safetensors --type q8_0 -o ..\models\text_encoders\gemma4-12b-with-proj-ltx-2.5-Q8_0.gguf
+```
+
 ## Examples
 
 ### LTX-2.3 dev T2V
@@ -74,4 +108,32 @@ By default, the hires refine pass uses the main sampler and scheduler, then trim
   src="../assets/ltx2/hires_i2v.webm"
   controls
   muted
-  style="max-width: 100%; height: auto;"></video>
+  style="max-width: 100%; height: auto;"></video>
+
+### LTX-2.5 dev T2V
+
+```
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\ltx-2.5-22b-dev-transformer-Q8_0.gguf --vae ..\models\vae\ltx-2.5-video-vae-conv-bf16.safetensors --audio-vae ..\models\vae\ltx-2.5-audio-vae-bf16.safetensors --llm ..\models\text_encoders\gemma4-12b-with-proj-ltx-2.5-Q8_0.gguf -p "A wide aerial shot of a red vintage convertible driving along a coastal cliff road at sunset, waves crashing below" --cfg-scale 3.0 --sampling-method euler -v -n "worst quality, low quality, blurry, distorted, artifacts" -W 1280 -H 720 --diffusion-fa --offload-to-cpu --video-frames 121 --fps 24 -o t2v.webm
+```
+
+### LTX-2.5 dev I2V
+
+```
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\ltx-2.5-22b-dev-transformer-Q8_0.gguf --vae ..\models\vae\ltx-2.5-video-vae-conv-bf16.safetensors --audio-vae ..\models\vae\ltx-2.5-audio-vae-bf16.safetensors --llm ..\models\text_encoders\gemma4-12b-with-proj-ltx-2.5-Q8_0.gguf -p "a lovely cat blinking slowly, gentle camera push in" --cfg-scale 3.0 --sampling-method euler -v -W 1280 -H 720 --diffusion-fa --offload-to-cpu --video-frames 121 -i ..\assets\ernie_image\turbo_example.png -o i2v.webm
+```
+
+### LTX-2.5 spatial latent upscale
+
+Works exactly like the LTX-2.3 upscaler described below; put
+`ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors` under `--hires-upscalers-dir` and
+pass its name without path or extension to `--hires-upscaler`.
+
+```
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\ltx-2.5-22b-dev-transformer-Q8_0.gguf --vae ..\models\vae\ltx-2.5-video-vae-conv-bf16.safetensors --audio-vae ..\models\vae\ltx-2.5-audio-vae-bf16.safetensors --llm ..\models\text_encoders\gemma4-12b-with-proj-ltx-2.5-Q8_0.gguf --hires-upscalers-dir ..\models\latent_upscale_models --hires-upscaler ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0 --hires --hires-steps 6 -p "a lovely cat" --cfg-scale 3.0 --sampling-method euler -v -W 640 -H 360 --diffusion-fa --offload-to-cpu --video-frames 121 -o hires_t2v.webm
+```
+
+## Not implemented
+
+- The diffusion video decoder (`ltx-2.5-video-vae-bf16.safetensors`). Use the conv VAE.
+- The temporal latent upscaler and the duration head (`--auto-duration`); pass
+  `--video-frames` explicitly.

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

@@ -14,8 +14,12 @@ Run by adding `--diffusion-fa` to the arguments and watch for:
 ```
 and the compute buffer shrink in the debug log:
 ```
-[DEBUG] ggml_extend.hpp:1004 - flux compute buffer size: 650.00 MB(VRAM)
+[DEBUG] ggml_runner.cpp:280 - flux compute buffer size: 650.00 MB(VRAM) on CUDA0 (peak across 1 segment)
 ```
+This reports the actual peak compute workspace capacity per backend, including
+CPU fallback. It excludes weights and cache buffers. Within a runner lifecycle,
+the summary is printed only on the first graph or when backend capacities or the
+segment count change.
 
 ## Offload weights to the CPU to save VRAM without reducing generation speed.
 
@@ -43,7 +47,7 @@ Use disk params to reduce both VRAM and RAM usage:
 --backend cuda0 --params-backend disk
 ```
 
-This reloads parameters from the model file on demand and releases them after use. It has the lowest memory residency, but can be slower because weights must be read again. `disk` is never selected implicitly; set it explicitly when RAM usage matters more than reload cost.
+This reloads parameters from the model file on demand, retains unpinned compute copies while space permits, and releases them under pressure or at module-run completion. It has the lowest source-memory residency, but can be slower because evicted weights must be read again. `disk` is never selected implicitly; set it explicitly when RAM usage matters more than reload cost.
 
 Per-module assignments can target only the largest modules:
 
@@ -53,25 +57,37 @@ Per-module assignments can target only the largest modules:
 
 See [backend selection](./backend.md) for full syntax.
 
-## Run models that don't fit in VRAM (CPU streaming).
+## Run models that don't fit in VRAM (automatic segmented execution).
 
-`--offload-to-cpu` alone keeps every parameter in system RAM and stages it to the runtime backend on first use, then leaves it resident there. If the diffusion model is larger than the runtime backend's free memory (e.g. Flux dev at bf16 on an 8 GiB GPU), that residency stops fitting during the sampling loop and generation fails. Two additional flags make it fit by trading a small amount of speed for room:
+`--offload-to-cpu` keeps the source parameters in system RAM and creates compute-side GPU replicas on demand. Unpinned replicas remain resident for reuse, but automatic graph-cut execution evicts them from the last segment backward when the next weight or compute allocation needs space. Disk-backed parameters follow the same policy without retaining a RAM source copy.
 
-- `--max-vram <GiB>` sets a VRAM budget the graph-cut segmenter respects. It cuts each forward pass into segments sized to fit the budget, running them in sequence and freeing intermediate activations between them. Negative values auto-detect free VRAM and spare the given amount (`--max-vram -1` uses most of the free VRAM and keeps ~1 GiB headroom), a positive value caps the budget, `0` disables segmentation.
-- `--stream-layers` streams the diffusion model's transformer blocks one at a time. Each block's parameters are copied from the CPU to the runtime backend just before it runs and evicted when the residency budget is reached. Prefetching hides most of the copy latency behind compute. This flag only takes effect when the diffusion params backend is CPU, so it must be combined with `--offload-to-cpu` (or an explicit `--params-backend diffusion=cpu`); a warning is logged and the flag is ignored otherwise.
+When a graph has cut markers and its missing weights plus incremental compute workspace exceed the available device headroom, it runs its fixed segment list in order. A reusable monolithic compute buffer is not counted as a new allocation. An explicit `--max-vram` budget deducts already-resident managed weights and compute/cache buffers registered by every runner sharing the device, so later graph runs remain segmented when the full graph exceeds the budget. The current segment's weights are pinned during compute, and the next parameter-bearing segment is prefetched when the device supports asynchronous transfer. No opt-in streaming flag is required.
 
-The three flags stack. The recommended shape for "biggest model my card can host":
+- `--max-vram <GiB>` optionally lowers the live-memory limit. A positive value is a managed per-device budget, `0` uses the device's current free memory without an explicit budget, and a negative value snapshots free memory at startup while reserving that many GiB (`--max-vram -1` reserves about 1 GiB). Driver contexts and unrelated external allocations remain outside the managed budget.
+- `--disable-prefetch` disables asynchronous next-segment prefetch while retaining synchronous loading, eviction, and segmented execution.
+- `--disable-segmented-compute` forces monolithic graph execution for diagnostics or compatibility, even when the automatic memory check would select segments.
+
+Single-device monolithic execution also reclaims unpinned weight replicas before
+loading weights or allocating compute workspace, including graphs without cut
+markers and runs with `--disable-segmented-compute`. It still respects the managed
+device budget and fails if the graph cannot fit after reclamation.
+
+Segment completion releases active workspace use while retaining the runner's
+allocator/scheduler capacity. Compatible gallocr reservations are reused across
+graphs; idle workspaces can be reclaimed under pressure and are freed at runner
+completion. Cross-graph caches survive individual graphs, but cut buffers do not.
+
+The recommended shape for "biggest model my card can host" is:
 
 ```shell
 sd-cli --diffusion-model flux1-dev.safetensors ... \
-       --offload-to-cpu --max-vram -1 --stream-layers
+       --offload-to-cpu --max-vram -1
 ```
 
 - `--offload-to-cpu`: params in RAM, staged as needed.
-- `--max-vram -1`: use most of the free VRAM as the compute budget, spare 1 GiB headroom, let the graph-cut segmenter split each forward pass to fit.
-- `--stream-layers`: on top of the segmenter, stream individual transformer blocks so their weights don't all need to be resident at once.
+- `--max-vram -1`: reserve about 1 GiB from the startup free-memory snapshot; live free memory can still lower the effective limit for every graph.
 
-Ordered from fastest to smallest-VRAM: no flags → `--offload-to-cpu` → `--offload-to-cpu --max-vram <N>` → `--offload-to-cpu --max-vram <N> --stream-layers`. Each step down costs a few percent of throughput to buy more room; combined they can run models roughly 3-4x larger than the raw VRAM would allow.
+Use `--params-backend diffusion=disk` instead when reducing system RAM residency is more important than avoiding repeated model-file reads.
 
 ## Use quantization to reduce memory 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/sd.md

@@ -3,7 +3,7 @@
 - download original weights(.ckpt or .safetensors). For example
     - Stable Diffusion v1.4 from https://huggingface.co/CompVis/stable-diffusion-v-1-4-original
     - Stable Diffusion v1.5 from https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5
-    - Stable Diffuison v2.1 from https://huggingface.co/Manojb/stable-diffusion-2-1-base
+    - Stable Diffusion v2.1 from https://huggingface.co/Manojb/stable-diffusion-2-1-base
     - Stable Diffusion 3 2B from https://huggingface.co/stabilityai/stable-diffusion-3-medium
 
 ### txt2img example

docs/sefi_image.md

@@ -44,7 +44,7 @@ The dispatcher picks `alpha` from the filename (`turbo` substring => 1.0, otherw
 ### 5B (needs streaming on 12 GiB VRAM)
 
 ```
-./build/bin/sd-cli --diffusion-model /path/to/sefi_5b_turbo.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_4b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 1.0 --steps 4 -W 1024 -H 1024 -s 42 --diffusion-fa --max-vram 8 --stream-layers --offload-to-cpu -o out.png
+./build/bin/sd-cli --diffusion-model /path/to/sefi_5b_turbo.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_4b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 1.0 --steps 4 -W 1024 -H 1024 -s 42 --diffusion-fa --max-vram 8 --offload-to-cpu -o out.png
 ```
 
 <img alt="SeFi-Image 5B turbo example" src="../assets/sefi_image/example.png" />

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

examples/cli/main.cpp

@@ -36,12 +36,13 @@ struct SDCliParams {
     SDMode mode             = IMG_GEN;
     std::string output_path = "output.png";
     int output_begin_idx    = -1;
+    int compression_quality = 90;
     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;
@@ -80,7 +81,7 @@ struct SDCliParams {
              &metadata_format},
             {"",
              "--preview-path",
-             "path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp",
+             "path to write preview image to (default: ./preview.png). For image generation, the filename can have %03d placeholder for sequential numbering. Multi-frame previews support .avi, .webm, and animated .webp",
              0,
              &preview_path},
             {"",
@@ -93,12 +94,16 @@ struct SDCliParams {
         options.int_options = {
             {"",
              "--preview-interval",
-             "interval in denoising steps between consecutive updates of the image preview file (default is 1, meaning updating at every step)",
+             "preview interval: in each sampling pass, positive N updates every Nth denoiser step and -N previews only completed logical step N; 0 previews the final completed step of the first pass (base-resolution or high-noise). Default: 1",
              &preview_interval},
             {"",
              "--output-begin-idx",
              "starting index for output image sequence, must be non-negative (default 0 if specified %d in output path, 1 otherwise)",
              &output_begin_idx},
+            {"",
+             "--compression-quality",
+             "compression quality of video and JPEG / WebP images (90 by default)",
+             &compression_quality},
         };
 
         options.bool_options = {
@@ -110,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",
@@ -215,6 +216,7 @@ struct SDCliParams {
              on_imatrix_in_arg},
         };
 
+        add_log_options(options, log_level);
         return options;
     };
 
@@ -264,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"
@@ -302,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) &&
@@ -318,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;
@@ -350,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);
@@ -372,27 +376,6 @@ bool load_images_from_dir(const std::string dir,
     return true;
 }
 
-void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy, void* data) {
-    (void)step;
-    (void)is_noisy;
-    SDCliParams* cli_params = (SDCliParams*)data;
-    // is_noisy is set to true if the preview corresponds to noisy latents, false if it's denoised latents
-    // unused in this app, it will either be always noisy or always denoised here
-    if (frame_count == 1) {
-        if (!write_image_to_file(cli_params->preview_path,
-                                 image->data,
-                                 image->width,
-                                 image->height,
-                                 image->channel)) {
-            LOG_ERROR("save preview image to '%s' failed", cli_params->preview_path.c_str());
-        }
-    } else {
-        if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, cli_params->preview_fps) != 0) {
-            LOG_ERROR("save preview video to '%s' failed", cli_params->preview_path.c_str());
-        }
-    }
-}
-
 std::string format_frame_idx(std::string pattern, int frame_idx) {
     std::smatch match;
     std::string result = pattern;
@@ -412,6 +395,36 @@ std::string format_frame_idx(std::string pattern, int frame_idx) {
     return result;
 }
 
+int continuous_preview_counter = 0;
+
+void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy, void* data) {
+    (void)step;
+    (void)is_noisy;
+    SDCliParams* cli_params = (SDCliParams*)data;
+    // is_noisy is set to true if the preview corresponds to noisy latents, false if it's denoised latents
+    // unused in this app, it will either be always noisy or always denoised here
+    if (frame_count == 1) {
+        fs::path path = cli_params->preview_path;
+        if (encoded_image_format_from_path(path.string()) == EncodedImageFormat::UNKNOWN)
+            path += ".png";
+        if (std::regex_search(path.string(), format_specifier_regex))
+            path = fs::path(format_frame_idx(path.string(), continuous_preview_counter++));
+        if (!write_image_to_file(path.string(),
+                                 image->data,
+                                 image->width,
+                                 image->height,
+                                 image->channel,
+                                 "",
+                                 cli_params->compression_quality)) {
+            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) {
+            LOG_ERROR("save preview video to '%s' failed", cli_params->preview_path.c_str());
+        }
+    }
+}
+
 static fs::path get_video_audio_sidecar_path(const SDCliParams& cli_params) {
     fs::path out_path     = cli_params.output_path;
     fs::path base_path    = out_path;
@@ -486,7 +499,7 @@ bool save_results(const SDCliParams& cli_params,
         std::string params          = gen_params.embed_image_metadata
                                           ? get_image_params(ctx_params, gen_params, metadata_seed, cli_params.mode)
                                           : "";
-        const bool ok               = write_image_to_file(path.string(), img.data, img.width, img.height, img.channel, params, 90);
+        const bool ok               = write_image_to_file(path.string(), img.data, img.width, img.height, img.channel, params, cli_params.compression_quality);
         LOG_INFO("save result image %d to '%s' (%s)", idx, path.string().c_str(), ok ? "success" : "failure");
         return ok;
     };
@@ -532,7 +545,7 @@ bool save_results(const SDCliParams& cli_params,
         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, 90, 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) == 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;
@@ -637,8 +650,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;
@@ -686,11 +697,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) &&
@@ -794,7 +805,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;
             }
@@ -876,8 +887,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;
         }
     }
@@ -888,8 +898,7 @@ int main(int argc, const char* argv[]) {
                                   gen_params.pm_id_images,
                                   0,
                                   0,
-                                  0,
-                                  cli_params.verbose)) {
+                                  0)) {
             return 1;
         }
     }

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;
         }
@@ -502,7 +531,7 @@ ArgOptions SDContextParams::get_options() {
          &rpc_servers},
         {"",
          "--max-vram",
-         "maximum VRAM budget in GiB for graph-cut segmented execution. Accepts a single value or assignments by backend/device, e.g. 6 or cuda0=6,vulkan0=4. 0 disables graph splitting; a negative value auto-detects free VRAM, sparing the specified value",
+         "optional per-device budget in GiB for managed weights and runner buffers during automatic graph-cut execution. Accepts a single value or assignments by backend/device, e.g. 6 or cuda0=6,vulkan0=4. 0 uses live free VRAM without an explicit budget; a negative value reserves that much free VRAM",
          0,
          &max_vram},
     };
@@ -517,19 +546,17 @@ ArgOptions SDContextParams::get_options() {
 
     options.bool_options = {
         {"",
-         "--stream-layers",
-         "enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram; defaults to false)",
-         true, &stream_layers},
+         "--disable-prefetch",
+         "disable asynchronous next-segment weight prefetch (defaults to false)",
+         true, &disable_prefetch},
+        {"",
+         "--disable-segmented-compute",
+         "force monolithic graph execution even when automatic graph cutting is needed (defaults to false)",
+         true, &disable_segmented_compute},
         {"",
          "--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",
@@ -572,6 +599,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;
@@ -643,6 +687,12 @@ ArgOptions SDContextParams::get_options() {
     };
 
     options.manual_options = {
+        {"",
+         "--auto-fit",
+         "on|off (default: on). Use one GPU for diffusion/te/vae computation and place weights on that GPU, "
+         "RAM, another GPU, or disk in that order, according to available memory (--max-vram limits GPU budgets). "
+         "Disabled by explicit --backend or --params-backend; uses automatic graph segmentation when needed",
+         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). "
@@ -831,7 +881,8 @@ std::string SDContextParams::to_string() const {
         << "  sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n"
         << "  offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n"
         << "  max_vram: \"" << max_vram << "\",\n"
-        << "  stream_layers: " << (stream_layers ? "true" : "false") << ",\n"
+        << "  disable_prefetch: " << (disable_prefetch ? "true" : "false") << ",\n"
+        << "  disable_segmented_compute: " << (disable_segmented_compute ? "true" : "false") << ",\n"
         << "  eager_load: " << (eager_load ? "true" : "false") << ",\n"
         << "  backend: \"" << backend << "\",\n"
         << "  params_backend: \"" << params_backend << "\",\n"
@@ -903,7 +954,8 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.force_sdxl_vae_conv_scale       = force_sdxl_vae_conv_scale;
     sd_ctx_params.vae_format                      = str_to_vae_format(vae_format);
     sd_ctx_params.max_vram                        = max_vram.c_str();
-    sd_ctx_params.stream_layers                   = stream_layers;
+    sd_ctx_params.disable_prefetch                = disable_prefetch;
+    sd_ctx_params.disable_segmented_compute       = disable_segmented_compute;
     sd_ctx_params.eager_load                      = eager_load;
     sd_ctx_params.backend                         = effective_backend.c_str();
     sd_ctx_params.params_backend                  = effective_params_backend.c_str();
@@ -1008,12 +1060,12 @@ ArgOptions SDGenerationParams::get_options() {
          &hires_upscaler},
         {"",
          "--extra-sample-args",
-         "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_divisions",
+         "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_max_order, lms_shift, lms_divisions",
          (int)',',
          &extra_sample_args},
         {"",
          "--extra-tiling-args",
-         "extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
+         "extra VAE tiling args, key=value list. Supported video VAEs accept temporal_tile_frames/temporal_tile_size (default: 4), temporal_tile_overlap (default: 1)",
          (int)',',
          &extra_tiling_args},
         {"",
@@ -1230,7 +1282,7 @@ ArgOptions SDGenerationParams::get_options() {
          &vae_tiling_params.enabled},
         {"",
          "--temporal-tiling",
-         "enable temporal tiling for LTX video VAE decode",
+         "enable temporal tiling for supported video VAE decode",
          true,
          &vae_tiling_params.temporal_tiling},
         {"",
@@ -1555,6 +1607,16 @@ ArgOptions SDGenerationParams::get_options() {
         return 1;
     };
 
+    std::string sample_methods = sample_method_to_str[0];
+    for (int i = 1; i < SAMPLE_METHOD_COUNT; i++) {
+        sample_methods += ", " + std::string(sample_method_to_str[i]);
+    }
+
+    std::string schedulers = scheduler_to_str[0];
+    for (int i = 1; i < SCHEDULER_COUNT; i++) {
+        schedulers += ", " + std::string(scheduler_to_str[i]);
+    }
+
     options.manual_options = {
         {"-s",
          "--seed",
@@ -1562,17 +1624,18 @@ ArgOptions SDGenerationParams::get_options() {
          on_seed_arg},
         {"",
          "--sampling-method",
-         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp, lms]"
-         "(default: euler for Flux/SD3/Wan, euler_a otherwise)",
+         "sampling method, one of [" + sample_methods + "], "
+                                                        "default: euler for Flux/SD3/Wan, euler_a otherwise",
          on_sample_method_arg},
         {"",
          "--high-noise-sampling-method",
-         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp, lms]"
-         " default: euler for Flux/SD3/Wan, euler_a otherwise",
+         "(high noise) sampling method, one of [" + sample_methods + "], "
+                                                                     "default: euler for Flux/SD3/Wan, euler_a otherwise",
          on_high_noise_sample_method_arg},
         {"",
          "--scheduler",
-         "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal, flux2, flux, beta], alias: normal=discrete, default: model-specific",
+         "denoiser sigma scheduler, one of [" + schedulers + "], "
+                                                             "alias: normal=discrete, default: model-specific",
          on_scheduler_arg},
         {"",
          "--sigmas",

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,
@@ -146,17 +147,18 @@ struct SDContextParams {
     std::map<std::string, std::string> embedding_map;
     std::vector<sd_embedding_t> embedding_vec;
 
-    rng_type_t rng_type         = CUDA_RNG;
-    rng_type_t sampler_rng_type = RNG_TYPE_COUNT;
-    bool offload_params_to_cpu  = false;
-    std::string max_vram        = "0";
-    bool stream_layers          = false;
-    bool eager_load             = false;
+    rng_type_t rng_type            = CUDA_RNG;
+    rng_type_t sampler_rng_type    = RNG_TYPE_COUNT;
+    bool offload_params_to_cpu     = false;
+    std::string max_vram           = "0";
+    bool disable_prefetch          = false;
+    bool disable_segmented_compute = false;
+    bool eager_load                = false;
     std::string backend;
     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;

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,10 +116,11 @@ 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);
     fflush(out_stream);
 }
@@ -109,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

@@ -835,6 +835,9 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
     const uint32_t audio_byte_rate       = has_audio ? static_cast<uint32_t>(audio->sample_rate * audio_block_align) : 0;
     const uint32_t audio_data_size       = has_audio ? static_cast<uint32_t>(audio_pcm.size()) : 0;
 
+    if (mjpg_quality != 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);
 

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

examples/server/api.md

@@ -518,7 +518,8 @@ Shared default fields used by both `img_gen` and `vid_gen`:
 | `output_format` | `string` |
 | `output_compression` | `integer` |
 
-`vae_tiling_params.extra_tiling_args` accepts a key=value list. For LTX video VAE temporal tiling, `temporal_tile_frames` defaults to `4` and `temporal_tile_overlap` defaults to `1`.
+`vae_tiling_params.extra_tiling_args` accepts a key=value list. Supported video VAEs accept `temporal_tile_frames` (alias `temporal_tile_size`, default `4`) and `temporal_tile_overlap` (default `1`).
+LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEHV use overlap blending. MiniMax H3 keeps its model-specific fixed temporal windows because its latent-to-frame mapping is non-linear.
 
 `img_gen`-specific default fields:
 

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

@@ -60,6 +60,8 @@ enum sample_method_t {
     SAMPLE_METHOD_COUNT
 };
 
+extern SD_API const char* sample_method_to_str[];
+
 enum scheduler_t {
     DISCRETE_SCHEDULER,
     KARRAS_SCHEDULER,
@@ -80,6 +82,8 @@ enum scheduler_t {
     SCHEDULER_COUNT
 };
 
+extern SD_API const char* scheduler_to_str[];
+
 enum prediction_t {
     EPS_PRED,
     V_PRED,
@@ -132,14 +136,18 @@ enum sd_type_t {
     // SD_TYPE_IQ4_NL_4_4 = 36,
     // SD_TYPE_IQ4_NL_4_8 = 37,
     // SD_TYPE_IQ4_NL_8_8 = 38,
-    SD_TYPE_MXFP4 = 39,  // MXFP4 (1 block)
-    SD_TYPE_NVFP4 = 40,  // NVFP4 (4 blocks, E4M3 scale)
-    SD_TYPE_Q1_0  = 41,
-    SD_TYPE_COUNT = 42,
+    SD_TYPE_MXFP4   = 39,  // MXFP4 (1 block)
+    SD_TYPE_NVFP4   = 40,  // NVFP4 (4 blocks, E4M3 scale)
+    SD_TYPE_Q1_0    = 41,
+    SD_TYPE_Q2_0    = 42,
+    SD_TYPE_F8_E4M3 = 43,
+    SD_TYPE_F8_E5M2 = 44,
+    SD_TYPE_COUNT   = 45,
 };
 
 enum sd_log_level_t {
     SD_LOG_DEBUG,
+    SD_LOG_VERBOSE,
     SD_LOG_INFO,
     SD_LOG_WARN,
     SD_LOG_ERROR
@@ -222,8 +230,8 @@ typedef struct {
     bool vae_conv_direct;
     bool force_sdxl_vae_conv_scale;
     enum sd_vae_format_t vae_format;
-    const char* max_vram;  // GiB budget or backend assignment spec for graph-cut segmented param offload (0 = disabled, -1 = auto)
-    bool stream_layers;  // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram)
+    const char* max_vram;  // Optional per-device GiB budget for managed weights and runner buffers; 0 uses live free VRAM without an explicit budget
+    bool disable_prefetch;  // Disable asynchronous next-segment weight prefetch
     bool eager_load;  // Load all params into the params backend at model-load time instead of lazily on first use
     const char* backend;
     const char* params_backend;
@@ -231,6 +239,7 @@ typedef struct {
     bool auto_fit;
     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
 } sd_ctx_params_t;
 
 typedef struct {
@@ -439,6 +448,9 @@ typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void*
 
 SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data);
 SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data);
+// In each sampling pass, a positive interval previews every Nth denoiser step, while a
+// negative interval previews only completed logical step -interval. Zero previews the final
+// completed step of the first sampling pass (base-resolution or high-noise).
 SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data);
 SD_API void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data);
 SD_API int32_t sd_get_num_physical_cores();

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 \

scripts/merge_safetensors.py

@@ -23,16 +23,16 @@ from typing import BinaryIO
 # Configuration
 # -----------------------------------------------------------------------------
 
-OUTPUT_PATH = Path(r"..\models\diffusion_models\minimax_h3_ref2va_pruned_bf16.safetensors")
+OUTPUT_PATH = Path(r".minimax_h3_fl2va_pruned_bf16.safetensors")
 
 SOURCE_RULES = [
     {
-        "path": Path(r"..\models\diffusion_models\minimax_h3_ref2va_bf16.safetensors"),
+        "path": Path(r".minimax_h3_fl2va_bf16.safetensors"),
         "include": [r".*"],
         "exclude": [r".*adaln_proj\.linear.*", r"time_embedder.*"],
     },
     {
-        "path": Path(r"..\models\diffusion_models\minimax_h3_ref2va_pruned_int8_convrot.safetensors"),
+        "path": Path(r".minimax_h3_fl2va_pruned_int8_convrot.safetensors"),
         "include": [r"^.*adaln_proj\.linear.*", "adaln_t_table"],
         "exclude": [],
     },

src/conditioning/conditioner.hpp

@@ -1,11 +1,13 @@
 #ifndef __SD_CONDITIONING_CONDITIONER_HPP__
 #define __SD_CONDITIONING_CONDITIONER_HPP__
 
+#include <cinttypes>
 #include <cmath>
 #include <iomanip>
 #include <limits>
 #include <optional>
 #include <sstream>
+#include "core/ggml_tensor_utils.h"
 
 #include "core/tensor_ggml.hpp"
 #include "core/util.h"
@@ -142,14 +144,13 @@ public:
     virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors)           = 0;
     virtual void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {}
     virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
-    virtual void set_stream_layers_enabled(bool enabled) {}
     virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
     virtual void set_graph_cut_layer_split_enabled(bool enabled) {}
     virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {}
     virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
     virtual void set_flash_attention_enabled(bool enabled) = 0;
     virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
-    virtual void runner_done() {}
+    virtual void runner_end() {}
 };
 
 // ldm.modules.encoders.modules.FrozenCLIPEmbedder
@@ -161,9 +162,9 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
     std::shared_ptr<CLIPTextModelRunner> text_model2;
 
     std::map<std::string, std::string> embedding_map;
-    int32_t num_custom_embeddings   = 0;
-    int32_t num_custom_embeddings_2 = 0;
+    int32_t num_custom_embeddings = 0;
     std::vector<uint8_t> token_embed_custom;
+    std::vector<uint8_t> token_embed_custom2;
     std::map<std::string, std::pair<int, int>> embedding_pos_map;
 
     FrozenCLIPEmbedderWithCustomWords(ggml_backend_t backend,
@@ -202,13 +203,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         }
     }
 
-    void set_stream_layers_enabled(bool enabled) override {
-        text_model->set_stream_layers_enabled(enabled);
-        if (sd_version_is_sdxl(version)) {
-            text_model2->set_stream_layers_enabled(enabled);
-        }
-    }
-
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
         text_model->set_runtime_backends(backends);
         if (sd_version_is_sdxl(version)) {
@@ -244,10 +238,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         }
     }
 
-    void runner_done() override {
-        text_model->runner_done();
+    void runner_end() override {
+        text_model->runner_end();
         if (sd_version_is_sdxl(version)) {
-            text_model2->runner_done();
+            text_model2->runner_end();
         }
     }
 
@@ -257,74 +251,93 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
             LOG_ERROR("embedding '%s' failed", embd_name.c_str());
             return false;
         }
-        auto iter = embedding_pos_map.find(embd_name);
-        if (iter != embedding_pos_map.end()) {
-            LOG_DEBUG("embedding already read in: %s", embd_name.c_str());
-            for (int i = iter->second.first; i < iter->second.second; i++) {
+        auto push_ids = [&](int pos_start, int pos_end, bool cached) {
+            for (int i = pos_start; i < pos_end; i++) {
                 bpe_tokens.push_back(text_model->model.vocab_size + i);
             }
+            if (!cached) {
+                LOG_VERBOSE("embedding '%s' applied: %i token(s), custom embeddings: %i", embd_name.c_str(), pos_end - pos_start, num_custom_embeddings);
+            }
+        };
+        auto iter = embedding_pos_map.find(embd_name);
+        if (iter != embedding_pos_map.end()) {
+            LOG_VERBOSE("embedding already read in: %s", embd_name.c_str());
+            push_ids(iter->second.first, iter->second.second, true);
             return true;
         }
         ggml_init_params params;
-        params.mem_size        = 100 * 1024 * 1024;  // max for custom embeddings 100 MB
-        params.mem_buffer      = nullptr;
-        params.no_alloc        = false;
-        ggml_context* embd_ctx = ggml_init(params);
-        ggml_tensor* embd      = nullptr;
-        ggml_tensor* embd2     = nullptr;
-        auto on_load           = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) {
-            if (tensor_storage.ne[0] != text_model->model.hidden_size) {
-                if (text_model2) {
-                    if (tensor_storage.ne[0] == text_model2->model.hidden_size) {
-                        embd2       = ggml_new_tensor_2d(embd_ctx, tensor_storage.type, text_model2->model.hidden_size, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1);
-                        *dst_tensor = embd2;
-                    } else {
-                        LOG_DEBUG("embedding wrong hidden size, got %i, expected %i or %i", tensor_storage.ne[0], text_model->model.hidden_size, text_model2->model.hidden_size);
-                        return false;
-                    }
-                } else {
-                    LOG_DEBUG("embedding wrong hidden size, got %i, expected %i", tensor_storage.ne[0], text_model->model.hidden_size);
+        params.mem_size       = 100 * 1024 * 1024;  // max for custom embeddings 100 MB
+        params.mem_buffer     = nullptr;
+        params.no_alloc       = false;
+        auto ggml_ctx_deleter = [](ggml_context* ctx) { ggml_free(ctx); };
+        auto embd_ctx         = std::unique_ptr<ggml_context, decltype(ggml_ctx_deleter)>(ggml_init(params), ggml_ctx_deleter);
+        if (!embd_ctx.get()) {
+            LOG_ERROR("ggml_init failed when loading embeddings file");
+            return false;
+        }
+        ggml_tensor* embd    = nullptr;
+        ggml_tensor* embd2   = nullptr;
+        ggml_type embd_type  = text_model->model.get_token_embed_weight()->type;
+        ggml_type embd2_type = text_model2 ? text_model2->model.get_token_embed_weight()->type : embd_type;
+        int64_t hidden_size  = text_model->model.hidden_size;
+        int64_t hidden_size2 = text_model2 ? text_model2->model.hidden_size : 0;
+        auto on_load         = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) {
+            if (tensor_storage.ne[0] == hidden_size) {
+                embd = ggml_new_tensor_2d(embd_ctx.get(), embd_type, hidden_size, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1);
+                if (embd == nullptr) {
                     return false;
                 }
-            } else {
-                embd        = ggml_new_tensor_2d(embd_ctx, tensor_storage.type, text_model->model.hidden_size, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1);
                 *dst_tensor = embd;
+            } else if (text_model2) {
+                if (tensor_storage.ne[0] == hidden_size2) {
+                    embd2 = ggml_new_tensor_2d(embd_ctx.get(), embd2_type, hidden_size2, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1);
+                    if (embd2 == nullptr) {
+                        return false;
+                    }
+                    *dst_tensor = embd2;
+                } else {
+                    LOG_VERBOSE("embedding skipped, wrong hidden size, got %i, expected %i or %i", tensor_storage.ne[0], hidden_size, hidden_size2);
+                }
+            } else {
+                LOG_VERBOSE("embedding skipped, wrong hidden size, got %i, expected %i", tensor_storage.ne[0], hidden_size);
             }
             return true;
         };
         model_loader.set_n_threads(1);
-        model_loader.load_tensors(on_load);
-        int pos_start = num_custom_embeddings;
-        if (embd) {
-            int64_t hidden_size = text_model->model.hidden_size;
-            token_embed_custom.resize(token_embed_custom.size() + ggml_nbytes(embd));
-            memcpy((void*)(token_embed_custom.data() + num_custom_embeddings * hidden_size * ggml_type_size(embd->type)),
-                   embd->data,
-                   ggml_nbytes(embd));
-            for (int i = 0; i < embd->ne[1]; i++) {
-                bpe_tokens.push_back(text_model->model.vocab_size + num_custom_embeddings);
-                // LOG_DEBUG("new custom token: %i", text_model.vocab_size + num_custom_embeddings);
-                num_custom_embeddings++;
-            }
-            LOG_DEBUG("embedding '%s' applied, custom embeddings: %i", embd_name.c_str(), num_custom_embeddings);
-        }
-        if (embd2) {
-            int64_t hidden_size = text_model2->model.hidden_size;
-            token_embed_custom.resize(token_embed_custom.size() + ggml_nbytes(embd2));
-            memcpy((void*)(token_embed_custom.data() + num_custom_embeddings_2 * hidden_size * ggml_type_size(embd2->type)),
-                   embd2->data,
-                   ggml_nbytes(embd2));
-            for (int i = 0; i < embd2->ne[1]; i++) {
-                bpe_tokens.push_back(text_model2->model.vocab_size + num_custom_embeddings_2);
-                // LOG_DEBUG("new custom token: %i", text_model.vocab_size + num_custom_embeddings);
-                num_custom_embeddings_2++;
-            }
-            LOG_DEBUG("embedding '%s' applied, custom embeddings: %i (text model 2)", embd_name.c_str(), num_custom_embeddings_2);
-        }
-        int pos_end = num_custom_embeddings;
-        if (pos_end == pos_start) {
+        if (!model_loader.load_tensors(on_load)) {
+            LOG_ERROR("embedding '%s' failed", embd_name.c_str());
             return false;
         }
+        if (!embd && !embd2) {
+            LOG_WARN("embedding '%s' has no usable tensor", embd_name.c_str());
+            return false;
+        }
+        int pos_start      = num_custom_embeddings;
+        int64_t embd_rows  = embd ? embd->ne[1] : 0;
+        int64_t embd2_rows = embd2 ? embd2->ne[1] : 0;
+        if (embd_rows < embd2_rows) {
+            LOG_WARN("embedding '%s' has fewer rows for text model 1, zero-padding", embd_name.c_str());
+        } else if (text_model2 && embd2_rows < embd_rows) {
+            LOG_WARN("embedding '%s' has fewer rows for text model 2, zero-padding", embd_name.c_str());
+        }
+        int64_t rows      = std::max(embd_rows, embd2_rows);
+        size_t embd_bytes = hidden_size * ggml_type_size(embd_type);
+        token_embed_custom.resize(token_embed_custom.size() + embd_bytes * rows);
+        if (embd) {
+            memcpy((void*)(token_embed_custom.data() + embd_bytes * num_custom_embeddings),
+                   embd->data, embd_bytes * embd_rows);
+        }
+        if (text_model2) {
+            size_t embd2_bytes = hidden_size2 * ggml_type_size(embd2_type);
+            token_embed_custom2.resize(token_embed_custom2.size() + embd2_bytes * rows);
+            if (embd2) {
+                memcpy((void*)(token_embed_custom2.data() + embd2_bytes * num_custom_embeddings),
+                       embd2->data, embd2_bytes * embd2_rows);
+            }
+        }
+        num_custom_embeddings += (int)rows;
+        int pos_end = num_custom_embeddings;
+        push_ids(pos_start, pos_end, false);
         embedding_pos_map[embd_name] = std::pair{pos_start, pos_end};
         return true;
     }
@@ -368,7 +381,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                 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());
         }
 
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
@@ -389,7 +402,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                 size_t padding_size = (75 - (current_size % 75)) % 75;  // Ensure no negative padding
 
                 if (padding_size > 0) {
-                    LOG_DEBUG("BREAK token encountered, padding current chunk by %zu tokens.", padding_size);
+                    LOG_VERBOSE("BREAK token encountered, padding current chunk by %zu tokens.", padding_size);
                     tokens.insert(tokens.end(), padding_size, tokenizer.EOS_TOKEN_ID);
                     weights.insert(weights.end(), padding_size, 1.0f);
                 }
@@ -461,21 +474,17 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                                                                max_token_idx,
                                                                false,
                                                                clip_skip,
-                                                               false,
-                                                               true,
-                                                               true);
+                                                               false);
                 GGML_ASSERT(!chunk_hidden_states.empty());
                 if (sd_version_is_sdxl(version)) {
                     auto chunk_hidden_states2 = text_model2->compute(n_threads,
                                                                      input_ids2,
                                                                      num_custom_embeddings,
-                                                                     token_embed_custom.data(),
+                                                                     token_embed_custom2.data(),
                                                                      max_token_idx,
                                                                      false,
                                                                      clip_skip,
-                                                                     false,
-                                                                     true,
-                                                                     true);
+                                                                     false);
                     GGML_ASSERT(!chunk_hidden_states2.empty());
                     chunk_hidden_states = sd::ops::concat(chunk_hidden_states, chunk_hidden_states2, 0);
 
@@ -483,18 +492,16 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                         pooled = text_model2->compute(n_threads,
                                                       input_ids2,
                                                       num_custom_embeddings,
-                                                      token_embed_custom.data(),
+                                                      token_embed_custom2.data(),
                                                       max_token_idx,
                                                       true,
                                                       clip_skip,
-                                                      false,
-                                                      true,
-                                                      true);
+                                                      false);
                         GGML_ASSERT(!pooled.empty());
                     }
                 }
                 int64_t t1 = ggml_time_ms();
-                LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
+                LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
 
                 chunk_hidden_states = apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
 
@@ -608,7 +615,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(pixel_values, return_pooled, clip_skip);
         };
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
+        return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
     }
 };
 
@@ -675,18 +682,6 @@ struct SD3CLIPEmbedder : public Conditioner {
         }
     }
 
-    void set_stream_layers_enabled(bool enabled) override {
-        if (clip_l) {
-            clip_l->set_stream_layers_enabled(enabled);
-        }
-        if (clip_g) {
-            clip_g->set_stream_layers_enabled(enabled);
-        }
-        if (t5) {
-            t5->set_stream_layers_enabled(enabled);
-        }
-    }
-
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
         if (clip_l) {
             clip_l->set_runtime_backends(backends);
@@ -753,15 +748,15 @@ struct SD3CLIPEmbedder : public Conditioner {
         }
     }
 
-    void runner_done() override {
+    void runner_end() override {
         if (clip_l) {
-            clip_l->runner_done();
+            clip_l->runner_end();
         }
         if (clip_g) {
-            clip_g->runner_done();
+            clip_g->runner_end();
         }
         if (t5) {
-            t5->runner_done();
+            t5->runner_end();
         }
     }
 
@@ -778,7 +773,7 @@ struct SD3CLIPEmbedder : public Conditioner {
                 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());
         }
 
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
@@ -881,9 +876,7 @@ struct SD3CLIPEmbedder : public Conditioner {
                                                         max_token_idx,
                                                         false,
                                                         clip_skip,
-                                                        false,
-                                                        true,
-                                                        true);
+                                                        false);
                 GGML_ASSERT(!chunk_hidden_states_l.empty());
                 chunk_hidden_states_l = ::apply_token_weights(std::move(chunk_hidden_states_l), chunk_weights);
 
@@ -897,9 +890,7 @@ struct SD3CLIPEmbedder : public Conditioner {
                                                     max_token_idx,
                                                     true,
                                                     clip_skip,
-                                                    false,
-                                                    true,
-                                                    true);
+                                                    false);
                     GGML_ASSERT(!pooled_l.empty());
                 }
             } else {
@@ -928,9 +919,7 @@ struct SD3CLIPEmbedder : public Conditioner {
                                                         max_token_idx,
                                                         false,
                                                         clip_skip,
-                                                        false,
-                                                        true,
-                                                        true);
+                                                        false);
                 GGML_ASSERT(!chunk_hidden_states_g.empty());
                 chunk_hidden_states_g = ::apply_token_weights(std::move(chunk_hidden_states_g), chunk_weights);
 
@@ -944,9 +933,7 @@ struct SD3CLIPEmbedder : public Conditioner {
                                                     max_token_idx,
                                                     true,
                                                     clip_skip,
-                                                    false,
-                                                    true,
-                                                    true);
+                                                    false);
                     GGML_ASSERT(!pooled_g.empty());
                 }
             } else {
@@ -969,9 +956,7 @@ struct SD3CLIPEmbedder : public Conditioner {
                 chunk_hidden_states_t5 = t5->compute(n_threads,
                                                      input_ids,
                                                      sd::Tensor<float>(),
-                                                     false,
-                                                     true,
-                                                     true);
+                                                     false);
                 GGML_ASSERT(!chunk_hidden_states_t5.empty());
                 chunk_hidden_states_t5 = ::apply_token_weights(std::move(chunk_hidden_states_t5), chunk_weights);
             } else {
@@ -996,7 +981,7 @@ struct SD3CLIPEmbedder : public Conditioner {
             }
 
             int64_t t1 = ggml_time_ms();
-            LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
+            LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
             if (zero_out_masked) {
                 chunk_hidden_states.fill_(0.0f);
             }
@@ -1079,15 +1064,6 @@ struct FluxCLIPEmbedder : public Conditioner {
         }
     }
 
-    void set_stream_layers_enabled(bool enabled) override {
-        if (clip_l) {
-            clip_l->set_stream_layers_enabled(enabled);
-        }
-        if (t5) {
-            t5->set_stream_layers_enabled(enabled);
-        }
-    }
-
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
         if (clip_l) {
             clip_l->set_runtime_backends(backends);
@@ -1139,12 +1115,12 @@ struct FluxCLIPEmbedder : public Conditioner {
         }
     }
 
-    void runner_done() override {
+    void runner_end() override {
         if (clip_l) {
-            clip_l->runner_done();
+            clip_l->runner_end();
         }
         if (t5) {
-            t5->runner_done();
+            t5->runner_end();
         }
     }
 
@@ -1160,7 +1136,7 @@ struct FluxCLIPEmbedder : public Conditioner {
                 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());
         }
 
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
@@ -1247,9 +1223,7 @@ struct FluxCLIPEmbedder : public Conditioner {
                                              max_token_idx,
                                              true,
                                              clip_skip,
-                                             false,
-                                             true,
-                                             true);
+                                             false);
                     GGML_ASSERT(!pooled.empty());
                 } else {
                     pooled = sd::Tensor<float>::zeros({768});
@@ -1268,9 +1242,7 @@ struct FluxCLIPEmbedder : public Conditioner {
                 chunk_hidden_states = t5->compute(n_threads,
                                                   input_ids,
                                                   sd::Tensor<float>(),
-                                                  false,
-                                                  true,
-                                                  true);
+                                                  false);
                 GGML_ASSERT(!chunk_hidden_states.empty());
                 chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
                 if (zero_out_masked) {
@@ -1281,7 +1253,7 @@ struct FluxCLIPEmbedder : public Conditioner {
             }
 
             int64_t t1 = ggml_time_ms();
-            LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
+            LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
             if (!hidden_states.empty()) {
                 hidden_states = sd::ops::concat(hidden_states, chunk_hidden_states, 1);
             } else {
@@ -1366,12 +1338,6 @@ struct T5CLIPEmbedder : public Conditioner {
         }
     }
 
-    void set_stream_layers_enabled(bool enabled) override {
-        if (t5) {
-            t5->set_stream_layers_enabled(enabled);
-        }
-    }
-
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
         if (t5) {
             t5->set_runtime_backends(backends);
@@ -1408,9 +1374,9 @@ struct T5CLIPEmbedder : public Conditioner {
         }
     }
 
-    void runner_done() override {
+    void runner_end() override {
         if (t5) {
-            t5->runner_done();
+            t5->runner_end();
         }
     }
 
@@ -1426,7 +1392,7 @@ struct T5CLIPEmbedder : public Conditioner {
                 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());
         }
 
         auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
@@ -1467,7 +1433,7 @@ struct T5CLIPEmbedder : public Conditioner {
                 ++num_pad;
             }
         }
-        // LOG_DEBUG("PAD: %d", num_pad);
+        // LOG_VERBOSE("PAD: %d", num_pad);
     }
 
     SDCondition get_learned_condition_common(int n_threads,
@@ -1508,9 +1474,7 @@ struct T5CLIPEmbedder : public Conditioner {
             auto chunk_hidden_states = t5->compute(n_threads,
                                                    input_ids,
                                                    t5_attn_mask_chunk,
-                                                   false,
-                                                   true,
-                                                   true);
+                                                   false);
             GGML_ASSERT(!chunk_hidden_states.empty());
             chunk_hidden_states = apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
 
@@ -1521,7 +1485,7 @@ struct T5CLIPEmbedder : public Conditioner {
             }
 
             int64_t t1 = ggml_time_ms();
-            LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
+            LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
 
             if (!hidden_states.empty()) {
                 hidden_states = sd::ops::concat(hidden_states, chunk_hidden_states, 1);
@@ -1582,12 +1546,6 @@ struct MiniT2IConditioner : public Conditioner {
         }
     }
 
-    void set_stream_layers_enabled(bool enabled) override {
-        if (t5) {
-            t5->set_stream_layers_enabled(enabled);
-        }
-    }
-
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
         if (t5) {
             t5->set_runtime_backends(backends);
@@ -1624,9 +1582,9 @@ struct MiniT2IConditioner : public Conditioner {
         }
     }
 
-    void runner_done() override {
+    void runner_end() override {
         if (t5) {
-            t5->runner_done();
+            t5->runner_end();
         }
     }
 
@@ -1657,9 +1615,7 @@ struct MiniT2IConditioner : public Conditioner {
         sd::Tensor<float> hidden_states = t5->compute(n_threads,
                                                       input_ids,
                                                       sd::Tensor<float>::from_vector(t5_mask),
-                                                      false,
-                                                      true,
-                                                      true);
+                                                      false);
         GGML_ASSERT(!hidden_states.empty());
         result.c_crossattn = std::move(hidden_states);
         result.c_vector    = sd::Tensor<float>::from_vector(mask);
@@ -1696,10 +1652,6 @@ struct AnimaConditioner : public Conditioner {
         llm->set_max_graph_vram_bytes(max_vram_bytes);
     }
 
-    void set_stream_layers_enabled(bool enabled) override {
-        llm->set_stream_layers_enabled(enabled);
-    }
-
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
         llm->set_runtime_backends(backends);
     }
@@ -1724,8 +1676,8 @@ struct AnimaConditioner : public Conditioner {
         llm->set_weight_adapter(adapter);
     }
 
-    void runner_done() override {
-        llm->runner_done();
+    void runner_end() override {
+        llm->runner_end();
     }
 
     std::tuple<std::vector<int>, std::vector<float>, std::vector<int>, std::vector<float>> tokenize(std::string text) {
@@ -1738,7 +1690,7 @@ struct AnimaConditioner : public Conditioner {
                 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> qwen_tokens;
@@ -1787,16 +1739,14 @@ struct AnimaConditioner : public Conditioner {
                                           {},
                                           {},
                                           false,
-                                          false,
-                                          true,
-                                          true);
+                                          false);
         GGML_ASSERT(!hidden_states.empty());
         hidden_states         = apply_token_weights(std::move(hidden_states), qwen_weights);
         auto t5_ids_tensor    = sd::Tensor<int32_t>::from_vector(t5_tokens);
         auto t5_weight_tensor = sd::Tensor<float>::from_vector(t5_weights);
 
         int64_t t1 = ggml_time_ms();
-        LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
+        LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
 
         SDCondition result;
         result.c_crossattn  = std::move(hidden_states);
@@ -1887,13 +1837,6 @@ struct LLMEmbedder : public Conditioner {
         }
     }
 
-    void set_stream_layers_enabled(bool enabled) override {
-        llm->set_stream_layers_enabled(enabled);
-        if (byt5) {
-            byt5->set_stream_layers_enabled(enabled);
-        }
-    }
-
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
         llm->set_runtime_backends(backends);
         if (byt5) {
@@ -1942,12 +1885,12 @@ struct LLMEmbedder : public Conditioner {
         }
     }
 
-    void runner_done() override {
+    void runner_end() override {
         if (llm) {
-            llm->runner_done();
+            llm->runner_end();
         }
         if (byt5) {
-            byt5->runner_done();
+            byt5->runner_end();
         }
     }
 
@@ -1984,7 +1927,7 @@ struct LLMEmbedder : public Conditioner {
                 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;
@@ -2051,8 +1994,6 @@ struct LLMEmbedder : public Conditioner {
                                           out_layers,
                                           false,
                                           false,
-                                          true,
-                                          true,
                                           deepstack_image_embeds,
                                           image_grids);
         GGML_ASSERT(!hidden_states.empty());
@@ -2220,9 +2161,7 @@ struct LLMEmbedder : public Conditioner {
                             prompt += "<Picture " + std::to_string(++picture_index) + ">: ";
                             add_vision_outputs(llm->encode_image_outputs(n_threads,
                                                                          resized,
-                                                                         false,
-                                                                         true,
-                                                                         true),
+                                                                         false),
                                                static_cast<int>(resized.shape()[1]) / patch_size,
                                                static_cast<int>(resized.shape()[0]) / patch_size);
                             continue;
@@ -2250,9 +2189,7 @@ struct LLMEmbedder : public Conditioner {
                             auto pair = sd::ops::concat(first.unsqueeze(2), second.unsqueeze(2), 2);
                             add_vision_outputs(llm->encode_video_block_outputs(n_threads,
                                                                                pair,
-                                                                               false,
-                                                                               true,
-                                                                               true),
+                                                                               false),
                                                static_cast<int>(first.shape()[1]) / patch_size,
                                                static_cast<int>(first.shape()[0]) / patch_size);
                         }
@@ -2263,9 +2200,7 @@ struct LLMEmbedder : public Conditioner {
                         prompt += "<Picture " + std::to_string(i + 1) + ">: ";
                         add_vision_outputs(llm->encode_image_outputs(n_threads,
                                                                      resized,
-                                                                     false,
-                                                                     true,
-                                                                     true),
+                                                                     false),
                                            static_cast<int>(resized.shape()[1]) / patch_size,
                                            static_cast<int>(resized.shape()[0]) / patch_size);
                     }
@@ -2310,7 +2245,7 @@ struct LLMEmbedder : public Conditioner {
                     prompt_template_encode_start_idx++;
                 }
             }
-            LOG_DEBUG("prompt_template_encode_start_idx %d", prompt_template_encode_start_idx);
+            LOG_VERBOSE("prompt_template_encode_start_idx %d", prompt_template_encode_start_idx);
 
             prompt = prompt_prefix;
             if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
@@ -2350,9 +2285,9 @@ struct LLMEmbedder : public Conditioner {
 
                     resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
 
-                    LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
+                    LOG_VERBOSE("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
                     auto resized_image = clip_preprocess(image, w_bar, h_bar);
-                    auto image_embed   = llm->encode_image(n_threads, resized_image, false, true, true);
+                    auto image_embed   = llm->encode_image(n_threads, resized_image, false);
                     GGML_ASSERT(!image_embed.empty());
 
                     std::string image_prefix = prompt + img_prompt + "<|vision_start|>";
@@ -2407,11 +2342,11 @@ struct LLMEmbedder : public Conditioner {
 
                     resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
 
-                    LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
+                    LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
 
                     auto resized_image = clip_preprocess(image, w_bar, h_bar);
 
-                    auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
+                    auto image_embed = llm->encode_image(n_threads, resized_image, false);
                     GGML_ASSERT(!image_embed.empty());
                     image_embeds.emplace_back(image_embed_idx, image_embed);
                     image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
@@ -2491,10 +2426,10 @@ struct LLMEmbedder : public Conditioner {
 
                     resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
 
-                    LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
+                    LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
 
                     auto resized_image = clip_preprocess(image, w_bar, h_bar);
-                    auto image_embed   = llm->encode_image(n_threads, resized_image, false, true, true);
+                    auto image_embed   = llm->encode_image(n_threads, resized_image, false);
                     GGML_ASSERT(!image_embed.empty());
 
                     std::string image_prefix = prompt_prefix + img_prompt + "<|vision_start|>";
@@ -2559,10 +2494,10 @@ struct LLMEmbedder : public Conditioner {
 
                     resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
 
-                    LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
+                    LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
 
                     auto resized_image = clip_preprocess(image, w_bar, h_bar);
-                    auto image_embed   = llm->encode_image(n_threads, resized_image, false, true, true);
+                    auto image_embed   = llm->encode_image(n_threads, resized_image, false);
                     GGML_ASSERT(!image_embed.empty());
 
                     std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
@@ -2622,10 +2557,10 @@ struct LLMEmbedder : public Conditioner {
 
                     resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
 
-                    LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
+                    LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
 
                     auto resized_image = clip_preprocess(image, w_bar, h_bar);
-                    auto image_embed   = llm->encode_image(n_threads, resized_image, false, true, true);
+                    auto image_embed   = llm->encode_image(n_threads, resized_image, false);
                     GGML_ASSERT(!image_embed.empty());
                     image_embeds.emplace_back(image_embed_idx, image_embed);
                     image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
@@ -2802,7 +2737,7 @@ struct LLMEmbedder : public Conditioner {
             }
 
             int64_t t1 = ggml_time_ms();
-            LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
+            LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
 
             SDCondition result;
             result.c_crossattn = std::move(hidden_states);
@@ -2857,9 +2792,7 @@ struct LLMEmbedder : public Conditioner {
                 auto byt5_hidden_states = byt5->compute(n_threads,
                                                         input_ids,
                                                         sd::Tensor<float>(),
-                                                        false,
-                                                        true,
-                                                        true);
+                                                        false);
                 GGML_ASSERT(!byt5_hidden_states.empty());
                 extra_hidden_states_vec.push_back(std::move(byt5_hidden_states));
             }
@@ -2879,7 +2812,7 @@ struct LLMEmbedder : public Conditioner {
         }
 
         int64_t t1 = ggml_time_ms();
-        LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
+        LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
         SDCondition result;
         result.c_crossattn        = std::move(hidden_states);
         result.extra_c_crossattns = std::move(extra_hidden_states_vec);
@@ -2960,13 +2893,11 @@ struct LTXAVTextProjectionRunner : public GGMLRunner {
 
     sd::Tensor<float> compute(int n_threads,
                               const sd::Tensor<float>& x,
-                              bool auto_free           = true,
-                              bool free_compute_buffer = true,
-                              bool free_compute_params = true) {
+                              bool auto_runner_end = true) {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(x);
         };
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
+        return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
     }
 };
 
@@ -2978,15 +2909,36 @@ struct LTXAVEmbedder : public Conditioner {
     std::shared_ptr<GemmaTokenizer> tokenizer;
     std::shared_ptr<LLM::LLMRunner> llm;
     std::shared_ptr<LTXAVTextProjectionRunner> projector;
+    std::string projector_prefix;
     bool dual_projection = false;
 
+    // Gemma 4 keeps a per-layer output scalar that no Gemma 3 checkpoint has, and widens its
+    // full-attention heads to 512 so their q_proj is twice a sliding layer's.
+    static LLM::LLMArch detect_gemma_arch(const String2TensorStorage& tensor_storage_map,
+                                          const std::string& llm_prefix) {
+        if (tensor_storage_map.find(llm_prefix + ".model.layers.0.layer_scalar") != tensor_storage_map.end()) {
+            return LLM::LLMArch::GEMMA4_12B;
+        }
+        auto global_q  = tensor_storage_map.find(llm_prefix + ".model.layers.5.self_attn.q_proj.weight");
+        auto sliding_q = tensor_storage_map.find(llm_prefix + ".model.layers.0.self_attn.q_proj.weight");
+        if (global_q != tensor_storage_map.end() &&
+            sliding_q != tensor_storage_map.end() &&
+            global_q->second.ne[1] == sliding_q->second.ne[1] * 2) {
+            return LLM::LLMArch::GEMMA4_12B;
+        }
+        return LLM::LLMArch::GEMMA3_12B;
+    }
+
     LTXAVEmbedder(ggml_backend_t backend,
                   const String2TensorStorage& tensor_storage_map      = {},
                   const std::string& llm_prefix                       = "text_encoders.llm",
                   const std::string& projector_prefix                 = "text_embedding_projection",
-                  std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) {
+                  std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : projector_prefix(projector_prefix) {
+        LLM::LLMArch arch = detect_gemma_arch(tensor_storage_map, llm_prefix);
+        LOG_INFO("ltxav text encoder: %s", arch == LLM::LLMArch::GEMMA4_12B ? "gemma 4" : "gemma 3");
         tokenizer       = std::make_shared<GemmaTokenizer>();
-        llm             = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::GEMMA3_12B,
+        llm             = std::make_shared<LLM::LLMRunner>(arch,
                                                backend,
                                                tensor_storage_map,
                                                llm_prefix,
@@ -3001,7 +2953,7 @@ struct LTXAVEmbedder : public Conditioner {
 
     void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         llm->get_param_tensors(tensors, "text_encoders.llm");
-        projector->get_param_tensors(tensors, "text_embedding_projection");
+        projector->get_param_tensors(tensors, projector_prefix);
     }
 
     void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
@@ -3039,9 +2991,9 @@ struct LTXAVEmbedder : public Conditioner {
         projector->set_weight_adapter(adapter);
     }
 
-    void runner_done() override {
-        llm->runner_done();
-        projector->runner_done();
+    void runner_end() override {
+        llm->runner_end();
+        projector->runner_end();
     }
 
     std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
@@ -3108,9 +3060,7 @@ struct LTXAVEmbedder : public Conditioner {
                                           {},
                                           {},
                                           true,
-                                          false,
-                                          true,
-                                          true);
+                                          false);
         GGML_ASSERT(!hidden_states.empty());
         hidden_states = apply_token_weights(std::move(hidden_states), weights);
 
@@ -3169,7 +3119,7 @@ struct LTXAVEmbedder : public Conditioner {
         }
 
         hidden_states.reshape_({kNumStates * kHiddenSize, valid_tokens});
-        return projector->compute(n_threads, hidden_states, false, true, true);
+        return projector->compute(n_threads, hidden_states, false);
     }
 
     SDCondition get_learned_condition(int n_threads,
@@ -3186,7 +3136,7 @@ struct LTXAVEmbedder : public Conditioner {
         GGML_ASSERT(!hidden_states.empty());
 
         int64_t t1 = ggml_time_ms();
-        LOG_DEBUG("computing LTXAV condition graph completed, taking %" PRId64 " ms", t1 - t0);
+        LOG_VERBOSE("computing LTXAV condition graph completed, taking %" PRId64 " ms", t1 - t0);
 
         SDCondition result;
         result.c_crossattn = std::move(hidden_states);

src/core/backend_fit.cpp

@@ -2,388 +2,409 @@
 
 #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 Plan {
+        bool valid         = false;
+        size_t main_device = SIZE_MAX;
+        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) {
+        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 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;
+            float gib = budgets.default_gib;
+            auto it   = budgets.backend_gib.find(budget_key(device.name));
+            if (it != budgets.backend_gib.end()) {
+                gib = it->second;
+            }
+            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;
+    }
 
-        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);
-            }
-            return out;
+    static int64_t available_ram_bytes() {
+#if defined(_WIN32)
+        MEMORYSTATUSEX status{};
+        status.dwLength = sizeof(status);
+        if (GlobalMemoryStatusEx(&status)) {
+            return (int64_t)status.ullAvailPhys;
         }
-
-        Plan compute_plan(const std::vector<Component>& components, const std::vector<Device>& devices) {
-            Plan plan;
-            if (devices.empty()) {
-                return plan;
+#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;
+    }
 
-            std::vector<size_t> order(components.size());
-            for (size_t i = 0; i < order.size(); i++) {
-                order[i] = i;
+    static Plan compute_plan(const std::vector<Component>& components,
+                             const std::vector<Device>& devices,
+                             int64_t ram_budget_bytes) {
+        Plan plan;
+        for (size_t di = 0; di < devices.size(); ++di) {
+            if (devices[di].budget_bytes > 0 &&
+                (plan.main_device == SIZE_MAX || devices[di].budget_bytes > devices[plan.main_device].budget_bytes)) {
+                plan.main_device = di;
             }
-            std::sort(order.begin(), order.end(), [&](size_t a, size_t b) {
-                return components[a].params_bytes > components[b].params_bytes;
-            });
-
-            {
-                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;
-                }
-            }
-
-            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) {
-                    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;
-                    }
-                }
-                if (best >= 0) {
-                    decision.device_idxs.push_back((size_t)best);
-                    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;
-                    }
-                }
-                decision.on_cpu = true;
-            }
-            plan.valid      = true;
-            plan.time_share = true;
+        }
+        if (plan.main_device == SIZE_MAX) {
             return plan;
         }
 
-        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());
-            }
+        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;
+        });
 
-        void append_assignment(std::string& spec, const char* key, const std::string& value) {
-            if (!spec.empty()) {
-                spec += ",";
-            }
-            spec += key;
-            spec += "=";
-            spec += value;
+        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());
 
-        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) {
+        for (size_t ci : order) {
+            const Component& comp = components[ci];
+            Decision& decision    = plan.decisions[ci];
+            if (comp.params_bytes == 0) {
+                continue;
+            }
+
+            // Higher-priority offloaded weights need GPU cache space across graph runs.
+            int64_t headroom = 0;
+            for (size_t other = 0; other < components.size(); ++other) {
+                if (components[other].params_bytes == 0) {
                     continue;
                 }
-                const Decision& decision = plan.decisions[ci];
-                if (decision.on_cpu) {
-                    append_assignment(runtime_spec, module_key, "cpu");
-                    return;
+                const bool resident          = other == ci || plan.decisions[other].params_location == ParamsLocation::MAIN_GPU;
+                const int64_t cached_weights = components[other].kind < comp.kind
+                                                   ? components[other].params_bytes
+                                                   : components[other].staging_bytes;
+                headroom                     = std::max(headroom, components[other].reserve_bytes +
+                                                                      (resident ? 0 : cached_weights));
+            }
+            int64_t& main_remaining = remaining[plan.main_device];
+            if (headroom <= main_remaining && comp.params_bytes <= main_remaining - headroom) {
+                decision.params_location = ParamsLocation::MAIN_GPU;
+                decision.params_device   = plan.main_device;
+                main_remaining -= comp.params_bytes;
+                continue;
+            }
+            if (comp.params_bytes <= ram_budget_bytes) {
+                decision.params_location = ParamsLocation::CPU;
+                ram_budget_bytes -= comp.params_bytes;
+                continue;
+            }
+
+            size_t best = SIZE_MAX;
+            for (size_t di = 0; di < devices.size(); ++di) {
+                if (di != plan.main_device && comp.params_bytes <= remaining[di] &&
+                    (best == SIZE_MAX || remaining[di] > remaining[best])) {
+                    best = di;
                 }
-                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;
+            }
+            if (best != SIZE_MAX) {
+                decision.params_location = ParamsLocation::OTHER_GPU;
+                decision.params_device   = best;
+                remaining[best] -= comp.params_bytes;
             }
         }
+        plan.valid = true;
+        return plan;
+    }
 
-    }  // namespace
+    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 (single-GPU compute on %s):", devices[plan.main_device].name.c_str());
+        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("  main-GPU weight cache priority: diffusion > te > vae");
+        LOG_INFO("  components (params: main GPU -> RAM -> other GPU -> disk):");
+        for (size_t ci = 0; ci < components.size(); ++ci) {
+            const Component& comp = components[ci];
+            if (comp.params_bytes == 0) {
+                continue;
+            }
+            const std::string params = params_backend_name(plan.decisions[ci], devices);
+            LOG_INFO("    %-12s params %6lld MiB, compute reserve %5lld MiB -> compute %s, params %s",
+                     comp.name, (long long)(comp.params_bytes / MiB), (long long)(comp.reserve_bytes / MiB),
+                     devices[plan.main_device].name.c_str(), params.c_str());
+        }
+    }
+
+    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 "";
+    }
 
     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;
-            }
-        }
-
-        auto components = estimate_components(loader, override_wtype);
-        auto devices    = enumerate_gpu_devices(budgets);
-        auto plan       = compute_plan(components, devices);
+        const auto components    = estimate_components(loader, override_wtype);
+        const auto devices       = enumerate_gpu_devices(budgets);
+        const int64_t free_ram   = available_ram_bytes();
+        const int64_t ram_budget = std::max<int64_t>(free_ram - std::max<int64_t>(2048 * MiB, free_ram / 10), 0);
+        const auto plan          = compute_plan(components, devices, ram_budget);
+        runtime_spec.clear();
+        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);
+            append_assignment(runtime_spec, key, devices[plan.main_device].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;
     }
 
     bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) {
-        if (prefer_temporal_tiling) {
-            if (tiling_params.temporal_tiling) {
-                return false;
-            }
+        const char* retry_mode = nullptr;
+        if (prefer_temporal_tiling && !tiling_params.temporal_tiling) {
             tiling_params.temporal_tiling = true;
-        } else {
-            if (tiling_params.enabled) {
-                return false;
-            }
+            retry_mode                    = tiling_params.enabled ? "spatial+temporal" : "temporal";
+        } else if (!tiling_params.enabled) {
             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;
             }
             if (tiling_params.tile_size_y <= 0) {
                 tiling_params.tile_size_y = 256;
             }
+            retry_mode = tiling_params.temporal_tiling ? "spatial+temporal" : "spatial";
+        } else {
+            return false;
         }
 
-        LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling",
-                 tiling_params.temporal_tiling ? "temporal" : "spatial");
+        LOG_WARN("VAE decode failed (likely out of memory); retrying with %s tiling",
+                 retry_mode);
         return true;
     }
 

src/core/compute_workspace.cpp

@@ -0,0 +1,257 @@
+#include "core/compute_workspace.h"
+
+#include <algorithm>
+#include <cstring>
+#include <map>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_graph_cut.h"
+#include "ggml-cpu.h"
+#include "ggml/src/ggml-impl.h"
+
+namespace sd {
+    ComputeWorkspace::~ComputeWorkspace() {
+        segment_end();
+        release();
+        ggml_backend_free(cpu_backend_);
+    }
+
+    void ComputeWorkspace::set_extra_backends(const std::vector<ggml_backend_t>& backends) {
+        if (extra_backends_ != backends) {
+            GGML_ASSERT(!active_);
+            release();
+            extra_backends_ = backends;
+        }
+    }
+
+    bool ComputeWorkspace::needs_scheduler(ggml_cgraph* graph) const {
+        if (!extra_backends_.empty()) {
+            return true;
+        }
+        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+            if (!ggml_backend_supports_op(backend_, ggml_graph_node(graph, i))) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    ggml_backend_sched_t ComputeWorkspace::make_scheduler(size_t graph_size) {
+        std::vector<ggml_backend_t> backends{backend_};
+        backends.insert(backends.end(), extra_backends_.begin(), extra_backends_.end());
+        if (!sd_backend_is_cpu(backend_)) {
+            if (cpu_backend_ == nullptr) {
+                cpu_backend_ = sd_backend_cpu_init();
+            }
+            if (cpu_backend_ == nullptr) {
+                return nullptr;
+            }
+            backends.push_back(cpu_backend_);
+        }
+        std::vector<ggml_backend_buffer_type_t> bufts;
+        for (auto backend : backends) {
+            auto buft = backend == cpu_backend_
+                            ? ggml_backend_dev_host_buffer_type(ggml_backend_get_device(backend_))
+                            : nullptr;
+            bufts.push_back(buft != nullptr ? buft : ggml_backend_get_default_buffer_type(backend));
+        }
+        return ggml_backend_sched_new(backends.data(), bufts.data(), static_cast<int>(backends.size()),
+                                      graph_size, false, false);
+    }
+
+    bool ComputeWorkspace::measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const {
+        return measurement.scheduler == needs_scheduler(graph);
+    }
+
+    bool ComputeWorkspace::prepare(const Measurement& measurement) {
+        GGML_ASSERT(!active_);
+        if (measurement.buffers.empty()) {
+            return false;
+        }
+        const bool grows = std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
+                                       [&](const BackendBufferSize& size) { return size.bytes > bytes(size.backend); });
+        if (measurement.scheduler != (scheduler_ != nullptr) || grows) {
+            release();
+        }
+        return true;
+    }
+
+    bool ComputeWorkspace::release_excess(const Measurement& measurement) {
+        return std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
+                           [&](const BackendBufferSize& size) { return bytes(size.backend) > size.bytes; }) &&
+               release();
+    }
+
+    bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) {
+        GGML_ASSERT(!active_);
+        const bool use_scheduler = needs_scheduler(graph);
+        if (use_scheduler) {
+            if (allocator_ != nullptr) {
+                release();
+            }
+            const size_t capacity = static_cast<size_t>(graph->n_nodes + graph->n_leafs) + 8;
+            if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
+                release();
+                scheduler_          = make_scheduler(capacity);
+                scheduler_capacity_ = capacity;
+            }
+            if (scheduler_ == nullptr) {
+                return false;
+            }
+            ggml_backend_sched_reset(scheduler_);
+            assign_nodes(scheduler_, graph);
+            // Scheduler allocation rewrites sources. Split the execution graph only once.
+            if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) {
+                release();
+                return false;
+            }
+        } else {
+            if (scheduler_ != nullptr) {
+                release();
+            }
+            if (allocator_ == nullptr) {
+                allocator_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
+            }
+            auto signature = ggml_graph_cut::graph_layout(graph, true);
+            if (signature != reservation_) {
+                if (!ggml_gallocr_reserve(allocator_, graph)) {
+                    release();
+                    return false;
+                }
+                reservation_ = std::move(signature);
+                ++reservations_;
+            }
+            if (!ggml_gallocr_alloc_graph(allocator_, graph)) {
+                release();
+                return false;
+            }
+        }
+        active_ = true;
+        return true;
+    }
+
+    ComputeWorkspace::Measurement ComputeWorkspace::measure(
+        ggml_cgraph* graph,
+        size_t direct_bytes,
+        const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
+        const AssignNodes& assign_nodes) {
+        if (!needs_scheduler(graph)) {
+            return {{{backend_, direct_bytes}}, false};
+        }
+        std::vector<const ggml_tensor*> tensors;
+        std::unordered_set<const ggml_tensor*> seen;
+        auto visit = [&](const ggml_tensor* tensor) {
+            if (tensor != nullptr && seen.insert(tensor).second) {
+                tensors.push_back(tensor);
+            }
+        };
+        for (int i = 0; i < graph->n_nodes; ++i) {
+            visit(graph->nodes[i]);
+        }
+        for (int i = 0; i < graph->n_leafs; ++i) {
+            visit(graph->leafs[i]);
+        }
+        for (size_t i = 0; i < tensors.size(); ++i) {
+            visit(tensors[i]->view_src);
+            for (auto source : tensors[i]->src) {
+                visit(source);
+            }
+        }
+        const size_t graph_size = tensors.size() + 8;
+        auto context            = ggml_init({tensors.size() * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false), nullptr, true});
+        if (context == nullptr) {
+            return {};
+        }
+        std::unordered_map<const ggml_tensor*, ggml_tensor*> copies;
+        std::map<ggml_backend_t, ggml_backend_buffer_t> external_buffers;
+        for (auto tensor : tensors) {
+            auto copy      = ggml_dup_tensor(context, tensor);
+            *copy          = *tensor;
+            copies[tensor] = copy;
+        }
+        for (const auto& entry : copies) {
+            auto source    = entry.first;
+            auto copy      = entry.second;
+            copy->view_src = source->view_src == nullptr ? nullptr : copies.at(source->view_src);
+            for (int i = 0; i < GGML_MAX_SRC; ++i) {
+                copy->src[i] = source->src[i] == nullptr ? nullptr : copies.at(source->src[i]);
+            }
+            auto external = external_backend(source);
+            if (external != nullptr && source->view_src == nullptr) {
+                auto& buffer = external_buffers[external];
+                if (buffer == nullptr) {
+                    buffer = ggml_backend_alloc_buffer(external, 0);
+                    GGML_ASSERT(buffer != nullptr);
+                    ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+                }
+                copy->buffer = buffer;
+                copy->data   = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
+                copy->extra  = nullptr;
+            }
+        }
+        auto copy_graph     = ggml_new_graph_custom(context, graph_size, false);
+        copy_graph->n_nodes = graph->n_nodes;
+        copy_graph->n_leafs = graph->n_leafs;
+        for (int i = 0; i < graph->n_nodes; ++i) {
+            copy_graph->nodes[i] = copies.at(graph->nodes[i]);
+        }
+        for (int i = 0; i < graph->n_leafs; ++i) {
+            copy_graph->leafs[i] = copies.at(graph->leafs[i]);
+        }
+        Measurement result;
+        result.scheduler = true;
+        auto scheduler   = make_scheduler(graph_size);
+        if (scheduler != nullptr) {
+            assign_nodes(scheduler, copy_graph);
+            std::vector<size_t> sizes(extra_backends_.size() + 2);
+            ggml_backend_sched_reserve_size(scheduler, copy_graph, sizes.data());
+            result.buffers.push_back({backend_, sizes[0]});
+            for (size_t i = 0; i < extra_backends_.size(); ++i) {
+                result.buffers.push_back({extra_backends_[i], sizes[i + 1]});
+            }
+            ggml_backend_sched_free(scheduler);
+        }
+        for (const auto& entry : external_buffers) {
+            ggml_backend_buffer_free(entry.second);
+        }
+        ggml_free(context);
+        return result;
+    }
+
+    void ComputeWorkspace::synchronize() const {
+        if (scheduler_ != nullptr) {
+            ggml_backend_sched_synchronize(scheduler_);
+        } else {
+            ggml_backend_synchronize(backend_);
+        }
+    }
+
+    void ComputeWorkspace::segment_end() {
+        if (active_) {
+            synchronize();
+            active_ = false;
+        }
+    }
+
+    bool ComputeWorkspace::release() {
+        if (active_) {
+            return false;
+        }
+        ggml_gallocr_free(allocator_);
+        allocator_ = nullptr;
+        ggml_backend_sched_free(scheduler_);
+        scheduler_          = nullptr;
+        scheduler_capacity_ = 0;
+        reservation_.clear();
+        return true;
+    }
+
+    size_t ComputeWorkspace::bytes(ggml_backend_t backend) const {
+        if (scheduler_ != nullptr) {
+            return ggml_backend_sched_get_buffer_size(scheduler_, backend);
+        }
+        return allocator_ != nullptr && backend == backend_ ? ggml_gallocr_get_buffer_size(allocator_, 0) : 0;
+    }
+}

src/core/compute_workspace.h

@@ -0,0 +1,64 @@
+#ifndef __SD_CORE_COMPUTE_WORKSPACE_H__
+#define __SD_CORE_COMPUTE_WORKSPACE_H__
+
+#include <functional>
+#include <vector>
+
+#include "ggml-alloc.h"
+#include "ggml-backend.h"
+
+namespace sd {
+    struct BackendBufferSize {
+        ggml_backend_t backend = nullptr;
+        size_t bytes           = 0;
+    };
+
+    class ComputeWorkspace {
+        ggml_backend_t backend_;
+        std::vector<ggml_backend_t> extra_backends_;
+        ggml_backend_t cpu_backend_     = nullptr;
+        ggml_gallocr_t allocator_       = nullptr;
+        ggml_backend_sched_t scheduler_ = nullptr;
+        size_t scheduler_capacity_      = 0;
+        std::vector<uint64_t> reservation_;
+        bool active_         = false;
+        size_t reservations_ = 0;
+
+        ggml_backend_sched_t make_scheduler(size_t graph_size);
+        bool needs_scheduler(ggml_cgraph* graph) const;
+
+    public:
+        struct Measurement {
+            std::vector<BackendBufferSize> buffers;
+            bool scheduler = false;
+        };
+        using AssignNodes = std::function<void(ggml_backend_sched_t, ggml_cgraph*)>;
+
+        explicit ComputeWorkspace(ggml_backend_t backend)
+            : backend_(backend) {}
+        ~ComputeWorkspace();
+        ComputeWorkspace(const ComputeWorkspace&)            = delete;
+        ComputeWorkspace& operator=(const ComputeWorkspace&) = delete;
+
+        void set_extra_backends(const std::vector<ggml_backend_t>& backends);
+        bool measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const;
+        bool prepare(const Measurement& measurement);
+        bool release_excess(const Measurement& measurement);
+        bool allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes);
+        Measurement measure(
+            ggml_cgraph* graph,
+            size_t direct_bytes,
+            const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
+            const AssignNodes& assign_nodes);
+        void synchronize() const;
+        void segment_end();
+        bool release();
+        bool active() const { return active_; }
+        ggml_backend_sched_t scheduler() const { return scheduler_; }
+        ggml_backend_t cpu_backend() const { return cpu_backend_; }
+        size_t bytes(ggml_backend_t backend) const;
+        size_t reservation_count() const { return reservations_; }
+    };
+}
+
+#endif  // __SD_CORE_COMPUTE_WORKSPACE_H__

src/core/ggml_extend.cpp

@@ -0,0 +1,727 @@
+#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);
+}
+
+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) {
+    if (ggml_n_dims(x) >= 3 && w != nullptr && b != nullptr) {
+        w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1);
+        b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+    }
+
+    const float eps = 1e-6f;  // default eps parameter
+    x               = ggml_group_norm(ctx, x, num_groups, eps);
+    if (w != nullptr && b != nullptr) {
+        x = ggml_mul_inplace(ctx, x, w);
+        // b = ggml_repeat(ctx, b, x);
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+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,235 @@
+#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);
+
+// 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);
+
+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

@@ -392,7 +392,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 +542,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,7 +587,7 @@ 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;
@@ -965,3 +965,35 @@ const char* sd_backend_module_name(SDBackendModule module) {
     }
     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

@@ -96,4 +96,6 @@ std::string sd_backend_resolve_name(const std::string& name);
 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

@@ -2,6 +2,7 @@
 
 #include <algorithm>
 #include <cctype>
+#include <climits>
 #include <cmath>
 #include <cstring>
 #include <map>
@@ -67,25 +68,6 @@ namespace sd::ggml_graph_cut {
         return -1;
     }
 
-    static Plan::InputShape input_shape(const ggml_tensor* tensor) {
-        Plan::InputShape shape;
-        if (tensor == nullptr) {
-            return shape;
-        }
-        shape.type = tensor->type;
-        for (int i = 0; i < GGML_MAX_DIMS; ++i) {
-            shape.ne[static_cast<size_t>(i)] = tensor->ne[i];
-        }
-        return shape;
-    }
-
-    static size_t graph_cut_segment_vram_bytes(const Segment& segment) {
-        return segment.compute_buffer_size +
-               segment.input_param_bytes +
-               segment.input_previous_cut_bytes +
-               segment.output_bytes;
-    }
-
     static std::string lower_ascii_copy(std::string value) {
         std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) {
             return static_cast<char>(std::tolower(c));
@@ -291,55 +273,6 @@ namespace sd::ggml_graph_cut {
         return max_vram_bytes_to_gib(resolve_auto_max_vram_bytes(-max_vram, backend));
     }
 
-    static bool is_segment_output_needed_after(const Plan& plan,
-                                               size_t end_segment_index,
-                                               int output_node_index) {
-        if (end_segment_index + 1 >= plan.segments.size()) {
-            return false;
-        }
-        for (size_t seg_idx = end_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) {
-            const auto& segment = plan.segments[seg_idx];
-            for (const auto& input_ref : segment.input_refs) {
-                if (input_ref.type == Segment::INPUT_PREVIOUS_CUT &&
-                    input_ref.node_index == output_node_index) {
-                    return true;
-                }
-            }
-        }
-        return false;
-    }
-
-    static Segment make_segment_seed(const Plan& plan,
-                                     size_t start_segment_index,
-                                     size_t end_segment_index) {
-        GGML_ASSERT(start_segment_index < plan.segments.size());
-        GGML_ASSERT(end_segment_index < plan.segments.size());
-        GGML_ASSERT(start_segment_index <= end_segment_index);
-
-        Segment seed;
-        const auto& start_segment  = plan.segments[start_segment_index];
-        const auto& target_segment = plan.segments[end_segment_index];
-        std::unordered_set<int> seen_output_node_indices;
-        for (size_t seg_idx = start_segment_index; seg_idx <= end_segment_index; ++seg_idx) {
-            const bool is_boundary_segment = seg_idx == end_segment_index;
-            for (int output_node_index : plan.segments[seg_idx].output_node_indices) {
-                if ((is_boundary_segment ||
-                     is_segment_output_needed_after(plan, end_segment_index, output_node_index)) &&
-                    seen_output_node_indices.insert(output_node_index).second) {
-                    seed.output_node_indices.push_back(output_node_index);
-                }
-            }
-        }
-        if (start_segment_index == end_segment_index) {
-            seed.group_name = target_segment.group_name;
-        } else {
-            seed.group_name = sd_format("%s..%s",
-                                        start_segment.group_name.c_str(),
-                                        target_segment.group_name.c_str());
-        }
-        return seed;
-    }
-
     static void build_segment(ggml_cgraph* gf,
                               Plan& plan,
                               Segment& segment,
@@ -416,31 +349,7 @@ namespace sd::ggml_graph_cut {
                       }
                       return a.display_name < b.display_name;
                   });
-        segment.input_refs = input_refs;
-        for (const auto& input : input_refs) {
-            ggml_tensor* current_input = input_tensor(gf, input);
-            size_t tensor_bytes        = current_input == nullptr
-                                             ? 0
-                                             : (input.type == Segment::INPUT_PREVIOUS_CUT
-                                                    ? cache_tensor_bytes(current_input)
-                                                    : ggml_nbytes(current_input));
-            switch (input.type) {
-                case Segment::INPUT_PREVIOUS_CUT:
-                    segment.input_previous_cut_bytes += tensor_bytes;
-                    break;
-                case Segment::INPUT_PARAM:
-                    segment.input_param_bytes += tensor_bytes;
-                    break;
-                case Segment::INPUT_EXTERNAL:
-                default:
-                    segment.input_external_bytes += tensor_bytes;
-                    break;
-            }
-        }
-        for (int output_node_index : segment.output_node_indices) {
-            ggml_tensor* output = ggml_graph_node(gf, output_node_index);
-            segment.output_bytes += cache_tensor_bytes(output);
-        }
+        segment.input_refs          = input_refs;
         segment.compute_buffer_size = measure_segment_compute_buffer(backend, gf, segment, log_desc);
 
         for (int output_node_index : segment.output_node_indices) {
@@ -449,15 +358,80 @@ namespace sd::ggml_graph_cut {
         plan.segments.push_back(std::move(segment));
     }
 
+    static bool validate_plan(ggml_cgraph* gf,
+                              const Plan& plan,
+                              std::string* validation_error) {
+        auto fail = [&](const std::string& reason) {
+            if (validation_error != nullptr) {
+                *validation_error = reason;
+            }
+            return false;
+        };
+        if (!plan.has_cuts) {
+            return true;
+        }
+        if (plan.segments.size() <= 1) {
+            return fail("fewer than two segments");
+        }
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        std::unordered_set<int> completed_outputs;
+        for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) {
+            const Segment& segment          = plan.segments[segment_index];
+            const std::string segment_label = "segment " + std::to_string(segment_index) +
+                                              " ('" + segment.group_name + "')";
+            if (segment.internal_node_indices.empty() || segment.output_node_indices.empty()) {
+                return fail(segment_label + " has no internal nodes or outputs");
+            }
+            for (const Segment::InputRef& input : segment.input_refs) {
+                if (input.type == Segment::INPUT_PREVIOUS_CUT) {
+                    if (input.node_index < 0 || input.node_index >= n_nodes ||
+                        completed_outputs.find(input.node_index) == completed_outputs.end()) {
+                        return fail(segment_label + " references an unavailable cut node " +
+                                    std::to_string(input.node_index));
+                    }
+                } else if (input.leaf_index < 0 || input.leaf_index >= gf->n_leafs) {
+                    return fail(segment_label + " references an invalid leaf " +
+                                std::to_string(input.leaf_index));
+                }
+            }
+            std::unordered_set<int> segment_nodes;
+            segment_nodes.reserve(segment.internal_node_indices.size());
+            for (int node_index : segment.internal_node_indices) {
+                if (node_index < 0 || node_index >= n_nodes) {
+                    return fail(segment_label + " contains an invalid node " +
+                                std::to_string(node_index));
+                }
+                if (!segment_nodes.insert(node_index).second) {
+                    return fail(segment_label + " contains duplicate node " +
+                                std::to_string(node_index));
+                }
+            }
+            for (int output_index : segment.output_node_indices) {
+                if (output_index < 0 || output_index >= n_nodes ||
+                    segment_nodes.find(output_index) == segment_nodes.end()) {
+                    return fail(segment_label + " has an output outside its node set: " +
+                                std::to_string(output_index));
+                }
+                if (completed_outputs.find(output_index) != completed_outputs.end()) {
+                    return fail(segment_label + " repeats output node " +
+                                std::to_string(output_index));
+                }
+                completed_outputs.insert(output_index);
+            }
+        }
+        return true;
+    }
+
     bool is_graph_cut_tensor(const ggml_tensor* tensor) {
         if (tensor == nullptr || tensor->name[0] == '\0') {
             return false;
         }
-        return std::strncmp(tensor->name, GGML_RUNNER_CUT_PREFIX, std::strlen(GGML_RUNNER_CUT_PREFIX)) == 0;
+        return starts_with(tensor->name, GGML_RUNNER_CUT_PREFIX) &&
+               ends_with(tensor->name, GGML_RUNNER_CUT_SUFFIX);
     }
 
     std::string make_graph_cut_name(const std::string& group, const std::string& output) {
-        return std::string(GGML_RUNNER_CUT_PREFIX) + group + "|" + output;
+        return std::string(GGML_RUNNER_CUT_PREFIX) + group + "|" + output + GGML_RUNNER_CUT_SUFFIX;
     }
 
     void mark_graph_cut(ggml_tensor* tensor, const std::string& group, const std::string& output) {
@@ -508,26 +482,87 @@ namespace sd::ggml_graph_cut {
         return ggml_nbytes(cache_src);
     }
 
-    bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) {
-        GGML_ASSERT(gf != nullptr);
-        if (ggml_graph_n_nodes(gf) != plan.n_nodes || gf->n_leafs != plan.n_leafs) {
-            return false;
-        }
-        for (const auto& input_shape_ref : plan.input_shapes) {
-            if (input_shape_ref.leaf_index < 0 || input_shape_ref.leaf_index >= gf->n_leafs) {
-                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;
+        auto add = [&](const ggml_tensor* tensor) {
+            if (tensor != nullptr && indices.emplace(tensor, tensors.size() + 1).second) {
+                tensors.push_back(tensor);
             }
-            ggml_tensor* leaf = gf->leafs[input_shape_ref.leaf_index];
-            if (leaf == nullptr || input_shape_ref.type != leaf->type) {
-                return false;
+        };
+        for (int i = 0; i < graph->n_leafs; ++i) {
+            add(graph->leafs[i]);
+        }
+        for (int i = 0; i < graph->n_nodes; ++i) {
+            add(graph->nodes[i]);
+        }
+        for (size_t i = 0; i < tensors.size(); ++i) {
+            add(tensors[i]->view_src);
+            for (auto source : tensors[i]->src) {
+                add(source);
+            }
+        }
+        std::vector<uint64_t> signature;
+        signature.reserve(tensors.size() * 24);
+        signature.push_back(graph->n_nodes);
+        signature.push_back(graph->n_leafs);
+        for (int i = 0; i < graph->n_leafs; ++i) {
+            signature.push_back(indices.at(graph->leafs[i]));
+        }
+        for (int i = 0; i < graph->n_nodes; ++i) {
+            signature.push_back(indices.at(graph->nodes[i]));
+        }
+        for (auto tensor : tensors) {
+            signature.push_back(tensor->op);
+            signature.push_back(tensor->type);
+            signature.push_back(tensor->flags);
+            signature.push_back(tensor->view_offs);
+            if (include_bindings) {
+                signature.push_back(tensor->data != nullptr);
+                auto buffer = tensor_buffer(tensor);
+                signature.push_back(reinterpret_cast<uintptr_t>(buffer == nullptr ? nullptr : ggml_backend_buffer_get_type(buffer)));
             }
             for (int d = 0; d < GGML_MAX_DIMS; ++d) {
-                if (input_shape_ref.ne[static_cast<size_t>(d)] != leaf->ne[d]) {
-                    return false;
-                }
+                signature.push_back(tensor->ne[d]);
+                signature.push_back(tensor->nb[d]);
+            }
+            signature.push_back(tensor->view_src == nullptr ? 0 : indices.at(tensor->view_src));
+            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));
             }
         }
-        return true;
+        return signature;
+    }
+
+    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 != layout) {
+            return false;
+        }
+        for (int i = 0; i < gf->n_leafs; ++i) {
+            if (plan.leaf_names[i] != gf->leafs[i]->name) {
+                return false;
+            }
+        }
+        std::vector<std::pair<int, std::string>> cut_markers;
+        for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) {
+            auto node = ggml_graph_node(gf, i);
+            if (is_graph_cut_tensor(node)) {
+                cut_markers.emplace_back(i, node->name);
+            }
+        }
+        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) {
@@ -577,33 +612,49 @@ namespace sd::ggml_graph_cut {
         return tensors;
     }
 
-    std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
-                                                               const Plan& plan,
-                                                               size_t current_segment_index) {
-        GGML_ASSERT(gf != nullptr);
-        std::unordered_set<std::string> future_input_names;
-        for (size_t seg_idx = current_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) {
-            const auto& segment = plan.segments[seg_idx];
-            for (const auto& input_ref : segment.input_refs) {
-                if (input_ref.type != Segment::INPUT_PREVIOUS_CUT) {
-                    continue;
-                }
-                ggml_tensor* current_input = input_tensor(gf, input_ref);
-                if (current_input != nullptr && current_input->name[0] != '\0') {
-                    future_input_names.insert(current_input->name);
-                }
-            }
-        }
-        return future_input_names;
-    }
-
     ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
                                      const Segment& segment,
                                      ggml_context** graph_ctx_out) {
         GGML_ASSERT(gf != nullptr);
         GGML_ASSERT(graph_ctx_out != nullptr);
 
-        const size_t graph_size = segment.internal_node_indices.size() + segment.input_refs.size() + 8;
+        // Collect leaf inputs and internal nodes, then any tensor they
+        // reference that is not already represented, notably the view_src of a
+        // view-typed input leaf. ggml_gallocr sizes its hash set from
+        // n_nodes + n_leafs (plus a 25% margin that rounds down to zero for a
+        // one-node segment), so every distinct tensor it will hash must be
+        // counted here or a tiny segment overflows the hash set and aborts.
+        std::vector<ggml_tensor*> leaves;
+        std::unordered_set<ggml_tensor*> represented;
+        for (const auto& input : segment.input_refs) {
+            ggml_tensor* current_input = input_tensor(gf, input);
+            if (current_input == nullptr) {
+                continue;
+            }
+            if (represented.insert(current_input).second) {
+                leaves.push_back(current_input);
+            }
+        }
+        for (int node_idx : segment.internal_node_indices) {
+            represented.insert(ggml_graph_node(gf, node_idx));
+        }
+        auto add_reference = [&](ggml_tensor* tensor) {
+            if (tensor != nullptr && represented.insert(tensor).second) {
+                leaves.push_back(tensor);
+            }
+        };
+        for (int node_idx : segment.internal_node_indices) {
+            ggml_tensor* node = ggml_graph_node(gf, node_idx);
+            for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
+                add_reference(node->src[src_idx]);
+            }
+            add_reference(node->view_src);
+        }
+        for (size_t i = 0; i < leaves.size(); ++i) {
+            add_reference(leaves[i]->view_src);
+        }
+
+        const size_t graph_size = segment.internal_node_indices.size() + leaves.size() + 8;
         ggml_init_params params = {
             /*.mem_size   =*/ggml_graph_overhead_custom(graph_size, false) + 1024,
             /*.mem_buffer =*/nullptr,
@@ -614,13 +665,9 @@ namespace sd::ggml_graph_cut {
         ggml_cgraph* segment_graph = ggml_new_graph_custom(graph_ctx, graph_size, false);
         GGML_ASSERT(segment_graph != nullptr);
 
-        for (const auto& input : segment.input_refs) {
-            ggml_tensor* current_input = input_tensor(gf, input);
-            if (current_input == nullptr) {
-                continue;
-            }
+        for (ggml_tensor* leaf : leaves) {
             GGML_ASSERT(segment_graph->n_leafs < segment_graph->size);
-            segment_graph->leafs[segment_graph->n_leafs++] = current_input;
+            segment_graph->leafs[segment_graph->n_leafs++] = leaf;
         }
 
         for (int output_node_index : segment.output_node_indices) {
@@ -629,6 +676,10 @@ namespace sd::ggml_graph_cut {
                 continue;
             }
             ggml_set_output(output);
+            if (output->view_src != nullptr) {
+                // A consumed output view does not keep its storage alive in gallocr.
+                ggml_set_output(output->view_src);
+            }
         }
         for (int node_idx : segment.internal_node_indices) {
             ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx));
@@ -683,6 +734,10 @@ namespace sd::ggml_graph_cut {
             if (output != nullptr && saved_output_flags.find(output) == saved_output_flags.end()) {
                 saved_output_flags[output] = output->flags;
             }
+            if (output != nullptr && output->view_src != nullptr &&
+                saved_output_flags.find(output->view_src) == saved_output_flags.end()) {
+                saved_output_flags[output->view_src] = output->view_src->flags;
+            }
         }
 
         ggml_context* graph_ctx    = nullptr;
@@ -711,6 +766,46 @@ namespace sd::ggml_graph_cut {
         return buffer_size;
     }
 
+    static size_t measure_graph_compute_buffer(
+        ggml_backend_t backend,
+        ggml_cgraph* gf,
+        const std::unordered_set<const ggml_tensor*>& params_tensor_set) {
+        struct TensorRuntimeBinding {
+            ggml_backend_buffer_t buffer = nullptr;
+            void* data                   = nullptr;
+            void* extra                  = nullptr;
+        };
+        std::unordered_map<ggml_tensor*, TensorRuntimeBinding> saved_bindings;
+        auto mark_external = [&](ggml_tensor* tensor) {
+            if (tensor == nullptr || saved_bindings.find(tensor) != saved_bindings.end()) {
+                return;
+            }
+            saved_bindings[tensor] = {tensor->buffer, tensor->data, tensor->extra};
+            tensor->data           = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
+        };
+        for (int i = 0; i < leaf_count(gf); ++i) {
+            ggml_tensor* leaf = leaf_tensor(gf, i);
+            if (!is_params_tensor(params_tensor_set, leaf)) {
+                continue;
+            }
+            mark_external(leaf);
+            mark_external(leaf->view_src);
+        }
+
+        ggml_gallocr_t allocr = ggml_gallocr_new(
+            ggml_backend_get_default_buffer_type(backend));
+        size_t sizes[1] = {0};
+        ggml_gallocr_reserve_n_size(allocr, gf, nullptr, nullptr, sizes);
+        ggml_gallocr_free(allocr);
+
+        for (const auto& kv : saved_bindings) {
+            kv.first->buffer = kv.second.buffer;
+            kv.first->data   = kv.second.data;
+            kv.first->extra  = kv.second.extra;
+        }
+        return sizes[0];
+    }
+
     Plan build_plan(ggml_backend_t backend,
                     ggml_cgraph* gf,
                     const std::unordered_set<const ggml_tensor*>& params_tensor_set,
@@ -723,24 +818,22 @@ namespace sd::ggml_graph_cut {
         if (n_nodes <= 0) {
             return plan;
         }
-        plan.n_nodes = n_nodes;
-        plan.n_leafs = gf->n_leafs;
+        plan.layout = graph_layout(gf, false);
         for (int i = 0; i < gf->n_leafs; ++i) {
-            ggml_tensor* leaf = gf->leafs[i];
-            if (is_params_tensor(params_tensor_set, leaf)) {
-                continue;
-            }
-            auto shape       = input_shape(leaf);
-            shape.leaf_index = i;
-            plan.input_shapes.push_back(shape);
+            plan.leaf_names.emplace_back(gf->leafs[i]->name);
         }
+        plan.compute_buffer_size =
+            measure_graph_compute_buffer(backend, gf, params_tensor_set);
 
         std::unordered_map<const ggml_tensor*, int> producer_index;
         producer_index.reserve(static_cast<size_t>(n_nodes));
         for (int i = 0; i < n_nodes; ++i) {
-            producer_index[ggml_graph_node(gf, i)] = i;
+            ggml_tensor* node    = ggml_graph_node(gf, i);
+            producer_index[node] = i;
+            if (is_graph_cut_tensor(node)) {
+                plan.cut_markers.push_back({i, node->name});
+            }
         }
-
         std::vector<Segment> grouped_segments;
         std::unordered_map<std::string, size_t> group_to_segment;
         for (int i = 0; i < n_nodes; ++i) {
@@ -751,7 +844,9 @@ namespace sd::ggml_graph_cut {
 
             plan.has_cuts = true;
             std::string full_name(node->name);
-            std::string payload = full_name.substr(std::strlen(GGML_RUNNER_CUT_PREFIX));
+            size_t prefix_len   = std::strlen(GGML_RUNNER_CUT_PREFIX);
+            size_t suffix_len   = std::strlen(GGML_RUNNER_CUT_SUFFIX);
+            std::string payload = full_name.substr(prefix_len, full_name.size() - prefix_len - suffix_len);
             size_t sep          = payload.find('|');
             std::string group   = sep == std::string::npos ? payload : payload.substr(0, sep);
 
@@ -789,11 +884,24 @@ namespace sd::ggml_graph_cut {
         if (final_output_index < 0) {
             final_output_index = n_nodes - 1;
         }
-        ggml_tensor* final_output = final_output_index >= 0 ? ggml_graph_node(gf, final_output_index) : nullptr;
-        if (final_output != nullptr && available_cut_output_node_indices.find(final_output_index) == available_cut_output_node_indices.end()) {
-            Segment final_segment;
-            final_segment.group_name = "ggml_runner.final";
+        Segment final_segment;
+        final_segment.group_name = "ggml_runner.final";
+        if (final_output_index >= 0 &&
+            available_cut_output_node_indices.find(final_output_index) ==
+                available_cut_output_node_indices.end()) {
             final_segment.output_node_indices.push_back(final_output_index);
+        }
+        for (int i = 0; i < n_nodes; ++i) {
+            ggml_tensor* node = ggml_graph_node(gf, i);
+            if (i == final_output_index || node == nullptr ||
+                (node->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 ||
+                available_cut_output_node_indices.find(i) !=
+                    available_cut_output_node_indices.end()) {
+                continue;
+            }
+            final_segment.output_node_indices.push_back(i);
+        }
+        if (!final_segment.output_node_indices.empty()) {
             build_segment(gf,
                           plan,
                           final_segment,
@@ -804,210 +912,59 @@ namespace sd::ggml_graph_cut {
                           log_desc);
         }
 
-        return plan;
-    }
-
-    Plan apply_max_vram_budget(ggml_cgraph* gf,
-                               const Plan& base_plan,
-                               size_t max_graph_vram_bytes,
-                               ggml_backend_t backend,
-                               const std::unordered_set<const ggml_tensor*>& params_tensor_set,
-                               const char* log_desc) {
-        GGML_ASSERT(backend != nullptr);
-        GGML_ASSERT(gf != nullptr);
-        int64_t t_budget_begin = ggml_time_ms();
-        if (max_graph_vram_bytes == 0 || !base_plan.has_cuts || base_plan.segments.size() <= 1) {
-            return base_plan;
-        }
-
-        const int n_nodes = ggml_graph_n_nodes(gf);
-        std::unordered_map<const ggml_tensor*, int> producer_index;
-        producer_index.reserve(static_cast<size_t>(n_nodes));
-        for (int i = 0; i < n_nodes; ++i) {
-            producer_index[ggml_graph_node(gf, i)] = i;
-        }
-
-        Plan merged_plan;
-        merged_plan.available = true;
-        merged_plan.has_cuts  = base_plan.has_cuts;
-        merged_plan.valid     = base_plan.valid;
-        merged_plan.n_nodes   = base_plan.n_nodes;
-        merged_plan.n_leafs   = base_plan.n_leafs;
-
-        std::unordered_set<int> available_cut_output_node_indices;
-        available_cut_output_node_indices.reserve(static_cast<size_t>(n_nodes));
-
-        size_t start_segment_index = 0;
-        while (start_segment_index < base_plan.segments.size()) {
-            Plan single_plan;
-            auto single_available_cut_output_node_indices = available_cut_output_node_indices;
-            auto single_seed                              = make_segment_seed(base_plan,
-                                                                              start_segment_index,
-                                                                              start_segment_index);
-            build_segment(gf,
-                          single_plan,
-                          single_seed,
-                          producer_index,
-                          single_available_cut_output_node_indices,
-                          backend,
-                          params_tensor_set,
-                          log_desc);
-            GGML_ASSERT(!single_plan.segments.empty());
-
-            size_t best_end_segment_index = start_segment_index;
-            bool can_merge_next_segment   = graph_cut_segment_vram_bytes(single_plan.segments.back()) <= max_graph_vram_bytes;
-
-            while (can_merge_next_segment && best_end_segment_index + 1 < base_plan.segments.size()) {
-                const size_t next_end_segment_index = best_end_segment_index + 1;
-                Plan candidate_plan;
-                auto candidate_available_cut_output_node_indices = available_cut_output_node_indices;
-                auto candidate_seed                              = make_segment_seed(base_plan,
-                                                                                     start_segment_index,
-                                                                                     next_end_segment_index);
-                build_segment(gf,
-                              candidate_plan,
-                              candidate_seed,
-                              producer_index,
-                              candidate_available_cut_output_node_indices,
-                              backend,
-                              params_tensor_set,
-                              log_desc);
-                GGML_ASSERT(!candidate_plan.segments.empty());
-
-                const auto& candidate_segment = candidate_plan.segments.back();
-                const size_t candidate_bytes  = graph_cut_segment_vram_bytes(candidate_segment);
-                if (candidate_bytes > max_graph_vram_bytes) {
-                    break;
+        std::unordered_set<std::string> future_cut_names;
+        for (auto segment = plan.segments.rbegin(); segment != plan.segments.rend(); ++segment) {
+            segment->future_cut_names = future_cut_names;
+            segment->live_cut_names   = future_cut_names;
+            for (const auto& input : segment->input_refs) {
+                if (input.type != Segment::INPUT_PREVIOUS_CUT) {
+                    continue;
                 }
-
-                best_end_segment_index = next_end_segment_index;
+                segment->live_cut_names.insert(input.display_name);
+                future_cut_names.insert(input.display_name);
             }
-
-            auto best_seed = make_segment_seed(base_plan,
-                                               start_segment_index,
-                                               best_end_segment_index);
-            build_segment(gf,
-                          merged_plan,
-                          best_seed,
-                          producer_index,
-                          available_cut_output_node_indices,
-                          backend,
-                          params_tensor_set,
-                          log_desc);
-            start_segment_index = best_end_segment_index + 1;
         }
 
-        if (log_desc != nullptr && merged_plan.segments.size() != base_plan.segments.size()) {
-            LOG_INFO("%s graph cut max_vram=%.2f MB merged %zu segments -> %zu segments",
+        std::string plan_validation_error;
+        plan.valid = validate_plan(gf, plan, &plan_validation_error);
+        if (!plan.valid && log_desc != nullptr) {
+            LOG_WARN("%s graph cut plan validation failed (%s); using monolithic execution",
                      log_desc,
-                     max_graph_vram_bytes / 1024.0 / 1024.0,
-                     base_plan.segments.size(),
-                     merged_plan.segments.size());
+                     plan_validation_error.c_str());
         }
 
-        if (log_desc != nullptr) {
-            LOG_DEBUG("%s graph cut max_vram budget merge took %lld ms",
-                      log_desc,
-                      ggml_time_ms() - t_budget_begin);
-        }
-
-        return merged_plan;
+        return plan;
     }
 
     Plan resolve_plan(ggml_backend_t backend,
                       ggml_cgraph* gf,
                       PlanCache* cache,
-                      size_t max_graph_vram_bytes,
                       const std::unordered_set<const ggml_tensor*>& params_tensor_set,
                       const char* log_desc) {
         GGML_ASSERT(backend != nullptr);
         GGML_ASSERT(gf != nullptr);
         GGML_ASSERT(cache != nullptr);
 
-        int64_t t_prepare_begin = ggml_time_ms();
-        Plan base_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();
-        if (cache->graph_cut_plan.available && plan_matches_graph(gf, cache->graph_cut_plan)) {
-            base_plan = cache->graph_cut_plan;
-        } else {
-            base_plan                                = build_plan(backend, gf, params_tensor_set, log_desc);
-            cache->graph_cut_plan                    = base_plan;
-            cache->graph_cut_plan.available          = true;
-            cache->budgeted_graph_cut_plan.available = false;
-            if (log_desc != nullptr) {
-                LOG_INFO("%s build cached graph cut plan done (taking %lld ms)", log_desc, ggml_time_ms() - t_plan_begin);
-            }
+        plans.push_front(build_plan(backend, gf, params_tensor_set, log_desc));
+        if (plans.size() > PlanCache::MAX_PLANS) {
+            plans.pop_back();
         }
-
-        Plan resolved_plan = base_plan;
-        if (max_graph_vram_bytes > 0 && base_plan.has_cuts) {
-            if (cache->budgeted_graph_cut_plan.available &&
-                cache->budgeted_graph_cut_plan_max_vram_bytes == max_graph_vram_bytes &&
-                plan_matches_graph(gf, cache->budgeted_graph_cut_plan)) {
-                resolved_plan = cache->budgeted_graph_cut_plan;
-            } else {
-                resolved_plan                                 = apply_max_vram_budget(gf,
-                                                                                      base_plan,
-                                                                                      max_graph_vram_bytes,
-                                                                                      backend,
-                                                                                      params_tensor_set,
-                                                                                      log_desc);
-                cache->budgeted_graph_cut_plan                = resolved_plan;
-                cache->budgeted_graph_cut_plan.available      = true;
-                cache->budgeted_graph_cut_plan_max_vram_bytes = max_graph_vram_bytes;
-            }
-        }
-        return resolved_plan;
-    }
-
-    void annotate_residency(Plan& plan, size_t max_graph_vram_bytes) {
-        // Cached plans may be reused with a smaller live budget.
-        for (auto& seg : plan.segments) {
-            seg.residency = SegmentResidency::STREAMED;
-        }
-        if (max_graph_vram_bytes == 0 || plan.segments.size() < 2) {
-            return;
-        }
-
-        bool any_param_bearing = false;
-        for (const auto& seg : plan.segments) {
-            if (seg.input_param_bytes > 0) {
-                any_param_bearing = true;
-                break;
-            }
-        }
-        if (!any_param_bearing) {
-            return;
-        }
-
-        // Leave room for the largest active streamed segment.
-        size_t worst_streamed_footprint = 0;
-        for (const auto& seg : plan.segments) {
-            const size_t seg_footprint = seg.input_param_bytes +
-                                         seg.compute_buffer_size +
-                                         seg.output_bytes +
-                                         seg.input_previous_cut_bytes +
-                                         seg.input_external_bytes;
-            if (seg_footprint > worst_streamed_footprint) {
-                worst_streamed_footprint = seg_footprint;
-            }
-        }
-        constexpr size_t safety = 512ull * 1024 * 1024;
-        const size_t reserved   = safety + worst_streamed_footprint;
-
-        if (max_graph_vram_bytes <= reserved) {
-            return;
-        }
-        const size_t available = max_graph_vram_bytes - reserved;
-
-        size_t cumulative = 0;
-        for (auto& seg : plan.segments) {
-            if (cumulative + seg.input_param_bytes > available) {
-                break;
-            }
-            seg.residency = SegmentResidency::RESIDENT;
-            cumulative += seg.input_param_bytes;
+        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 plans.front();
     }
 
 }  // namespace sd::ggml_graph_cut

src/core/ggml_graph_cut.h

@@ -3,22 +3,17 @@
 
 #include <array>
 #include <cstdint>
+#include <list>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
+#include <utility>
 #include <vector>
 
 #include "ggml-backend.h"
 #include "ggml.h"
 
 namespace sd::ggml_graph_cut {
-
-    // Streaming residency for a segment's params.
-    enum class SegmentResidency : uint8_t {
-        STREAMED = 0,
-        RESIDENT = 1,
-    };
-
     struct Segment {
         enum InputType {
             INPUT_EXTERNAL = 0,
@@ -33,41 +28,33 @@ namespace sd::ggml_graph_cut {
             int node_index = -1;
         };
 
-        size_t compute_buffer_size      = 0;
-        size_t output_bytes             = 0;
-        size_t input_external_bytes     = 0;
-        size_t input_previous_cut_bytes = 0;
-        size_t input_param_bytes        = 0;
+        size_t compute_buffer_size = 0;
         std::string group_name;
         std::vector<int> internal_node_indices;
         std::vector<int> output_node_indices;
         std::vector<InputRef> input_refs;
-        SegmentResidency residency = SegmentResidency::STREAMED;
+        std::unordered_set<std::string> future_cut_names;
+        std::unordered_set<std::string> live_cut_names;
     };
 
     struct Plan {
-        struct InputShape {
-            int leaf_index                        = -1;
-            ggml_type type                        = GGML_TYPE_COUNT;
-            std::array<int64_t, GGML_MAX_DIMS> ne = {0, 0, 0, 0};
-        };
-
-        bool available = false;
-        bool has_cuts  = false;
-        bool valid     = true;
-        int n_nodes    = 0;
-        int n_leafs    = 0;
-        std::vector<InputShape> input_shapes;
+        bool available             = false;
+        bool has_cuts              = false;
+        bool valid                 = true;
+        size_t compute_buffer_size = 0;
+        std::vector<uint64_t> layout;
+        std::vector<std::string> leaf_names;
+        std::vector<std::pair<int, std::string>> cut_markers;
         std::vector<Segment> segments;
     };
 
     struct PlanCache {
-        Plan graph_cut_plan;
-        Plan budgeted_graph_cut_plan;
-        size_t budgeted_graph_cut_plan_max_vram_bytes = 0;
+        static constexpr size_t MAX_PLANS = 4;
+        std::list<Plan> graph_cut_plans;
     };
 
     static constexpr const char* GGML_RUNNER_CUT_PREFIX = "ggml_runner_cut:";
+    static constexpr const char* GGML_RUNNER_CUT_SUFFIX = "|";
 
     struct MaxVramAssignment {
         float default_gib = 0.f;
@@ -88,13 +75,12 @@ namespace sd::ggml_graph_cut {
     ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor);
     ggml_tensor* cache_source_tensor(ggml_tensor* tensor);
     size_t cache_tensor_bytes(const ggml_tensor* tensor);
+    // Plans ignore runtime bindings; allocator reservations must include them.
+    std::vector<uint64_t> graph_layout(ggml_cgraph* graph, bool include_bindings);
     bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan);
     ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index);
     ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref);
     std::vector<ggml_tensor*> param_tensors(ggml_cgraph* gf, const Segment& segment);
-    std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
-                                                               const Plan& plan,
-                                                               size_t current_segment_index);
     ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
                                      const Segment& segment,
                                      ggml_context** graph_ctx_out);
@@ -108,21 +94,12 @@ namespace sd::ggml_graph_cut {
                     ggml_cgraph* gf,
                     const std::unordered_set<const ggml_tensor*>& params_tensor_set,
                     const char* log_desc);
-    Plan apply_max_vram_budget(ggml_cgraph* gf,
-                               const Plan& base_plan,
-                               size_t max_graph_vram_bytes,
-                               ggml_backend_t backend,
-                               const std::unordered_set<const ggml_tensor*>& params_tensor_set,
-                               const char* log_desc);
     Plan resolve_plan(ggml_backend_t backend,
                       ggml_cgraph* gf,
                       PlanCache* cache,
-                      size_t max_graph_vram_bytes,
                       const std::unordered_set<const ggml_tensor*>& params_tensor_set,
                       const char* log_desc);
 
-    // Mark leading segments resident when they fit after streamed-segment headroom.
-    void annotate_residency(Plan& plan, size_t max_graph_vram_bytes);
 }  // namespace sd::ggml_graph_cut
 
 #endif  // __SD_CORE_GGML_GRAPH_CUT_H__

src/core/ggml_runner.cpp

@@ -0,0 +1,955 @@
+#include <algorithm>
+#include <map>
+#include <utility>
+
+#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;
+
+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.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);
+        }
+    }
+    auto output = execute_graph(graph, n_threads, no_return, read_outputs);
+    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;
+}
+
+ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, size_t direct_bytes) {
+    auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t {
+        if (!params_tensor_set_.count(tensor)) {
+            return nullptr;
+        }
+        auto placement = graph_cut_layer_split_assignments_.find(tensor);
+        return placement == graph_cut_layer_split_assignments_.end() ? runtime_backend : placement->second;
+    };
+    auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) {
+        pin_multi_device_nodes(scheduler, copy, graph);
+    };
+    return workspace_.measure(graph, direct_bytes, external_backend, assign_nodes);
+}
+
+std::vector<DeviceMemoryRequest> GGMLRunner::memory_requests(
+    const std::vector<BackendBufferSize>& sizes,
+    size_t pending_cache_bytes) const {
+    std::vector<DeviceMemoryRequest> requests;
+    for (const auto& size : sizes) {
+        const size_t retained    = retained_runtime_buffer_bytes(size.backend);
+        const size_t reusable    = workspace_.bytes(size.backend);
+        const size_t cache_bytes = size.backend == runtime_backend ? pending_cache_bytes : 0;
+        const size_t pending     = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes);
+        size_t limit             = max_graph_vram_bytes;
+        if (is_multi_device()) {
+            size_t index = 0;
+            if (size.backend != runtime_backend) {
+                auto position = std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), size.backend);
+                index         = static_cast<size_t>(position - extra_runtime_backends.begin()) + 1;
+            }
+            if (index < graph_cut_layer_split_backend_vram_limits_.size()) {
+                limit = graph_cut_layer_split_backend_vram_limits_[index];
+            }
+        }
+        requests.push_back({size.backend, reinterpret_cast<uintptr_t>(this), pending,
+                            retained, limit});
+    }
+    return requests;
+}
+
+bool GGMLRunner::fits(const std::vector<DeviceMemoryRequest>& requests,
+                      const std::vector<ggml_tensor*>& params) const {
+    auto manager = residency_manager.lock();
+    if (manager == nullptr) {
+        return params.empty();
+    }
+    for (const auto& request : requests) {
+        if (!manager->fits_compute_backend_capacity(request, params)) {
+            return false;
+        }
+    }
+    return true;
+}
+
+bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) {
+    if (sd_backend_is_cpu(runtime_backend)) {
+        sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
+    }
+    if (workspace_.cpu_backend() != nullptr) {
+        sd_backend_cpu_set_n_threads(workspace_.cpu_backend(), n_threads);
+    }
+    auto scheduler = workspace_.scheduler();
+    ggml_status status;
+    if (scheduler != nullptr) {
+        if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
+            LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring", get_desc().c_str());
+            multi_device_eval_callback_warned = true;
+        }
+        status = ggml_backend_sched_graph_compute(scheduler, graph);
+    } else {
+        status = sd_backend_graph_compute_with_eval_callback(runtime_backend, graph,
+                                                             sd_get_backend_eval_callback(),
+                                                             sd_get_backend_eval_callback_data());
+    }
+    workspace_.synchronize();
+    if (status != GGML_STATUS_SUCCESS) {
+        LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
+        return false;
+    }
+    const std::string description = get_desc();
+    if (!debug_tensors.empty()) {
+        std::unordered_set<const ggml_tensor*> graph_tensors;
+        const int leaf_count = ggml_graph_cut::leaf_count(graph);
+        const int node_count = ggml_graph_n_nodes(graph);
+        graph_tensors.reserve(static_cast<size_t>(leaf_count + node_count));
+        for (int index = 0; index < leaf_count; ++index) {
+            graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index));
+        }
+        for (int index = 0; index < node_count; ++index) {
+            graph_tensors.insert(ggml_graph_node(graph, index));
+        }
+
+        for (const auto& entry : debug_tensors) {
+            ggml_tensor* tensor = entry.first;
+            if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) {
+                continue;
+            }
+            ggml_backend_buffer_t buffer =
+                tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
+            if (buffer == nullptr) {
+                LOG_WARN("%s skip debug tensor '%s': tensor buffer not set",
+                         description.c_str(),
+                         entry.second.c_str());
+                continue;
+            }
+            if (tensor->type != GGML_TYPE_F32) {
+                LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s",
+                         description.c_str(),
+                         entry.second.c_str(),
+                         ggml_type_name(tensor->type));
+                continue;
+            }
+            auto debug_tensor = make_sd_tensor_from_ggml<float>(tensor);
+            print_sd_tensor(debug_tensor, false, entry.second.c_str());
+        }
+    }
+
+    return true;
+}
+
+std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs) {
+    if (!assign_graph_cut_layer_split_backends(graph)) {
+        return std::nullopt;
+    }
+    const auto params = collect_used_param_tensors(graph);
+    ggml_graph_cut::Plan plan;
+    if (!resolve_graph_cut_plan(graph, &plan)) {
+        return std::nullopt;
+    }
+    const auto full_measurement = measure(graph, plan.compute_buffer_size);
+    if (full_measurement.buffers.empty()) {
+        return std::nullopt;
+    }
+    auto manager         = residency_manager.lock();
+    const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
+                           manager != nullptr && manager->segmented_compute_enabled() &&
+                           plan.valid && plan.has_cuts && plan.segments.size() > 1 &&
+                           !fits(memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)), params);
+    if (!segmented) {
+        ggml_graph_cut::Segment segment;
+        segment.group_name          = "graph";
+        segment.compute_buffer_size = plan.compute_buffer_size;
+        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+            segment.internal_node_indices.push_back(i);
+        }
+        for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) {
+            auto tensor = ggml_graph_cut::leaf_tensor(graph, i);
+            ggml_graph_cut::Segment::InputRef input;
+            input.leaf_index = i;
+            input.type       = canonical_param_tensor(tensor) != nullptr
+                                   ? ggml_graph_cut::Segment::INPUT_PARAM
+                                   : ggml_graph_cut::Segment::INPUT_EXTERNAL;
+            segment.input_refs.push_back(input);
+        }
+        plan.segments = {std::move(segment)};
+    }
+    const bool segments_changed = plan.segments.size() != logged_segment_count_;
+    if (segments_changed && (segmented || logged_segment_count_ > 1)) {
+        LOG_VERBOSE("%s using %zu segment%s", get_desc().c_str(),
+                    plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
+    }
+    SegmentGraphBindings bindings(cut_cache_, plan, graph);
+    SegmentWeightPipeline weights(manager, runtime_backend, reinterpret_cast<uintptr_t>(this),
+                                  graph, plan, params_tensor_set_,
+                                  segmented && manager != nullptr && manager->prefetch_enabled());
+
+    std::map<ggml_backend_t, size_t> peak_compute_bytes;
+    auto track_compute_buffer = [&](ggml_backend_t backend) {
+        if (backend != nullptr) {
+            auto& peak = peak_compute_bytes[backend];
+            peak       = std::max(peak, workspace_.bytes(backend));
+        }
+    };
+    std::optional<Tensor<float>> output = Tensor<float>();
+    for (size_t index = 0; index < plan.segments.size(); ++index) {
+        const auto& segment = plan.segments[index];
+        const bool last     = index + 1 == plan.segments.size();
+        auto fail_segment   = [&](const char* phase) {
+            LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", get_desc().c_str(),
+                        index + 1, plan.segments.size(), segment.group_name.c_str(), phase);
+            return std::nullopt;
+        };
+        cut_cache_.prune(segment.live_cut_names);
+        bindings.reset(segment);
+        if (!bindings.bind_cached_inputs(segment, get_desc().c_str())) {
+            return fail_segment("input binding");
+        }
+        ggml_context* segment_context = nullptr;
+        auto segment_graph            = segmented
+                                            ? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context)
+                                            : graph;
+        struct SegmentCleanup {
+            GGMLRunner& runner;
+            SegmentWeightPipeline& weights;
+            SegmentGraphBindings& bindings;
+            ggml_context* context;
+            ~SegmentCleanup() {
+                runner.workspace_.segment_end();
+                bindings.restore();
+                weights.segment_end();
+                ggml_free(context);
+                runner.sync_runtime_residency();
+            }
+        } segment_cleanup{*this, weights, bindings, segment_context};
+
+        auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement;
+        if (!workspace_.prepare(measurement)) {
+            return fail_segment("workspace preparation");
+        }
+        const size_t cut_bytes       = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment);
+        const size_t new_cache_bytes = add_bytes(cut_bytes, cache_.pending_bytes(segment_graph));
+        auto ensure_capacity         = [&]() {
+            sync_runtime_residency();
+            auto requests = memory_requests(measurement.buffers, new_cache_bytes);
+            if (!fits(requests, weights.params(index)) && workspace_.release_excess(measurement)) {
+                sync_runtime_residency();
+                requests = memory_requests(measurement.buffers, new_cache_bytes);
+            }
+            return weights.ensure_segment_capacity(index, requests);
+        };
+        if (!weights.segment_start(index, ensure_capacity)) {
+            return fail_segment("weight preparation");
+        }
+        // Preparing weights can execute LoRA graphs and reclaim an idle workspace.
+        if (!workspace_.measurement_matches(segment_graph, measurement)) {
+            measurement = measure(segment_graph, segment.compute_buffer_size);
+        }
+        if (!workspace_.prepare(measurement) || !ensure_capacity()) {
+            return fail_segment("workspace capacity check");
+        }
+        if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
+                pin_multi_device_nodes(scheduler, current);
+            })) {
+            return fail_segment("workspace allocation");
+        }
+        for (const auto& size : measurement.buffers) {
+            track_compute_buffer(size.backend);
+        }
+        if (workspace_.scheduler() != nullptr) {
+            track_compute_buffer(workspace_.cpu_backend());
+        }
+        if (!ensure_capacity()) {
+            return fail_segment("allocated capacity check");
+        }
+        copy_data_to_backend_tensor(segment_graph, false);
+        auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes);
+        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())) {
+            return fail_segment("execution or output caching");
+        }
+        sync_runtime_residency();
+        if (last) {
+            if (read_outputs && !read_outputs()) {
+                return fail_segment("output finalization");
+            }
+            if (!no_return) {
+                auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str());
+                output      = read_graph_tensor(result, "output");
+                if (!output.has_value()) {
+                    return fail_segment("output readback");
+                }
+            }
+        }
+        // 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_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();
+    }
+    return output;
+}

src/core/ggml_runner.h

@@ -0,0 +1,350 @@
+#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;
+    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);
+    }
+};
+
+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;
+    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_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

@@ -145,19 +145,24 @@ namespace sd {
         std::vector<int64_t> backend_capacities = graph_cut_layer_split_backend_capacities(split_backends,
                                                                                            backend_vram_limits,
                                                                                            primary_backend_vram_limit);
+        // Existing placements may already occupy the reported free VRAM. Reuse
+        // them; execution checks missing weights and reclaims memory as needed.
+        const bool reuse_assignments = std::all_of(seen_params.begin(), seen_params.end(), [&](ggml_tensor* param) {
+            return param_assignments.count(param) != 0;
+        });
 
         std::vector<ggml_backend_t> backend_by_segment(plan.segments.size(), split_backends[0]);
         size_t current_backend = 0;
         int64_t current_used   = 0;
         for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
             int64_t bytes = segment_param_bytes[seg_idx];
-            while (current_backend + 1 < split_backends.size() &&
+            while (!reuse_assignments && current_backend + 1 < split_backends.size() &&
                    bytes > 0 &&
                    current_used + bytes > backend_capacities[current_backend]) {
                 current_backend++;
                 current_used = 0;
             }
-            if (bytes > 0 && current_used + bytes > backend_capacities[current_backend]) {
+            if (!reuse_assignments && bytes > 0 && current_used + bytes > backend_capacities[current_backend]) {
                 LOG_ERROR("%s graph-cut layer split: segment %zu needs %.1f MB on %s, but only %.1f MB is available under current VRAM limits",
                           desc,
                           seg_idx,
@@ -167,7 +172,6 @@ namespace sd {
                 return false;
             }
             current_used += bytes;
-            backend_by_segment[seg_idx] = split_backends[current_backend];
 
             for (ggml_tensor* param : segment_params[seg_idx]) {
                 ggml_backend_t target_backend = split_backends[current_backend];
@@ -186,12 +190,16 @@ namespace sd {
                               ggml_get_name(param));
                     return false;
                 }
-                size_t backend_idx                               = (size_t)std::distance(split_backends.begin(), backend_it);
+                size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it);
+                if (reuse_assignments) {
+                    current_backend = backend_idx;
+                }
                 assignment.first_segment_by_backend[backend_idx] = std::min(assignment.first_segment_by_backend[backend_idx], seg_idx);
                 assignment.last_segment_by_backend[backend_idx]  = std::max(assignment.last_segment_by_backend[backend_idx], seg_idx + 1);
                 assignment.tensors_by_backend[backend_idx].push_back(param);
                 assignment.bytes_by_backend[backend_idx] += (int64_t)ggml_nbytes(param);
             }
+            backend_by_segment[seg_idx] = split_backends[current_backend];
         }
 
         const int n_nodes = ggml_graph_n_nodes(gf);
@@ -243,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/runner_cache.cpp

@@ -0,0 +1,211 @@
+#include "core/runner_cache.h"
+
+#include <algorithm>
+#include <iterator>
+#include <unordered_set>
+
+#include "core/ggml_graph_cut.h"
+#include "core/util.h"
+
+namespace sd {
+    static std::unordered_set<const ggml_tensor*> cache_graph_tensors(ggml_cgraph* graph) {
+        std::unordered_set<const ggml_tensor*> tensors;
+        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+            tensors.insert(ggml_graph_node(graph, i));
+        }
+        for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) {
+            tensors.insert(ggml_graph_cut::leaf_tensor(graph, i));
+        }
+        return tensors;
+    }
+
+    CachedTensor::~CachedTensor() {
+        ggml_backend_buffer_free(buffer);
+        ggml_free(context);
+    }
+
+    std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend,
+                                                     const std::string& name,
+                                                     ggml_tensor* source) {
+        if (ggml_graph_cut::tensor_buffer(source) == nullptr) {
+            return nullptr;
+        }
+        auto entry     = std::make_unique<CachedTensor>();
+        entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true});
+        if (entry->context == nullptr) {
+            return nullptr;
+        }
+        entry->tensor = ggml_dup_tensor(entry->context, source);
+        // Cut views are rebound with their original strides and offsets.
+        std::copy(std::begin(source->nb), std::end(source->nb), std::begin(entry->tensor->nb));
+        ggml_set_name(entry->tensor, name.c_str());
+        entry->buffer = ggml_backend_alloc_ctx_tensors(entry->context, backend);
+        if (entry->buffer == nullptr) {
+            return nullptr;
+        }
+        if (source->view_src != nullptr || !ggml_is_contiguous(source) || source->buffer == nullptr) {
+            std::vector<uint8_t> data(ggml_nbytes(source));
+            ggml_backend_tensor_get(source, data.data(), 0, data.size());
+            ggml_backend_tensor_set(entry->tensor, data.data(), 0, data.size());
+        } else {
+            ggml_backend_tensor_copy(source, entry->tensor);
+        }
+        return entry;
+    }
+
+    static ggml_tensor* cached_tensor(const CachedTensors& tensors, const std::string& name) {
+        auto entry = tensors.find(name);
+        return entry == tensors.end() ? nullptr : entry->second->tensor;
+    }
+
+    static size_t resident_bytes(const CachedTensors& tensors, ggml_backend_dev_t device) {
+        size_t bytes = 0;
+        for (const auto& entry : tensors) {
+            auto buffer = entry.second->buffer;
+            if (!ggml_backend_buffer_is_host(buffer) &&
+                ggml_backend_buft_get_device(ggml_backend_buffer_get_type(buffer)) == device) {
+                const size_t size = ggml_backend_buffer_get_size(buffer);
+                bytes             = size > SIZE_MAX - bytes ? SIZE_MAX : bytes + size;
+            }
+        }
+        return bytes;
+    }
+
+    ggml_tensor* RunnerCache::get(const std::string& name) const {
+        return cached_tensor(committed_, name);
+    }
+
+    void RunnerCache::stage(const std::string& name, ggml_tensor* tensor) {
+        if (tensor != nullptr) {
+            ggml_set_output(tensor);
+            outputs_[name] = tensor;
+        }
+    }
+
+    size_t RunnerCache::pending_bytes(ggml_cgraph* graph) const {
+        if (outputs_.empty()) {
+            return 0;
+        }
+        auto tensors = cache_graph_tensors(graph);
+        auto buft    = ggml_backend_get_default_buffer_type(backend_);
+        size_t bytes = 0;
+        for (const auto& output : outputs_) {
+            if (pending_.count(output.first) || !tensors.count(output.second)) {
+                continue;
+            }
+            const size_t size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, output.second),
+                                         ggml_backend_buft_get_alignment(buft));
+            bytes             = size > SIZE_MAX - bytes ? SIZE_MAX : bytes + size;
+        }
+        return bytes;
+    }
+
+    size_t RunnerCache::resident_bytes(ggml_backend_dev_t device) const {
+        const size_t committed = sd::resident_bytes(committed_, device);
+        const size_t pending   = sd::resident_bytes(pending_, device);
+        return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending;
+    }
+
+    bool RunnerCache::capture(ggml_cgraph* graph) {
+        if (outputs_.empty()) {
+            return true;
+        }
+        const auto tensors = cache_graph_tensors(graph);
+        for (const auto& output : outputs_) {
+            if (pending_.count(output.first) || !tensors.count(output.second)) {
+                continue;
+            }
+            GGML_ASSERT(ggml_is_contiguous(output.second));
+            auto entry = CachedTensor::copy(backend_, output.first, output.second);
+            if (entry == nullptr) {
+                return false;
+            }
+            pending_[output.first] = std::move(entry);
+        }
+        ggml_backend_synchronize(backend_);
+        return true;
+    }
+
+    void RunnerCache::graph_end(bool success) {
+        // Graph inputs can still reference the previous generation until graph end.
+        if (success) {
+            for (auto& entry : pending_) {
+                committed_[entry.first] = std::move(entry.second);
+            }
+        }
+        pending_.clear();
+        outputs_.clear();
+    }
+
+    void RunnerCache::clear() {
+        graph_end(false);
+        committed_.clear();
+    }
+
+    ggml_tensor* GraphCutTensorCache::get(const std::string& name) const {
+        return cached_tensor(tensors_, name);
+    }
+
+    size_t GraphCutTensorCache::resident_bytes(ggml_backend_dev_t device) const {
+        return sd::resident_bytes(tensors_, device);
+    }
+
+    size_t GraphCutTensorCache::estimate_output_bytes(
+        ggml_cgraph* graph,
+        const ggml_graph_cut::Segment& segment) const {
+        ggml_backend_buffer_type_t buffer_type =
+            ggml_backend_get_default_buffer_type(backend_);
+        if (buffer_type == nullptr) {
+            return SIZE_MAX;
+        }
+        const size_t alignment = ggml_backend_buft_get_alignment(buffer_type);
+        size_t total_size      = 0;
+        for (size_t output_idx = 0; output_idx < segment.output_node_indices.size(); ++output_idx) {
+            ggml_tensor* output = ggml_graph_cut::output_tensor(graph, segment, output_idx);
+            if (output == nullptr || !ggml_graph_cut::is_graph_cut_tensor(output) ||
+                !segment.future_cut_names.count(output->name)) {
+                continue;
+            }
+            ggml_tensor* source      = ggml_graph_cut::cache_source_tensor(output);
+            const size_t tensor_size = GGML_PAD(
+                ggml_backend_buft_get_alloc_size(buffer_type, source), alignment);
+            total_size = tensor_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + tensor_size;
+        }
+        return total_size;
+    }
+
+    void GraphCutTensorCache::prune(const std::unordered_set<std::string>& keep_names) {
+        for (auto it = tensors_.begin(); it != tensors_.end();) {
+            it = keep_names.count(it->first) ? std::next(it) : tensors_.erase(it);
+        }
+    }
+
+    bool GraphCutTensorCache::capture(ggml_cgraph* graph,
+                                      const ggml_graph_cut::Segment& segment,
+                                      const char* log_desc) {
+        size_t copied_bytes = 0;
+        size_t copied_count = 0;
+        for (int index : segment.output_node_indices) {
+            auto output = ggml_graph_node(graph, index);
+            if (!ggml_graph_cut::is_graph_cut_tensor(output) ||
+                !segment.future_cut_names.count(output->name)) {
+                continue;
+            }
+            auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output));
+            if (entry == nullptr) {
+                LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name);
+                return false;
+            }
+            const size_t size = ggml_backend_buffer_get_size(entry->buffer);
+            copied_bytes      = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size;
+            ++copied_count;
+            tensors_[output->name] = std::move(entry);
+        }
+        ggml_backend_synchronize(backend_);
+        if (copied_count > 0) {
+            LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)",
+                      log_desc, copied_bytes / (1024.f * 1024.f), copied_count);
+        }
+        return true;
+    }
+}

src/core/runner_cache.h

@@ -0,0 +1,66 @@
+#ifndef __SD_CORE_RUNNER_CACHE_H__
+#define __SD_CORE_RUNNER_CACHE_H__
+
+#include <map>
+#include <memory>
+#include <string>
+#include <unordered_set>
+
+#include "ggml-backend.h"
+
+namespace sd::ggml_graph_cut {
+    struct Segment;
+}
+
+namespace sd {
+    struct CachedTensor {
+        ggml_context* context        = nullptr;
+        ggml_backend_buffer_t buffer = nullptr;
+        ggml_tensor* tensor          = nullptr;
+        ~CachedTensor();
+        static std::unique_ptr<CachedTensor> copy(ggml_backend_t backend,
+                                                  const std::string& name,
+                                                  ggml_tensor* source);
+    };
+    using CachedTensors = std::map<std::string, std::unique_ptr<CachedTensor>>;
+
+    class RunnerCache {
+        ggml_backend_t backend_;
+        CachedTensors committed_;
+        CachedTensors pending_;
+        std::map<std::string, ggml_tensor*> outputs_;
+
+    public:
+        explicit RunnerCache(ggml_backend_t backend)
+            : backend_(backend) {}
+        RunnerCache(const RunnerCache&)            = delete;
+        RunnerCache& operator=(const RunnerCache&) = delete;
+
+        ggml_tensor* get(const std::string& name) const;
+        void stage(const std::string& name, ggml_tensor* tensor);
+        const std::map<std::string, ggml_tensor*>& outputs() const { return outputs_; }
+        size_t pending_bytes(ggml_cgraph* graph) const;
+        size_t resident_bytes(ggml_backend_dev_t device) const;
+        bool capture(ggml_cgraph* graph);
+        void graph_end(bool success);
+        void clear();
+    };
+
+    class GraphCutTensorCache {
+        ggml_backend_t backend_;
+        CachedTensors tensors_;
+
+    public:
+        explicit GraphCutTensorCache(ggml_backend_t backend)
+            : backend_(backend) {}
+        ggml_tensor* get(const std::string& name) const;
+        size_t resident_bytes(ggml_backend_dev_t device) const;
+        size_t estimate_output_bytes(ggml_cgraph* graph,
+                                     const ggml_graph_cut::Segment& segment) const;
+        bool capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc);
+        void prune(const std::unordered_set<std::string>& keep_names);
+        void clear() { tensors_.clear(); }
+    };
+}
+
+#endif  // __SD_CORE_RUNNER_CACHE_H__

src/core/segment_graph_bindings.cpp

@@ -0,0 +1,137 @@
+#include "core/segment_graph_bindings.h"
+
+#include <algorithm>
+#include <iterator>
+
+#include "core/ggml_graph_cut.h"
+#include "core/runner_cache.h"
+#include "core/util.h"
+#include "ggml.h"
+
+namespace sd {
+    SegmentGraphBindings::SegmentGraphBindings(GraphCutTensorCache& tensor_cache,
+                                               const ggml_graph_cut::Plan& plan,
+                                               ggml_cgraph* graph)
+        : tensor_cache_(tensor_cache),
+          graph_(graph) {
+        GGML_ASSERT(graph_ != nullptr);
+        for (int i = 0; i < ggml_graph_n_nodes(graph_); ++i) {
+            ggml_tensor* tensor = ggml_graph_node(graph_, i);
+            Topology topology{tensor->op, {}, tensor->view_src, tensor->flags};
+            std::copy(std::begin(tensor->src), std::end(tensor->src), topology.sources.begin());
+            topology_[tensor] = topology;
+        }
+        for (const auto& segment : plan.segments) {
+            for (const auto& input : segment.input_refs) {
+                if (input.type != ggml_graph_cut::Segment::INPUT_EXTERNAL) {
+                    continue;
+                }
+                ggml_tensor* tensor = ggml_graph_cut::input_tensor(graph_, input);
+                if (tensor == nullptr || tensor->buffer == nullptr) {
+                    continue;
+                }
+                external_bindings_[tensor] = {tensor->buffer, tensor->data, tensor->extra};
+            }
+        }
+    }
+
+    void SegmentGraphBindings::reset(const ggml_graph_cut::Segment& segment) {
+        restore();
+        for (const auto& input : segment.input_refs) {
+            ggml_tensor* tensor = ggml_graph_cut::input_tensor(graph_, input);
+            if (tensor == nullptr) {
+                continue;
+            }
+            switch (input.type) {
+                case ggml_graph_cut::Segment::INPUT_PREVIOUS_CUT:
+                    tensor->buffer = nullptr;
+                    tensor->data   = nullptr;
+                    tensor->extra  = nullptr;
+                    break;
+                case ggml_graph_cut::Segment::INPUT_EXTERNAL: {
+                    auto binding = external_bindings_.find(tensor);
+                    if (binding != external_bindings_.end()) {
+                        tensor->buffer = binding->second.buffer;
+                        tensor->data   = binding->second.data;
+                        tensor->extra  = binding->second.extra;
+                    } else {
+                        tensor->buffer = nullptr;
+                        tensor->data   = nullptr;
+                        tensor->extra  = nullptr;
+                    }
+                    break;
+                }
+                case ggml_graph_cut::Segment::INPUT_PARAM:
+                    break;
+            }
+        }
+
+        for (int node_index : segment.internal_node_indices) {
+            ggml_tensor* node = ggml_graph_node(graph_, node_index);
+            if (node == nullptr) {
+                continue;
+            }
+            node->buffer = nullptr;
+            node->data   = nullptr;
+            node->extra  = nullptr;
+        }
+    }
+
+    void SegmentGraphBindings::restore() {
+        for (const auto& entry : topology_) {
+            entry.first->op       = entry.second.op;
+            entry.first->view_src = entry.second.view_source;
+            entry.first->flags    = entry.second.flags;
+            std::copy(entry.second.sources.begin(), entry.second.sources.end(), std::begin(entry.first->src));
+        }
+    }
+
+    bool SegmentGraphBindings::bind_cached_inputs(
+        const ggml_graph_cut::Segment& segment,
+        const char* log_desc) {
+        std::unordered_map<ggml_tensor*, ggml_tensor*> cached_view_sources;
+        for (const auto& input : segment.input_refs) {
+            if (input.type != ggml_graph_cut::Segment::INPUT_PREVIOUS_CUT) {
+                continue;
+            }
+            ggml_tensor* input_tensor = ggml_graph_cut::input_tensor(graph_, input);
+            if (input_tensor == nullptr) {
+                continue;
+            }
+            ggml_tensor* cached_tensor = tensor_cache_.get(input.display_name);
+            if (cached_tensor == nullptr) {
+                LOG_ERROR("%s missing graph cut cache tensor: %s",
+                          log_desc,
+                          input.display_name.c_str());
+                return false;
+            }
+            if (input_tensor->view_src != nullptr) {
+                cached_view_sources[topology_.at(input_tensor).view_source] = cached_tensor;
+                input_tensor->view_src                                      = cached_tensor;
+                input_tensor->buffer                                        = nullptr;
+                input_tensor->data                                          = cached_tensor->data == nullptr
+                                                                                  ? nullptr
+                                                                                  : static_cast<void*>(static_cast<char*>(cached_tensor->data) +
+                                                              input_tensor->view_offs);
+                input_tensor->extra                                         = cached_tensor->extra;
+            } else {
+                input_tensor->buffer = cached_tensor->buffer;
+                input_tensor->data   = cached_tensor->data;
+                input_tensor->extra  = cached_tensor->extra;
+            }
+            for (int source_index = 0; source_index < GGML_MAX_SRC; ++source_index) {
+                input_tensor->src[source_index] = nullptr;
+            }
+            input_tensor->op = GGML_OP_NONE;
+        }
+        // ggml flattens view chains, so descendants also need the cached root.
+        for (int node_index : segment.internal_node_indices) {
+            ggml_tensor* node  = ggml_graph_node(graph_, node_index);
+            auto cached_source = cached_view_sources.find(topology_.at(node).view_source);
+            if (cached_source != cached_view_sources.end()) {
+                node->view_src = cached_source->second;
+            }
+        }
+        return true;
+    }
+}

src/core/segment_graph_bindings.h

@@ -0,0 +1,52 @@
+#ifndef __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__
+#define __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__
+
+#include <array>
+#include <unordered_map>
+#include "ggml.h"
+
+struct ggml_backend_buffer;
+struct ggml_cgraph;
+struct ggml_tensor;
+
+namespace sd {
+    class GraphCutTensorCache;
+
+    namespace ggml_graph_cut {
+        struct Plan;
+        struct Segment;
+    }
+
+    class SegmentGraphBindings {
+    public:
+        SegmentGraphBindings(GraphCutTensorCache& tensor_cache,
+                             const ggml_graph_cut::Plan& plan,
+                             ggml_cgraph* graph);
+
+        void reset(const ggml_graph_cut::Segment& segment);
+        void restore();
+        ~SegmentGraphBindings() { restore(); }
+        bool bind_cached_inputs(const ggml_graph_cut::Segment& segment,
+                                const char* log_desc);
+
+    private:
+        struct ExternalBinding {
+            ggml_backend_buffer* buffer = nullptr;
+            void* data                  = nullptr;
+            void* extra                 = nullptr;
+        };
+
+        GraphCutTensorCache& tensor_cache_;
+        ggml_cgraph* graph_ = nullptr;
+        std::unordered_map<ggml_tensor*, ExternalBinding> external_bindings_;
+        struct Topology {
+            ggml_op op;
+            std::array<ggml_tensor*, GGML_MAX_SRC> sources;
+            ggml_tensor* view_source;
+            int flags;
+        };
+        std::unordered_map<ggml_tensor*, Topology> topology_;
+    };
+}
+
+#endif  // __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__

src/core/segment_weight_pipeline.cpp

@@ -0,0 +1,205 @@
+#include "core/segment_weight_pipeline.h"
+
+#include <utility>
+
+#include "core/ggml_graph_cut.h"
+#include "device_residency_manager.h"
+
+namespace sd {
+    static ggml_tensor* canonical_param(
+        ggml_tensor* tensor,
+        const std::unordered_set<const ggml_tensor*>& params) {
+        for (ggml_tensor* current = tensor; current != nullptr; current = current->view_src) {
+            if (params.find(current) != params.end()) {
+                return current;
+            }
+        }
+        return nullptr;
+    }
+
+    SegmentWeightPipeline::SegmentWeightPipeline(
+        const std::shared_ptr<DeviceResidencyManager>& residency_manager,
+        ggml_backend_t compute_backend,
+        uintptr_t owner_id,
+        ggml_cgraph* graph,
+        const ggml_graph_cut::Plan& plan,
+        const std::unordered_set<const ggml_tensor*>& params,
+        bool enabled)
+        : residency_manager_(residency_manager),
+          compute_backend_(compute_backend),
+          owner_id_(owner_id),
+          enabled_(enabled && residency_manager != nullptr) {
+        segment_params_.resize(plan.segments.size());
+        for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) {
+            std::unordered_set<ggml_tensor*> seen;
+            for (ggml_tensor* tensor :
+                 ggml_graph_cut::param_tensors(graph, plan.segments[segment_index])) {
+                ggml_tensor* param = canonical_param(tensor, params);
+                if (param != nullptr && seen.insert(param).second) {
+                    segment_params_[segment_index].push_back(param);
+                }
+            }
+        }
+    }
+
+    SegmentWeightPipeline::~SegmentWeightPipeline() {
+        segment_end();
+        clear();
+    }
+
+    size_t SegmentWeightPipeline::next_parameter_segment(size_t segment_index) const {
+        for (size_t next = segment_index + 1; next < segment_params_.size(); ++next) {
+            if (!segment_params_[next].empty()) {
+                return next;
+            }
+        }
+        return SIZE_MAX;
+    }
+
+    std::vector<std::vector<ggml_tensor*>> SegmentWeightPipeline::preferred_eviction_order() const {
+        return {segment_params_.rbegin(), segment_params_.rend()};
+    }
+
+    void SegmentWeightPipeline::disable() {
+        clear();
+        enabled_ = false;
+    }
+
+    void SegmentWeightPipeline::activate(size_t segment_index) {
+        if (!enabled_ || queued_segment_ == SIZE_MAX || queued_segment_ != segment_index) {
+            return;
+        }
+
+        auto manager = residency_manager_.lock();
+        if (manager == nullptr ||
+            !manager->activate_prefetched_params(owner_id_, queued_params_)) {
+            disable();
+            return;
+        }
+        queued_params_.clear();
+        queued_segment_ = SIZE_MAX;
+    }
+
+    bool SegmentWeightPipeline::ensure_segment_capacity(
+        size_t segment_index,
+        const std::vector<DeviceMemoryRequest>& requests) {
+        if (segment_index >= segment_params_.size()) {
+            return false;
+        }
+        auto manager = residency_manager_.lock();
+        if (manager == nullptr) {
+            return segment_params_[segment_index].empty();
+        }
+        std::vector<ggml_tensor*> protected_params = segment_params_[segment_index];
+        protected_params.insert(protected_params.end(), queued_params_.begin(), queued_params_.end());
+        for (const auto& request : requests) {
+            if (!manager->ensure_compute_backend_capacity(request, segment_params_[segment_index],
+                                                          preferred_eviction_order(), protected_params)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    bool SegmentWeightPipeline::segment_start(size_t segment_index, const std::function<bool()>& ensure_capacity) {
+        GGML_ASSERT(pinned_params_.empty());
+        activate(segment_index);
+        if (!ensure_capacity()) {
+            return false;
+        }
+        auto manager = residency_manager_.lock();
+        if (manager == nullptr) {
+            return segment_params_[segment_index].empty();
+        }
+        if (!manager->prepare_params(segment_params_[segment_index])) {
+            return false;
+        }
+        pinned_params_ = segment_params_[segment_index];
+        return true;
+    }
+
+    void SegmentWeightPipeline::segment_end() {
+        if (auto manager = residency_manager_.lock()) {
+            manager->release_compute_backend_params(pinned_params_);
+        }
+        pinned_params_.clear();
+    }
+
+    void SegmentWeightPipeline::enqueue_next(
+        size_t segment_index,
+        const DeviceMemoryRequest& request) {
+        if (!enabled_ || queued_segment_ != SIZE_MAX) {
+            return;
+        }
+
+        const size_t next_segment = next_parameter_segment(segment_index);
+        if (next_segment == SIZE_MAX) {
+            return;
+        }
+
+        std::unordered_set<ggml_tensor*> active_params(
+            segment_params_[segment_index].begin(),
+            segment_params_[segment_index].end());
+        std::vector<ggml_tensor*> params;
+        params.reserve(segment_params_[next_segment].size());
+        for (ggml_tensor* param : segment_params_[next_segment]) {
+            if (active_params.find(param) == active_params.end()) {
+                params.push_back(param);
+            }
+        }
+        if (params.empty()) {
+            return;
+        }
+
+        auto manager = residency_manager_.lock();
+        if (manager == nullptr) {
+            disable();
+            return;
+        }
+        const WeightResidencyInfo residency =
+            manager->inspect_compute_backend_params(params);
+        if (residency.missing_bytes == 0) {
+            return;
+        }
+        if (!residency.async_prefetch_supported) {
+            disable();
+            return;
+        }
+
+        DeviceMemoryRequest backend_request        = request;
+        backend_request.compute_backend            = compute_backend_;
+        backend_request.owner_id                   = owner_id_;
+        std::vector<ggml_tensor*> protected_params = segment_params_[segment_index];
+        protected_params.insert(protected_params.end(), params.begin(), params.end());
+        if (!manager->ensure_compute_backend_capacity(backend_request,
+                                                      params,
+                                                      preferred_eviction_order(),
+                                                      protected_params)) {
+            disable();
+            return;
+        }
+        switch (manager->prefetch_params(owner_id_, params)) {
+            case WeightPrefetchResult::Scheduled:
+                queued_params_  = std::move(params);
+                queued_segment_ = next_segment;
+                return;
+            case WeightPrefetchResult::AlreadyResident:
+                return;
+            case WeightPrefetchResult::Unsupported:
+                disable();
+                return;
+            case WeightPrefetchResult::Failed:
+                disable();
+                return;
+        }
+        disable();
+    }
+
+    void SegmentWeightPipeline::clear() {
+        if (auto manager = residency_manager_.lock()) {
+            manager->clear_prefetched_params(owner_id_);
+        }
+        queued_params_.clear();
+        queued_segment_ = SIZE_MAX;
+    }
+}

src/core/segment_weight_pipeline.h

@@ -0,0 +1,62 @@
+#ifndef __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__
+#define __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <unordered_set>
+#include <vector>
+
+#include "ggml-backend.h"
+
+struct DeviceMemoryRequest;
+struct DeviceResidencyManager;
+struct ggml_cgraph;
+struct ggml_tensor;
+
+namespace sd::ggml_graph_cut {
+    struct Plan;
+}
+
+namespace sd {
+    class SegmentWeightPipeline {
+    private:
+        std::weak_ptr<DeviceResidencyManager> residency_manager_;
+        ggml_backend_t compute_backend_ = nullptr;
+        uintptr_t owner_id_             = 0;
+        std::vector<std::vector<ggml_tensor*>> segment_params_;
+        std::vector<ggml_tensor*> queued_params_;
+        std::vector<ggml_tensor*> pinned_params_;
+        size_t queued_segment_ = SIZE_MAX;
+        bool enabled_          = true;
+
+        size_t next_parameter_segment(size_t segment_index) const;
+        std::vector<std::vector<ggml_tensor*>> preferred_eviction_order() const;
+        void disable();
+        void activate(size_t segment_index);
+        void clear();
+
+    public:
+        SegmentWeightPipeline(
+            const std::shared_ptr<DeviceResidencyManager>& residency_manager,
+            ggml_backend_t compute_backend,
+            uintptr_t owner_id,
+            ggml_cgraph* graph,
+            const ggml_graph_cut::Plan& plan,
+            const std::unordered_set<const ggml_tensor*>& params,
+            bool enabled = true);
+        ~SegmentWeightPipeline();
+
+        const std::vector<ggml_tensor*>& params(size_t index) const { return segment_params_[index]; }
+        bool ensure_segment_capacity(size_t segment_index,
+                                     const std::vector<DeviceMemoryRequest>& requests);
+        bool segment_start(size_t segment_index, const std::function<bool()>& ensure_capacity);
+        void segment_end();
+        // Prefetch is best effort; segment_start falls back to synchronous loading.
+        void enqueue_next(size_t segment_index,
+                          const DeviceMemoryRequest& request);
+    };
+}
+
+#endif  // __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__

src/core/tensor_ggml.hpp

@@ -9,6 +9,7 @@
 #include <type_traits>
 
 #include "core/tensor.hpp"
+#include "ggml-backend.h"
 #include "ggml.h"
 
 namespace sd {
@@ -54,10 +55,28 @@ namespace sd {
             GGML_ABORT("ggml tensor type does not match sd::Tensor type");
         }
         Tensor<T> result(shape_from_ggml(tensor));
-        if (tensor->buffer != nullptr) {
-            ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor));
+        std::vector<uint8_t> strided_data;
+        void* destination = result.data();
+        if (!ggml_is_contiguous(tensor)) {
+            strided_data.resize(ggml_nbytes(tensor));
+            destination = strided_data.data();
+        }
+        auto buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
+        if (buffer != nullptr) {
+            ggml_backend_tensor_get(tensor, destination, 0, ggml_nbytes(tensor));
         } else {
-            std::memcpy(result.data(), tensor->data, ggml_nbytes(tensor));
+            std::memcpy(destination, tensor->data, ggml_nbytes(tensor));
+        }
+        if (!strided_data.empty()) {
+            for (int64_t i = 0; i < result.numel(); ++i) {
+                int64_t index = i;
+                size_t offset = 0;
+                for (int d = 0; d < GGML_MAX_DIMS; ++d) {
+                    offset += static_cast<size_t>(index % tensor->ne[d]) * tensor->nb[d];
+                    index /= tensor->ne[d];
+                }
+                std::memcpy(result.data() + i, strided_data.data() + offset, sizeof(T));
+            }
         }
         return result;
     }

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;

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
@@ -715,7 +716,7 @@ std::vector<ADetailerDetection> ADetailerGGML::predict(sd_image_t image,
     LetterboxInput input  = make_letterbox_input(image, params.input_size);
     int64_t start         = ggml_time_ms();
     sd::Tensor<float> raw = detector->compute(n_threads, input.tensor);
-    detector->free_compute_buffer();
+    detector->runner_end();
     if (raw.empty()) {
         LOG_ERROR("YOLOv8 detector inference failed");
         return {};

src/device_residency_manager.h

@@ -0,0 +1,77 @@
+#ifndef __DEVICE_RESIDENCY_MANAGER_H__
+#define __DEVICE_RESIDENCY_MANAGER_H__
+
+#include <cstdint>
+#include <functional>
+#include <vector>
+
+#include "ggml-backend.h"
+
+struct ggml_tensor;
+
+enum class WeightPrefetchResult {
+    Scheduled,
+    AlreadyResident,
+    Unsupported,
+    Failed,
+};
+
+struct WeightResidencyInfo {
+    bool async_prefetch_supported = false;
+    size_t missing_bytes          = 0;
+};
+
+struct DeviceMemoryRequest {
+    ggml_backend_t compute_backend  = nullptr;
+    uintptr_t owner_id              = 0;
+    size_t pending_allocation_bytes = 0;
+    size_t runtime_resident_bytes   = 0;
+    size_t max_backend_bytes        = 0;
+
+    // Runtime buffers only; the manager accounts for weights separately.
+    size_t runtime_peak_bytes() const {
+        return pending_allocation_bytes > SIZE_MAX - runtime_resident_bytes
+                   ? SIZE_MAX
+                   : runtime_resident_bytes + pending_allocation_bytes;
+    }
+};
+
+struct DeviceResidencyManager {
+    virtual ~DeviceResidencyManager() = default;
+
+    virtual bool segmented_compute_enabled() const                                          = 0;
+    virtual bool prefetch_enabled() const                                                   = 0;
+    virtual void set_workspace_reclaimer(uintptr_t owner_id, std::function<bool()> reclaim) = 0;
+    virtual void remove_runtime_owner(uintptr_t owner_id)                                   = 0;
+    // Capacity requests select their backend's weights; protection spans all backends.
+    virtual bool fits_compute_backend_capacity(const DeviceMemoryRequest& request,
+                                               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 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,
+                                          ggml_backend_t compute_backend,
+                                          size_t resident_bytes) = 0;
+    virtual bool ensure_compute_backend_capacity(
+        const DeviceMemoryRequest& request,
+        const std::vector<ggml_tensor*>& required_params,
+        const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
+        const std::vector<ggml_tensor*>& protected_params) = 0;
+    virtual WeightPrefetchResult prefetch_params(
+        uintptr_t owner_id,
+        const std::vector<ggml_tensor*>& tensors)                                     = 0;
+    virtual bool activate_prefetched_params(uintptr_t owner_id,
+                                            const std::vector<ggml_tensor*>& tensors) = 0;
+    virtual void clear_prefetched_params(uintptr_t owner_id)                          = 0;
+};
+
+// Transitional alias for model constructors that have not yet adopted the
+// residency-oriented name. It does not introduce a second implementation.
+using RunnerWeightManager = DeviceResidencyManager;
+
+#endif  // __DEVICE_RESIDENCY_MANAGER_H__

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;
     }
@@ -49,7 +50,7 @@ struct GenerationExtension {
     virtual void get_param_tensors(std::map<std::string, ggml_tensor*>&) {}
     virtual void collect_loras(std::vector<ModelManager::LoraSpec>&) {}
     virtual void add_ignore_tensors(std::set<std::string>&) const {}
-    virtual void runner_done() {}
+    virtual void runner_end() {}
     virtual void reset_runtime_condition() {}
     virtual bool prepare_condition(GenerationExtensionConditionContext&) {
         return false;

src/extensions/photomaker_extension.cpp

@@ -1,3 +1,4 @@
+#include <cinttypes>
 #include "extensions/generation_extension.h"
 
 #include <algorithm>
@@ -108,8 +109,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 +119,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,13 +128,7 @@ 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,
@@ -175,9 +170,9 @@ struct PhotoMakerExtension : public GenerationExtension {
         ignore_tensors.insert("pmid.unet.");
     }
 
-    void runner_done() override {
+    void runner_end() override {
         if (pmid_model != nullptr) {
-            pmid_model->runner_done();
+            pmid_model->runner_end();
         }
     }
 

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/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 {
@@ -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, true, 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,96 +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_done();
-        free_compute_buffer();
-        model_manager.reset();
+        runner_end();
+        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>();
-        weight_manager = model_manager;
         lora_tensors.clear();
+        scalar_values.clear();
         original_tensor_to_final_tensor.clear();
         applied_lora_tensors.clear();
         skipped_incompatible_lora_tensors.clear();
@@ -236,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);
                 }
             }
@@ -390,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);
             }
@@ -503,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);
             }
@@ -664,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;
                 }
 
@@ -791,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;
@@ -944,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,
@@ -952,20 +933,24 @@ struct LoraModel : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_lora_graph(model_tensors, model_tensor_names, version);
         };
-        GGMLRunner::compute<float>(get_graph, n_threads, false, false, false, true);
-        stat(!warn_unused);
-        for (auto item : original_tensor_to_final_tensor) {
-            ggml_tensor* original_tensor = item.first;
-            ggml_tensor* final_tensor    = item.second;
-
-            ggml_backend_tensor_copy(final_tensor, original_tensor);
+        auto read_outputs = [&]() {
+            for (const auto& item : original_tensor_to_final_tensor) {
+                ggml_backend_tensor_copy(item.second, item.first);
+            }
+            return true;
+        };
+        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();
-        GGMLRunner::free_compute_buffer();
+        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) {
@@ -1072,6 +1057,34 @@ public:
         return out;
     }
 
+    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,
+                                    WeightAdapter::ForwardParams forward_params) override {
+        for (auto& lora_model : lora_models) {
+            ggml_tensor* weight_diff = lora_model->get_weight_diff(prefix + "weight", backend, ctx, w, false);
+            if (weight_diff != nullptr) {
+                GGML_ASSERT(forward_params.op_type == ForwardParams::op_type_t::OP_LINEAR);
+                ggml_tensor* out_diff = ggml_ext_linear(ctx,
+                                                        x,
+                                                        weight_diff,
+                                                        nullptr,
+                                                        forward_params.linear.force_prec_f32,
+                                                        forward_params.linear.scale);
+                output                = ggml_add_inplace(ctx, output, out_diff);
+            }
+
+            ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight");
+            if (out_diff != nullptr) {
+                output = ggml_add_inplace(ctx, output, out_diff);
+            }
+        }
+        return output;
+    }
+
     size_t get_extra_graph_size() override {
         size_t lora_tensor_num = 0;
         for (auto& lora_model : lora_models) {

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, true, 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:

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:
@@ -268,10 +270,11 @@ public:
                 int64_t dim_out,
                 int64_t mult          = 4,
                 Activation activation = Activation::GEGLU,
-                bool precision_fix    = false) {
+                bool precision_fix    = false,
+                bool bias             = true) {
         int64_t inner_dim = dim * mult;
         if (activation == Activation::GELU) {
-            blocks["net.0"] = std::shared_ptr<GGMLBlock>(new GELU(dim, inner_dim));
+            blocks["net.0"] = std::shared_ptr<GGMLBlock>(new GELU(dim, inner_dim, bias));
         } else {
             blocks["net.0"] = std::shared_ptr<GGMLBlock>(new GEGLU(dim, inner_dim));
         }
@@ -285,7 +288,7 @@ public:
         // The purpose of the scale here is to prevent NaN issues in certain situations.
         // For example, when using Vulkan without enabling force_prec_f32,
         // or when using CUDA but the weights are k-quants.
-        blocks["net.2"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim_out, true, false, force_prec_f32, scale));
+        blocks["net.2"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim_out, bias, false, force_prec_f32, scale));
     }
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
@@ -339,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;
         }
@@ -369,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]

src/model/common/ggml_block.hpp

@@ -0,0 +1,880 @@
+#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"];
+        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 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);
+    }
+};
+
+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->ggml_ctx, ctx->backend, 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 {

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, false, 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;
         }
@@ -717,7 +718,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, false, 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 {

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, false, 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);
     }
 
@@ -423,54 +418,28 @@ struct ControlNet : public GGMLRunner {
             return build_graph(x, hint, timesteps, context, y);
         };
 
-        auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false, false, false, true);
+        auto read_outputs = [&]() {
+            controls.clear();
+            controls.reserve(control_outputs_ggml.size());
+            for (ggml_tensor* control : control_outputs_ggml) {
+                auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(control), 4);
+                if (control_host.empty()) {
+                    return false;
+                }
+                controls.push_back(std::move(control_host));
+            }
+            return true;
+        };
+        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()) {
+            controls.clear();
             return std::nullopt;
         }
-
         guided_hint_cached = get_cache_tensor_by_name(guided_hint_cache_name()) != nullptr;
-        controls.clear();
-        controls.reserve(control_outputs_ggml.size());
-        for (ggml_tensor* control : control_outputs_ggml) {
-            auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(control), 4);
-            GGML_ASSERT(!control_host.empty());
-            controls.push_back(std::move(control_host));
-        }
         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>(weight_manager.lock());
-        if (manager == nullptr) {
-            owned_model_manager = std::make_shared<ModelManager>();
-            weight_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;
         }
     };
@@ -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, false, 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, false, 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

@@ -325,13 +325,11 @@ namespace HiDreamO1 {
 
         sd::Tensor<float> compute(int n_threads,
                                   const sd::Tensor<float>& image,
-                                  bool auto_free           = true,
-                                  bool free_compute_buffer = true,
-                                  bool free_compute_params = true) {
+                                  bool auto_runner_end = true) {
             auto get_graph = [&]() {
                 return build_graph(image);
             };
-            auto output = GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params);
+            auto output = GGMLRunner::compute(get_graph, n_threads, auto_runner_end);
             return output.has_value() ? std::move(output.value()) : sd::Tensor<float>();
         }
     };
@@ -459,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, false, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -510,8 +508,8 @@ namespace HiDreamO1 {
             vision_runner->set_weight_adapter(adapter);
         }
 
-        void runner_done() override {
-            vision_runner->runner_done();
+        void runner_end() override {
+            vision_runner->runner_end();
         }
 
         SDCondition get_learned_condition(int n_threads,
@@ -659,7 +657,7 @@ namespace HiDreamO1 {
             result.c_vinput_mask  = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
             result.c_image_embeds.reserve(vlm_images.size());
             for (const auto& vlm_image : vlm_images) {
-                auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false, true, true);
+                auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false);
                 if (image_embed.empty()) {
                     LOG_ERROR("hidream_o1 conditioner: encode VLM image failed");
                     return SDCondition();

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"
@@ -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, false, 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;
         }
@@ -142,7 +146,7 @@ namespace Ideogram4 {
     __STATIC_INLINE__ std::shared_ptr<Linear> make_linear(int64_t in_features,
                                                           int64_t out_features,
                                                           bool bias = true) {
-        return std::make_shared<Linear>(in_features, out_features, bias, false, false, 1.f, true);
+        return std::make_shared<Linear>(in_features, out_features, bias);
     }
 
     __STATIC_INLINE__ std::vector<float> gen_ideogram4_pe(int grid_h,
@@ -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, false, 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;
         }
     };
@@ -775,7 +778,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, false, 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, false, 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, false, 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"
@@ -129,6 +132,10 @@ namespace LTXV {
         bool self_attention_gated  = false;
         bool cross_attention_gated = false;
 
+        bool ff_bias                         = true;
+        bool audio_ff_bias                   = true;
+        bool use_keyframes_abs_pos_embedding = false;
+
         static std::pair<int64_t, int64_t> infer_attention_layout(int64_t hidden_size,
                                                                   int64_t preferred_heads = -1) {
             if (preferred_heads > 0 && hidden_size % preferred_heads == 0) {
@@ -207,6 +214,19 @@ namespace LTXV {
                 tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_gate_logits.weight") != tensor_storage_map.end()) {
                 config.cross_attention_gated = true;
             }
+            // LTX 2.5 sets ff_bias=false but leaves audio_ff_bias at its default, so the two
+            // branches must be detected separately; older checkpoints ship both sets of biases.
+            if (tensor_storage_map.find(prefix + ".transformer_blocks.0.ff.net.0.proj.bias") == tensor_storage_map.end() &&
+                tensor_storage_map.find(prefix + ".transformer_blocks.0.ff.net.2.bias") == tensor_storage_map.end()) {
+                config.ff_bias = false;
+            }
+            if (tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_ff.net.0.proj.bias") == tensor_storage_map.end() &&
+                tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_ff.net.2.bias") == tensor_storage_map.end()) {
+                config.audio_ff_bias = false;
+            }
+            if (tensor_storage_map.find(prefix + ".keyframes_abs_pos_embedding") != tensor_storage_map.end()) {
+                config.use_keyframes_abs_pos_embedding = true;
+            }
             if (tensor_storage_map.find(prefix + ".caption_projection.linear_1.weight") == tensor_storage_map.end() &&
                 tensor_storage_map.find(prefix + ".caption_projection.linear_2.weight") == tensor_storage_map.end()) {
                 config.use_caption_projection = false;
@@ -257,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;
         }
     };
@@ -874,8 +894,7 @@ namespace LTXV {
                          const String2TensorStorage& tensor_storage_map = {},
                          const std::string prefix                       = "") override {
             if (num_learnable_registers > 0) {
-                ggml_type wtype               = get_type(prefix + "learnable_registers", tensor_storage_map, GGML_TYPE_F32);
-                params["learnable_registers"] = ggml_new_tensor_2d(ctx, wtype, hidden_size, num_learnable_registers);
+                params["learnable_registers"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_learnable_registers);
             }
         }
 
@@ -1130,7 +1149,9 @@ namespace LTXV {
                                 int64_t a_context_dim,
                                 bool apply_gated_attention,
                                 bool cross_attention_adaln,
-                                bool video_rope_interleaved)
+                                bool video_rope_interleaved,
+                                bool ff_bias       = true,
+                                bool audio_ff_bias = true)
             : v_dim(v_dim),
               a_dim(a_dim),
               cross_attention_adaln(cross_attention_adaln) {
@@ -1140,8 +1161,8 @@ namespace LTXV {
             blocks["audio_attn2"]         = std::make_shared<CrossAttention>(a_dim, a_context_dim, a_heads, ad_head, apply_gated_attention, false);
             blocks["audio_to_video_attn"] = std::make_shared<CrossAttention>(v_dim, a_dim, a_heads, ad_head, apply_gated_attention, false);
             blocks["video_to_audio_attn"] = std::make_shared<CrossAttention>(a_dim, v_dim, a_heads, ad_head, apply_gated_attention, false);
-            blocks["ff"]                  = std::make_shared<FeedForward>(v_dim, v_dim, 4, FeedForward::Activation::GELU);
-            blocks["audio_ff"]            = std::make_shared<FeedForward>(a_dim, a_dim, 4, FeedForward::Activation::GELU);
+            blocks["ff"]                  = std::make_shared<FeedForward>(v_dim, v_dim, 4, FeedForward::Activation::GELU, false, ff_bias);
+            blocks["audio_ff"]            = std::make_shared<FeedForward>(a_dim, a_dim, 4, FeedForward::Activation::GELU, false, audio_ff_bias);
         }
 
         std::vector<ggml_tensor*> get_ada_values(GGMLRunnerContext* ctx,
@@ -1320,6 +1341,12 @@ namespace LTXV {
                                                                    get_type(prefix + "audio_scale_shift_table", tensor_storage_map, GGML_TYPE_F32),
                                                                    config.audio_hidden_size,
                                                                    2);
+            if (config.use_keyframes_abs_pos_embedding) {
+                params["keyframes_abs_pos_embedding"] = ggml_new_tensor_2d(ctx,
+                                                                           get_type(prefix + "keyframes_abs_pos_embedding", tensor_storage_map, GGML_TYPE_F32),
+                                                                           config.hidden_size,
+                                                                           1);
+            }
         }
 
         LTXAVModelBlock(const LTXAVConfig& config)
@@ -1386,7 +1413,9 @@ namespace LTXV {
                                                                                                               config.audio_cross_attention_dim,
                                                                                                               config.self_attention_gated || config.cross_attention_gated,
                                                                                                               config.cross_attention_adaln,
-                                                                                                              config.video_rope_interleaved);
+                                                                                                              config.video_rope_interleaved,
+                                                                                                              config.ff_bias,
+                                                                                                              config.audio_ff_bias);
             }
 
             blocks["norm_out"]       = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
@@ -1534,6 +1563,38 @@ namespace LTXV {
             return {v_context, a_context};
         }
 
+        // The video encoder is causal, so the first latent frame covers a single pixel frame while
+        // every later one covers temporal_scale_factor. LTX 2.5 marks that token class with a
+        // learned embedding added right after patchify_proj.
+        ggml_tensor* apply_keyframes_abs_pos_embedding(GGMLRunnerContext* ctx,
+                                                       ggml_tensor* vx,
+                                                       int64_t tokens_per_latent_frame) {
+            if (!config.use_keyframes_abs_pos_embedding || params.count("keyframes_abs_pos_embedding") == 0) {
+                return vx;
+            }
+            int64_t tokens = vx->ne[1];
+            if (tokens_per_latent_frame <= 0 || tokens_per_latent_frame > tokens) {
+                return vx;
+            }
+            auto embedding = params["keyframes_abs_pos_embedding"];
+            auto first     = ggml_cont(ctx->ggml_ctx,
+                                       ggml_view_3d(ctx->ggml_ctx, vx, vx->ne[0], tokens_per_latent_frame, vx->ne[2], vx->nb[1], vx->nb[2], 0));
+            first          = ggml_add(ctx->ggml_ctx, first, embedding);
+            if (tokens_per_latent_frame == tokens) {
+                return first;
+            }
+            auto rest = ggml_cont(ctx->ggml_ctx,
+                                  ggml_view_3d(ctx->ggml_ctx,
+                                               vx,
+                                               vx->ne[0],
+                                               tokens - tokens_per_latent_frame,
+                                               vx->ne[2],
+                                               vx->nb[1],
+                                               vx->nb[2],
+                                               tokens_per_latent_frame * vx->nb[1]));
+            return ggml_concat(ctx->ggml_ctx, first, rest, 1);
+        }
+
         std::vector<ggml_tensor*> get_output_scale_shift(GGMLRunnerContext* ctx,
                                                          ggml_tensor* table,
                                                          ggml_tensor* embedded_timestep,
@@ -1575,6 +1636,7 @@ namespace LTXV {
 
             vx = patchify_video(ctx, vx, n);
             vx = patchify_proj->forward(ctx, vx);
+            vx = apply_keyframes_abs_pos_embedding(ctx, vx, width * height);
             if (ax != nullptr && ggml_nelements(ax) > 0 && audio_time > 0) {
                 ax = patchify_audio(ctx, ax);
                 ax = audio_patchify_proj->forward(ctx, ax);
@@ -1939,7 +2001,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, false, false), x.dim());
+            auto out = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
             return out;
         }
 
@@ -2011,7 +2073,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,
@@ -2025,8 +2087,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;
@@ -2045,7 +2107,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, false, 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;
         }
     };
@@ -123,13 +125,6 @@ namespace MiniMaxH3 {
         return to_shift * base / (1.f + (to_shift - 1.f) * base);
     }
 
-    static float time_shift_slope(float sigma, float from_shift, float to_shift) {
-        float base = sigma / (from_shift + sigma * (1.f - from_shift));
-        float a    = 1.f + (from_shift - 1.f) * base;
-        float b    = 1.f + (to_shift - 1.f) * base;
-        return to_shift * a * a / (from_shift * b * b);
-    }
-
     struct TimeEmbedder : public GGMLBlock {
         TimeEmbedder(int64_t input_dim, int64_t hidden_dim, int64_t output_dim) {
             blocks["proj_in"]  = std::make_shared<Linear>(input_dim, hidden_dim, true, true);
@@ -594,8 +589,7 @@ namespace MiniMaxH3 {
                                                       const std::vector<TokenModulationSpan>& segments,
                                                       const std::vector<SequenceSegment>& sequence_segments,
                                                       const TokenModulationSpan& video_segment,
-                                                      const TokenModulationSpan& audio_segment,
-                                                      float audio_slope) {
+                                                      const TokenModulationSpan& audio_segment) {
             auto video_proj = std::dynamic_pointer_cast<Linear>(blocks["video_patch_proj"]);
             auto audio_proj = std::dynamic_pointer_cast<Linear>(blocks["audio_patch_proj"]);
 
@@ -715,7 +709,7 @@ namespace MiniMaxH3 {
                                                audio->ne[2]);
             audio_out        = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, audio_out, 1, 2, 0, 3));
             video_out        = ggml_ext_scale(ctx->ggml_ctx, video_out, -1.f);
-            audio_out        = ggml_ext_scale(ctx->ggml_ctx, audio_out, -audio_slope);
+            audio_out        = ggml_ext_scale(ctx->ggml_ctx, audio_out, -1.f);
             return {video_out, audio_out};
         }
     };
@@ -1040,9 +1034,9 @@ namespace MiniMaxH3 {
             GGML_ASSERT(!audio_input_cache.empty());
             GGML_ASSERT(!context_tensor.empty());
 
-            auto video   = make_input(video_input_cache);
-            auto audio   = make_input(audio_input_cache);
-            auto context = make_input(context_tensor);
+            auto video         = make_input(video_input_cache);
+            auto audio_carrier = make_input(audio_input_cache);
+            auto context       = make_input(context_tensor);
             std::vector<ggml_tensor*> condition_inputs;
             condition_inputs.reserve(condition_videos.size());
             for (const auto& condition : condition_videos) {
@@ -1054,21 +1048,26 @@ namespace MiniMaxH3 {
                 audio_condition_inputs.push_back(make_input(condition));
             }
 
-            float sigma_v = std::clamp(timestep[0] / 1000.f, 1e-6f, 1.f);
-            float t_v     = 1.f - sigma_v;
-            float t_a     = 1.f - time_shift_sigma(sigma_v, video_shift, audio_shift);
-            auto layout   = build_layout(context_tensor.shape()[1],
-                                         video_input_cache.shape()[2],
-                                         video_input_cache.shape()[1],
-                                         video_input_cache.shape()[0],
-                                         audio_length,
-                                         condition_videos,
-                                         condition_audios,
-                                         keyframe_indices,
-                                         reference_blocks,
-                                         text_tags,
-                                         t_v,
-                                         t_a);
+            float sigma_v     = std::clamp(timestep[0] / 1000.f, 1e-6f, 1.f);
+            float sigma_a     = time_shift_sigma(sigma_v, video_shift, audio_shift);
+            float audio_scale = video_shift / audio_shift;
+            float t_v         = 1.f - sigma_v;
+            float t_a         = 1.f - sigma_a;
+            // The sampler carries c_a = (sigma_v / sigma_a) * x_a so the packed
+            // latent follows one sigma schedule. Restore x_a for the H3 network.
+            auto audio  = ggml_ext_scale(compute_ctx, audio_carrier, sigma_a / sigma_v);
+            auto layout = build_layout(context_tensor.shape()[1],
+                                       video_input_cache.shape()[2],
+                                       video_input_cache.shape()[1],
+                                       video_input_cache.shape()[0],
+                                       audio_length,
+                                       condition_videos,
+                                       condition_audios,
+                                       keyframe_indices,
+                                       reference_blocks,
+                                       text_tags,
+                                       t_v,
+                                       t_a);
 
             position_input_cache = sd::Tensor<float>(
                 {3, static_cast<int64_t>(layout.positions.size() / 3)},
@@ -1129,10 +1128,15 @@ namespace MiniMaxH3 {
                                             layout.segments,
                                             layout.sequence_segments,
                                             layout.video_segment,
-                                            layout.audio_segment,
-                                            time_shift_slope(sigma_v, video_shift, audio_shift));
-            auto merged     = merge_av_latents(compute_ctx, output.first, output.second);
-            auto graph      = new_graph_custom(H3_GRAPH_SIZE);
+                                            layout.audio_segment);
+            // Convert the model's audio velocity to d(c_a) / d(sigma_v).
+            output.second = ggml_add(compute_ctx,
+                                     ggml_ext_scale(compute_ctx, audio, 1.f - audio_scale),
+                                     ggml_ext_scale(compute_ctx,
+                                                    output.second,
+                                                    1.f + (audio_scale - 1.f) * sigma_a));
+            auto merged   = merge_av_latents(compute_ctx, output.first, output.second);
+            auto graph    = new_graph_custom(H3_GRAPH_SIZE);
             ggml_build_forward_expand(graph, merged);
             return graph;
         }
@@ -1164,11 +1168,9 @@ namespace MiniMaxH3 {
                                    extra->video_sigma_shift,
                                    extra->audio_sigma_shift);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
-                                                                              n_threads,
-                                                                              false,
-                                                                              false,
-                                                                              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, false, 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;
     }
@@ -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, false, 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"

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, false, 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, false, 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/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, false, 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,12 @@
 #ifndef __SD_MODEL_DIFFUSION_WAN_HPP__
 #define __SD_MODEL_DIFFUSION_WAN_HPP__
 
+#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"
@@ -75,13 +78,13 @@ namespace WAN {
                     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;
         }
     };
@@ -909,7 +912,7 @@ namespace WAN {
                                            config.theta,
                                            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);
@@ -950,7 +953,7 @@ namespace WAN {
                 return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -1007,7 +1010,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 +1020,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 +1040,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/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, false, 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]
         }
@@ -568,13 +575,11 @@ struct CLIPTextModelRunner : public GGMLRunner {
                               size_t max_token_idx,
                               bool return_pooled,
                               int clip_skip,
-                              bool auto_free           = true,
-                              bool free_compute_buffer = true,
-                              bool free_compute_params = true) {
+                              bool auto_runner_end = true) {
         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_free, free_compute_buffer, free_compute_params);
+        auto result = GGMLRunner::compute(get_graph, n_threads, auto_runner_end);
         if (return_pooled) {
             return take_or_empty(std::move(result));
         }