inferencelayer 0.2.3

//! GLM-OCR vision tower (`GlmOcrForConditionalGeneration.model.visual`, the GLM-4.6V family) —
//! native-resolution ViT on the CPU, in pure Rust.
//!
//! Shares the Qwen3.5-VL tower's scaffolding ([`crate::vision`]: the integer-exact PIL-bicubic
//! resize, block-major patchify, 2-D rope tables, prepacked GEMMs) but every block-internal module
//! differs, which is why it is its own tower rather than a flag on [`crate::vision::VisionTower`]:
//!
//! * **No learned position embedding at all** — rotary only. (Qwen interpolates a 48×48 table.)
//! * Block norms are weight-only **RMSNorm (eps 1e-5)**, not LayerNorm.
//! * Attention has **per-head q/k RMSNorm (dim 64), applied BEFORE the rope** — Qwen's tower has none.
//! * The MLP is a **gated SwiGLU with biases** (1024→4096→1024), not a plain 2-layer fc.
//! * The merge head is completely different: `post_layernorm` (RMS) → a **2×2/2 Conv2d downsample**
//!   1024→1536 over each merged block (channel-major — loaded here as a permuted Linear over the
//!   4-consecutive-token concat the block-major order provides) → `proj` 1536→1536 → LayerNorm(+bias)
//!   → exact GELU → SwiGLU 1536→4608→1536 (bias-less).
//! * Preprocessing: factor **28** smart-resize with GLM's rules (small-side pre-upscale with `int()`
//!   truncation; the pixel budget folds in the temporal factor 2), CLIP mean/std normalization.
//!
//! Gated stage-by-stage against `tests/fixtures/export_glmocr.py`'s oracle (transformers glm_ocr),
//! so the first mismatch names the layer that is wrong.

use anyhow::{Context, Result};
use rayon::prelude::*;
use std::path::Path;

use crate::encoder_weights::Act;
use crate::vision::{
    ImagePatches, Linear, Norm, VisionConfig, act_inplace, layer_norm, patchify_with, resize_rgb8,
    rope_tables_2d,
};
use crate::weights::LazySt;

/// OpenAI-CLIP normalization — what `Glm46VImageProcessor` applies.
pub const CLIP_MEAN: [f32; 3] = [0.481_454_66, 0.457_827_5, 0.408_210_73];
pub const CLIP_STD: [f32; 3] = [0.268_629_54, 0.261_302_6, 0.275_777_1];

/// GLM's `smart_resize` (`image_processing_glm46v.py`): differs from Qwen's in three load-bearing
/// ways — a small-side pre-upscale with **`int()` truncation** (not a clamp), the pixel budget
/// comparing **`t_bar·h_bar·w_bar`** (the temporal factor 2 halves the effective spatial budget for
/// stills), and beta computed from the (possibly upscaled) ORIGINAL sides.
pub fn glm_smart_resize(
    height: usize,
    width: usize,
    factor: usize,
    temporal: usize,
    min_pixels: usize,
    max_pixels: usize,
) -> Result<(usize, usize)> {
    let f = factor as f64;
    let (mut hf, mut wf) = (height as f64, width as f64);
    // Small-side pre-upscale: scale = max(f/h, f/w); h = int(h*scale) — TRUNCATION, per the reference.
    if height < factor || width < factor {
        let scale = (f / hf).max(f / wf);
        hf = (hf * scale).trunc();
        wf = (wf * scale).trunc();
    }
    let ratio = hf.max(wf) / hf.min(wf);
    anyhow::ensure!(
        ratio <= 200.0,
        "absolute aspect ratio must be <= 200, got {ratio:.1}"
    );
    // Python round is half-to-even (same hazard as the Qwen port, factor 28 here).
    let round_half_even = |x: f64| -> f64 {
        let lo = x.floor();
        match (x - lo).partial_cmp(&0.5).expect("finite") {
            std::cmp::Ordering::Greater => lo + 1.0,
            std::cmp::Ordering::Less => lo,
            std::cmp::Ordering::Equal if (lo as i64) % 2 == 0 => lo,
            std::cmp::Ordering::Equal => lo + 1.0,
        }
    };
    let mut h = (round_half_even(hf / f) * f) as usize;
    let mut w = (round_half_even(wf / f) * f) as usize;
    let t_bar = temporal; // round(t/t)*t = t for stills duplicated to `temporal` frames

    if t_bar * h * w > max_pixels {
        let beta = ((t_bar as f64) * hf * wf / max_pixels as f64).sqrt();
        h = (((hf / beta) / f).floor() as usize * factor).max(factor);
        w = (((wf / beta) / f).floor() as usize * factor).max(factor);
    } else if t_bar * h * w < min_pixels {
        let beta = (min_pixels as f64 / ((t_bar as f64) * hf * wf)).sqrt();
        h = ((hf * beta) / f).ceil() as usize * factor;
        w = ((wf * beta) / f).ceil() as usize * factor;
    }
    Ok((h, w))
}

