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    // `CMF_INIT_LATENT=<path>` takes the starting noise from a file of raw
77    // little-endian f32 instead of drawing it. Comparing this engine's
78    // image against the diffusers reference needs the SAME sample: the two
79    // draw from different generators, and two different draws differ far
80    // more than any arithmetic between them, so a pixel comparison without
81    // this measures the RNGs.
82    if let Ok(path) = std::env::var("CMF_INIT_LATENT") {
83        match std::fs::read(&path) {
84            Ok(b) if b.len() == n * 4 => {
85                return b
86                    .chunks_exact(4)
87                    .map(|c| f32::from_le_bytes(c.try_into().unwrap()))
88                    .collect();
89            }
90            Ok(b) => panic!("{path}: {} floats, the latent needs {n}", b.len() / 4),
91            Err(e) => panic!("{path}: {e}"),
92        }
93    }
94    let mut rng = SplitMix64::new(seed);
95    let mut u = || (rng.next_u64() >> 11) as f64 / (1u64 << 53) as f64;
96    let mut out = Vec::with_capacity(n);
97    while out.len() < n {
98        // Box-Muller; guard log(0).
99        let (a, b) = (u().max(1e-300), u());
100        let r = (-2.0 * a.ln()).sqrt();
101        let ang = 2.0 * std::f64::consts::PI * b;
102        out.push((r * ang.cos()) as f32);
103        if out.len() < n {
104            out.push((r * ang.sin()) as f32);
105        }
106    }
107    out
108}
109
110/// (features, token count) of the conditional prompt plus the same
111/// pair for the "" uncond prompt when CFG runs.
112type CapFeats = (Vec<f32>, usize, Option<(Vec<f32>, usize)>);
113
114/// Prompt → caption features (cond + optional uncond) through Gemma;
115/// the encoder is dropped on return.
116fn encode_prompt(
117    tok: &Tokenizer,
118    enc: &GemmaEncoder,
119    prompt: &str,
120    p: &GenParams,
121    want_uncond: bool,
122) -> CapFeats {
123    let sys = p.system_prompt.as_deref().unwrap_or(DEFAULT_SYSTEM_PROMPT);
124    let full = format!("{sys} <Prompt Start> {prompt}");
125    let mut ids = tok.with_bos(tok.encode(&full));
126    ids.truncate(p.max_tokens);
127    // diffusers takes hidden_states[-2] — the stream entering the
128    // last layer.
129    let (_, streams) = enc.encode(&ids, true);
130    let cap = streams[streams.len() - 2].clone();
131    let cap_u = if want_uncond {
132        let uncond_ids = tok.with_bos(tok.encode(""));
133        let (_, s) = enc.encode(&uncond_ids, true);
134        Some((s[s.len() - 2].clone(), uncond_ids.len()))
135    } else {
136        None
137    };
138    (cap, ids.len(), cap_u)
139}
140
141/// The flow-matching Euler loop over the Next-DiT.
142#[allow(clippy::too_many_arguments)]
143fn denoise(
144    dit: &NextDit,
145    cap: &[f32],
146    cap_n: usize,
147    cap_u: Option<&(Vec<f32>, usize)>,
148    lh: usize,
149    lw: usize,
150    p: &GenParams,
151    progress: &mut dyn FnMut(usize, usize),
152) -> Vec<f32> {
153    let mut latents = gauss_latent(dit.in_channels * lh * lw, p.seed);
154    let sg = sigmas(p.steps, 6.0);
155    // The caption embedding and its refiner blocks depend on the prompt
156    // alone — not on the timestep, not on the latents. Refined once here
157    // instead of inside every model call: at 30 steps under CFG that is
158    // 2 evaluations where the loop used to do 60.
159    let cap_r = dit.refine_caption(cap, cap_n);
160    let cap_u_r = cap_u.map(|(cu, un)| (dit.refine_caption(cu, *un), *un));
161    let mut uncond_slot: Option<Vec<f32>> = None;
162    for i in 0..p.steps {
163        let t = (1.0 - sg[i]) as f32;
164        let cfg_on = cap_u_r.is_some() && (i + 1) as f32 / p.steps as f32 <= p.cfg_trunc_ratio;
165        // Both CFG branches in ONE pass when nothing better is on offer:
166        // the weights are read once for the pair (the whole cost on a
167        // CPU or a phone) and the image branch runs once. A fused
168        // whole-block device path beats it — that one wants a single
169        // sequence — so Metal keeps the two-call shape.
170        let batched = cfg_on && !crate::gpu::fused_dit_block_available();
171        let mut pred = if batched {
172            let (cu, un) = cap_u_r.as_ref().map(|(v, n)| (v.as_slice(), *n)).unwrap();
173            let (pc, pu) = dit.forward_cfg_pair(&latents, lh, lw, &cap_r, cap_n, cu, un, t);
174            uncond_slot = Some(pu);
175            pc
176        } else {
177            dit.forward_with_cap(&latents, lh, lw, &cap_r, cap_n, t)
178        };
179        if let Some((cu, un)) = &cap_u_r {
180            // CFG truncation: past the ratio only cond runs.
181            if cfg_on {
182                let uncond = match uncond_slot.take() {
183                    Some(u) => u,
184                    None => dit.forward_with_cap(&latents, lh, lw, cu, *un, t),
185                };
186                let gs = p.guidance_scale;
187                let mut comb: Vec<f32> = uncond
188                    .iter()
189                    .zip(&pred)
190                    .map(|(&u, &c)| u + gs * (c - u))
191                    .collect();
192                if p.cfg_normalization {
193                    // ‖cond‖/‖comb‖ per row (last dim), as diffusers.
194                    for (cr, gr) in pred.chunks_exact(lw).zip(comb.chunks_exact_mut(lw)) {
195                        let cn = cr.iter().map(|&v| v * v).sum::<f32>().sqrt();
196                        let gn = gr.iter().map(|&v| v * v).sum::<f32>().sqrt();
197                        if gn > 0.0 {
198                            let f = cn / gn;
199                            for v in gr.iter_mut() {
200                                *v *= f;
201                            }
202                        }
203                    }
204                }
205                pred = comb;
206            }
207        }
208        // Lumina predicts toward the image (t=1); the scheduler
209        // steps σ→0, hence the sign flip: x += (σ₊ − σ)·(−pred).
210        let d = (sg[i + 1] - sg[i]) as f32;
211        for (x, &v) in latents.iter_mut().zip(&pred) {
212            *x -= d * v;
213        }
214        progress(i + 1, p.steps);
215    }
216    latents
217}
218
219fn to_rgb01(img: Vec<f32>) -> Vec<f32> {
220    img.iter()
221        .map(|&v| (v / 2.0 + 0.5).clamp(0.0, 1.0))
222        .collect()
223}
224
225/// Generate an RGB image `[3, height, width]` in [0, 1] from `root`:
226/// either a diffusers Lumina-Image 2.0 directory (tokenizer/
227/// text_encoder/ transformer/ vae/, exact f32) or a packaged .cmf
228/// file (`cortiq imagine-pack`, quantized + mmap). `progress` is
229/// called after every denoise step.
230pub fn generate(
231    root: &Path,
232    prompt: &str,
233    p: &GenParams,
234    mut progress: impl FnMut(usize, usize),
235) -> Result<Vec<f32>, String> {
236    if p.height % 16 != 0 || p.width % 16 != 0 {
237        return Err("height/width must be multiples of 16".into());
238    }
239    let (lh, lw) = (p.height / 8, p.width / 8);
240    let want_uncond = p.guidance_scale > 1.0;
241
242    if root.is_file() {
243        // ── packaged .cmf: one mmap, components stay quantized ──
244        let model = std::sync::Arc::new(
245            cortiq_core::CmfModel::open(root).map_err(|e| format!("{}: {e}", root.display()))?,
246        );
247        let vocab = model
248            .vocab
249            .as_deref()
250            .ok_or("packaged .cmf has no embedded tokenizer")?;
251        let tok = Tokenizer::from_bytes(vocab).map_err(|e| format!("tokenizer: {e}"))?;
252        let (cap, cap_n, cap_u) = {
253            let enc = GemmaEncoder::from_cmf(&model)?;
254            encode_prompt(&tok, &enc, prompt, p, want_uncond)
255        };
256        let latents = {
257            let dit = NextDit::from_cmf(&model)?;
258            denoise(&dit, &cap, cap_n, cap_u.as_ref(), lh, lw, p, &mut progress)
259        };
260        let vae = VaeDecoder::from_cmf(&model)?;
261        return Ok(to_rgb01(vae.decode(&latents, lh, lw)));
262    }
263
264    // ── diffusers directory: exact f32, stages load-and-drop ──
265    let tok = Tokenizer::from_file(root.join("tokenizer").join("tokenizer.json"))
266        .map_err(|e| format!("tokenizer: {e}"))?;
267    let (cap, cap_n, cap_u) = {
268        let enc = GemmaEncoder::load_dir(&root.join("text_encoder"))?;
269        encode_prompt(&tok, &enc, prompt, p, want_uncond)
270    };
271    let latents = {
272        let dit = NextDit::load_dir(&root.join("transformer"))?;
273        denoise(&dit, &cap, cap_n, cap_u.as_ref(), lh, lw, p, &mut progress)
274    };
275    let vae = VaeDecoder::load_dir(&root.join("vae"))?;
276    Ok(to_rgb01(vae.decode(&latents, lh, lw)))
277}
cortiq_engine/imagegen.rs

cortiq_engine/
imagegen.rs
