refactor: align PixArt weights with upstream layout (#2061)

README.md

@@ -64,6 +64,7 @@ API and command-line option may change frequently.***
     - [HiDream-O1-Image](./docs/hidream_o1_image.md)
     - [Ideogram4](./docs/ideogram4.md)
     - [LLaDA-Image](./docs/llada_image.md)
+    - [PixArt](./docs/pixart.md)
   - [Image Edit Models](./docs/edit.md)
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)

docs/performance.md

@@ -21,6 +21,54 @@ 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.
 
+## Use VAE tiling to reduce encode and decode memory usage.
+
+`--vae-tiling` enables spatial tiling for both VAE encoding and decoding. The
+default tile size is 256x256 **image pixels**, independent of the VAE scale factor:
+
+```shell
+--vae-tiling --vae-tile-size 256x256 --vae-tile-overlap 0.5
+```
+
+`--vae-tile-size` accepts one size or `WIDTHxHEIGHT`. A zero dimension uses the
+256-pixel default. Sizes are rounded down to a multiple of the VAE scale factor
+and capped at the current input dimensions. Explicit sizes below four latent
+pixels per axis (or the full axis when it is smaller) are rejected. Encoding and
+decoding use the same spatial sizes, without an additional encoding multiplier.
+Inputs that fit within a tile are processed as one tile.
+
+For a 512x512 image with the default 50% overlap, both encoding and decoding use
+3x3 tiles. A 256-pixel tile corresponds to 32 latent pixels for an 8x VAE, 16 for
+a 16x VAE, and 8 for a 32x VAE. Smaller tiles reduce each graph's memory demand,
+but overlapping work can increase processing time and tiling can affect image
+quality, especially during encoding. Use larger tiles when more context is needed.
+
+`--vae-relative-tile-size` overrides the absolute size on each axis with a positive
+value. Values up to and including 1 specify a fraction of the current input size;
+values greater than 1 specify a target number of tiles per axis, accounting for
+overlap. For example, `0.5x0.5` uses half the width and height in both encode and
+decode. The target overlap is clamped to 0 through 0.5 and the actual overlap is
+adjusted to fit the image. Size and overlap options require `--vae-tiling`.
+
+**Migration:** `--vae-tile-size` and the C/JSON fields `tile_size_w` and
+`tile_size_h` now use image pixels instead of latent units. The C/JSON fields
+`tile_size_x/y` have been renamed to `tile_size_w/h`, and `rel_size_x/y` to
+`rel_size_w/h`. The command-line option names are unchanged. For example, an old
+decode tile size of 32 corresponds to 256 pixels for an 8x VAE or 512 pixels for a
+16x VAE. Encoding no longer enlarges explicit or relative tile sizes.
+
+The main VAE decode path retries allocation failures with smaller tiles, even
+without `--vae-tiling`. Supported video VAEs first try temporal tiling; spatial
+retries use at most 256-pixel tiles initially and then halve the effective tile
+dimensions down to the minimum size. Each spatial retry must reduce the effective
+tile size. These runtime adjustments do not change the caller's parameters.
+Execution failures are not retried, and encoding has no automatic OOM retry.
+
+`--temporal-tiling` remains independent of spatial tiling. MiniMax H3 always uses
+spatial tiling (256x256 pixels and 25% overlap by default) and its own temporal
+windows. With `--vae-tiling`, its overlap follows `--vae-tile-overlap`; explicit
+spatial sizes are honored.
+
 ## Offload weights to the CPU to save VRAM without reducing generation speed.
 
 Using `--offload-to-cpu` allows you to offload weights to the CPU, saving VRAM without reducing generation speed.
@@ -63,6 +111,12 @@ See [backend selection](./backend.md) for full syntax.
 
 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.
 
+When choosing between monolithic and segmented execution, the runner requires
+an additional 128 MiB of headroom in both available device memory and any explicit
+managed budget. This planning headroom absorbs small allocation estimate changes;
+subsequent capacity checks can consume it while still preserving the 512 MiB device
+scratch reserve and respecting the managed budget.
+
 - `--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.

docs/pixart.md

@@ -0,0 +1,46 @@
+# How to Use
+
+You can run PixArt-α / PixArt-Σ with stable-diffusion.cpp.
+
+PixArt is a DiT-based text-to-image model family conditioned by a T5-XXL text
+encoder and a 4-channel VAE: SDXL-style for PixArt-Σ and SD1.x-style for PixArt-α.
+
+## Download weights
+
+- Download the transformer (diffusion model)
+    - PixArt-Σ XL-2 1024-MS: https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS/tree/main/transformer
+    - PixArt-α XL-2 1024-MS: https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS/tree/main/transformer
+- Download the T5-XXL text encoder
+    - safetensors: https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS/tree/main/text_encoder
+- Download the VAE
+    - PixArt-Σ: https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS/tree/main/vae
+    - PixArt-α: https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS/tree/main/vae
+    - Use the VAE matching the checkpoint's latent space. For TAE decoding or
+      preview, use TAESDXL for PixArt-Σ and TAESD for PixArt-α.
+- Tokenizer: the T5 vocabulary is embedded; no extra tokenizer file is needed.
+
+## Examples
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\pixart_sigma_xl2_1024_ms.safetensors --t5xxl ..\models\text_encoders\t5xxl.safetensors --vae ..\models\vae\pixart_vae.safetensors -p "a lovely cat" --cfg-scale 4.5 -W 1024 -H 1024 --steps 20 -v
+```
+
+## Notes
+
+- The VAE scaling factor defaults to `0.13025` for PixArt-Σ. PixArt-α
+  checkpoints with resolution micro-condition weights use `0.18215`.
+  PixArt-α 512 has the same tensor layout as PixArt-Σ, so it requires an
+  explicit override: `--model-args "pixart_vae_scale_factor=0.18215"`.
+  This argument can also override the scale for other compatible checkpoints.
+- PixArt-Σ checkpoints compute 2D sincos positional embeddings at runtime;
+  the trained grid is 64x64 patches with an interpolation scale of 2.
+  For checkpoints trained at a different resolution, the positional embedding
+  parameters can be adjusted via model args:
+  `--model-args "pixart_pos_embed_base_size=<trained grid>,pixart_interpolation_scale=<scale>"`
+  (e.g. `pixart_pos_embed_base_size=32,pixart_interpolation_scale=1,pixart_vae_scale_factor=0.18215` for
+  PixArt-α XL-2 512).
+- Checkpoints carrying resolution/aspect-ratio micro-condition weights are
+  detected but those conditions are not applied yet; a warning is logged and
+  generation proceeds with the timestep embedding only.
+- The transformer predicts 8 channels (noise + learned variance); only the
+  noise half is used for sampling, matching the reference implementation.

docs/qwen_image_2.1.md

@@ -23,7 +23,7 @@ Run the following commands from the build directory. Use image dimensions divisi
 ### Text to image
 
 ```powershell
-.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf -p "a lovely cat holding a sign says 'qwen2.1.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu -o qwen_image_2.1.png
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf -p "a lovely cat holding a sign says 'qwen2.1.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu --fa -o qwen_image_2.1.png
 ```
 
 <img alt="Qwen Image 2.1 example" src="../assets/qwen/qwen_image_2.1.png" />
@@ -35,7 +35,7 @@ To use GGUF diffusion weights, set `--diffusion-model` to the path of a file suc
 Pass the reference image with `-r` and describe the edit in `-p`. Vision weights are required; the example below loads them separately with `--llm_vision`.
 
 ```powershell
-.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3VL-8B-Instruct-mmproj-BF16.gguf -r ..\assets\qwen\qwen_image_2.1.png -p "change 'qwen2.1.cpp' to 'sd.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu -o qwen_image_2.1_edit.png
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3VL-8B-Instruct-mmproj-BF16.gguf -r ..\assets\qwen\qwen_image_2.1.png -p "change 'qwen2.1.cpp' to 'sd.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu --fa -o qwen_image_2.1_edit.png
 ```
 
 For multiple reference images, repeat `-r` in the desired order, for example `-r first.png -r second.png`.