/// GLM preprocessing of one interleaved RGB8 image: GLM smart-resize → the shared integer-exact
/// bicubic (u8) → CLIP normalize → block-major patchify.
pub fn glm_preprocess_rgb8(
    rgb: &[u8],
    width: usize,
    height: usize,
    cfg: &VisionConfig,
) -> Result<ImagePatches> {
    anyhow::ensure!(
        rgb.len() == width * height * cfg.in_channels,
        "rgb buffer is {} bytes, expected {}×{}×{}",
        rgb.len(),
        width,
        height,
        cfg.in_channels
    );
    let factor = cfg.patch * cfg.merge;
    let (rh, rw) = glm_smart_resize(
        height,
        width,
        factor,
        cfg.temporal,
        cfg.min_pixels,
        cfg.max_pixels,
    )?;
    let resized = if (rh, rw) == (height, width) {
        rgb.to_vec()
    } else {
        resize_rgb8(rgb, width, height, rw, rh)
    };
    Ok(patchify_with(&resized, rh, rw, cfg, CLIP_MEAN, CLIP_STD))
}

/// Decode encoded image bytes and GLM-preprocess. Pure Rust (the `image` crate).
pub fn glm_preprocess_bytes(bytes: &[u8], cfg: &VisionConfig) -> Result<ImagePatches> {
    let img = image::load_from_memory(bytes).context("decode image")?;
    let rgb = img.to_rgb8();
    let (w, h) = (rgb.width() as usize, rgb.height() as usize);
    glm_preprocess_rgb8(rgb.as_raw(), w, h, cfg)
}

/// Weight-only RMSNorm over rows of width `h`, in place.
fn rms_norm(x: &mut [f32], h: usize, w: &[f32], eps: f32) {
    x.par_chunks_mut(h).for_each(|row| {
        let ms = row.iter().map(|v| v * v).sum::<f32>() / h as f32;
        let inv = 1.0 / (ms + eps).sqrt();
        for (j, v) in row.iter_mut().enumerate() {
            *v = *v * inv * w[j];
        }
    });
}

pub(crate) struct GlmBlock {
    pub(crate) norm1_w: Vec<f32>,
    pub(crate) qkv: Linear,
    pub(crate) q_norm_w: Vec<f32>, // [head_dim]
    pub(crate) k_norm_w: Vec<f32>,
    pub(crate) proj: Linear,
    pub(crate) norm2_w: Vec<f32>,
    pub(crate) gate: Linear,
    pub(crate) up: Linear,
    pub(crate) down: Linear,
}

/// The GLM-OCR vision tower.
pub struct GlmVisionTower {
    pub cfg: VisionConfig,
    pub(crate) patch: Linear,
    pub(crate) blocks: Vec<GlmBlock>,
    pub(crate) post_ln_w: Vec<f32>,
    /// The 2×2/2 conv downsample, loaded as a Linear over the 4-consecutive-token concat:
    /// `W_lin[o, (i·2+j)·hidden + c] = W_conv[o, c, i, j]`.
    pub(crate) downsample: Linear,
    pub(crate) merger_proj: Linear,
    pub(crate) merger_post_norm: Norm,
    pub(crate) merger_gate: Linear,
    pub(crate) merger_up: Linear,
    pub(crate) merger_down: Linear,
}

