Struct VaeDecoder 

pub struct VaeDecoder {
    pub latent_channels: usize,
    pub scaling_factor: f32,
    pub shift_factor: f32,
    /* private fields */
}

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.

Fields

latent_channels: usize

scaling_factor: f32

shift_factor: f32

Implementations

impl VaeDecoder

pub fn chain_args(&self) -> VaeChainArgs<'_>

The borrowed device-chain view (see VaeChainArgs).

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

decode with the resident device chain first: de-normalise on the host, try gpu::vae_decode_chain, and fall back to decode (the unchanged Lumina path) when the backend declines. Same contract as decode: model-scale latents in, RGB [3, 8h, 8w] in ≈[-1, 1] out.

Examples found in repository

examples/zimage_vaecheck.rs (line 39)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = cortiq_core::CmfModel::open(&a[1]).unwrap();
    let o = read_st(&a[2]);
    let reps: usize = a.get(3).and_then(|v| v.parse().ok()).unwrap_or(3);
    let (zs, z) = &o["z"];
    let (_, img) = &o["img"];
    let (h, w) = (zs[zs.len() - 2], zs[zs.len() - 1]);
    let vae = cortiq_engine::vae::VaeDecoder::from_cmf(&model).unwrap();
    for r in 0..reps {
        let t = std::time::Instant::now();
        let got = vae.decode_fast(z, h, w);
        let dt = t.elapsed().as_secs_f64();
        let (mut d, mut n, mut se) = (0f64, 0f64, 0f64);
        for (x, y) in got.iter().zip(img) {
            d += (*x as f64 - *y as f64).powi(2);
            n += (*y as f64).powi(2);
            let q = |v: f32| ((v / 2.0 + 0.5).clamp(0.0, 1.0) * 255.0).round_ties_even() as f64;
            se += (q(*x) - q(*y)).powi(2);
        }
        let psnr = 10.0 * (255f64.powi(2) / (se / got.len() as f64).max(1e-12)).log10();
        println!("rep {r}: {:.3} s   img rel {:.3e}   u8 PSNR {:.2} dB", dt, (d / n).sqrt(), psnr);
    }
}

impl VaeDecoder

pub fn load_dir(dir: &Path) -> Result<Self, String>

Load from a diffusers vae/ directory (config.json + diffusion_pytorch_model.safetensors).

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

Load from a packaged imagegen .cmf (vae.* tensors, stored f16/f32, + vae.config_json).

Examples found in repository

examples/zimage_vaecheck.rs (line 36)

fn main() {
    let a: Vec<String> = std::env::args().collect();
    let model = cortiq_core::CmfModel::open(&a[1]).unwrap();
    let o = read_st(&a[2]);
    let reps: usize = a.get(3).and_then(|v| v.parse().ok()).unwrap_or(3);
    let (zs, z) = &o["z"];
    let (_, img) = &o["img"];
    let (h, w) = (zs[zs.len() - 2], zs[zs.len() - 1]);
    let vae = cortiq_engine::vae::VaeDecoder::from_cmf(&model).unwrap();
    for r in 0..reps {
        let t = std::time::Instant::now();
        let got = vae.decode_fast(z, h, w);
        let dt = t.elapsed().as_secs_f64();
        let (mut d, mut n, mut se) = (0f64, 0f64, 0f64);
        for (x, y) in got.iter().zip(img) {
            d += (*x as f64 - *y as f64).powi(2);
            n += (*y as f64).powi(2);
            let q = |v: f32| ((v / 2.0 + 0.5).clamp(0.0, 1.0) * 255.0).round_ties_even() as f64;
            se += (q(*x) - q(*y)).powi(2);
        }
        let psnr = 10.0 * (255f64.powi(2) / (se / got.len() as f64).max(1e-12)).log10();
        println!("rep {r}: {:.3} s   img rel {:.3e}   u8 PSNR {:.2} dB", dt, (d / n).sqrt(), psnr);
    }
}

pub fn decode(&self, z: &[f32], h: usize, w: usize) -> Vec<f32>

Decode latents [latent_channels, h, w] (model scale, i.e. as produced by the diffusion loop) into RGB [3, 8h, 8w] in [-1, 1].