@@ -44,7 +44,27 @@ For multiple reference images, repeat `-r` in the desired order, for example `-r
 
 By default, the first denoising call for each fixed condition saves the text and reference-image keys and values from every transformer layer. Later calls only compute the target-image tokens. Positive and negative conditions use separate caches, which are released when sampling ends.
 
-The cache uses FP32 on all attention backends. For the default 32-layer model, a prefix of 4096 tokens takes about 4 GiB per condition, in addition to weights and working buffers. The runner accounts for the cache when checking the memory budget. If a cached execution runs out of memory, it releases the prefix caches, disables caching for the rest of that sampling run, and retries the full sequence once. Per-step conditioning extensions currently use the full-sequence path.
+Set `qwen_image_2_1_prefix_cache_type` in `--model-args` to `auto` or a type name using the same parser and case-sensitive names as `--type`:
+
+- `auto` (default): use FP16 only when Flash Attention is enabled, Sage Attention is disabled, the attention scale is unchanged, and every attention operation in the cache-writing or cache-reading graph selects Flash Attention after backend support checks. If an operation falls back, rebuild the prefix in FP32 before executing and keep FP32 for the rest of that sampling run.
+- `f32`: always store FP32 keys and values.
+- `f16`: always store FP16 keys and values, including with ordinary attention or custom attention scaling. This saves cache memory but can introduce additional rounding error.
+- Other types, such as `bf16`, `q4_1`, `q5_0`, `q5_1`, `q8_0`, `q4_K`, `q6_K`, `iq4_nl`, and `iq4_xs`: use the requested storage type if the ggml build provides runtime conversion to and from FP32. Quantization is lossy and must be selected explicitly; `auto` never selects a quantized type.
+
+Cache data is packed into contiguous rows of `hidden_size` elements before conversion, so 256-element quantization blocks work with the model's 128-element attention heads without padding. The type's block size must divide `hidden_size`. Unknown types, types lacking runtime conversion (for example `q8_1` and several IQ formats), and incompatible block sizes are ignored with a warning, leaving the previous setting or the default `auto` unchanged.
+
+For example, use `--model-args qwen_image_2_1_prefix_cache_type=q8_0` to enable 8-bit cache storage. Cached keys and values are converted back to the attention input type before concatenating with the current target tokens. This reduces persistent cache memory; attention working buffers still use floating-point values, and conversion adds work on each step. Backends without the required conversion operations use the existing CPU fallback.
+
+For the default 32-layer model, a prefix of 4096 tokens takes approximately the following memory per condition, excluding weights, working buffers, and allocation overhead:
+
+| Cache type | Memory |
+| --- | ---: |
+| `f32` | 4 GiB |
+| `f16` | 2 GiB |
+| `q8_0` | 1.0625 GiB |
+| `q4_0` | 0.5625 GiB |
+
+The runner accounts for the cache when checking the memory budget. If a cached execution runs out of memory, it releases the prefix caches, disables caching for the rest of that sampling run, and retries the full sequence once. Per-step conditioning extensions currently use the full-sequence path.
 
 Disable this optimization with `--model-args qwen_image_2_1_prefix_cache=false`. It reuses step-independent activations; numerical results can still differ slightly because the matrix sizes change.
 

examples/common/common.cpp

@@ -518,7 +518,8 @@ ArgOptions SDContextParams::get_options() {
         {"",
          "--model-args",
          "extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, "
-         "chroma_t5_mask_pad, qwen_image_zero_cond_t, qwen_image_2_1_prefix_cache",
+         "chroma_t5_mask_pad, qwen_image_zero_cond_t, qwen_image_2_1_prefix_cache, "
+         "qwen_image_2_1_prefix_cache_type (auto or a type name from --type)",
          (int)',',
          &model_args},
         {"",
@@ -1340,7 +1341,7 @@ ArgOptions SDGenerationParams::get_options() {
          &embed_image_metadata},
         {"",
          "--vae-tiling",
-         "process vae in tiles to reduce memory usage",
+         "process vae encode and decode in spatial tiles to reduce memory usage (default: 256x256 image pixels)",
          true,
          &vae_tiling_params.enabled},
         {"",
@@ -1604,12 +1605,12 @@ ArgOptions SDGenerationParams::get_options() {
         size_t x_pos              = tile_size_str.find('x');
         try {
             if (x_pos != std::string::npos) {
-                std::string tile_x_str        = tile_size_str.substr(0, x_pos);
-                std::string tile_y_str        = tile_size_str.substr(x_pos + 1);
-                vae_tiling_params.tile_size_x = std::stoi(tile_x_str);
-                vae_tiling_params.tile_size_y = std::stoi(tile_y_str);
+                std::string tile_w_str        = tile_size_str.substr(0, x_pos);
+                std::string tile_h_str        = tile_size_str.substr(x_pos + 1);
+                vae_tiling_params.tile_size_w = std::stoi(tile_w_str);
+                vae_tiling_params.tile_size_h = std::stoi(tile_h_str);
             } else {
-                vae_tiling_params.tile_size_x = vae_tiling_params.tile_size_y = std::stoi(tile_size_str);
+                vae_tiling_params.tile_size_w = vae_tiling_params.tile_size_h = std::stoi(tile_size_str);
             }
         } catch (const std::invalid_argument&) {
             return -1;
@@ -1627,12 +1628,12 @@ ArgOptions SDGenerationParams::get_options() {
         size_t x_pos             = rel_size_str.find('x');
         try {
             if (x_pos != std::string::npos) {
-                std::string rel_x_str        = rel_size_str.substr(0, x_pos);
-                std::string rel_y_str        = rel_size_str.substr(x_pos + 1);
-                vae_tiling_params.rel_size_x = std::stof(rel_x_str);
-                vae_tiling_params.rel_size_y = std::stof(rel_y_str);
+                std::string rel_w_str        = rel_size_str.substr(0, x_pos);
+                std::string rel_h_str        = rel_size_str.substr(x_pos + 1);
+                vae_tiling_params.rel_size_w = std::stof(rel_w_str);
+                vae_tiling_params.rel_size_h = std::stof(rel_h_str);
             } else {
-                vae_tiling_params.rel_size_x = vae_tiling_params.rel_size_y = std::stof(rel_size_str);
+                vae_tiling_params.rel_size_w = vae_tiling_params.rel_size_h = std::stof(rel_size_str);
             }
         } catch (const std::invalid_argument&) {
             return -1;
@@ -1762,11 +1763,11 @@ ArgOptions SDGenerationParams::get_options() {
          on_scm_policy_arg},
         {"",
          "--vae-tile-size",
-         "tile size for vae tiling in latent units, not image pixels, format [X]x[Y] (default: 32x32)",
+         "tile size for vae encode and decode in image pixels, format [W]x[H] or [S] (default: 256x256; requires --vae-tiling)",
          on_tile_size_arg},
         {"",
          "--vae-relative-tile-size",
-         "relative tile size for vae tiling, format [X]x[Y], in fraction of image size if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)",
+         "relative tile size for vae encode and decode, format [W]x[H] or [S]: <=1 is a dimension fraction, >1 a target tile count (overrides --vae-tile-size; requires --vae-tiling)",
          on_relative_tile_size_arg},
         {"",
          "--prompt-file",
@@ -2223,20 +2224,20 @@ bool SDGenerationParams::from_json_str(
         if (tiling_json.contains("temporal_tiling") && tiling_json["temporal_tiling"].is_boolean()) {
             vae_tiling_params.temporal_tiling = tiling_json["temporal_tiling"];
         }
-        if (tiling_json.contains("tile_size_x") && tiling_json["tile_size_x"].is_number_integer()) {
-            vae_tiling_params.tile_size_x = tiling_json["tile_size_x"];
+        if (tiling_json.contains("tile_size_w") && tiling_json["tile_size_w"].is_number_integer()) {
+            vae_tiling_params.tile_size_w = tiling_json["tile_size_w"];
         }
-        if (tiling_json.contains("tile_size_y") && tiling_json["tile_size_y"].is_number_integer()) {
-            vae_tiling_params.tile_size_y = tiling_json["tile_size_y"];
+        if (tiling_json.contains("tile_size_h") && tiling_json["tile_size_h"].is_number_integer()) {
+            vae_tiling_params.tile_size_h = tiling_json["tile_size_h"];
         }
         if (tiling_json.contains("target_overlap") && tiling_json["target_overlap"].is_number()) {
             vae_tiling_params.target_overlap = tiling_json["target_overlap"];
         }
-        if (tiling_json.contains("rel_size_x") && tiling_json["rel_size_x"].is_number()) {
-            vae_tiling_params.rel_size_x = tiling_json["rel_size_x"];
+        if (tiling_json.contains("rel_size_w") && tiling_json["rel_size_w"].is_number()) {
+            vae_tiling_params.rel_size_w = tiling_json["rel_size_w"];
         }
-        if (tiling_json.contains("rel_size_y") && tiling_json["rel_size_y"].is_number()) {
-            vae_tiling_params.rel_size_y = tiling_json["rel_size_y"];
+        if (tiling_json.contains("rel_size_h") && tiling_json["rel_size_h"].is_number()) {
+            vae_tiling_params.rel_size_h = tiling_json["rel_size_h"];
         }
         if (tiling_json.contains("extra_tiling_args") && tiling_json["extra_tiling_args"].is_string()) {
             extra_tiling_args = tiling_json["extra_tiling_args"].get<std::string>();
@@ -2933,11 +2934,11 @@ std::string SDGenerationParams::to_string() const {
         << "  vae_tiling_params: { "
         << vae_tiling_params.enabled << ", "
         << vae_tiling_params.temporal_tiling << ", "
-        << vae_tiling_params.tile_size_x << ", "
-        << vae_tiling_params.tile_size_y << ", "
+        << vae_tiling_params.tile_size_w << ", "
+        << vae_tiling_params.tile_size_h << ", "
         << vae_tiling_params.target_overlap << ", "
-        << vae_tiling_params.rel_size_x << ", "
-        << vae_tiling_params.rel_size_y << ", "
+        << vae_tiling_params.rel_size_w << ", "
+        << vae_tiling_params.rel_size_h << ", "
         << "\"" << extra_tiling_args << "\" },\n"
         << "}";
     return oss.str();
@@ -3139,11 +3140,11 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
         root["vae_tiling"] = {
             {"enabled", gen_params.vae_tiling_params.enabled},
             {"temporal_tiling", gen_params.vae_tiling_params.temporal_tiling},
-            {"tile_size_x", gen_params.vae_tiling_params.tile_size_x},
-            {"tile_size_y", gen_params.vae_tiling_params.tile_size_y},
+            {"tile_size_w", gen_params.vae_tiling_params.tile_size_w},
+            {"tile_size_h", gen_params.vae_tiling_params.tile_size_h},
             {"target_overlap", gen_params.vae_tiling_params.target_overlap},
-            {"rel_size_x", gen_params.vae_tiling_params.rel_size_x},
-            {"rel_size_y", gen_params.vae_tiling_params.rel_size_y},
+            {"rel_size_w", gen_params.vae_tiling_params.rel_size_w},
+            {"rel_size_h", gen_params.vae_tiling_params.rel_size_h},
             {"extra_tiling_args", gen_params.extra_tiling_args},
         };
     }

examples/server/CMakeLists.txt

@@ -21,21 +21,21 @@ if(SD_SERVER_BUILD_FRONTEND AND EXISTS "${FRONTEND_DIR}")
         set(HAVE_FRONTEND_BUILD ON)
 
         add_custom_target(${TARGET}_frontend_install
-            COMMAND "${PNPM_EXECUTABLE}" -C "${FRONTEND_DIR}" install
+            COMMAND "${PNPM_EXECUTABLE}" --ignore-workspace -C "${FRONTEND_DIR}" install
             WORKING_DIRECTORY "${FRONTEND_DIR}"
             COMMENT "Installing frontend dependencies"
             VERBATIM
         )
 
         add_custom_target(${TARGET}_frontend_build
-            COMMAND "${PNPM_EXECUTABLE}" -C "${FRONTEND_DIR}" run build
+            COMMAND "${PNPM_EXECUTABLE}" --ignore-workspace -C "${FRONTEND_DIR}" run build
             WORKING_DIRECTORY "${FRONTEND_DIR}"
             COMMENT "Building frontend"
             VERBATIM
         )
 
         add_custom_target(${TARGET}_frontend_header
-            COMMAND "${PNPM_EXECUTABLE}" -C "${FRONTEND_DIR}" run build:header
+            COMMAND "${PNPM_EXECUTABLE}" --ignore-workspace -C "${FRONTEND_DIR}" run build:header
             WORKING_DIRECTORY "${FRONTEND_DIR}"
             COMMENT "Generating gen_index_html.h"
             VERBATIM

examples/server/api.md

@@ -524,11 +524,11 @@ Shared default fields used by both `img_gen` and `vid_gen`:
 | `vae_tiling_params` | `object` |
 | `vae_tiling_params.enabled` | `boolean` |
 | `vae_tiling_params.temporal_tiling` | `boolean` |
-| `vae_tiling_params.tile_size_x` | `integer` |
-| `vae_tiling_params.tile_size_y` | `integer` |
+| `vae_tiling_params.tile_size_w` | `integer` |
+| `vae_tiling_params.tile_size_h` | `integer` |
 | `vae_tiling_params.target_overlap` | `number` |
-| `vae_tiling_params.rel_size_x` | `number` |
-| `vae_tiling_params.rel_size_y` | `number` |
+| `vae_tiling_params.rel_size_w` | `number` |
+| `vae_tiling_params.rel_size_h` | `number` |
 | `vae_tiling_params.extra_tiling_args` | `string` |
 | `cache_mode` | `string` |
 | `cache_option` | `string` |
@@ -537,6 +537,8 @@ Shared default fields used by both `img_gen` and `vid_gen`:
 | `output_format` | `string` |
 | `output_compression` | `integer` |
 
+`vae_tiling_params.tile_size_w` and `tile_size_h` are in **image pixels**, with `0` selecting the 256-pixel default. Both encode and decode use these sizes without an encoding multiplier. Positive `rel_size_w`/`rel_size_h` values override the corresponding absolute size: values up to 1 are dimension fractions, and values greater than 1 are target tile counts. Set `enabled` to use spatial tiling. Sizes are aligned down to the VAE scale factor and capped at the input dimensions; explicit sizes below the minimum supported tile size are rejected. These fields previously used latent units; see [VAE tiling](../../docs/performance.md#use-vae-tiling-to-reduce-encode-and-decode-memory-usage) for migration and OOM retry behavior.
+
 `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.
 
@@ -767,11 +769,11 @@ Example:
   "vae_tiling_params": {
     "enabled": false,
     "temporal_tiling": false,
-    "tile_size_x": 0,
-    "tile_size_y": 0,
+    "tile_size_w": 0,
+    "tile_size_h": 0,
     "target_overlap": 0.5,
-    "rel_size_x": 0.0,
-    "rel_size_y": 0.0,
+    "rel_size_w": 0.0,
+    "rel_size_h": 0.0,
     "extra_tiling_args": ""
   },
 
@@ -900,11 +902,11 @@ Other native fields:
 | `vae_tiling_params` | `object` |
 | `vae_tiling_params.enabled` | `boolean` |
 | `vae_tiling_params.temporal_tiling` | `boolean` |
-| `vae_tiling_params.tile_size_x` | `integer` |
-| `vae_tiling_params.tile_size_y` | `integer` |
+| `vae_tiling_params.tile_size_w` | `integer` |
+| `vae_tiling_params.tile_size_h` | `integer` |
 | `vae_tiling_params.target_overlap` | `number` |
-| `vae_tiling_params.rel_size_x` | `number` |
-| `vae_tiling_params.rel_size_y` | `number` |
+| `vae_tiling_params.rel_size_w` | `number` |
+| `vae_tiling_params.rel_size_h` | `number` |
 | `vae_tiling_params.extra_tiling_args` | `string` |
 | `cache_mode` | `string` |
 | `cache_option` | `string` |
@@ -1115,11 +1117,11 @@ Example:
   "vae_tiling_params": {
     "enabled": false,
     "temporal_tiling": false,
-    "tile_size_x": 0,
-    "tile_size_y": 0,
+    "tile_size_w": 0,
+    "tile_size_h": 0,
     "target_overlap": 0.5,
-    "rel_size_x": 0.0,
-    "rel_size_y": 0.0,
+    "rel_size_w": 0.0,
+    "rel_size_h": 0.0,
     "extra_tiling_args": ""
   },
 
@@ -1240,11 +1242,11 @@ Other native fields:
 | `vae_tiling_params` | `object` |
 | `vae_tiling_params.enabled` | `boolean` |
 | `vae_tiling_params.temporal_tiling` | `boolean` |
-| `vae_tiling_params.tile_size_x` | `integer` |
-| `vae_tiling_params.tile_size_y` | `integer` |
+| `vae_tiling_params.tile_size_w` | `integer` |
+| `vae_tiling_params.tile_size_h` | `integer` |
 | `vae_tiling_params.target_overlap` | `number` |
-| `vae_tiling_params.rel_size_x` | `number` |
-| `vae_tiling_params.rel_size_y` | `number` |
+| `vae_tiling_params.rel_size_w` | `number` |
+| `vae_tiling_params.rel_size_h` | `number` |
 | `vae_tiling_params.extra_tiling_args` | `string` |
 | `cache_mode` | `string` |
 | `cache_option` | `string` |

examples/server/routes_sdcpp.cpp

@@ -78,11 +78,11 @@ static json make_vae_tiling_json(const sd_tiling_params_t& params) {
     return {
         {"enabled", params.enabled},
         {"temporal_tiling", params.temporal_tiling},
-        {"tile_size_x", params.tile_size_x},
-        {"tile_size_y", params.tile_size_y},
+        {"tile_size_w", params.tile_size_w},
+        {"tile_size_h", params.tile_size_h},
         {"target_overlap", params.target_overlap},
-        {"rel_size_x", params.rel_size_x},
-        {"rel_size_y", params.rel_size_y},
+        {"rel_size_w", params.rel_size_w},
+        {"rel_size_h", params.rel_size_h},
         {"extra_tiling_args", params.extra_tiling_args ? params.extra_tiling_args : ""},
     };
 }

include/stable-diffusion.h

@@ -173,11 +173,13 @@ enum lora_apply_mode_t {
 typedef struct {
     bool enabled;
     bool temporal_tiling;
-    int tile_size_x;
-    int tile_size_y;
+    // Spatial tile dimensions in image pixels for both encode and decode; 0 uses 256.
+    int tile_size_w;
+    int tile_size_h;
     float target_overlap;
-    float rel_size_x;
-    float rel_size_y;
+    // Positive values override tile_size: <= 1 is a dimension fraction, > 1 a target tile count.
+    float rel_size_w;
+    float rel_size_h;
     const char* extra_tiling_args;
 } sd_tiling_params_t;
 

src/core/backend_fit.cpp

@@ -478,7 +478,12 @@ namespace sd::backend_fit {
         return true;
     }
 
-    bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling, ggml_status status) {
+    bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params,
+                                         bool prefer_temporal_tiling,
+                                         ggml_status status,
+                                         int latent_tile_size_w,
+                                         int latent_tile_size_h,
+                                         int scale_factor) {
         // Execution failures can leave the device unusable; tiling only helps with allocation failures.
         if (status != GGML_STATUS_ALLOC_FAILED) {
             return false;
@@ -487,19 +492,31 @@ namespace sd::backend_fit {
         if (prefer_temporal_tiling && !tiling_params.temporal_tiling) {
             tiling_params.temporal_tiling = true;
             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;
+            if (latent_tile_size_w <= 0 || latent_tile_size_h <= 0 || scale_factor <= 0) {
+                return false;
+            }
+            auto smaller_tile = [&](int size) {
+                int next_size = size / 2;
+                if (!tiling_params.enabled) {
+                    next_size = std::min(next_size, 256 / scale_factor);
+                }
+                return std::min(size, std::max(4, next_size));
+            };
+            const int tile_size_w = smaller_tile(latent_tile_size_w);
+            const int tile_size_h = smaller_tile(latent_tile_size_h);
+            if (tile_size_w == latent_tile_size_w && tile_size_h == latent_tile_size_h) {
+                return false;
+            }
+            tiling_params.enabled     = true;
+            tiling_params.rel_size_w  = 0.0f;
+            tiling_params.rel_size_h  = 0.0f;
+            tiling_params.tile_size_w = tile_size_w * scale_factor;
+            tiling_params.tile_size_h = tile_size_h * scale_factor;
+            retry_mode                = tiling_params.temporal_tiling ? "spatial+temporal" : "spatial";
+            LOG_WARN("Reducing VAE decode tiles from %dx%d to %dx%d image pixels",
+                     latent_tile_size_w * scale_factor, latent_tile_size_h * scale_factor,
+                     tiling_params.tile_size_w, tiling_params.tile_size_h);
         }
 
         LOG_WARN("VAE decode ran out of memory; retrying with %s tiling",

src/core/backend_fit.h

@@ -17,7 +17,10 @@ namespace sd::backend_fit {
 
     bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params,
                                          bool prefer_temporal_tiling,
-                                         ggml_status status);
+                                         ggml_status status,
+                                         int latent_tile_size_w,
+                                         int latent_tile_size_h,
+                                         int scale_factor);
 
 }  // namespace sd::backend_fit
 

src/core/ggml_extend.cpp

@@ -465,7 +465,11 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
                               int d1,
                               int d2,
                               bool force_prec_f32,
-                              bool direct) {
+                              bool direct,
+                              float scale) {
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, scale);
+    }
     if (direct) {
         int64_t OC = w->ne[3] / IC;
         int64_t N  = x->ne[3] / IC;
@@ -502,6 +506,9 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
         }
     }
 
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, 1.f / scale);
+    }
     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);
@@ -623,7 +630,11 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
                                     bool skip_reshape,
                                     bool flash_attn,
                                     float kv_scale,
