metrale_model_layers/

video_preprocess.rs

// SPDX-License-Identifier: MIT OR Apache-2.0

//! 2026-09-25: Video to patch tensor, the temporal counterpart of
//! [`crate::vision_preprocess`].
//!
//! Owner: model-layers (vision input).
//! Invariants:
//! - Every group of a `PreprocessedVideo` has the same `grid_h × grid_w`
//!   plane, from one target size chosen from the first kept frame.
//! - `preprocess_video` returns at least one full temporal group or an error.
//!
//! The flattened patch dimension is `C × temporal_patch_size × patch²`, so a
//! patch spans `tp` frames. A still image fills that axis by repeating itself
//! `tp` times (`preprocess_image`); a video of `n` sampled frames becomes
//! `grid_t = n / tp` temporal groups, each a full patch plane built from `tp`
//! consecutive frames. The encoder takes each group as one image. `grid_t` is
//! carried, not flattened away, because MRoPE advances T once per group of a
//! video and holds it constant across an image.
//!
//! The decoder is chosen by magic bytes, not the declared MIME: GIF decodes
//! in-process; everything else goes to the opt-in ffmpeg subprocess
//! (`video_decode_ffmpeg`).

use anyhow::{Context, Result, ensure};
use image::RgbImage;
use metrale_config::VisionConfig;

use crate::vision_preprocess::{MEAN, STD, decode_data_uri_bytes, target_size_for};

/// 2026-09-25: Frames per second to sample at when the caller gives no usable
/// rate. The Qwen3-VL-30B-A3B checkpoint's `video_preprocessor_config.json`
/// declares `fps: 2`.
pub const DEFAULT_FPS: f32 = 2.0;

/// 2026-09-25: Sampling floor and ceiling, as `min_frames: 4` and
/// `max_frames: 768` in the same `video_preprocessor_config.json`.
pub const DEFAULT_MIN_FRAMES: usize = 4;
pub const DEFAULT_MAX_FRAMES: usize = 768;

/// 2026-09-25: A decoded, ready-to-encode video.
pub struct PreprocessedVideo {
    /// 2026-09-25: One entry per temporal group, each shaped like a preprocessed
    /// still: `[grid_h * grid_w, C * tp * patch * patch]`.
    pub groups: Vec<Vec<f32>>,
    pub grid_t: usize,
    pub grid_h: usize,
    pub grid_w: usize,
}

/// 2026-09-25: Summarised rather than derived: the payload is megabytes of f32.
impl std::fmt::Debug for PreprocessedVideo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "PreprocessedVideo {{ grid_t: {}, grid_h: {}, grid_w: {}, groups: {} x {} f32 }}",
            self.grid_t,
            self.grid_h,
            self.grid_w,
            self.groups.len(),
            self.groups.first().map_or(0, Vec::len)
        )
    }
}

impl PreprocessedVideo {
    /// 2026-09-25: Merged tokens this video contributes: one per `merge × merge`
    /// block of patches, per temporal group.
    pub fn pad_count(&self, spatial_merge_size: usize) -> usize {
        let sms = spatial_merge_size.max(1);
        self.grid_t * (self.grid_h / sms) * (self.grid_w / sms)
    }
}

