Struct Flux2VaeDecoder 

pub struct Flux2VaeDecoder { /* private fields */ }

Decoder wrapper for Diffusers AutoencoderKLFlux2 as used by MOSS-Image. Its denoiser emits 128 channels at a 16×-smaller grid; Flux2 first restores the BatchNorm statistics, unpatches each four channel group into a 2×2 spatial tile, optionally applies post_quant_conv, and only then enters the ordinary decoder.

Keeping this wrapper beside VaeDecoder is intentional: Conv2d is private implementation detail, while the staged image pipelines need the exact Flux2 pre-decode contract without duplicating the decoder.

Implementations

impl Flux2VaeDecoder

pub fn from_cmf(model: &CmfModel) -> Result<Self, String>

Load a Flux2 VAE packaged below vae. in a CMF file.

pub fn prewarm(&self) -> bool

Compile/upload the reusable decoder chain while another image stage runs. The small Flux2 post_quant_conv stays on the host before the chain.

pub fn decode_fast( &self, packed: &[f32], h: usize, w: usize, ) -> Result<Vec<f32>, String>

Decode a channels-first Flux2 packed latent [4*C, h, w] into RGB. h/w are the denoiser grid; RGB is 16h × 16w.