feat: add PixArt model family support (#2047)

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

@@ -63,6 +63,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.

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

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

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]

src/core/ggml_runner.cpp

@@ -837,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/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/pixart.hpp

@@ -0,0 +1,390 @@
+#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("pos_embed.proj.weight")) {
+                config.hidden_size = w->ne[3];
+                config.in_channels = w->ne[2];
+                config.patch_size  = w->ne[0];
+            }
+            if (auto w = find("proj_out.weight")) {
+                config.out_channels = w->ne[1] / (config.patch_size * config.patch_size);
+            }
+            if (auto w = find("caption_projection.linear_1.weight")) {
+                config.caption_channels = w->ne[0];
+            }
+            if (auto w = find("transformer_blocks.0.attn2.to_k.weight")) {
+                config.cross_attention_dim = w->ne[0];
+            }
+            if (auto w = find("transformer_blocks.0.ff.net.0.proj.weight")) {
+                config.ffn_dim = w->ne[1];
+            }
+            if (find("adaln_single.emb.resolution_embedder.linear_1.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 + ".transformer_blocks.";
+            for (const auto& [name, _] : weights) {
+                if (starts_with(name, block_prefix)) {
+                    layers = std::max(layers, atoi(name.substr(block_prefix.size()).c_str()) + 1);
+                }
+            }
+            if (layers > 0) {
+                config.num_layers = layers;
+                LOG_VERBOSE("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["linear_1"] = std::make_shared<Linear>(in_channels, out_dim);
+            blocks["linear_2"] = std::make_shared<Linear>(out_dim, out_dim);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            x = std::dynamic_pointer_cast<Linear>(blocks["linear_1"])->forward(ctx, x);
+            x = ggml_silu(ctx->ggml_ctx, x);
+            return std::dynamic_pointer_cast<Linear>(blocks["linear_2"])->forward(ctx, x);
+        }
+    };
+
+    class PixArtAttention : public GGMLBlock {
+        int64_t num_heads;
+
+    public:
+        PixArtAttention(int64_t dim, int64_t num_heads, int64_t context_dim)
+            : num_heads(num_heads) {
+            blocks["to_q"]     = std::make_shared<Linear>(dim, dim);
+            blocks["to_k"]     = std::make_shared<Linear>(context_dim, dim);
+            blocks["to_v"]     = std::make_shared<Linear>(context_dim, dim);
+            blocks["to_out.0"] = std::make_shared<Linear>(dim, dim);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* context, ggml_tensor* mask = nullptr) {
+            // x: [N, n_token, dim], context: [N, n_context, context_dim]
+            auto q = std::dynamic_pointer_cast<Linear>(blocks["to_q"])->forward(ctx, x);
+            auto k = std::dynamic_pointer_cast<Linear>(blocks["to_k"])->forward(ctx, context);
+            auto v = std::dynamic_pointer_cast<Linear>(blocks["to_v"])->forward(ctx, context);
+
+            auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled);
+            return std::dynamic_pointer_cast<Linear>(blocks["to_out.0"])->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["attn1"]         = std::make_shared<PixArtAttention>(dim, num_heads, dim);
+            blocks["attn2"]         = std::make_shared<PixArtAttention>(dim, num_heads, context_dim);
+            blocks["ff.net.0.proj"] = std::make_shared<Linear>(dim, ffn_dim);
+            blocks["ff.net.2"]      = 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["attn1"]);
+            auto attn2 = std::dynamic_pointer_cast<PixArtAttention>(blocks["attn2"]);
+            auto proj  = std::dynamic_pointer_cast<Linear>(blocks["ff.net.0.proj"]);
+            auto fc2   = std::dynamic_pointer_cast<Linear>(blocks["ff.net.2"]);
+
+            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 + "scale_shift_table", tensor_storage_map, GGML_TYPE_F32);
+            params["scale_shift_table"] = ggml_new_tensor_2d(ctx, wtype, config.hidden_size, 2);
+        }
+
+    public:
+        PixArtModel() = default;
+        PixArtModel(const PixArtConfig& config)
+            : config(config) {
+            blocks["pos_embed.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["adaln_single.emb.timestep_embedder"] = std::make_shared<PixArtTimestepEmbedding>(ADALN_EMBED_DIM, config.hidden_size);
+            blocks["adaln_single.linear"]                = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size);
+            blocks["caption_projection.linear_1"]        = std::make_shared<Linear>(config.caption_channels, config.hidden_size);
+            blocks["caption_projection.linear_2"]        = std::make_shared<Linear>(config.hidden_size, config.cross_attention_dim);
+            for (int i = 0; i < config.num_layers; ++i) {
+                blocks["transformer_blocks." + std::to_string(i)] =
+                    std::make_shared<PixArtBlock>(config.hidden_size, config.num_heads, config.cross_attention_dim, config.ffn_dim);
+            }
+            blocks["norm_out"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
+            blocks["proj_out"] = 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["pos_embed.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["adaln_single.emb.timestep_embedder"])->forward(ctx, t);
+
+            auto mod = std::dynamic_pointer_cast<Linear>(blocks["adaln_single.linear"])
+                           ->forward(ctx, ggml_silu(ctx->ggml_ctx, emb));  // [N, 6 * hidden]
+
+            auto ctx_emb = std::dynamic_pointer_cast<Linear>(blocks["caption_projection.linear_1"])->forward(ctx, context);
+            ctx_emb      = ggml_ext_gelu(ctx->ggml_ctx, ctx_emb, true);
+            ctx_emb      = std::dynamic_pointer_cast<Linear>(blocks["caption_projection.linear_2"])->forward(ctx, ctx_emb);
+
+            for (int i = 0; i < config.num_layers; ++i) {
+                auto block = std::dynamic_pointer_cast<PixArtBlock>(blocks["transformer_blocks." + std::to_string(i)]);
+                h          = block->forward(ctx, h, mod, ctx_emb, context_mask);
+                sd::ggml_graph_cut::mark_graph_cut(h, "pixart.transformer_blocks." + std::to_string(i), "h");
+            }
+
+            // scale_shift_table + emb -> (shift, scale) for the affine-free final norm
+            auto tail_table = params["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["norm_out"])->forward(ctx, h);
+            h          = modulate(ctx->ggml_ctx, h, parts[0], parts[1]);
+            h          = std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->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/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/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/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,

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;
         }
     }
@@ -527,6 +533,10 @@ SDVersion ModelLoader::get_sd_version() const {
             return VERSION_ERNIE_IMAGE;
         }
         if (tensor_storage.name.find("model.diffusion_model.adaln_single.emb.timestep_embedder.linear_1.bias") != std::string::npos) {
+            // PixArt shares this signature with LTX-AV; pos_embed.proj is PixArt-only.
+            if (tensor_storage_map.find("model.diffusion_model.pos_embed.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 &&
@@ -1591,6 +1601,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/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;
         }

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.adaln_single.emb.resolution_embedder.linear_1.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/tokenizers/t5_unigram_tokenizer.cpp

@@ -295,6 +295,12 @@ 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);
+    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