Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/supported-models.md
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ Implemented VLM variants include:
- FastVLM (`llava_qwen2` / `fastvlm`): Apple's low-latency VLM. A FastViTHD hybrid encoder runs entirely on channels-last maps: a conv stem, three RepMixer stages (depthwise token mixing plus a BatchNorm ConvFFN), two attention stages with channel LayerNorm and `head_dim` 32 self-attention, inter-stage large-kernel PatchEmbed downsamples, RepCPE position encoders, and a squeeze-excite `conv_exp` head. Each 1024x1024 pad-to-square image becomes a `(16, 16, 3072)` map flattened to 256 tokens, projected to the Qwen2 decoder width by an `mlp2x_gelu` MLP. The `<image>` placeholder is the fixed `-200` sentinel (not a vocabulary token); the runtime splices one sentinel per image, expands it to 256 tokens, and scatters the image embeddings (LLaVA merge). The text decoder is stock Qwen2, reused unchanged. The loader accepts both the genuine (`apple/FastVLM-0.5B`) and converted (`mlx-community/FastVLM-0.5B-bf16`) weight layouts. Best for fast image description and grounded chat.
- GLM-OCR (`glm_ocr`): document-OCR sibling of GLM-4V. A 24-block ViT (3D patch embedding, per-head q/k RMSNorm on the packed `cu_seqlens` attention, 2D vision RoPE, Conv2d spatial downsample, SwiGLU patch merger) feeds a 16-layer GLM-4 text decoder driven by full-width even/odd MRoPE (`rope_parameters` with `mrope_section [16, 24, 24]`, `partial_rotary_factor 1.0`). The tower has no learned position embedding or post-conv norm, and the loader drops the next-n prediction (MTP) layer. Patches are reordered from the processor's raster order into spatial-merge-window order so the rotary, downsample, and merged-token grid stay spatially aligned (OCR reads scrambled patches wrong). Best for plain text, tables, and formula recognition.
- Youtu-VL
- Kimi-VL and Kimi-VL 2.5 (`kimi_vl` / `kimi_k25`): MoonViT native-resolution vision encoder (Conv2d patch embedding, learned plus bicubically-interpolated 2D position embedding, a shared 2D rotary embedding, block-diagonal cross-image attention, and `spatial_merge_size` patch merging) feeding a `LayerNorm -> Linear -> GELU -> Linear` connector into a DeepSeek-V3-style MoE text backbone. Detected, loaded, and served end to end: the safetensors directory loader wires the MoonViT tower, the connector, and the DeepSeek-V3 MoE backbone; the native-resolution processor patchifies each image; and the runtime expands each `<|media_pad|>` placeholder into `(h/merge) * (w/merge)` tokens before the merged vision features replace them. Image path only; the Kimi-VL 2.5 3D MoonViT video (image plus video patch embedding) path is tracked as a separate follow-up.
- Kimi-VL and Kimi-VL 2.5 (`kimi_vl` / `kimi_k25`): MoonViT native-resolution vision encoder (Conv2d patch embedding, learned plus bicubically-interpolated 2D position embedding, a shared 2D rotary embedding, block-diagonal cross-image attention, and `spatial_merge_size` patch merging) feeding a `LayerNorm -> Linear -> GELU -> Linear` connector into a DeepSeek-V3-style MoE text backbone. Detected, loaded, and served end to end: the safetensors directory loader wires the MoonViT tower, the connector, and the DeepSeek-V3 MoE backbone; the native-resolution processor patchifies each image; and the runtime expands each `<|media_pad|>` placeholder into `(h/merge) * (w/merge)` tokens before the merged vision features replace them. Video is also supported, on both the CLI (`--video`) and the server (`video_url` content blocks): frames are extracted with `ffmpeg` (uniform sampling, default 2.0 fps, matching the Gemma 4 video path) and patchified per frame, and the encoder treats a clip as a 3D `(t, h, w)` patch grid rather than `t` independent images. A computed (not learned) temporal sinusoid is added to the tiled 2D position table, the 2D rotary tables repeat per frame, attention over a clip is one block-diagonal segment spanning all its frames, and the patch merger mean-pools over the temporal axis before the usual `spatial_merge_size` 2x2 merge, so the placeholder expansion `(h/merge) * (w/merge)` is unchanged by `t`. A single-frame clip is equivalent to the same frame processed as an image once the frame-0 temporal constant folds into the patch-embed bias.
- SmolVLM, SmolVLM2, and Idefics3 (`smolvlm` / `idefics3`): one shared Idefics3-family runtime, a SigLIP vision tower + pixel-shuffle token compression + a Llama text backbone (SmolLM2 for SmolVLM, Llama-3 for `Idefics3-8B-Llama3`). Both on-disk layouts load: `SmolVLMForConditionalGeneration` (`model.text_model.*` + a top-level `lm_head`) and `Idefics3ForConditionalGeneration` (the whole Llama-with-head nested under `language_model.*`). Images are tiled into square patches with a global thumbnail tile, and each `<image>` placeholder expands to `num_image_token` compressed feature tokens per tile.
- Idefics2 (`idefics2`): SigLIP vision tower + a perceiver-resampler connector (a `modality_projection` SwiGLU MLP into the text hidden size, then `n_latents` learned query slots refined by grouped-query cross-attention over the image patches) + a Mistral text backbone. Each image contributes a fixed `n_latents` (64) compressed feature tokens regardless of the patch grid, and the `<image>` placeholder is framed by `<fake_token_around_image>` and expanded to those tokens. This first port feeds a single full-resolution square tile per image; the optional `do_image_splitting` high-resolution tiling is a documented follow-up.
- MiniCPM-O
Expand Down
86 changes: 85 additions & 1 deletion src/commands/generate_vlm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,17 @@ fn print_preparation_summary(summary: VlmPreparationSummary) {
image_blocks, total_image_tokens
);
}
VlmPreparationSummary::KimiVLVideo {
media_blocks,
video_count,
frame_slots,
total_image_tokens,
} => {
println!(
"Kimi-VL: expanded {} media placeholder(s) across {} video(s) ({} sampled frame(s), {} total media tokens)",
media_blocks, video_count, frame_slots, total_image_tokens
);
}
VlmPreparationSummary::GraniteVision {
image_blocks,
total_image_tokens,
Expand Down Expand Up @@ -319,8 +330,17 @@ pub(crate) fn compute_vlm_embeddings(
target_fps,
);
}
if let LoadedModel::KimiVL(kimi) = model {
return compute_kimi_vl_video_embeddings(
kimi,
prompt_tokens,
image_paths,
video_paths,
target_fps,
);
}
return Err(anyhow::anyhow!(
"--video input is currently only supported by Gemma4 VLMs"
"--video input is currently only supported by Gemma 4 and Kimi-VL VLMs"
));
}

