Add Qwen vision loading and Metal inference

This commit is contained in:
Georg Bauer
2026-09-11 18:32:41 +02:00
parent 48c2f751b4
commit e32954ac1b
30 changed files with 1687 additions and 100 deletions
+89
View File
@@ -0,0 +1,89 @@
# Qwen vision verification — 2026-09-11
Qwen3.8 Flash Next now accepts the existing PNG/JPEG attachments. Its optional
vision encoder has a separate Model Manager entry; downloading, validating or
deleting it does not alter the text-model artifact set. The existing keep-vision-
weights-loaded preference also applies to Qwen. Application preprocessing,
loading and inference are Rust, using the existing Metal runtime kernels.
The four pinned vision files (897,900,287 bytes total) were downloaded and SHA-256
verified in `~/Library/Application Support/de.rfc1437.ds4server/models/qwen3.8-flash-next`.
The artifact revision is `74559cdf34fbfc0b593de72d17e93f37fd4f9ea7` of
`Youssofal/Qwen3.8-Flash-Next-MTPLX-Bare-Speed`; the text manifest remains unchanged.
The video processor configuration is part of that artifact set; this change
implements still images.
## Grounding check
Every primary description run used exactly `Describe this image`. The supplied
1024×1024 PNG was copied unchanged to
`local-eval-results/qwen-vision-input/image.png`. Its SHA-256 is
`b7568d90f4df6180d9af14a824dd553cd995457561d29167354c1ba66b728347`.
Only the image bytes and prompt enter the model; neither the source filename nor
the neutral filename is included in its text input.
Rust/Metal generation succeeded both with MTP and with ordinary autoregressive
decoding. Qwen described a golden-tan cartoon llama/alpaca, large eyes, upright
ears, an open smiling mouth, mountains, a sunset and a grainy poster texture.
These details are visible in the supplied image. The direct cold-session result
begins:
> This is a stylized, cartoon-style illustration of a llama (or alpaca) shown from the neck up, set against a sunset landscape.
Controls used the same prompt:
| Input / session | Observed result |
| --- | --- |
| No image, Rust and MTPLX | Reports no attached image and requests one |
| Solid blue image, same dimensions, after the animal image | Describes a uniform blue field; zero cached prompt tokens |
| Same image repeated | Same description; all 1,069 prompt tokens reused |
| Saved checkpoint, reset, restore, follow-up | Correct animal description; 1,446 cached tokens out of 1,460 |
This demonstrates image-dependent descriptions for these inputs, not a general
guarantee against hallucinations.
## Oracle and reproducibility
The oracle is local MTPLX reference `e652d55` with MLX 0.32.2. Python scripts under
`tools/qwen-vision*-reference.py` run only that reference, never the application.
The Rust tower matches its exported values exactly at patch embedding, position
embedding, rotary positions, blocks 0 and 26, and the final merger. Both the
1024×1024 input (2,621,440 final values) and a small non-square fixture (168,960
final values) had zero differing values. CPU resize/preprocessing also matches
three Pillow/MTPLX golden hashes. This is exact encoder agreement; full generated
token-sequence parity is not claimed.
Local evidence is retained under `local-eval-results/`:
- `qwen-vision-rust-mtp.jsonl`, `qwen-vision-rust-ar.jsonl`,
`qwen-vision-rust-no-image.jsonl`: complete application runs.
- `qwen-vision-lifecycle.jsonl`: cold, repeat, restored and changed-image runs.
- `qwen-vision-chat-reference.jsonl`: oracle image/no-image runs.
- `qwen-vision-image/`, `qwen-vision-small/`: exported oracle arrays.
- `qwen-vision-small-rust.log`: small-fixture exact comparison.
Example application invocation (empty YAML config avoids an unrelated system
prompt):
```sh
target/release/ds4-server model-eval \
--model qwen3.8-flash-next --config /tmp/qwen-vision-config.yaml \
--prompt 'Describe this image' \
--image-file local-eval-results/qwen-vision-input/image.png \
--context 8192 --max-tokens 1024 --reasoning low \
--temperature 0 --top-p 0.95 --seed 1 --acceleration on \
--prefill-chunk 2048 --warmup off --canary on --max-memory-gib 108
```
The ignored GPU tests `qwen_vision_tower_matches_mtplx_image` and
`qwen_vision_chat_checkpoint_preserves_image_identity` are runnable with
`DS4_QWEN38_ARTIFACTS` pointing to the model directory and respectively
`DS4_QWEN_VISION_REFERENCE` pointing to exported arrays or
`DS4_QWEN_VISION_IMAGE` pointing to the neutral input. Run one GPU model process
at a time under `test-supervisor` with an appropriate memory limit.
Final checks: `cargo fmt --all -- --check`, Clippy with all targets/features and
warnings denied, `RUST_TEST_THREADS=1 cargo test --all-features` (315 passed,
204 opt-in tests ignored), and `make bundle` all succeeded. The two encoder
comparisons and image lifecycle test were additionally executed explicitly with
GPU access. The updated, signed application is `target/release/DS4Server.app`.
+1 -1
View File
@@ -1402,7 +1402,7 @@ impl App {
} }
Message::ComposerAction(action) => self.composer.perform(action), Message::ComposerAction(action) => self.composer.perform(action),
Message::ChooseVisionImage => { Message::ChooseVisionImage => {
if self.config.model != ModelChoice::Glm53Flash || self.generating { if !self.config.model.supports_vision() || self.generating {
return Task::none(); return Task::none();
} }
return Task::perform( return Task::perform(
+2 -2
View File
@@ -604,7 +604,7 @@ impl App {
return; return;
} }
self.config = config; self.config = config;
if self.config.model != ModelChoice::Glm53Flash { if !self.config.model.supports_vision() {
self.pending_vision_image = None; self.pending_vision_image = None;
} }
self.context_limit = self.config.active_generation().context_tokens.max(0) as u32; self.context_limit = self.config.active_generation().context_tokens.max(0) as u32;
@@ -889,7 +889,7 @@ impl App {
} }
Message::PreferenceKeepVisionLoadedChanged(value) => { Message::PreferenceKeepVisionLoadedChanged(value) => {
self.preference_draft.keep_vision_loaded = self.preference_draft.keep_vision_loaded =
self.preference_draft.acceleration_model == ModelChoice::Glm53Flash && value; self.preference_draft.acceleration_model.supports_vision() && value;
self.preference_error = None; self.preference_error = None;
} }
Message::PreferenceDsparkConfidenceChanged(value) => { Message::PreferenceDsparkConfidenceChanged(value) => {
+1 -1
View File
@@ -11,7 +11,7 @@ impl App {
self.config = config; self.config = config;
self.preference_draft = PreferenceDraft::from_saved(&self.config); self.preference_draft = PreferenceDraft::from_saved(&self.config);
self.context_limit = self.config.active_generation().context_tokens.max(0) as u32; self.context_limit = self.config.active_generation().context_tokens.max(0) as u32;
if model != ModelChoice::Glm53Flash { if !model.supports_vision() {
self.pending_vision_image = None; self.pending_vision_image = None;
} }
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
+5 -3
View File
@@ -337,8 +337,8 @@ impl App {
.padding([2, 6]) .padding([2, 6])
.into() .into()
}; };
let vision_ready = self.config.model == ModelChoice::Glm53Flash let vision_ready = self.config.model.supports_vision()
&& model::engine_artifacts(ModelChoice::Glm53Flash, false, &models_path()) && model::engine_artifacts(self.config.model, false, &models_path())
.vision .vision
.is_some(); .is_some();
let attach = if vision_ready && !self.generating && self.pending_vision_image.is_none() let attach = if vision_ready && !self.generating && self.pending_vision_image.is_none()
@@ -356,8 +356,10 @@ impl App {
attach, attach,
container(text(if vision_ready { container(text(if vision_ready {
"Attach PNG or JPEG" "Attach PNG or JPEG"
} else if self.config.model.supports_vision() {
"Download the vision encoder in Model Manager"
} else { } else {
"GLM 5.3 Flash vision sidecar is not ready" "This model does not support image attachments"
})) }))
.padding(8) .padding(8)
.style(preference_group_style), .style(preference_group_style),
+2 -2
View File
@@ -25,7 +25,7 @@ impl App {
.supports_integrated_mtp() .supports_integrated_mtp()
.then_some(Message::PreferenceGlmMtpTimingChanged); .then_some(Message::PreferenceGlmMtpTimingChanged);
let keep_vision_loaded_toggle: Option<fn(bool) -> Message> = let keep_vision_loaded_toggle: Option<fn(bool) -> Message> =
(self.preference_draft.acceleration_model == ModelChoice::Glm53Flash) (self.preference_draft.acceleration_model.supports_vision())
.then_some(Message::PreferenceKeepVisionLoadedChanged); .then_some(Message::PreferenceKeepVisionLoadedChanged);
let dspark_strict_toggle: Option<fn(bool) -> Message> = self let dspark_strict_toggle: Option<fn(bool) -> Message> = self
.preference_draft .preference_draft
@@ -623,7 +623,7 @@ impl App {
), ),
hint( hint(
toggle(self.preference_draft.keep_vision_loaded) toggle(self.preference_draft.keep_vision_loaded)
.label("Keep GLM 5.3 vision weights loaded") .label("Keep vision weights loaded")
.on_toggle_maybe(keep_vision_loaded_toggle), .on_toggle_maybe(keep_vision_loaded_toggle),
"Keeps the vision encoder mapped between image turns for lower image latency. Off releases it after encoding all images in a turn, leaving more memory for long text contexts.", "Keeps the vision encoder mapped between image turns for lower image latency. Off releases it after encoding all images in a turn, leaving more memory for long text contexts.",
), ),
+1 -3
View File
@@ -6,6 +6,7 @@ mod kvstore;
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
mod metal; mod metal;
mod qwen; mod qwen;
mod qwen_vision;
#[cfg(test)] #[cfg(test)]
#[path = "../tools/sampler-replay-benchmark.rs"] #[path = "../tools/sampler-replay-benchmark.rs"]
mod sampler_replay_benchmark; mod sampler_replay_benchmark;
@@ -1504,9 +1505,6 @@ impl Generator {
.iter() .iter()
.any(|message| message.content.contains(VISION_DATA_START)); .any(|message| message.content.contains(VISION_DATA_START));
if let metal::Executor::QwenMtplx(executor) = &mut self.executor { if let metal::Executor::QwenMtplx(executor) = &mut self.executor {
if has_vision {
return Err("vision input requires GLM 5.3 Flash".into());
}
return executor.generate(messages, settings, cancelled, &self.metrics, progress, emit); return executor.generate(messages, settings, cancelled, &self.metrics, progress, emit);
} }
let (tokens, overlays) = if has_vision { let (tokens, overlays) = if has_vision {
+2
View File
@@ -70,6 +70,8 @@ mod stream;
mod submission; mod submission;
mod unary; mod unary;
mod views; mod views;
mod vision;
mod vision_kernels;
mod weights; mod weights;
#[cfg(test)] #[cfg(test)]
pub(super) use encoder::Scope; pub(super) use encoder::Scope;
@@ -549,6 +549,7 @@ fn equivalent_operation(a: &Operation, b: &Operation, x: &Array, y: &Array) -> b
| (Softmax(a), Softmax(b)) | (Softmax(a), Softmax(b))
| (SearchSorted(a), SearchSorted(b)) | (SearchSorted(a), SearchSorted(b))
| (CumsumLastF32(a), CumsumLastF32(b)) => a == b, | (CumsumLastF32(a), CumsumLastF32(b)) => a == b,
(Vision(a), Vision(b)) => a == b,
(RmsNorm(a), RmsNorm(b)) => a == b, (RmsNorm(a), RmsNorm(b)) => a == b,
( (
SdpaVector { SdpaVector {
+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> { 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 = out.layout().dtype();
let dtype_name = dtype.kernel_name(); let dtype_name = dtype.kernel_name();
if a.layout().size() == 0 || b.layout().size() == 0 { 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 ab = a.buffer();
let bb = b.buffer(); let bb = b.buffer();
let cb = bias.map(Array::buffer);
let ob = out.buffer(); let ob = out.buffer();
let ai = |index| { let ai = |index| {
tensor(index, &ab) tensor(index, &ab)
@@ -480,15 +490,33 @@ pub(super) fn evaluate(a: &Array, b: &Array, out: &Array) -> Result<(), String>
}, },
[32, wn, wm], [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( dispatch_geometry(
&format!("steel_gemm_splitk_accum_{dtype_name}_float32"), &format!(
&[ "steel_gemm_splitk_accum_{dtype_name}_float32{}",
tensor(0, &intermediate).input(parts as usize * stride as usize), if bias.is_some() { "_axbpy" } else { "" }
output(1), ),
bytes(2, &parts), &accum,
bytes(3, &stride),
bytes(4, &n),
],
&[], &[],
[n, m, 1], [n, m, 1],
super::block_dims([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, padding: 0,
}; };
let mut bindings = vec![ai(0), bi(1), output(3), bytes(4, &params)]; 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<_>>(); let batch_strides = a_batch.iter().chain(&b_batch).copied().collect::<Vec<_>>();
if batch_ndim > 1 { if batch_ndim > 1 {
bindings.extend([ bindings.extend([
@@ -535,7 +586,7 @@ pub(super) fn evaluate(a: &Array, b: &Array, out: &Array) -> Result<(), String>
&bindings, &bindings,
&constants(&[ &constants(&[
(10, batch_ndim > 1), (10, batch_ndim > 1),
(100, false), (100, bias.is_some()),
(110, false), (110, false),
(200, m.is_multiple_of(64)), (200, m.is_multiple_of(64)),
(201, n.is_multiple_of(bn)), (201, n.is_multiple_of(bn)),
+2
View File
@@ -216,6 +216,7 @@ impl Options {
pub(super) struct Execution { pub(super) struct Execution {
pub(super) model: BoundModel, 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) cache: Vec<Option<Cache>>,
pub(super) mtp_cache: [Option<Cache>; 1], pub(super) mtp_cache: [Option<Cache>; 1],
pub(super) streams: Streams, pub(super) streams: Streams,
@@ -324,6 +325,7 @@ impl Execution {
let mtp_cache = model.borrow_dependent().make_mtp_cache(4); let mtp_cache = model.borrow_dependent().make_mtp_cache(4);
Ok(Self { Ok(Self {
model, model,
vision_prompt: None,
cache, cache,
mtp_cache, mtp_cache,
streams, streams,
+1
View File
@@ -47,6 +47,7 @@ pub(super) fn operator(op: &Operation) -> Option<&'static str> {
U::Rsqrt => "Rsqrt", U::Rsqrt => "Rsqrt",
U::Sin => "Sin", U::Sin => "Sin",
U::Cos => "Cos", U::Cos => "Cos",
U::Erf => "Erf",
}, },
Operation::Binary(op) => match op { Operation::Binary(op) => match op {
B::Equal => "Equal", B::Equal => "Equal",
+29 -3
View File
@@ -89,6 +89,7 @@ pub(super) struct TextModel<'a> {
pub(super) mixer: hyper::GatedResidual<'a>, pub(super) mixer: hyper::GatedResidual<'a>,
pub(super) head: Option<Linear<'a>>, // None means tied embedding weights. pub(super) head: Option<Linear<'a>>, // None means tied embedding weights.
pub(super) mtp: Option<super::mtp::Mtp<'a>>, pub(super) mtp: Option<super::mtp::Mtp<'a>>,
pub(super) vision: Option<super::vision::Input>,
pub(super) last_widened: Option<Array>, pub(super) last_widened: Option<Array>,
policy: CompilePolicy, policy: CompilePolicy,
decode_runs: Option<Vec<Run>>, decode_runs: Option<Vec<Run>>,
@@ -111,7 +112,7 @@ pub(super) struct Output {
pub(super) hidden: Option<Array>, 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 { match layer {
Linear::Dense(weight) => indexing::take_rows(weight, ids, stream), Linear::Dense(weight) => indexing::take_rows(weight, ids, stream),
Linear::Quantized(layer) => linear::dequantize( Linear::Quantized(layer) => linear::dequantize(
@@ -171,6 +172,7 @@ impl<'a> TextModel<'a> {
head, head,
parameter_order: Vec::new(), parameter_order: Vec::new(),
mtp: None, mtp: None,
vision: None,
last_widened: None, last_widened: None,
policy, policy,
decode_runs: None, decode_runs: None,
@@ -213,6 +215,8 @@ impl<'a> TextModel<'a> {
streams: &Streams, streams: &Streams,
stream: Stream, stream: Stream,
) -> Result<Output, String> { ) -> 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 emb = embedding(self.embedding, ids, stream)?;
let mtp = self let mtp = self
.mtp .mtp
@@ -242,7 +246,18 @@ impl<'a> TextModel<'a> {
streams: &Streams, streams: &Streams,
stream: Stream, stream: Stream,
) -> Result<Array, String> { ) -> 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(), Some(emb) => emb.clone(),
None => embedding(self.embedding, ids, stream)?, None => embedding(self.embedding, ids, stream)?,
}; };
@@ -440,7 +455,18 @@ impl<'a> TextModel<'a> {
streams: &Streams, streams: &Streams,
stream: Stream, stream: Stream,
) -> Result<Output, String> { ) -> 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(), Some(h) => h.clone(),
None => embedding(self.embedding, ids, stream)?, None => embedding(self.embedding, ids, stream)?,
}; };
+5 -1
View File
@@ -6,6 +6,7 @@ use super::stream::{Device, Stream, Streams};
use super::{eval, views}; use super::{eval, views};
pub(super) enum Operation { pub(super) enum Operation {
Vision(super::vision_kernels::Kernel),
Compiled(super::fused::Expression), Compiled(super::fused::Expression),
Load(super::load::Load), Load(super::load::Load),
Reshape, Reshape,
@@ -208,7 +209,7 @@ impl Operation {
| Self::ArangeU32 | Self::ArangeU32
| Self::ArangeF32 { .. } | Self::ArangeF32 { .. }
| Self::ArangeI64 => Err("invalid indexing arity".into()), | 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::QuantizedLinearBf16 { .. } => Err("quantized linear requires four inputs".into()),
Self::GatherQmmBf16 { .. } => Err("GatherQMM requires six inputs".into()), Self::GatherQmmBf16 { .. } => Err("GatherQMM requires six inputs".into()),
Self::AffineQuantize { .. } => Err("quantize requires three sibling outputs".into()), Self::AffineQuantize { .. } => Err("quantize requires three sibling outputs".into()),
@@ -257,6 +258,9 @@ pub(super) fn evaluate(
.operation() .operation()
.ok_or("array has no executable primitive")?; .ok_or("array has no executable primitive")?;
match &*operation { match &*operation {
Operation::Vision(k) => {
return super::vision_kernels::evaluate(*k, inputs, &outputs[0]);
}
Operation::Compiled(expression) => { Operation::Compiled(expression) => {
if stream.device() != Device::Gpu { if stream.device() != Device::Gpu {
return Err("compiled expression requires Metal".into()); return Err("compiled expression requires Metal".into());
+174
View File
@@ -15,6 +15,9 @@ use std::sync::atomic::AtomicBool;
pub(in crate::engine) struct QwenExecutor { pub(in crate::engine) struct QwenExecutor {
// All snapshot arrays must die before the executor's native Metal context. // All snapshot arrays must die before the executor's native Metal context.
live: LiveSession, live: LiveSession,
vision_tower: Option<super::vision::Tower>,
vision_path: Option<std::path::PathBuf>,
keep_vision_loaded: bool,
metadata: QwenMetadata, metadata: QwenMetadata,
identity: Identity, identity: Identity,
policy: TurnOptions<'static>, policy: TurnOptions<'static>,
@@ -43,6 +46,13 @@ impl QwenExecutor {
let (execution, policy) = Execution::open(settings)?; let (execution, policy) = Execution::open(settings)?;
Ok(Self { Ok(Self {
live: LiveSession::default(), 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, metadata,
identity, identity,
policy, policy,
@@ -67,6 +77,43 @@ impl QwenExecutor {
if settings.context_tokens != self.identity.context as i32 { if settings.context_tokens != self.identity.context as i32 {
return Err("Qwen turn context differs from the loaded engine".into()); 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 baseline = self.stats;
let stats = &mut self.stats; let stats = &mut self.stats;
let result = self.live.generate( let result = self.live.generate(
@@ -163,6 +210,10 @@ impl QwenExecutor {
fn clear_execution_cache(&mut self) -> Result<(), String> { fn clear_execution_cache(&mut self) -> Result<(), String> {
self.execution.streams.synchronize(self.execution.stream)?; 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(); let model = self.execution.model.borrow_dependent();
self.execution.cache = model.make_cache(4); self.execution.cache = model.make_cache(4);
self.execution.mtp_cache = model.make_mtp_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); 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();
}
+20 -6
View File
@@ -95,6 +95,8 @@ pub(in crate::engine) struct LiveSession {
pub(super) current: Option<SessionSnapshot>, pub(super) current: Option<SessionSnapshot>,
pub(super) prefix: Option<SessionSnapshot>, pub(super) prefix: Option<SessionSnapshot>,
pub(super) tag: [u8; 32], 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 { impl LiveSession {
@@ -129,7 +131,10 @@ impl LiveSession {
emit: impl FnMut(bool, String), emit: impl FnMut(bool, String),
observe: impl FnMut(TurnProgress) -> Result<(), String>, observe: impl FnMut(TurnProgress) -> Result<(), String>,
) -> Result<TextTurn, 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) let context = u32::try_from(settings.context_tokens)
.ok() .ok()
.filter(|&n| n > 0) .filter(|&n| n > 0)
@@ -154,11 +159,18 @@ impl LiveSession {
o.stops = &stop_ids; o.stops = &stop_ids;
} }
} }
let saved = [&self.current, &self.prefix] let saved = [
.into_iter() (&self.current, &self.current_images),
.flatten() (&self.prefix, &self.prefix_images),
.filter(|s| prompt.starts_with(&s.token_ids)) ]
.max_by_key(|s| s.token_ids.len()); .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 mut captured_prefix = None;
let result = execution.generate_text( let result = execution.generate_text(
&prompt, &prompt,
@@ -182,6 +194,7 @@ impl LiveSession {
// errors. Never tag those partial turns as a completed conversation. // errors. Never tag those partial turns as a completed conversation.
if let Some(prefix) = captured_prefix { if let Some(prefix) = captured_prefix {
self.prefix = Some(prefix); self.prefix = Some(prefix);
self.prefix_images = images.clone();
} }
let mut result = result?; let mut result = result?;
if let Some(snapshot) = result if let Some(snapshot) = result
@@ -194,6 +207,7 @@ impl LiveSession {
self.tag = self.tag =
conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed); conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed);
self.current = Some(snapshot); self.current = Some(snapshot);
self.current_images = images;
} }
Ok(result) Ok(result)
} }
@@ -27,6 +27,8 @@ pub(super) struct Identity {
struct Manifest { struct Manifest {
identity: Identity, identity: Identity,
tag: [u8; 32], tag: [u8; 32],
#[serde(default)]
images: Vec<crate::engine::qwen_vision::Identity>,
tokens: Vec<i32>, tokens: Vec<i32>,
lazy_kv: bool, lazy_kv: bool,
payload: serde_json::Value, payload: serde_json::Value,
@@ -111,6 +113,7 @@ impl LiveSession {
let manifest = serde_json::to_vec(&Manifest { let manifest = serde_json::to_vec(&Manifest {
identity, identity,
tag, tag,
images: self.current_images.clone(),
tokens: saved.token_ids.clone(), tokens: saved.token_ids.clone(),
lazy_kv: saved.lazy_kv, lazy_kv: saved.lazy_kv,
payload, payload,
@@ -230,6 +233,8 @@ impl LiveSession {
self.current = Some(saved); self.current = Some(saved);
self.prefix = None; self.prefix = None;
self.tag = manifest.tag; self.tag = manifest.tag;
self.current_images = manifest.images;
self.prefix_images.clear();
Ok(true) Ok(true)
} }
} }
@@ -258,6 +263,8 @@ fn qwen_checkpoint_roundtrip_and_failed_restore_preserve_live_state() {
.unwrap(); .unwrap();
let state = Some(State::Recurrent(vec![Some(leaf.clone()), None])); let state = Some(State::Recurrent(vec![Some(leaf.clone()), None]));
let mut live = LiveSession { let mut live = LiveSession {
current_images: Vec::new(),
prefix_images: Vec::new(),
current: Some(SessionSnapshot { current: Some(SessionSnapshot {
token_ids: vec![1, 2, 3, 4], token_ids: vec![1, 2, 3, 4],
logits: leaf.clone(), logits: leaf.clone(),
+1 -1
View File
@@ -15,12 +15,12 @@ pub(super) enum Unary {
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))] #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
Log, Log,
Sigmoid, Sigmoid,
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
Tanh, Tanh,
Sqrt, Sqrt,
Rsqrt, Rsqrt,
Sin, Sin,
Cos, Cos,
Erf,
} }
pub(super) fn unary(input: &Array, op: Unary, stream: Stream) -> Result<Array, String> { pub(super) fn unary(input: &Array, op: Unary, stream: Stream) -> Result<Array, String> {
+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,
))
}
@@ -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],
)
}
}
}
+4 -1
View File
@@ -34,7 +34,10 @@ pub(super) fn load_shards(
let entry = entry.map_err(|e| e.to_string())?; let entry = entry.map_err(|e| e.to_string())?;
let name = entry.file_name(); let name = entry.file_name();
let Some(name) = name.to_str() else { continue }; 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)?); parameters.extend(super::load::safetensors(&entry.path(), streams)?);
files += 1; files += 1;
} }
+3 -1
View File
@@ -200,7 +200,9 @@ impl QwenMetadata {
) )
.ok_or("Qwen artifact size overflows")?; .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)
} }
} }
+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));
}
+33 -14
View File
@@ -17,7 +17,7 @@ pub(crate) const MODEL_CHOICES: [ModelChoice; 5] = [
ModelChoice::Glm53Flash, ModelChoice::Glm53Flash,
ModelChoice::Qwen38FlashNext, ModelChoice::Qwen38FlashNext,
]; ];
pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 7] = [ pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 8] = [
ManagedArtifactId::DeepSeekV4Flash0731, ManagedArtifactId::DeepSeekV4Flash0731,
ManagedArtifactId::DeepSeekV4Flash0731Dspark, ManagedArtifactId::DeepSeekV4Flash0731Dspark,
ManagedArtifactId::DeepSeekV4Pro, ManagedArtifactId::DeepSeekV4Pro,
@@ -25,6 +25,7 @@ pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 7] = [
ManagedArtifactId::Glm53Flash, ManagedArtifactId::Glm53Flash,
ManagedArtifactId::Glm53FlashVision, ManagedArtifactId::Glm53FlashVision,
ManagedArtifactId::Qwen38FlashNext, ManagedArtifactId::Qwen38FlashNext,
ManagedArtifactId::Qwen38FlashNextVision,
]; ];
const DEEPSEEK_REPOSITORY: &str = "antirez/deepseek-v4-gguf"; const DEEPSEEK_REPOSITORY: &str = "antirez/deepseek-v4-gguf";
@@ -128,6 +129,10 @@ impl ModelChoice {
self == Self::Qwen38FlashNext self == Self::Qwen38FlashNext
} }
pub(crate) fn supports_vision(self) -> bool {
matches!(self, Self::Glm53Flash | Self::Qwen38FlashNext)
}
pub(crate) fn supports_integrated_mtp(self) -> bool { pub(crate) fn supports_integrated_mtp(self) -> bool {
self.is_glm() || self.is_qwen38() self.is_glm() || self.is_qwen38()
} }
@@ -189,7 +194,8 @@ pub(crate) fn engine_artifacts(
return EngineArtifacts { return EngineArtifacts {
model: qwen::root(models_path), model: qwen::root(models_path),
support: None, support: None,
vision: None, vision: qwen::vision_is_installed(&qwen::root(models_path))
.then(|| qwen::root(models_path).join("model-vision.safetensors")),
}; };
} }
EngineArtifacts { EngineArtifacts {
@@ -218,8 +224,14 @@ pub(crate) fn validate_engine_artifacts(
return Err(format!("DSpark is not compatible with {model}")); return Err(format!("DSpark is not compatible with {model}"));
} }
if model.is_qwen38() { if model.is_qwen38() {
if artifacts.support.is_some() || artifacts.vision.is_some() { if artifacts.support.is_some() {
return Err("Qwen3.8 does not accept GGUF support or vision artifacts".into()); return Err("Qwen3.8 does not accept GGUF support artifacts".into());
}
if let Some(path) = &artifacts.vision
&& (path != &artifacts.model.join("model-vision.safetensors")
|| !qwen::vision_is_installed(&artifacts.model))
{
return Err("Qwen vision artifacts have not passed verification".into());
} }
return qwen::validate_installed(&artifacts.model); return qwen::validate_installed(&artifacts.model);
} }
@@ -258,9 +270,14 @@ pub(crate) enum ManagedArtifactId {
Glm53Flash, Glm53Flash,
Glm53FlashVision, Glm53FlashVision,
Qwen38FlashNext, Qwen38FlashNext,
Qwen38FlashNextVision,
} }
impl ManagedArtifactId { impl ManagedArtifactId {
fn is_qwen(self) -> bool {
self.model().is_qwen38()
}
pub(crate) fn model(self) -> ModelChoice { pub(crate) fn model(self) -> ModelChoice {
match self { match self {
Self::DeepSeekV4Flash0731 | Self::DeepSeekV4Flash0731Dspark => { Self::DeepSeekV4Flash0731 | Self::DeepSeekV4Flash0731Dspark => {
@@ -269,7 +286,7 @@ impl ManagedArtifactId {
Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro, Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro,
Self::Glm52 => ModelChoice::Glm52, Self::Glm52 => ModelChoice::Glm52,
Self::Glm53Flash | Self::Glm53FlashVision => ModelChoice::Glm53Flash, Self::Glm53Flash | Self::Glm53FlashVision => ModelChoice::Glm53Flash,
Self::Qwen38FlashNext => ModelChoice::Qwen38FlashNext, Self::Qwen38FlashNext | Self::Qwen38FlashNextVision => ModelChoice::Qwen38FlashNext,
} }
} }
@@ -281,15 +298,17 @@ impl ManagedArtifactId {
Self::Glm52 => &GLM, Self::Glm52 => &GLM,
Self::Glm53Flash => &GLM53_FLASH, Self::Glm53Flash => &GLM53_FLASH,
Self::Glm53FlashVision => &GLM53_FLASH_VISION, Self::Glm53FlashVision => &GLM53_FLASH_VISION,
Self::Qwen38FlashNext => unreachable!("Qwen uses a pinned artifact set"), Self::Qwen38FlashNext | Self::Qwen38FlashNextVision => {
unreachable!("Qwen uses a pinned artifact set")
}
} }
} }
} }
impl fmt::Display for ManagedArtifactId { impl fmt::Display for ManagedArtifactId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
if *self == Self::Qwen38FlashNext { if self.is_qwen() {
formatter.write_str(qwen::LABEL) formatter.write_str(qwen::label(*self))
} else { } else {
formatter.write_str(self.artifact().label) formatter.write_str(self.artifact().label)
} }
@@ -480,8 +499,8 @@ pub(crate) fn managed_artifacts(models_path: &Path) -> Vec<ManagedArtifact> {
MANAGED_ARTIFACTS MANAGED_ARTIFACTS
.into_iter() .into_iter()
.map(|id| { .map(|id| {
if id == ManagedArtifactId::Qwen38FlashNext { if id.is_qwen() {
return qwen::managed_artifact(models_path); return qwen::managed_artifact(id, models_path);
} }
let model = id.model(); let model = id.model();
let artifact = id.artifact(); let artifact = id.artifact();
@@ -527,8 +546,8 @@ pub(crate) fn artifact_download_progress(
id: ManagedArtifactId, id: ManagedArtifactId,
models_path: &Path, models_path: &Path,
) -> DownloadProgress { ) -> DownloadProgress {
if id == ManagedArtifactId::Qwen38FlashNext { if id.is_qwen() {
return qwen::download_progress(models_path); return qwen::download_progress(id, models_path);
} }
let model = id.model(); let model = id.model();
let artifact = id.artifact(); let artifact = id.artifact();
@@ -560,8 +579,8 @@ pub(crate) fn artifact_verification_progress(
id: ManagedArtifactId, id: ManagedArtifactId,
verified: u64, verified: u64,
) -> DownloadProgress { ) -> DownloadProgress {
if id == ManagedArtifactId::Qwen38FlashNext { if id.is_qwen() {
return qwen::verification_progress(verified); return qwen::verification_progress(id, verified);
} }
let artifact = id.artifact(); let artifact = id.artifact();
DownloadProgress { DownloadProgress {
+90 -26
View File
@@ -22,6 +22,7 @@ struct Manifest {
format: u32, format: u32,
source: Source, source: Source,
files: Vec<Artifact>, files: Vec<Artifact>,
excluded: Vec<Artifact>,
} }
#[derive(Deserialize)] #[derive(Deserialize)]
@@ -33,6 +34,7 @@ struct Source {
#[derive(Deserialize)] #[derive(Deserialize)]
pub(super) struct Artifact { pub(super) struct Artifact {
pub(super) path: String, pub(super) path: String,
#[serde(default)]
pub(super) role: String, pub(super) role: String,
pub(super) size: u64, pub(super) size: u64,
pub(super) sha256: String, pub(super) sha256: String,
@@ -61,6 +63,23 @@ pub(super) fn artifacts() -> &'static [Artifact] {
&manifest().files &manifest().files
} }
pub(super) fn selected(id: ManagedArtifactId) -> &'static [Artifact] {
match id {
ManagedArtifactId::Qwen38FlashNext => artifacts(),
ManagedArtifactId::Qwen38FlashNextVision => &manifest().excluded,
_ => unreachable!("not a Qwen artifact set"),
}
}
pub(super) fn label(id: ManagedArtifactId) -> &'static str {
match id {
ManagedArtifactId::Qwen38FlashNext => LABEL,
ManagedArtifactId::Qwen38FlashNextVision => "Qwen3.8 Flash Next vision encoder",
_ => unreachable!("not a Qwen artifact set"),
}
}
#[cfg(test)]
pub(super) fn total_bytes() -> u64 { pub(super) fn total_bytes() -> u64 {
artifacts().iter().map(|artifact| artifact.size).sum() artifacts().iter().map(|artifact| artifact.size).sum()
} }
@@ -106,6 +125,12 @@ pub(super) fn is_installed(root: &Path) -> bool {
.all(|artifact| artifact_is_installed(root, artifact)) .all(|artifact| artifact_is_installed(root, artifact))
} }
pub(super) fn vision_is_installed(root: &Path) -> bool {
selected(ManagedArtifactId::Qwen38FlashNextVision)
.iter()
.all(|a| artifact_is_installed(root, a))
}
pub(super) fn validate_installed(root: &Path) -> Result<(), String> { pub(super) fn validate_installed(root: &Path) -> Result<(), String> {
for artifact in artifacts() { for artifact in artifacts() {
let artifact_path = path(root, artifact); let artifact_path = path(root, artifact);
@@ -136,8 +161,8 @@ pub(super) fn validate_installed(root: &Path) -> Result<(), String> {
Ok(()) Ok(())
} }
fn stored_bytes(root: &Path) -> u64 { fn stored_bytes(id: ManagedArtifactId, root: &Path) -> u64 {
artifacts() selected(id)
.iter() .iter()
.flat_map(|artifact| [path(root, artifact), partial_path(root, artifact)]) .flat_map(|artifact| [path(root, artifact), partial_path(root, artifact)])
.filter_map(|path| path.metadata().ok()) .filter_map(|path| path.metadata().ok())
@@ -145,8 +170,8 @@ fn stored_bytes(root: &Path) -> u64 {
.sum() .sum()
} }
fn downloaded_bytes(root: &Path) -> u64 { fn downloaded_bytes(id: ManagedArtifactId, root: &Path) -> u64 {
artifacts() selected(id)
.iter() .iter()
.map(|artifact| { .map(|artifact| {
if artifact_is_installed(root, artifact) { if artifact_is_installed(root, artifact) {
@@ -161,8 +186,8 @@ fn downloaded_bytes(root: &Path) -> u64 {
.sum() .sum()
} }
fn complete_payloads(root: &Path) -> bool { fn complete_payloads(id: ManagedArtifactId, root: &Path) -> bool {
artifacts().iter().all(|artifact| { selected(id).iter().all(|artifact| {
path(root, artifact) path(root, artifact)
.metadata() .metadata()
.or_else(|_| partial_path(root, artifact).metadata()) .or_else(|_| partial_path(root, artifact).metadata())
@@ -170,12 +195,12 @@ fn complete_payloads(root: &Path) -> bool {
}) })
} }
pub(super) fn managed_artifact(models_path: &Path) -> ManagedArtifact { pub(super) fn managed_artifact(id: ManagedArtifactId, models_path: &Path) -> ManagedArtifact {
let root = root(models_path); let root = root(models_path);
let stored = stored_bytes(&root); let stored = stored_bytes(id, &root);
let state = if is_installed(&root) { let state = if selected(id).iter().all(|a| artifact_is_installed(&root, a)) {
ManagedArtifactState::Ready ManagedArtifactState::Ready
} else if complete_payloads(&root) { } else if complete_payloads(id, &root) {
ManagedArtifactState::NeedsVerification ManagedArtifactState::NeedsVerification
} else if stored > 0 { } else if stored > 0 {
ManagedArtifactState::Partial ManagedArtifactState::Partial
@@ -183,44 +208,45 @@ pub(super) fn managed_artifact(models_path: &Path) -> ManagedArtifact {
ManagedArtifactState::Missing ManagedArtifactState::Missing
}; };
ManagedArtifact { ManagedArtifact {
id: ManagedArtifactId::Qwen38FlashNext, id,
stored, stored,
expected: total_bytes(), expected: selected(id).iter().map(|a| a.size).sum(),
state, state,
} }
} }
pub(super) fn download_progress(models_path: &Path) -> DownloadProgress { pub(super) fn download_progress(id: ManagedArtifactId, models_path: &Path) -> DownloadProgress {
let root = root(models_path); let root = root(models_path);
let downloaded = downloaded_bytes(&root); let downloaded = downloaded_bytes(id, &root);
let installed = is_installed(&root); let installed = selected(id).iter().all(|a| artifact_is_installed(&root, a));
let verification = (!installed && complete_payloads(&root)).then_some(VerificationProgress { let verification =
verified: 0, (!installed && complete_payloads(id, &root)).then_some(VerificationProgress {
total: total_bytes(), verified: 0,
}); total: selected(id).iter().map(|a| a.size).sum(),
});
let phase = if installed { let phase = if installed {
DownloadPhase::Complete DownloadPhase::Complete
} else if verification.is_some() { } else if verification.is_some() {
DownloadPhase::Verifying(LABEL) DownloadPhase::Verifying(label(id))
} else if downloaded > 0 { } else if downloaded > 0 {
DownloadPhase::Downloading(LABEL) DownloadPhase::Downloading(label(id))
} else { } else {
DownloadPhase::Pending(LABEL) DownloadPhase::Pending(label(id))
}; };
DownloadProgress { DownloadProgress {
downloaded, downloaded,
total: total_bytes(), total: selected(id).iter().map(|a| a.size).sum(),
phase, phase,
verification, verification,
} }
} }
pub(super) fn verification_progress(verified: u64) -> DownloadProgress { pub(super) fn verification_progress(id: ManagedArtifactId, verified: u64) -> DownloadProgress {
let total = total_bytes(); let total = selected(id).iter().map(|a| a.size).sum();
DownloadProgress { DownloadProgress {
downloaded: total, downloaded: total,
total, total,
phase: DownloadPhase::Verifying(LABEL), phase: DownloadPhase::Verifying(label(id)),
verification: Some(VerificationProgress { verification: Some(VerificationProgress {
verified: verified.min(total), verified: verified.min(total),
total, total,
@@ -336,3 +362,41 @@ mod tests {
fs::remove_dir_all(root).unwrap(); fs::remove_dir_all(root).unwrap();
} }
} }
#[test]
fn qwen_vision_is_an_independent_managed_artifact_set() {
let id = ManagedArtifactId::Qwen38FlashNextVision;
let root = std::env::temp_dir().join(format!("qwen-vision-catalog-{}", std::process::id()));
let model_root = self::root(&root);
fs::create_dir_all(&model_root).unwrap();
assert_eq!(selected(id).len(), 4);
assert_eq!(
selected(id).iter().map(|a| a.size).sum::<u64>(),
897_900_287
);
assert!(
selected(id)
.iter()
.all(|v| artifacts().iter().all(|a| a.path != v.path))
);
for artifact in selected(id) {
File::create(path(&model_root, artifact))
.unwrap()
.set_len(artifact.size)
.unwrap();
fs::write(verification_path(&model_root, artifact), &artifact.sha256).unwrap();
}
let text = model_root.join("model-00001-of-00017.safetensors");
fs::write(&text, b"preserve text model").unwrap();
assert_eq!(
managed_artifact(id, &root).state,
ManagedArtifactState::Ready
);
assert_eq!(download_progress(id, &root).phase, DownloadPhase::Complete);
assert!(!is_installed(&model_root));
assert!(vision_is_installed(&model_root));
super::delete_managed_artifact(id, &root).unwrap();
assert_eq!(fs::read(&text).unwrap(), b"preserve text model");
assert!(!vision_is_installed(&model_root));
fs::remove_dir_all(root).unwrap();
}
+48 -24
View File
@@ -14,8 +14,8 @@ pub(crate) fn download_managed_artifact(
cancel: &AtomicBool, cancel: &AtomicBool,
verified_bytes: &AtomicU64, verified_bytes: &AtomicU64,
) -> Result<DownloadOutcome, String> { ) -> Result<DownloadOutcome, String> {
if id == ManagedArtifactId::Qwen38FlashNext { if id.is_qwen() {
return download_qwen(models_path, cancel, verified_bytes); return download_qwen(id, models_path, cancel, verified_bytes);
} }
download_artifact_with_cancel( download_artifact_with_cancel(
id.model(), id.model(),
@@ -32,8 +32,8 @@ pub(crate) fn validate_managed_artifact(
cancel: &AtomicBool, cancel: &AtomicBool,
verified_bytes: &AtomicU64, verified_bytes: &AtomicU64,
) -> Result<DownloadOutcome, String> { ) -> Result<DownloadOutcome, String> {
if id == ManagedArtifactId::Qwen38FlashNext { if id.is_qwen() {
return validate_qwen(models_path, cancel, verified_bytes); return validate_qwen(id, models_path, cancel, verified_bytes);
} }
let model = id.model(); let model = id.model();
let artifact = id.artifact(); let artifact = id.artifact();
@@ -73,8 +73,8 @@ pub(crate) fn delete_managed_artifact(
id: ManagedArtifactId, id: ManagedArtifactId,
models_path: &Path, models_path: &Path,
) -> Result<(), String> { ) -> Result<(), String> {
if id == ManagedArtifactId::Qwen38FlashNext { if id.is_qwen() {
return delete_qwen(models_path); return delete_qwen(id, models_path);
} }
let model = id.model(); let model = id.model();
let artifact = id.artifact(); let artifact = id.artifact();
@@ -93,6 +93,7 @@ pub(crate) fn delete_managed_artifact(
} }
fn download_qwen( fn download_qwen(
id: ManagedArtifactId,
models_path: &Path, models_path: &Path,
cancel: &AtomicBool, cancel: &AtomicBool,
verified_bytes: &AtomicU64, verified_bytes: &AtomicU64,
@@ -100,7 +101,7 @@ fn download_qwen(
let root = qwen::root(models_path); let root = qwen::root(models_path);
fs::create_dir_all(&root).map_err(|error| format!("{}: {error}", root.display()))?; fs::create_dir_all(&root).map_err(|error| format!("{}: {error}", root.display()))?;
verified_bytes.store(0, Ordering::Relaxed); verified_bytes.store(0, Ordering::Relaxed);
for artifact in qwen::artifacts() { for artifact in qwen::selected(id) {
if cancel.load(Ordering::Relaxed) { if cancel.load(Ordering::Relaxed) {
return Ok(DownloadOutcome::Stopped); return Ok(DownloadOutcome::Stopped);
} }
@@ -145,26 +146,29 @@ fn download_qwen(
fs::write(qwen::verification_path(&root, artifact), &artifact.sha256) fs::write(qwen::verification_path(&root, artifact), &artifact.sha256)
.map_err(|error| error.to_string())?; .map_err(|error| error.to_string())?;
} }
if let Err(error) = crate::engine::validate_qwen_artifacts(&root) { if id == ManagedArtifactId::Qwen38FlashNext
clear_qwen_markers(&root)?; && let Err(error) = crate::engine::validate_qwen_artifacts(&root)
{
clear_qwen_markers(id, &root)?;
return Err(error); return Err(error);
} }
Ok(DownloadOutcome::Complete) Ok(DownloadOutcome::Complete)
} }
fn validate_qwen( fn validate_qwen(
id: ManagedArtifactId,
models_path: &Path, models_path: &Path,
cancel: &AtomicBool, cancel: &AtomicBool,
verified_bytes: &AtomicU64, verified_bytes: &AtomicU64,
) -> Result<DownloadOutcome, String> { ) -> Result<DownloadOutcome, String> {
let root = qwen::root(models_path); let root = qwen::root(models_path);
verified_bytes.store(0, Ordering::Relaxed); verified_bytes.store(0, Ordering::Relaxed);
for artifact in qwen::artifacts() { for artifact in qwen::selected(id) {
if !qwen::path(&root, artifact).exists() && !qwen::partial_path(&root, artifact).exists() { if !qwen::path(&root, artifact).exists() && !qwen::partial_path(&root, artifact).exists() {
return Err(format!("{} is not downloaded", artifact.path)); return Err(format!("{} is not downloaded", artifact.path));
} }
} }
for artifact in qwen::artifacts() { for artifact in qwen::selected(id) {
let destination = qwen::path(&root, artifact); let destination = qwen::path(&root, artifact);
let partial = qwen::partial_path(&root, artifact); let partial = qwen::partial_path(&root, artifact);
let (source, promote) = if destination.exists() { let (source, promote) = if destination.exists() {
@@ -189,15 +193,17 @@ fn validate_qwen(
fs::write(qwen::verification_path(&root, artifact), &artifact.sha256) fs::write(qwen::verification_path(&root, artifact), &artifact.sha256)
.map_err(|error| error.to_string())?; .map_err(|error| error.to_string())?;
} }
if let Err(error) = crate::engine::validate_qwen_artifacts(&root) { if id == ManagedArtifactId::Qwen38FlashNext
clear_qwen_markers(&root)?; && let Err(error) = crate::engine::validate_qwen_artifacts(&root)
{
clear_qwen_markers(id, &root)?;
return Err(error); return Err(error);
} }
Ok(DownloadOutcome::Complete) Ok(DownloadOutcome::Complete)
} }
fn clear_qwen_markers(root: &Path) -> Result<(), String> { fn clear_qwen_markers(id: ManagedArtifactId, root: &Path) -> Result<(), String> {
for artifact in qwen::artifacts() { for artifact in qwen::selected(id) {
let path = qwen::verification_path(root, artifact); let path = qwen::verification_path(root, artifact);
match fs::remove_file(&path) { match fs::remove_file(&path) {
Ok(()) => {} Ok(()) => {}
@@ -208,9 +214,9 @@ fn clear_qwen_markers(root: &Path) -> Result<(), String> {
Ok(()) Ok(())
} }
fn delete_qwen(models_path: &Path) -> Result<(), String> { fn delete_qwen(id: ManagedArtifactId, models_path: &Path) -> Result<(), String> {
let root = qwen::root(models_path); let root = qwen::root(models_path);
for artifact in qwen::artifacts() { for artifact in qwen::selected(id) {
for path in [ for path in [
qwen::path(&root, artifact), qwen::path(&root, artifact),
qwen::partial_path(&root, artifact), qwen::partial_path(&root, artifact),
@@ -225,7 +231,14 @@ fn delete_qwen(models_path: &Path) -> Result<(), String> {
} }
match fs::remove_dir(&root) { match fs::remove_dir(&root) {
Ok(()) => Ok(()), Ok(()) => Ok(()),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), Err(error)
if matches!(
error.kind(),
std::io::ErrorKind::NotFound | std::io::ErrorKind::DirectoryNotEmpty
) =>
{
Ok(())
}
Err(error) => Err(format!("{}: {error}", root.display())), Err(error) => Err(format!("{}: {error}", root.display())),
} }
} }
@@ -480,7 +493,7 @@ mod tests {
); );
assert!(ModelChoice::from_id("unknown").is_none()); assert!(ModelChoice::from_id("unknown").is_none());
assert_eq!(MODEL_CHOICES.len(), 5); assert_eq!(MODEL_CHOICES.len(), 5);
assert_eq!(MANAGED_ARTIFACTS.len(), 7); assert_eq!(MANAGED_ARTIFACTS.len(), 8);
assert_eq!( assert_eq!(
ModelChoice::from_id("qwen3.8-flash-next"), ModelChoice::from_id("qwen3.8-flash-next"),
Some(ModelChoice::Qwen38FlashNext) Some(ModelChoice::Qwen38FlashNext)
@@ -600,20 +613,31 @@ mod tests {
)) ))
.unwrap(); .unwrap();
assert!( assert!(
validate_qwen(&models, &AtomicBool::new(false), &AtomicU64::new(0)) validate_qwen(
.unwrap_err() ManagedArtifactId::Qwen38FlashNext,
.contains("mtp.safetensors is not downloaded") &models,
&AtomicBool::new(false),
&AtomicU64::new(0)
)
.unwrap_err()
.contains("mtp.safetensors is not downloaded")
); );
fs::hard_link(source.join("mtp.safetensors"), root.join("mtp.safetensors")).unwrap(); fs::hard_link(source.join("mtp.safetensors"), root.join("mtp.safetensors")).unwrap();
let verified = AtomicU64::new(0); let verified = AtomicU64::new(0);
assert_eq!( assert_eq!(
validate_qwen(&models, &AtomicBool::new(false), &verified).unwrap(), validate_qwen(
ManagedArtifactId::Qwen38FlashNext,
&models,
&AtomicBool::new(false),
&verified
)
.unwrap(),
DownloadOutcome::Complete DownloadOutcome::Complete
); );
assert_eq!(verified.load(Ordering::Relaxed), qwen::total_bytes()); assert_eq!(verified.load(Ordering::Relaxed), qwen::total_bytes());
assert!(qwen::is_installed(&root)); assert!(qwen::is_installed(&root));
assert!(crate::engine::validate_qwen_artifacts(&root).is_ok()); assert!(crate::engine::validate_qwen_artifacts(&root).is_ok());
delete_qwen(&models).unwrap(); delete_qwen(ManagedArtifactId::Qwen38FlashNext, &models).unwrap();
assert!(!root.exists()); assert!(!root.exists());
fs::remove_dir_all(models).unwrap(); fs::remove_dir_all(models).unwrap();
} }
+29
View File
@@ -38,6 +38,7 @@ struct Options {
prompt: String, prompt: String,
additional_prompts: Vec<String>, additional_prompts: Vec<String>,
input_file: Option<PathBuf>, input_file: Option<PathBuf>,
image_file: Option<PathBuf>,
models_path: PathBuf, models_path: PathBuf,
config_path: PathBuf, config_path: PathBuf,
plain_chat: bool, plain_chat: bool,
@@ -68,6 +69,7 @@ impl Options {
let mut prompt_overridden = false; let mut prompt_overridden = false;
let mut additional_prompts = Vec::new(); let mut additional_prompts = Vec::new();
let mut input_file = None; let mut input_file = None;
let mut image_file = None;
let mut models_path = crate::app::models_path(); let mut models_path = crate::app::models_path();
let mut config_path = crate::app::config_path(); let mut config_path = crate::app::config_path();
let mut context = None; let mut context = None;
@@ -120,6 +122,7 @@ impl Options {
prompt_overridden = true; prompt_overridden = true;
} }
} }
"--image-file" => image_file = Some(value()?.into()),
"--input-file" => input_file = Some(value()?.into()), "--input-file" => input_file = Some(value()?.into()),
"--models-dir" => models_path = value()?.into(), "--models-dir" => models_path = value()?.into(),
"--config" => config_path = value()?.into(), "--config" => config_path = value()?.into(),
@@ -190,6 +193,7 @@ impl Options {
prompt, prompt,
additional_prompts, additional_prompts,
input_file, input_file,
image_file,
models_path, models_path,
config_path, config_path,
plain_chat, plain_chat,
@@ -249,6 +253,7 @@ Models: deepseek-v4-flash-0731, deepseek-v4-pro, glm-5.2, glm-5.3-flash, qwen3.8
Options:\n\ Options:\n\
--prompt TEXT Chat turn; repeat for one ongoing conversation\n\ --prompt TEXT Chat turn; repeat for one ongoing conversation\n\
(default: autonomous short story)\n\ (default: autonomous short story)\n\
--image-file PATH Attach PNG/JPEG pixels to the first turn (no filename sent)\n\
--input-file PATH Append a UTF-8 file to the first turn (maximum: 16 MiB)\n\ --input-file PATH Append a UTF-8 file to the first turn (maximum: 16 MiB)\n\
--context TOKENS Context size; otherwise use the model profile\n\ --context TOKENS Context size; otherwise use the model profile\n\
--max-tokens TOKENS Override the model profile generation limit\n\ --max-tokens TOKENS Override the model profile generation limit\n\
@@ -937,6 +942,30 @@ fn stop_worker(
} }
fn prompts_with_input_file(options: &Options) -> Result<Vec<String>, String> { fn prompts_with_input_file(options: &Options) -> Result<Vec<String>, String> {
let mut prompts = text_prompts_with_input_file(options)?;
if let Some(path) = &options.image_file {
use base64::Engine as _;
if !options.model.supports_vision() {
return Err("selected model does not support images".into());
}
if fs::metadata(path).map_err(|e| e.to_string())?.len() > 50 * 1024 * 1024 {
return Err("image exceeds 50 MiB".into());
}
let bytes = fs::read(path).map_err(|e| e.to_string())?;
let mime = match image::guess_format(&bytes).map_err(|e| e.to_string())? {
image::ImageFormat::Png => "image/png",
image::ImageFormat::Jpeg => "image/jpeg",
_ => return Err("image must be PNG or JPEG".into()),
};
prompts[0].push_str(&crate::engine::vision_data_marker(&format!(
"data:{mime};base64,{}",
base64::engine::general_purpose::STANDARD.encode(bytes)
)));
}
Ok(prompts)
}
fn text_prompts_with_input_file(options: &Options) -> Result<Vec<String>, String> {
let mut prompts = Vec::with_capacity(1 + options.additional_prompts.len()); let mut prompts = Vec::with_capacity(1 + options.additional_prompts.len());
let Some(path) = &options.input_file else { let Some(path) = &options.input_file else {
prompts.push(options.prompt.clone()); prompts.push(options.prompt.clone());
+2 -2
View File
@@ -51,8 +51,8 @@ impl SpeculativePreferences {
if !model.supports_integrated_mtp() && (self.glm_mtp || self.glm_mtp_timing) { if !model.supports_integrated_mtp() && (self.glm_mtp || self.glm_mtp_timing) {
return Err("Integrated MTP is unavailable for the selected model.".into()); return Err("Integrated MTP is unavailable for the selected model.".into());
} }
if self.keep_vision_loaded && model != ModelChoice::Glm53Flash { if self.keep_vision_loaded && !model.supports_vision() {
return Err("Persistent vision weights are available only for GLM 5.3 Flash.".into()); return Err("Persistent vision weights are available only for GLM 5.3 Flash and Qwen3.8 Flash Next.".into());
} }
if self.dspark_enabled && !model.supports_dspark() { if self.dspark_enabled && !model.supports_dspark() {
return Err("DSpark is not available for the selected model.".into()); return Err("DSpark is not available for the selected model.".into());
+32
View File
@@ -0,0 +1,32 @@
"""Describe a neutral image and run a no-image control through pinned MTPLX.
Usage: test-supervisor ... --command python tools/qwen-vision-chat-reference.py MODEL_DIR IMAGE
Uses the reference runtime only as an oracle; application inference is Rust/Metal.
"""
import os,sys,json,time,argparse
from pathlib import Path
root=Path(sys.argv[1]);image=Path(sys.argv[2])
from mtplx.server.openai import _server_runtime_env_overrides
os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(root),generation_mode='mtp'),None))
os.environ['MTPLX_SUSTAINED_PREFILL']='1'
os.environ['MTPLX_PREFILL_CHUNK_SIZE']='2048'
import mlx.core as mx
from mlx_lm.utils import load_model
from mtplx.models import qwen4_exp as qwen
from mtplx.runtime import _load_tokenizer_resilient,MTPLXRuntime
from mtplx.mtp_patch import MTPContract
from mtplx.sampling import SamplerConfig
from mtplx.generation import generate_mtpk
from mtplx.server.openai import ChatMessage,_encode_messages_uncached,_materialize_vision_splice
print(json.dumps(dict(event='loading')),flush=True)
tokenizer=_load_tokenizer_resilient(root,json.loads((root/'config.json').read_text()))
model,_=load_model(root,lazy=False,strict=True,get_model_classes=lambda **_:(qwen.Model,qwen.ModelArgs))
model.post_weight_load(root);assert model.attach_mtp(root)
rt=MTPLXRuntime(model,tokenizer,root,True,MTPContract())
for with_image in [True,False]:
text='Describe this image'+('<|vision_start|><|image_pad|><|vision_end|>' if with_image else '')
ids=_encode_messages_uncached(tokenizer,[ChatMessage(role='user',content=text)],enable_thinking=True,reasoning_effort='low',preserve_reasoning_history=True,tools=None)
splice=None
if with_image:ids,splice=_materialize_vision_splice(argparse.Namespace(args=argparse.Namespace(model=str(root))),[image.read_bytes()],ids)
print(json.dumps(dict(event='input',with_image=with_image,prompt='Describe this image',ids=ids)),flush=True)
out=generate_mtpk(rt,ids,max_tokens=1024,sampler=SamplerConfig(temperature=0.,top_p=.95,top_k=20),speculative_depth=3,seed=1,stop_token_ids=set(tokenizer.eos_token_ids),mtp_history_policy='committed',verify_strategy='batched',vision_splice=splice,capture_final_state=True,prefill_callback=lambda d:print(json.dumps(dict(event='prefill',**d)),flush=True),token_callback=lambda ids:print(json.dumps(dict(event='tokens',ids=ids)),flush=True))
print(json.dumps(dict(event='result',with_image=with_image,text=out.text,tokens=out.tokens,finish_reason=out.finish_reason)),flush=True)
+24
View File
@@ -0,0 +1,24 @@
"""Export pinned MTPLX vision stages; oracle only, never application code.
Usage: python tools/qwen-vision-reference.py MODEL_DIR IMAGE OUTPUT_DIR
Run in the pinned MTPLX reference environment with GPU access.
"""
import sys,json,hashlib,pathlib
import mlx.core as mx
import numpy as np
from mtplx.vision import load_vision_tower
from mtplx.vision.processing import decode_image,preprocess_images
root=pathlib.Path(sys.argv[1]); image=pathlib.Path(sys.argv[2]); out=pathlib.Path(sys.argv[3]);out.mkdir(exist_ok=True)
def save(name,x):
mx.eval(x); a=np.asarray(x.astype(mx.float32));a.tofile(out/(name+'.f32')); print(json.dumps(dict(stage=name,shape=list(x.shape),dtype=str(x.dtype),min=float(a.min()),max=float(a.max()),sha256=hashlib.sha256(a.tobytes()).hexdigest())),flush=True)
pixels,grids=preprocess_images([decode_image(image.read_bytes())],json.loads((root/'preprocessor_config.json').read_text()))
(out/'grid.json').write_text(json.dumps(grids[0]));save('pixels',pixels); print(json.dumps(dict(grids=grids)),flush=True)
tower=load_vision_tower(root)
h=tower.patch_embed(pixels.astype(tower.patch_embed.proj.weight.dtype));save('patch',h)
p=tower.fast_pos_embed_interpolate(grids);save('position',p)
h=h+p;r=tower.rot_pos_emb(grids);save('rotary',r)
for i,b in enumerate(tower.blocks):
h=b(h,[],r)
if i in [0,26]:save('block'+str(i),h)
else:mx.eval(h)
print(json.dumps(dict(block=i)),flush=True)
h=tower.merger(h);save('embeddings',h)