hugo/DS4Server
1
0 Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

Add Qwen vision loading and Metal inference

Browse Source
This commit is contained in:
Georg Bauer committed 2026-09-11 18:32:41 +02:00
1 parent 48c2f751b4
commit e32954ac1b
30 files changed
+1687 -100
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/engine/metal/qwen_mtplx.rs
+2
View File
@@ -70,6 +70,8 @@ mod stream;
mod submission;
mod unary;
mod views;
mod vision;
mod vision_kernels;
mod weights;
#[cfg(test)]
pub(super) use encoder::Scope;
src/engine/metal/qwen_mtplx/compile_graph.rs
+1
View File
@@ -549,6 +549,7 @@ fn equivalent_operation(a: &Operation, b: &Operation, x: &Array, y: &Array) -> b
| (Softmax(a), Softmax(b))
| (SearchSorted(a), SearchSorted(b))
| (CumsumLastF32(a), CumsumLastF32(b)) => a == b,
(Vision(a), Vision(b)) => a == b,
(RmsNorm(a), RmsNorm(b)) => a == b,
(
SdpaVector {
src/engine/metal/qwen_mtplx/dense.rs
+60 -9
View File
@@ -181,6 +181,15 @@ fn dot(a: &Array, b: &Array, out: &Array, k: u32) -> Result<(), String> {
}
pub(super) fn evaluate(a: &Array, b: &Array, out: &Array) -> Result<(), String> {
evaluate_bias(a, b, None, out)
}
pub(super) fn evaluate_bias(
a: &Array,
b: &Array,
bias: Option<&Array>,
out: &Array,
) -> Result<(), String> {
let dtype = out.layout().dtype();
let dtype_name = dtype.kernel_name();
if a.layout().size() == 0 || b.layout().size() == 0 {
@@ -247,6 +256,7 @@ pub(super) fn evaluate(a: &Array, b: &Array, out: &Array) -> Result<(), String>
}
let ab = a.buffer();
let bb = b.buffer();
let cb = bias.map(Array::buffer);
let ob = out.buffer();
let ai = |index| {
tensor(index, &ab)
@@ -480,15 +490,33 @@ pub(super) fn evaluate(a: &Array, b: &Array, out: &Array) -> Result<(), String>
},
[32, wn, wm],
)?;
let zero = 0_i32;
let one = 1_i32;
let scale = 1_f32;
let mut accum = vec![
tensor(0, &intermediate).input(parts as usize * stride as usize),
output(1),
bytes(2, &parts),
bytes(3, &stride),
bytes(4, &n),
];
if let (Some(c), Some(cb)) = (bias, cb.as_ref()) {
accum.extend([
tensor(5, cb)
.at_byte_offset(c.offset())
.input(c.data_size()),
bytes(6, &zero),
bytes(7, &one),
bytes(8, &scale),
bytes(9, &scale),
]);
}
dispatch_geometry(
&format!("steel_gemm_splitk_accum_{dtype_name}_float32"),
&[
tensor(0, &intermediate).input(parts as usize * stride as usize),
output(1),
bytes(2, &parts),
bytes(3, &stride),
bytes(4, &n),
],
&format!(
"steel_gemm_splitk_accum_{dtype_name}_float32{}",
if bias.is_some() { "_axbpy" } else { "" }
),
&accum,
&[],
[n, m, 1],
super::block_dims([n, m, 1]),
@@ -520,6 +548,29 @@ pub(super) fn evaluate(a: &Array, b: &Array, out: &Array) -> Result<(), String>
padding: 0,
};
let mut bindings = vec![ai(0), bi(1), output(3), bytes(4, &params)];
#[repr(C)]
struct Add {
ldc: i32,
fdc: i32,
stride: i64,
alpha: f32,
beta: f32,
}
let add = Add {
ldc: 0,
fdc: 1,
stride: 0,
alpha: 1.,
beta: 1.,
};
if let (Some(c), Some(cb)) = (bias, cb.as_ref()) {
bindings.extend([
tensor(2, cb)
.at_byte_offset(c.offset())
.input(c.data_size()),
bytes(5, &add),
]);
}
let batch_strides = a_batch.iter().chain(&b_batch).copied().collect::<Vec<_>>();
if batch_ndim > 1 {
bindings.extend([
@@ -535,7 +586,7 @@ pub(super) fn evaluate(a: &Array, b: &Array, out: &Array) -> Result<(), String>
&bindings,
&constants(&[
(10, batch_ndim > 1),
(100, false),
(100, bias.is_some()),
(110, false),
(200, m.is_multiple_of(64)),
(201, n.is_multiple_of(bn)),
src/engine/metal/qwen_mtplx/execution.rs
+2
View File
@@ -216,6 +216,7 @@ impl Options {
pub(super) struct Execution {
pub(super) model: BoundModel,
pub(super) vision_prompt: Option<(Vec<i32>, Vec<crate::engine::qwen_vision::Identity>)>,
pub(super) cache: Vec<Option<Cache>>,
pub(super) mtp_cache: [Option<Cache>; 1],
pub(super) streams: Streams,
@@ -324,6 +325,7 @@ impl Execution {
let mtp_cache = model.borrow_dependent().make_mtp_cache(4);
Ok(Self {
model,
vision_prompt: None,
cache,
mtp_cache,
streams,
src/engine/metal/qwen_mtplx/fused.rs
+1
View File
@@ -47,6 +47,7 @@ pub(super) fn operator(op: &Operation) -> Option<&'static str> {
U::Rsqrt => "Rsqrt",
U::Sin => "Sin",
U::Cos => "Cos",
U::Erf => "Erf",
},
Operation::Binary(op) => match op {
B::Equal => "Equal",
src/engine/metal/qwen_mtplx/model.rs
+29 -3
View File
@@ -89,6 +89,7 @@ pub(super) struct TextModel<'a> {
pub(super) mixer: hyper::GatedResidual<'a>,
pub(super) head: Option<Linear<'a>>, // None means tied embedding weights.
pub(super) mtp: Option<super::mtp::Mtp<'a>>,
pub(super) vision: Option<super::vision::Input>,
pub(super) last_widened: Option<Array>,
policy: CompilePolicy,
decode_runs: Option<Vec<Run>>,
@@ -111,7 +112,7 @@ pub(super) struct Output {
pub(super) hidden: Option<Array>,
}
fn embedding(layer: Linear<'_>, ids: &Array, stream: Stream) -> Result<Array, String> {
pub(super) fn embedding(layer: Linear<'_>, ids: &Array, stream: Stream) -> Result<Array, String> {
match layer {
Linear::Dense(weight) => indexing::take_rows(weight, ids, stream),
Linear::Quantized(layer) => linear::dequantize(
@@ -171,6 +172,7 @@ impl<'a> TextModel<'a> {
head,
parameter_order: Vec::new(),
mtp: None,
vision: None,
last_widened: None,
policy,
decode_runs: None,
@@ -213,6 +215,8 @@ impl<'a> TextModel<'a> {
streams: &Streams,
stream: Stream,
) -> Result<Output, String> {
let image = self.vision.clone();
let vision = vision.or_else(|| image.as_ref().map(|v| (Some(&v.table), v.delta)));
let emb = embedding(self.embedding, ids, stream)?;
let mtp = self
.mtp
@@ -242,7 +246,18 @@ impl<'a> TextModel<'a> {
streams: &Streams,
stream: Stream,
) -> Result<Array, String> {
let emb = match input_embeddings {
let image = self.vision.clone();
let vision = vision.or_else(|| image.as_ref().map(|v| (Some(&v.table), v.delta)));
let spliced = match (&image, input_embeddings) {
(Some(v), None) => v.embeddings(
self.embedding,
ids,
super::vision::cache_position(cache.as_deref()) + 1,
stream,
)?,
_ => None,
};
let emb = match input_embeddings.or(spliced.as_ref()) {
Some(emb) => emb.clone(),
None => embedding(self.embedding, ids, stream)?,
};
@@ -440,7 +455,18 @@ impl<'a> TextModel<'a> {
streams: &Streams,
stream: Stream,
) -> Result<Output, String> {
let h = match input_embeddings {
let image = self.vision.clone();
let vision = vision.or_else(|| image.as_ref().map(|v| (Some(&v.table), v.delta)));
let spliced = match (&image, input_embeddings) {
(Some(v), None) => v.embeddings(
self.embedding,
ids,
super::vision::cache_position(cache.as_deref()),
stream,
)?,
_ => None,
};
let h = match input_embeddings.or(spliced.as_ref()) {
Some(h) => h.clone(),
None => embedding(self.embedding, ids, stream)?,
};
src/engine/metal/qwen_mtplx/ops.rs
+5 -1
View File
@@ -6,6 +6,7 @@ use super::stream::{Device, Stream, Streams};
use super::{eval, views};
pub(super) enum Operation {
Vision(super::vision_kernels::Kernel),
Compiled(super::fused::Expression),
Load(super::load::Load),
Reshape,
@@ -208,7 +209,7 @@ impl Operation {
| Self::ArangeU32
| Self::ArangeF32 { .. }
| Self::ArangeI64 => Err("invalid indexing arity".into()),
Self::RmsNorm(_) => Err("RMSNorm requires two inputs".into()),
Self::Vision(_) | Self::RmsNorm(_) => Err("RMSNorm requires two inputs".into()),
Self::QuantizedLinearBf16 { .. } => Err("quantized linear requires four inputs".into()),
Self::GatherQmmBf16 { .. } => Err("GatherQMM requires six inputs".into()),
Self::AffineQuantize { .. } => Err("quantize requires three sibling outputs".into()),
@@ -257,6 +258,9 @@ pub(super) fn evaluate(
.operation()
.ok_or("array has no executable primitive")?;
match &*operation {
Operation::Vision(k) => {
return super::vision_kernels::evaluate(*k, inputs, &outputs[0]);
}
Operation::Compiled(expression) => {
if stream.device() != Device::Gpu {
return Err("compiled expression requires Metal".into());
src/engine/metal/qwen_mtplx/owner.rs
+174
View File
@@ -15,6 +15,9 @@ use std::sync::atomic::AtomicBool;
pub(in crate::engine) struct QwenExecutor {
// All snapshot arrays must die before the executor's native Metal context.
live: LiveSession,
vision_tower: Option<super::vision::Tower>,
vision_path: Option<std::path::PathBuf>,
keep_vision_loaded: bool,
metadata: QwenMetadata,
identity: Identity,
policy: TurnOptions<'static>,
@@ -43,6 +46,13 @@ impl QwenExecutor {
let (execution, policy) = Execution::open(settings)?;
Ok(Self {
live: LiveSession::default(),
vision_tower: None,
vision_path: settings
.artifacts
.vision
.as_ref()
.map(|_| settings.artifacts.model.clone()),
keep_vision_loaded: settings.speculative.keep_vision_loaded,
metadata,
identity,
policy,
@@ -67,6 +77,43 @@ impl QwenExecutor {
if settings.context_tokens != self.identity.context as i32 {
return Err("Qwen turn context differs from the loaded engine".into());
}
self.execution.vision_prompt = None;
self.execution
.model
.with_dependent_mut(|_, model| model.vision = None);
if messages
.iter()
.any(|m| m.content.contains(crate::engine::VISION_DATA_START))
{
let root = self
.vision_path
.as_ref()
.ok_or("Qwen vision encoder is not installed; download it in Model Manager")?;
if self.vision_tower.is_none() {
self.vision_tower = Some(super::vision::Tower::load(
root,
&self.execution.streams,
self.execution.stream,
)?);
}
let prepared = super::vision::prepare(
cancelled,
&self.metadata,
messages,
settings,
self.vision_tower.as_ref().unwrap(),
&self.execution.streams,
self.execution.stream,
);
if !self.keep_vision_loaded {
self.vision_tower = None;
}
let (ids, input, identities) = prepared?;
self.execution
.model
.with_dependent_mut(|_, model| model.vision = Some(input));
self.execution.vision_prompt = Some((ids, identities));
}
let baseline = self.stats;
let stats = &mut self.stats;
let result = self.live.generate(
@@ -163,6 +210,10 @@ impl QwenExecutor {
fn clear_execution_cache(&mut self) -> Result<(), String> {
self.execution.streams.synchronize(self.execution.stream)?;
self.execution.vision_prompt = None;
self.execution
.model
.with_dependent_mut(|_, model| model.vision = None);
let model = self.execution.model.borrow_dependent();
self.execution.cache = model.make_cache(4);
self.execution.mtp_cache = model.make_mtp_cache(4);
@@ -866,3 +917,126 @@ fn qwen_product_generator_chat_matches_reference() {
std::panic::resume_unwind(error);
}
}
#[test]
#[ignore = "requires installed Qwen, Metal, and DS4_QWEN_VISION_IMAGE; run under test-supervisor"]
fn qwen_vision_chat_checkpoint_preserves_image_identity() {
use crate::settings::{
GenerationPreferences, ReasoningMode, RuntimePreferences, SpeculativePreferences,
};
use base64::Engine as _;
let root = std::path::PathBuf::from(std::env::var("DS4_QWEN38_ARTIFACTS").unwrap());
let image = std::fs::read(std::env::var("DS4_QWEN_VISION_IMAGE").unwrap()).unwrap();
let directory = std::env::temp_dir().join(format!("qwen-vision-chat-{}", std::process::id()));
std::fs::create_dir_all(&directory).unwrap();
let settings = crate::settings::effective_settings(
crate::model::ModelChoice::Qwen38FlashNext,
&GenerationPreferences {
context_tokens: 8192,
max_generated_tokens: 1024,
system_prompt: String::new(),
temperature: Some(0.),
top_p: Some(0.95),
min_p: Some(0.),
seed: Some(1),
reasoning_mode: ReasoningMode::Low,
},
&RuntimePreferences {
speculative: SpeculativePreferences {
glm_mtp: true,
..Default::default()
},
..Default::default()
},
root.parent().unwrap(),
)
.unwrap();
eprintln!("vision lifecycle: loading Qwen");
let mut executor = QwenExecutor::open(&settings.engine).unwrap();
let metrics = Metrics::new(&directory);
let cancelled = AtomicBool::new(false);
let user = |content: String| ChatTurn {
user: true,
tool: false,
system: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: false,
content,
};
let marker = crate::engine::vision_data_marker(&format!(
"data:image/png;base64,{}",
base64::engine::general_purpose::STANDARD.encode(image)
));
let mut messages = vec![user(format!("Describe this image{marker}"))];
let run = |executor: &mut QwenExecutor, messages: &[ChatTurn], label: &str| {
let (output, complete) = executor
.generate(
messages,
&settings.turn,
&cancelled,
&metrics,
|p, _, _| {
if p % 128 == 0 {
eprintln!("{label}: position {p}");
}
},
|_, _| {},
)
.unwrap();
assert!(complete);
assert_eq!(output.finish_reason, "stop");
println!(
"{}",
serde_json::json!({"event":"vision_lifecycle","phase":label,"content":output.message.content,"reasoning":output.message.reasoning,"prompt":output.prompt_tokens,"cached":output.cached_tokens})
);
output
};
let first = run(&mut executor, &messages, "cold");
assert_eq!(first.prompt_tokens, 1069);
assert_eq!(first.cached_tokens, 0);
let content = first.message.content.to_lowercase();
assert!(content.contains("llama") || content.contains("alpaca"));
assert!(content.contains("mountain"));
assert!(
executor.vision_tower.is_none(),
"default policy must release vision weights"
);
assert_eq!(executor.live.current_images.len(), 1);
let identities = executor.live.current_images.clone();
let repeat = run(&mut executor, &messages, "same_image");
assert_eq!(repeat.cached_tokens, first.prompt_tokens);
assert_eq!(repeat.message.content, first.message.content);
messages.push(repeat.message);
let tag = crate::engine::conversation_tag("", ReasoningMode::Low, &messages);
let checkpoint = directory.join("vision.kv");
executor
.save_checkpoint(&checkpoint, tag, &mut |_| {})
.unwrap();
executor.reset().unwrap();
assert!(executor.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
assert_eq!(executor.live.current_images, identities);
messages.push(user("Describe this image".into()));
let followup = run(&mut executor, &messages, "restored_followup");
assert!(followup.cached_tokens >= first.prompt_tokens);
let text = followup.message.content.to_lowercase();
assert!(text.contains("llama") || text.contains("alpaca"));
// Same text and same image-pad span, different pixels must miss the old KV.
let blue = image::RgbImage::from_pixel(1024, 1024, image::Rgb([0, 0, 255]));
let mut encoded = std::io::Cursor::new(Vec::new());
blue.write_to(&mut encoded, image::ImageFormat::Png)
.unwrap();
let marker = crate::engine::vision_data_marker(&format!(
"data:image/png;base64,{}",
base64::engine::general_purpose::STANDARD.encode(encoded.get_ref())
));
let changed = run(
&mut executor,
&[user(format!("Describe this image{marker}"))],
"different_image",
);
assert_eq!(changed.cached_tokens, 0);
assert!(changed.message.content.to_lowercase().contains("blue"));
assert_ne!(changed.message.content, first.message.content);
std::fs::remove_dir_all(directory).unwrap();
}
src/engine/metal/qwen_mtplx/request.rs
+20 -6
View File
@@ -95,6 +95,8 @@ pub(in crate::engine) struct LiveSession {
pub(super) current: Option<SessionSnapshot>,
pub(super) prefix: Option<SessionSnapshot>,
pub(super) tag: [u8; 32],
pub(super) current_images: Vec<crate::engine::qwen_vision::Identity>,
pub(super) prefix_images: Vec<crate::engine::qwen_vision::Identity>,
}
impl LiveSession {
@@ -129,7 +131,10 @@ impl LiveSession {
emit: impl FnMut(bool, String),
observe: impl FnMut(TurnProgress) -> Result<(), String>,
) -> Result<TextTurn, String> {
let prompt = self.prompt(model, messages, settings);
let (prompt, images) = execution
.vision_prompt
.clone()
.unwrap_or_else(|| (self.prompt(model, messages, settings), Vec::new()));
let context = u32::try_from(settings.context_tokens)
.ok()
.filter(|&n| n > 0)
@@ -154,11 +159,18 @@ impl LiveSession {
o.stops = &stop_ids;
}
}
let saved = [&self.current, &self.prefix]
.into_iter()
.flatten()
.filter(|s| prompt.starts_with(&s.token_ids))
.max_by_key(|s| s.token_ids.len());
let saved = [
(&self.current, &self.current_images),
(&self.prefix, &self.prefix_images),
]
.into_iter()
.filter_map(|(s, old)| {
s.as_ref().filter(|s| {
crate::engine::qwen_vision::same_prefix(old, &images, s.token_ids.len())
})
})
.filter(|s| prompt.starts_with(&s.token_ids))
.max_by_key(|s| s.token_ids.len());
let mut captured_prefix = None;
let result = execution.generate_text(
&prompt,
@@ -182,6 +194,7 @@ impl LiveSession {
// errors. Never tag those partial turns as a completed conversation.
if let Some(prefix) = captured_prefix {
self.prefix = Some(prefix);
self.prefix_images = images.clone();
}
let mut result = result?;
if let Some(snapshot) = result
@@ -194,6 +207,7 @@ impl LiveSession {
self.tag =
conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed);
self.current = Some(snapshot);
self.current_images = images;
}
Ok(result)
}
src/engine/metal/qwen_mtplx/session_checkpoint.rs
+7
View File
@@ -27,6 +27,8 @@ pub(super) struct Identity {
struct Manifest {
identity: Identity,
tag: [u8; 32],
#[serde(default)]
images: Vec<crate::engine::qwen_vision::Identity>,
tokens: Vec<i32>,
lazy_kv: bool,
payload: serde_json::Value,
@@ -111,6 +113,7 @@ impl LiveSession {
let manifest = serde_json::to_vec(&Manifest {
identity,
tag,
images: self.current_images.clone(),
tokens: saved.token_ids.clone(),
lazy_kv: saved.lazy_kv,
payload,
@@ -230,6 +233,8 @@ impl LiveSession {
self.current = Some(saved);
self.prefix = None;
self.tag = manifest.tag;
self.current_images = manifest.images;
self.prefix_images.clear();
Ok(true)
}
}
@@ -258,6 +263,8 @@ fn qwen_checkpoint_roundtrip_and_failed_restore_preserve_live_state() {
.unwrap();
let state = Some(State::Recurrent(vec![Some(leaf.clone()), None]));
let mut live = LiveSession {
current_images: Vec::new(),
prefix_images: Vec::new(),
current: Some(SessionSnapshot {
token_ids: vec![1, 2, 3, 4],
logits: leaf.clone(),
src/engine/metal/qwen_mtplx/unary.rs
+1 -1
View File
@@ -15,12 +15,12 @@ pub(super) enum Unary {
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
Log,
Sigmoid,
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
Tanh,
Sqrt,
Rsqrt,
Sin,
Cos,
Erf,
}
pub(super) fn unary(input: &Array, op: Unary, stream: Stream) -> Result<Array, String> {
src/engine/metal/qwen_mtplx/vision.rs
+544
View File
@@ -0,0 +1,544 @@
//! Qwen3-VL tower, following MTPLX e652d55 vision/qwen3_vl_tower.py.
use super::array::{Array, Dtype};
use super::binary::{Binary, binary};
use super::stream::{Stream, Streams};
use super::{
indexing, ops,
unary::{Unary, unary},
vision_kernels::{self, Kernel},
};
use std::path::Path;
pub(super) fn f32_array(shape: &[i32], data: &[f32]) -> Result<Array, String> {
Array::new(
shape,
Dtype::F32,
super::scalar_buffer(
&data
.iter()
.flat_map(|x| x.to_le_bytes())
.collect::<Vec<_>>(),
)?,
)
}
fn i32_array(shape: &[i32], data: &[i32]) -> Result<Array, String> {
Array::new(
shape,
Dtype::I32,
super::scalar_buffer(
&data
.iter()
.flat_map(|x| x.to_le_bytes())
.collect::<Vec<_>>(),
)?,
)
}
fn scalar(x: f32, dtype: Dtype, s: Stream) -> Result<Array, String> {
ops::astype(&f32_array(&[], &[x])?, dtype, false, s)
}
fn add(a: &Array, b: &Array, s: Stream) -> Result<Array, String> {
binary(a, b, Binary::Add, s)
}
fn mul(a: &Array, b: &Array, s: Stream) -> Result<Array, String> {
binary(a, b, Binary::Multiply, s)
}
fn pad_last(a: &Array, extra: i32, s: Stream) -> Result<Array, String> {
let mut shape = a.layout().shape().to_vec();
*shape.last_mut().unwrap() = extra;
ops::concatenate(
&[a.clone(), ops::zeros(&shape, a.layout().dtype(), s)?],
-1,
s,
)
}
fn contiguous(a: &Array, s: Stream) -> Result<Array, String> {
ops::contiguous(a, false, s)
}
pub(super) struct Tower {
weights: super::weights::Parameters,
}
impl Tower {
pub(super) fn load(root: &Path, streams: &Streams, s: Stream) -> Result<Self, String> {
let weights = super::load::safetensors(&root.join("model-vision.safetensors"), streams)?;
if weights.len() != 333 || weights.keys().any(|k| !k.starts_with("vision_tower.")) {
return Err("Qwen vision tower must contain its 333 pinned tensors".into());
}
let mut tower = Self { weights };
let key = "vision_tower.patch_embed.proj.weight";
let w = tower
.weights
.get(key)
.ok_or("missing Qwen vision patch weight")?;
let w = match w.layout().shape() {
[1152, 3, 2, 16, 16] => ops::transpose(w, &[0, 2, 3, 4, 1], s)?,
[1152, 2, 16, 16, 3] => w.clone(),
_ => return Err("unexpected Qwen vision patch layout".into()),
};
tower
.weights
.insert(key.into(), contiguous(&pad_last(&w, 13, s)?, s)?);
ops::evaluate(
streams,
&tower.weights.values().cloned().collect::<Vec<_>>(),
s,
false,
)?;
Ok(tower)
}
fn weight(&self, name: &str) -> Result<&Array, String> {
self.weights
.get(&format!("vision_tower.{name}"))
.ok_or_else(|| format!("missing Qwen vision tensor {name}"))
}
fn linear(&self, x: &Array, name: &str, s: Stream) -> Result<Array, String> {
let w = self.weight(&format!("{name}.weight"))?;
let b = self.weight(&format!("{name}.bias"))?;
let rows = x.layout().dim(0)?;
if rows < 16
|| x.layout().dtype() != Dtype::BF16
|| x.layout().shape().len() != 2
|| x.layout().dim(1)? != w.layout().dim(1)?
{
return Err("invalid Qwen vision projection layout".into());
}
vision_kernels::call(
Kernel::Linear,
&[contiguous(x, s)?, ops::transpose(w, &[1, 0], s)?, b.clone()],
&[rows, w.layout().dim(0)?],
s,
)
}
fn norm(&self, x: &Array, name: &str, s: Stream) -> Result<Array, String> {
vision_kernels::call(
Kernel::LayerNorm,
&[
contiguous(x, s)?,
self.weight(&format!("{name}.weight"))?.clone(),
self.weight(&format!("{name}.bias"))?.clone(),
],
x.layout().shape(),
s,
)
}
fn positions(&self, h: i32, w: i32, s: Stream) -> Result<(Array, Array), String> {
let mut indices = [Vec::new(), Vec::new(), Vec::new(), Vec::new()];
let mut weights = [Vec::new(), Vec::new(), Vec::new(), Vec::new()];
let mut rotary = Vec::new();
for br in 0..h / 2 {
for bc in 0..w / 2 {
for ir in 0..2 {
for ic in 0..2 {
let (r, c) = (br * 2 + ir, bc * 2 + ic);
// mx.linspace(0,47,n) uses a double step cast to F32.
let y = (r as f64 * 47.0 / f64::from(h - 1)) as f32;
let x = (c as f64 * 47.0 / f64::from(w - 1)) as f32;
let (yf, xf) = (y as i32, x as i32);
let (yc, xc) = ((yf + 1).min(47), (xf + 1).min(47));
let (dy, dx) = (y - yf as f32, x - xf as f32);
for (i, (idx, v)) in [
(yf * 48 + xf, (1. - dy) * (1. - dx)),
(yf * 48 + xc, (1. - dy) * dx),
(yc * 48 + xf, dy * (1. - dx)),
(yc * 48 + xc, dy * dx),
]
.into_iter()
.enumerate()
{
indices[i].push(idx);
weights[i].push(v)
}
rotary.extend([r as f32, c as f32]);
}
}
}
}
let mut parts = Vec::new();
for i in 0..4 {
let p = indexing::take_rows(
self.weight("pos_embed.weight")?,
&i32_array(&[h * w], &indices[i])?,
s,
)?;
let weights =
ops::astype(&f32_array(&[h * w, 1], &weights[i])?, Dtype::BF16, false, s)?;
parts.push(mul(&p, &weights, s)?);
}
let position = add(
&add(&add(&parts[0], &parts[1], s)?, &parts[2], s)?,
&parts[3],
s,
)?;
let exponent = binary(
&f32_array(&[18], &(0..18).map(|i| (i * 2) as f32).collect::<Vec<_>>())?,
&f32_array(&[], &[36.])?,
Binary::Divide,
s,
)?;
let powers = binary(&f32_array(&[], &[10000.])?, &exponent, Binary::Power, s)?;
let inv = binary(&f32_array(&[], &[1.])?, &powers, Binary::Divide, s)?;
let rotary = ops::reshape(
&mul(&f32_array(&[h * w, 2, 1], &rotary)?, &inv, s)?,
&[h * w, 36],
s,
)?;
Ok((position, rotary))
}
fn attention(&self, x: &Array, name: &str, rotary: &Array, s: Stream) -> Result<Array, String> {
let n = x.layout().dim(0)?;
let qkv = self.linear(x, &format!("{name}.qkv"), s)?;
let qkv = ops::transpose(&ops::reshape(&qkv, &[n, 3, 16, 72], s)?, &[1, 0, 2, 3], s)?;
let mut parts = ops::split(&qkv, &[1, 2], 0, s)?;
let angles = ops::reshape(rotary, &[1, n, 1, 36], s)?;
let cos = unary(&angles, Unary::Cos, s)?;
let sin = unary(&angles, Unary::Sin, s)?;
let cos = ops::concatenate(&[cos.clone(), cos], -1, s)?;
let sin = ops::concatenate(&[sin.clone(), sin], -1, s)?;
for p in &mut parts[..2] {
let halves = ops::split(p, &[36], -1, s)?;
let rotated = ops::concatenate(
&[unary(&halves[1], Unary::Negative, s)?, halves[0].clone()],
-1,
s,
)?;
*p = ops::astype(
&add(&mul(p, &cos, s)?, &mul(&rotated, &sin, s)?, s)?,
Dtype::BF16,
false,
s,
)?;
}
let inputs = parts
.iter()
.map(|p| contiguous(&pad_last(&ops::transpose(p, &[0, 2, 1, 3], s)?, 8, s)?, s))
.collect::<Result<Vec<_>, _>>()?;
let output = vision_kernels::call(Kernel::Attention, &inputs, &[1, 16, n, 80], s)?;
let output = ops::slice(&output, &[0, 0, 0, 0], &[1, 16, n, 72], &[1; 4], s)?;
let output = ops::reshape(&ops::transpose(&output, &[0, 2, 1, 3], s)?, &[n, 1152], s)?;
self.linear(&output, &format!("{name}.proj"), s)
}
pub(super) fn encode(
&self,
pixels: &Array,
h: i32,
w: i32,
streams: &Streams,
s: Stream,
mut observe: impl FnMut(&str, &Array) -> Result<(), String>,
) -> Result<Array, String> {
if h < 2 || w < 2 || h % 2 != 0 || w % 2 != 0 || pixels.layout().shape() != [h * w, 1536] {
return Err("invalid Qwen vision patch grid".into());
}
let pixels = ops::astype(pixels, Dtype::BF16, false, s)?;
let patches = ops::transpose(
&ops::reshape(&pixels, &[h * w, 3, 2, 16, 16], s)?,
&[0, 2, 3, 4, 1],
s,
)?;
let patches = contiguous(&pad_last(&patches, 13, s)?, s)?;
let mut x = vision_kernels::call(
Kernel::PatchConv,
&[patches, self.weight("patch_embed.proj.weight")?.clone()],
&[h * w, 1152],
s,
)?;
x = add(&x, self.weight("patch_embed.proj.bias")?, s)?;
observe("patch", &x)?;
let (position, rotary) = self.positions(h, w, s)?;
observe("position", &position)?;
observe("rotary", &rotary)?;
x = add(&x, &position, s)?;
for i in 0..27 {
observe("progress", &x)?;
let name = format!("blocks.{i}");
let a = self.attention(
&self.norm(&x, &format!("{name}.norm1"), s)?,
&format!("{name}.attn"),
&rotary,
s,
)?;
x = add(&x, &a, s)?;
let mlp = self.linear(
&self.norm(&x, &format!("{name}.norm2"), s)?,
&format!("{name}.mlp.linear_fc1"),
s,
)?;
let mlp = gelu(&mlp, true, s)?;
x = add(
&x,
&self.linear(&mlp, &format!("{name}.mlp.linear_fc2"), s)?,
s,
)?;
ops::evaluate(streams, std::slice::from_ref(&x), s, false)?;
if i == 0 || i == 26 {
observe(&format!("block{i}"), &x)?;
}
}
let x = ops::reshape(&self.norm(&x, "merger.norm", s)?, &[h * w / 4, 4608], s)?;
let x = gelu(&self.linear(&x, "merger.linear_fc1", s)?, false, s)?;
let x = self.linear(&x, "merger.linear_fc2", s)?;
ops::evaluate(streams, std::slice::from_ref(&x), s, false)?;
observe("embeddings", &x)?;
Ok(x)
}
}
fn gelu(x: &Array, approx: bool, s: Stream) -> Result<Array, String> {
let inputs = std::slice::from_ref(x);
let graph = super::compiled::Replay::trace_compiled(inputs, |args| {
let x = &args[0];
let dtype = x.layout().dtype();
let c = |v| scalar(v, dtype, s);
let result = if approx {
let cubic = binary(x, &c(3.)?, Binary::Power, s)?;
let inner = mul(
&c((2_f64 / std::f64::consts::PI).sqrt() as f32)?,
&add(x, &mul(&c(0.044715)?, &cubic, s)?, s)?,
s,
)?;
mul(
&mul(&c(0.5)?, x, s)?,
&add(&c(1.)?, &unary(&inner, Unary::Tanh, s)?, s)?,
s,
)?
} else {
let erf = unary(
&binary(x, &c(2_f64.sqrt() as f32)?, Binary::Divide, s)?,
Unary::Erf,
s,
)?;
binary(
&mul(x, &add(&c(1.)?, &erf, s)?, s)?,
&c(2.)?,
Binary::Divide,
s,
)?
};
Ok(vec![result])
})?;
Ok(graph.run(inputs)?.remove(0))
}
#[test]
#[ignore = "requires Metal and pinned Qwen vision artifacts/reference values"]
fn qwen_vision_tower_matches_mtplx_image() {
super::configure_sources().unwrap();
let _context = super::super::gpu::Context::open_qwen(0).unwrap();
let streams = Streams::new(Some(super::allocator::Allocator::new().unwrap()));
let s = streams.default_stream(super::stream::Device::Gpu).unwrap();
let root = std::path::PathBuf::from(std::env::var("DS4_QWEN38_ARTIFACTS").unwrap());
let dir = std::path::PathBuf::from(std::env::var("DS4_QWEN_VISION_REFERENCE").unwrap());
let read = |name: &str| {
std::fs::read(dir.join(format!("{name}.f32")))
.unwrap()
.chunks_exact(4)
.map(|b| f32::from_le_bytes(b.try_into().unwrap()))
.collect::<Vec<_>>()
};
let grid: Vec<i32> =
serde_json::from_slice(&std::fs::read(dir.join("grid.json")).unwrap()).unwrap();
let (h, w) = (grid[1], grid[2]);
let pixels = f32_array(&[h * w, 1536], &read("pixels")).unwrap();
let tower = Tower::load(&root, &streams, s).unwrap();
tower
.encode(&pixels, h, w, &streams, s, |name, a| {
if name == "progress" {
return Ok(());
}
let a = ops::astype(a, Dtype::F32, false, s)?;
ops::evaluate(&streams, std::slice::from_ref(&a), s, false)?;
let got = ops::read_f32_as_f64(&a, &streams, s)?;
let expected = read(name);
let max = got
.iter()
.zip(&expected)
.map(|(a, b)| (a - f64::from(*b)).abs())
.fold(0., f64::max);
let different = got
.iter()
.zip(&expected)
.filter(|(a, b)| **a != f64::from(**b))
.count();
eprintln!(
"{name}: values={} different={different} max_error={max}",
got.len()
);
assert_eq!(got.len(), expected.len());
assert_eq!(different, 0, "{name} differs from MTPLX (max error {max})");
Ok(())
})
.unwrap();
}
#[derive(Clone)]
pub(super) struct Input {
pub(super) table: Array,
pub(super) delta: i32,
pub(super) images: Vec<(std::ops::Range<usize>, Array)>,
}
impl Input {
pub(super) fn embeddings(
&self,
layer: super::mlp::Linear<'_>,
ids: &Array,
start: usize,
s: Stream,
) -> Result<Option<Array>, String> {
let n = ids.layout().dim(1)? as usize;
let end = start.checked_add(n).ok_or("Qwen image window overflow")?;
if !self
.images
.iter()
.any(|(r, _)| start < r.end && end > r.start)
{
return Ok(None);
}
let mut embedded = super::model::embedding(layer, ids, s)?;
for (span, rows) in &self.images {
let a = start.max(span.start);
let b = end.min(span.end);
if a >= b {
continue;
}
let rows = ops::slice(
rows,
&[(a - span.start) as i32, 0],
&[(b - span.start) as i32, 2560],
&[1, 1],
s,
)?;
let rows = ops::reshape(&rows, &[1, (b - a) as i32, 2560], s)?;
embedded = ops::slice_update(
&embedded,
&rows,
&[0, (a - start) as i32, 0],
&[1, (b - start) as i32, 2560],
&[1; 3],
s,
)?;
}
Ok(Some(embedded))
}
}
pub(super) fn cache_position(cache: Option<&[Option<super::decoder::Cache>]>) -> usize {
cache
.into_iter()
.flatten()
.flatten()
.find_map(|c| match &c.attention {
super::decoder::AttentionCache::Qsa(c) => Some(c.kv.offset as usize),
_ => None,
})
.unwrap_or(0)
}
pub(super) fn prepare(
cancelled: &std::sync::atomic::AtomicBool,
model: &crate::engine::qwen::QwenMetadata,
messages: &[crate::engine::ChatTurn],
settings: &crate::settings::TurnSettings,
tower: &Tower,
streams: &Streams,
s: Stream,
) -> Result<(Vec<i32>, Input, Vec<crate::engine::qwen_vision::Identity>), String> {
use crate::engine::{VISION_DATA_END, VISION_DATA_START};
use base64::Engine as _;
use sha2::{Digest, Sha256};
let check = || {
if cancelled.load(std::sync::atomic::Ordering::Relaxed) {
Err("Qwen vision input cancelled".to_owned())
} else {
Ok(())
}
};
check()?;
let mut rendered = messages.to_vec();
let mut decoded = Vec::new();
let mut total = 0;
for message in &mut rendered {
if !message.content.contains(VISION_DATA_START) {
continue;
}
if !message.user && !message.tool {
return Err("vision input is allowed only in user or tool messages".into());
}
let mut content = String::new();
let mut rest = message.content.as_str();
while let Some(start) = rest.find(VISION_DATA_START) {
content.push_str(&rest[..start]);
let image = &rest[start + VISION_DATA_START.len()..];
let end = image
.find(VISION_DATA_END)
.ok_or("unterminated image input")?;
let uri = &image[..end];
let payload = uri
.strip_prefix("data:image/png;base64,")
.or_else(|| uri.strip_prefix("data:image/jpeg;base64,"))
.ok_or("image input must be an inline PNG or JPEG data URI")?;
if decoded.len() >= 16 || payload.len() > 90 * 1024 * 1024 {
return Err("image request exceeds limits".into());
}
check()?;
let raw = base64::engine::general_purpose::STANDARD
.decode(payload)
.map_err(|_| "invalid image base64")?;
total += raw.len();
if total > 64 * 1024 * 1024 {
return Err("image inputs exceed the 64 MiB request limit".into());
}
let digest: [u8; 32] = Sha256::digest(&raw).into();
let image = crate::engine::qwen_vision::decode(&raw)?;
decoded.push((image, digest));
content.push_str("<|vision_start|><|image_pad|><|vision_end|>");
rest = &rest[start + VISION_DATA_START.len() + end + VISION_DATA_END.len()..];
}
content.push_str(rest);
message.content = content;
}
let plain =
model.render_conversation(&settings.system_prompt, &rendered, settings.reasoning_mode);
let mut ids = Vec::new();
let mut spans = Vec::new();
let mut index = 0;
for id in plain {
if id == 248056 {
let (image, digest) = decoded
.get(index)
.ok_or("image pad has no supplied image")?;
let start = ids.len();
ids.extend(std::iter::repeat_n(id, (image.h * image.w / 4) as usize));
spans.push(crate::engine::qwen_vision::Identity {
start,
end: ids.len(),
digest: *digest,
});
index += 1;
} else {
ids.push(id)
}
if ids.len() >= settings.context_tokens as usize {
return Err("image prompt exceeds configured context".into());
}
}
if index != decoded.len() {
return Err("rendered Qwen prompt lost an image placeholder".into());
}
let grids = decoded.iter().map(|(i, _)| (i.h, i.w)).collect::<Vec<_>>();
let (table, delta) = crate::engine::qwen_vision::positions(&ids, &grids)?;
let table = i32_array(&[3, ids.len() as i32], &table)?;
let mut images = Vec::new();
for ((image, _), span) in decoded.iter().zip(&spans) {
check()?;
let pixels = f32_array(&[image.h * image.w, 1536], &image.pixels)?;
let rows = tower.encode(&pixels, image.h, image.w, streams, s, |_, _| check())?;
images.push((span.start..span.end, rows));
}
Ok((
ids,
Input {
table,
delta,
images,
},
spans,
))
}
src/engine/metal/qwen_mtplx/vision_kernels.rs
+191
View File
@@ -0,0 +1,191 @@
//! Vision-only dispatches from pinned MLX 0.32.2 normalization, conv and SDPA.
//! Rust supplies the same layouts and parameters to the existing Metal library.
use super::super::gpu::{Buffer, MetalConstant};
use super::array::{Array, Dtype, Layout};
use super::{bytes, dispatch_geometry, dispatch_specialized, ops, stream::Stream, tensor};
#[derive(Clone, Copy, PartialEq)]
pub(super) enum Kernel {
LayerNorm,
PatchConv,
Attention,
Linear,
}
pub(super) fn call(
kernel: Kernel,
inputs: &[Array],
shape: &[i32],
stream: Stream,
) -> Result<Array, String> {
Array::make_operation_with_inputs(
stream,
inputs,
Layout::new(shape, Dtype::BF16)?,
ops::Operation::Vision(kernel),
)
}
pub(super) fn evaluate(kernel: Kernel, inputs: &[Array], out: &Array) -> Result<(), String> {
if kernel == Kernel::Linear {
return super::dense::evaluate_bias(&inputs[0], &inputs[1], Some(&inputs[2]), out);
}
let nbytes = out.layout().nbytes() as u64;
out.set_data(Buffer::mtplx_bytes(nbytes)?)?;
let buffers = inputs.iter().map(Array::buffer).collect::<Vec<_>>();
let ob = out.buffer();
let mut bindings = inputs
.iter()
.zip(&buffers)
.enumerate()
.map(|(i, (a, b))| {
tensor(i as u32, b)
.at_byte_offset(a.offset())
.input(a.data_size())
})
.collect::<Vec<_>>();
bindings.push(tensor(inputs.len() as u32, &ob).output(out.data_size()));
match kernel {
Kernel::Linear => {
super::dense::evaluate_bias(&inputs[0], &inputs[1], Some(&inputs[2]), out)
}
Kernel::LayerNorm => {
let axis = out.layout().dim(-1)? as u32;
let rows = out.layout().size() as u32 / axis;
let threads = axis.div_ceil(8).div_ceil(32) * 32;
let eps = 1e-6_f32;
let stride = 1_u32;
bindings.extend([
bytes(4, &eps),
bytes(5, &axis),
bytes(6, &stride),
bytes(7, &stride),
]);
dispatch_geometry(
"layer_normbfloat16",
&bindings,
&[],
[rows * threads, 1, 1],
[threads, 1, 1],
true,
)
}
Kernel::PatchConv => {
// Inputs are contiguous [patches,2,16,16,16], weights
// [1152,2,16,16,16]. Original Conv3d pads RGB channels to 16.
#[repr(C)]
struct Params {
n: i32,
c: i32,
o: i32,
input: [i32; 3],
weight: [i32; 3],
output: [i32; 3],
stride: [i32; 3],
pad: [i32; 3],
kdil: [i32; 3],
idil: [i32; 3],
is: [i64; 5],
ws: [i64; 5],
os: [i64; 5],
groups: i32,
flip: u8,
padding: [u8; 3],
}
let n = out.layout().dim(0)?;
let params = Params {
n,
c: 16,
o: 1152,
input: [2, 16, 16],
weight: [2, 16, 16],
output: [1, 1, 1],
stride: [2, 16, 16],
pad: [0; 3],
kdil: [1; 3],
idil: [1; 3],
is: [8192, 4096, 256, 16, 1],
ws: [8192, 4096, 256, 16, 1],
os: [1152, 1152, 1152, 1152, 1],
groups: 1,
flip: 0,
padding: [0; 3],
};
let gemm = [n, 1152, 8192, 512, 16, 16, 16, -8160, 18, (n + 31) / 32, 0];
bindings.extend([bytes(3, &params), bytes(4, &gemm)]);
dispatch_specialized(
"implicit_gemm_conv_3d_bfloat16_bm32_bn64_bk16_wm2_wn2_filter_s",
&bindings,
&[],
[18, ((n + 31) / 32) as u32, 1],
[32, 2, 2],
)
}
Kernel::Attention => {
#[repr(C)]
struct Params {
b: i32,
h: i32,
d: i32,
ql: i32,
kl: i32,
gqa: i32,
scale: f32,
nq: i32,
nk: i32,
aq: i32,
ak: i32,
rq: i32,
rk: i32,
off: i32,
strides: [[i64; 3]; 4],
}
let shape = out.layout().shape().to_vec();
let n = shape[2];
let mut strides = [[0; 3]; 4];
for (s, a) in strides
.iter_mut()
.zip(inputs.iter().chain(std::iter::once(out)))
{
s.copy_from_slice(&a.layout().strides()[..3]);
}
let params = Params {
b: 1,
h: 16,
d: 80,
ql: n,
kl: n,
gqa: 1,
scale: (72_f64.powf(-0.5)) as f32,
nq: ((n + 31) / 32),
nk: ((n + 31) / 32),
aq: n / 32,
ak: n / 32,
rq: n % 32,
rk: n % 32,
off: 0,
strides,
};
let constants = [
(200, n % 32 == 0),
(201, n % 32 == 0),
(300, false),
(301, false),
(302, false),
]
.map(|(index, v)| MetalConstant {
index,
value: u32::from(v),
kind: 0,
});
bindings.push(bytes(4, &params));
dispatch_specialized(
"steel_attention_bfloat16_bq32_bk32_bd80_wm4_wn1_maskbfloat16",
&bindings,
&constants,
[((n + 31) / 32) as u32, 16, 1],
[32, 4, 1],
)
}
}
}
src/engine/metal/qwen_mtplx/weights.rs
+4 -1
View File
@@ -34,7 +34,10 @@ pub(super) fn load_shards(
let entry = entry.map_err(|e| e.to_string())?;
let name = entry.file_name();
let Some(name) = name.to_str() else { continue };
if name.starts_with("model") && name.ends_with(".safetensors") {
if name.starts_with("model")
&& name.ends_with(".safetensors")
&& name != "model-vision.safetensors"
{
parameters.extend(super::load::safetensors(&entry.path(), streams)?);
files += 1;
}
src/engine/qwen.rs
+3 -1
View File
@@ -200,7 +200,9 @@ impl QwenMetadata {
)
.ok_or("Qwen artifact size overflows")?;
}
Ok(Self::new(tokenizer, mapped_bytes, tensor_count))
let mut metadata = Self::new(tokenizer, mapped_bytes, tensor_count);
metadata.summary.vision_loaded = root.join("model-vision.safetensors").is_file();
Ok(metadata)
}
}
src/engine/qwen_vision.rs
+284
View File
@@ -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));
}