Module vae 

FLUX-class VAE decoder (diffusers AutoencoderKL, the Lumina-Image 2.0 latent decoder): latents [16, h, w] → RGB [3, 8h, 8w].

First increment of the image-generation runtime (docs/GENERATIVE.ru.md): a plain-Rust NCHW decoder — conv2d, GroupNorm(32), SiLU, spatial self-attention, nearest ×2 upsampling — loaded straight from a diffusers vae/ directory (config.json + safetensors). The CMF packaging comes with the Lumina converter; parity is pinned by a numpy reference (python/vae_ref.py + tests/vae_parity.rs).

Convs run parallel over output channels via scoped threads — naive kernels, good enough to validate the pipeline end-to-end; the im2col

• GEMM path rides later on the existing matmul kernels.

Structs

Conv2d

2-D convolution, stride 1, square kernel, symmetric padding

GroupNorm

GroupNorm over channel groups (eps 1e-6, affine), NCHW in place.

StTensor

One tensor from a .safetensors file, dequantized to f32.

VaeAttnRef

The mid-block single-head spatial attention: GroupNorm, then token-major q/k/v Linear [c, c] (row-major out×in) + bias, softmax(q·kᵀ/√c)·v, out Linear + bias, residual add onto the block input.

VaeChainArgs

The whole decoder as borrowed slices, in execution order: conv_in → mid_res1 → mid_attn → mid_res2 → ups[0..] → norm_out+SiLU → conv_out. Consumed by gpu::vae_decode_chain; key is stable for the life of the VaeDecoder, so a backend caches its uploaded weights by it instead of fingerprinting every conv on every call.

VaeConvRef

One conv of the decoder as the device chain sees it: weight [oc, ic, k, k] row-major f32, bias [oc], stride 1, pad k/2.

VaeDecoder

The full decoder: conv_in → mid(res, attn, res) → up-blocks → GroupNorm + SiLU → conv_out, plus the diffusers latent de-normalization z/scaling_factor + shift_factor.

VaeNormRef

GroupNorm (eps 1e-6, affine) weight/bias [c] over groups groups.

VaeResnetRef

One resnet: x + conv2(silu(norm2(conv1(silu(norm1(x)))))), with the skip through the 1×1 shortcut when in/out channels differ.

VaeUpRef

One up block: its resnets in order, then (if present) nearest-2× upsample followed by upsample conv.

Functions

read_safetensors

All tensors of one .safetensors file, keyed by name.

read_safetensors_each

Stream every tensor of one .safetensors file through f as f32, one at a time — a 9 GB shard costs one raw blob plus the single tensor in flight, not a second full-file f32 copy.