Add Qwen vision loading and Metal inference

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

No files matched your search

+1 -1
View File
@@ -1402,7 +1402,7 @@ impl App {
}
Message::ComposerAction(action) => self.composer.perform(action),
Message::ChooseVisionImage => {
if self.config.model != ModelChoice::Glm53Flash || self.generating {
if !self.config.model.supports_vision() || self.generating {
return Task::none();
}
return Task::perform(
+2 -2
View File
@@ -604,7 +604,7 @@ impl App {
return;
}
self.config = config;
if self.config.model != ModelChoice::Glm53Flash {
if !self.config.model.supports_vision() {
self.pending_vision_image = None;
}
self.context_limit = self.config.active_generation().context_tokens.max(0) as u32;
@@ -889,7 +889,7 @@ impl App {
}
Message::PreferenceKeepVisionLoadedChanged(value) => {
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;
}
Message::PreferenceDsparkConfidenceChanged(value) => {
+1 -1
View File
@@ -11,7 +11,7 @@ impl App {
self.config = config;
self.preference_draft = PreferenceDraft::from_saved(&self.config);
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;
}
#[cfg(target_os = "macos")]
+5 -3
View File
@@ -337,8 +337,8 @@ impl App {
.padding([2, 6])
.into()
};
let vision_ready = self.config.model == ModelChoice::Glm53Flash
&& model::engine_artifacts(ModelChoice::Glm53Flash, false, &models_path())
let vision_ready = self.config.model.supports_vision()
&& model::engine_artifacts(self.config.model, false, &models_path())
.vision
.is_some();
let attach = if vision_ready && !self.generating && self.pending_vision_image.is_none()
@@ -356,8 +356,10 @@ impl App {
attach,
container(text(if vision_ready {
"Attach PNG or JPEG"
} else if self.config.model.supports_vision() {
"Download the vision encoder in Model Manager"
} else {
"GLM 5.3 Flash vision sidecar is not ready"
"This model does not support image attachments"
}))
.padding(8)
.style(preference_group_style),
+2 -2
View File
@@ -25,7 +25,7 @@ impl App {
.supports_integrated_mtp()
.then_some(Message::PreferenceGlmMtpTimingChanged);
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);
let dspark_strict_toggle: Option<fn(bool) -> Message> = self
.preference_draft
@@ -623,7 +623,7 @@ impl App {
),
hint(
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),
"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")]
mod metal;
mod qwen;
mod qwen_vision;
#[cfg(test)]
#[path = "../tools/sampler-replay-benchmark.rs"]
mod sampler_replay_benchmark;
@@ -1504,9 +1505,6 @@ impl Generator {
.iter()
.any(|message| message.content.contains(VISION_DATA_START));
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);
}
let (tokens, overlays) = if has_vision {
+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;
@@ -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 {
+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)),
+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,
+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",
+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)?,
};
+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());
+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();
}
+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)
}
@@ -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(),
+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> {
+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 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;
}
+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)
}
}
+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::Qwen38FlashNext,
];
pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 7] = [
pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 8] = [
ManagedArtifactId::DeepSeekV4Flash0731,
ManagedArtifactId::DeepSeekV4Flash0731Dspark,
ManagedArtifactId::DeepSeekV4Pro,
@@ -25,6 +25,7 @@ pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 7] = [
ManagedArtifactId::Glm53Flash,
ManagedArtifactId::Glm53FlashVision,
ManagedArtifactId::Qwen38FlashNext,
ManagedArtifactId::Qwen38FlashNextVision,
];
const DEEPSEEK_REPOSITORY: &str = "antirez/deepseek-v4-gguf";
@@ -128,6 +129,10 @@ impl ModelChoice {
self == Self::Qwen38FlashNext
}
pub(crate) fn supports_vision(self) -> bool {
matches!(self, Self::Glm53Flash | Self::Qwen38FlashNext)
}
pub(crate) fn supports_integrated_mtp(self) -> bool {
self.is_glm() || self.is_qwen38()
}
@@ -189,7 +194,8 @@ pub(crate) fn engine_artifacts(
return EngineArtifacts {
model: qwen::root(models_path),
support: None,
vision: None,
vision: qwen::vision_is_installed(&qwen::root(models_path))
.then(|| qwen::root(models_path).join("model-vision.safetensors")),
};
}
EngineArtifacts {
@@ -218,8 +224,14 @@ pub(crate) fn validate_engine_artifacts(
return Err(format!("DSpark is not compatible with {model}"));
}
if model.is_qwen38() {
if artifacts.support.is_some() || artifacts.vision.is_some() {
return Err("Qwen3.8 does not accept GGUF support or vision artifacts".into());
if artifacts.support.is_some() {
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);
}
@@ -258,9 +270,14 @@ pub(crate) enum ManagedArtifactId {
Glm53Flash,
Glm53FlashVision,
Qwen38FlashNext,
Qwen38FlashNextVision,
}
impl ManagedArtifactId {
fn is_qwen(self) -> bool {
self.model().is_qwen38()
}
pub(crate) fn model(self) -> ModelChoice {
match self {
Self::DeepSeekV4Flash0731 | Self::DeepSeekV4Flash0731Dspark => {
@@ -269,7 +286,7 @@ impl ManagedArtifactId {
Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro,
Self::Glm52 => ModelChoice::Glm52,
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::Glm53Flash => &GLM53_FLASH,
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 {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
if *self == Self::Qwen38FlashNext {
formatter.write_str(qwen::LABEL)
if self.is_qwen() {
formatter.write_str(qwen::label(*self))
} else {
formatter.write_str(self.artifact().label)
}
@@ -480,8 +499,8 @@ pub(crate) fn managed_artifacts(models_path: &Path) -> Vec<ManagedArtifact> {
MANAGED_ARTIFACTS
.into_iter()
.map(|id| {
if id == ManagedArtifactId::Qwen38FlashNext {
return qwen::managed_artifact(models_path);
if id.is_qwen() {
return qwen::managed_artifact(id, models_path);
}
let model = id.model();
let artifact = id.artifact();
@@ -527,8 +546,8 @@ pub(crate) fn artifact_download_progress(
id: ManagedArtifactId,
models_path: &Path,
) -> DownloadProgress {
if id == ManagedArtifactId::Qwen38FlashNext {
return qwen::download_progress(models_path);
if id.is_qwen() {
return qwen::download_progress(id, models_path);
}
let model = id.model();
let artifact = id.artifact();
@@ -560,8 +579,8 @@ pub(crate) fn artifact_verification_progress(
id: ManagedArtifactId,
verified: u64,
) -> DownloadProgress {
if id == ManagedArtifactId::Qwen38FlashNext {
return qwen::verification_progress(verified);
if id.is_qwen() {
return qwen::verification_progress(id, verified);
}
let artifact = id.artifact();
DownloadProgress {
+90 -26
View File
@@ -22,6 +22,7 @@ struct Manifest {
format: u32,
source: Source,
files: Vec<Artifact>,
excluded: Vec<Artifact>,
}
#[derive(Deserialize)]
@@ -33,6 +34,7 @@ struct Source {
#[derive(Deserialize)]
pub(super) struct Artifact {
pub(super) path: String,
#[serde(default)]
pub(super) role: String,
pub(super) size: u64,
pub(super) sha256: String,
@@ -61,6 +63,23 @@ pub(super) fn artifacts() -> &'static [Artifact] {
&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 {
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))
}
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> {
for artifact in artifacts() {
let artifact_path = path(root, artifact);
@@ -136,8 +161,8 @@ pub(super) fn validate_installed(root: &Path) -> Result<(), String> {
Ok(())
}
fn stored_bytes(root: &Path) -> u64 {
artifacts()
fn stored_bytes(id: ManagedArtifactId, root: &Path) -> u64 {
selected(id)
.iter()
.flat_map(|artifact| [path(root, artifact), partial_path(root, artifact)])
.filter_map(|path| path.metadata().ok())
@@ -145,8 +170,8 @@ fn stored_bytes(root: &Path) -> u64 {
.sum()
}
fn downloaded_bytes(root: &Path) -> u64 {
artifacts()
fn downloaded_bytes(id: ManagedArtifactId, root: &Path) -> u64 {
selected(id)
.iter()
.map(|artifact| {
if artifact_is_installed(root, artifact) {
@@ -161,8 +186,8 @@ fn downloaded_bytes(root: &Path) -> u64 {
.sum()
}
fn complete_payloads(root: &Path) -> bool {
artifacts().iter().all(|artifact| {
fn complete_payloads(id: ManagedArtifactId, root: &Path) -> bool {
selected(id).iter().all(|artifact| {
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 stored = stored_bytes(&root);
let state = if is_installed(&root) {
let stored = stored_bytes(id, &root);
let state = if selected(id).iter().all(|a| artifact_is_installed(&root, a)) {
ManagedArtifactState::Ready
} else if complete_payloads(&root) {
} else if complete_payloads(id, &root) {
ManagedArtifactState::NeedsVerification
} else if stored > 0 {
ManagedArtifactState::Partial
@@ -183,44 +208,45 @@ pub(super) fn managed_artifact(models_path: &Path) -> ManagedArtifact {
ManagedArtifactState::Missing
};
ManagedArtifact {
id: ManagedArtifactId::Qwen38FlashNext,
id,
stored,
expected: total_bytes(),
expected: selected(id).iter().map(|a| a.size).sum(),
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 downloaded = downloaded_bytes(&root);
let installed = is_installed(&root);
let verification = (!installed && complete_payloads(&root)).then_some(VerificationProgress {
verified: 0,
total: total_bytes(),
});
let downloaded = downloaded_bytes(id, &root);
let installed = selected(id).iter().all(|a| artifact_is_installed(&root, a));
let verification =
(!installed && complete_payloads(id, &root)).then_some(VerificationProgress {
verified: 0,
total: selected(id).iter().map(|a| a.size).sum(),
});
let phase = if installed {
DownloadPhase::Complete
} else if verification.is_some() {
DownloadPhase::Verifying(LABEL)
DownloadPhase::Verifying(label(id))
} else if downloaded > 0 {
DownloadPhase::Downloading(LABEL)
DownloadPhase::Downloading(label(id))
} else {
DownloadPhase::Pending(LABEL)
DownloadPhase::Pending(label(id))
};
DownloadProgress {
downloaded,
total: total_bytes(),
total: selected(id).iter().map(|a| a.size).sum(),
phase,
verification,
}
}
pub(super) fn verification_progress(verified: u64) -> DownloadProgress {
let total = total_bytes();
pub(super) fn verification_progress(id: ManagedArtifactId, verified: u64) -> DownloadProgress {
let total = selected(id).iter().map(|a| a.size).sum();
DownloadProgress {
downloaded: total,
total,
phase: DownloadPhase::Verifying(LABEL),
phase: DownloadPhase::Verifying(label(id)),
verification: Some(VerificationProgress {
verified: verified.min(total),
total,
@@ -336,3 +362,41 @@ mod tests {
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,
verified_bytes: &AtomicU64,
) -> Result<DownloadOutcome, String> {
if id == ManagedArtifactId::Qwen38FlashNext {
return download_qwen(models_path, cancel, verified_bytes);
if id.is_qwen() {
return download_qwen(id, models_path, cancel, verified_bytes);
}
download_artifact_with_cancel(
id.model(),
@@ -32,8 +32,8 @@ pub(crate) fn validate_managed_artifact(
cancel: &AtomicBool,
verified_bytes: &AtomicU64,
) -> Result<DownloadOutcome, String> {
if id == ManagedArtifactId::Qwen38FlashNext {
return validate_qwen(models_path, cancel, verified_bytes);
if id.is_qwen() {
return validate_qwen(id, models_path, cancel, verified_bytes);
}
let model = id.model();
let artifact = id.artifact();
@@ -73,8 +73,8 @@ pub(crate) fn delete_managed_artifact(
id: ManagedArtifactId,
models_path: &Path,
) -> Result<(), String> {
if id == ManagedArtifactId::Qwen38FlashNext {
return delete_qwen(models_path);
if id.is_qwen() {
return delete_qwen(id, models_path);
}
let model = id.model();
let artifact = id.artifact();
@@ -93,6 +93,7 @@ pub(crate) fn delete_managed_artifact(
}
fn download_qwen(
id: ManagedArtifactId,
models_path: &Path,
cancel: &AtomicBool,
verified_bytes: &AtomicU64,
@@ -100,7 +101,7 @@ fn download_qwen(
let root = qwen::root(models_path);
fs::create_dir_all(&root).map_err(|error| format!("{}: {error}", root.display()))?;
verified_bytes.store(0, Ordering::Relaxed);
for artifact in qwen::artifacts() {
for artifact in qwen::selected(id) {
if cancel.load(Ordering::Relaxed) {
return Ok(DownloadOutcome::Stopped);
}
@@ -145,26 +146,29 @@ fn download_qwen(
fs::write(qwen::verification_path(&root, artifact), &artifact.sha256)
.map_err(|error| error.to_string())?;
}
if let Err(error) = crate::engine::validate_qwen_artifacts(&root) {
clear_qwen_markers(&root)?;
if id == ManagedArtifactId::Qwen38FlashNext
&& let Err(error) = crate::engine::validate_qwen_artifacts(&root)
{
clear_qwen_markers(id, &root)?;
return Err(error);
}
Ok(DownloadOutcome::Complete)
}
fn validate_qwen(
id: ManagedArtifactId,
models_path: &Path,
cancel: &AtomicBool,
verified_bytes: &AtomicU64,
) -> Result<DownloadOutcome, String> {
let root = qwen::root(models_path);
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() {
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 partial = qwen::partial_path(&root, artifact);
let (source, promote) = if destination.exists() {
@@ -189,15 +193,17 @@ fn validate_qwen(
fs::write(qwen::verification_path(&root, artifact), &artifact.sha256)
.map_err(|error| error.to_string())?;
}
if let Err(error) = crate::engine::validate_qwen_artifacts(&root) {
clear_qwen_markers(&root)?;
if id == ManagedArtifactId::Qwen38FlashNext
&& let Err(error) = crate::engine::validate_qwen_artifacts(&root)
{
clear_qwen_markers(id, &root)?;
return Err(error);
}
Ok(DownloadOutcome::Complete)
}
fn clear_qwen_markers(root: &Path) -> Result<(), String> {
for artifact in qwen::artifacts() {
fn clear_qwen_markers(id: ManagedArtifactId, root: &Path) -> Result<(), String> {
for artifact in qwen::selected(id) {
let path = qwen::verification_path(root, artifact);
match fs::remove_file(&path) {
Ok(()) => {}
@@ -208,9 +214,9 @@ fn clear_qwen_markers(root: &Path) -> Result<(), String> {
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);
for artifact in qwen::artifacts() {
for artifact in qwen::selected(id) {
for path in [
qwen::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) {
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())),
}
}
@@ -480,7 +493,7 @@ mod tests {
);
assert!(ModelChoice::from_id("unknown").is_none());
assert_eq!(MODEL_CHOICES.len(), 5);
assert_eq!(MANAGED_ARTIFACTS.len(), 7);
assert_eq!(MANAGED_ARTIFACTS.len(), 8);
assert_eq!(
ModelChoice::from_id("qwen3.8-flash-next"),
Some(ModelChoice::Qwen38FlashNext)
@@ -600,20 +613,31 @@ mod tests {
))
.unwrap();
assert!(
validate_qwen(&models, &AtomicBool::new(false), &AtomicU64::new(0))
.unwrap_err()
.contains("mtp.safetensors is not downloaded")
validate_qwen(
ManagedArtifactId::Qwen38FlashNext,
&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();
let verified = AtomicU64::new(0);
assert_eq!(
validate_qwen(&models, &AtomicBool::new(false), &verified).unwrap(),
validate_qwen(
ManagedArtifactId::Qwen38FlashNext,
&models,
&AtomicBool::new(false),
&verified
)
.unwrap(),
DownloadOutcome::Complete
);
assert_eq!(verified.load(Ordering::Relaxed), qwen::total_bytes());
assert!(qwen::is_installed(&root));
assert!(crate::engine::validate_qwen_artifacts(&root).is_ok());
delete_qwen(&models).unwrap();
delete_qwen(ManagedArtifactId::Qwen38FlashNext, &models).unwrap();
assert!(!root.exists());
fs::remove_dir_all(models).unwrap();
}
+29
View File
@@ -38,6 +38,7 @@ struct Options {
prompt: String,
additional_prompts: Vec<String>,
input_file: Option<PathBuf>,
image_file: Option<PathBuf>,
models_path: PathBuf,
config_path: PathBuf,
plain_chat: bool,
@@ -68,6 +69,7 @@ impl Options {
let mut prompt_overridden = false;
let mut additional_prompts = Vec::new();
let mut input_file = None;
let mut image_file = None;
let mut models_path = crate::app::models_path();
let mut config_path = crate::app::config_path();
let mut context = None;
@@ -120,6 +122,7 @@ impl Options {
prompt_overridden = true;
}
}
"--image-file" => image_file = Some(value()?.into()),
"--input-file" => input_file = Some(value()?.into()),
"--models-dir" => models_path = value()?.into(),
"--config" => config_path = value()?.into(),
@@ -190,6 +193,7 @@ impl Options {
prompt,
additional_prompts,
input_file,
image_file,
models_path,
config_path,
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\
--prompt TEXT Chat turn; repeat for one ongoing conversation\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\
--context TOKENS Context size; otherwise use the model profile\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> {
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 Some(path) = &options.input_file else {
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) {
return Err("Integrated MTP is unavailable for the selected model.".into());
}
if self.keep_vision_loaded && model != ModelChoice::Glm53Flash {
return Err("Persistent vision weights are available only for GLM 5.3 Flash.".into());
if self.keep_vision_loaded && !model.supports_vision() {
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() {
return Err("DSpark is not available for the selected model.".into());