/// 2026-09-25: Pick which frame indices to keep so the clip plays at `fps`.
///
/// `native_fps` is the container's rate. Indices are the nearest to a uniform
/// grid over the clip, not every Nth frame, which quantises badly when the
/// ratio is not an integer. The count is clamped into `[min_frames,
/// max_frames]`, never above the frames that exist, and then down to a
/// multiple of `temporal_patch_size` (at least one group), because a partial
/// group cannot be encoded. A clip shorter than one group returns all of its
/// frames. A non-finite or non-positive rate counts as `DEFAULT_FPS`.
pub fn sample_indices(
    n_frames: usize,
    native_fps: f32,
    target_fps: f32,
    min_frames: usize,
    max_frames: usize,
    temporal_patch_size: usize,
) -> Vec<usize> {
    if n_frames == 0 {
        return Vec::new();
    }
    let tp = temporal_patch_size.max(1);
    let native_fps = if native_fps.is_finite() && native_fps > 0.0 {
        native_fps
    } else {
        DEFAULT_FPS
    };
    let target_fps = if target_fps.is_finite() && target_fps > 0.0 {
        target_fps
    } else {
        DEFAULT_FPS
    };

    let duration = n_frames as f32 / native_fps;
    let wanted = (duration * target_fps).round().max(1.0) as usize;

    // 2026-09-25: Never more frames than exist: repeating frames would add tokens
    // with no new information.
    let max_frames = max_frames.max(1);
    let min_frames = min_frames.max(1).min(max_frames);
    let wanted = wanted.clamp(min_frames, max_frames).min(n_frames);

    let wanted = (wanted / tp).max(1) * tp;
    let wanted = wanted.min((n_frames / tp).max(1) * tp).min(n_frames);

    if wanted >= n_frames {
        return (0..n_frames).collect();
    }
    // 2026-09-25: Uniform positions across the clip, nearest index.
    let mut out = Vec::with_capacity(wanted);
    for i in 0..wanted {
        let pos = if wanted == 1 {
            0.0
        } else {
            (i as f32) * ((n_frames - 1) as f32) / ((wanted - 1) as f32)
        };
        out.push((pos.round() as usize).min(n_frames - 1));
    }
    out
}

/// 2026-09-25: Decode every frame of a container, choosing a backend by what the
/// bytes are.
///
/// Returns the frames and the rate they represent. GIF is decoded in-process
/// and reports the container's average rate, so the caller still samples it.
/// ffmpeg samples during decode and reports `target_fps`, so the caller's
/// sampling applies only the frame-count clamps.
///
/// Dispatch is on magic bytes, not the declared MIME: an mp4 labelled
/// `video/gif` or `application/octet-stream` still reaches ffmpeg, and a GIF
/// labelled mp4 does not spawn a process.
pub fn decode_frames(
    data_uri: &str,
    target_fps: f32,
    ffmpeg: &crate::video_decode_ffmpeg::FfmpegPolicy,
) -> Result<(Vec<RgbImage>, f32)> {
    let (mime, bytes) = decode_data_uri_bytes(data_uri)?;
    ensure!(!bytes.is_empty(), "the video payload is empty");

    if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
        return decode_gif(&bytes);
    }

    // 2026-09-25: Everything else goes to the subprocess backend. Its error is
    // wrapped with the sniffed container name.
    let kind = sniff_container(&bytes, &mime);
    crate::video_decode_ffmpeg::decode_frames(&bytes, target_fps, ffmpeg)
        .with_context(|| format!("decoding {kind}"))
        .map(|f| (f, target_fps))
}

/// 2026-09-25: Best-effort container name for error messages only; nothing
/// branches on it.
fn sniff_container(bytes: &[u8], mime: &str) -> String {
    let by_magic = if bytes.len() > 12 && &bytes[4..8] == b"ftyp" {
        Some("an MP4/MOV container")
    } else if bytes.starts_with(&[0x1a, 0x45, 0xdf, 0xa3]) {
        Some("a Matroska/WebM container")
    } else if bytes.starts_with(b"RIFF") {
        Some("an AVI container")
    } else {
        None
    };
    match (by_magic, mime.is_empty()) {
        (Some(k), _) => k.to_string(),
        (None, false) => format!("a {mime} payload"),
        (None, true) => "an unrecognized container".to_string(),
    }
}