impl GlmVisionTower {
    /// Load `model.visual.*` from a GLM-OCR checkpoint directory (the HF snapshot, read directly).
    pub fn load(dir: &Path) -> Result<Self> {
        let cfg_json: serde_json::Value = serde_json::from_slice(
            &std::fs::read(dir.join("config.json"))
                .with_context(|| format!("read {}", dir.join("config.json").display()))?,
        )?;
        let v = cfg_json
            .get("vision_config")
            .context("config.json vision_config")?;
        let g = |k: &str| -> Result<usize> {
            v.get(k)
                .and_then(|x| x.as_u64())
                .map(|x| x as usize)
                .with_context(|| format!("vision_config.{k}"))
        };
        let mut cfg = VisionConfig {
            depth: g("depth")?,
            hidden: g("hidden_size")?,
            heads: g("num_heads")?,
            intermediate: g("intermediate_size")?,
            out_hidden: g("out_hidden_size")?,
            patch: g("patch_size")?,
            merge: g("spatial_merge_size")?,
            temporal: g("temporal_patch_size")?,
            in_channels: 3,
            grid_side: 0, // no learned position table
            eps: v
                .get("rms_norm_eps")
                .and_then(|x| x.as_f64())
                .unwrap_or(1e-5) as f32,
            act: Act::Silu,
            rope_theta: 10_000.0,
            min_pixels: 12_544,
            max_pixels: 9_633_792,
        };
        // GLM ships `preprocessor_config.json` with the bounds at `size.{shortest,longest}_edge`.
        let pcfg = dir.join("preprocessor_config.json");
        if pcfg.exists() {
            let p: serde_json::Value = serde_json::from_slice(&std::fs::read(&pcfg)?)?;
            if let Some(size) = p.get("size") {
                let get = |k: &str| size.get(k).and_then(|x| x.as_u64()).map(|x| x as usize);
                if let (Some(lo), Some(hi)) = (get("shortest_edge"), get("longest_edge")) {
                    cfg.min_pixels = lo;
                    cfg.max_pixels = hi;
                }
            }
        }

        let st = LazySt::open(dir)?;
        let p = "model.visual";
        let hid = cfg.hidden;

        let mut blocks = Vec::with_capacity(cfg.depth);
        for i in 0..cfg.depth {
            let b = format!("{p}.blocks.{i}");
            blocks.push(GlmBlock {
                norm1_w: st.tensor_f32(&format!("{b}.norm1.weight"))?,
                qkv: Linear::load(&st, &format!("{b}.attn.qkv"))?,
                q_norm_w: st.tensor_f32(&format!("{b}.attn.q_norm.weight"))?,
                k_norm_w: st.tensor_f32(&format!("{b}.attn.k_norm.weight"))?,
                proj: Linear::load(&st, &format!("{b}.attn.proj"))?,
                norm2_w: st.tensor_f32(&format!("{b}.norm2.weight"))?,
                gate: Linear::load(&st, &format!("{b}.mlp.gate_proj"))?,
                up: Linear::load(&st, &format!("{b}.mlp.up_proj"))?,
                down: Linear::load(&st, &format!("{b}.mlp.down_proj"))?,
            });
        }

        // Downsample conv [out, hid, 2, 2] → Linear [out, 4·hid] over the (i,j)-major concat.
        let conv_w = st.tensor_f32(&format!("{p}.downsample.weight"))?;
        let conv_b = st.tensor_f32(&format!("{p}.downsample.bias"))?;
        let out_h = conv_b.len();
        let m = cfg.merge;
        anyhow::ensure!(
            conv_w.len() == out_h * hid * m * m,
            "downsample weight shape mismatch"
        );
        let mut lin = vec![0f32; out_h * m * m * hid];
        for o in 0..out_h {
            for c in 0..hid {
                for i in 0..m {
                    for j in 0..m {
                        lin[o * (m * m * hid) + (i * m + j) * hid + c] =
                            conv_w[((o * hid + c) * m + i) * m + j];
                    }
                }
            }
        }
        let downsample = Linear::from_parts(lin, Some(conv_b), out_h, m * m * hid);

        Ok(Self {
            patch: Linear::load(&st, &format!("{p}.patch_embed.proj"))?,
            post_ln_w: st.tensor_f32(&format!("{p}.post_layernorm.weight"))?,
            downsample,
            merger_proj: Linear::load_shaped(&st, &format!("{p}.merger.proj"), cfg.out_hidden)?,
            merger_post_norm: Norm::load(&st, &format!("{p}.merger.post_projection_norm"))?,
            merger_gate: Linear::load_shaped(
                &st,
                &format!("{p}.merger.gate_proj"),
                cfg.out_hidden * 3,
            )?,
            merger_up: Linear::load_shaped(
                &st,
                &format!("{p}.merger.up_proj"),
                cfg.out_hidden * 3,
            )?,
            merger_down: Linear::load_shaped(
                &st,
                &format!("{p}.merger.down_proj"),
                cfg.out_hidden,
            )?,
            blocks,
            cfg,
        })
    }