-                                    bool sage_attn) {  // avoid overflow
+                                    bool sage_attn,
+                                    bool* used_flash_attn) {  // avoid overflow
+    if (used_flash_attn != nullptr) {
+        *used_flash_attn = false;
+    }
     int64_t L_q;
     int64_t L_k;
     int64_t C;
@@ -755,6 +766,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
         if (can_use_flash_attn) {
             kqv = build_kqv(q, k, v, mask);
             if (kqv != nullptr) {
+                if (used_flash_attn != nullptr) {
+                    *used_flash_attn = true;
+                }
                 kqv = ggml_view_4d(ctx,
                                    kqv,
                                    d_head,

src/core/ggml_extend.h

@@ -154,7 +154,8 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
                               int d1              = 1,
                               int d2              = 1,
                               bool force_prec_f32 = false,
-                              bool direct         = false);
+                              bool direct         = false,
+                              float scale         = 1.f);
 
 // w: [OC，IC, KD, 1 * 1]
 // x: [N, IC, ID, IH*IW]
@@ -217,11 +218,12 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
                                     ggml_tensor* k,
                                     ggml_tensor* v,
                                     int64_t n_head,
-                                    ggml_tensor* mask = nullptr,
-                                    bool skip_reshape = false,
-                                    bool flash_attn   = false,
-                                    float kv_scale    = 1.0f,
-                                    bool sage_attn    = false);
+                                    ggml_tensor* mask     = nullptr,
+                                    bool skip_reshape     = false,
+                                    bool flash_attn       = false,
+                                    float kv_scale        = 1.0f,
+                                    bool sage_attn        = false,
+                                    bool* used_flash_attn = nullptr);
 
 ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
                                  ggml_tensor* x,

src/core/ggml_runner.cpp

@@ -21,11 +21,12 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
                                     ggml_tensor* mask,
                                     bool skip_reshape,
                                     bool flash_attn,
-                                    float kv_scale) {
+                                    float kv_scale,
+                                    bool* used_flash_attn) {
     if (ctx->attn_scale > 0.f) {
         kv_scale = ctx->attn_scale;
     }
-    return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled);
+    return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled, used_flash_attn);
 }
 
 void GGMLRunner::alloc_params_ctx() {
@@ -515,9 +516,10 @@ GGMLRunner::~GGMLRunner() {
     free_params_ctx();
 }
 
-GGMLRunnerContext GGMLRunner::get_context() {
+GGMLRunnerContext GGMLRunner::get_context(ggml_cgraph* graph) {
     GGMLRunnerContext runner_ctx;
     runner_ctx.ggml_ctx              = compute_ctx;
+    runner_ctx.graph                 = graph;
     runner_ctx.backend               = runtime_backend;
     runner_ctx.flash_attn_enabled    = flash_attn_enabled;
     runner_ctx.sage_attn_enabled     = sage_attn_enabled;
@@ -532,8 +534,8 @@ GGMLRunnerContext GGMLRunner::get_context() {
     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.cache_tensor = [this, graph](const std::string& name, ggml_tensor* tensor) {
+        this->cache(name, tensor, graph);
     };
     runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) {
         this->set_backend_tensor_data(tensor, data);
@@ -575,7 +577,7 @@ ggml_tensor* GGMLRunner::to_backend(ggml_tensor* tensor) {
     }
 }
 
-void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
+void GGMLRunner::cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph) {
     if (tensor != nullptr && tensor->view_src != nullptr) {
         tensor = ggml_cont(compute_ctx, tensor);
     }
@@ -583,6 +585,10 @@ void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
         ggml_set_output(tensor);
     }
     cache_.stage(name, tensor);
+    if (graph != nullptr && tensor != nullptr) {
+        // Schedule the cache output here so its source can be reused before graph end.
+        ggml_build_forward_expand(graph, tensor);
+    }
 }
 
 std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
@@ -831,11 +837,20 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
         last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
         return std::nullopt;
     }
+    auto fits_monolithic = [&]() {
+        // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits.
+        constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL;
+        auto requests                      = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph));
+        for (auto& request : requests) {
+            request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom);
+        }
+        return fits(requests, params);
+    };
     auto manager         = residency_manager.lock();
     const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
                            manager != nullptr && manager->segmented_compute_enabled() &&
                            cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1 &&