/// 2026-09-25: In-process GIF decode. The rate is frames over the summed
/// per-frame delays; when the delays sum to zero, the rate is `DEFAULT_FPS`.
fn decode_gif(bytes: &[u8]) -> Result<(Vec<RgbImage>, f32)> {
    use image::AnimationDecoder;
    use image::codecs::gif::GifDecoder;
    let decoder =
        GifDecoder::new(std::io::Cursor::new(bytes.to_vec())).context("not a decodable GIF")?;
    let frames = decoder
        .into_frames()
        .collect_frames()
        .context("failed to decode animation frames")?;
    ensure!(!frames.is_empty(), "the container decoded to zero frames");

    let total_ms: f64 = frames
        .iter()
        .map(|f| {
            let (num, den) = f.delay().numer_denom_ms();
            if den == 0 {
                0.0
            } else {
                num as f64 / den as f64
            }
        })
        .sum();
    let fps = if total_ms > 0.0 {
        (frames.len() as f64 * 1000.0 / total_ms) as f32
    } else {
        DEFAULT_FPS
    };

    let rgb: Vec<RgbImage> = frames
        .into_iter()
        .map(|f| image::DynamicImage::ImageRgba8(f.into_buffer()).to_rgb8())
        .collect();
    Ok((rgb, fps))
}

/// 2026-09-25: Full pipeline: a base64 `data:` URI holding an animated container
/// becomes temporal groups of patches.
pub fn preprocess_video(
    data_uri: &str,
    vcfg: &VisionConfig,
    max_pixels: Option<usize>,
    target_fps: f32,
    ffmpeg: &crate::video_decode_ffmpeg::FfmpegPolicy,
) -> Result<PreprocessedVideo> {
    ensure!(
        vcfg.patch_size > 0 && vcfg.spatial_merge_size > 0 && vcfg.temporal_patch_size > 0,
        "vision_config geometry is invalid (patch/merge/temporal size is 0)"
    );
    let (frames, native_fps) = decode_frames(data_uri, target_fps, ffmpeg)?;
    let tp = vcfg.temporal_patch_size;

    let keep = sample_indices(
        frames.len(),
        native_fps,
        target_fps,
        DEFAULT_MIN_FRAMES,
        DEFAULT_MAX_FRAMES,
        tp,
    );
    ensure!(!keep.is_empty(), "frame sampling selected no frames");

    // 2026-09-25: A clip shorter than one temporal group is refused rather than
    // padded into a still.
    ensure!(
        keep.len() >= tp,
        "video has {} usable frame(s) but temporal_patch_size is {tp}; a clip must carry at \
         least one full temporal group",
        keep.len()
    );

    // 2026-09-25: The target size is decided once, from the first kept frame, and
    // applied to every frame: the groups share one pad run whose token count
    // assumes a single grid.
    let first = &frames[keep[0]];
    let grid_unit = (vcfg.patch_size * vcfg.spatial_merge_size) as u32;
    let (th, tw) = target_size_for(first.height(), first.width(), grid_unit, max_pixels);

    let ps = vcfg.patch_size;
    let grid_h = (th as usize) / ps;
    let grid_w = (tw as usize) / ps;
    let grid_t = keep.len() / tp;
    let patch_dim = 3 * tp * ps * ps;
    let plane = grid_h * grid_w;

    let mut groups = Vec::with_capacity(grid_t);
    for g in 0..grid_t {
        // 2026-09-25: Resize each of this group's `tp` frames once, before the
        // patch loop reads them.
        let resized: Vec<RgbImage> = (0..tp)
            .map(|k| {
                let f = &frames[keep[g * tp + k]];
                image::imageops::resize(f, tw, th, image::imageops::FilterType::CatmullRom)
            })
            .collect();

        let mut pixels = vec![0.0f32; plane * patch_dim];
        for ph in 0..grid_h {
            for pw in 0..grid_w {
                let patch_idx = ph * grid_w + pw;
                for c in 0..3usize {
                    for (t, frame) in resized.iter().enumerate() {
                        for py in 0..ps {
                            for px in 0..ps {
                                let raw = frame
                                    .get_pixel((pw * ps + px) as u32, (ph * ps + py) as u32)[c]
                                    as f32
                                    / 255.0;
                                let off = c * (tp * ps * ps) + t * (ps * ps) + py * ps + px;
                                pixels[patch_idx * patch_dim + off] = (raw - MEAN[c]) / STD[c];
                            }
                        }
                    }
                }
            }
        }
        groups.push(pixels);
    }

    Ok(PreprocessedVideo {
        groups,
        grid_t,
        grid_h,
        grid_w,
    })
}

#[cfg(test)]
#[path = "video_preprocess_tests.rs"]
mod tests;