    /// Run the tower over preprocessed images → `[total_merged_tokens, out_hidden]`.
    pub fn forward(&self, images: &[ImagePatches]) -> Result<Vec<f32>> {
        let mut out = Vec::new();
        for img in images {
            out.extend(self.run(img, false)?.1);
        }
        Ok(out)
    }

    /// [`Self::forward`] for one image, capturing intermediates for the parity gate: `[0]` is the
    /// patch embed, `[1+i]` block `i`'s output, then (post_layernorm, downsample) separately.
    #[allow(clippy::type_complexity)]
    pub fn debug_forward(
        &self,
        img: &ImagePatches,
    ) -> Result<(Vec<Vec<f32>>, Vec<f32>, Vec<f32>, Vec<f32>)> {
        let (states, out) = self.run(img, true)?;
        // run() with capture=true also stashes post_ln and downsample at the END of states.
        let mut states = states;
        let ds = states.pop().expect("downsample state");
        let pl = states.pop().expect("post_ln state");
        Ok((states, pl, ds, out))
    }

    fn run(&self, img: &ImagePatches, capture: bool) -> Result<(Vec<Vec<f32>>, Vec<f32>)> {
        let cfg = &self.cfg;
        let (hid, heads, hd) = (cfg.hidden, cfg.heads, cfg.head_dim());
        let n = img.num_patches();
        anyhow::ensure!(
            img.patches.len() == n * cfg.patch_dim(),
            "patch buffer {} does not match grid {:?}",
            img.patches.len(),
            img.grid
        );

        // Patch embed (Conv3d as Linear over the flattened patch). NO position embedding.
        let mut x = vec![0f32; n * hid];
        self.patch.forward(&mut x, &img.patches);

        let (cos, sin) = rope_tables_2d(cfg, img.grid);

        let mut qkv = vec![0f32; n * 3 * hid];
        let mut attn = vec![0f32; n * hid];
        let mut gate = vec![0f32; n * cfg.intermediate];
        let mut up = vec![0f32; n * cfg.intermediate];
        let mut normed = vec![0f32; n * hid];
        let mut states: Vec<Vec<f32>> = Vec::new();
        if capture {
            states.push(x.clone());
        }

        for (bi, blk) in self.blocks.iter().enumerate() {
            // ---- attention (pre-RMSNorm, per-head qk-norm BEFORE rope) ----
            normed.copy_from_slice(&x);
            rms_norm(&mut normed, hid, &blk.norm1_w, cfg.eps);
            if capture && bi == 0 {
                states.push(normed.clone()); // b0_norm1
            }
            blk.qkv.forward(&mut qkv, &normed);
            glm_attention(
                &mut attn,
                &qkv,
                &cos,
                &sin,
                n,
                heads,
                hd,
                &blk.q_norm_w,
                &blk.k_norm_w,
                cfg.eps,
            );
            blk.proj.forward(&mut normed, &attn);
            if capture && bi == 0 {
                states.push(normed.clone()); // b0_attn (module output, pre-residual)
            }
            x.par_iter_mut()
                .zip(normed.par_iter())
                .for_each(|(a, b)| *a += b);

            // ---- gated MLP ----
            normed.copy_from_slice(&x);
            rms_norm(&mut normed, hid, &blk.norm2_w, cfg.eps);
            blk.gate.forward(&mut gate, &normed);
            blk.up.forward(&mut up, &normed);
            act_inplace(&mut gate, Act::Silu);
            gate.par_iter_mut()
                .zip(up.par_iter())
                .for_each(|(g, u)| *g *= u);
            blk.down.forward(&mut normed, &gate);
            if capture && bi == 0 {
                states.push(normed.clone()); // b0_mlp (module output, pre-residual)
            }
            x.par_iter_mut()
                .zip(normed.par_iter())
                .for_each(|(a, b)| *a += b);

            if capture {
                states.push(x.clone());
            }
        }

        // ---- post_layernorm (RMS) ----
        rms_norm(&mut x, hid, &self.post_ln_w, cfg.eps);
        if capture {
            states.push(x.clone());
        }

        // ---- downsample: 4 consecutive (block-major) patches = one 2×2 block → conv-as-linear ----
        let unit = cfg.merge_unit();
        anyhow::ensure!(
            n.is_multiple_of(unit),
            "patch count {n} not a multiple of merge²"
        );
        let tokens = n / unit;
        let mut ds = vec![0f32; tokens * cfg.out_hidden];
        self.downsample.forward(&mut ds, &x); // [n, hid] row-major IS [n/unit, unit·hid]
        if capture {
            states.push(ds.clone());
        }

        // ---- GLM merger: proj → LayerNorm(+bias) → exact GELU → SwiGLU ----
        let oh = cfg.out_hidden;
        let mut mproj = vec![0f32; tokens * oh];
        self.merger_proj.forward(&mut mproj, &ds);
        layer_norm(&mut mproj, oh, &self.merger_post_norm, 1e-5);
        act_inplace(&mut mproj, Act::GeluErf);
        let inner = self.merger_gate.n;
        let mut mg = vec![0f32; tokens * inner];
        let mut mu = vec![0f32; tokens * inner];
        self.merger_gate.forward(&mut mg, &mproj);
        self.merger_up.forward(&mut mu, &mproj);
        act_inplace(&mut mg, Act::Silu);
        mg.par_iter_mut()
            .zip(mu.par_iter())
            .for_each(|(g, u)| *g *= u);
        let mut out = vec![0f32; tokens * oh];
        self.merger_down.forward(&mut out, &mg);
        Ok((states, out))
    }
}

