oxibonsai_image/
lib.rs

1//! # oxibonsai-image
2//!
3//! FLUX.2 DiT (`bonsai-image`) GGUF **weight loader and configuration** for the
4//! OxiBonsai text-to-image port.
5//!
6//! This crate (Phase 2 of the OxiBonsai → Bonsai-Image port) provides:
7//!
8//! - [`DitConfig`] — the diffusion-transformer architecture configuration,
9//!   parsed entirely from the `bonsai-image.*` GGUF metadata namespace written
10//!   by the `oxibonsai-model` MLX→GGUF converter.
11//! - [`DitWeights`] — a flat, typed registry over every tensor in a
12//!   `bonsai-image` DiT GGUF file: quantized linears as ternary
13//!   [`oxibonsai_core::quant_ternary::BlockTQ2_0_g128`] blocks (looked up by
14//!   diffusers base name), and plain tensors as BF16 (a borrowed raw-byte
15//!   slice plus on-demand `u16`-bits / `f32` decoders that allocate).
16//!
17//! It does **not** implement the forward pass or build the nested transformer
18//! block hierarchy; that is a later phase, gated on a golden-reference oracle.
19//! The registry here is designed so block structs can be built on top of it.
20//!
21//! ## Storage conventions honoured
22//!
23//! 1. Quantized linear → looked up under its base name (no `.weight`), GGUF
24//!    type `TQ2_0_g128`.
25//! 2. Plain tensor → looked up under its full name, GGUF type `BF16`.
26//! 3. All tensors are stored with their logical shape **reversed**; the
27//!    accessors recover the logical shape (and, for quantized linears, the
28//!    `(out, in)` dimensions).
29//!
30//! ## Example
31//!
32//! ```no_run
33//! use std::path::Path;
34//! use oxibonsai_image::DitWeights;
35//!
36//! let weights = DitWeights::open(Path::new("/path/to/bonsai-image-dit.gguf"))
37//!     .expect("load DiT");
38//! let cfg = weights.config();
39//! assert_eq!(cfg.hidden_size(), cfg.num_attention_heads * cfg.attention_head_dim);
40//!
41//! // Quantized linear by diffusers base name.
42//! let q = weights
43//!     .quantized_linear("transformer_blocks.0.attn.to_q")
44//!     .expect("to_q present");
45//! assert_eq!(q.blocks.len() as u64, q.expected_block_count());
46//!
47//! // Plain BF16 tensor by full name, decoded on demand.
48//! let emb = weights.bf16_tensor("x_embedder.weight").expect("x_embedder present");
49//! let _values: Vec<f32> = emb.to_f32_vec();
50//! ```
51
52pub mod blocks;
53pub mod config;
54/// GPU (CUDA) backend for the DiT ternary matmuls + joint flash-attention. The
55/// `target_os`-disjoint sibling of [`gpu`]; gated on `cfg(all(feature =
56/// "native-cuda", any(target_os = "linux", target_os = "windows")))`. Reuses the
57/// same `OXI_DIT_GPU` / `OXI_DIT_ATTN_GPU` env vars as the Metal path.
58#[cfg(all(
59    feature = "native-cuda",
60    any(target_os = "linux", target_os = "windows")
61))]
62pub mod cuda_gpu;
63pub mod error;
64pub mod forward;
65pub mod gemm;
66#[cfg(all(feature = "metal", target_os = "macos"))]
67pub mod gpu;
68pub mod math;
69/// End-to-end text-to-image orchestration ([`pipeline::text_to_image`]): the
70/// single library entry point shared by the `generate` example and the
71/// `oxibonsai image` CLI subcommand (DiT → VAE → PNG, native or golden-parity).
72pub mod pipeline;
73pub mod png;
74/// FLUX.2 native sampling scaffolding — byte-exact Pure-Rust port of MLX's
75/// Threefry random generator ([`sample::mlx_rng`]) plus the FLUX.2 initial-noise,
76/// position-id, and flow-match schedule generators. Makes the text-to-image
77/// pipeline self-sufficient (no golden `.npy` dump for scaffolding).
78pub mod sample;
79/// Resident multi-prompt session: load the pipeline once, render many prompts.
80pub mod session;
81pub mod te;
82pub mod vae;
83pub mod weights;
84
85pub use config::{DitConfig, ARCHITECTURE, DEFAULT_EPS};
86pub use error::{DitError, DitResult};
87pub use forward::{DitForward, ForwardTaps, QkvNorm, Stage0, StepTap};
88pub use pipeline::{
89    text_to_image, GoldenOverride, PipelineError, TeSource, TextToImageCfg, TextToImageOut,
90};
91pub use png::{encode_rgb8, PngError, PngResult};
92pub use session::{ImageSession, RenderOutcome, RenderParams, StageTimings};
93pub use weights::{Bf16Tensor, DitWeights, QuantizedLinear};
oxibonsai_image/lib.rs

oxibonsai_image/
lib.rs
