Crate oxibonsai_image 

oxibonsai-image

FLUX.2 DiT (bonsai-image) GGUF weight loader and configuration for the OxiBonsai text-to-image port.

This crate (Phase 2 of the OxiBonsai → Bonsai-Image port) provides:

• DitConfig — the diffusion-transformer architecture configuration, parsed entirely from the bonsai-image.* GGUF metadata namespace written by the oxibonsai-model MLX→GGUF converter.
• DitWeights — a flat, typed registry over every tensor in a bonsai-image DiT GGUF file: quantized linears as ternary oxibonsai_core::quant_ternary::BlockTQ2_0_g128 blocks (looked up by diffusers base name), and plain tensors as BF16 (a borrowed raw-byte slice plus on-demand u16-bits / f32 decoders that allocate).

It does not implement the forward pass or build the nested transformer block hierarchy; that is a later phase, gated on a golden-reference oracle. The registry here is designed so block structs can be built on top of it.

Storage conventions honoured

1. Quantized linear → looked up under its base name (no .weight), GGUF type TQ2_0_g128.
2. Plain tensor → looked up under its full name, GGUF type BF16.
3. All tensors are stored with their logical shape reversed; the accessors recover the logical shape (and, for quantized linears, the (out, in) dimensions).

Example

use std::path::Path;
use oxibonsai_image::DitWeights;

let weights = DitWeights::open(Path::new("/path/to/bonsai-image-dit.gguf"))
    .expect("load DiT");
let cfg = weights.config();
assert_eq!(cfg.hidden_size(), cfg.num_attention_heads * cfg.attention_head_dim);

// Quantized linear by diffusers base name.
let q = weights
    .quantized_linear("transformer_blocks.0.attn.to_q")
    .expect("to_q present");
assert_eq!(q.blocks.len() as u64, q.expected_block_count());

// Plain BF16 tensor by full name, decoded on demand.
let emb = weights.bf16_tensor("x_embedder.weight").expect("x_embedder present");
let _values: Vec<f32> = emb.to_f32_vec();

Re-exports

pub use config::DitConfig;

pub use config::ARCHITECTURE;

pub use config::DEFAULT_EPS;

pub use error::DitError;

pub use error::DitResult;

pub use forward::DitForward;

pub use forward::ForwardTaps;

pub use forward::QkvNorm;

pub use forward::Stage0;

pub use forward::StepTap;

pub use pipeline::text_to_image;

pub use pipeline::GoldenOverride;

pub use pipeline::PipelineError;

pub use pipeline::TeSource;

pub use pipeline::TextToImageCfg;

pub use pipeline::TextToImageOut;

pub use png::encode_rgb8;

pub use png::PngError;

pub use png::PngResult;

pub use session::ImageSession;

pub use session::RenderOutcome;

pub use session::RenderParams;

pub use session::StageTimings;

pub use weights::Bf16Tensor;

pub use weights::DitWeights;

pub use weights::QuantizedLinear;

Modules

blocks

The two FLUX.2 Klein DiT transformer-block types.

config

FLUX.2 DiT (bonsai-image) configuration parsed from GGUF metadata.

error

Error type for the oxibonsai-image DiT loader.

forward

Top-level FLUX.2 Klein 4B DiT forward pass (Pure Rust, f32).

gemm

Fast f32 GEMM for the DiT forward: out[m, n] = Σ_k in[m, k] * w[n, k].

math

Numeric primitives for the FLUX.2 Klein DiT forward pass.

pipeline

End-to-end text-to-image orchestration (pipeline::text_to_image): the single library entry point shared by the generate example and the oxibonsai image CLI subcommand (DiT → VAE → PNG, native or golden-parity). End-to-end Pure-Rust text-to-image orchestration.

png

A minimal, Pure-Rust 8-bit-RGB PNG encoder.

sample

FLUX.2 native sampling scaffolding — byte-exact Pure-Rust port of MLX’s Threefry random generator (sample::mlx_rng) plus the FLUX.2 initial-noise, position-id, and flow-match schedule generators. Makes the text-to-image pipeline self-sufficient (no golden .npy dump for scaffolding). FLUX.2 native sampling scaffolding: initial noise, position ids, and the flow-match Euler sigma/timestep schedule — all Pure-Rust and byte-exact against the MLX (mflux-prism) FLUX.2 pipeline.

session

Resident multi-prompt session: load the pipeline once, render many prompts. Resident image-generation session.

te

Pure-Rust Qwen3-4B text encoder for Bonsai-Image (FLUX.2 Klein).

vae

Pure-Rust FLUX.2 SMALL VAE decoder (latent → RGB), validated stage-by-stage against golden MLX tensors.

weights

Flat, typed weight registry for a FLUX.2 DiT (bonsai-image) GGUF file.