Add Qwen vision loading and Metal inference
This commit is contained in:
@@ -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
@@ -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(
|
||||||
|
|||||||
@@ -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
@@ -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")]
|
||||||
|
|||||||
@@ -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),
|
||||||
|
|||||||
@@ -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
@@ -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 {
|
||||||
|
|||||||
@@ -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 {
|
||||||
|
|||||||
@@ -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],
|
||||||
)?;
|
)?;
|
||||||
dispatch_geometry(
|
let zero = 0_i32;
|
||||||
&format!("steel_gemm_splitk_accum_{dtype_name}_float32"),
|
let one = 1_i32;
|
||||||
&[
|
let scale = 1_f32;
|
||||||
|
let mut accum = vec![
|
||||||
tensor(0, &intermediate).input(parts as usize * stride as usize),
|
tensor(0, &intermediate).input(parts as usize * stride as usize),
|
||||||
output(1),
|
output(1),
|
||||||
bytes(2, &parts),
|
bytes(2, &parts),
|
||||||
bytes(3, &stride),
|
bytes(3, &stride),
|
||||||
bytes(4, &n),
|
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{}",
|
||||||
|
if bias.is_some() { "_axbpy" } else { "" }
|
||||||
|
),
|
||||||
|
&accum,
|
||||||
&[],
|
&[],
|
||||||
[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, ¶ms)];
|
let mut bindings = vec![ai(0), bi(1), output(3), bytes(4, ¶ms)];
|
||||||
|
#[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)),
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
@@ -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)?,
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -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());
|
||||||
|
|||||||
@@ -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();
|
||||||
|
}
|
||||||
|
|||||||
@@ -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,9 +159,16 @@ impl LiveSession {
|
|||||||
o.stops = &stop_ids;
|
o.stops = &stop_ids;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
let saved = [&self.current, &self.prefix]
|
let saved = [
|
||||||
|
(&self.current, &self.current_images),
|
||||||
|
(&self.prefix, &self.prefix_images),
|
||||||
|
]
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.flatten()
|
.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))
|
.filter(|s| prompt.starts_with(&s.token_ids))
|
||||||
.max_by_key(|s| s.token_ids.len());
|
.max_by_key(|s| s.token_ids.len());
|
||||||
let mut captured_prefix = None;
|
let mut captured_prefix = None;
|
||||||
@@ -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(),
|
||||||
|
|||||||
@@ -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> {
|
||||||
|
|||||||
@@ -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, ¶ms), 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, ¶ms));
|
||||||
|
dispatch_specialized(
|
||||||
|
"steel_attention_bfloat16_bq32_bk32_bd80_wm4_wn1_maskbfloat16",
|
||||||
|
&bindings,
|
||||||
|
&constants,
|
||||||
|
[((n + 31) / 32) as u32, 16, 1],
|
||||||
|
[32, 4, 1],
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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
@@ -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)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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
@@ -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 {
|
||||||
|
|||||||
+88
-24
@@ -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 =
|
||||||
|
(!installed && complete_payloads(id, &root)).then_some(VerificationProgress {
|
||||||
verified: 0,
|
verified: 0,
|
||||||
total: total_bytes(),
|
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();
|
||||||
|
}
|
||||||
|
|||||||
+46
-22
@@ -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(
|
||||||
|
ManagedArtifactId::Qwen38FlashNext,
|
||||||
|
&models,
|
||||||
|
&AtomicBool::new(false),
|
||||||
|
&AtomicU64::new(0)
|
||||||
|
)
|
||||||
.unwrap_err()
|
.unwrap_err()
|
||||||
.contains("mtp.safetensors is not downloaded")
|
.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();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
@@ -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());
|
||||||
|
|||||||
@@ -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)
|
||||||
@@ -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)
|
||||||
Reference in New Issue
Block a user