pictor_image/
session.rs

1//! Resident image-generation session.
2//!
3//! [`ImageSession`] loads the DiT, VAE, text encoder, and tokenizer **once** and
4//! holds them in memory, so a long-running front-end (e.g. the `pictor repl`
5//! command) can render many prompts without re-paying the per-call load and
6//! 4-bit dequant cost. The text encoder is opened in *resident* mode
7//! ([`TeWeights::set_resident`][crate::te::TeWeights::set_resident]): its dequantised f32 weights (~16 GB) stay
8//! cached across renders instead of being thrown away after each forward. That
9//! trades RAM for speed and is intended for high-memory machines.
10//!
11//! The per-render math is identical to the native (non-golden) path of
12//! [`crate::pipeline::text_to_image`] — both share
13//! `decoded_chw_to_rgb8` and the same `sample`/`forward`
14//! stages, so a session render and a one-shot render of the same
15//! `(prompt, seed, steps, size)` produce byte-identical PNGs.
16
17use std::path::Path;
18use std::time::{Duration, Instant};
19
20use crate::pipeline::{
21    decoded_chw_to_rgb8, latent_seq_to_packed_nchw, PipelineError, TeSource, TextToImageOut,
22};
23use crate::png::encode_rgb8;
24use crate::sample;
25use crate::te::{Qwen3Tokenizer, TeWeights, TextEncoder};
26use crate::vae::{VaeDecoder, VaeWeights};
27use crate::{DitForward, DitWeights};
28
29/// The fixed text sequence length the DiT conditioning expects (tokenizer pad).
30const SEQ_TXT: usize = 512;
31
32/// Per-render knobs. Mirrors the relevant fields of
33/// [`crate::pipeline::TextToImageCfg`], minus the model paths (those are fixed
34/// for the life of the session) and golden-parity overrides.
35#[derive(Debug, Clone)]
36pub struct RenderParams {
37    /// The text prompt.
38    pub prompt: String,
39    /// RNG seed for the initial noise.
40    pub seed: u64,
41    /// Number of Euler sampler steps.
42    pub steps: usize,
43    /// Target width in pixels.
44    pub width: usize,
45    /// Target height in pixels.
46    pub height: usize,
47    /// Guidance scale (surfaced for parity; the DiT forward is unconditional).
48    pub guidance: f32,
49}
50
51impl Default for RenderParams {
52    fn default() -> Self {
53        Self {
54            prompt: String::new(),
55            seed: 42,
56            steps: 4,
57            width: 512,
58            height: 512,
59            guidance: 1.0,
60        }
61    }
62}
63
64/// Wall-clock split of a single [`ImageSession::render`], so a front-end can
65/// show where the time went (the same stages the `PICTOR_IMAGE_TIMING` taps print).
66#[derive(Debug, Clone, Copy)]
67pub struct StageTimings {
68    /// Tokenize + text-encoder forward.
69    pub te_encode: Duration,
70    /// DiT flow-matching sampler (all steps).
71    pub dit_sample: Duration,
72    /// VAE decode.
73    pub vae_decode: Duration,
74    /// CHW→RGB + PNG encode.
75    pub png_encode: Duration,
76    /// End-to-end render time.
77    pub total: Duration,
78}
79
80/// The result of one render: the encoded image plus its stage timings.
81pub struct RenderOutcome {
82    /// The encoded PNG and its dimensions.
83    pub image: TextToImageOut,
84    /// Per-stage wall-clock split.
85    pub timings: StageTimings,
86}
87
88/// A loaded, resident text-to-image pipeline. Build once with [`Self::load`],
89/// then call [`Self::render`] per prompt.
90pub struct ImageSession {
91    dit_weights: DitWeights,
92    vae: VaeDecoder,
93    te_weights: TeWeights,
94    tokenizer: Qwen3Tokenizer,
95    in_channels: usize,
96    joint_dim: usize,
97}
98
99impl ImageSession {
100    /// Load every model asset and keep it resident.
101    ///
102    /// `te_source` selects the 4-bit MLX safetensors or the f32 `.npy` dir, same
103    /// as [`crate::pipeline::TextToImageCfg`]. The text encoder is put in
104    /// resident mode so its dequantised weights persist across renders.
105    ///
106    /// # Errors
107    /// [`PipelineError`] wrapping whichever asset failed to load.
108    pub fn load(
109        dit_gguf: &Path,
110        vae_weights: &Path,
111        te_source: &TeSource,
112        tokenizer_dir: &Path,
113    ) -> Result<Self, PipelineError> {
114        let dit_weights = DitWeights::open(dit_gguf)?;
115        let dcfg = dit_weights.config();
116        let in_channels = dcfg.in_channels as usize;
117        let joint_dim = dcfg.joint_attention_dim as usize;
118
119        // VaeDecoder owns its weights, so the loader's buffers can be released.
120        let vae_loaded = VaeWeights::open(vae_weights)?;
121        let vae = VaeDecoder::from_weights(&vae_loaded)?;
122        drop(vae_loaded);
123
124        let te_weights = match te_source {
125            TeSource::Mlx4bit(p) => TeWeights::open_mlx_4bit(p)?,
126            TeSource::NpyDir(d) => TeWeights::open(d)?,
127        };
128        te_weights.set_resident(true);
129
130        let tokenizer = Qwen3Tokenizer::open(tokenizer_dir)?;
131
132        Ok(Self {
133            dit_weights,
134            vae,
135            te_weights,
136            tokenizer,
137            in_channels,
138            joint_dim,
139        })
140    }
141
142    /// Populate the resident text-encoder cache so the *first* real render runs
143    /// at warm speed. A single short encode touches every layer's weights, so
144    /// one pass dequantises the whole encoder. Returns how long warming took.
145    ///
146    /// DiT and VAE weights are memory-mapped and page in lazily; this only warms
147    /// the encoder, which is the dominant amortizable (re-dequant) cost.
148    ///
149    /// # Errors
150    /// [`PipelineError`] if the warm-up encode fails.
151    pub fn warm(&self) -> Result<Duration, PipelineError> {
152        let t = Instant::now();
153        let _ = self.encode("warmup")?;
154        Ok(t.elapsed())
155    }
156
157    /// Tokenize and text-encode `prompt` into the `[SEQ_TXT * joint_dim]`
158    /// conditioning vector.
159    fn encode(&self, prompt: &str) -> Result<Vec<f32>, PipelineError> {
160        let toks = self.tokenizer.tokenize(prompt, SEQ_TXT)?;
161        let encoder = TextEncoder::new(&self.te_weights);
162        let out = encoder.forward(&toks.input_ids, &toks.attention_mask)?;
163        let cond = out.cond_7680()?;
164        let need = SEQ_TXT * self.joint_dim;
165        if cond.len() != need {
166            return Err(PipelineError::Shape(format!(
167                "TE cond len {} != SEQ_TXT*joint_dim ({SEQ_TXT}*{})",
168                cond.len(),
169                self.joint_dim
170            )));
171        }
172        Ok(cond)
173    }
174
175    /// Render one image. Reuses the resident weights; nothing is re-loaded.
176    ///
177    /// # Errors
178    /// [`PipelineError`] wrapping whichever stage failed.
179    pub fn render(&self, params: &RenderParams) -> Result<RenderOutcome, PipelineError> {
180        let t_total = Instant::now();
181
182        let (lat_h, lat_w) = sample::latent_grid(params.height, params.width);
183        let seq_img = lat_h * lat_w;
184
185        // ── 1. Conditioning ──
186        let t_te = Instant::now();
187        let cond_vec = self.encode(&params.prompt)?;
188        let te_encode = t_te.elapsed();
189
190        // ── 2. DiT sampling ──
191        let t_dit = Instant::now();
192        let fwd = DitForward::new(&self.dit_weights);
193        let init = sample::create_noise(params.seed, params.height, params.width);
194        if init.len() != seq_img * self.in_channels {
195            return Err(PipelineError::Shape(format!(
196                "native noise len {} != {seq_img}*{} (in_channels vs PACKED_CHANNELS {})",
197                init.len(),
198                self.in_channels,
199                sample::PACKED_CHANNELS
200            )));
201        }
202        let img_ids = sample::img_ids(lat_h, lat_w);
203        let txt_ids = sample::txt_ids(SEQ_TXT);
204        let (timesteps, sigmas) = sample::flow_match_schedule(seq_img, params.steps);
205        let latent = fwd.sample(
206            &init, &cond_vec, &img_ids, &txt_ids, seq_img, SEQ_TXT, &timesteps, &sigmas, None,
207        )?;
208        let dit_sample = t_dit.elapsed();
209
210        // ── 3. VAE decode ──
211        let t_vae = Instant::now();
212        let packed = latent_seq_to_packed_nchw(&latent, seq_img, self.in_channels, lat_h, lat_w)?;
213        let decoded = self
214            .vae
215            .decode_packed_latents(&packed, lat_h, lat_w, None)?;
216        let vae_decode = t_vae.elapsed();
217
218        // ── 4. Pixels + PNG ──
219        let t_png = Instant::now();
220        let (w, h, rgb) = decoded_chw_to_rgb8(decoded.c, decoded.h, decoded.w, &decoded.data)?;
221        let png = encode_rgb8(w, h, &rgb)?;
222        let png_encode = t_png.elapsed();
223
224        Ok(RenderOutcome {
225            image: TextToImageOut {
226                png,
227                width: w,
228                height: h,
229                stage_cosines: Vec::new(),
230            },
231            timings: StageTimings {
232                te_encode,
233                dit_sample,
234                vae_decode,
235                png_encode,
236                total: t_total.elapsed(),
237            },
238        })
239    }
240}
pictor_image/session.rs

pictor_image/
session.rs
