diff --git a/README.md b/README.md index f184c51b..877cf7b9 100644 --- a/README.md +++ b/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) diff --git a/docs/pixart.md b/docs/pixart.md new file mode 100644 index 00000000..f9386abf --- /dev/null +++ b/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=,pixart_interpolation_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. diff --git a/src/model.h b/src/model.h index fcba8620..055feb94 100644 --- a/src/model.h +++ b/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; diff --git a/src/model/diffusion/pixart.hpp b/src/model/diffusion/pixart.hpp new file mode 100644 index 00000000..221800f5 --- /dev/null +++ b/src/model/diffusion/pixart.hpp @@ -0,0 +1,390 @@ +#ifndef __SD_MODEL_DIFFUSION_PIXART_HPP__ +#define __SD_MODEL_DIFFUSION_PIXART_HPP__ + +#include +#include +#include +#include + +#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 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 out(static_cast(gh) * gw * dim); + int64_t quarter = dim / 4; + for (int64_t h = 0; h < gh; ++h) { + float pos_h = static_cast(h) / (static_cast(gh) / base_size) / interpolation_scale; + for (int64_t w = 0; w < gw; ++w) { + float pos_w = static_cast(w) / (static_cast(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(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(in_channels, out_dim); + blocks["linear_2"] = std::make_shared(out_dim, out_dim); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + x = std::dynamic_pointer_cast(blocks["linear_1"])->forward(ctx, x); + x = ggml_silu(ctx->ggml_ctx, x); + return std::dynamic_pointer_cast(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(dim, dim); + blocks["to_k"] = std::make_shared(context_dim, dim); + blocks["to_v"] = std::make_shared(context_dim, dim); + blocks["to_out.0"] = std::make_shared(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(blocks["to_q"])->forward(ctx, x); + auto k = std::dynamic_pointer_cast(blocks["to_k"])->forward(ctx, context); + auto v = std::dynamic_pointer_cast(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(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(dim, num_heads, dim); + blocks["attn2"] = std::make_shared(dim, num_heads, context_dim); + blocks["ff.net.0.proj"] = std::make_shared(dim, ffn_dim); + blocks["ff.net.2"] = std::make_shared(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(blocks["attn1"]); + auto attn2 = std::dynamic_pointer_cast(blocks["attn2"]); + auto proj = std::dynamic_pointer_cast(blocks["ff.net.0.proj"]); + auto fc2 = std::dynamic_pointer_cast(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(config.in_channels, + config.hidden_size, + std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}, + std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}); + blocks["adaln_single.emb.timestep_embedder"] = std::make_shared(ADALN_EMBED_DIM, config.hidden_size); + blocks["adaln_single.linear"] = std::make_shared(config.hidden_size, 6 * config.hidden_size); + blocks["caption_projection.linear_1"] = std::make_shared(config.caption_channels, config.hidden_size); + blocks["caption_projection.linear_2"] = std::make_shared(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(config.hidden_size, config.num_heads, config.cross_attention_dim, config.ffn_dim); + } + blocks["norm_out"] = std::make_shared(config.hidden_size, 1e-6f, false); + blocks["proj_out"] = std::make_shared(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(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(blocks["adaln_single.emb.timestep_embedder"])->forward(ctx, t); + + auto mod = std::dynamic_pointer_cast(blocks["adaln_single.linear"]) + ->forward(ctx, ggml_silu(ctx->ggml_ctx, emb)); // [N, 6 * hidden] + + auto ctx_emb = std::dynamic_pointer_cast(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(blocks["caption_projection.linear_2"])->forward(ctx, ctx_emb); + + for (int i = 0; i < config.num_layers; ++i) { + auto block = std::dynamic_pointer_cast(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(blocks["norm_out"])->forward(ctx, h); + h = modulate(ctx->ggml_ctx, h, parts[0], parts[1]); + h = std::dynamic_pointer_cast(blocks["proj_out"])->forward(ctx, h); // [N, hp*wp, p*p*out_ch] + h = DiT::unpatchify(ctx->ggml_ctx, h, hp, wp, static_cast(p), static_cast(p), false); + return h; // [N, out_channels, H, W] + } + }; + + struct PixArtRunner : public DiffusionModelRunner { + PixArtConfig config; + PixArtModel model; + std::vector pos_vec; + + PixArtRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr 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& tensors, const std::string& prefix) override { + model.get_param_tensors(tensors, prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor, + const sd::Tensor& 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 compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context, + const sd::Tensor& 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 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__ diff --git a/src/model/vae/auto_encoder_kl.hpp b/src/model/vae/auto_encoder_kl.hpp index 115ff345..3a473387 100644 --- a/src/model/vae/auto_encoder_kl.hpp +++ b/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)) { diff --git a/src/model/vae/tae.hpp b/src/model/vae/tae.hpp index e9257b7c..173b40c5 100644 --- a/src/model/vae/tae.hpp +++ b/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) { diff --git a/src/model_loader.cpp b/src/model_loader.cpp index ac97cb7e..aa9ea56b 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -533,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 && @@ -1597,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. diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 7d380fa5..3eba5d71 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/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 struct has_set_runtime_backends : std::false_type {}; template @@ -666,6 +679,10 @@ void StableDiffusionGGML::refresh_compvis_denoiser_sigmas() { std::vector 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; } diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index 9b8819a0..a48959b6 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -1,5 +1,6 @@ #include "model_builders.h" +#include #include #include @@ -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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + true, + 0, + false, + weight_manager, + sd_ctx_params->model_args); + result.diffusion = std::make_shared(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(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; diff --git a/src/tokenizers/t5_unigram_tokenizer.cpp b/src/tokenizers/t5_unigram_tokenizer.cpp index 36526086..d0eb6fe6 100644 --- a/src/tokenizers/t5_unigram_tokenizer.cpp +++ b/src/tokenizers/t5_unigram_tokenizer.cpp @@ -295,6 +295,12 @@ bool T5UniGramTokenizer::encode(const std::string& input, std::vector& resu std::vector tokens; std::vector 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