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.

What it does

pictor implements the complete Bonsai-Image text-to-image pipeline in Pure Rust. Given a text prompt, it runs:

prompt ──▶ Text Encoder ──▶ DiT ──▶ VAE decoder ──▶ PNG
           (Qwen3-4B,        (ternary FLUX.2    (AutoencoderKLFlux2)
            4-bit MLX)        transformer,
                              TQ2_0_g128)

Every stage is parity-validated against the MLX golden reference:

Feature Flags

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:

[dependencies]
pictor = { version = "0.2.2", features = ["metal"] }  # or "native-cuda"

Build with the matching GPU feature:

# Apple Silicon
cargo build --release --features metal

# NVIDIA
cargo build --release --features native-cuda

# CPU-only
cargo build --release

Library Usage

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:

# 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):

See 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.

# 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.

Performance

Measured at 512×512, 4 Euler steps, FP32 accumulate throughout (no TF32/FP16-MAC shortcuts, preserving cosine ≥ 0.999 parity):

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 and NOTICE.