cortiq_engine/
imagegen.rs

1//! End-to-end Lumina-Image 2.0 text→image: Gemma-2 prompt encode →
2//! flow-matching Next-DiT loop → FLUX-VAE decode.
3//!
4//! Fourth increment of the image-generation runtime
5//! (docs/GENERATIVE.ru.md). Mirrors diffusers Lumina2Pipeline: the
6//! system-prompt template `"<sys> <Prompt Start> <prompt>"`, Gemma
7//! hidden_states[-2] as caption features, FlowMatchEulerDiscrete with
8//! static shift 6 (σ' = 6σ/(1+5σ) over linspace(1, 1/N, N), terminal
9//! 0), the model called at t = 1−σ, CFG with per-row norm
10//! rescaling and the sign flip before the Euler step. Loads stages
11//! sequentially and drops each when done — peak RSS is one component
12//! (Gemma 10.4 GB f32), not the sum.
13
14use crate::dit::NextDit;
15use crate::sampler::SplitMix64;
16use crate::textenc::GemmaEncoder;
17use crate::tokenizer::Tokenizer;
18use crate::vae::VaeDecoder;
19use std::path::Path;
20
21pub const DEFAULT_SYSTEM_PROMPT: &str = "You are an assistant designed to generate superior \
22     images with the superior degree of image-text alignment based on textual prompts or user \
23     prompts.";
24
25pub struct GenParams {
26    pub height: usize,
27    pub width: usize,
28    pub steps: usize,
29    /// ≤1 disables classifier-free guidance (single forward per step).
30    pub guidance_scale: f32,
31    /// Fraction of steps that run with CFG; past it only cond runs.
32    pub cfg_trunc_ratio: f32,
33    /// Rescale the guided prediction to the cond-branch norm per row.
34    pub cfg_normalization: bool,
35    pub seed: u64,
36    pub system_prompt: Option<String>,
37    /// Gemma prompt token cap (diffusers max_sequence_length).
38    pub max_tokens: usize,
39}
40
41impl Default for GenParams {
42    fn default() -> Self {
43        Self {
44            height: 512,
45            width: 512,
46            steps: 30,
47            guidance_scale: 4.0,
48            cfg_trunc_ratio: 1.0,
49            cfg_normalization: true,
50            seed: 42,
51            system_prompt: None,
52            max_tokens: 256,
53        }
54    }
55}
56
57/// σ schedule: linspace(1, 1/N, N) through the static shift, plus the
58/// terminal 0.
59fn sigmas(steps: usize, shift: f64) -> Vec<f64> {
60    let n = steps;
61    let mut out: Vec<f64> = (0..n)
62        .map(|i| {
63            let s = if n == 1 {
64                1.0
65            } else {
66                1.0 - i as f64 * (1.0 - 1.0 / n as f64) / (n - 1) as f64
67            };
68            shift * s / (1.0 + (shift - 1.0) * s)
69        })
70        .collect();
71    out.push(0.0);
72    out
73}
74
75fn gauss_latent(n: usize, seed: u64) -> Vec<f32> {
76    let mut rng = SplitMix64::new(seed);
77    let mut u = || (rng.next_u64() >> 11) as f64 / (1u64 << 53) as f64;
78    let mut out = Vec::with_capacity(n);
79    while out.len() < n {
80        // Box-Muller; guard log(0).
81        let (a, b) = (u().max(1e-300), u());
82        let r = (-2.0 * a.ln()).sqrt();
83        let ang = 2.0 * std::f64::consts::PI * b;
84        out.push((r * ang.cos()) as f32);
85        if out.len() < n {
86            out.push((r * ang.sin()) as f32);
87        }
88    }
89    out
90}
91
92/// (features, token count) of the conditional prompt plus the same
93/// pair for the "" uncond prompt when CFG runs.
94type CapFeats = (Vec<f32>, usize, Option<(Vec<f32>, usize)>);
95
96/// Prompt → caption features (cond + optional uncond) through Gemma;
97/// the encoder is dropped on return.
98fn encode_prompt(
99    tok: &Tokenizer,
100    enc: &GemmaEncoder,
101    prompt: &str,
102    p: &GenParams,
103    want_uncond: bool,
104) -> CapFeats {
105    let sys = p.system_prompt.as_deref().unwrap_or(DEFAULT_SYSTEM_PROMPT);
106    let full = format!("{sys} <Prompt Start> {prompt}");
107    let mut ids = tok.with_bos(tok.encode(&full));
108    ids.truncate(p.max_tokens);
109    // diffusers takes hidden_states[-2] — the stream entering the
110    // last layer.
111    let (_, streams) = enc.encode(&ids, true);
112    let cap = streams[streams.len() - 2].clone();
113    let cap_u = if want_uncond {
114        let uncond_ids = tok.with_bos(tok.encode(""));
115        let (_, s) = enc.encode(&uncond_ids, true);
116        Some((s[s.len() - 2].clone(), uncond_ids.len()))
117    } else {
118        None
119    };
120    (cap, ids.len(), cap_u)
121}
122
123/// The flow-matching Euler loop over the Next-DiT.
124#[allow(clippy::too_many_arguments)]
125fn denoise(
126    dit: &NextDit,
127    cap: &[f32],
128    cap_n: usize,
129    cap_u: Option<&(Vec<f32>, usize)>,
130    lh: usize,
131    lw: usize,
132    p: &GenParams,
133    progress: &mut dyn FnMut(usize, usize),
134) -> Vec<f32> {
135    let mut latents = gauss_latent(dit.in_channels * lh * lw, p.seed);
136    let sg = sigmas(p.steps, 6.0);
137    // The caption embedding and its refiner blocks depend on the prompt
138    // alone — not on the timestep, not on the latents. Refined once here
139    // instead of inside every model call: at 30 steps under CFG that is
140    // 2 evaluations where the loop used to do 60.
141    let cap_r = dit.refine_caption(cap, cap_n);
142    let cap_u_r = cap_u.map(|(cu, un)| (dit.refine_caption(cu, *un), *un));
143    for i in 0..p.steps {
144        let t = (1.0 - sg[i]) as f32;
145        let mut pred = dit.forward_with_cap(&latents, lh, lw, &cap_r, cap_n, t);
146        if let Some((cu, un)) = &cap_u_r {
147            // CFG truncation: past the ratio only cond runs.
148            if (i + 1) as f32 / p.steps as f32 <= p.cfg_trunc_ratio {
149                let uncond = dit.forward_with_cap(&latents, lh, lw, cu, *un, t);
150                let gs = p.guidance_scale;
151                let mut comb: Vec<f32> = uncond
152                    .iter()
153                    .zip(&pred)
154                    .map(|(&u, &c)| u + gs * (c - u))
155                    .collect();
156                if p.cfg_normalization {
157                    // ‖cond‖/‖comb‖ per row (last dim), as diffusers.
158                    for (cr, gr) in pred.chunks_exact(lw).zip(comb.chunks_exact_mut(lw)) {
159                        let cn = cr.iter().map(|&v| v * v).sum::<f32>().sqrt();
160                        let gn = gr.iter().map(|&v| v * v).sum::<f32>().sqrt();
161                        if gn > 0.0 {
162                            let f = cn / gn;
163                            for v in gr.iter_mut() {
164                                *v *= f;
165                            }
166                        }
167                    }
168                }
169                pred = comb;
170            }
171        }
172        // Lumina predicts toward the image (t=1); the scheduler
173        // steps σ→0, hence the sign flip: x += (σ₊ − σ)·(−pred).
174        let d = (sg[i + 1] - sg[i]) as f32;
175        for (x, &v) in latents.iter_mut().zip(&pred) {
176            *x -= d * v;
177        }
178        progress(i + 1, p.steps);
179    }
180    latents
181}
182
183fn to_rgb01(img: Vec<f32>) -> Vec<f32> {
184    img.iter()
185        .map(|&v| (v / 2.0 + 0.5).clamp(0.0, 1.0))
186        .collect()
187}
188
189/// Generate an RGB image `[3, height, width]` in [0, 1] from `root`:
190/// either a diffusers Lumina-Image 2.0 directory (tokenizer/
191/// text_encoder/ transformer/ vae/, exact f32) or a packaged .cmf
192/// file (`cortiq imagine-pack`, quantized + mmap). `progress` is
193/// called after every denoise step.
194pub fn generate(
195    root: &Path,
196    prompt: &str,
197    p: &GenParams,
198    mut progress: impl FnMut(usize, usize),
199) -> Result<Vec<f32>, String> {
200    if p.height % 16 != 0 || p.width % 16 != 0 {
201        return Err("height/width must be multiples of 16".into());
202    }
203    let (lh, lw) = (p.height / 8, p.width / 8);
204    let want_uncond = p.guidance_scale > 1.0;
205
206    if root.is_file() {
207        // ── packaged .cmf: one mmap, components stay quantized ──
208        let model = std::sync::Arc::new(
209            cortiq_core::CmfModel::open(root).map_err(|e| format!("{}: {e}", root.display()))?,
210        );
211        let vocab = model
212            .vocab
213            .as_deref()
214            .ok_or("packaged .cmf has no embedded tokenizer")?;
215        let tok = Tokenizer::from_bytes(vocab).map_err(|e| format!("tokenizer: {e}"))?;
216        let (cap, cap_n, cap_u) = {
217            let enc = GemmaEncoder::from_cmf(&model)?;
218            encode_prompt(&tok, &enc, prompt, p, want_uncond)
219        };
220        let latents = {
221            let dit = NextDit::from_cmf(&model)?;
222            denoise(&dit, &cap, cap_n, cap_u.as_ref(), lh, lw, p, &mut progress)
223        };
224        let vae = VaeDecoder::from_cmf(&model)?;
225        return Ok(to_rgb01(vae.decode(&latents, lh, lw)));
226    }
227
228    // ── diffusers directory: exact f32, stages load-and-drop ──
229    let tok = Tokenizer::from_file(root.join("tokenizer").join("tokenizer.json"))
230        .map_err(|e| format!("tokenizer: {e}"))?;
231    let (cap, cap_n, cap_u) = {
232        let enc = GemmaEncoder::load_dir(&root.join("text_encoder"))?;
233        encode_prompt(&tok, &enc, prompt, p, want_uncond)
234    };
235    let latents = {
236        let dit = NextDit::load_dir(&root.join("transformer"))?;
237        denoise(&dit, &cap, cap_n, cap_u.as_ref(), lh, lw, p, &mut progress)
238    };
239    let vae = VaeDecoder::load_dir(&root.join("vae"))?;
240    Ok(to_rgb01(vae.decode(&latents, lh, lw)))
241}
cortiq_engine/imagegen.rs

cortiq_engine/
imagegen.rs
