455 lines
14 KiB
Rust
455 lines
14 KiB
Rust
//! DeepSeek V4 Flash Vision-Exp image layout and prefill masks.
|
|
//! Behavioral reference: DS4 0aaea5a238fb41a35106a551e73c8409dfb751ac.
|
|
use image::{DynamicImage, ImageDecoder, ImageReader};
|
|
use sha2::{Digest, Sha256};
|
|
use std::io::Cursor;
|
|
|
|
pub(super) const SOURCE_REVISION: &[u8] = b"e46e16bf6035c6f317eb2ac7458eb0362926d402";
|
|
pub(super) const START: i32 = 0;
|
|
pub(super) const PAD: i32 = 1;
|
|
pub(super) const IMAGE: i32 = 2;
|
|
pub(super) const NEWLINE: i32 = 3;
|
|
pub(super) const END: i32 = 4;
|
|
|
|
pub(super) struct Patches {
|
|
pub values: Vec<f32>,
|
|
pub fingerprint: [u8; 32],
|
|
pub width: u32,
|
|
pub height: u32,
|
|
pub content_width: u32,
|
|
pub content_height: u32,
|
|
pub grid_width: u32,
|
|
pub grid_height: u32,
|
|
}
|
|
|
|
pub(super) fn decode(encoded: &[u8]) -> Result<Patches, String> {
|
|
if encoded.is_empty() || encoded.len() > 64 * 1024 * 1024 {
|
|
return Err("image is empty or exceeds the 64 MiB encoded limit".into());
|
|
}
|
|
let reader = ImageReader::new(Cursor::new(encoded))
|
|
.with_guessed_format()
|
|
.map_err(|e| e.to_string())?;
|
|
let mut decoder = reader.into_decoder().map_err(|e| e.to_string())?;
|
|
let (width, height) = decoder.dimensions();
|
|
if width == 0
|
|
|| height == 0
|
|
|| width > 16384
|
|
|| height > 16384
|
|
|| u64::from(width) * u64::from(height) > 64 * 1024 * 1024
|
|
{
|
|
return Err("invalid DeepSeek image dimensions".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 rgb = image.into_rgb8();
|
|
preprocess(rgb.as_raw(), rgb.width(), rgb.height())
|
|
}
|
|
|
|
fn grid_tokens(height: u32, width: u32) -> u32 {
|
|
let h = (height / 14).div_ceil(3);
|
|
let w = (width / 14).div_ceil(3);
|
|
(h + h % 2) * (w + 1) + 2 + (h.div_ceil(2) * (w + 1) % 2) * 2
|
|
}
|
|
|
|
fn resize_shape(height: u32, width: u32) -> Result<(u32, u32), String> {
|
|
let mut pw = width.min(height * 8);
|
|
let mut ph = height;
|
|
if u64::from(pw) * u64::from(ph) < 147456 {
|
|
let scale = (147456.0 / (f64::from(pw) * f64::from(ph))).sqrt();
|
|
pw = ((f64::from(pw) * scale) as u32).max(1);
|
|
ph = ((f64::from(ph) * scale) as u32).max(1);
|
|
}
|
|
let mut bw = pw.div_ceil(14) * 14;
|
|
let mut bh = ph.div_ceil(14) * 14;
|
|
let mut budget = 381;
|
|
while grid_tokens(bh, bw) > 381 {
|
|
if budget <= 4 {
|
|
return Err("image cannot fit the DeepSeek vision token budget".into());
|
|
}
|
|
let ratio = f64::from(ph) / f64::from(pw);
|
|
let mw = ((f64::from(budget) - 2.0) / ratio + 0.25).sqrt() - 0.5;
|
|
let mh = mw * ratio;
|
|
if mw < 1.0 {
|
|
bw = 42;
|
|
bh = (((budget - 2) / 2) & !1) * 42;
|
|
} else if mh < 2.0 {
|
|
bh = 84;
|
|
bw = ((budget - 2) / 2 - 1) * 42;
|
|
if bw <= 42 {
|
|
return Err("image cannot fit the DeepSeek vision token budget".into());
|
|
}
|
|
} else {
|
|
let mw = mw.floor() as u32;
|
|
let mh = mh.floor() as u32 & !1;
|
|
let scale =
|
|
(f64::from(mw) * 42.0 / f64::from(pw)).min(f64::from(mh) * 42.0 / f64::from(ph));
|
|
bw = (f64::from(pw) * scale / 14.0).floor() as u32 * 14;
|
|
bh = (f64::from(ph) * scale / 14.0).floor() as u32 * 14;
|
|
}
|
|
if bw == 0 || bh == 0 {
|
|
return Err("invalid DeepSeek resized image dimensions".into());
|
|
}
|
|
budget -= 1;
|
|
}
|
|
Ok((bh, bw))
|
|
}
|
|
|
|
pub(super) fn preprocess(rgb: &[u8], width: u32, height: u32) -> Result<Patches, String> {
|
|
if width == 0
|
|
|| height == 0
|
|
|| width > 16384
|
|
|| height > 16384
|
|
|| u64::from(width) * u64::from(height) > 64 * 1024 * 1024
|
|
|| rgb.len() != width as usize * height as usize * 3
|
|
{
|
|
return Err("invalid DeepSeek image dimensions".into());
|
|
}
|
|
let (bh, bw) = resize_shape(height, width)?;
|
|
if bw > 16384 || bh > 16384 || u64::from(bw) * u64::from(bh) > 64 * 1024 * 1024 {
|
|
return Err("resized DeepSeek image exceeds the decoded image limits".into());
|
|
}
|
|
let (cw, ch, ox, oy) = if width >= height * 8 {
|
|
(bw, bh, 0, 0)
|
|
} else {
|
|
let scale = (f64::from(bw) / f64::from(width)).min(f64::from(bh) / f64::from(height));
|
|
let cw = (f64::from(width) * scale)
|
|
.round_ties_even()
|
|
.clamp(1.0, f64::from(bw)) as u32;
|
|
let ch = (f64::from(height) * scale)
|
|
.round_ties_even()
|
|
.clamp(1.0, f64::from(bh)) as u32;
|
|
(
|
|
cw,
|
|
ch,
|
|
(f64::from(bw - cw) * 0.5).round_ties_even() as u32,
|
|
(f64::from(bh - ch) * 0.5).round_ties_even() as u32,
|
|
)
|
|
};
|
|
let mut canvas = vec![127.0_f32; bw as usize * bh as usize * 3];
|
|
let offset = (oy as usize * bw as usize + ox as usize) * 3;
|
|
super::metal::vision::resize_bicubic(rgb, width, height, &mut canvas[offset..], cw, ch, bw);
|
|
for value in &mut canvas {
|
|
*value = *value / 127.5 - 1.0;
|
|
}
|
|
let gw = bw / 14;
|
|
let gh = bh / 14;
|
|
let mut values = Vec::with_capacity(gw as usize * gh as usize * 588);
|
|
for py in 0..gh {
|
|
for px in 0..gw {
|
|
for c in 0..3 {
|
|
for y in 0..14 {
|
|
for x in 0..14 {
|
|
values.push(
|
|
canvas[((py * 14 + y) as usize * bw as usize + (px * 14 + x) as usize)
|
|
* 3
|
|
+ c],
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
let mut hash = Sha256::new();
|
|
hash.update(width.to_le_bytes());
|
|
hash.update(height.to_le_bytes());
|
|
hash.update(rgb);
|
|
Ok(Patches {
|
|
fingerprint: hash.finalize().into(),
|
|
values,
|
|
width,
|
|
height,
|
|
content_width: cw,
|
|
content_height: ch,
|
|
grid_width: gw,
|
|
grid_height: gh,
|
|
})
|
|
}
|
|
|
|
pub(super) struct Layout {
|
|
pub types: Vec<i32>,
|
|
pub permutation: Vec<usize>,
|
|
}
|
|
|
|
pub(super) fn layout(height: u32, width: u32, position: u32) -> Result<Layout, String> {
|
|
if height == 0 || width == 0 {
|
|
return Err("invalid DeepSeek image token grid".into());
|
|
}
|
|
let rows = u64::from(height) + u64::from(height % 2);
|
|
let row_length = u64::from(width) + 1;
|
|
let leading = 3 - position % 4;
|
|
let trailing = ((rows / 2 * row_length) % 2) * 2;
|
|
let count = u64::from(leading) + 2 + rows * row_length + trailing;
|
|
if count > 384 {
|
|
return Err("DeepSeek image token layout exceeds its model limit".into());
|
|
}
|
|
let mut types = vec![PAD; leading as usize];
|
|
let mut permutation = Vec::with_capacity(height as usize * width as usize);
|
|
types.push(START);
|
|
for pair in 0..rows / 2 {
|
|
for column in 0..row_length {
|
|
for row in pair * 2..pair * 2 + 2 {
|
|
if row < u64::from(height) && column < u64::from(width) {
|
|
types.push(IMAGE);
|
|
permutation.push((row * u64::from(width) + column) as usize);
|
|
} else {
|
|
types.push(if row < u64::from(height) {
|
|
NEWLINE
|
|
} else {
|
|
PAD
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
types.extend(std::iter::repeat_n(PAD, trailing as usize));
|
|
types.push(END);
|
|
Ok(Layout { types, permutation })
|
|
}
|
|
|
|
#[derive(Debug, PartialEq)]
|
|
pub(super) struct Span {
|
|
pub block: usize,
|
|
pub start: usize,
|
|
pub end: usize,
|
|
}
|
|
|
|
pub(super) fn spans(tokens: &[i32], vocab: i32) -> Result<Vec<Span>, String> {
|
|
let mut result = Vec::new();
|
|
let mut cursor = 0;
|
|
while cursor < tokens.len() {
|
|
if (0..vocab).contains(&tokens[cursor]) {
|
|
cursor += 1;
|
|
continue;
|
|
}
|
|
let block = cursor;
|
|
while cursor < tokens.len() && tokens[cursor] == vocab + PAD {
|
|
cursor += 1;
|
|
}
|
|
if tokens.get(cursor) != Some(&(vocab + START)) {
|
|
return Err("malformed DeepSeek image block".into());
|
|
}
|
|
let start = cursor;
|
|
cursor += 1;
|
|
loop {
|
|
let token = *tokens
|
|
.get(cursor)
|
|
.ok_or("unterminated DeepSeek image block")?;
|
|
if token <= vocab + START || token > vocab + END {
|
|
return Err("malformed DeepSeek image block".into());
|
|
}
|
|
if token == vocab + END {
|
|
break;
|
|
}
|
|
cursor += 1;
|
|
}
|
|
result.push(Span {
|
|
block,
|
|
start,
|
|
end: cursor,
|
|
});
|
|
cursor += 1;
|
|
}
|
|
Ok(result)
|
|
}
|
|
|
|
pub(super) fn prefill_chunk(tokens: &[i32], vocab: i32, cap: u32) -> Result<usize, String> {
|
|
let mut limit = tokens.len().min(cap as usize);
|
|
if limit == 0 {
|
|
return Err("empty DeepSeek prefill chunk".into());
|
|
}
|
|
for span in spans(tokens, vocab)? {
|
|
if span.block >= limit {
|
|
break;
|
|
}
|
|
if span.end + 1 > limit {
|
|
if span.end + 1 > cap as usize {
|
|
if span.block == 0 {
|
|
return Err("DeepSeek image block exceeds prefill capacity".into());
|
|
}
|
|
limit = span.block;
|
|
break;
|
|
}
|
|
limit = span.end + 1;
|
|
}
|
|
}
|
|
Ok(limit)
|
|
}
|
|
|
|
pub(super) fn attention_bounds(
|
|
tokens: &[i32],
|
|
vocab: i32,
|
|
pos: u32,
|
|
n_raw: u32,
|
|
window: u32,
|
|
) -> Result<Vec<u32>, String> {
|
|
let end = pos
|
|
.checked_add(
|
|
tokens
|
|
.len()
|
|
.try_into()
|
|
.map_err(|_| "vision attention length overflow")?,
|
|
)
|
|
.ok_or("vision attention position overflow")?;
|
|
if tokens.is_empty() || n_raw == 0 || n_raw > end {
|
|
return Err("invalid DeepSeek vision attention extent".into());
|
|
}
|
|
let first = end - n_raw;
|
|
let mut bounds = Vec::with_capacity(tokens.len() * 2);
|
|
for q in 0..tokens.len() as u32 {
|
|
let p = pos + q;
|
|
bounds.extend([
|
|
first.max(if window == 0 {
|
|
0
|
|
} else {
|
|
(p + 1).saturating_sub(window)
|
|
}),
|
|
p.min(end - 1),
|
|
]);
|
|
}
|
|
for span in spans(tokens, vocab)? {
|
|
for q in span.start..=span.end {
|
|
let p = pos + q as u32;
|
|
let left = (q - span.start).min(383) as u32;
|
|
let right = (span.end - q).min(384) as u32;
|
|
let add = if window != 0 && left >= window {
|
|
left - (window - 1)
|
|
} else {
|
|
0
|
|
};
|
|
bounds[2 * q] = first.max(if window == 0 {
|
|
0
|
|
} else {
|
|
p.saturating_sub(window - 1 + add)
|
|
});
|
|
bounds[2 * q + 1] = (p + right).min(end - 1);
|
|
}
|
|
}
|
|
Ok(bounds)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
#[test]
|
|
fn issue100_ds4_image_layout_and_attention_vectors() {
|
|
for (h, w, p, types, perm) in [
|
|
(
|
|
2,
|
|
3,
|
|
0,
|
|
vec![1, 1, 1, 0, 2, 2, 2, 2, 2, 2, 3, 3, 4],
|
|
vec![0, 3, 1, 4, 2, 5],
|
|
),
|
|
(
|
|
3,
|
|
2,
|
|
5,
|
|
vec![1, 1, 0, 2, 2, 2, 2, 3, 3, 2, 1, 2, 1, 3, 1, 4],
|
|
vec![0, 2, 1, 3, 4, 5],
|
|
),
|
|
(1, 1, 3, vec![0, 2, 1, 3, 1, 4], vec![0]),
|
|
] {
|
|
let actual = layout(h, w, p).unwrap();
|
|
assert_eq!(actual.types, types);
|
|
assert_eq!(actual.permutation, perm);
|
|
}
|
|
let tokens = [7, 101, 101, 100, 102, 103, 104, 8, 100, 101, 104, 9];
|
|
assert_eq!(
|
|
spans(&tokens, 100).unwrap(),
|
|
vec![
|
|
Span {
|
|
block: 1,
|
|
start: 3,
|
|
end: 6
|
|
},
|
|
Span {
|
|
block: 8,
|
|
start: 8,
|
|
end: 10
|
|
}
|
|
]
|
|
);
|
|
assert_eq!(prefill_chunk(&tokens, 100, 5).unwrap(), 1);
|
|
assert_eq!(prefill_chunk(&tokens[1..], 100, 6).unwrap(), 6);
|
|
assert!(prefill_chunk(&tokens[1..], 100, 5).is_err());
|
|
assert!(prefill_chunk(&tokens[4..], 100, 8).is_err());
|
|
for malformed in [
|
|
[101, 4, 5, 6],
|
|
[100, 4, 104, 6],
|
|
[100, 100, 104, 6],
|
|
[100, 102, 6, 7],
|
|
[100, 105, 6, 7],
|
|
[-1, 100, 104, 6],
|
|
] {
|
|
assert!(spans(&malformed, 100).is_err());
|
|
}
|
|
assert_eq!(
|
|
attention_bounds(
|
|
&[7, 101, 100, 102, 102, 102, 102, 102, 104, 8],
|
|
100,
|
|
10,
|
|
20,
|
|
4
|
|
)
|
|
.unwrap(),
|
|
vec![
|
|
7, 10, 8, 11, 9, 18, 10, 18, 11, 18, 12, 18, 12, 18, 12, 18, 12, 18, 16, 19
|
|
]
|
|
);
|
|
}
|
|
#[test]
|
|
fn issue100_ds4_preprocess_dimensions_and_limits() {
|
|
let mut rgb = Vec::new();
|
|
for y in 0..9 {
|
|
for x in 0..17 {
|
|
rgb.extend([
|
|
(x * 13 + y * 3) as u8,
|
|
(x * 5 + y * 17) as u8,
|
|
(x * 7 + y * 11) as u8,
|
|
]);
|
|
}
|
|
}
|
|
let p = preprocess(&rgb, 17, 9).unwrap();
|
|
assert_eq!(
|
|
(
|
|
p.grid_width,
|
|
p.grid_height,
|
|
p.content_width,
|
|
p.content_height
|
|
),
|
|
(38, 20, 529, 280)
|
|
);
|
|
assert_eq!(p.values.len(), 760 * 588);
|
|
assert!(
|
|
p.values
|
|
.iter()
|
|
.all(|v| v.is_finite() && (-1.0..=1.0).contains(v))
|
|
);
|
|
for (h, w) in [
|
|
(1, 1),
|
|
(1, 8),
|
|
(1, 16384),
|
|
(16384, 1),
|
|
(384, 384),
|
|
(1024, 1024),
|
|
(1080, 1920),
|
|
(4096, 4096),
|
|
] {
|
|
let (bh, bw) = resize_shape(h, w).unwrap();
|
|
assert!(grid_tokens(bh, bw) <= 381);
|
|
for start in 0..4 {
|
|
assert!(
|
|
layout((bh / 14).div_ceil(3), (bw / 14).div_ceil(3), start)
|
|
.unwrap()
|
|
.types
|
|
.len()
|
|
<= 384
|
|
);
|
|
}
|
|
}
|
|
assert!(preprocess(&[], 0, 1).is_err());
|
|
assert!(layout(384, 384, 0).is_err());
|
|
}
|
|
}
|