Mojo24x7/PixArt-Sigma-512-AX650

PixArt-Σ 512 for the AXERA AX650 / AX8850 NPU

Text-to-image on an edge NPU in 27 seconds. PixArt-Σ XL-2 512×512 compiled for the AXERA AX650 family with Pulsar2 7.0, running fully on-device on a Radxa AICore AX-M1 (AX8850) attached to an RK3588 host.

No prior AXERA build of PixArt-Σ existed. Code, conversion scripts and the full write-up: https://github.com/Mojo24x7/pixart-sigma-axm1

stage	where it runs	time
T5-XXL text encoder (int8, ONNX Runtime)	host CPU	~8.5 s
PixArt-Σ DiT, 20 steps, CFG 4.5	AX-M1 NPU	~27 s
SDXL VAE decoder	AX-M1 NPU	~1 s
total, 512×512		~36 s

Output is clean and photorealistic — see the samples below. U16 quantisation does not grain this model, which is not true of every diffusion model on this hardware (see Findings).

Samples

Generated on the AX-M1 itself — 512×512, 20 DPM steps, CFG 4.5, ~36 s each, no post-processing or upscaling.

macro photograph of fresh strawberries with water droplets on dark slate, studio lighting	a close-up photograph of a tabby cat in golden hour sunlight, sharp focus on the eyes, soft bokeh
an F1 race car on a mountain road, snowy peaks, photorealistic	a starry night over a flower meadow, folk-art

Look at the strawberry seed texture and the water droplets: no graining, no cracked texture. The same U16 pipeline visibly wrecks Z-Image-Turbo and SD1.5 on this hardware — see Findings.

In fairness: U16 is not artifact-free. Roughly one image in two carries a few small isolated colour speckles, most often against flat sky or dark backgrounds. They are a world away from the graining that makes Z-Image unusable here, but they are visible if you look for them.

Files

file	size	what
dit.axmodel	720 MB	PixArt-Σ transformer, Pulsar2 7.0, U16, single NPU subgraph
vae.axmodel	71 MB	SDXL VAE decoder, 64×64 → 512×512
t5/t5_enc_int8_ext.onnx + .data	4.77 GB	T5-XXL encoder, int8 dynamic, ONNX Runtime, external data
t5/tokenizer/	—	T5 tokenizer
neg_emb.npy, neg_mask.npy	4.9 MB	pre-encoded empty negative prompt, so only one T5 pass per image is needed

Requirements

Hardware

accelerator	AXERA AX650 / AX8850. Tested on a Radxa AICore AX-M1 (M.2, 8 GB)
host	any AXCL host with an M.2 slot — tested on a Radxa ROCK 5B+ (RK3588), kernel 6.1.84
host RAM	~6 GB free for the int8 T5 encoder
card memory	~0.8 GB of the AX-M1's 7 GB usable CMM for DiT + VAE
disk	~6 GB for the model download

1. AXCL host driver — required, and not a pip package

The accelerator needs Radxa's AXCL host stack installed on the host (docs). On our board:

axclhost 3.10.2 · axclhost-dkms 3.6.5-1 · axclhost-firmware 3.6.5-1 · task-axclhost 3.6.5-1

Verify before going further — if this does not work, nothing else will:

axcl-smi                          # must list the AX8850 and its memory
ls /dev/axcl_host /dev/ax_mmb_dev # both must exist
ls /usr/lib/axcl/                 # the runtime .so files axengine loads

2. Python packages

★ axengine is not on PyPI — pip install axengine fails with No matching distribution found. It is a wheel published on AXERA's GitHub releases, and it is a thin ctypes wrapper that loads /usr/lib/axcl at runtime, which is why step 1 has to be working first.

python3 -V     # 3.11 tested

pip install https://github.com/AXERA-TECH/pyaxengine/releases/download/0.1.3.rc3/axengine-0.1.3-py3-none-any.whl
pip install onnxruntime transformers diffusers torch pillow numpy huggingface_hub

torch is needed only for the scheduler and the RNG; the CPU wheel is fine and no CUDA is involved anywhere.

3. Check the stack before downloading 6 GB

python3 -c "import axengine, onnxruntime, torch, diffusers, transformers, PIL, numpy; print('deps OK')"
python3 -c "from axengine import InferenceSession; print('axengine can reach the card')"

Conversion host — only if you rebuild the models yourself

Not needed to run the published models.

• x86_64 with Docker and Pulsar2 7.0 (from AXERA-TECH/Pulsar2).
• ≥ 64 GB RAM for the fp32 DiT ONNX export.
• Python 3.10/3.11 with torch (CPU), diffusers, transformers, onnx, onnxruntime, onnxsim, numpy, huggingface_hub.
• Budget ~2 h 22 m for the DiT build — see Findings on why.

Usage

# deps first - see Requirements above (axengine comes from a GitHub wheel, not PyPI)

git clone https://github.com/Mojo24x7/pixart-sigma-axm1
hf download Mojo24x7/PixArt-Sigma-512-AX650 --local-dir ./pixart-ax