Expand Down Expand Up @@ -838,6 +858,70 @@ fn compute_gemma4_video_embeddings(
Ok(Some(embeddings))
}

/// Compute video embeddings (and optional companion image embeddings) for the
/// Kimi-VL / Kimi-VL 2.5 (`kimi_k25`) MoonViT model.
///
/// Decodes each video via `mlxcel::video::load_videos` (subprocess `ffmpeg`,
/// default 2.0 fps), then routes the decoded frames and any `--image` companions
/// through the shared [`mlxcel::vlm_runtime::compute_kimi_vl_media_embeddings`]
/// helper, which patchifies per frame, builds the mixed `(t, h, w)` / `(h, w)`
/// media-grid list, expands `<|media_pad|>` placeholders in media order, and
/// runs the MoonViT tower with temporal position embeddings + temporal pooling.
fn compute_kimi_vl_video_embeddings(
kimi: &mlxcel::vision::KimiVLModel,
prompt_tokens: &mut Vec<i32>,
image_paths: &[PathBuf],
video_paths: &[PathBuf],
target_fps: f64,
) -> Result<Option<InputEmbeddings>> {
if !video::ffmpeg_available() {
return Err(anyhow::anyhow!(
"Video input requires `ffmpeg` on PATH. Install ffmpeg (e.g. `brew install ffmpeg` \
on macOS or `apt install ffmpeg` on Linux) and retry."
));
}

// Decode the videos. `target_fps == 0` is rejected by `smart_nframes`,
// so guard here with a clean error.
let videos = video::load_videos(video_paths, Some(target_fps), None)
.map_err(|err| anyhow::anyhow!("Failed to load video(s): {}", err))?;
let frame_slots: usize = videos.iter().map(Vec::len).sum();
println!(
"Loaded {} video(s) ({} total frames after sampling).",
videos.len(),
frame_slots
);

// Optional companion images (e.g. user passes both --image and --video).
let images: Vec<image::DynamicImage> = image_paths
.iter()
.map(|path| {
image::open(path).map_err(|e| anyhow::anyhow!("Failed to load image {:?}: {}", path, e))
})
.collect::<Result<Vec<_>>>()?;
if !images.is_empty() {
println!("Loaded {} image(s).", images.len());
}

let (embeddings, stats) = mlxcel::vlm_runtime::compute_kimi_vl_media_embeddings(
kimi,
prompt_tokens,
&images,
&videos,
)?;

if let Some(stats) = stats {
print_preparation_summary(VlmPreparationSummary::KimiVLVideo {
media_blocks: stats.image_blocks,
video_count: videos.len(),
frame_slots,
total_image_tokens: stats.total_image_tokens,
});
}

Ok(Some(embeddings))
}

