1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
//! Pure-Rust FLUX.2 SMALL VAE **decoder** (latent → RGB), validated
//! stage-by-stage against golden MLX tensors.
//!
//! This is the decode-only half of the FLUX.2 (`bonsai-image`) VAE: it takes the
//! DiT's packed latent `[1, 128, 32, 32]`, denormalises it with stored
//! BatchNorm statistics, unpatchifies it to `[1, 32, 64, 64]`, and runs the
//! convolutional decoder (`post_quant_conv → conv_in → mid_block → up_blocks →
//! conv_norm_out → conv_out`) to produce `[1, 3, 512, 512]`. Everything is
//! computed in f32 on flat NCHW (batch-1) buffers.
//!
//! Weights are loaded from the per-tensor `.npy` files exported by
//! `/tmp/bonsai_vae_export_weights.py` via [`weights::VaeWeights`]. The conv
//! weight layout is the MLX `[out, kH, kW, in]` (see [`weights`]).
//!
//! The decoder is **untiled** (single forward pass). Tiling + cosine-blend (for
//! a byte-exact match to the pipeline's auto-tiled output at ≥256px) is
//! deliberately out of scope here.
//!
//! # Example
//!
//! ```no_run
//! use std::path::Path;
//! use oxibonsai_image::vae::{VaeWeights, VaeDecoder};
//!
//! let weights = VaeWeights::open(Path::new("/tmp/bonsai_golden/vae/weights"))
//! .expect("open vae weights");
//! let decoder = VaeDecoder::from_weights(&weights).expect("build decoder");
//! // packed: flat NCHW [1,128,32,32] latent from the DiT.
//! let packed = vec![0.0f32; 128 * 32 * 32];
//! let rgb = decoder
//! .decode_packed_latents(&packed, 32, 32, None)
//! .expect("decode");
//! assert_eq!((rgb.c, rgb.h, rgb.w), (3, 512, 512));
//! ```
/// GPU (CUDA) f32 backend for the VAE decoder convs/norms/silu/upsample. The
/// `target_os`-disjoint sibling of [`gpu`]; gated on `cfg(all(feature =
/// "native-cuda", any(target_os = "linux", target_os = "windows")))`. Reuses the
/// same `OXI_VAE_GPU` env var as the Metal path.
/// GPU (Metal) f32 backend for the VAE decoder convs/norms/silu/upsample. Gated
/// on `cfg(all(feature = "metal", target_os = "macos"))`; opt-in via
/// `OXI_VAE_GPU=1`.
/// Pure-Rust loader for the FLUX.2 `AutoencoderKLFlux2` VAE weights read directly
/// from the canonical diffusers `*.safetensors` (bf16→f32, conv-weight
/// transpose, `to_out.0` un-nesting). The self-serve alternative to the `.npy`
/// golden dump; selected automatically by [`weights::VaeWeights::open`] for a
/// `.safetensors` path.
pub use ;
pub use ;
pub use ;