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; #[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
258pub 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 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
385fn 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
412pub 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 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(); 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 if frames.len() == 1 {
464 let f0 = frames[0].clone();
465 frames.push((f0.0, f0.1));
466 }
467 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) .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
496pub const IMG_MAX_RAW_BYTES: usize = 12 * 1024 * 1024;
503
504pub(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
518pub 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 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 let (h, w) = smart_resize(64, 64).unwrap();
552 assert!(h * w >= MIN_PIXELS);
553 let (h, w) = smart_resize(8000, 12000).unwrap();
555 assert!(h * w <= MAX_PIXELS);
556 assert!(smart_resize(10, 4000).is_err()); }
558
559 fn bmp_header_claiming(w: u32, h: u32) -> Vec<u8> {
564 let mut b = Vec::new();
565 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());
571 b.extend_from_slice(&(h as i32).to_le_bytes());
572 b.extend_from_slice(&1u16.to_le_bytes()); b.extend_from_slice(&24u16.to_le_bytes()); b.extend_from_slice(&[0u8; 24]); b
576 }
577
578 #[test]
579 fn decode_bomb_refuses_pre_decode() {
580 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 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 assert_eq!(image_header_dims(&bomb).unwrap(), (16_000, 16_000));
603 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 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 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 assert!(prep.patches.iter().all(|v| (-1.0..=1.0).contains(v)));
634 }
635
636 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); b.push(0); b.push(0); b.extend_from_slice(&[0, 0, 0, 0xFF, 0xFF, 0xFF]); for _ in 0..frames {
650 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); }
660 b.push(0x3B); b
662 }
663
664 #[test]
665 fn gif_decode_bomb_is_refused_before_full_expansion() {
666 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 let ok = crafted_gif(2000, 2000, 4);
682 let vid = prep_video_gif(&ok).unwrap();
683 assert_eq!(vid.groups.len(), 2); let err = expect_err(&crafted_gif(8, 8, GIF_MAX_FRAMES + 8));
687 assert!(err.contains("decode budget"), "{err}");
688
689 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 #[test]
737 fn data_uri_per_image_raw_cap() {
738 let cap_chars = IMG_MAX_RAW_BYTES / 3 * 4;
739 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 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 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}