1use crate::vision::{V_MERGE, V_PATCH, V_PATCH_IN, V_TEMPORAL};
16use base64::Engine as _;
17use image::RgbImage;
18use image::imageops::FilterType;
19
20pub const MIN_PIXELS: usize = 65536;
22pub const MAX_PIXELS: usize = 16_777_216;
23const FACTOR: usize = V_PATCH * V_MERGE; pub struct PreppedImage {
26 pub patches: Vec<f32>,
28 pub gh: usize,
29 pub gw: usize,
30}
31
32impl PreppedImage {
33 pub fn n_tokens(&self) -> usize {
35 n_tokens_for_grid(self.gh, self.gw)
36 }
37}
38
39pub fn n_tokens_for_grid(gh: usize, gw: usize) -> usize {
42 gh * gw / (V_MERGE * V_MERGE)
43}
44
45fn 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
58pub 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
84pub const IMG_MAX_DECODE_PIXELS: usize = 1 << 26;
92
93pub 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
103pub 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
119fn 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
141fn 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
162pub 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
173pub fn prep_data_uri(uri: &str) -> Result<PreppedImage, String> {
175 let bytes = decode_data_uri(uri)?;
176 prep_image_bytes(&bytes)
177}
178
179pub struct VisionUnit {
183 pub prep: PreppedImage,
184 pub video: Option<usize>,
186}
187
188pub struct PreppedVideo {
193 pub groups: Vec<PreppedImage>,
194 pub timestamps: Vec<f32>,
195}
196
197pub 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;
206pub const VID_MAX_FRAMES: usize = 32;
208
209pub const GIF_MAX_FRAMES: usize = 512;
221pub const GIF_MAX_TOTAL_PIXELS: usize = 1 << 26; fn smart_resize_video(frames: usize, h: usize, w: usize) -> Result<(usize, usize), String> {
225 if h < 2 || w < 2 {
226 return Err(format!("frame too small: {w}x{h}"));
227 }
228 let ar = h.max(w) as f64 / h.min(w) as f64;
229 if ar > 200.0 {
230 return Err(format!("aspect ratio {ar:.0} exceeds 200"));
231 }
232 let f = FACTOR as f64;
233 let (hf, wf) = (h as f64, w as f64);
234 let t_bar = ((frames as f64 / V_TEMPORAL as f64).round() * V_TEMPORAL as f64).max(2.0);
235 let mut h_bar = (round_half_even(hf / f) * f).max(f);
236 let mut w_bar = (round_half_even(wf / f) * f).max(f);
237 let (min_px, max_px) = (VID_MIN_PIXELS as f64, video_max_pixels() as f64);
238 if t_bar * h_bar * w_bar > max_px {
239 let beta = (frames as f64 * hf * wf / max_px).sqrt();
240 h_bar = ((hf / beta / f).floor() * f).max(f);
241 w_bar = ((wf / beta / f).floor() * f).max(f);
242 } else if t_bar * h_bar * w_bar < min_px {
243 let beta = (min_px / (frames as f64 * hf * wf)).sqrt();
244 h_bar = (hf * beta / f).ceil() * f;
245 w_bar = (wf * beta / f).ceil() * f;
246 }
247 Ok((h_bar as usize, w_bar as usize))
248}
249
250pub fn prep_video_gif(bytes: &[u8]) -> Result<PreppedVideo, String> {
255 use image::AnimationDecoder;
256 use image::ImageDecoder as _;
257 let dec = image::codecs::gif::GifDecoder::new(std::io::Cursor::new(bytes))
258 .map_err(|e| format!("gif decode: {e}"))?;
259 let (cw, ch) = dec.dimensions();
265 let canvas_px = (cw as usize) * (ch as usize);
266 if canvas_px == 0 {
267 return Err("gif has an empty canvas".into());
268 }
269 let max_frames = GIF_MAX_FRAMES.min(GIF_MAX_TOTAL_PIXELS / canvas_px);
270 if max_frames == 0 {
271 return Err(format!(
272 "gif canvas {cw}x{ch} exceeds the decode budget ({GIF_MAX_TOTAL_PIXELS} px)"
273 ));
274 }
275 let mut frames: Vec<(RgbImage, f32)> = Vec::new(); let mut t = 0f32;
277 for fr in dec.into_frames() {
278 if frames.len() >= max_frames {
279 return Err(format!(
280 "gif exceeds the decode budget: more than {max_frames} frames at {cw}x{ch} \
281 (ceiling {GIF_MAX_FRAMES} frames / {GIF_MAX_TOTAL_PIXELS} total px)"
282 ));
283 }
284 let fr = fr.map_err(|e| format!("gif frame: {e}"))?;
285 let (num, den) = fr.delay().numer_denom_ms();
286 let dt = if den == 0 {
287 100.0
288 } else {
289 num as f32 / den as f32
290 } / 1000.0;
291 frames.push((
292 image::DynamicImage::ImageRgba8(fr.into_buffer()).to_rgb8(),
293 t,
294 ));
295 t += dt.max(0.01);
296 }
297 if frames.is_empty() {
298 return Err("gif has no frames".into());
299 }
300 if frames.len() == 1 {
302 let f0 = frames[0].clone();
303 frames.push((f0.0, f0.1));
304 }
305 let total = frames.len();
307 let take = total.min(VID_MAX_FRAMES) & !1;
308 let picked: Vec<usize> = (0..take)
309 .map(|i| i * total / take) .collect();
311 let (h, w) = (frames[0].0.height() as usize, frames[0].0.width() as usize);
312 let (rh, rw) = smart_resize_video(take, h, w)?;
313 let (gh, gw) = (rh / V_PATCH, rw / V_PATCH);
314 let mut groups = Vec::with_capacity(take / 2);
315 let mut timestamps = Vec::with_capacity(take / 2);
316 for g in 0..take / 2 {
317 let (a, b) = (picked[2 * g], picked[2 * g + 1]);
318 let mut patches = vec![0f32; gh * gw * V_PATCH_IN];
319 for (slot, idx) in [(0usize, a), (1usize, b)] {
320 let resized = image::imageops::resize(
321 &frames[idx].0,
322 rw as u32,
323 rh as u32,
324 FilterType::CatmullRom,
325 );
326 fill_slot(&mut patches, &resized, gh, gw, slot);
327 }
328 groups.push(PreppedImage { patches, gh, gw });
329 timestamps.push(frames[a].1);
330 }
331 Ok(PreppedVideo { groups, timestamps })
332}
333
334pub fn decode_data_uri(uri: &str) -> Result<Vec<u8>, String> {
336 let rest = uri
337 .strip_prefix("data:")
338 .ok_or_else(|| "expected data: URI (http fetch requires MEMRA_FETCH_URLS=1)".to_string())?;
339 let (meta, payload) = rest
340 .split_once(',')
341 .ok_or_else(|| "malformed data URI: no comma".to_string())?;
342 if !meta.ends_with(";base64") {
343 return Err("data URI must be base64-encoded".into());
344 }
345 base64::engine::general_purpose::STANDARD
346 .decode(payload.trim())
347 .map_err(|e| format!("base64 decode: {e}"))
348}
349
350#[cfg(test)]
351mod tests {
352 use super::*;
353
354 #[test]
355 fn smart_resize_multiples_and_budget() {
356 let (h, w) = smart_resize(1080, 1920).unwrap();
358 assert_eq!(h % 32, 0);
359 assert_eq!(w % 32, 0);
360 assert!(h * w >= MIN_PIXELS && h * w <= MAX_PIXELS);
361 let (h, w) = smart_resize(64, 64).unwrap();
363 assert!(h * w >= MIN_PIXELS);
364 let (h, w) = smart_resize(8000, 12000).unwrap();
366 assert!(h * w <= MAX_PIXELS);
367 assert!(smart_resize(10, 4000).is_err()); }
369
370 fn bmp_header_claiming(w: u32, h: u32) -> Vec<u8> {
375 let mut b = Vec::new();
376 b.extend_from_slice(b"BM"); b.extend_from_slice(&54u32.to_le_bytes()); b.extend_from_slice(&0u32.to_le_bytes()); b.extend_from_slice(&54u32.to_le_bytes()); b.extend_from_slice(&40u32.to_le_bytes()); b.extend_from_slice(&(w as i32).to_le_bytes());
382 b.extend_from_slice(&(h as i32).to_le_bytes());
383 b.extend_from_slice(&1u16.to_le_bytes()); b.extend_from_slice(&24u16.to_le_bytes()); b.extend_from_slice(&[0u8; 24]); b
387 }
388
389 #[test]
390 fn decode_bomb_refuses_pre_decode() {
391 let bomb = bmp_header_claiming(16_000, 16_000);
396 let err = plan_image_bytes(&bomb).unwrap_err();
397 assert!(
398 err.contains("exceeds the decode budget"),
399 "want the named pre-decode refusal, got: {err}"
400 );
401 let err = match prep_image_bytes(&bomb) {
405 Ok(_) => panic!("bomb must not prep"),
406 Err(e) => e,
407 };
408 assert!(
409 err.contains("exceeds the decode budget"),
410 "prep must refuse at admission, not at decode: {err}"
411 );
412 assert_eq!(image_header_dims(&bomb).unwrap(), (16_000, 16_000));
414 let img = RgbImage::new(64, 64);
416 let mut buf = std::io::Cursor::new(Vec::new());
417 img.write_to(&mut buf, image::ImageFormat::Png).unwrap();
418 let planned = plan_image_bytes(buf.get_ref()).unwrap();
419 let prep = prep_image_bytes(buf.get_ref()).unwrap();
420 assert_eq!(planned, (prep.gh, prep.gw), "planned grid == decoded grid");
421 }
422
423 #[test]
424 fn patchify_shape_and_order() {
425 let mut img = RgbImage::new(64, 64);
427 for (x, y, p) in img.enumerate_pixels_mut() {
428 *p = image::Rgb([x as u8, y as u8, 200]);
429 }
430 let mut buf = std::io::Cursor::new(Vec::new());
431 img.write_to(&mut buf, image::ImageFormat::Png).unwrap();
432 let prep = prep_image_bytes(buf.get_ref()).unwrap();
433 assert_eq!(prep.patches.len(), prep.gh * prep.gw * V_PATCH_IN);
434 assert_eq!(prep.gh % V_MERGE, 0);
435 assert_eq!(prep.gw % V_MERGE, 0);
436 let row = &prep.patches[0..V_PATCH_IN];
438 let slot = V_PATCH * V_PATCH;
439 for c in 0..3 {
440 let b = c * V_TEMPORAL * slot;
441 assert_eq!(row[b..b + slot], row[b + slot..b + 2 * slot]);
442 }
443 assert!(prep.patches.iter().all(|v| (-1.0..=1.0).contains(v)));
445 }
446
447 fn crafted_gif(w: u16, h: u16, frames: usize) -> Vec<u8> {
452 let mut b = Vec::new();
453 b.extend_from_slice(b"GIF89a");
454 b.extend_from_slice(&w.to_le_bytes());
455 b.extend_from_slice(&h.to_le_bytes());
456 b.push(0x80); b.push(0); b.push(0); b.extend_from_slice(&[0, 0, 0, 0xFF, 0xFF, 0xFF]); for _ in 0..frames {
461 b.push(0x2C); b.extend_from_slice(&0u16.to_le_bytes()); b.extend_from_slice(&0u16.to_le_bytes()); b.extend_from_slice(&1u16.to_le_bytes()); b.extend_from_slice(&1u16.to_le_bytes()); b.push(0); b.push(0x02); b.extend_from_slice(&[0x02, 0x44, 0x01]); b.push(0x00); }
471 b.push(0x3B); b
473 }
474
475 #[test]
476 fn gif_decode_bomb_is_refused_before_full_expansion() {
477 fn expect_err(bytes: &[u8]) -> String {
481 match prep_video_gif(bytes) {
482 Err(e) => e,
483 Ok(_) => panic!("decode-bomb GIF was accepted"),
484 }
485 }
486 let bomb = crafted_gif(2000, 2000, 64);
487 assert!(bomb.len() < 2048, "the bomb itself is tiny on the wire");
488 let err = expect_err(&bomb);
489 assert!(err.contains("decode budget"), "{err}");
490
491 let ok = crafted_gif(2000, 2000, 4);
493 let vid = prep_video_gif(&ok).unwrap();
494 assert_eq!(vid.groups.len(), 2); let err = expect_err(&crafted_gif(8, 8, GIF_MAX_FRAMES + 8));
498 assert!(err.contains("decode budget"), "{err}");
499
500 let err = expect_err(&crafted_gif(0xFFFF, 0xFFFF, 1));
502 assert!(err.contains("exceeds the decode budget"), "{err}");
503 }
504
505 #[test]
506 fn data_uri_roundtrip() {
507 let png = {
508 let img = RgbImage::new(32, 32);
509 let mut buf = std::io::Cursor::new(Vec::new());
510 img.write_to(&mut buf, image::ImageFormat::Png).unwrap();
511 buf.into_inner()
512 };
513 let uri = format!(
514 "data:image/png;base64,{}",
515 base64::engine::general_purpose::STANDARD.encode(&png)
516 );
517 let prep = prep_data_uri(&uri).unwrap();
518 assert_eq!(prep.n_tokens(), prep.gh * prep.gw / 4);
519 assert!(decode_data_uri("http://x/y.png").is_err());
520 }
521}