Add Qwen vision loading and Metal inference
This commit is contained in:
Georg Bauer committed
1 parent
48c2f751b4
commit
e32954ac1b
30 files changed
+1687
-100
No files matched your search
@@ -0,0 +1,284 @@
|
||||
//! 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],
|
||||
}
|
||||
|
||||
pub(super) fn same_prefix(a: &[Identity], b: &[Identity], tokens: usize) -> bool {
|
||||
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],
|
||||
}];
|
||||
let changed = vec![Identity {
|
||||
start: 10,
|
||||
end: 14,
|
||||
digest: [2; 32],
|
||||
}];
|
||||
assert!(same_prefix(&old, &old, 20));
|
||||
assert!(!same_prefix(&old, &changed, 20));
|
||||
assert!(same_prefix(&old, &changed, 10));
|
||||
assert!(!same_prefix(&old, &[], 20));
|
||||
}
|
||||
Reference in new issue
Block a user