candle_transformers/models/z_image/
sampling.rs

1//! Sampling utilities for Z-Image model.
2
3use candle::{DType, Device, Result, Tensor};
4
5/// Generate initial Gaussian noise
6///
7/// # Arguments
8/// * `batch_size` - Batch size
9/// * `channels` - Number of channels (typically 16, VAE latent channels)
10/// * `height` - Height (latent space, i.e., image_height / 16)
11/// * `width` - Width (latent space)
12/// * `device` - Compute device
13///
14/// # Returns
15/// Noise tensor of shape (batch_size, channels, height, width)
16pub fn get_noise(
17    batch_size: usize,
18    channels: usize,
19    height: usize,
20    width: usize,
21    device: &Device,
22) -> Result<Tensor> {
23    Tensor::randn(0f32, 1.0, (batch_size, channels, height, width), device)
24}
25
26/// Get linear time schedule with shift
27///
28/// # Arguments
29/// * `num_steps` - Number of inference steps
30/// * `mu` - Time shift parameter (from calculate_shift)
31///
32/// # Returns
33/// Time points from 1.0 to 0.0 (num_steps+1 points)
34pub fn get_schedule(num_steps: usize, mu: f64) -> Vec<f64> {
35    let timesteps: Vec<f64> = (0..=num_steps)
36        .map(|v| v as f64 / num_steps as f64)
37        .rev()
38        .collect();
39
40    // Apply time shift (for Flow Matching)
41    timesteps
42        .into_iter()
43        .map(|t| {
44            if t <= 0.0 || t >= 1.0 {
45                t // boundary case
46            } else {
47                let e = mu.exp();
48                e / (e + (1.0 / t - 1.0))
49            }
50        })
51        .collect()
52}
53
54/// Post-process image from VAE output
55/// Converts from [-1, 1] to [0, 255] u8 image
56pub fn postprocess_image(image: &Tensor) -> Result<Tensor> {
57    let image = image.clamp(-1.0, 1.0)?;
58    let image = ((image + 1.0)? * 127.5)?;
59    image.to_dtype(DType::U8)
60}
61
62/// CFG configuration
63#[derive(Debug, Clone)]
64pub struct CfgConfig {
65    /// Guidance scale (typically 5.0)
66    pub guidance_scale: f64,
67    /// CFG truncation threshold (1.0 = full CFG, 0.0 = no CFG)
68    pub cfg_truncation: f64,
69    /// Whether to normalize CFG output
70    pub cfg_normalization: bool,
71}
72
73impl Default for CfgConfig {
74    fn default() -> Self {
75        Self {
76            guidance_scale: 5.0,
77            cfg_truncation: 1.0,
78            cfg_normalization: false,
79        }
80    }
81}
82
83/// Apply Classifier-Free Guidance
84///
85/// # Arguments
86/// * `pos_pred` - Positive (conditional) prediction
87/// * `neg_pred` - Negative (unconditional) prediction
88/// * `cfg` - CFG configuration
89/// * `t_norm` - Normalized time [0, 1]
90pub fn apply_cfg(
91    pos_pred: &Tensor,
92    neg_pred: &Tensor,
93    cfg: &CfgConfig,
94    t_norm: f64,
95) -> Result<Tensor> {
96    // CFG truncation: disable CFG in late sampling
97    let current_scale = if t_norm > cfg.cfg_truncation {
98        0.0
99    } else {
100        cfg.guidance_scale
101    };
102
103    if current_scale <= 0.0 {
104        return Ok(pos_pred.clone());
105    }
106
107    // CFG formula: pred = pos + scale * (pos - neg)
108    let diff = (pos_pred - neg_pred)?;
109    let pred = (pos_pred + (diff * current_scale)?)?;
110
111    // Optional: CFG normalization (limit output norm)
112    if cfg.cfg_normalization {
113        let ori_norm = pos_pred.sqr()?.sum_all()?.sqrt()?;
114        let new_norm = pred.sqr()?.sum_all()?.sqrt()?;
115        let ori_norm_val = ori_norm.to_scalar::<f32>()?;
116        let new_norm_val = new_norm.to_scalar::<f32>()?;
117
118        if new_norm_val > ori_norm_val {
119            let scale = ori_norm_val / new_norm_val;
120            return pred * scale as f64;
121        }
122    }
123
124    Ok(pred)
125}
126
127/// Scale latents to initial noise level
128///
129/// For flow matching, the initial sample should be pure noise.
130/// This function scales the noise by the initial sigma.
131pub fn scale_noise(noise: &Tensor, sigma: f64) -> Result<Tensor> {
132    noise * sigma
133}
candle_transformers/models/z_image/sampling.rs

candle_transformers/models/z_image/
sampling.rs
