//! Media generation: run `cortiq imagine` (text→image, Lumina) and //! `cortiq animate` (text→video+audio, MiniMax-H3) over local `.cmf` files as //! tracked background jobs, the same shape as import jobs so the admin UI can //! poll progress and fetch the result. //! //! These models are NOT served by `cortiq serve` — each generation is one CLI //! run. The gateway owns the process, parses its progress, and keeps the output //! under `/media//`. use serde::Serialize; use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::process::Stdio; use std::sync::{Arc, Mutex}; use std::time::{SystemTime, UNIX_EPOCH}; use tokio::io::{AsyncBufReadExt, BufReader}; use tokio::sync::Notify; fn now() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_secs() } /// Random job id — media jobs are never deduplicated (every run is a new work). fn gen_media_id() -> String { let mut b = [0u8; 8]; let _ = getrandom::getrandom(&mut b); b.iter().map(|x| format!("{x:02x}")).collect() } /// What a `.cmf` under models_dir can do, decided by `cortiq info`. #[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize)] #[serde(rename_all = "lowercase")] pub enum Modality { Chat, Image, Video, } /// `cortiq info` output → modality. The arch line names the family; the media /// packs also carry their pipeline components (DiT/VAE/vocoder) in the text. fn classify_info(info: &str) -> Modality { let lower = info.to_ascii_lowercase(); let arch = lower .lines() .find_map(|l| l.trim().strip_prefix("arch:")) .unwrap_or("") .trim() .to_string(); // generic arch suffixes first (e.g. "lumina2-image", "mmh3-video"), then // known family names as a fallback for packs without the suffix if arch.contains("video") || arch.contains("mmh3") || arch.contains("minimax") { return Modality::Video; } if arch.contains("image") || arch.contains("lumina") || arch.contains("diffusion") { return Modality::Image; } if lower.contains("text-to-video") { return Modality::Video; } if lower.contains("text-to-image") { return Modality::Image; } Modality::Chat } /// One local `.cmf` with its detected modality. #[derive(Clone, Serialize)] pub struct MediaModel { pub name: String, // file stem pub path: String, pub size_bytes: u64, pub modality: Modality, } #[derive(Default)] struct InfoCache(Mutex>); // path → (mtime, size, modality) /// Modality of one `.cmf`, via `cortiq info` — cached by (path, mtime, size), /// so the process runs once per file version. pub async fn classify_file(path: &Path, cortiq_bin: &str) -> Modality { static CACHE: std::sync::OnceLock = std::sync::OnceLock::new(); let cache = CACHE.get_or_init(InfoCache::default); let size = crate::import::path_size(path); let mtime = std::fs::metadata(path) .ok() .and_then(|m| m.modified().ok()) .and_then(|t| t.duration_since(UNIX_EPOCH).ok()) .map(|d| d.as_secs()) .unwrap_or(0); let key = path.to_string_lossy().to_string(); if let Some(k) = cache .0 .lock() .unwrap() .get(&key) .filter(|(m, s, _)| *m == mtime && *s == size) .map(|(_, _, k)| *k) { return k; } let output = tokio::process::Command::new(cortiq_bin) .arg("info") .arg(path) .output() .await; let text = output .map(|o| { format!( "{}{}", String::from_utf8_lossy(&o.stdout), String::from_utf8_lossy(&o.stderr) ) }) .unwrap_or_default(); let k = classify_info(&text); cache.0.lock().unwrap().insert(key, (mtime, size, k)); k } /// Scan models_dir and classify every `.cmf` (file or sharded dir). pub async fn list_media_models(models_dir: &str, cortiq_bin: &str) -> Vec { let mut out = Vec::new(); let Ok(rd) = std::fs::read_dir(models_dir) else { return out; }; for entry in rd.filter_map(|e| e.ok()) { let path = entry.path(); let name = path .file_name() .unwrap_or_default() .to_string_lossy() .to_string(); if !name.to_ascii_lowercase().ends_with(".cmf") { continue; } let size = crate::import::path_size(&path); if size == 0 { continue; // empty or still being written } let modality = classify_file(&path, cortiq_bin).await; out.push(MediaModel { name: name .trim_end_matches(".cmf") .trim_end_matches(".CMF") .to_string(), path: path.to_string_lossy().to_string(), size_bytes: size, modality, }); } out.sort_by(|a, b| a.name.cmp(&b.name)); out } /// Parameters for one generation, straight from the playground form. #[derive(Clone, Debug, serde::Deserialize)] pub struct MediaParams { pub model: String, // file stem under models_dir pub kind: String, // image | video pub prompt: String, #[serde(default)] pub width: Option, #[serde(default)] pub height: Option, #[serde(default)] pub steps: Option, #[serde(default)] pub cfg: Option, // image only #[serde(default)] pub seed: Option, #[serde(default)] pub frames: Option, // video only #[serde(default)] pub quality: Option, // video only: AVI JPEG quality /// Optional first/last keyframe for video, base64-encoded binary P6 PPM. #[serde(default)] pub first_frame_b64: Option, #[serde(default)] pub last_frame_b64: Option, } /// Live state of one generation. #[derive(Clone, Serialize, serde::Deserialize)] pub struct MediaJob { pub id: String, pub kind: String, // image | video pub model: String, pub prompt: String, pub width: u32, pub height: u32, pub steps: u32, pub seed: u64, pub frames: Option, pub state: String, // queued | running | done | error | cancelled pub progress: Option, pub phase: Option, pub log: Vec, pub started: u64, pub finished: Option, /// Files this job produced: name → relative kind ("image", "video", "audio", "mp4"). pub outputs: Vec, pub out_bytes: Option, } /// Everything needed to actually launch a queued job when its turn comes. struct Prepared { bin: String, args: Vec, gpu: bool, dir: PathBuf, out_path: PathBuf, kind: String, } #[derive(Default)] pub struct MediaStore { jobs: Mutex>, cancels: Mutex>>, /// One generation runs at a time — these models take the whole GPU/RAM /// budget — but new requests queue up instead of being rejected. queue: Mutex>, current: Mutex>, prepared: Mutex>, persist: Mutex>, } impl MediaStore { pub fn new() -> Arc { Arc::new(Self::default()) } /// Load saved generations (their files are still under media_dir) and start /// persisting. Jobs that were queued/running when the gateway died settle /// as errors — the process is gone. pub fn attach_persistence(&self, path: PathBuf) { if let Ok(text) = std::fs::read_to_string(&path) { if let Ok(mut jobs) = serde_json::from_str::>(&text) { let mut g = self.jobs.lock().unwrap(); for mut j in jobs.drain(..) { if j.state == "running" || j.state == "queued" { j.state = "error".into(); j.log.push("✗ interrupted by a gateway restart".into()); j.finished = Some(now()); } g.insert(j.id.clone(), j); } } } *self.persist.lock().unwrap() = Some(path); } fn save(&self) { let Some(path) = self.persist.lock().unwrap().clone() else { return; }; let jobs: Vec = self.jobs.lock().unwrap().values().cloned().collect(); if let Ok(json) = serde_json::to_vec(&jobs) { let tmp = path.with_extension("json.tmp"); if std::fs::write(&tmp, json).is_ok() { let _ = std::fs::rename(&tmp, &path); } } } /// How many jobs sit in the queue ahead of `id` (0 = next up). pub fn queue_position(&self, id: &str) -> Option { self.queue.lock().unwrap().iter().position(|x| x == id) } pub fn list(&self) -> Vec { let mut v: Vec = self.jobs.lock().unwrap().values().cloned().collect(); v.sort_by_key(|j| std::cmp::Reverse(j.started)); v.truncate(20); v } pub fn get(&self, id: &str) -> Option { self.jobs.lock().unwrap().get(id).cloned() } pub fn cancel(self: &Arc, id: &str, media_dir: &Path) -> bool { let state = self .jobs .lock() .unwrap() .get(id) .map(|j| j.state.clone()) .unwrap_or_default(); match state.as_str() { "queued" => { self.queue.lock().unwrap().retain(|x| x != id); self.prepared.lock().unwrap().remove(id); let _ = std::fs::remove_dir_all(media_dir.join(id)); self.set_cancelled(id); true } "running" => { if let Some(n) = self.cancels.lock().unwrap().get(id) { n.notify_one(); true } else { false } } _ => false, } } /// Remove a settled job and its output directory. Err = still running. pub fn delete(&self, id: &str, media_dir: &Path) -> Result { let job = self.jobs.lock().unwrap().get(id).cloned(); let Some(job) = job else { return Ok(false) }; if job.state == "running" || job.state == "queued" { return Err("cancel the running generation first".into()); } let _ = std::fs::remove_dir_all(media_dir.join(id)); self.jobs.lock().unwrap().remove(id); self.save(); Ok(true) } fn push_line(&self, id: &str, line: String) { if let Some(j) = self.jobs.lock().unwrap().get_mut(id) { j.log.push(line); let n = j.log.len(); if n > 120 { j.log.drain(0..n - 120); } } } fn set_progress(&self, id: &str, frac: f32, phase: Option) { if let Some(j) = self.jobs.lock().unwrap().get_mut(id) { let f = frac.clamp(0.0, 1.0); if j.progress.is_none() || f > j.progress.unwrap() { j.progress = Some(f); } if let Some(p) = phase { j.phase = Some(p); } } } fn finish(self: &Arc, id: &str, ok: bool, outputs: Vec, out_bytes: u64) { let mut g = self.jobs.lock().unwrap(); if let Some(j) = g.get_mut(id) { if j.state == "cancelled" { drop(g); self.after_slot_freed(id); return; } j.finished = Some(now()); j.state = if ok { "done" } else { "error" }.into(); if ok { j.progress = Some(1.0); j.outputs = outputs; j.out_bytes = Some(out_bytes); } } drop(g); self.after_slot_freed(id); } fn set_cancelled(self: &Arc, id: &str) { if let Some(j) = self.jobs.lock().unwrap().get_mut(id) { j.state = "cancelled".into(); j.finished = Some(now()); } self.after_slot_freed(id); } fn after_slot_freed(self: &Arc, id: &str) { { let mut cur = self.current.lock().unwrap(); if cur.as_deref() == Some(id) { *cur = None; } } self.cancels.lock().unwrap().remove(id); self.save(); promote_next(self.clone()); } } /// Start the next queued generation, if the slot is free. fn promote_next(store: Arc) { let next = { let mut cur = store.current.lock().unwrap(); if cur.is_some() { return; } let Some(id) = store.queue.lock().unwrap().pop_front() else { return; }; *cur = Some(id.clone()); id }; let Some(p) = store.prepared.lock().unwrap().remove(&next) else { // job was cancelled while queued and slipped through — free the slot *store.current.lock().unwrap() = None; return; }; if let Some(j) = store.jobs.lock().unwrap().get_mut(&next) { j.state = "running".into(); j.phase = Some("starting".into()); j.started = now(); } store.save(); let cancel = Arc::new(Notify::new()); store .cancels .lock() .unwrap() .insert(next.clone(), cancel.clone()); let s = store.clone(); tokio::spawn(async move { run_generation( s, next, p.bin, p.args, p.gpu, p.dir, p.out_path, p.kind, cancel, ) .await; }); } /// Progress markers in `cortiq imagine`/`animate` stdout. The exact wording has /// varied between releases, so accept anything that carries an X/Y pair near a /// step/frame word, a bare percentage, or the converter-style `@PROGRESS`. fn parse_media_progress(line: &str) -> Option<(f32, Option)> { let l = line.trim(); if let Some(rest) = l.strip_prefix("@PROGRESS ") { let mut it = rest.splitn(2, ' '); let frac: f32 = it.next()?.trim().parse().ok()?; let phase = it .next() .map(|s| s.trim().to_string()) .filter(|s| !s.is_empty()); return Some((frac.clamp(0.0, 1.0), phase)); } let lower = l.to_ascii_lowercase(); // "step 12/30", "denoise 3/4", "frame 7/39" — any word followed by x/y for word in ["step", "frame", "sigma", "denois"] { if let Some(pos) = lower.find(word) { let tail = &lower[pos..]; if let Some((x, y)) = extract_ratio(tail) { return Some((x / y, Some(l.to_string()))); } } } // bare "42%" (avoid matching sizes like "42%faster") if let Some(idx) = lower.find('%') { let head: String = lower[..idx] .chars() .rev() .take_while(|c| c.is_ascii_digit() || *c == '.') .collect(); let num: String = head.chars().rev().collect(); if let Ok(p) = num.parse::() { if (0.0..=100.0).contains(&p) { return Some((p / 100.0, None)); } } } None } /// First "N/M" pair in the string with M > 0. fn extract_ratio(s: &str) -> Option<(f32, f32)> { let bytes = s.as_bytes(); let mut i = 0; while i < bytes.len() { if bytes[i].is_ascii_digit() { let start = i; while i < bytes.len() && bytes[i].is_ascii_digit() { i += 1; } if i < bytes.len() && bytes[i] == b'/' { let x: f32 = s[start..i].parse().ok()?; let ystart = i + 1; let mut j = ystart; while j < bytes.len() && bytes[j].is_ascii_digit() { j += 1; } if j > ystart { if let Ok(y) = s[ystart..j].parse::() { if y > 0.0 && x <= y { return Some((x, y)); } } } } } else { i += 1; } } None } /// Root for generated files: `/media`. pub fn media_dir() -> PathBuf { crate::config::data_dir().join("media") } const B64: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; /// Minimal base64 decode (standard alphabet, padding optional) — the keyframes /// are the only base64 input the gateway takes; not worth a dependency. fn b64_decode(s: &str) -> Option> { let mut lut = [255u8; 256]; for (i, c) in B64.iter().enumerate() { lut[*c as usize] = i as u8; } let mut out = Vec::with_capacity(s.len() * 3 / 4); let mut buf = 0u32; let mut bits = 0u32; for c in s.bytes() { if c == b'=' || c == b'\n' || c == b'\r' { continue; } let v = lut[c as usize]; if v == 255 { return None; } buf = (buf << 6) | v as u32; bits += 6; if bits >= 8 { bits -= 8; out.push((buf >> bits) as u8); } } Some(out) } /// Kick off a generation. Returns the job id, or an error when another one is /// already running / the model file is missing. pub fn start_generation( store: Arc, cmf: &crate::config::CmfCfg, p: MediaParams, ) -> Result { if p.prompt.trim().is_empty() { return Err("empty prompt".into()); } let kind = p.kind.as_str(); if kind != "image" && kind != "video" { return Err(format!("unknown kind '{kind}'")); } let model_path = Path::new(&cmf.models_dir).join(format!("{}.cmf", p.model)); if !model_path.exists() { return Err(format!("model file not found: {}", model_path.display())); } let id = gen_media_id(); let dir = media_dir().join(&id); if let Err(e) = std::fs::create_dir_all(&dir) { return Err(format!("cannot create media dir {}: {e}", dir.display())); } // Optional keyframes (video): decode to PPM files next to the output. let mut first_frame: Option = None; let mut last_frame: Option = None; for (b64, slot, name) in [ (&p.first_frame_b64, &mut first_frame, "first.ppm"), (&p.last_frame_b64, &mut last_frame, "last.ppm"), ] { if let Some(data) = b64.as_deref().filter(|s| !s.is_empty()) { let bytes = b64_decode(data).ok_or("invalid keyframe base64")?; if !bytes.starts_with(b"P6") { return Err("keyframe must be a binary P6 PPM".into()); } let path = dir.join(name); std::fs::write(&path, bytes).map_err(|e| format!("write keyframe: {e}"))?; *slot = Some(path); } } let width = p.width.unwrap_or(512).clamp(64, 2048); let height = p .height .unwrap_or(if kind == "video" { 288 } else { 512 }) .clamp(64, 2048); let steps = p .steps .unwrap_or(if kind == "video" { 4 } else { 30 }) .clamp(1, 200); let seed = p.seed.unwrap_or(42); let frames = p.frames.unwrap_or(39).clamp(1, 1000); let out_name = if kind == "image" { "out.ppm" } else { "out.avi" }; let out_path = dir.join(out_name); let mut args: Vec = if kind == "image" { vec![ "imagine".into(), model_path.to_string_lossy().to_string(), "--prompt".into(), p.prompt.clone(), "--width".into(), width.to_string(), "--height".into(), height.to_string(), "--steps".into(), steps.to_string(), "--cfg".into(), format!("{}", p.cfg.unwrap_or(4.0)), "--seed".into(), seed.to_string(), "--out".into(), out_path.to_string_lossy().to_string(), ] } else { let mut a = vec![ "animate".into(), model_path.to_string_lossy().to_string(), "--prompt".into(), p.prompt.clone(), "--width".into(), width.to_string(), "--height".into(), height.to_string(), "--frames".into(), frames.to_string(), "--steps".into(), steps.to_string(), "--seed".into(), seed.to_string(), "--quality".into(), p.quality.unwrap_or(92).clamp(10, 100).to_string(), "--out".into(), out_path.to_string_lossy().to_string(), ]; if let Some(f) = &first_frame { a.push("--first-frame".into()); a.push(f.to_string_lossy().to_string()); } if let Some(f) = &last_frame { a.push("--last-frame".into()); a.push(f.to_string_lossy().to_string()); } a }; // keep the model path first for readability in logs let _ = &mut args; store.jobs.lock().unwrap().insert( id.clone(), MediaJob { id: id.clone(), kind: kind.to_string(), model: p.model.clone(), prompt: p.prompt.clone(), width, height, steps, seed, frames: (kind == "video").then_some(frames), state: "queued".into(), progress: None, phase: None, log: vec![format!("→ cortiq {}", args.join(" "))], started: now(), finished: None, outputs: Vec::new(), out_bytes: None, }, ); store.prepared.lock().unwrap().insert( id.clone(), Prepared { bin: cmf.cortiq_bin.clone(), args, gpu: cmf.gpu, dir, out_path, kind: kind.to_string(), }, ); store.queue.lock().unwrap().push_back(id.clone()); store.save(); promote_next(store.clone()); Ok(id) } #[allow(clippy::too_many_arguments)] async fn run_generation( store: Arc, id: String, bin: String, args: Vec, gpu: bool, dir: PathBuf, out_path: PathBuf, kind: String, cancel: Arc, ) { let mut cmd = tokio::process::Command::new(&bin); cmd.args(&args) .env("CMF_GPU", if gpu { "1" } else { "0" }) .stdout(Stdio::piped()) .stderr(Stdio::piped()); let mut child = match cmd.spawn() { Ok(c) => c, Err(e) => { store.push_line(&id, format!("✗ failed to start cortiq: {e}")); store.finish(&id, false, Vec::new(), 0); return; } }; let mut readers = Vec::new(); if let Some(o) = child.stdout.take() { readers.push(tokio::spawn(read_lines(store.clone(), id.clone(), o))); } if let Some(e) = child.stderr.take() { readers.push(tokio::spawn(read_lines(store.clone(), id.clone(), e))); } let status = tokio::select! { s = child.wait() => s, _ = cancel.notified() => { let _ = child.start_kill(); let _ = child.wait().await; for r in readers { r.abort(); } let _ = std::fs::remove_dir_all(&dir); store.push_line(&id, "✗ cancelled".into()); store.set_cancelled(&id); return; } }; for r in readers { let _ = r.await; } let ok = status.map(|s| s.success()).unwrap_or(false); if !ok || !out_path.exists() { store.push_line(&id, "✗ generation failed — see the log above".into()); store.finish(&id, false, Vec::new(), 0); return; } let mut outputs = vec![if kind == "image" { "image" } else { "video" }.to_string()]; if kind == "video" { // `animate` writes a .wav next to the .avi if dir.join("out.wav").exists() { outputs.push("audio".into()); } // opportunistic browser-friendly remux when ffmpeg is present if let Ok(p) = which_ffmpeg() { store.push_line(&id, "→ remuxing to mp4 (ffmpeg found)".into()); let mp4 = dir.join("out.mp4"); let mut f = tokio::process::Command::new(p); f.arg("-y").arg("-i").arg(dir.join("out.avi")); if dir.join("out.wav").exists() { f.arg("-i") .arg(dir.join("out.wav")) .args(["-map", "0:v", "-map", "1:a", "-c:a", "aac"]); } f.args(["-c:v", "libx264", "-pix_fmt", "yuv420p"]).arg(&mp4); f.stdout(Stdio::null()).stderr(Stdio::null()); if matches!(f.status().await, Ok(s) if s.success()) && mp4.exists() { outputs.push("mp4".into()); } } } let bytes = crate::import::path_size(&out_path); store.push_line( &id, format!("✓ done — {}", crate::import::format_bytes(bytes)), ); store.finish(&id, true, outputs, bytes); } /// Strip ANSI color escapes that `tracing` writes even into a pipe. fn strip_ansi(s: &str) -> String { let mut out = String::with_capacity(s.len()); let mut chars = s.chars(); while let Some(c) = chars.next() { if c == '\u{1b}' { for d in chars.by_ref() { if d.is_ascii_alphabetic() { break; } } } else { out.push(c); } } out } async fn read_lines(store: Arc, id: String, stream: impl tokio::io::AsyncRead + Unpin) { let mut lines = BufReader::new(stream).lines(); while let Ok(Some(l)) = lines.next_line().await { let l = strip_ansi(l.trim()); if l.is_empty() { continue; } // the Metal backend logs one line per pipeline — dozens per run, all identical if l.contains("Metal GPU path") { continue; } if let Some((frac, phase)) = parse_media_progress(&l) { store.set_progress(&id, frac, phase); } else { store.push_line(&id, l); } } } fn which_ffmpeg() -> Result { for p in [ "ffmpeg", "/usr/local/bin/ffmpeg", "/opt/homebrew/bin/ffmpeg", ] { if std::process::Command::new(p) .arg("-version") .stdout(Stdio::null()) .stderr(Stdio::null()) .status() .map(|s| s.success()) .unwrap_or(false) { return Ok(p.to_string()); } } Err(()) } /// (relative file, content-type) for an output kind of a finished job. pub fn output_file(kind: &str) -> Option<(&'static str, &'static str)> { match kind { "image" => Some(("out.ppm", "image/x-portable-pixmap")), "video" => Some(("out.avi", "video/x-msvideo")), "audio" => Some(("out.wav", "audio/wav")), "mp4" => Some(("out.mp4", "video/mp4")), _ => None, } } #[cfg(test)] mod tests { use super::*; #[test] fn progress_parses_common_shapes() { assert_eq!(parse_media_progress("step 15/30").unwrap().0, 0.5); assert_eq!( parse_media_progress("Denoising step 3/4 ...").unwrap().0, 0.75 ); assert_eq!(parse_media_progress("frame 39/39 decoded").unwrap().0, 1.0); assert_eq!(parse_media_progress("@PROGRESS 0.25 vae").unwrap().0, 0.25); assert!(parse_media_progress("loaded 2361 tensors").is_none()); } #[test] fn classify_by_arch() { assert_eq!(classify_info(" Arch: lumina2\n"), Modality::Image); assert_eq!(classify_info(" Arch: mmh3\n"), Modality::Video); assert_eq!(classify_info(" Arch: qwen3\n"), Modality::Chat); } #[test] fn b64_roundtrip() { assert_eq!(b64_decode("UDYgd2g=").unwrap(), b"P6 wh"); assert!(b64_decode("!!").is_none()); } }