Skip to main content

memra_engine/
vision_pre.rs

1//! Host preprocessor for vision input (lane/vision): bytes -> ViT patch rows.
2//!
3//! Qwen2VLImageProcessorFast semantics: smart_resize to multiples of
4//! factor = patch(16) * merge(2) = 32 with the pixel-area budget, rescale 1/255,
5//! normalize mean/std 0.5 -> [-1, 1], patchify to [gh*gw, 3*2*16*16 = 1536] rows in
6//! row-major grid order with (c, t, ph, pw) inner order — the flatten of the conv
7//! weight [1152, 3, 2, 16, 16], so `VisionTower::forward` consumes rows directly.
8//! Images duplicate their frame across temporal_patch 2; videos fill the pair with
9//! consecutive sampled frames.
10//!
11//! Resize filter: CatmullRom (Keys bicubic a=-0.5, PIL-BICUBIC family). The HF fast
12//! processor runs torch bicubic (a=-0.75) antialias — close but not bit-equal; the
13//! merger-cosine parity gate arbitrates whether the difference matters.
14
15use crate::vision::{V_MERGE, V_PATCH, V_PATCH_IN, V_TEMPORAL};
16use base64::Engine as _;
17use image::RgbImage;
18use image::imageops::FilterType;
19
20/// Area budget (pixels) from preprocessor_config: shortest_edge / longest_edge.
21pub const MIN_PIXELS: usize = 65536;
22pub const MAX_PIXELS: usize = 16_777_216;
23const FACTOR: usize = V_PATCH * V_MERGE; // 32
24
25pub struct PreppedImage {
26    /// [gh*gw, 1536] row-major grid order.
27    pub patches: Vec<f32>,
28    pub gh: usize,
29    pub gw: usize,
30}
31
32impl PreppedImage {
33    /// Trunk tokens this image occupies (after 2x2 merge).
34    pub fn n_tokens(&self) -> usize {
35        n_tokens_for_grid(self.gh, self.gw)
36    }
37}
38
39/// Trunk tokens a `(gh, gw)` patch grid occupies after the 2x2 merge — the planned twin
40/// of `PreppedImage::n_tokens`, usable from `plan_image_bytes` BEFORE any decode.
41pub fn n_tokens_for_grid(gh: usize, gw: usize) -> usize {
42    gh * gw / (V_MERGE * V_MERGE)
43}
44
45/// HF's smart_resize uses Python round(), which is round-half-EVEN: a side landing
46/// exactly on .5 factors (e.g. 336/32 = 10.5) rounds to the even multiple (320, not
47/// 352). Rust's f64::round is half-away-from-zero and diverged there — caught by the
48/// ornith15 parity gate on a 448x336 probe (grid 22x28 vs HF 20x28).
49fn round_half_even(x: f64) -> f64 {
50    let r = x.round();
51    if (x - x.trunc()).abs() == 0.5 && r % 2.0 != 0.0 {
52        r - x.signum()
53    } else {
54        r
55    }
56}
57
58/// smart_resize (HF): round each side to a multiple of 32 preserving aspect ratio,
59/// then scale into the [MIN_PIXELS, MAX_PIXELS] area budget.
60pub fn smart_resize(h: usize, w: usize) -> Result<(usize, usize), String> {
61    if h < 2 || w < 2 {
62        return Err(format!("image too small: {w}x{h}"));
63    }
64    let ar = h.max(w) as f64 / h.min(w) as f64;
65    if ar > 200.0 {
66        return Err(format!("aspect ratio {ar:.0} exceeds 200"));
67    }
68    let f = FACTOR as f64;
69    let (hf, wf) = (h as f64, w as f64);
70    let mut h_bar = (round_half_even(hf / f) * f).max(f);
71    let mut w_bar = (round_half_even(wf / f) * f).max(f);
72    if h_bar * w_bar > MAX_PIXELS as f64 {
73        let beta = (hf * wf / MAX_PIXELS as f64).sqrt();
74        h_bar = ((hf / beta / f).floor() * f).max(f);
75        w_bar = ((wf / beta / f).floor() * f).max(f);
76    } else if h_bar * w_bar < MIN_PIXELS as f64 {
77        let beta = (MIN_PIXELS as f64 / (hf * wf)).sqrt();
78        h_bar = (hf * beta / f).ceil() * f;
79        w_bar = (wf * beta / f).ceil() * f;
80    }
81    Ok((h_bar as usize, w_bar as usize))
82}
83
84/// Still-image DECODE ceiling (hermes finding, fixed 2026-08-23 — the GIF bomb's
85/// sibling): `load_from_memory` + `to_rgb8` expanded the FULL canvas in host RAM before
86/// smart_resize's MAX_PIXELS check ever ran, so a small crafted file claiming huge
87/// dimensions allocated GBs per request, pre-admission. The budget is now admitted from
88/// the HEADER, before any pixel decodes — same ceiling family as `GIF_MAX_TOTAL_PIXELS`
89/// (67.1M px = 192 MiB retained RGB), 4x the resize budget so every legitimately sized
90/// image still decodes.
91pub const IMG_MAX_DECODE_PIXELS: usize = 1 << 26;
92
93/// Image dimensions from the container HEADER — no pixel decode, no canvas allocation.
94pub fn image_header_dims(bytes: &[u8]) -> Result<(usize, usize), String> {
95    let (w, h) = image::ImageReader::new(std::io::Cursor::new(bytes))
96        .with_guessed_format()
97        .map_err(|e| format!("image container: {e}"))?
98        .into_dimensions()
99        .map_err(|e| format!("image header: {e}"))?;
100    Ok((w as usize, h as usize))
101}
102
103/// PRE-DECODE admission for one still image: header dims -> decode-budget check ->
104/// smart_resize (min-size / aspect-ratio / area budget). Returns the patch grid
105/// `(gh, gw)` the decoded image WILL produce — `n_tokens` and pad runs derive from it,
106/// so budget admission can price a vision request before any canvas expands.
107pub fn plan_image_bytes(bytes: &[u8]) -> Result<(usize, usize), String> {
108    let (w, h) = image_header_dims(bytes)?;
109    if w.saturating_mul(h) > IMG_MAX_DECODE_PIXELS {
110        return Err(format!(
111            "image {w}x{h} exceeds the decode budget ({IMG_MAX_DECODE_PIXELS} px) — \
112             refused before decode"
113        ));
114    }
115    let (rh, rw) = smart_resize(h, w)?;
116    Ok((rh / V_PATCH, rw / V_PATCH))
117}
118
119/// Decode + resize one image to its target grid. Returns the resized RGB frame and
120/// the patch grid (gh, gw) in 16px patches (both even — factor 32 guarantees it).
121/// Admission runs FIRST (`plan_image_bytes`, header-only), and the decoder itself is
122/// capped to the admitted dimensions so a header lying small cannot expand past them.
123fn decode_frame(bytes: &[u8]) -> Result<(RgbImage, usize, usize), String> {
124    plan_image_bytes(bytes)?;
125    let (hw, hh) = image_header_dims(bytes)?;
126    let mut reader = image::ImageReader::new(std::io::Cursor::new(bytes))
127        .with_guessed_format()
128        .map_err(|e| format!("image container: {e}"))?;
129    let mut limits = image::Limits::default();
130    limits.max_image_width = Some(hw as u32);
131    limits.max_image_height = Some(hh as u32);
132    reader.limits(limits);
133    let img = reader.decode().map_err(|e| format!("image decode: {e}"))?;
134    let rgb = img.to_rgb8();
135    let (w, h) = (rgb.width() as usize, rgb.height() as usize);
136    let (rh, rw) = smart_resize(h, w)?;
137    let resized = image::imageops::resize(&rgb, rw as u32, rh as u32, FilterType::CatmullRom);
138    Ok((resized, rh / V_PATCH, rw / V_PATCH))
139}
140
141/// Fill patch rows for one temporal slot `t` from a frame. Rows are row-major over
142/// the (gh, gw) grid; inner order (c, t, ph, pw).
143fn fill_slot(rows: &mut [f32], frame: &RgbImage, gh: usize, gw: usize, t: usize) {
144    let inv = 1.0f32 / 127.5;
145    for py in 0..gh {
146        for px in 0..gw {
147            let row = &mut rows[(py * gw + px) * V_PATCH_IN..(py * gw + px + 1) * V_PATCH_IN];
148            for c in 0..3 {
149                let base = c * V_TEMPORAL * V_PATCH * V_PATCH + t * V_PATCH * V_PATCH;
150                for ph in 0..V_PATCH {
151                    for pw in 0..V_PATCH {
152                        let p =
153                            frame.get_pixel((px * V_PATCH + pw) as u32, (py * V_PATCH + ph) as u32);
154                        row[base + ph * V_PATCH + pw] = p.0[c] as f32 * inv - 1.0;
155                    }
156                }
157            }
158        }
159    }
160}
161
162/// Image bytes (png/jpeg/webp/gif/bmp) -> patch rows. The single frame fills both
163/// temporal slots (HF: images are tiled to temporal_patch_size).
164pub fn prep_image_bytes(bytes: &[u8]) -> Result<PreppedImage, String> {
165    let (frame, gh, gw) = decode_frame(bytes)?;
166    let mut patches = vec![0f32; gh * gw * V_PATCH_IN];
167    for t in 0..V_TEMPORAL {
168        fill_slot(&mut patches, &frame, gh, gw, t);
169    }
170    Ok(PreppedImage { patches, gh, gw })
171}
172
173/// `data:image/...;base64,<payload>` -> patch rows.
174pub fn prep_data_uri(uri: &str) -> Result<PreppedImage, String> {
175    let bytes = decode_data_uri(uri)?;
176    prep_image_bytes(&bytes)
177}
178
179/// One pad-run unit crossing the API boundary: a standalone image, or one temporal
180/// group of a video. Units with the same `video` index are consecutive and forward
181/// TOGETHER through `forward_seq` (one attention span per video).
182pub struct VisionUnit {
183    pub prep: PreppedImage,
184    /// Some(video_idx) for video groups; None for standalone images.
185    pub video: Option<usize>,
186}
187
188/// One prepared VIDEO: temporal groups as PreppedImage units (each = one pad run of
189/// `gh*gw/4` tokens) + per-group timestamps for the HF placeholder format
190/// (`<t.t seconds>` before each group's pad run). Groups forward TOGETHER through
191/// `VisionTower::forward_seq` — one attention span per video, the HF cu_seqlens law.
192pub struct PreppedVideo {
193    pub groups: Vec<PreppedImage>,
194    pub timestamps: Vec<f32>,
195}
196
197/// Serving cap on total video patches (groups*gh*gw): sdpa_naive keys the whole span in
198/// shared memory, so the pixel budget stays well under the HF default. Env-tunable.
199pub fn video_max_pixels() -> usize {
200    std::env::var("MEMRA_VIDEO_MAX_PIXELS")
201        .ok()
202        .and_then(|v| v.parse().ok())
203        .unwrap_or(2_097_152)
204}
205pub const VID_MIN_PIXELS: usize = 4096;
206/// Sampled frame cap (2 frames per temporal group).
207pub const VID_MAX_FRAMES: usize = 32;
208
209/// Decode ceilings for `prep_video_gif` (hermes finding, fixed 2026-08-19): the loop used
210/// to expand EVERY frame to full-canvas RGB in host RAM before the `VID_MAX_FRAMES`
211/// sample — and it runs in the HTTP handler pre-admission, so a small crafted GIF (big
212/// canvas x many frames; LZW expands ~1000x) allocated GBs per request. The canvas
213/// dimensions come from the GIF header, so `frames x canvas pixels` is checked against
214/// the pixel ceiling AS DECODE PROCEEDS and the request is refused (clean 4xx at the
215/// handler) the moment the budget would cross — retained RAM is bounded by
216/// `GIF_MAX_TOTAL_PIXELS` RGB (192 MiB) plus at most one transient canvas, no matter
217/// what the stream claims. 512 frames / 67.1M px comfortably cover legitimate clips
218/// (a 480p GIF may run ~370 frames, ~12 s at 30 fps) — the serve path samples down to
219/// 32 frames and ~2M px right after this anyway.
220pub const GIF_MAX_FRAMES: usize = 512;
221pub const GIF_MAX_TOTAL_PIXELS: usize = 1 << 26; // 67.1M px = 192 MiB retained RGB
222
223/// Header-only video plan used by the HTTP admission path. It carries exactly the information
224/// needed to price/render the pad runs; frame pixels are not materialized until the request has
225/// passed budget and concurrency admission.
226#[derive(Debug, Clone)]
227pub struct PlannedVideoGroup {
228    pub gh: usize,
229    pub gw: usize,
230    pub timestamp: f32,
231}
232
233#[derive(Debug, Clone)]
234pub struct PlannedVideo {
235    pub groups: Vec<PlannedVideoGroup>,
236}
237
238fn gif_need(bytes: &[u8], pos: usize, len: usize, what: &str) -> Result<(), String> {
239    if pos.checked_add(len).is_none_or(|end| end > bytes.len()) {
240        return Err(format!("truncated GIF {what}"));
241    }
242    Ok(())
243}
244
245fn gif_skip_subblocks(bytes: &[u8], pos: &mut usize) -> Result<(), String> {
246    loop {
247        gif_need(bytes, *pos, 1, "sub-block length")?;
248        let len = bytes[*pos] as usize;
249        *pos += 1;
250        if len == 0 {
251            return Ok(());
252        }
253        gif_need(bytes, *pos, len, "sub-block payload")?;
254        *pos += len;
255    }
256}
257
258/// Parse GIF structure, frame count, delays, and canvas dimensions without decoding a single
259/// pixel. This mirrors the sampling arithmetic in `prep_video_gif`, but keeps the expensive
260/// composited frame buffers behind the HTTP admission gates.
261pub fn plan_video_gif(bytes: &[u8]) -> Result<PlannedVideo, String> {
262    if bytes.len() < 13 || !(bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a")) {
263        return Err("invalid GIF header".into());
264    }
265    let cw = u16::from_le_bytes([bytes[6], bytes[7]]) as usize;
266    let ch = u16::from_le_bytes([bytes[8], bytes[9]]) as usize;
267    let canvas_px = cw
268        .checked_mul(ch)
269        .ok_or_else(|| "gif canvas dimensions overflow".to_string())?;
270    if canvas_px == 0 {
271        return Err("gif has an empty canvas".into());
272    }
273    let max_frames = GIF_MAX_FRAMES.min(GIF_MAX_TOTAL_PIXELS / canvas_px);
274    if max_frames == 0 {
275        return Err(format!(
276            "gif canvas {cw}x{ch} exceeds the decode budget ({GIF_MAX_TOTAL_PIXELS} px)"
277        ));
278    }
279
280    let mut pos = 13usize;
281    let packed = bytes[10];
282    if packed & 0x80 != 0 {
283        let table_len = 3usize
284            .checked_mul(1usize << ((packed & 0x07) as usize + 1))
285            .ok_or_else(|| "GIF color table length overflow".to_string())?;
286        gif_need(bytes, pos, table_len, "global color table")?;
287        pos += table_len;
288    }
289
290    let mut timestamps = Vec::new();
291    let mut elapsed = 0f32;
292    let mut next_delay = 0.01f32;
293    let mut trailer_seen = false;
294    while pos < bytes.len() {
295        match bytes[pos] {
296            0x3B => {
297                trailer_seen = true;
298                break;
299            }
300            0x21 => {
301                gif_need(bytes, pos, 2, "extension label")?;
302                let label = bytes[pos + 1];
303                pos += 2;
304                if label == 0xF9 {
305                    gif_need(bytes, pos, 1, "graphic-control block size")?;
306                    let block_len = bytes[pos] as usize;
307                    pos += 1;
308                    if block_len != 4 {
309                        return Err(format!(
310                            "unsupported GIF graphic-control block length {block_len}"
311                        ));
312                    }
313                    gif_need(bytes, pos, block_len + 1, "graphic-control block")?;
314                    let delay_cs = u16::from_le_bytes([bytes[pos + 1], bytes[pos + 2]]);
315                    // image::Delay exposes the same centiseconds as milliseconds
316                    // (delay_cs * 10) before dividing by 1000. Mirror that exact
317                    // operation order so the metadata-only planner and decoder
318                    // accumulate identical f32 timestamps over long GIFs.
319                    next_delay = ((delay_cs as f32 * 10.0) / 1000.0).max(0.01);
320                    pos += block_len;
321                    if bytes[pos] != 0 {
322                        return Err("GIF graphic-control block is not terminated".into());
323                    }
324                    pos += 1;
325                } else {
326                    gif_skip_subblocks(bytes, &mut pos)?;
327                }
328            }
329            0x2C => {
330                gif_need(bytes, pos, 10, "image descriptor")?;
331                let fw = u16::from_le_bytes([bytes[pos + 5], bytes[pos + 6]]) as usize;
332                let fh = u16::from_le_bytes([bytes[pos + 7], bytes[pos + 8]]) as usize;
333                if fw == 0 || fh == 0 {
334                    return Err("GIF frame has an empty rectangle".into());
335                }
336                if timestamps.len() >= max_frames {
337                    return Err(format!(
338                        "gif exceeds the decode budget: more than {max_frames} frames at {cw}x{ch} \
339                         (ceiling {GIF_MAX_FRAMES} frames / {GIF_MAX_TOTAL_PIXELS} total px)"
340                    ));
341                }
342                let frame_packed = bytes[pos + 9];
343                pos += 10;
344                if frame_packed & 0x80 != 0 {
345                    let table_len = 3usize
346                        .checked_mul(1usize << ((frame_packed & 0x07) as usize + 1))
347                        .ok_or_else(|| "GIF local color table length overflow".to_string())?;
348                    gif_need(bytes, pos, table_len, "local color table")?;
349                    pos += table_len;
350                }
351                gif_need(bytes, pos, 1, "LZW minimum code size")?;
352                pos += 1;
353                gif_skip_subblocks(bytes, &mut pos)?;
354                timestamps.push(elapsed);
355                elapsed += next_delay;
356                next_delay = 0.01;
357            }
358            other => return Err(format!("unsupported GIF block 0x{other:02x}")),
359        }
360    }
361    if !trailer_seen {
362        return Err("GIF is missing its trailer".into());
363    }
364    if timestamps.is_empty() {
365        return Err("gif has no frames".into());
366    }
367    if timestamps.len() == 1 {
368        timestamps.push(timestamps[0]);
369    }
370    let total = timestamps.len();
371    let take = total.min(VID_MAX_FRAMES) & !1;
372    let picked: Vec<usize> = (0..take).map(|i| i * total / take).collect();
373    let (rh, rw) = smart_resize_video(take, ch, cw)?;
374    let (gh, gw) = (rh / V_PATCH, rw / V_PATCH);
375    let groups = (0..take / 2)
376        .map(|g| PlannedVideoGroup {
377            gh,
378            gw,
379            timestamp: timestamps[picked[2 * g]],
380        })
381        .collect();
382    Ok(PlannedVideo { groups })
383}
384
385/// HF Qwen3VL video smart_resize: the pixel budget covers t_bar*h*w — ALL frames.
386fn smart_resize_video(frames: usize, h: usize, w: usize) -> Result<(usize, usize), String> {
387    if h < 2 || w < 2 {
388        return Err(format!("frame too small: {w}x{h}"));
389    }
390    let ar = h.max(w) as f64 / h.min(w) as f64;
391    if ar > 200.0 {
392        return Err(format!("aspect ratio {ar:.0} exceeds 200"));
393    }
394    let f = FACTOR as f64;
395    let (hf, wf) = (h as f64, w as f64);
396    let t_bar = ((frames as f64 / V_TEMPORAL as f64).round() * V_TEMPORAL as f64).max(2.0);
397    let mut h_bar = (round_half_even(hf / f) * f).max(f);
398    let mut w_bar = (round_half_even(wf / f) * f).max(f);
399    let (min_px, max_px) = (VID_MIN_PIXELS as f64, video_max_pixels() as f64);
400    if t_bar * h_bar * w_bar > max_px {
401        let beta = (frames as f64 * hf * wf / max_px).sqrt();
402        h_bar = ((hf / beta / f).floor() * f).max(f);
403        w_bar = ((wf / beta / f).floor() * f).max(f);
404    } else if t_bar * h_bar * w_bar < min_px {
405        let beta = (min_px / (frames as f64 * hf * wf)).sqrt();
406        h_bar = (hf * beta / f).ceil() * f;
407        w_bar = (wf * beta / f).ceil() * f;
408    }
409    Ok((h_bar as usize, w_bar as usize))
410}
411
412/// Animated GIF -> prepared video: decode frames + delays, uniform-sample to an even
413/// count <= VID_MAX_FRAMES, resize on the total-pixel budget, patchify CONSECUTIVE
414/// frame pairs into temporal groups (frame 2g fills t=0, 2g+1 fills t=1). Timestamps
415/// come from the GIF's own delays at the sampled indices (HF `_calculate_timestamps`).
416pub fn prep_video_gif(bytes: &[u8]) -> Result<PreppedVideo, String> {
417    use image::AnimationDecoder;
418    use image::ImageDecoder as _;
419    let dec = image::codecs::gif::GifDecoder::new(std::io::Cursor::new(bytes))
420        .map_err(|e| format!("gif decode: {e}"))?;
421    // Decode budget from the HEADER, before any frame expands: every decoded frame
422    // composites to the full canvas, so canvas pixels bound the per-frame cost and
423    // frames x canvas is checked against the ceiling as decode proceeds (at most one
424    // transient frame past the cap ever exists). Over-limit refuses cleanly — the
425    // handler surfaces it as a 4xx — instead of expanding the whole stream in host RAM.
426    let (cw, ch) = dec.dimensions();
427    let canvas_px = (cw as usize) * (ch as usize);
428    if canvas_px == 0 {
429        return Err("gif has an empty canvas".into());
430    }
431    let max_frames = GIF_MAX_FRAMES.min(GIF_MAX_TOTAL_PIXELS / canvas_px);
432    if max_frames == 0 {
433        return Err(format!(
434            "gif canvas {cw}x{ch} exceeds the decode budget ({GIF_MAX_TOTAL_PIXELS} px)"
435        ));
436    }
437    let mut frames: Vec<(RgbImage, f32)> = Vec::new(); // (frame, start_seconds)
438    let mut t = 0f32;
439    for fr in dec.into_frames() {
440        if frames.len() >= max_frames {
441            return Err(format!(
442                "gif exceeds the decode budget: more than {max_frames} frames at {cw}x{ch} \
443                 (ceiling {GIF_MAX_FRAMES} frames / {GIF_MAX_TOTAL_PIXELS} total px)"
444            ));
445        }
446        let fr = fr.map_err(|e| format!("gif frame: {e}"))?;
447        let (num, den) = fr.delay().numer_denom_ms();
448        let dt = if den == 0 {
449            100.0
450        } else {
451            num as f32 / den as f32
452        } / 1000.0;
453        frames.push((
454            image::DynamicImage::ImageRgba8(fr.into_buffer()).to_rgb8(),
455            t,
456        ));
457        t += dt.max(0.01);
458    }
459    if frames.is_empty() {
460        return Err("gif has no frames".into());
461    }
462    // still gif: duplicate the frame so one temporal group forms
463    if frames.len() == 1 {
464        let f0 = frames[0].clone();
465        frames.push((f0.0, f0.1));
466    }
467    // uniform sample to an even count <= VID_MAX_FRAMES
468    let total = frames.len();
469    let take = total.min(VID_MAX_FRAMES) & !1;
470    let picked: Vec<usize> = (0..take)
471        .map(|i| i * total / take) // floor spacing, strictly increasing for take <= total
472        .collect();
473    let (h, w) = (frames[0].0.height() as usize, frames[0].0.width() as usize);
474    let (rh, rw) = smart_resize_video(take, h, w)?;
475    let (gh, gw) = (rh / V_PATCH, rw / V_PATCH);
476    let mut groups = Vec::with_capacity(take / 2);
477    let mut timestamps = Vec::with_capacity(take / 2);
478    for g in 0..take / 2 {
479        let (a, b) = (picked[2 * g], picked[2 * g + 1]);
480        let mut patches = vec![0f32; gh * gw * V_PATCH_IN];
481        for (slot, idx) in [(0usize, a), (1usize, b)] {
482            let resized = image::imageops::resize(
483                &frames[idx].0,
484                rw as u32,
485                rh as u32,
486                FilterType::CatmullRom,
487            );
488            fill_slot(&mut patches, &resized, gh, gw, slot);
489        }
490        groups.push(PreppedImage { patches, gh, gw });
491        timestamps.push(frames[a].1);
492    }
493    Ok(PreppedVideo { groups, timestamps })
494}
495
496/// Per-image RAW byte cap, enforced AT the data-URI decode (hermes review finding
497/// 48f96cb4cd37e436): the decode runs in the HTTP content walkers BEFORE slot admission,
498/// and until this cap only `MAX_BODY_BYTES` (192 MiB) bounded it, so one oversized image
499/// could expand ~144 MiB of host bytes pre-check. 12 MiB raw is the per-image line item
500/// the memra-server `MAX_BODY_BYTES` budget has always itemized (8 images x 12 MiB raw
501/// x 4/3 base64); this constant makes that line item enforced rather than aspirational.
502pub const IMG_MAX_RAW_BYTES: usize = 12 * 1024 * 1024;
503
504/// The exact base64 length ceiling for `IMG_MAX_RAW_BYTES` raw bytes: base64 emits 4
505/// chars per 3 raw bytes and the cap is a multiple of 3, so any longer payload decodes
506/// past the cap. Checking the payload LENGTH refuses oversized input before ANY decode
507/// allocation. Shared by the qwen and gemma data-URI decoders.
508pub(crate) fn data_uri_payload_over_cap(payload: &str) -> Option<String> {
509    if payload.len() > IMG_MAX_RAW_BYTES / 3 * 4 {
510        return Some(format!(
511            "image data exceeds {} MiB (per-image raw limit, refused before decode)",
512            IMG_MAX_RAW_BYTES / (1024 * 1024)
513        ));
514    }
515    None
516}
517
518/// Parse a base64 data URI into raw bytes (any `data:*;base64,` media type). Refuses a
519/// payload past `IMG_MAX_RAW_BYTES` (by encoded length, before decoding anything).
520pub fn decode_data_uri(uri: &str) -> Result<Vec<u8>, String> {
521    let rest = uri
522        .strip_prefix("data:")
523        .ok_or_else(|| "expected data: URI (http fetch requires MEMRA_FETCH_URLS=1)".to_string())?;
524    let (meta, payload) = rest
525        .split_once(',')
526        .ok_or_else(|| "malformed data URI: no comma".to_string())?;
527    if !meta.ends_with(";base64") {
528        return Err("data URI must be base64-encoded".into());
529    }
530    let payload = payload.trim();
531    if let Some(err) = data_uri_payload_over_cap(payload) {
532        return Err(err);
533    }
534    base64::engine::general_purpose::STANDARD
535        .decode(payload)
536        .map_err(|e| format!("base64 decode: {e}"))
537}
538
539#[cfg(test)]
540mod tests {
541    use super::*;
542
543    #[test]
544    fn smart_resize_multiples_and_budget() {
545        // typical photo
546        let (h, w) = smart_resize(1080, 1920).unwrap();
547        assert_eq!(h % 32, 0);
548        assert_eq!(w % 32, 0);
549        assert!(h * w >= MIN_PIXELS && h * w <= MAX_PIXELS);
550        // tiny icon scales UP to the floor
551        let (h, w) = smart_resize(64, 64).unwrap();
552        assert!(h * w >= MIN_PIXELS);
553        // huge pano scales DOWN under the cap
554        let (h, w) = smart_resize(8000, 12000).unwrap();
555        assert!(h * w <= MAX_PIXELS);
556        assert!(smart_resize(10, 4000).is_err()); // ar > 200
557    }
558
559    /// Minimal BMP whose HEADER claims `w x h` — the pixel payload is absent, so any
560    /// path that survives past the header check would fail loudly at decode, and any
561    /// path that ALLOCATES the claimed canvas before checking would try to expand
562    /// w*h*3 bytes. The tooth's decode-bomb stand-in.
563    fn bmp_header_claiming(w: u32, h: u32) -> Vec<u8> {
564        let mut b = Vec::new();
565        b.extend_from_slice(b"BM"); // signature
566        b.extend_from_slice(&54u32.to_le_bytes()); // file size (lie, irrelevant)
567        b.extend_from_slice(&0u32.to_le_bytes()); // reserved
568        b.extend_from_slice(&54u32.to_le_bytes()); // pixel data offset
569        b.extend_from_slice(&40u32.to_le_bytes()); // BITMAPINFOHEADER size
570        b.extend_from_slice(&(w as i32).to_le_bytes());
571        b.extend_from_slice(&(h as i32).to_le_bytes());
572        b.extend_from_slice(&1u16.to_le_bytes()); // planes
573        b.extend_from_slice(&24u16.to_le_bytes()); // bpp
574        b.extend_from_slice(&[0u8; 24]); // compression..colors_important
575        b
576    }
577
578    #[test]
579    fn decode_bomb_refuses_pre_decode() {
580        // TOOTH (hermes findings: still-image decode bomb + full-canvas expansion
581        // before the pixel budget; fixed 2026-08-23): a tiny request whose header
582        // claims a 768-megapixel canvas must refuse at ADMISSION — named decode-budget
583        // error from the header dims, before load/to_rgb8 can expand anything.
584        let bomb = bmp_header_claiming(16_000, 16_000);
585        let err = plan_image_bytes(&bomb).unwrap_err();
586        assert!(
587            err.contains("exceeds the decode budget"),
588            "want the named pre-decode refusal, got: {err}"
589        );
590        // The full prep path refuses with the same admission error (it must not reach
591        // the decoder at all — an absent pixel payload would produce a decode error
592        // instead, which would mean the canvas was attempted).
593        let err = match prep_image_bytes(&bomb) {
594            Ok(_) => panic!("bomb must not prep"),
595            Err(e) => e,
596        };
597        assert!(
598            err.contains("exceeds the decode budget"),
599            "prep must refuse at admission, not at decode: {err}"
600        );
601        // Header dims really are read without pixel decode.
602        assert_eq!(image_header_dims(&bomb).unwrap(), (16_000, 16_000));
603        // Positive control: an in-budget image plans to the same grid decode produces.
604        let img = RgbImage::new(64, 64);
605        let mut buf = std::io::Cursor::new(Vec::new());
606        img.write_to(&mut buf, image::ImageFormat::Png).unwrap();
607        let planned = plan_image_bytes(buf.get_ref()).unwrap();
608        let prep = prep_image_bytes(buf.get_ref()).unwrap();
609        assert_eq!(planned, (prep.gh, prep.gw), "planned grid == decoded grid");
610    }
611
612    #[test]
613    fn patchify_shape_and_order() {
614        // 2x2-patch (32x32 px) synthetic image, distinct channel values
615        let mut img = RgbImage::new(64, 64);
616        for (x, y, p) in img.enumerate_pixels_mut() {
617            *p = image::Rgb([x as u8, y as u8, 200]);
618        }
619        let mut buf = std::io::Cursor::new(Vec::new());
620        img.write_to(&mut buf, image::ImageFormat::Png).unwrap();
621        let prep = prep_image_bytes(buf.get_ref()).unwrap();
622        assert_eq!(prep.patches.len(), prep.gh * prep.gw * V_PATCH_IN);
623        assert_eq!(prep.gh % V_MERGE, 0);
624        assert_eq!(prep.gw % V_MERGE, 0);
625        // temporal slots identical for still images
626        let row = &prep.patches[0..V_PATCH_IN];
627        let slot = V_PATCH * V_PATCH;
628        for c in 0..3 {
629            let b = c * V_TEMPORAL * slot;
630            assert_eq!(row[b..b + slot], row[b + slot..b + 2 * slot]);
631        }
632        // values in [-1, 1]
633        assert!(prep.patches.iter().all(|v| (-1.0..=1.0).contains(v)));
634    }
635
636    /// Hand-crafted minimal GIF: `frames` one-black-pixel frames on a `w x h` canvas.
637    /// Each frame is the canonical 35-byte smallest-GIF image block (LZW: clear, index 0,
638    /// end -> bytes 0x44 0x01), so a high frame count costs ~18 bytes/frame on the wire
639    /// while every DECODED frame composites to the full canvas — the decode-bomb shape.
640    fn crafted_gif(w: u16, h: u16, frames: usize) -> Vec<u8> {
641        let mut b = Vec::new();
642        b.extend_from_slice(b"GIF89a");
643        b.extend_from_slice(&w.to_le_bytes());
644        b.extend_from_slice(&h.to_le_bytes());
645        b.push(0x80); // global color table, 2 entries
646        b.push(0); // background color index
647        b.push(0); // aspect ratio
648        b.extend_from_slice(&[0, 0, 0, 0xFF, 0xFF, 0xFF]); // GCT: black, white
649        for _ in 0..frames {
650            b.push(0x2C); // image descriptor
651            b.extend_from_slice(&0u16.to_le_bytes()); // left
652            b.extend_from_slice(&0u16.to_le_bytes()); // top
653            b.extend_from_slice(&1u16.to_le_bytes()); // width 1
654            b.extend_from_slice(&1u16.to_le_bytes()); // height 1
655            b.push(0); // no local color table
656            b.push(0x02); // LZW min code size
657            b.extend_from_slice(&[0x02, 0x44, 0x01]); // sub-block: clear, idx 0, end
658            b.push(0x00); // block terminator
659        }
660        b.push(0x3B); // trailer
661        b
662    }
663
664    #[test]
665    fn gif_decode_bomb_is_refused_before_full_expansion() {
666        // 2000x2000 canvas = 4M px/frame -> the 67.1M px budget admits 16 frames; a
667        // 64-frame stream (~1.3 KB on the wire, ~1 GiB decoded) must refuse at the
668        // budget, not expand: pre-fix this test allocated 64 x 16 MB RGBA canvases.
669        fn expect_err(bytes: &[u8]) -> String {
670            match prep_video_gif(bytes) {
671                Err(e) => e,
672                Ok(_) => panic!("decode-bomb GIF was accepted"),
673            }
674        }
675        let bomb = crafted_gif(2000, 2000, 64);
676        assert!(bomb.len() < 2048, "the bomb itself is tiny on the wire");
677        let err = expect_err(&bomb);
678        assert!(err.contains("decode budget"), "{err}");
679
680        // same canvas, frame count within budget: decodes fine.
681        let ok = crafted_gif(2000, 2000, 4);
682        let vid = prep_video_gif(&ok).unwrap();
683        assert_eq!(vid.groups.len(), 2); // 4 frames -> 2 temporal groups
684
685        // frame-count bomb on a tiny canvas: trips the flat frame ceiling.
686        let err = expect_err(&crafted_gif(8, 8, GIF_MAX_FRAMES + 8));
687        assert!(err.contains("decode budget"), "{err}");
688
689        // canvas alone past the pixel budget: refused straight from the header.
690        let err = expect_err(&crafted_gif(0xFFFF, 0xFFFF, 1));
691        assert!(err.contains("exceeds the decode budget"), "{err}");
692    }
693
694    #[test]
695    fn gif_plan_reads_metadata_without_materializing_frames() {
696        let bytes = crafted_gif(64, 64, 4);
697        let plan = plan_video_gif(&bytes).unwrap();
698        assert_eq!(plan.groups.len(), 2);
699        assert!(plan.groups.iter().all(|group| group.gh > 0 && group.gw > 0));
700        let prepared = prep_video_gif(&bytes).unwrap();
701        assert_eq!(
702            plan.groups
703                .iter()
704                .map(|group| (group.gh, group.gw))
705                .collect::<Vec<_>>(),
706            prepared
707                .groups
708                .iter()
709                .map(|group| (group.gh, group.gw))
710                .collect::<Vec<_>>()
711        );
712        assert!(plan_video_gif(&crafted_gif(2000, 2000, 64)).is_err());
713    }
714
715    #[test]
716    fn data_uri_roundtrip() {
717        let png = {
718            let img = RgbImage::new(32, 32);
719            let mut buf = std::io::Cursor::new(Vec::new());
720            img.write_to(&mut buf, image::ImageFormat::Png).unwrap();
721            buf.into_inner()
722        };
723        let uri = format!(
724            "data:image/png;base64,{}",
725            base64::engine::general_purpose::STANDARD.encode(&png)
726        );
727        let prep = prep_data_uri(&uri).unwrap();
728        assert_eq!(prep.n_tokens(), prep.gh * prep.gw / 4);
729        assert!(decode_data_uri("http://x/y.png").is_err());
730    }
731
732    /// The per-image raw cap is enforced at the decode, BY ENCODED LENGTH, so an
733    /// oversized data URI is refused by name before a single byte is decoded (hermes
734    /// review finding 48f96cb4cd37e436: only MAX_BODY_BYTES bounded this, letting one
735    /// image expand ~144 MiB of host bytes pre-admission).
736    #[test]
737    fn data_uri_per_image_raw_cap() {
738        let cap_chars = IMG_MAX_RAW_BYTES / 3 * 4;
739        // One base64 quad past the cap: refused, and the error names the limit.
740        let over = format!("data:image/png;base64,{}", "A".repeat(cap_chars + 4));
741        let err = decode_data_uri(&over).unwrap_err();
742        assert!(
743            err.contains("12 MiB"),
744            "cap refusal must name the limit: {err}"
745        );
746        // Exactly at the cap: the size gate admits and the decode yields cap bytes
747        // (the image decoder then judges the content, as before).
748        let at = format!("data:image/png;base64,{}", "A".repeat(cap_chars));
749        assert_eq!(decode_data_uri(&at).unwrap().len(), IMG_MAX_RAW_BYTES);
750        // The gemma mirror carries the identical cap.
751        let gerr = crate::vision_gemma::gemma_decode_data_uri(&over).unwrap_err();
752        assert!(
753            gerr.contains("12 MiB"),
754            "gemma cap refusal must name the limit: {gerr}"
755        );
756    }
757}