# pictor
Pure-Rust text-to-image pipeline: FLUX.2-Klein DiT (TQ2_0_g128 ternary) + AutoencoderKLFlux2 VAE + Qwen3-4B 4-bit text encoder + PNG output, all parity-validated against the MLX reference at cosine ≥ 0.999.
**Version:** 0.2.2
Part of [pictor](https://github.com/vbasky/pictor).
---
## What it does
`pictor` implements the complete Bonsai-Image text-to-image pipeline in
Pure Rust. Given a text prompt, it runs:
```text
prompt ──▶ Text Encoder ──▶ DiT ──▶ VAE decoder ──▶ PNG
(Qwen3-4B, (ternary FLUX.2 (AutoencoderKLFlux2)
4-bit MLX) transformer,
TQ2_0_g128)
```
| Text Encoder | Qwen3-4B, 4-bit | MLX 4-bit `.safetensors` (≈2.1 GB) |
| DiT | FLUX.2-Klein ternary transformer | GGUF, `TQ2_0_g128` |
| VAE decoder | AutoencoderKLFlux2 | FLUX.2 `.safetensors` |
| Tokenizer | Qwen3 BPE | `tokenizer.json` |
| PNG encode | oxiarc-deflate | PNG |
Every stage is parity-validated against the MLX golden reference:
| `te_parity` | Text-encoder output cosine ≥ 0.999 |
| `dit_parity` | DiT forward across all 59 reference taps, each cosine ≥ 0.999 |
| `vae_parity` | VAE decode across all 11 reference taps, each cosine ≥ 0.999 |
| `vae_safetensors_parity` | Native safetensors loader vs `.npy` reference, bit-identical weights |
---
## Feature Flags
| `metal` | Apple Silicon GPU (Metal) | on (when built with `--features metal`) | Enables TQ2_0_g128 DiT matmuls + VAE + flash-attention on macOS via `pictor-kernels`. GPU is **default-on** at runtime once the feature is compiled in. |
| `native-cuda` | NVIDIA GPU (CUDA) | off | Linux/Windows only. Enables cudarc-backed TQ2 GEMM, warp-cooperative flash-attention, and stage-0 GPU context embedding. |
| *(none)* | CPU only | — | Pure-Rust Rayon+NEON fallback on every stage. Always available regardless of feature selection. |
`metal` and `native-cuda` are mutually exclusive by target platform — the Metal
path is `cfg(target_os = "macos")` and the CUDA path is
`cfg(any(target_os = "linux", target_os = "windows"))`.
---
## Quick Start
Add to your `Cargo.toml`:
```toml
[dependencies]
pictor = { version = "0.2.2", features = ["metal"] } # or "native-cuda"
```
Build with the matching GPU feature:
```bash
# Apple Silicon
cargo build --release --features metal
# NVIDIA
cargo build --release --features native-cuda
# CPU-only
cargo build --release
```
### Library Usage
```rust,no_run
use std::path::PathBuf;
use pictor::{
pipeline::{text_to_image, TeSource, TextToImageCfg},
};
let cfg = TextToImageCfg {
prompt: "a tiny bonsai tree in a ceramic pot".to_string(),
seed: 42,
steps: 4,
width: 512,
height: 512,
guidance: 1.0,
dit_gguf: PathBuf::from("./bonsai-dit.gguf"),
vae_weights_dir: PathBuf::from("./bonsai-vae/vae/diffusion_pytorch_model.safetensors"),
te_source: TeSource::Mlx4bit(PathBuf::from("./bonsai-te/text_encoder-mlx-4bit/model.safetensors")),
tokenizer_dir: PathBuf::from("./bonsai-te/text_encoder-mlx-4bit"),
golden_override: None,
};
let out = text_to_image(cfg).expect("pipeline failed");
std::fs::write("bonsai.png", &out.png).expect("write png");
println!("Generated {}×{} PNG", out.width, out.height);
```
`text_to_image` returns a [`TextToImageOut`] containing the PNG byte stream,
image dimensions, and per-stage cosine similarities vs the golden (when
[`GoldenOverride`] is set). All errors are surfaced as [`PipelineError`] — no
`unwrap` / `panic`.
---
## Environment Variables
Asset paths can be set via a `.env` file or shell environment variables (CLI
flags take precedence). Create a `.env` in your working directory:
```dotenv
# DiT GGUF (produced by mlx_image_convert from the ternary safetensors)
PICTOR_DIT_GGUF=./bonsai-dit.gguf
# Text encoder: 4-bit MLX model.safetensors (≈2.1 GB)
PICTOR_TE_4BIT=./bonsai-te/text_encoder-mlx-4bit/model.safetensors
# Tokenizer directory containing tokenizer.json
# (defaults to the TE .safetensors parent when omitted)
PICTOR_TE_TOKENIZER_DIR=./bonsai-te/text_encoder-mlx-4bit
# VAE weights: .safetensors file or legacy .npy directory
PICTOR_VAE_WEIGHTS=./bonsai-vae/vae/diffusion_pytorch_model.safetensors
```
**GPU stage toggles** (default on; set to `"0"` to opt out):
| `PICTOR_DIT_ATTN_GPU` | DiT joint flash-attention (Metal / CUDA) |
| `PICTOR_VAE_GPU` | VAE decode (Metal / CUDA) |
| `PICTOR_TE_GPU` | Text-encoder GEMM (Metal; dormant — set `PICTOR_TE_GPU=1`) |
See [`docs/IMAGEN.md`](../../docs/IMAGEN.md) for the full environment-variable
and flag reference, including the complete asset-acquisition walkthrough.
---
## Asset Acquisition
You need three model assets plus a tokenizer. Downloads use the HuggingFace CLI
(`pip install huggingface_hub`); all conversion and inference are Pure Rust.
```bash
# 1. DiT: download ternary MLX checkpoint and convert to GGUF
# (hf download keeps the repo subfolder, so files land under ./bonsai-dit/transformer-packed-mflux/)
hf download prism-ml/bonsai-image-ternary-4B-mlx-2bit \
transformer-packed-mflux/diffusion_pytorch_model.safetensors --local-dir ./bonsai-dit
cargo run -p pictor-model --example mlx_image_convert --release -- \
./bonsai-dit/transformer-packed-mflux/diffusion_pytorch_model.safetensors ./bonsai-dit.gguf tq2_0_g128
# 2. Text encoder + tokenizer (same repo; no conversion — native 4-bit Rust loader)
hf download prism-ml/bonsai-image-ternary-4B-mlx-2bit \
text_encoder-mlx-4bit/model.safetensors text_encoder-mlx-4bit/tokenizer.json \
--local-dir ./bonsai-te
# 3. VAE (no conversion — native safetensors Rust loader).
# Option A (simplest, non-gated): the VAE bundled in the same PrismML repo
hf download prism-ml/bonsai-image-ternary-4B-mlx-2bit \
vae/diffusion_pytorch_model.safetensors --local-dir ./bonsai-vae
# Option B (canonical, gated — needs huggingface-cli login + license):
# hf download black-forest-labs/FLUX.2-dev \
# vae/diffusion_pytorch_model.safetensors --local-dir ./flux2
```
For the full step-by-step walkthrough see [`docs/IMAGEN.md`](../../docs/IMAGEN.md).
---
## Performance
Measured at 512×512, 4 Euler steps, FP32 accumulate throughout (no TF32/FP16-MAC
shortcuts, preserving cosine ≥ 0.999 parity):
| Apple Silicon (M3-class) | Metal (default-on GPU) | ≈ 52–62 s |
| NVIDIA A4000-class | CUDA | ≈ 31.7 s |
| Any | CPU only (Rayon + NEON) | ≈ 10–15 min |
GPU acceleration is composed of three independently validated kernels: v10 TQ2
ternary GEMM (≈3.8× over v9), joint flash-attention (≈5.47× over CPU, simdgroup
f32 MACs, flash-v2 online softmax), and implicit-GEMM im2col-free conv for the
VAE (≈3.2× over CPU). All three default to on; the CPU fallback is always
available.
---
## Pure Rust Declaration
`pictor-image` is C/C++/Fortran-free
— zero FFI in the default build; every dependency is Pure Rust.
---
## License
Apache-2.0 — derived from oxibonsai (COOLJAPAN OU). See [LICENSE](../../LICENSE) and [NOTICE](../../NOTICE).