-                           !fits(memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)), params);
+                           !fits_monolithic();
     ggml_graph_cut::Plan monolithic_plan;
     if (!segmented) {
         monolithic_plan.segments.emplace_back();

src/core/ggml_runner.h

@@ -67,6 +67,7 @@ struct WeightAdapter {
 struct GGMLRunnerContext {
     ggml_backend_t backend                                           = nullptr;
     ggml_context* ggml_ctx                                           = nullptr;
+    ggml_cgraph* graph                                               = nullptr;
     bool flash_attn_enabled                                          = false;
     bool sage_attn_enabled                                           = false;
     float linear_scale                                               = 0.f;
@@ -102,6 +103,12 @@ struct GGMLRunnerContext {
         return get_cache_tensor(name);
     }
 
+    void expand_graph(ggml_tensor* tensor) const {
+        if (graph != nullptr && tensor != nullptr) {
+            ggml_build_forward_expand(graph, tensor);
+        }
+    }
+
     void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const {
         if (!cache_tensor || tensor == nullptr) {
             return;
@@ -122,10 +129,11 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
                                     ggml_tensor* k,
                                     ggml_tensor* v,
                                     int64_t n_head,
-                                    ggml_tensor* mask = nullptr,
-                                    bool skip_reshape = false,
-                                    bool flash_attn   = false,
-                                    float kv_scale    = 1.f);
+                                    ggml_tensor* mask     = nullptr,
+                                    bool skip_reshape     = false,
+                                    bool flash_attn       = false,
+                                    float kv_scale        = 1.f,
+                                    bool* used_flash_attn = nullptr);
 
 struct GGMLRunner {
 private:
@@ -289,7 +297,8 @@ public:
 
     virtual ~GGMLRunner();
 
-    virtual GGMLRunnerContext get_context();
+    // Binding a graph schedules cache outputs at registration instead of graph end.
+    virtual GGMLRunnerContext get_context(ggml_cgraph* graph = nullptr);
 
     void reset_compute_ctx();
 
@@ -324,7 +333,7 @@ public:
 
     ggml_tensor* to_backend(ggml_tensor* tensor);
 
-    void cache(const std::string name, ggml_tensor* tensor);
+    void cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph = nullptr);
 
     ggml_tensor* get_cache_tensor_by_name(const std::string& name) {
         return cache_.get(name);

src/model.h

@@ -62,6 +62,7 @@ enum SDVersion {
     VERSION_SENSENOVA_U1_5,
     VERSION_LLADA_IMAGE,
     VERSION_ESRGAN,
+    VERSION_PIXART,
     VERSION_COUNT,
 };
 
@@ -252,6 +253,10 @@ static inline bool sd_version_is_sensenova_u1(SDVersion version) {
     return version == VERSION_SENSENOVA_U1_5;
 }
 
+static inline bool sd_version_is_pixart(SDVersion version) {
+    return version == VERSION_PIXART;
+}
+
 static inline bool sd_version_supports_video_generation(SDVersion version) {
     return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version);
 }
@@ -320,7 +325,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_sefi_image(version) ||
         sd_version_is_krea2(version) ||
         sd_version_is_mage_flow(version) ||
-        sd_version_is_sensenova_u1(version)) {
+        sd_version_is_sensenova_u1(version) ||
+        sd_version_is_pixart(version)) {
         return true;
     }
     return false;

src/model/adapter/lora.hpp

@@ -163,6 +163,27 @@ struct LoraModel : public GGMLRunner {
 
             lora_tensors = std::move(new_lora_tensors);
         }
+
+        std::unordered_map<std::string, ggml_tensor*> new_lora_tensors;
+        for (const auto& [old_name, tensor] : lora_tensors) {
+            std::string new_name = old_name;
+            if (starts_with(old_name, "lora.model.diffusion_model.transformer_blocks.")) {
+                // Qwen Image 2.1 stores the gate before the projection in fused MLP weights.
+                for (const auto& suffix : {std::string(".img_mlp.gate_layer.weight."), std::string(".img_mlp.proj.weight.")}) {
+                    size_t pos = old_name.find(suffix);
+                    if (pos == std::string::npos) {
+                        continue;
+                    }
+                    std::string fused_name = old_name.substr(5, pos - 5) + ".img_mlp.gate_up.weight";
+                    if (model_tensor_names.find(fused_name) != model_tensor_names.end()) {
+                        new_name = "lora." + fused_name + (suffix == ".img_mlp.proj.weight." ? ".1." : ".") + old_name.substr(pos + suffix.size());
+                    }
+                    break;
+                }
+            }
+            new_lora_tensors[new_name] = tensor;
+        }
+        lora_tensors = std::move(new_lora_tensors);
     }
 
     ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {

src/model/diffusion/control.hpp

@@ -376,7 +376,7 @@ struct ControlNet : public GGMLRunner {
             hint = make_input(hint_tensor);
         }
 
-        auto runner_ctx = get_context();
+        auto runner_ctx = get_context(gf);
 
         auto outs = control_net.forward(&runner_ctx,
                                         x,
@@ -389,8 +389,7 @@ struct ControlNet : public GGMLRunner {
         if (guided_hint_input == nullptr && !outs.empty()) {
             guided_hint_output_ggml = outs[0];
             ggml_set_output(guided_hint_output_ggml);
-            cache(guided_hint_cache_name(), guided_hint_output_ggml);
-            ggml_build_forward_expand(gf, guided_hint_output_ggml);
+            runner_ctx.persist_cache_tensor(guided_hint_cache_name(), guided_hint_output_ggml);
         }
 
         control_outputs_ggml.reserve(outs.size() > 0 ? outs.size() - 1 : 0);

src/model/diffusion/pixart.hpp

@@ -0,0 +1,399 @@
+#ifndef __SD_MODEL_DIFFUSION_PIXART_HPP__
+#define __SD_MODEL_DIFFUSION_PIXART_HPP__
+
+#include <algorithm>
+#include <cinttypes>
+#include <cmath>
+#include <vector>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/mmdit.hpp"
+#include "model/diffusion/model.hpp"
+#include "model_loader.h"
+
+// Ref: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/pixart_transformer_2d.py
+// Ref: https://github.com/PixArt-alpha/PixArt-sigma
+
+namespace PixArt {
+    constexpr int PIXART_GRAPH_SIZE = 20480;
+    constexpr int ADALN_EMBED_DIM   = 256;
+
+    struct PixArtConfig {
+        int64_t in_channels         = 4;
+        int64_t out_channels        = 8;  // learn_sigma: noise prediction + learned variance
+        int64_t hidden_size         = 1152;
+        int64_t cross_attention_dim = 1152;
+        int64_t caption_channels    = 4096;
+        int64_t num_heads           = 16;
+        int64_t patch_size          = 2;
+        int64_t ffn_dim             = 4608;
+        int64_t pos_embed_base_size = 64;
+        float interpolation_scale   = 2.f;
+        int num_layers              = 28;
+
+        static PixArtConfig detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) {
+            PixArtConfig config;
+            auto find = [&](const std::string& suffix) -> const TensorStorage* {
+                auto it = weights.find(prefix + "." + suffix);
+                return it == weights.end() ? nullptr : &it->second;
+            };
+            if (auto w = find("x_embedder.proj.weight")) {
+                config.hidden_size = w->ne[3];
+                config.in_channels = w->ne[2];
+                config.patch_size  = w->ne[0];
+            }
+            if (auto w = find("final_layer.linear.weight")) {
+                config.out_channels = w->ne[1] / (config.patch_size * config.patch_size);
+            }
+            if (auto w = find("y_embedder.y_proj.fc1.weight")) {
+                config.caption_channels = w->ne[0];
+            }
+            if (auto w = find("blocks.0.cross_attn.kv_linear.weight")) {
+                config.cross_attention_dim = w->ne[0];
+            }
+            if (auto w = find("blocks.0.mlp.fc1.weight")) {
+                config.ffn_dim = w->ne[1];
+            }
+            if (find("csize_embedder.mlp.0.weight") != nullptr) {
+                LOG_WARN("pixart: resolution/aspect-ratio micro conditions are not supported; output may differ from the reference");
+            }
+            int layers                     = 0;
+            const std::string block_prefix = prefix + ".blocks.";
+            for (const auto& [name, _] : weights) {
+                if (starts_with(name, block_prefix)) {
+                    layers = std::max(layers, atoi(name.substr(block_prefix.size()).c_str()) + 1);
+                }
+            }
+            if (layers > 0) {
+                config.num_layers = layers;
+                LOG_VERBOSE("pixart: layers = %d, hidden_size = %" PRId64,
+                            layers, config.hidden_size);
+            }
+            return config;
+        }
+    };
+
+    // Mirrors diffusers get_2d_sincos_pos_embed for a (gh, gw) patch grid.
+    static std::vector<float> gen_2d_sincos_pos_embed(int64_t dim,
+                                                      int64_t gh,
+                                                      int64_t gw,
+                                                      int64_t base_size,
+                                                      float interpolation_scale) {
+        // diffusers: meshgrid(grid_w, grid_h, indexing="xy") -> grid[0]=w, grid[1]=h,
+        // embedding = concat(sincos(w), sincos(h))
+        std::vector<float> out(static_cast<size_t>(gh) * gw * dim);
+        int64_t quarter = dim / 4;
+        for (int64_t h = 0; h < gh; ++h) {
+            float pos_h = static_cast<float>(h) / (static_cast<float>(gh) / base_size) / interpolation_scale;
+            for (int64_t w = 0; w < gw; ++w) {
+                float pos_w  = static_cast<float>(w) / (static_cast<float>(gw) / base_size) / interpolation_scale;
+                float* dst_w = out.data() + (h * gw + w) * dim;
+                float* dst_h = dst_w + dim / 2;
+                for (int64_t i = 0; i < quarter; ++i) {
+                    float omega        = 1.f / powf(10000.f, static_cast<float>(i) / quarter);
+                    dst_w[i]           = sinf(pos_w * omega);
+                    dst_w[i + quarter] = cosf(pos_w * omega);
+                    dst_h[i]           = sinf(pos_h * omega);
+                    dst_h[i + quarter] = cosf(pos_h * omega);
+                }
+            }
+        }
+        return out;
+    }
+
+    class PixArtTimestepEmbedding : public GGMLBlock {
+    public:
+        PixArtTimestepEmbedding(int64_t in_channels, int64_t out_dim) {
+            blocks["mlp.0"] = std::make_shared<Linear>(in_channels, out_dim);
+            blocks["mlp.2"] = std::make_shared<Linear>(out_dim, out_dim);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            x = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"])->forward(ctx, x);
+            x = ggml_silu(ctx->ggml_ctx, x);
+            return std::dynamic_pointer_cast<Linear>(blocks["mlp.2"])->forward(ctx, x);
+        }
+    };
+
+    class PixArtAttention : public GGMLBlock {
+        int64_t num_heads;
+        bool self_attention;
+
+    public:
+        PixArtAttention(int64_t dim, int64_t num_heads, int64_t context_dim, bool self_attention)
+            : num_heads(num_heads), self_attention(self_attention) {
+            if (self_attention) {
+                blocks["qkv"] = std::make_shared<Linear>(dim, 3 * dim);
+            } else {
+                blocks["q_linear"]  = std::make_shared<Linear>(dim, dim);
+                blocks["kv_linear"] = std::make_shared<Linear>(context_dim, 2 * dim);
+            }
+            blocks["proj"] = std::make_shared<Linear>(dim, dim);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* context, ggml_tensor* mask = nullptr) {
+            std::vector<ggml_tensor*> qkv;
+            if (self_attention) {
+                auto projected = std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x);
+                qkv            = ggml_ext_chunk(ctx->ggml_ctx, projected, 3, 0);
+            } else {
+                auto q     = std::dynamic_pointer_cast<Linear>(blocks["q_linear"])->forward(ctx, x);
+                auto kv    = std::dynamic_pointer_cast<Linear>(blocks["kv_linear"])->forward(ctx, context);
+                auto parts = ggml_ext_chunk(ctx->ggml_ctx, kv, 2, 0);
+                qkv        = {q, parts[0], parts[1]};
+            }
+            auto out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, mask, false, ctx->flash_attn_enabled);
+            return std::dynamic_pointer_cast<Linear>(blocks["proj"])->forward(ctx, out);
+        }
+    };
+
+    class PixArtBlock : public GGMLBlock {
+        int64_t dim;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            ggml_type wtype             = get_type(prefix + "scale_shift_table", tensor_storage_map, GGML_TYPE_F32);
+            params["scale_shift_table"] = ggml_new_tensor_2d(ctx, wtype, dim, 6);
+        }
+
+    public:
+        PixArtBlock(int64_t dim, int64_t num_heads, int64_t context_dim, int64_t ffn_dim)
+            : dim(dim) {
+            blocks["attn"]       = std::make_shared<PixArtAttention>(dim, num_heads, dim, true);
+            blocks["cross_attn"] = std::make_shared<PixArtAttention>(dim, num_heads, context_dim, false);
+            blocks["mlp.fc1"]    = std::make_shared<Linear>(dim, ffn_dim);
+            blocks["mlp.fc2"]    = std::make_shared<Linear>(ffn_dim, dim);
+        }
+
+        static ggml_tensor* norm(ggml_context* ctx, ggml_tensor* x) {
+            return ggml_norm(ctx, x, 1e-6f);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* mod, ggml_tensor* context, ggml_tensor* context_mask) {
+            // x: [N, n_token, dim]
+            // mod: [N, 6 * dim], shared adaLN-single output
+            int64_t N = x->ne[2];
+
+            auto table = params["scale_shift_table"];
+            if (table->type != GGML_TYPE_F32) {
+                table = ggml_cast(ctx->ggml_ctx, table, GGML_TYPE_F32);
+            }
+            table   = ggml_reshape_3d(ctx->ggml_ctx, table, dim, 6, 1);
+            auto m  = ggml_add(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, mod, dim, 6, N), table);
+            auto mv = ggml_ext_chunk(ctx->ggml_ctx, ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, m), dim * 6, N), 6, 0);
+
+            auto attn1 = std::dynamic_pointer_cast<PixArtAttention>(blocks["attn"]);
+            auto attn2 = std::dynamic_pointer_cast<PixArtAttention>(blocks["cross_attn"]);
+            auto proj  = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
+            auto fc2   = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"]);
+
+            auto gate = [&](ggml_tensor* y, ggml_tensor* g) {
+                g = ggml_reshape_3d(ctx->ggml_ctx, g, dim, 1, N);
+                return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, y, g));
+            };
+
+            auto h = modulate(ctx->ggml_ctx, norm(ctx->ggml_ctx, x), mv[0], mv[1]);
+            x      = gate(attn1->forward(ctx, h, h), mv[2]);
+            // ada_norm_single: no norm before cross-attention (PixArtMS.py)
+            x = ggml_add(ctx->ggml_ctx, x, attn2->forward(ctx, x, context, context_mask));
+            h = modulate(ctx->ggml_ctx, norm(ctx->ggml_ctx, x), mv[3], mv[4]);
+            h = proj->forward(ctx, h);
+            h = ggml_ext_gelu(ctx->ggml_ctx, h, true);
+            h = fc2->forward(ctx, h);
+            return gate(h, mv[5]);
+        }
+    };
+
+    class PixArtModel : public GGMLBlock {
+        PixArtConfig config;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            ggml_type wtype                         = get_type(prefix + "final_layer.scale_shift_table", tensor_storage_map, GGML_TYPE_F32);
+            params["final_layer.scale_shift_table"] = ggml_new_tensor_2d(ctx, wtype, config.hidden_size, 2);
+        }
+
+    public:
+        PixArtModel() = default;
+        PixArtModel(const PixArtConfig& config)
+            : config(config) {
+            blocks["x_embedder.proj"]       = std::make_shared<Conv2d>(config.in_channels,
+                                                                 config.hidden_size,
+                                                                 std::pair<int, int>{static_cast<int>(config.patch_size), static_cast<int>(config.patch_size)},
+                                                                 std::pair<int, int>{static_cast<int>(config.patch_size), static_cast<int>(config.patch_size)});
+            blocks["t_embedder"]            = std::make_shared<PixArtTimestepEmbedding>(ADALN_EMBED_DIM, config.hidden_size);
+            blocks["t_block.1"]             = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size);
+            blocks["y_embedder.y_proj.fc1"] = std::make_shared<Linear>(config.caption_channels, config.hidden_size);
+            blocks["y_embedder.y_proj.fc2"] = std::make_shared<Linear>(config.hidden_size, config.cross_attention_dim);
+            for (int i = 0; i < config.num_layers; ++i) {
+                blocks["blocks." + std::to_string(i)] =
+                    std::make_shared<PixArtBlock>(config.hidden_size, config.num_heads, config.cross_attention_dim, config.ffn_dim);
+            }
+            blocks["final_layer.norm_final"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
+            blocks["final_layer.linear"]     = std::make_shared<Linear>(config.hidden_size,
+                                                                    config.patch_size * config.patch_size * config.out_channels);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timesteps,
+                             ggml_tensor* context,
+                             ggml_tensor* pos_embed,
+                             ggml_tensor* context_mask) {
+            // x: [N, C, H, W] latent, context: [N, n_ctx, caption_channels]
+            int64_t W  = x->ne[0];
+            int64_t H  = x->ne[1];
+            int64_t N  = x->ne[3];
+            int64_t p  = config.patch_size;
+            int64_t wp = W / p;
+            int64_t hp = H / p;
+
+            auto h = std::dynamic_pointer_cast<Conv2d>(blocks["x_embedder.proj"])->forward(ctx, x);  // [N, hidden, hp, wp]
+            h      = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 1, 2, 0, 3));       // [N, hp, wp, hidden] -> [N, hp*wp, hidden]
+            h      = ggml_reshape_3d(ctx->ggml_ctx, h, config.hidden_size, wp * hp, N);              // [N, hp*wp, hidden]
+            h      = ggml_add(ctx->ggml_ctx, h, pos_embed);
+
+            auto t   = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, ADALN_EMBED_DIM, 10000);
+            auto emb = std::dynamic_pointer_cast<PixArtTimestepEmbedding>(blocks["t_embedder"])->forward(ctx, t);
+
+            auto mod = std::dynamic_pointer_cast<Linear>(blocks["t_block.1"])
+                           ->forward(ctx, ggml_silu(ctx->ggml_ctx, emb));  // [N, 6 * hidden]
+
+            auto ctx_emb = std::dynamic_pointer_cast<Linear>(blocks["y_embedder.y_proj.fc1"])->forward(ctx, context);
+            ctx_emb      = ggml_ext_gelu(ctx->ggml_ctx, ctx_emb, true);
+            ctx_emb      = std::dynamic_pointer_cast<Linear>(blocks["y_embedder.y_proj.fc2"])->forward(ctx, ctx_emb);
+
+            for (int i = 0; i < config.num_layers; ++i) {
+                auto block = std::dynamic_pointer_cast<PixArtBlock>(blocks["blocks." + std::to_string(i)]);
+                h          = block->forward(ctx, h, mod, ctx_emb, context_mask);
+                sd::ggml_graph_cut::mark_graph_cut(h, "pixart.blocks." + std::to_string(i), "h");
+            }
+
+            // scale_shift_table + emb -> (shift, scale) for the affine-free final norm
+            auto tail_table = params["final_layer.scale_shift_table"];
+            if (tail_table->type != GGML_TYPE_F32) {
+                tail_table = ggml_cast(ctx->ggml_ctx, tail_table, GGML_TYPE_F32);
+            }
+            auto ss    = ggml_add(ctx->ggml_ctx,
+                                  ggml_reshape_3d(ctx->ggml_ctx, tail_table, config.hidden_size, 2, 1),
+                                  ggml_reshape_3d(ctx->ggml_ctx, emb, config.hidden_size, 1, N));  // [2, hidden, N]
+            auto parts = ggml_ext_chunk(ctx->ggml_ctx,
+                                        ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, ss), config.hidden_size * 2, N),
+                                        2, 0);
+            h          = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer.norm_final"])->forward(ctx, h);
+            h          = modulate(ctx->ggml_ctx, h, parts[0], parts[1]);
+            h          = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, h);  // [N, hp*wp, p*p*out_ch]
+            h          = DiT::unpatchify(ctx->ggml_ctx, h, hp, wp, static_cast<int>(p), static_cast<int>(p), false);
+            return h;  // [N, out_channels, H, W]
+        }
+    };
+
+    struct PixArtRunner : public DiffusionModelRunner {
+        PixArtConfig config;
+        PixArtModel model;
+        std::vector<float> pos_vec;
+
+        PixArtRunner(ggml_backend_t backend,
+                     const String2TensorStorage& tensor_storage_map      = {},
+                     const std::string prefix                            = "",
+                     std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                     const char* model_args                              = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(PixArtConfig::detect_from_weights(tensor_storage_map, prefix)) {
+            for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
+                if (key == "pixart_pos_embed_base_size") {
+                    int parsed = 0;
+                    if (parse_strict_int(value, parsed)) {
+                        config.pos_embed_base_size = parsed;
+                    } else {
+                        LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                } else if (key == "pixart_interpolation_scale") {
+                    float parsed = 0.f;
+                    if (parse_strict_float(value, parsed)) {
+                        config.interpolation_scale = parsed;
+                    } else {
+                        LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                }
+            }
+            model = PixArtModel(config);
+            model.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "pixart";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            model.get_param_tensors(tensors, prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor,
+                                 const sd::Tensor<float>& mask_tensor) {
+            ggml_cgraph* gf        = new_graph_custom(PIXART_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
+            GGML_ASSERT(x->ne[3] == 1);
+            GGML_ASSERT(!context_tensor.empty());
+            ggml_tensor* context = make_input(context_tensor);
+
+            ggml_tensor* context_mask = nullptr;
+            if (!mask_tensor.empty()) {
+                // additive attention bias over context tokens: 0 keep / -inf discard
+                context_mask = ggml_reshape_4d(compute_ctx, make_input(mask_tensor), mask_tensor.shape()[0], 1, 1, 1);
+            }
+
+            int64_t W  = x->ne[0];
+            int64_t H  = x->ne[1];
+            int64_t wp = W / config.patch_size;
+            int64_t hp = H / config.patch_size;
+
+            pos_vec  = gen_2d_sincos_pos_embed(config.hidden_size, hp, wp,
+                                               config.pos_embed_base_size, config.interpolation_scale);
+            auto pos = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.hidden_size, wp * hp, 1);
+            set_backend_tensor_data(pos, pos_vec.data());
+
+            auto runner_ctx  = get_context();
+            ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pos, context_mask);
+            // learn_sigma: keep the noise prediction half of the output channels
+            out = ggml_ext_slice(compute_ctx, out, 2, 0, config.in_channels);
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context,
+                                  const sd::Tensor<float>& context_mask) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context, context_mask);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            auto context     = tensor_or_empty(diffusion_params.context);
+            auto context_msk = tensor_or_empty(diffusion_params.y);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           context,
+                           context_msk);
+        }
+    };
+}  // namespace PixArt
+
+#endif  // __SD_MODEL_DIFFUSION_PIXART_HPP__

src/model/diffusion/qwen_image_2_1.hpp

@@ -134,6 +134,8 @@ namespace Qwen {
         std::string name;
         std::string cut_group;
         int64_t prefix_length = 0;
+        ggml_type type        = GGML_TYPE_F32;
+        bool* flash_attn_used = nullptr;
     };
 
     class QwenImage21ZeroCenterRMSNorm : public RMSNorm {
@@ -189,10 +191,15 @@ namespace Qwen {
             q      = Rope::apply_rope(ctx->ggml_ctx, q, pe);
             k      = Rope::apply_rope(ctx->ggml_ctx, k, pe);
             if (cache.mode == QwenImage21PrefixCache::Mode::STORE) {
+                // Preserve query-first attention evaluation while writing each layer's
+                // prefix before its full-sequence K/V can accumulate across layers.
+                ctx->expand_graph(q);
                 auto persist = [&](ggml_tensor* tensor, int axis, const char* name) {
                     auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length);
-                    auto copy = ggml_new_tensor(ctx->ggml_ctx, GGML_TYPE_F32, 4, part->ne);
-                    copy      = ggml_cpy(ctx->ggml_ctx, part, copy);
+                    // Pack the contiguous data into wider rows so quantization blocks
+                    // can exceed head_dim without padding or changing element order.
+                    part      = ggml_reshape_2d(ctx->ggml_ctx, part, x->ne[0], cache.prefix_length);
+                    auto copy = ggml_cast(ctx->ggml_ctx, part, cache.type);
                     // Keep the copy in this layer's segment so graph cuts do not
                     // retain or recompute the full-sequence K/V in the final segment.
                     sd::ggml_graph_cut::mark_graph_cut(copy, cache.cut_group, name);
@@ -201,21 +208,37 @@ namespace Qwen {
                 persist(k, 1, "k");
                 persist(v, 2, "v");
             }
+            auto attend = [&](ggml_tensor* aq, ggml_tensor* ak, ggml_tensor* av, ggml_tensor* mask) {
+                bool used_flash_attn = false;
+                auto out             = ggml_ext_attention_ext(ctx, aq, ak, av, heads, mask, true, ctx->flash_attn_enabled, 1.f, &used_flash_attn);
+                if (cache.flash_attn_used != nullptr) {
+                    *cache.flash_attn_used &= used_flash_attn;
+                }
+                return out;
+            };
             ggml_tensor* result = nullptr;
             if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
                 auto prefix_k = ctx->load_cache_tensor(cache.name + ".k");
                 auto prefix_v = ctx->load_cache_tensor(cache.name + ".v");
                 GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr);
-                k      = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1);
-                v      = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
-                result = ggml_ext_attention_ext(ctx, q, k, v, heads, nullptr, true, ctx->flash_attn_enabled);
+                if (prefix_k->type != k->type) {
+                    prefix_k = ggml_cast(ctx->ggml_ctx, prefix_k, k->type);
+                }
+                if (prefix_v->type != v->type) {
+                    prefix_v = ggml_cast(ctx->ggml_ctx, prefix_v, v->type);
+                }
+                prefix_k = ggml_reshape_4d(ctx->ggml_ctx, prefix_k, dim_head, cache.prefix_length, heads, k->ne[3]);
+                prefix_v = ggml_reshape_4d(ctx->ggml_ctx, prefix_v, dim_head, heads, cache.prefix_length, v->ne[3]);
+                k        = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1);
+                v        = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
+                result   = attend(q, k, v, nullptr);
             } else {
                 for (size_t i = 0; i < segments.size(); ++i) {
                     const auto& segment = segments[i];
                     auto sq             = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end);
                     auto sk             = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end);
                     auto sv             = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end);