/// Full self-attention over one image's patches: per-head RMS q/k-norm, then the 2-D rope
/// (contiguous-halves rotation), then softmax(QKᵀ/√d)·V. Mirrors `vision.rs`'s attention with the
/// qk-norm inserted.
#[allow(clippy::too_many_arguments)]
fn glm_attention(
    out: &mut [f32],
    qkv: &[f32],
    cos: &[f32],
    sin: &[f32],
    n: usize,
    heads: usize,
    hd: usize,
    q_norm_w: &[f32],
    k_norm_w: &[f32],
    eps: f32,
) {
    let hid = heads * hd;
    let scale = 1.0 / (hd as f32).sqrt();
    let half = hd / 2;

    let head_rms = |v: &mut [f32], w: &[f32]| {
        let ms = v.iter().map(|x| x * x).sum::<f32>() / hd as f32;
        let inv = 1.0 / (ms + eps).sqrt();
        for (j, x) in v.iter_mut().enumerate() {
            *x = *x * inv * w[j];
        }
    };
    let rope = |vec: &mut [f32], tok: usize| {
        let (c, s) = (
            &cos[tok * hd..(tok + 1) * hd],
            &sin[tok * hd..(tok + 1) * hd],
        );
        for h in 0..heads {
            let v = &mut vec[h * hd..(h + 1) * hd];
            let orig: Vec<f32> = v.to_vec();
            for j in 0..hd {
                let rot = if j < half {
                    -orig[j + half]
                } else {
                    orig[j - half]
                };
                v[j] = orig[j] * c[j] + rot * s[j];
            }
        }
    };

    let mut q = vec![0f32; n * hid];
    let mut k = vec![0f32; n * hid];
    let mut v = vec![0f32; n * hid];
    q.par_chunks_mut(hid)
        .zip(k.par_chunks_mut(hid))
        .zip(v.par_chunks_mut(hid))
        .enumerate()
        .for_each(|(tok, ((qr, kr), vr))| {
            let row = &qkv[tok * 3 * hid..(tok + 1) * 3 * hid];
            qr.copy_from_slice(&row[..hid]);
            kr.copy_from_slice(&row[hid..2 * hid]);
            vr.copy_from_slice(&row[2 * hid..]);
            for h in 0..heads {
                head_rms(&mut qr[h * hd..(h + 1) * hd], q_norm_w);
                head_rms(&mut kr[h * hd..(h + 1) * hd], k_norm_w);
            }
            rope(qr, tok);
            rope(kr, tok);
        });

    let mut heads_out: Vec<Vec<f32>> = vec![Vec::new(); heads];
    heads_out.par_iter_mut().enumerate().for_each(|(h, slot)| {
        let mut scores = vec![0f32; n * n];
        // SAFETY: disjoint per-head output; q/k read-only. Strides view [n, hid] as this head's
        // [n, hd] block.
        unsafe {
            gemm::gemm(
                n,
                n,
                hd,
                scores.as_mut_ptr(),
                1,
                n as isize,
                false,
                q.as_ptr().add(h * hd),
                1,
                hid as isize,
                k.as_ptr().add(h * hd),
                hid as isize,
                1,
                0.0,
                scale,
                false,
                false,
                false,
                gemm::Parallelism::None,
            );
        }
        for row in scores.chunks_mut(n) {
            let max = row.iter().copied().fold(f32::NEG_INFINITY, f32::max);
            let mut sum = 0.0;
            for s in row.iter_mut() {
                *s = (*s - max).exp();
                sum += *s;
            }
            let inv = 1.0 / sum;
            for s in row.iter_mut() {
                *s *= inv;
            }
        }
        let mut ctx = vec![0f32; n * hd];
        // SAFETY: as above.
        unsafe {
            gemm::gemm(
                n,
                hd,
                n,
                ctx.as_mut_ptr(),
                1,
                hd as isize,
                false,
                scores.as_ptr(),
                1,
                n as isize,
                v.as_ptr().add(h * hd),
                1,
                hid as isize,
                0.0,
                1.0,
                false,
                false,
                false,
                gemm::Parallelism::None,
            );
        }
        *slot = ctx;
    });

    out.par_chunks_mut(hid).enumerate().for_each(|(tok, dst)| {
        for (h, ho) in heads_out.iter().enumerate() {
            dst[h * hd..(h + 1) * hd].copy_from_slice(&ho[tok * hd..(tok + 1) * hd]);
        }
    });
}

