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

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.

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

Requirements

Hardware

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.

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

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.