305 lines
10 KiB
Rust
305 lines
10 KiB
Rust
//! PNG/JPEG preprocessing and image positions from MTPLX's Qwen3-VL path.
|
|
use image::{DynamicImage, ImageDecoder, ImageReader, RgbImage};
|
|
use std::io::Cursor;
|
|
|
|
pub(super) struct Image {
|
|
pub(super) pixels: Vec<f32>,
|
|
pub(super) h: i32,
|
|
pub(super) w: i32,
|
|
}
|
|
|
|
pub(super) fn decode(bytes: &[u8]) -> Result<Image, String> {
|
|
if bytes.is_empty() || bytes.len() > 50 * 1024 * 1024 {
|
|
return Err("Qwen image exceeds the 50 MiB encoded limit".into());
|
|
}
|
|
let reader = ImageReader::new(Cursor::new(bytes))
|
|
.with_guessed_format()
|
|
.map_err(|e| e.to_string())?;
|
|
if !matches!(
|
|
reader.format(),
|
|
Some(image::ImageFormat::Png | image::ImageFormat::Jpeg)
|
|
) {
|
|
return Err("Qwen images must be PNG or JPEG".into());
|
|
}
|
|
let mut decoder = reader.into_decoder().map_err(|e| e.to_string())?;
|
|
let (w, h) = decoder.dimensions();
|
|
if w == 0 || h == 0 || w > 8000 || h > 8000 {
|
|
return Err("Qwen images must have dimensions between 1 and 8000".into());
|
|
}
|
|
let orientation = decoder.orientation().map_err(|e| e.to_string())?;
|
|
let mut image = DynamicImage::from_decoder(decoder).map_err(|e| e.to_string())?;
|
|
image.apply_orientation(orientation);
|
|
let mut image = image.into_rgb8();
|
|
let (w, h) = image.dimensions();
|
|
let (rh, rw) = smart_resize(h, w)?;
|
|
if (w, h) != (rw, rh) {
|
|
image = resize(&image, rw, rh);
|
|
}
|
|
let (h, w) = ((rh / 16) as i32, (rw / 16) as i32);
|
|
let mut pixels = Vec::with_capacity(h as usize * w as usize * 1536);
|
|
for br in 0..h / 2 {
|
|
for bc in 0..w / 2 {
|
|
for ir in 0..2 {
|
|
for ic in 0..2 {
|
|
for c in 0..3 {
|
|
for _t in 0..2 {
|
|
for y in 0..16 {
|
|
for x in 0..16 {
|
|
let v = image.get_pixel(
|
|
((bc * 2 + ic) * 16 + x) as u32,
|
|
((br * 2 + ir) * 16 + y) as u32,
|
|
)[c];
|
|
pixels.push((f32::from(v) * (1_f32 / 255.) - 0.5) / 0.5);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Ok(Image { pixels, h, w })
|
|
}
|
|
fn smart_resize(h: u32, w: u32) -> Result<(u32, u32), String> {
|
|
if f64::from(h.max(w)) / f64::from(h.min(w)) > 200. {
|
|
return Err("Qwen image aspect ratio exceeds 200".into());
|
|
}
|
|
let (hf, wf) = (f64::from(h), f64::from(w));
|
|
let (mut rh, mut rw) = (
|
|
(hf / 32.).round_ties_even() as u32 * 32,
|
|
(wf / 32.).round_ties_even() as u32 * 32,
|
|
);
|
|
if u64::from(rh) * u64::from(rw) > 16777216 {
|
|
let beta = (hf * wf / 16777216.).sqrt();
|
|
rh = ((hf / beta / 32.).floor() as u32 * 32).max(32);
|
|
rw = ((wf / beta / 32.).floor() as u32 * 32).max(32);
|
|
} else if u64::from(rh) * u64::from(rw) < 65536 {
|
|
let beta = (65536. / (hf * wf)).sqrt();
|
|
rh = (hf * beta / 32.).ceil() as u32 * 32;
|
|
rw = (wf * beta / 32.).ceil() as u32 * 32;
|
|
}
|
|
Ok((rh, rw))
|
|
}
|
|
// Pillow Resample.c: bicubic a=-0.5, normalized 22-bit coefficients,
|
|
// horizontal then vertical passes with rounding/clipping after each pass.
|
|
fn coefficients(input: u32, output: u32) -> Vec<(usize, Vec<i32>)> {
|
|
let scale = f64::from(input) / f64::from(output);
|
|
let filter_scale = scale.max(1.);
|
|
let support = 2. * filter_scale;
|
|
(0..output)
|
|
.map(|i| {
|
|
let center = (f64::from(i) + 0.5) * scale;
|
|
let first = ((center - support + 0.5) as i64).max(0) as usize;
|
|
let last = ((center + support + 0.5) as i64).min(i64::from(input)) as usize;
|
|
let ws = (first..last)
|
|
.map(|j| {
|
|
let x = ((j as f64 - center + 0.5) / filter_scale).abs();
|
|
if x < 1. {
|
|
((1.5 * x - 2.5) * x) * x + 1.
|
|
} else if x < 2. {
|
|
(((x - 5.) * x + 8.) * x - 4.) * (-0.5)
|
|
} else {
|
|
0.
|
|
}
|
|
})
|
|
.collect::<Vec<_>>();
|
|
let sum = ws.iter().sum::<f64>();
|
|
(
|
|
first,
|
|
ws.iter()
|
|
.map(|w| (w / sum * 4194304.).round() as i32)
|
|
.collect(),
|
|
)
|
|
})
|
|
.collect()
|
|
}
|
|
fn resize(input: &RgbImage, w: u32, h: u32) -> RgbImage {
|
|
let horizontal = if w == input.width() {
|
|
input.clone()
|
|
} else {
|
|
let coeff = coefficients(input.width(), w);
|
|
RgbImage::from_fn(w, input.height(), |x, y| {
|
|
let (start, ws) = &coeff[x as usize];
|
|
image::Rgb(std::array::from_fn(|c| {
|
|
let sum = ws.iter().enumerate().fold(1_i64 << 21, |sum, (i, k)| {
|
|
sum + i64::from(*k) * i64::from(input.get_pixel((*start + i) as u32, y)[c])
|
|
});
|
|
(sum >> 22).clamp(0, 255) as u8
|
|
}))
|
|
})
|
|
};
|
|
if h == horizontal.height() {
|
|
return horizontal;
|
|
}
|
|
let coeff = coefficients(horizontal.height(), h);
|
|
RgbImage::from_fn(w, h, |x, y| {
|
|
let (start, ws) = &coeff[y as usize];
|
|
image::Rgb(std::array::from_fn(|c| {
|
|
let sum = ws.iter().enumerate().fold(1_i64 << 21, |sum, (i, k)| {
|
|
sum + i64::from(*k) * i64::from(horizontal.get_pixel(x, (*start + i) as u32)[c])
|
|
});
|
|
(sum >> 22).clamp(0, 255) as u8
|
|
}))
|
|
})
|
|
}
|
|
|
|
pub(super) fn positions(ids: &[i32], grids: &[(i32, i32)]) -> Result<(Vec<i32>, i32), String> {
|
|
let n = ids.len();
|
|
let mut axes = vec![0; n * 3];
|
|
let (mut start, mut next) = (0, 0_i32);
|
|
for &(h, w) in grids {
|
|
let h = h / 2;
|
|
let w = w / 2;
|
|
if h <= 0 || w <= 0 {
|
|
return Err("invalid Qwen image grid".into());
|
|
}
|
|
let end = start
|
|
+ ids[start..]
|
|
.iter()
|
|
.position(|&t| t == 248056)
|
|
.ok_or("missing Qwen image pad")?;
|
|
for i in start..end {
|
|
for axis in 0..3 {
|
|
axes[axis * n + i] = next + (i - start) as i32;
|
|
}
|
|
}
|
|
next += (end - start) as i32;
|
|
let count = (h * w) as usize;
|
|
if end + count > n || ids[end..end + count].iter().any(|&t| t != 248056) {
|
|
return Err("Qwen image pad count differs from grid".into());
|
|
}
|
|
for i in 0..count {
|
|
axes[end + i] = next;
|
|
axes[n + end + i] = next + i as i32 / w;
|
|
axes[n * 2 + end + i] = next + i as i32 % w;
|
|
}
|
|
next += h.max(w).max(1);
|
|
start = end + count;
|
|
}
|
|
if ids[start..].contains(&248056) {
|
|
return Err("Qwen image pads exceed image grids".into());
|
|
}
|
|
for i in start..n {
|
|
for axis in 0..3 {
|
|
axes[axis * n + i] = next + (i - start) as i32;
|
|
}
|
|
}
|
|
next += (n - start) as i32;
|
|
Ok((axes, next - n as i32))
|
|
}
|
|
|
|
#[test]
|
|
fn qwen_vision_positions_and_resize_contract() {
|
|
assert_eq!(smart_resize(1024, 1024).unwrap(), (1024, 1024));
|
|
assert_eq!(smart_resize(16, 16).unwrap(), (256, 256));
|
|
assert!(smart_resize(1, 201).is_err());
|
|
let (table, delta) = positions(&[1, 248056, 248056, 248056, 248056, 2], &[(4, 4)]).unwrap();
|
|
assert_eq!(
|
|
table,
|
|
vec![0, 1, 1, 1, 1, 3, 0, 1, 1, 2, 2, 3, 0, 1, 2, 1, 2, 3]
|
|
);
|
|
assert_eq!(delta, -2);
|
|
assert!(positions(&[248056], &[(4, 4)]).is_err());
|
|
}
|
|
|
|
#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
|
pub(super) struct Identity {
|
|
pub(super) start: usize,
|
|
pub(super) end: usize,
|
|
pub(super) digest: [u8; 32],
|
|
/// Missing on checkpoints written before image-aware sparse QSA.
|
|
#[serde(default)]
|
|
pub(super) policy: u32,
|
|
}
|
|
|
|
pub(super) fn cache_policy() -> u32 {
|
|
u32::from(std::env::var("MTPLX_QWEN4_VISION_QSA").as_deref() != Ok("0"))
|
|
}
|
|
|
|
pub(super) fn compatible(images: &[Identity]) -> bool {
|
|
images.iter().all(|image| image.policy == cache_policy())
|
|
}
|
|
|
|
pub(super) fn same_prefix(a: &[Identity], b: &[Identity], tokens: usize) -> bool {
|
|
// Image pads do not identify partial image embeddings. MTPLX only restores
|
|
// before an image or after its complete span, never into the middle.
|
|
!b.iter()
|
|
.any(|image| image.start < tokens && tokens < image.end)
|
|
&& a.iter()
|
|
.filter(|i| i.start < tokens)
|
|
.eq(b.iter().filter(|i| i.start < tokens))
|
|
}
|
|
|
|
#[test]
|
|
fn qwen_image_preprocessing_matches_pillow_mtplx() {
|
|
use sha2::{Digest, Sha256};
|
|
for (w, h, gh, gw, expected) in [
|
|
(
|
|
17,
|
|
31,
|
|
22,
|
|
12,
|
|
"339e4ac4e02bb8fcab6d521194200191e4c77fbcce224de98b40d8d48393930f",
|
|
),
|
|
(
|
|
333,
|
|
177,
|
|
12,
|
|
22,
|
|
"3fc0890cb34e0c37bf32c5ce2a8bfa9b178ae3564c8e97b8e6496e71a66b2644",
|
|
),
|
|
(
|
|
1537,
|
|
257,
|
|
16,
|
|
96,
|
|
"96bcea68014063c4edc51b422af82930e1bd43984d57a41e12f44865eed45368",
|
|
),
|
|
] {
|
|
let image = RgbImage::from_fn(w, h, |x, y| {
|
|
image::Rgb(std::array::from_fn(|c| {
|
|
((x * 13 + y * 7 + c as u32 * 31) % 256) as u8
|
|
}))
|
|
});
|
|
let mut bytes = Cursor::new(Vec::new());
|
|
image.write_to(&mut bytes, image::ImageFormat::Png).unwrap();
|
|
let input = decode(bytes.get_ref()).unwrap();
|
|
assert_eq!((input.h, input.w), (gh, gw));
|
|
let hash = Sha256::digest(
|
|
input
|
|
.pixels
|
|
.iter()
|
|
.flat_map(|v| v.to_le_bytes())
|
|
.collect::<Vec<_>>(),
|
|
);
|
|
assert_eq!(
|
|
hash.iter().map(|b| format!("{b:02x}")).collect::<String>(),
|
|
expected,
|
|
"{w}x{h}"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn qwen_image_cache_keys_distinguish_pixels_and_preserve_prior_text() {
|
|
let old = vec![Identity {
|
|
start: 10,
|
|
end: 14,
|
|
digest: [1; 32],
|
|
policy: cache_policy(),
|
|
}];
|
|
let changed = vec![Identity {
|
|
start: 10,
|
|
end: 14,
|
|
digest: [2; 32],
|
|
policy: cache_policy(),
|
|
}];
|
|
assert!(same_prefix(&old, &old, 20));
|
|
assert!(!same_prefix(&old, &changed, 20));
|
|
assert!(same_prefix(&old, &changed, 10));
|
|
assert!(!same_prefix(&old, &[], 20));
|
|
assert!(!same_prefix(&old, &old, 11));
|
|
assert!(!same_prefix(&old, &old, 13));
|
|
assert!(same_prefix(&old, &old, 14));
|
|
}
|