-                    auto out            = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled);
+                    auto out            = attend(sq, sk, sv, masks[i]);
                     result              = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
                 }
             }
@@ -351,8 +374,35 @@ namespace Qwen {
         QwenImage21Model model;
         std::vector<float> pe_data;
         std::vector<sd::Tensor<float>> mask_data;
-        bool prefix_cache_enabled  = true;
-        bool prefix_cache_disabled = false;
+        ggml_type prefix_cache_type = GGML_TYPE_COUNT;
+        bool prefix_cache_enabled   = true;
+        bool prefix_cache_disabled  = false;
+        bool prefix_cache_auto_f32  = false;
+
+        static bool supports_prefix_cache_type(ggml_type type) {
+            if (type == GGML_TYPE_F32) {
+                return true;
+            }
+            const auto* traits = ggml_get_type_traits(type);
+            if (traits->from_float_ref == nullptr || traits->to_float == nullptr) {
+                return false;
+            }
+            auto cpu = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+            if (cpu == nullptr) {
+                return false;
+            }
+            auto ctx = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
+                ggml_init({3 * ggml_tensor_overhead(), nullptr, true}), ggml_free);
+            if (ctx == nullptr) {
+                return false;
+            }
+            // Some reference quantizers have no runtime copy support. Query the
+            // device through the registry so dynamically loaded CPU backends work.
+            auto source  = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, ggml_blck_size(type));
+            auto encoded = ggml_cast(ctx.get(), source, type);
+            auto decoded = ggml_cast(ctx.get(), encoded, GGML_TYPE_F32);
+            return ggml_backend_dev_supports_op(cpu, encoded) && ggml_backend_dev_supports_op(cpu, decoded);
+        }
 
         QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr, const char* model_args = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager),
@@ -361,6 +411,22 @@ namespace Qwen {
             for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
                 if (key == "qwen_image_2_1_prefix_cache" && !parse_strict_bool(value, prefix_cache_enabled)) {
                     LOG_WARN("ignoring invalid Qwen Image 2.1 model arg '%s=%s'", key.c_str(), value.c_str());
+                } else if (key == "qwen_image_2_1_prefix_cache_type") {
+                    if (value == "auto") {
+                        prefix_cache_type = GGML_TYPE_COUNT;
+                        continue;
+                    }
+                    const auto type = sd_type_to_ggml_type(str_to_sd_type(value.c_str()));
+                    if (type == GGML_TYPE_COUNT) {
+                        LOG_WARN("ignoring unknown Qwen Image 2.1 cache type '%s'", value.c_str());
+                    } else if (!supports_prefix_cache_type(type)) {
+                        LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': runtime conversion to and from F32 is unavailable", value.c_str());
+                    } else if (config.hidden_size % ggml_blck_size(type) != 0) {
+                        LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': block size %" PRId64 " does not divide hidden size %" PRId64,
+                                 value.c_str(), ggml_blck_size(type), config.hidden_size);
+                    } else {
+                        prefix_cache_type = type;
+                    }
                 }
             }
             model.init(params_ctx, weights, prefix);
@@ -377,11 +443,9 @@ namespace Qwen {
                 const auto name = cache.name + "." + std::to_string(i);
                 auto k          = get_cache_tensor_by_name(name + ".k");
                 auto v          = get_cache_tensor_by_name(name + ".v");
-                if (k == nullptr || v == nullptr || k->type != GGML_TYPE_F32 || v->type != GGML_TYPE_F32 ||
-                    k->ne[0] != config.head_dim || k->ne[1] != cache.prefix_length ||
-                    k->ne[2] != config.hidden_size / config.head_dim || k->ne[3] != 1 ||
-                    v->ne[0] != config.head_dim || v->ne[1] != config.hidden_size / config.head_dim ||
-                    v->ne[2] != cache.prefix_length || v->ne[3] != 1) {
+                if (k == nullptr || v == nullptr || k->type != cache.type || v->type != cache.type ||
+                    k->ne[0] != config.hidden_size || k->ne[1] != cache.prefix_length || k->ne[2] != 1 || k->ne[3] != 1 ||
+                    v->ne[0] != config.hidden_size || v->ne[1] != cache.prefix_length || v->ne[2] != 1 || v->ne[3] != 1) {
                     return false;
                 }
             }
@@ -418,15 +482,28 @@ namespace Qwen {
             }
             if (!runner_started()) {
                 prefix_cache_disabled = false;
+                prefix_cache_auto_f32 = false;
             }
             QwenImage21PrefixCache cache;
             if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) {
                 cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) +
                              ".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled);
                 cache.prefix_length = layout.prefix_length;
-                cache.mode          = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
+                if (prefix_cache_type != GGML_TYPE_COUNT) {
+                    cache.type = prefix_cache_type;
+                } else if (!prefix_cache_auto_f32 && flash_attn_enabled && !sage_attn_enabled &&
+                           (attn_scale <= 0.f || attn_scale == 1.f)) {
+                    cache.type = GGML_TYPE_F16;
+                }
+                cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
             }
-            auto run = [&](const QwenImage21PrefixCache& active_cache) {
+            bool flash_attn_used = true;
+            auto run             = [&](const QwenImage21PrefixCache& active_cache) {
+                flash_attn_used    = true;
+                auto checked_cache = active_cache;
+                if (prefix_cache_type == GGML_TYPE_COUNT && active_cache.type == GGML_TYPE_F16) {
+                    checked_cache.flash_attn_used = &flash_attn_used;
+                }
                 const bool cached         = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
                 const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0);
                 Rope::Embedding embedding;
@@ -459,7 +536,7 @@ namespace Qwen {
                 auto build = [&]() {
                     auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2);
                     auto pe    = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2,
-                                                    layout.positions.size() - (cached ? layout.prefix_length : 0));
+                                                                layout.positions.size() - (cached ? layout.prefix_length : 0));
                     set_backend_tensor_data(pe, pe_data.data());
                     std::vector<ggml_tensor*> masks, ref_inputs;
                     for (const auto& mask : mask_data) {
@@ -470,15 +547,28 @@ namespace Qwen {
                             ref_inputs.push_back(make_input(ref));
                         }
                     }
-                    auto ctx = get_context();
+                    auto ctx = get_context(graph);
                     auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), cached ? nullptr : make_input(context),
-                                             ref_inputs, pe, layout, masks, active_cache);
+                                                         ref_inputs, pe, layout, masks, checked_cache);
+                    if (!flash_attn_used) {
+                        return static_cast<ggml_cgraph*>(nullptr);
+                    }
                     ggml_build_forward_expand(graph, out);
                     return graph;
                 };
                 return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim());
             };
             auto result = run(cache);
+            if (result.empty() && !flash_attn_used) {
+                // Casting an F16 cache back to F32 cannot recover its original values.
+                // Recompute the prefix before executing a graph that falls back from FA.
+                free_cache_ctx_and_buffer();
+                prefix_cache_auto_f32 = true;
+                cache.type            = GGML_TYPE_F32;
+                cache.mode            = QwenImage21PrefixCache::Mode::STORE;
+                LOG_DEBUG("Qwen Image 2.1: Flash Attention unavailable; using F32 prefix caching for this sampling run");
+                result = run(cache);
+            }
             if (result.empty() && last_compute_status() == GGML_STATUS_ALLOC_FAILED &&
                 (cache.mode != QwenImage21PrefixCache::Mode::NONE || !cache_.empty())) {
                 // The failed graph has ended before persistent inputs are released.
@@ -493,7 +583,7 @@ namespace Qwen {
                     prefix_cache_disabled = true;
                     LOG_WARN("Qwen Image 2.1: incomplete prefix cache; disabling it for this sampling run");
                 } else {
-                    LOG_DEBUG("Qwen Image 2.1: cached prefix %" PRIu64 " (%" PRId64 " tokens)", extra->prefix_id, layout.prefix_length);
+                    LOG_DEBUG("Qwen Image 2.1: cached prefix %" PRIu64 " (%" PRId64 " tokens, %s)", extra->prefix_id, layout.prefix_length, ggml_type_name(cache.type));
                 }
             }
             return result;

src/model/diffusion/sensenova_u1.h