/// GLM twin of [`crate::vision::prepare_prompt`]: positions + tower forward + packaging for the
/// serving loop. The M-RoPE position RULE is identical to Qwen's (an image occupies `llm_h·llm_w`
/// slots and advances the counter by `max(llm_h, llm_w)`), so [`crate::vision::mrope_positions`] is
/// reused as-is; only the tower differs.
pub fn glm_prepare_prompt(
    tower: &GlmVisionTower,
    tokens: &[u32],
    image_token_id: u32,
    images: &[ImagePatches],
) -> Result<crate::server::VisionPrompt> {
    let embeds = tower.forward(images)?;
    glm_prompt_from_embeds(&tower.cfg, embeds, tokens, image_token_id, images)
}

/// [`glm_prepare_prompt`] with the tower on the GPU ([`crate::vision_glm_gpu::GlmVisionGpu`], iso to
/// the CPU tower) — the measured OCR bottleneck was the CPU tower's ViT.
pub fn glm_prepare_prompt_gpu(
    ctx: &crate::GpuCtx,
    tower: &crate::vision_glm_gpu::GlmVisionGpu,
    tokens: &[u32],
    image_token_id: u32,
    images: &[ImagePatches],
) -> Result<crate::server::VisionPrompt> {
    let embeds = tower.forward(ctx, images)?;
    glm_prompt_from_embeds(&tower.cfg, embeds, tokens, image_token_id, images)
}

/// [`glm_prepare_prompt_gpu`] with the CUDA tower arm — same embeds contract, no wgpu ctx.
#[cfg(feature = "cudarc")]
pub fn glm_prepare_prompt_cuda(
    tower: &crate::cuda_tower::CudaGlmTower,
    cfg: &VisionConfig,
    tokens: &[u32],
    image_token_id: u32,
    images: &[ImagePatches],
) -> Result<crate::server::VisionPrompt> {
    let embeds = tower.forward(images)?;
    glm_prompt_from_embeds(cfg, embeds, tokens, image_token_id, images)
}

fn glm_prompt_from_embeds(
    cfg: &VisionConfig,
    embeds: Vec<f32>,
    tokens: &[u32],
    image_token_id: u32,
    images: &[ImagePatches],
) -> Result<crate::server::VisionPrompt> {
    let grids: Vec<[u32; 3]> = images.iter().map(|i| i.grid).collect();
    let (mpos, next_pos) =
        crate::vision::mrope_positions(tokens, image_token_id, &grids, cfg.merge)?;
    let rows = embeds.len() / cfg.out_hidden;
    let want: usize = images.iter().map(|i| i.num_tokens(cfg)).sum();
    anyhow::ensure!(
        rows == want,
        "the tower produced {rows} tokens but the images need {want}"
    );

    let mut embed_index = vec![-1i32; tokens.len()];
    let mut next_row = 0i32;
    for (i, &t) in tokens.iter().enumerate() {
        if t == image_token_id {
            embed_index[i] = next_row;
            next_row += 1;
        }
    }
    anyhow::ensure!(
        next_row as usize == rows,
        "the prompt has {next_row} image placeholders but the tower produced {rows} rows"
    );

    Ok(crate::server::VisionPrompt {
        embeds,
        mpos,
        embed_index,
        next_pos,
    })
}