python3 pixart-sigma-axm1/runtime/pixart_full_ax.py \
    -m ./pixart-ax \
    -p "macro photograph of fresh strawberries with water droplets on dark slate, studio lighting, high detail, photorealistic" \
    -o strawberries.png --seed 7 --neg-cache

[t5]  8.2s
[npu] models loaded 11.1s
  step 20/20  25.83s
[dit] 25.8s  (1.29s/step)
[vae] 1.1s
saved strawberries.png  (seed 7, 20 steps, cfg 4.5)

That command reproduces samples/strawberries.png in this repository exactly — same prompt, same seed, same 20 steps. If your output matches, your setup is correct.

flag	default	what
-m, --model-dir	.	directory with dit.axmodel, vae.axmodel, t5/
-p, --prompt	required	what to generate
-n, --negative	cartoon, blurry, …	negative prompt
-o, --out	out.png	output file
-s, --steps	20	denoise steps — each costs ~1.3 s
-c, --cfg	4.5	guidance scale
--seed	1234	-1 for random
--neg-cache	off	use the shipped neg_emb.npy and skip one T5 pass (−8.5 s)

Steps scale linearly at ~1.3 s each: -s 10 ≈ 13 s of denoise, -s 30 ≈ 39 s. 20 is the quality/speed knee.

runtime/pixart_service.py is a warm HTTP service that loads all three models once — every request after the first is ~36 s with no reload, which is what you want for anything interactive.

Timing, measured

stage	where	time
T5-XXL int8 encode	host CPU, 8 threads	8.5 s per prompt
model load	AX-M1	12 s, once
DiT denoise, 20 steps	AX-M1 NPU	26.1 s — 1.30 s/step, 2 DiT calls per step for CFG
VAE decode 64→512	AX-M1 NPU	1.0 s
total cold / warm		~47 s / ~36 s

How it was converted

Full scripts in the GitHub repository. Three problems had to be solved.

1. Pulsar2 has no Einsum, and the AX NPU rejects 6-D Transpose

PixArt's un-patchify is a single Einsum with equation nhwpqc->nchpwq. Pulsar2 has no Einsum backend. Replacing it with the equivalent 6-D Transpose fails too — AxTranspose on 6 dimensions raises OpBuildException at NPU codegen.

The working substitution keeps everything on the NPU:

Reshape → Transpose(perm=[0,3,1,2]) → DepthToSpace(blocksize=2, mode=DCR)

which is mathematically identical to nhwpqc->nchpwq and uses only supported ops. DepthToSpace was verified to compile on the AX NPU with a standalone op test before being used. Scripts: pixart_unpatch_fix.py, pixart_d2s_surgery.py, d2s_test.py.

2. Affine-less LayerNorm

norm_final has empty Scale/Bias inputs, which the ONNX exporter emits but the converter will not accept. pixart_ln_fix.py injects explicit ones and zeros for all 57 norms (DIM 1152).

3. T5-XXL int8 above the 2 GB protobuf limit

A >2 GB quantised ONNX must be written with use_external_data_format=True. A single 4.77 GB .onnx fails with Protobuf parsing failed; with external data it becomes a 706 KB graph plus a 4.77 GB .data file that ONNX Runtime loads without issue. ~8.5 s per encode on four Cortex-A76 cores, and int8 does not degrade image quality here. Scripts: t5_export_int8.py, t5_requant_ext.py.

Findings

Measured, and not documented elsewhere as far as we can find:

• U16 quantisation is clean for PixArt-Σ. The same U16 pipeline visibly grains both Z-Image-Turbo and SD1.5 on this hardware — cracked skin, speckled texture. PixArt (DiT width 1152, DPM scheduler) quantises cleanly. "AX-M1 U16 grains diffusion models" is model-specific, not universal.
• PixArt-Σ compiles as a single NPU subgraph. Z-Image needed a 34-way split because its 3840-wide DiT overflows the ~11.5 MB on-chip memory on the RoPE op. PixArt's 1152 width stays inside it, so there is no split and no streaming — a large part of why it is 27 s rather than minutes.
• Whole-transformer Pulsar2 builds are pathologically slow. The NPU scheduler tiled this DiT into 725,929 jobs and took 2 h 22 m single-threaded. CNNs of similar size compile in minutes. If iterating, split into subgraphs and compile them in parallel; a whole build is only worth it for the simpler runtime.
• The AX-M1 beats RK3588's own NPU here by ~6.6×. The same DiT converted to RKNN fp16 runs on the RK3588 NPU at ~180 s versus 27 s on the AX-M1.

Credits

• PixArt-Σ — Chen et al., PixArt-alpha. Weights under the original licence.
• AXERA — Pulsar2 toolchain and the axengine runtime.
• Radxa — AICore AX-M1 hardware.
• Conversion, the un-patchify substitution, the int8 T5 pipeline and the measurements in this card: Mojo24x7.

Licence

Model weights follow PixArt-Σ's original licence (OpenRAIL++). The conversion scripts and runtime are MIT — see the GitHub repository.