cortiq_engine/
mimo_ingress.rs

1//! Shared MiMo CLI/server ingress. Media stays in ordered content blocks until
2//! strict template rendering; only pad rows are replaced by tower embeddings.
3//! Text-only requests neither discover nor load multimodal companions.
4use crate::{
5    Pipeline, media,
6    mimo_audio::MimoAudio,
7    mimo_mm::{MimoMm, MimoTokenIds},
8    mimo_vision::{self, MimoProcessorConfig, MimoVit, VideoSource},
9    tokenizer::Tokenizer,
10};
11use cortiq_core::CmfModel;
12use serde_json::{Value, json};
13use std::{
14    path::{Path, PathBuf},
15    sync::Arc,
16};
17
18#[derive(Default)]
19pub struct MediaOptions {
20    pub companion: Option<PathBuf>,
21    pub image_max_pixels: Option<usize>,
22    pub video_fps: Option<f64>,
23}
24
25pub struct PreparedMimoPrompt {
26    pub token_ids: Vec<u32>,
27    /// None for pure text: preserve ordinary prefill and prefix reuse.
28    pub rows: Option<Vec<f32>>,
29}
30
31#[derive(Debug)]
32enum Item {
33    Image(Value),
34    Video(Value),
35    Audio(Value),
36}
37
38fn record(v: &Value) -> Value {
39    if let Some(url) = v.as_str() {
40        json!({"url":url})
41    } else {
42        v.clone()
43    }
44}
45
46/// Normalize aliases without flattening content or inserting newlines. Reject
47/// unknown blocks instead of silently rendering a prompt without its media.
48fn normalize(messages: &[Value]) -> Result<(Vec<Value>, Vec<Item>), String> {
49    let mut messages = messages.to_vec();
50    let mut items = Vec::new();
51    for message in &mut messages {
52        let Some(content) = message.get_mut("content") else {
53            continue;
54        };
55        if content.is_string() || content.is_null() {
56            continue;
57        }
58        let parts = content
59            .as_array_mut()
60            .ok_or("message content must be text or an array")?;
61        for part in parts {
62            if part.is_string() {
63                continue;
64            }
65            let typ = part.get("type").and_then(Value::as_str).unwrap_or("");
66            if typ == "image"
67                || typ == "image_url"
68                || part.get("image").is_some()
69                || part.get("image_url").is_some()
70            {
71                let source = part
72                    .get("image_url")
73                    .or_else(|| part.get("image"))
74                    .unwrap_or(part);
75                items.push(Item::Image(record(source)));
76                *part = json!({"type":"image"});
77            } else if matches!(typ, "audio" | "input_audio" | "audio_url")
78                || part.get("input_audio").is_some()
79                || part.get("audio_url").is_some()
80                || part.get("audio").is_some()
81            {
82                let source = part
83                    .get("input_audio")
84                    .or_else(|| part.get("audio_url"))
85                    .or_else(|| part.get("audio"))
86                    .unwrap_or(part);
87                if source
88                    .get("format")
89                    .and_then(Value::as_str)
90                    .is_some_and(|s| s != "wav")
91                {
92                    return Err("MiMo input_audio currently supports WAV only".into());
93                }
94                items.push(Item::Audio(record(source)));
95                *part = json!({"type":"audio"});
96            } else if matches!(typ, "video" | "video_url")
97                || part.get("video").is_some()
98                || part.get("video_url").is_some()
99            {
100                let source = part
101                    .get("video")
102                    .or_else(|| part.get("video_url"))
103                    .unwrap_or(part);
104                items.push(Item::Video(record(source)));
105                *part = json!({"type":"video"});
106            } else if part.get("text").and_then(Value::as_str).is_none() {
107                return Err(format!("unsupported MiMo content block type '{typ}'"));
108            }
109        }
110    }
111    Ok((messages, items))
112}
113
114pub fn has_media(messages: &[Value]) -> Result<bool, String> {
115    Ok(!normalize(messages)?.1.is_empty())
116}
117
118/// Render text-only multipart requests exactly as MiMo's template specifies.
119pub fn text_ids(
120    tok: &Tokenizer,
121    messages: &[Value],
122    tools: Option<&[Value]>,
123    thinking: Option<bool>,
124) -> Result<Vec<u32>, String> {
125    let (messages, items) = normalize(messages)?;
126    if !items.is_empty() {
127        return Err("media requires a MiMo multimodal companion".into());
128    }
129    if tok.chat_template.is_none() {
130        return Err("MiMo requires its embedded chat template".into());
131    }
132    tok.try_apply_chat_template_json(&messages, tools, thinking)
133}
134
135/// Server entry point: use the exact model owned by the checked-out pipeline.
136pub fn prepare_for_pipeline(
137    pipeline: &Pipeline,
138    messages: &[Value],
139    tools: Option<&[Value]>,
140    thinking: Option<bool>,
141    options: &MediaOptions,
142) -> Result<PreparedMimoPrompt, String> {
143    let model = pipeline
144        .model
145        .as_ref()
146        .ok_or("MiMo media requires a file-backed pipeline")?;
147    prepare(model, pipeline, messages, tools, thinking, options)
148}
149
150pub fn prepare(
151    model: &Arc<CmfModel>,
152    pipeline: &Pipeline,
153    messages: &[Value],
154    tools: Option<&[Value]>,
155    thinking: Option<bool>,
156    options: &MediaOptions,
157) -> Result<PreparedMimoPrompt, String> {
158    let (messages, items) = normalize(messages)?;
159    let tok = &pipeline.tokenizer;
160    if tok.chat_template.is_none() {
161        return Err("MiMo requires its embedded chat template".into());
162    }
163    let ids = tok.try_apply_chat_template_json(&messages, tools, thinking)?;
164    if items.is_empty() {
165        return Ok(PreparedMimoPrompt {
166            token_ids: ids,
167            rows: None,
168        });
169    }
170    let mm = MimoMm::attach(model, options.companion.as_deref())?.ok_or(
171        "this is a text-only MiMo package; install its .mm.cmf companion or pass --mm PATH",
172    )?;
173    mm.validate_text(model, tok)?;
174    let cfg = mimo_vision::read_mm_config(mm.model())?;
175    let processor = MimoProcessorConfig::from_config(&cfg)?;
176    let (mut vision, mut audio) = (None, None);
177    let (mut images, mut videos, mut audio_tokens) = (Vec::new(), Vec::new(), Vec::new());
178    let (mut image_rows, mut video_rows, mut audio_rows) = (Vec::new(), Vec::new(), Vec::new());
179    for item in items {
180        match item {
181            Item::Image(source) => {
182                let frame = media::decode_rgb(&media::load_image_bytes(&source)?)?;
183                let input =
184                    mimo_vision::prepare_image(&frame, &processor, options.image_max_pixels)?;
185                if vision.is_none() {
186                    vision = Some(MimoVit::from_model(mm.model())?);
187                }
188                let rows = vision.as_ref().unwrap().forward(&input)?;
189                check_rows(&rows, input.tokens(), pipeline.hidden_size)?;
190                image_rows.extend(rows);
191                images.push(input);
192            }
193            Item::Video(source) => {
194                let path = source
195                    .get("url")
196                    .or_else(|| source.get("path"))
197                    .and_then(Value::as_str)
198                    .ok_or("video requires a local Y4M file or frame directory path")?;
199                if path.starts_with("http:")
200                    || path.starts_with("https:")
201                    || path.starts_with("data:")
202                {
203                    return Err(
204                        "MiMo video supports local Y4M files or frame directories only".into(),
205                    );
206                }
207                let fps = source
208                    .get("fps")
209                    .and_then(Value::as_f64)
210                    .or(options.video_fps);
211                let src = VideoSource::open(
212                    Path::new(path.strip_prefix("file://").unwrap_or(path)),
213                    fps,
214                )?;
215                let input = mimo_vision::prepare_video(&src, &processor, options.image_max_pixels)?;
216                if vision.is_none() {
217                    vision = Some(MimoVit::from_model(mm.model())?);
218                }
219                let rows = vision.as_ref().unwrap().forward(&input)?;
220                check_rows(&rows, input.tokens(), pipeline.hidden_size)?;
221                video_rows.extend(rows);
222                videos.push(input);
223            }
224            Item::Audio(source) => {
225                if audio.is_none() {
226                    audio = Some(MimoAudio::from_model(mm.model())?);
227                }
228                let embedded = audio
229                    .as_ref()
230                    .unwrap()
231                    .embed_wav(&media::load_image_bytes(&source)?)?;
232                if embedded.dim != pipeline.hidden_size {
233                    return Err("audio/text hidden dimension mismatch".into());
234                }
235                check_rows(&embedded.rows, embedded.n_tokens, pipeline.hidden_size)?;
236                audio_tokens.push(embedded.n_tokens);
237                audio_rows.extend(embedded.rows);
238            }
239        }
240    }
241    let ids = mimo_vision::expand_prompt_ids(
242        &ids,
243        &images.iter().collect::<Vec<_>>(),
244        &videos.iter().collect::<Vec<_>>(),
245        &audio_tokens,
246        tok,
247    )?;
248    let rows = inject_rows(
249        &ids,
250        pipeline.hidden_size,
251        [&image_rows, &video_rows, &audio_rows],
252        |id| pipeline.embed_id(id),
253    )?;
254    Ok(PreparedMimoPrompt {
255        token_ids: ids,
256        rows: Some(rows),
257    })
258}
259
260fn check_rows(rows: &[f32], n: usize, h: usize) -> Result<(), String> {
261    if n == 0 || n.checked_mul(h) != Some(rows.len()) || !rows.iter().all(|v| v.is_finite()) {
262        return Err("invalid or non-finite MiMo media embeddings".into());
263    }
264    Ok(())
265}
266
267fn inject_rows(
268    ids: &[u32],
269    h: usize,
270    media: [&[f32]; 3],
271    mut embed: impl FnMut(u32) -> Vec<f32>,
272) -> Result<Vec<f32>, String> {
273    let cap = ids
274        .len()
275        .checked_mul(h)
276        .ok_or("prompt embedding size overflow")?;
277    let mut rows = Vec::with_capacity(cap);
278    let mut offsets = [0usize; 3];
279    let t = MimoTokenIds::PINNED;
280    for &id in ids {
281        let which = [t.image_pad, t.video_pad, t.audio_pad]
282            .iter()
283            .position(|&pad| pad == id);
284        if let Some(i) = which {
285            let slice = media[i]
286                .get(offsets[i]..offsets[i] + h)
287                .ok_or("media pad/embedding count mismatch")?;
288            rows.extend_from_slice(slice);
289            offsets[i] += h;
290        } else {
291            let row = embed(id);
292            if row.len() != h {
293                return Err("invalid text embedding width".into());
294            }
295            rows.extend(row);
296        }
297    }
298    if offsets.iter().zip(media).any(|(&n, m)| n != m.len()) {
299        return Err("unused MiMo media embedding rows".into());
300    }
301    Ok(rows)
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307    #[test]
308    fn multipart_text_does_not_gain_newlines_or_fall_back() {
309        let mut tok = Tokenizer::byte_level();
310        tok.chat_template = Some("{% for m in messages %}{% if m.content is string %}{{ m.content }}{% else %}{% for p in m.content %}{{ p.text }}{% endfor %}{% endif %}{% endfor %}".into());
311        let messages = [json!({"role":"user", "content":[
312            {"type":"text", "text":"before"}, {"type":"text", "text":"after"}]})];
313        assert_eq!(
314            text_ids(&tok, &messages, None, None).unwrap(),
315            tok.encode("beforeafter")
316        );
317        tok.chat_template = Some("{{ raise_exception('bad template') }}".into());
318        assert!(text_ids(&tok, &messages, None, None).is_err());
319        tok.chat_template = None;
320        assert!(text_ids(&tok, &messages, None, None).is_err());
321    }
322
323    #[test]
324    fn multipart_order_and_aliases_are_preserved() {
325        let msgs = vec![
326            json!({"role":"user","content":[{"type":"text","text":"before"},
327            {"type":"image_url","image_url":{"url":"x.png"}}, {"type":"text","text":"after"},
328            {"type":"input_audio","input_audio":{"data":"AA==","format":"wav"}},
329            {"type":"video","video":{"path":"frames","fps":1.0}}]}),
330        ];
331        let (m, items) = normalize(&msgs).unwrap();
332        assert_eq!(m[0]["content"][0]["text"], "before");
333        assert_eq!(m[0]["content"][2]["text"], "after");
334        assert!(matches!(
335            &items[..],
336            [Item::Image(_), Item::Audio(_), Item::Video(_)]
337        ));
338        assert!(normalize(&[json!({"content":[{"type":"unknown"}]})]).is_err());
339        assert!(
340            normalize(&[
341                json!({"content":[{"type":"input_audio","input_audio":{"format":"mp3"}}]})
342            ])
343            .is_err()
344        );
345    }
346    #[test]
347    fn injection_consumes_each_modality_in_prompt_order() {
348        let t = MimoTokenIds::PINNED;
349        let ids = [1, t.audio_pad, t.image_pad, 2, t.video_pad, t.image_pad];
350        let media: [&[f32]; 3] = [&[10., 11., 12., 13.], &[20., 21.], &[30., 31.]];
351        assert_eq!(
352            inject_rows(&ids, 2, media, |id| vec![id as f32; 2]).unwrap(),
353            vec![1., 1., 30., 31., 10., 11., 2., 2., 20., 21., 12., 13.]
354        );
355        assert!(inject_rows(&ids[..5], 2, media, |_| vec![0.; 2]).is_err());
356        assert!(inject_rows(&ids, 3, media, |_| vec![0.; 3]).is_err());
357    }
358}
cortiq_engine/mimo_ingress.rs

cortiq_engine/
mimo_ingress.rs