/// Compute video embeddings (and optional preceding image embeddings)
/// for the Gemma 4 Unified (encoder-free) model.
///
Expand Down
111 changes: 96 additions & 15 deletions src/multimodal/kimi_vl_prompt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,21 @@
//! in place.
//!
//! When the rendered prompt carries no `<|media_pad|>` placeholder (a bare CLI
//! image), one expanded run per image is spliced in after the first token,
//! mirroring the Youtu-VL / InternVL fallback.
//! image or video), one expanded run per media item is spliced in after the
//! first token, mirroring the Youtu-VL / InternVL fallback.
//!
//! Video clips (Kimi-VL 2.5) use the same `<|media_pad|>` placeholder and the
//! same per-item count `(h / merge) * (w / merge)`. The count is independent of
//! the frame count `t` because the MoonViT merger mean-pools all frames of a
//! clip to one spatial token map before the spatial merge. A request may mix
//! images and videos; grid entries expand in media order.

use crate::vision::encoders::kimi_vl::KimiMediaGrid;

/// Statistics describing the Kimi-VL media-token insertion/expansion.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct InsertedKimiVlTokens {
/// Number of media items (images + video clips) that were expanded.
pub image_blocks: usize,
pub total_image_tokens: i32,
}
Expand All @@ -47,31 +56,61 @@ pub fn insert_kimi_vl_image_tokens(
spatial_merge_size: i32,
media_placeholder_token_id: i32,
) -> Option<InsertedKimiVlTokens> {
if prompt_tokens.is_empty() || grid_shapes.is_empty() || spatial_merge_size <= 0 {
if grid_shapes.is_empty() {
return None;
}
let media_grids: Vec<KimiMediaGrid> = grid_shapes
.iter()
.map(|&(h, w)| KimiMediaGrid::Image { h, w })
.collect();
insert_kimi_vl_media_tokens(
prompt_tokens,
&media_grids,
spatial_merge_size,
media_placeholder_token_id,
)
}

/// Insert (or expand) Kimi-VL media-placeholder runs for a mixed list of image
/// and video grids (media order).
///
/// The per-item token count is `(h / merge) * (w / merge)` for both images and
/// videos (the video path collapses `t` frames to one spatial map before the
/// spatial merge, so the count does not depend on `t`).
///
/// Returns `None` when there is nothing to do (empty prompt, no media, or a
/// zero merge size).
pub fn insert_kimi_vl_media_tokens(
prompt_tokens: &mut Vec<i32>,
media_grids: &[KimiMediaGrid],
spatial_merge_size: i32,
media_placeholder_token_id: i32,
) -> Option<InsertedKimiVlTokens> {
if prompt_tokens.is_empty() || media_grids.is_empty() || spatial_merge_size <= 0 {
return None;
}

let per_image_counts: Vec<i32> = grid_shapes
let per_item_counts: Vec<i32> = media_grids
.iter()
.map(|&(h, w)| (h / spatial_merge_size) * (w / spatial_merge_size))
.map(|grid| grid.merged_count(spatial_merge_size))
.collect();
let total_image_tokens: i32 = per_image_counts.iter().sum();
let image_blocks = grid_shapes.len();
let total_image_tokens: i32 = per_item_counts.iter().sum();
let image_blocks = media_grids.len();

// Case 1: the prompt already carries one <|media_pad|> per image. Expand
// each in place into its full run of placeholder tokens.
// Case 1: the prompt already carries one <|media_pad|> per media item.
// Expand each in place into its full run of placeholder tokens.
let placeholder_count = prompt_tokens
.iter()
.filter(|&&t| t == media_placeholder_token_id)
.count();
if placeholder_count > 0 {
let mut expanded = Vec::with_capacity(prompt_tokens.len() + total_image_tokens as usize);
let mut image_idx = 0usize;
let mut item_idx = 0usize;
for &token in prompt_tokens.iter() {
if token == media_placeholder_token_id && image_idx < per_image_counts.len() {
let count = per_image_counts[image_idx].max(0) as usize;
if token == media_placeholder_token_id && item_idx < per_item_counts.len() {
let count = per_item_counts[item_idx].max(0) as usize;
expanded.extend(std::iter::repeat_n(media_placeholder_token_id, count));
image_idx += 1;
item_idx += 1;
} else {
expanded.push(token);
}
Expand All @@ -83,9 +122,10 @@ pub fn insert_kimi_vl_image_tokens(
});
}

// Case 2: no placeholder — splice one run per image after the first token.
// Case 2: no placeholder — splice one run per media item after the first
// token, in media order.
let mut runs: Vec<i32> = Vec::with_capacity(total_image_tokens as usize);
for &count in &per_image_counts {
for &count in &per_item_counts {
runs.extend(std::iter::repeat_n(
media_placeholder_token_id,
count.max(0) as usize,
Expand Down Expand Up @@ -148,4 +188,45 @@ mod tests {
let mut prompt = vec![1, 2];
assert!(insert_kimi_vl_image_tokens(&mut prompt, &[], 2, MEDIA_PAD).is_none());
}

#[test]
fn video_grid_expands_independent_of_t() {
// A video (t, h, w) expands its single <|media_pad|> to (h/2)*(w/2)
// tokens regardless of t. Try t = 1 and t = 5 for the same (4, 4) grid;
// both must yield 4 media tokens.
for t in [1, 5, 8] {
let mut prompt = vec![1i32, MEDIA_PAD, 300];
let stats = insert_kimi_vl_media_tokens(
&mut prompt,
&[KimiMediaGrid::Video { t, h: 4, w: 4 }],
2,
MEDIA_PAD,
)
.unwrap();
assert_eq!(stats.image_blocks, 1);
assert_eq!(stats.total_image_tokens, 4, "t={t} must give 4 tokens");
assert_eq!(prompt[0], 1);
assert_eq!(&prompt[1..5], &[MEDIA_PAD; 4]);
assert_eq!(prompt[5], 300);
}
}

#[test]
fn mixed_image_and_video_expand_in_media_order() {
// Media order: image (4,4) -> 4 tokens, then video (t=3, 2, 2) -> 1
// token. Two placeholders expand in order.
let mut prompt = vec![1i32, MEDIA_PAD, MEDIA_PAD, 9];
let grids = [
KimiMediaGrid::Image { h: 4, w: 4 },
KimiMediaGrid::Video { t: 3, h: 2, w: 2 },
];
let stats = insert_kimi_vl_media_tokens(&mut prompt, &grids, 2, MEDIA_PAD).unwrap();
assert_eq!(stats.image_blocks, 2);
// image: (4/2)*(4/2)=4, video: (2/2)*(2/2)=1, total 5.
assert_eq!(stats.total_image_tokens, 5);
assert_eq!(prompt[0], 1);
assert_eq!(&prompt[1..5], &[MEDIA_PAD; 4]); // image run
assert_eq!(prompt[5], MEDIA_PAD); // video run (1 token)
assert_eq!(prompt[6], 9);
}
}
Loading