@@ -442,16 +442,9 @@ namespace SenseNovaU1 {
                 k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 2);
                 v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
             } else {
-                // Keep dedicated graph outputs alive until the runner copies them
-                // into its persistent cache buffer after graph execution.
-                auto cache_k = ggml_dup_tensor(ctx->ggml_ctx, k);
-                cache_k      = ggml_cpy(ctx->ggml_ctx, k, cache_k);
-                ggml_set_output(cache_k);
-                auto cache_v = ggml_dup_tensor(ctx->ggml_ctx, v);
-                cache_v      = ggml_cpy(ctx->ggml_ctx, v, cache_v);
-                ggml_set_output(cache_v);
-                ctx->persist_cache_tensor(layer_cache + ".k", cache_k);
-                ctx->persist_cache_tensor(layer_cache + ".v", cache_v);
+                ctx->expand_graph(q);
+                ctx->persist_cache_tensor(layer_cache + ".k", k);
+                ctx->persist_cache_tensor(layer_cache + ".v", v);
             }
 
             q = ggml_cont(ctx->ggml_ctx,
@@ -687,7 +680,7 @@ namespace SenseNovaU1 {
             ggml_set_name(attention_mask, "snu15.prefix.attention_mask");
             set_backend_tensor_data(attention_mask, attention_mask_vec.data());
 
-            auto runner_ctx = get_context();
+            auto runner_ctx = get_context(graph);
             auto text_model = model.text_model();
             auto hidden     = text_model->embed(&runner_ctx, ids);
             hidden          = text_model->forward(&runner_ctx,

src/model/vae/auto_encoder_kl.hpp

@@ -544,7 +544,7 @@ public:
                        const String2TensorStorage& tensor_storage_map = {},
                        const std::string& prefix                      = "")
         : version(version), decode_only(decode_only), use_video_decoder(use_video_decoder) {
-        if (sd_version_is_dit(version)) {
+        if (sd_version_is_dit(version) && version != VERSION_PIXART) {
             if (sd_version_uses_flux2_vae(version)) {
                 dd_config.z_channels = 32;
                 embed_dim            = 32;
@@ -678,7 +678,7 @@ struct AutoEncoderKL : public VAE {
         if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
             scale_factor = 0.18215f;
             shift_factor = 0.f;
-        } else if (sd_version_is_sdxl(version)) {
+        } else if (sd_version_is_sdxl(version) || sd_version_is_pixart(version)) {
             scale_factor = 0.13025f;
             shift_factor = 0.f;
         } else if (sd_version_is_sd3(version)) {

src/model/vae/ltx_vae.hpp

@@ -1300,7 +1300,7 @@ struct LTXVideoVAE : public VAE {
             feat_map[feat_idx] = get_cache_tensor_by_name(temporal_feat_cache_name(feat_idx));
         }
 
-        auto runner_ctx  = get_context();
+        auto runner_ctx  = get_context(gf);
         int feat_count   = 0;
         ggml_tensor* out = vae.decode_tiled_chunk(&runner_ctx,
                                                   z,
@@ -1313,8 +1313,7 @@ struct LTXVideoVAE : public VAE {
         for (int feat_idx = 0; feat_idx < feat_count && feat_idx < static_cast<int>(feat_map.size()); ++feat_idx) {
             ggml_tensor* feat_cache = feat_map[static_cast<size_t>(feat_idx)];
             if (feat_cache != nullptr) {
-                cache(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
-                ggml_build_forward_expand(gf, feat_cache);
+                runner_ctx.persist_cache_tensor(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
             }
         }
 

src/model/vae/minimax_h3_vae.hpp

@@ -556,12 +556,18 @@ namespace MiniMaxH3VAE {
                                    tensor.shape()[3]});
         }
 
-        static sd_tiling_params_t h3_tiling(sd_tiling_params_t params) {
+        sd_tiling_params_t resolve_tiling_params(sd_tiling_params_t params) const override {
+            if (!params.enabled) {
+                params.target_overlap = 0.25f;
+            }
+            if (params.tile_size_w == 0 && params.rel_size_w == 0.f) {
+                params.tile_size_w = 256;
+            }
+            if (params.tile_size_h == 0 && params.rel_size_h == 0.f) {
+                params.tile_size_h = 256;
+            }
             params.enabled         = true;
             params.temporal_tiling = false;
-            params.tile_size_x     = 16;
-            params.tile_size_y     = 16;
-            params.target_overlap  = 0.25f;
             return params;
         }
 
@@ -605,7 +611,7 @@ namespace MiniMaxH3VAE {
                                  bool circular_x = false,
                                  bool circular_y = false) override {
             auto input  = ensure_video_shape(x);
-            auto tiling = h3_tiling(tiling_params);
+            auto tiling = resolve_tiling_params(tiling_params);
             if (input.shape()[2] == 1) {
                 auto encoded = VAE::encode(n_threads, input, tiling, circular_x, circular_y);
                 if (!encoded.empty() && encoded.shape()[2] > 1) {
@@ -646,7 +652,7 @@ namespace MiniMaxH3VAE {
                                  bool circular_y   = false,
                                  bool silent       = false) override {
             auto input  = ensure_video_shape(x);
-            auto tiling = h3_tiling(tiling_params);
+            auto tiling = resolve_tiling_params(tiling_params);
             if (input.shape()[2] == 1) {
                 auto decoded = VAE::decode(n_threads,
                                            input,

src/model/vae/tae.hpp

@@ -701,7 +701,7 @@ public:
         bool use_midblock_gn = false;
         taef2                = sd_version_uses_flux2_vae(version);
 
-        if (sd_version_is_dit(version)) {
+        if (sd_version_is_dit(version) && !sd_version_is_pixart(version)) {
             z_channels = 16;
         }
         if (taef2) {

src/model/vae/vae.hpp

@@ -1,6 +1,9 @@
 #ifndef __SD_MODEL_VAE_VAE_HPP__
 #define __SD_MODEL_VAE_VAE_HPP__
 
+#include <cmath>
+#include <limits>
+
 #include "core/tensor_ggml.hpp"
 #include "model/common/block.hpp"
 #include "model/vae/vae_tiling.hpp"
@@ -117,8 +120,8 @@ protected:
                                     int output_width,
                                     int output_height,
                                     int scale,
-                                    int p_tile_size_x,
-                                    int p_tile_size_y,
+                                    int p_tile_size_w,
+                                    int p_tile_size_h,
                                     float tile_overlap_factor,
                                     bool circular_x,
                                     bool circular_y,
@@ -138,17 +141,28 @@ protected:
             }
             return output_tile;
         };
-        return ::process_tiles_2d(input,
-                                  output_width,
-                                  output_height,
-                                  scale,
-                                  p_tile_size_x,
-                                  p_tile_size_y,
-                                  tile_overlap_factor,
-                                  circular_x,
-                                  circular_y,
-                                  on_processing,
-                                  silent);
+        const bool original_circular_x = circular_x_enabled;
+        const bool original_circular_y = circular_y_enabled;
+        const int64_t latent_width     = decode_graph ? input.shape()[0] : output_width;
+        const int64_t latent_height    = decode_graph ? input.shape()[1] : output_height;
+        circular_x                     = circular_x || original_circular_x;
+        circular_y                     = circular_y || original_circular_y;
+        // Full-width axes wrap in convolutions; split axes wrap between tiles.
+        set_circular_axes(circular_x && p_tile_size_w >= latent_width,
+                          circular_y && p_tile_size_h >= latent_height);
+        auto output = ::process_tiles_2d(input,
+                                         output_width,
+                                         output_height,
+                                         scale,
+                                         p_tile_size_w,
+                                         p_tile_size_h,
+                                         tile_overlap_factor,
+                                         circular_x && p_tile_size_w < latent_width,
+                                         circular_y && p_tile_size_h < latent_height,
+                                         on_processing,
+                                         silent);
+        set_circular_axes(original_circular_x, original_circular_y);
+        return output;
     }
 
 public:
@@ -178,33 +192,48 @@ public:
         return supports_temporal_tiling(VAETemporalDirection::DECODE);
     }
 
-    void get_tile_sizes(int& tile_size_x,
-                        int& tile_size_y,
+    virtual sd_tiling_params_t resolve_tiling_params(sd_tiling_params_t params) const {
+        return params;
+    }
+
+    bool get_tile_sizes(int& tile_size_w,
+                        int& tile_size_h,
                         float& tile_overlap,
                         const sd_tiling_params_t& params,
-                        int64_t latent_x,
-                        int64_t latent_y,
-                        float encoding_factor = 1.0f) {
-        tile_overlap       = std::max(std::min(params.target_overlap, 0.5f), 0.0f);
-        auto get_tile_size = [&](int requested_size, float factor, int64_t latent_size) {
-            const int default_tile_size  = 32;
-            const int min_tile_dimension = 4;
-            int tile_size                = default_tile_size;
-            // factor <= 1 means simple fraction of the latent dimension
-            // factor > 1 means number of tiles across that dimension
-            if (factor > 0.f) {
-                if (factor > 1.0)
-                    factor = 1 / (factor - factor * tile_overlap + tile_overlap);
-                tile_size = static_cast<int>(std::round(latent_size * factor));
-            } else if (requested_size >= min_tile_dimension) {
-                tile_size = requested_size;
+                        int64_t latent_w,
+                        int64_t latent_h) {
+        const auto tiling = resolve_tiling_params(params);
+        if (latent_w <= 0 || latent_h <= 0 ||
+            latent_w > std::numeric_limits<int>::max() || latent_h > std::numeric_limits<int>::max() ||
+            !std::isfinite(tiling.target_overlap)) {
+            LOG_ERROR("invalid VAE tiling dimensions or overlap");
+            return false;
+        }
+        const int scale_factor = get_scale_factor();
+        tile_overlap           = std::max(std::min(tiling.target_overlap, 0.5f), 0.0f);
+        auto get_tile_size     = [&](int requested_size, double factor, int64_t latent_size, int& tile_size) {
+            if (requested_size < 0 || !std::isfinite(factor) || factor < 0.0) {
+                LOG_ERROR("VAE tile sizes and relative sizes must be finite and non-negative");
+                return false;
             }
-            tile_size = static_cast<int>(tile_size * encoding_factor);
-            return std::max(std::min(tile_size, static_cast<int>(latent_size)), min_tile_dimension);
+            const int min_tile_dimension = std::min(4, static_cast<int>(latent_size));
+            double size                  = (requested_size > 0 ? requested_size : 256) / scale_factor;
+            if (factor > 0.0) {
+                if (factor > 1.0) {
+                    factor = 1.0 / (factor * (1.0 - tile_overlap) + tile_overlap);
+                }
+                size = std::floor(static_cast<double>(latent_size) * factor);
+            }
+            if (size < min_tile_dimension && (requested_size > 0 || factor > 0.0)) {
+                LOG_ERROR("VAE tile size must be at least %d image pixels on this axis", min_tile_dimension * scale_factor);
+                return false;
+            }
+            tile_size = static_cast<int>(std::min(static_cast<double>(latent_size), std::max<double>(min_tile_dimension, size)));
+            return true;
         };
 
-        tile_size_x = get_tile_size(params.tile_size_x, params.rel_size_x, latent_x);
-        tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
+        return get_tile_size(tiling.tile_size_w, tiling.rel_size_w, latent_w, tile_size_w) &&
+               get_tile_size(tiling.tile_size_h, tiling.rel_size_h, latent_h, tile_size_h);
     }
 
     virtual sd::Tensor<float> encode(int n_threads,
@@ -213,6 +242,7 @@ public:
                                      bool circular_x = false,
                                      bool circular_y = false) {
         int64_t t0              = ggml_time_ms();
+        tiling_params           = resolve_tiling_params(tiling_params);
         sd::Tensor<float> input = x;
         sd::Tensor<float> output;
         if (scale_input) {
@@ -224,21 +254,19 @@ public:
             int64_t W              = input.shape()[0] / scale_factor;
             int64_t H              = input.shape()[1] / scale_factor;
             float tile_overlap;
-            int tile_size_x, tile_size_y;
-            // Image VAE encode is more sensitive to tile boundary context than decode.
-            // Keep the smaller legacy factor for video VAEs, but default image encode
-            // tiles to 64 latent pixels so a 512px SD image is encoded as one tile.
-            const float encode_tile_factor = sd_version_is_minimax_h3(version) ? 1.f : (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f
-                                                                                                                                                                                            : 2.0f;
-            get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor);
-            LOG_VERBOSE("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
+            int tile_size_w, tile_size_h;
+            if (!get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params, W, H)) {
+                return {};
+            }
+            LOG_VERBOSE("VAE encode tile size: %dx%d pixels (%dx%d latent)",
+                        tile_size_w * scale_factor, tile_size_h * scale_factor, tile_size_w, tile_size_h);
             output = tiled_compute(input,
                                    n_threads,
                                    static_cast<int>(W),
                                    static_cast<int>(H),
                                    scale_factor,
-                                   tile_size_x,
-                                   tile_size_y,
+                                   tile_size_w,
+                                   tile_size_h,
                                    tile_overlap,
                                    circular_x,
                                    circular_y,
@@ -271,6 +299,7 @@ public:
                                      bool circular_y   = false,
                                      bool silent       = false) {
         int64_t t0              = ggml_time_ms();
+        tiling_params           = resolve_tiling_params(tiling_params);
         sd::Tensor<float> input = x;
         sd::Tensor<float> output;
 
@@ -279,10 +308,13 @@ public:
             int64_t W              = input.shape()[0] * scale_factor;
             int64_t H              = input.shape()[1] * scale_factor;
             float tile_overlap;
-            int tile_size_x, tile_size_y;
-            get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, input.shape()[0], input.shape()[1]);
+            int tile_size_w, tile_size_h;
+            if (!get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params, input.shape()[0], input.shape()[1])) {
+                return {};
+            }
             if (!silent) {
-                LOG_VERBOSE("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
+                LOG_VERBOSE("VAE decode tile size: %dx%d pixels (%dx%d latent)",
+                            tile_size_w * scale_factor, tile_size_h * scale_factor, tile_size_w, tile_size_h);
             }
             output = tiled_compute(
                 input,
@@ -290,8 +322,8 @@ public:
                 static_cast<int>(W),
                 static_cast<int>(H),
                 scale_factor,
-                tile_size_x,
-                tile_size_y,
+                tile_size_w,
+                tile_size_h,
                 tile_overlap,
                 circular_x,
                 circular_y,

src/model/vae/wan_vae.hpp

@@ -24,6 +24,7 @@ namespace WAN {
         std::tuple<int, int, int> padding;
         std::tuple<int, int, int> dilation;
         bool bias;
+        float scale = 1.f;
 
         void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
             auto weight = tensor_storage_map.find(prefix + "weight");
@@ -60,6 +61,10 @@ namespace WAN {
               dilation(std::move(dilation)),
               bias(bias) {}
 
+        void set_scale(float scale_value) {
+            scale = scale_value;
+        }
+
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cache_x = nullptr) {
             // x: [N*IC, ID, IH, IW]
             // result: x: [N*OC, ID, IH, IW]
@@ -93,14 +98,14 @@ namespace WAN {
                 x2              = ggml_ext_conv_2d(ctx->ggml_ctx, x2, w2, b,
                                                    std::get<2>(stride), std::get<1>(stride), 0, 0,
                                                    std::get<2>(dilation), std::get<1>(dilation),
-                                                   ctx->conv2d_direct_enabled);
+                                                   ctx->conv2d_direct_enabled, false, false, scale);
                 return ggml_reshape_4d(ctx->ggml_ctx, x2, x2->ne[0], x2->ne[1], 1, out_channels);
             }
             return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
                                     std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                     0, 0, 0,
                                     std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
-                                    false, ctx->conv3d_direct_enabled);
+                                    false, ctx->conv3d_direct_enabled, scale);
         }
     };
 
@@ -1117,6 +1122,19 @@ namespace WAN {
             } else {
                 blocks["conv2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1}));
             }
+            if (version == VERSION_QWEN_IMAGE_2_1) {
+                // Keep large VAE activations within the FP16 convolution range.
+                const float conv_scale = 1.f / 128.f;
+                std::vector<GGMLBlock*> all_blocks;
+                get_all_blocks(all_blocks);
+                for (auto block : all_blocks) {
+                    if (auto conv = dynamic_cast<Conv2d*>(block)) {
+                        conv->set_scale(conv_scale);
+                    } else if (auto conv = dynamic_cast<CausalConv3d*>(block)) {
+                        conv->set_scale(conv_scale);
+                    }
+                }
+            }
         }
 
         static ggml_tensor* patchify(ggml_context* ctx,
@@ -1443,15 +1461,14 @@ namespace WAN {
 
             ggml_tensor* z = make_input(z_tensor);
 
-            auto runner_ctx = get_context();
+            auto runner_ctx = get_context(gf);
 
             ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx);
 
             for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) {
                 ggml_tensor* feat_cache = ae._feat_map[feat_idx];
                 if (feat_cache != nullptr) {
-                    cache("feat_idx:" + std::to_string(feat_idx), feat_cache);
-                    ggml_build_forward_expand(gf, feat_cache);
+                    runner_ctx.persist_cache_tensor("feat_idx:" + std::to_string(feat_idx), feat_cache);
                 }
             }
 

src/model_loader.cpp

@@ -322,7 +322,13 @@ bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path,
     }
 
     for (const std::string& shard_path : shard_paths) {
-        if (!parse_file(shard_path, prefix)) {
+        FileStamp stamp;
+        if (!read_file_stamp(shard_path, stamp)) {
+            return false;
+        }
+        parsed_dependencies_.push_back(stamp);
+        LOG_INFO("load %s using safetensors format", shard_path.c_str());
+        if (!init_from_safetensors_file(shard_path, prefix)) {
             return false;
         }
     }
@@ -526,7 +532,18 @@ SDVersion ModelLoader::get_sd_version() const {
         if (tensor_storage.name.find("model.diffusion_model.layers.0.adaLN_sa_ln.weight") != std::string::npos) {
             return VERSION_ERNIE_IMAGE;
         }
+        if (tensor_storage.name.find("model.diffusion_model.t_block.1.weight") != std::string::npos &&
+            tensor_storage_map.find("model.diffusion_model.x_embedder.proj.weight") != tensor_storage_map.end() &&
+            tensor_storage_map.find("model.diffusion_model.audio_patchify_proj.weight") == tensor_storage_map.end()) {
+            return VERSION_PIXART;
+        }
         if (tensor_storage.name.find("model.diffusion_model.adaln_single.emb.timestep_embedder.linear_1.bias") != std::string::npos) {
+            // PixArt shares this timestep embedding with LTX-AV.
+            if (tensor_storage_map.find("model.diffusion_model.pos_embed.proj.weight") != tensor_storage_map.end() &&
+                tensor_storage_map.find("model.diffusion_model.adaln_single.linear.weight") != tensor_storage_map.end() &&
+                tensor_storage_map.find("model.diffusion_model.audio_patchify_proj.weight") == tensor_storage_map.end()) {
+                return VERSION_PIXART;
+            }
             return VERSION_LTXAV;
         }
         if (tensor_storage.name.find("model.diffusion_model.video_patch_proj.weight") != std::string::npos &&
@@ -1084,10 +1101,12 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
         if (tensors_to_process.empty()) {
             continue;
         }
-        LOG_VERBOSE("loading %zu/%zu tensors from %s",
-                    tensors_to_process.size(),
-                    file_tensors.size(),
-                    file_path.c_str());
+        if (log_progress) {
+            LOG_VERBOSE("loading %zu/%zu tensors from %s",
+                        tensors_to_process.size(),
+                        file_tensors.size(),
+                        file_path.c_str());
+        }
 
         bool is_zip = fdata.is_zip;
 
@@ -1591,6 +1610,9 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage
             // Pass, do not convert. For Unet
         } else if (contains(name, "embedding")) {
             // Pass, do not convert embedding
+        } else if (contains(name, "scale_shift_table")) {
+            // Pass, do not convert. adaLN modulation tables (PixArt, LTXV) are sliced
+            // element-wise, which is invalid on quantized block layouts.
         } else if (ends_with(name, "_pad_token")) {
             // Pass, do not convert. LLaDA-Image stores its pad tokens far outside the f16
             // range, so any format with an f16 scale or payload turns them into inf.

src/model_manager.cpp

@@ -246,17 +246,13 @@ bool ModelManager::register_param_tensors(ModelComponent component,
         }
         ggml_set_name(tensor, name.c_str());
 
-        auto state            = std::make_unique<TensorState>();
-        state->name           = name;
-        state->tensor         = tensor;
-        state->component      = component;
-        state->source_file    = source_file;
-        state->source_version = source_version;
-        auto source           = sources.find(name);
-        if (source != sources.end()) {
-            state->source     = source->second;
-            state->has_source = true;
-        }
+        auto state             = std::make_unique<TensorState>();
+        state->name            = name;
+        state->tensor          = tensor;
+        state->component       = component;
+        state->source_file     = source_file;
+        state->source_version  = source_version;
+        state->sources         = find_tensor_sources(*state, sources);
         state->residency_mode  = residency_mode;
         state->compute_backend = compute_backend;
         state->params_backend  = params_backend;
@@ -757,12 +753,33 @@ bool ModelManager::validate_tensor(const TensorState& state) const {
         return true;
     }
 
-    if (!state.has_source) {
+    if (state.sources.empty()) {
         LOG_ERROR("%s tensor '%s' not in model metadata", model_component_name(state.component), state.name.c_str());
         return false;
     }
 
-    const TensorStorage& tensor_storage = state.source;
+    TensorStorage tensor_storage = state.sources.front();
+    if (state.sources.size() > 1) {
+        const int dim = tensor_storage.n_dims - 1;
+        if (dim < 0 || dim >= GGML_MAX_DIMS) {
+            return false;
+        }
+        tensor_storage.ne[dim] = 0;
+        for (const auto& part : state.sources) {
+            if (part.n_dims != tensor_storage.n_dims || part.ne[dim] < 0 ||
+                part.ne[dim] > state.tensor->ne[dim] - tensor_storage.ne[dim]) {
+                LOG_ERROR("invalid tensor part '%s' for '%s'", part.name.c_str(), state.name.c_str());
+                return false;
+            }
+            for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+                if (i != dim && part.ne[i] != tensor_storage.ne[i]) {
+                    LOG_ERROR("incompatible tensor part '%s' for '%s'", part.name.c_str(), state.name.c_str());
+                    return false;
+                }
+            }
+            tensor_storage.ne[dim] += part.ne[dim];
+        }
+    }
     if (state.tensor->ne[0] != tensor_storage.ne[0] ||
         state.tensor->ne[1] != tensor_storage.ne[1] ||
         state.tensor->ne[2] != tensor_storage.ne[2] ||
@@ -831,7 +848,7 @@ bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
 }
 
 bool ModelManager::can_mmap_storage(const TensorState& state) const {
-    if (state.source_file != 0 || !enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
+    if (state.sources.size() > 1 || state.source_file != 0 || !enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
         return false;
     }
     if (state.compute_backend == nullptr || state.params_backend == nullptr) {
@@ -941,6 +958,62 @@ bool ModelManager::alloc_params_buffers(const std::vector<TensorState*>& states,
     return true;
 }
 
+bool ModelManager::load_tensor_parts(TensorState& state) {
+    auto ctx = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
+        ggml_init({state.sources.size() * ggml_tensor_overhead(), nullptr, true}), ggml_free);
+    if (!ctx) {
+        return false;
+    }
+    const size_t size = ggml_nbytes(state.tensor);
+    std::vector<uint8_t> buffer;
+    void* data = state.tensor->data;
+    if (!ggml_backend_buffer_is_host(state.tensor->buffer)) {
+        buffer.resize(size);
+        data = buffer.data();
+    }
+    std::map<std::string, ggml_tensor*> parts;
+    std::set<std::string> names;
+    size_t offset = 0;
+    for (const auto& source : state.sources) {
+        auto part              = ggml_new_tensor(ctx.get(), state.tensor->type, source.n_dims, source.ne);
+        const size_t part_size = ggml_nbytes(part);
+        if (part_size > size - offset) {
+            return false;
+        }
+        part->data         = static_cast<uint8_t*>(data) + offset;
+        parts[source.name] = part;
+        names.insert(source.name);
+        offset += part_size;
+    }
+    if (offset != size) {
+        return false;
+    }
+    std::set<std::string> loaded;
+    std::mutex mutex;
+    auto callback = [&](const TensorStorage& source, ggml_tensor** dst) {
+        *dst      = nullptr;
+        auto part = parts.find(source.name);
+        if (part != parts.end()) {
+            *dst = part->second;
+            std::lock_guard<std::mutex> lock(mutex);
+            loaded.insert(source.name);
+        }
+        return true;
+    };
+    const bool success = state.source_file == 0
+                             ? model_loader_.load_tensors(callback, enable_mmap_, &names, false)
+                             : model_loader_.load_file_tensors(state.source_file, state.source_version, callback, names, enable_mmap_);
+    if (!success || loaded != names) {
+        return false;
+    }
+    if (!buffer.empty()) {
+        // Upload the assembled tensor once, including for row-split backend buffers.
+        ggml_backend_tensor_set(state.tensor, buffer.data(), 0, size);
+    }
+    state.loaded_to_params_backend = true;
+    return true;
+}
+
 bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
     using ReadGroup = std::pair<ModelLoader::FileId, SDVersion>;
     using ReadBatch = std::map<std::string, std::vector<TensorState*>>;
@@ -948,6 +1021,12 @@ bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
     for (auto* state : states) {
         if (state == nullptr)
             continue;
+        if (state->sources.size() > 1) {
+            if (!load_tensor_parts(*state)) {
+                return false;
+            }
+            continue;
+        }
         auto& batches = groups[{state->source_file, state->source_version}];
         // The loader supplies one destination per name; only conflicting types need another batch.
         auto batch = std::find_if(batches.begin(), batches.end(), [&](const ReadBatch& candidate) {

src/model_manager.h

@@ -38,8 +38,7 @@ private:
         std::string name;
         ggml_tensor* tensor      = nullptr;
         ModelComponent component = ModelComponent::Count;
-        TensorStorage source;
-        bool has_source                 = false;
+        std::vector<TensorStorage> sources;
         ModelLoader::FileId source_file = 0;
         SDVersion source_version        = VERSION_COUNT;
 
@@ -138,6 +137,8 @@ private:
     bool apply_loras_to_params(const std::vector<TensorState*>& states);
     bool mmap_params(const std::vector<TensorState*>& states,
                      std::vector<ParamsStorageBlock*>& created_storage_blocks);
+    static std::vector<TensorStorage> find_tensor_sources(const TensorState& state, const String2TensorStorage& sources);
+    bool load_tensor_parts(TensorState& state);
     bool can_mmap_storage(const TensorState& state) const;
     bool alloc_params_buffers(const std::vector<TensorState*>& states,
                               std::vector<ParamsStorageBlock*>& created_storage_blocks);

src/model_manager_files.cpp

@@ -15,6 +15,26 @@ static bool same_tensor_source(const TensorStorage& a, const TensorStorage& b) {
            a.int8_convrot_group_size == b.int8_convrot_group_size;
 }
 
+std::vector<TensorStorage> ModelManager::find_tensor_sources(const TensorState& state, const String2TensorStorage& sources) {
+    auto first = sources.find(state.name);
+    if (first == sources.end()) {
+        return {};
+    }
+    std::vector<TensorStorage> result{first->second};
+    if (state.component == ModelComponent::LoRA ||
+        std::equal(first->second.ne, first->second.ne + GGML_MAX_DIMS, state.tensor->ne)) {
+        return result;
+    }
+    for (size_t i = 1;; ++i) {
+        auto part = sources.find(state.name + "." + std::to_string(i));
+        if (part == sources.end()) {
+            break;
+        }
+        result.push_back(part->second);
+    }
+    return result;
+}
+
 void ModelManager::invalidate_sources(const std::unordered_set<TensorState*>& states) {
     auto affected = states;
     for (const auto& block : params_storage_blocks_) {
@@ -76,20 +96,16 @@ bool ModelManager::set_loader(ModelLoader loader) {
     }
     std::unordered_set<TensorState*> changed;
     for (const auto& state : tensor_states_) {
-        const auto& sources = sources_for(*state);
-        auto source         = sources.find(state->name);
-        const bool found    = source != sources.end();
-        if (found != state->has_source || (found && !same_tensor_source(state->source, source->second)) ||
+        const auto sources = find_tensor_sources(*state, sources_for(*state));
+        if (sources.size() != state->sources.size() ||
+            !std::equal(sources.begin(), sources.end(), state->sources.begin(), same_tensor_source) ||
             (lora_changed && state->component != ModelComponent::LoRA && state->applied_lora_epoch != UINT64_MAX)) {
             changed.insert(state.get());
         }
     }
     invalidate_sources(changed);
     for (auto* state : changed) {
-        const auto& sources = sources_for(*state);
-        auto source         = sources.find(state->name);
-        state->has_source   = source != sources.end();
-        state->source       = state->has_source ? source->second : TensorStorage{};
+        state->sources = find_tensor_sources(*state, sources_for(*state));
     }
     if (lora_changed) {
         ++current_lora_epoch_;
@@ -134,9 +150,8 @@ ModelLoader::FileVersions ModelManager::source_versions(const std::set<ModelComp
             versions[state->source_file] = loader.file_revision(state->source_file);
             continue;
         }
-        auto source = sources.find(state->name);
-        if (source != sources.end()) {
-            versions[source->second.file_id] = source->second.file_revision;
+        for (const auto& source : find_tensor_sources(*state, sources)) {
+            versions[source.file_id] = source.file_revision;
         }
     }
     return versions;

src/name_conversion.cpp

@@ -932,6 +932,68 @@ static bool is_diffusers_controlnet_name(const std::string& name) {
     return false;
 }
 
+static std::string convert_diffusers_dit_to_original_pixart(std::string name) {
+    static const std::vector<std::pair<std::string, std::string>> prefix_map = {
+        {"pos_embed.proj.", "x_embedder.proj."},
+        {"adaln_single.emb.timestep_embedder.linear_1.", "t_embedder.mlp.0."},
+        {"adaln_single.emb.timestep_embedder.linear_2.", "t_embedder.mlp.2."},
+        {"adaln_single.emb.resolution_embedder.linear_1.", "csize_embedder.mlp.0."},
+        {"adaln_single.emb.resolution_embedder.linear_2.", "csize_embedder.mlp.2."},
+        {"adaln_single.emb.aspect_ratio_embedder.linear_1.", "ar_embedder.mlp.0."},
+        {"adaln_single.emb.aspect_ratio_embedder.linear_2.", "ar_embedder.mlp.2."},
+        {"adaln_single.linear.", "t_block.1."},
+        {"caption_projection.linear_1.", "y_embedder.y_proj.fc1."},
+        {"caption_projection.linear_2.", "y_embedder.y_proj.fc2."},
+        {"proj_out.", "final_layer.linear."},
+    };
+    for (const auto& entry : prefix_map) {
+        if (starts_with(name, entry.first)) {
+            return entry.second + name.substr(entry.first.size());
+        }
+    }
+    if (name == "scale_shift_table") {
+        return "final_layer.scale_shift_table";
+    }
+    const std::string block_prefix = "transformer_blocks.";
+    if (!starts_with(name, block_prefix)) {
+        return name;
+    }
+    const size_t block_end = name.find('.', block_prefix.size());
+    if (block_end == std::string::npos) {
+        return name;
+    }
+    const std::string prefix                                                = "blocks." + name.substr(block_prefix.size(), block_end - block_prefix.size()) + ".";
+    name                                                                    = name.substr(block_end + 1);
+    static const std::vector<std::pair<std::string, std::string>> block_map = {
+        {"attn1.to_q.", "attn.qkv."},
+        {"attn1.to_out.0.", "attn.proj."},
+        {"attn2.to_q.", "cross_attn.q_linear."},
+        {"attn2.to_k.", "cross_attn.kv_linear."},
+        {"attn2.to_out.0.", "cross_attn.proj."},
+        {"ff.net.0.proj.", "mlp.fc1."},
+        {"ff.net.2.", "mlp.fc2."},
+    };
+    for (const auto& entry : block_map) {
+        if (starts_with(name, entry.first)) {
+            return prefix + entry.second + name.substr(entry.first.size());
+        }
+    }
+    static const std::vector<std::pair<std::string, std::string>> part_map = {
+        {"attn1.to_k.weight", "attn.qkv.weight.1"},
+        {"attn1.to_k.bias", "attn.qkv.bias.1"},
+        {"attn1.to_v.weight", "attn.qkv.weight.2"},
+        {"attn1.to_v.bias", "attn.qkv.bias.2"},
+        {"attn2.to_v.weight", "cross_attn.kv_linear.weight.1"},
+        {"attn2.to_v.bias", "cross_attn.kv_linear.bias.1"},
+    };
+    for (const auto& entry : part_map) {
+        if (name == entry.first || starts_with(name, entry.first + ".")) {
+            return prefix + entry.second + name.substr(entry.first.size());
+        }
+    }
+    return prefix + name;
+}
+
 std::string convert_diffusion_model_name(std::string name, std::string prefix, SDVersion version) {
     if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
         name = convert_diffusers_unet_to_original_sd1(name);
@@ -951,6 +1013,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
         name = convert_other_dit_to_original_anima(name);
     } else if (sd_version_is_krea2(version)) {
         name = convert_diffusers_dit_to_original_krea2(name);
+    } else if (sd_version_is_pixart(version)) {
+        name = convert_diffusers_dit_to_original_pixart(name);
     }
     return name;
 }

src/pipeline/diffusion_engine.cpp

@@ -105,6 +105,7 @@ const char* model_version_to_str[] = {
     "SenseNova U1.5",
     "LLaDA-Image",
     "ESRGAN",
+    "PixArt",
 };
 
 static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]),
@@ -125,6 +126,18 @@ void calculate_alphas_cumprod(float* alphas_cumprod,
     }
 }
 
+void calculate_alphas_cumprod_linear_beta(float* alphas_cumprod,
+                                          float beta_start,
+                                          float beta_end,
+                                          int timesteps = TIMESTEPS) {
+    float product = 1.0f;
+    for (int i = 0; i < timesteps; i++) {
+        float beta = beta_start + (beta_end - beta_start) * ((float)i / (timesteps - 1));
+        product *= 1.0f - beta;
+        alphas_cumprod[i] = product;
+    }
+}
+
 template <typename T, typename = void>
 struct has_set_runtime_backends : std::false_type {};
 template <typename T>
@@ -666,6 +679,10 @@ void StableDiffusionGGML::refresh_compvis_denoiser_sigmas() {
     std::vector<float> alphas_cumprod(TIMESTEPS);
     if (file_alphas_cumprod.size() == TIMESTEPS) {
         alphas_cumprod = file_alphas_cumprod;
+    } else if (sd_version_is_pixart(version)) {
+        // PixArt checkpoints train with a linear beta schedule (0.0001 -> 0.02)
+        // instead of the scaled_linear schedule used by SD1.x/SDXL.
+        calculate_alphas_cumprod_linear_beta(alphas_cumprod.data(), 0.0001f, 0.02f);
     } else {
         calculate_alphas_cumprod(alphas_cumprod.data());
     }
@@ -2731,7 +2748,7 @@ int StableDiffusionGGML::get_diffusion_model_down_factor() {
     if (sd_version_is_dit(version)) {
         if (sd_version_is_sensenova_u1(version)) {
             down_factor = 32;
-        } else if (version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version)) {
+        } else if (version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version) || sd_version_is_pixart(version)) {
             down_factor = 2;
         } else {
             down_factor = 1;
@@ -2769,6 +2786,8 @@ int StableDiffusionGGML::get_latent_channel() {
             latent_channel = 128;
         } else if (sd_version_is_mage_flow(version)) {
             latent_channel = 128;
+        } else if (sd_version_is_pixart(version)) {
+            latent_channel = 4;
         } else {
             latent_channel = 16;
         }
@@ -2864,14 +2883,26 @@ sd::Tensor<float> StableDiffusionGGML::decode_first_stage(const sd::Tensor<float
         return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f);
     }
     auto latents                      = first_stage_model->diffusion_to_vae_latents(x);
-    auto decoded                      = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
-    const bool prefer_temporal_tiling = decode_video && first_stage_model->can_temporal_tile_decode();
-    while (decoded.empty() &&
-           sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling,
-                                                            first_stage_model->last_compute_status())) {
-        decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
+    auto tiling_params                = first_stage_model->resolve_tiling_params(vae_tiling_params);
+    const bool prefer_temporal_tiling = decode_video && latents.dim() == 5 && latents.shape()[2] > 1 &&
+                                        first_stage_model->can_temporal_tile_decode();
+    for (;;) {
+        int tile_size_w = static_cast<int>(latents.shape()[0]);
+        int tile_size_h = static_cast<int>(latents.shape()[1]);
+        float tile_overlap;
+        if (tiling_params.enabled &&
+            !first_stage_model->get_tile_sizes(tile_size_w, tile_size_h, tile_overlap, tiling_params,
+                                               latents.shape()[0], latents.shape()[1])) {
+            return {};
+        }
+        auto decoded = first_stage_model->decode(n_threads, latents, tiling_params, decode_video, circular_x, circular_y);
+        if (!decoded.empty() ||
+            !sd::backend_fit::prepare_vae_decode_retry_tiling(tiling_params, prefer_temporal_tiling,
+                                                              first_stage_model->last_compute_status(),
+                                                              tile_size_w, tile_size_h, first_stage_model->get_scale_factor())) {
+            return decoded;
+        }
     }
-    return decoded;
 }
 
 sd::Tensor<float> StableDiffusionGGML::normalize_ltx_video_latents(const sd::Tensor<float>& x) {

src/pipeline/image.cpp

@@ -35,19 +35,21 @@ namespace sd::pipeline {
             return original_axes;
         }
 
-        int tile_size_x, tile_size_y;
+        int tile_size_w, tile_size_h;
         float overlap;
-        int latent_size_x = request.width / request.vae_scale_factor;
-        int latent_size_y = request.height / request.vae_scale_factor;
-        sd->first_stage_model->get_tile_sizes(tile_size_x,
-                                              tile_size_y,
-                                              overlap,
-                                              sd_img_gen_params->vae_tiling_params,
-                                              latent_size_x,
-                                              latent_size_y);
+        int latent_size_w = request.width / request.vae_scale_factor;
+        int latent_size_h = request.height / request.vae_scale_factor;
+        if (!sd->first_stage_model->get_tile_sizes(tile_size_w,
+                                                   tile_size_h,
+                                                   overlap,
+                                                   sd_img_gen_params->vae_tiling_params,
+                                                   latent_size_w,
+                                                   latent_size_h)) {
+            return original_axes;
+        }
 
-        sd->circular_x = sd->circular_x && (tile_size_x >= latent_size_x);
-        sd->circular_y = sd->circular_y && (tile_size_y >= latent_size_y);
+        sd->circular_x = sd->circular_x && (tile_size_w >= latent_size_w);
+        sd->circular_y = sd->circular_y && (tile_size_h >= latent_size_h);
 
         if (sd->first_stage_model) {
             sd->first_stage_model->set_circular_axes(sd->circular_x, sd->circular_y);
@@ -56,8 +58,8 @@ namespace sd::pipeline {
             sd->preview_vae->set_circular_axes(sd->circular_x, sd->circular_y);
         }
 
-        sd->circular_x = original_axes.circular_x && (tile_size_x < latent_size_x);
-        sd->circular_y = original_axes.circular_y && (tile_size_y < latent_size_y);
+        sd->circular_x = original_axes.circular_x && (tile_size_w < latent_size_w);
+        sd->circular_y = original_axes.circular_y && (tile_size_h < latent_size_h);
 
         return original_axes;
     }

src/pipeline/model_builders.cpp

@@ -1,5 +1,6 @@
 #include "model_builders.h"
 
+#include <cmath>
 #include <cstring>
 #include <utility>
 
@@ -28,6 +29,7 @@
 #include "model/diffusion/mmdit.hpp"
 #include "model/diffusion/model.hpp"
 #include "model/diffusion/pid.hpp"
+#include "model/diffusion/pixart.hpp"
 #include "model/diffusion/qwen_image.hpp"
 #include "model/diffusion/qwen_image_2_1.hpp"
 #include "model/diffusion/sensenova_u1.h"
@@ -301,6 +303,19 @@ namespace sd::model_builders {
                                                                            weight_manager,
                                                                            sd_ctx_params->model_args);
             }
+        } else if (version == VERSION_PIXART) {
+            result.conditioner = std::make_shared<T5CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                  tensor_storage_map,
+                                                                  true,
+                                                                  0,
+                                                                  false,
+                                                                  weight_manager,
+                                                                  sd_ctx_params->model_args);
+            result.diffusion   = std::make_shared<PixArt::PixArtRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                      tensor_storage_map,
+                                                                      "model.diffusion_model",
+                                                                      weight_manager,
+                                                                      sd_ctx_params->model_args);
         } else if (sd_version_is_mage_flow(version)) {
             result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                tensor_storage_map,
@@ -562,6 +577,23 @@ namespace sd::model_builders {
                                                              false,
                                                              vae_version,
                                                              weight_manager);
+                if (sd_version_is_pixart(version)) {
+                    // Alpha-512 and Sigma share tensor layouts; Alpha-512 needs an explicit scale override.
+                    if (tensor_storage_map.count("model.diffusion_model.csize_embedder.mlp.0.weight") != 0) {
+                        model->scale_factor = 0.18215f;
+                    }
+                    for (const auto& [key, value] : parse_key_value_args(sd_ctx_params->model_args, "model arg")) {
+                        if (key == "pixart_vae_scale_factor") {
+                            float parsed = 0.f;
+                            if (parse_strict_float(value, parsed) && std::isfinite(parsed) && parsed > 0.f) {
+                                model->scale_factor = parsed;
+                            } else {
+                                LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str());
+                            }
+                        }
+                    }
+                    LOG_VERBOSE("pixart: VAE scale factor = %.5f", model->scale_factor);
+                }
                 if (sd_version_is_sdxl(version) &&
                     (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || options.external_vae_is_invalid)) {
                     float vae_conv_2d_scale = 1.f / 32.f;

src/runtime/denoiser.hpp

@@ -728,12 +728,20 @@ inline float flux_time_shift(float mu, float sigma, float t) {
 
 // https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py#L289
 struct FluxScheduler : SigmaScheduler {
-    int image_seq_len = 0;
-    float base_shift  = 0.5f;
-    float max_shift   = 1.15f;
+    int image_seq_len      = 0;
+    int base_image_seq_len = 256;
+    int max_image_seq_len  = 4096;
+    float base_shift       = 0.5f;
+    float max_shift        = 1.15f;
+    float shift_terminal   = 0.0f;
 
-    explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr)
+    FluxScheduler(int image_seq_len, SDVersion version, const char* extra_sample_args = nullptr)
         : image_seq_len(image_seq_len) {
+        if (version == VERSION_QWEN_IMAGE_2_1) {
+            max_image_seq_len = 8192;
+            max_shift         = 0.9f;
+            shift_terminal    = 0.02f;
+        }
         parse_extra_sample_args(extra_sample_args);
     }
 
@@ -752,10 +760,8 @@ struct FluxScheduler : SigmaScheduler {
     }
 
     float compute_mu() const {
-        constexpr float base_shift_anchor = 256.0f;
-        constexpr float max_shift_anchor  = 4096.0f;
-        float m                           = (max_shift - base_shift) / (max_shift_anchor - base_shift_anchor);
-        float b                           = base_shift - m * base_shift_anchor;
+        float m = (max_shift - base_shift) / static_cast<float>(max_image_seq_len - base_image_seq_len);
+        float b = base_shift - m * static_cast<float>(base_image_seq_len);
         return static_cast<float>(image_seq_len) * m + b;
     }
 
@@ -764,7 +770,7 @@ struct FluxScheduler : SigmaScheduler {
         sigmas.reserve(n + 1);
 
         float mu = compute_mu();
-        LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu);
+        LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f, shift_terminal=%.3f", image_seq_len, n, mu, shift_terminal);
 
         if (n == 0) {
             sigmas.push_back(1.0f);
@@ -780,6 +786,16 @@ struct FluxScheduler : SigmaScheduler {
             }
         }
 
+        if (shift_terminal > 0.0f && n > 1) {
+            // The terminal shift applies to the last model evaluation, not the final zero sigma.
+            float scale_factor = (1.0f - sigmas[n - 1]) / (1.0f - shift_terminal);
+            if (std::isfinite(scale_factor) && scale_factor > 0.0f) {
+                for (uint32_t i = 0; i < n; ++i) {
+                    sigmas[i] = 1.0f - (1.0f - sigmas[i]) / scale_factor;
+                }
+            }
+        }
+
         sigmas[n] = 0.0f;
         return sigmas;
     }
@@ -1178,7 +1194,7 @@ struct Denoiser {
             }
             case FLUX_SCHEDULER: {
                 LOG_INFO("get_sigmas with Flux scheduler");
-                scheduler = std::make_shared<FluxScheduler>(image_seq_len, extra_sample_args);
+                scheduler = std::make_shared<FluxScheduler>(image_seq_len, version, extra_sample_args);
                 break;
             }
             default:

src/runtime/tiling.cpp

@@ -142,8 +142,8 @@ sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
                                    int output_width,
                                    int output_height,
                                    int scale,
-                                   int p_tile_size_x,
-                                   int p_tile_size_y,
+                                   int p_tile_size_w,
+                                   int p_tile_size_h,
                                    float tile_overlap_factor,
                                    bool circular_x,
                                    bool circular_y,
@@ -168,28 +168,28 @@ sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
 
     int num_tiles_x;
     float tile_overlap_factor_x;
-    sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_x, tile_overlap_factor, circular_x);
+    sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_w, tile_overlap_factor, circular_x);
 
     int num_tiles_y;
     float tile_overlap_factor_y;
-    sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_y, tile_overlap_factor, circular_y);
+    sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_h, tile_overlap_factor, circular_y);
 
-    int tile_overlap_x     = static_cast<int32_t>(p_tile_size_x * tile_overlap_factor_x);
-    int non_tile_overlap_x = p_tile_size_x - tile_overlap_x;
-    int tile_overlap_y     = static_cast<int32_t>(p_tile_size_y * tile_overlap_factor_y);
-    int non_tile_overlap_y = p_tile_size_y - tile_overlap_y;
-    int tile_size_x        = p_tile_size_x < small_width ? p_tile_size_x : small_width;
-    int tile_size_y        = p_tile_size_y < small_height ? p_tile_size_y : small_height;
-    int input_tile_size_x  = tile_size_x;
-    int input_tile_size_y  = tile_size_y;
-    int output_tile_size_x = tile_size_x;
-    int output_tile_size_y = tile_size_y;
+    int tile_overlap_x     = static_cast<int32_t>(p_tile_size_w * tile_overlap_factor_x);
+    int non_tile_overlap_x = p_tile_size_w - tile_overlap_x;
+    int tile_overlap_y     = static_cast<int32_t>(p_tile_size_h * tile_overlap_factor_y);
+    int non_tile_overlap_y = p_tile_size_h - tile_overlap_y;
+    int tile_size_w        = p_tile_size_w < small_width ? p_tile_size_w : small_width;
+    int tile_size_h        = p_tile_size_h < small_height ? p_tile_size_h : small_height;
+    int input_tile_size_w  = tile_size_w;
+    int input_tile_size_h  = tile_size_h;
+    int output_tile_size_w = tile_size_w;
+    int output_tile_size_h = tile_size_h;
     if (decode) {
-        output_tile_size_x *= scale;
-        output_tile_size_y *= scale;
+        output_tile_size_w *= scale;
+        output_tile_size_h *= scale;
     } else {
-        input_tile_size_x *= scale;
-        input_tile_size_y *= scale;
+        input_tile_size_w *= scale;
+        input_tile_size_h *= scale;
     }
 
     int num_tiles   = num_tiles_x * num_tiles_y;
@@ -205,9 +205,9 @@ sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
     }
     for (int y = 0; y < small_height && !last_y; y += non_tile_overlap_y) {
         int dy = 0;
-        if (!circular_y && y + tile_size_y >= small_height) {
+        if (!circular_y && y + tile_size_h >= small_height) {
             int original_y = y;
-            y              = small_height - tile_size_y;
+            y              = small_height - tile_size_h;
             dy             = original_y - y;
             if (decode) {
                 dy *= scale;
@@ -216,9 +216,9 @@ sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
         }
         for (int x = 0; x < small_width && !last_x; x += non_tile_overlap_x) {
             int dx = 0;
-            if (!circular_x && x + tile_size_x >= small_width) {
+            if (!circular_x && x + tile_size_w >= small_width) {
                 int original_x = x;
-                x              = small_width - tile_size_x;
+                x              = small_width - tile_size_w;
                 dx             = original_x - x;
                 if (decode) {
                     dx *= scale;
@@ -235,12 +235,12 @@ sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
             int overlap_y_out = decode ? tile_overlap_y * scale : tile_overlap_y;
 
             int64_t t1       = ggml_time_ms();
-            auto input_tile  = sd_tensor_split_2d(input, input_tile_size_x, input_tile_size_y, x_in, y_in);
+            auto input_tile  = sd_tensor_split_2d(input, input_tile_size_w, input_tile_size_h, x_in, y_in);
             auto output_tile = on_processing(input_tile);
             if (output_tile.empty()) {
                 return {};
             }
-            GGML_ASSERT(output_tile.shape()[0] == output_tile_size_x && output_tile.shape()[1] == output_tile_size_y);
+            GGML_ASSERT(output_tile.shape()[0] == output_tile_size_w && output_tile.shape()[1] == output_tile_size_h);
             if (output.empty()) {
                 std::vector<int64_t> output_shape = output_tile.shape();
                 output_shape[0]                   = output_width;

src/runtime/tiling.h

@@ -11,8 +11,8 @@ sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
                                    int output_width,
                                    int output_height,
                                    int scale,
-                                   int p_tile_size_x,
-                                   int p_tile_size_y,
+                                   int p_tile_size_w,
+                                   int p_tile_size_h,
                                    float tile_overlap_factor,
                                    bool circular_x,
                                    bool circular_y,

src/tokenizers/t5_unigram_tokenizer.cpp

@@ -295,7 +295,13 @@ bool T5UniGramTokenizer::encode(const std::string& input, std::vector<int>& resu
     std::vector<int32_t> tokens;
     std::vector<std::string> token_strs;
     std::string normalized = normalize(input);
-    auto splited_texts     = split_with_special_tokens(normalized, special_tokens);
+    if (normalized.empty()) {
+        // HF reference tokenizers emit no pieces for empty input; pad_tokens
+        // still appends EOS so the sequence becomes [EOS] + padding.
+        result = std::move(tokens);
+        return true;
+    }
+    auto splited_texts = split_with_special_tokens(normalized, special_tokens);
     if (splited_texts.empty()) {
         splited_texts.push_back(normalized);  // for empty string
     }