feat(playground): image & video generation tabs + real chat metrics
The infosave collection now ships media models (Lumina-Image-2.0, MiniMax-H3 Turbo) that cortiq runs via `imagine`/`animate` CLI, not `cortiq serve`. The gateway grows a media API: local .cmf files are classified by `cortiq info` (chat/image/video, cached), generations run as tracked jobs with step progress parsed from stdout, one at a time, cancellable, outputs under <data>/media. The playground becomes three tabs. Image: width/height/steps/CFG/seed, live progress, result rendered from P6 PPM to a canvas in the browser (PNG export via toBlob). Video: frames/steps/quality/seed plus optional first/last keyframes (rasterized to PPM client-side); the MJPEG+PCM AVI is decoded in JS and played on a canvas synced to the wav — no ffmpeg on either side (mp4 remux happens opportunistically when ffmpeg exists). A metrics card shows resolution, steps, seed, total time, time per step/frame and file size. Chat metrics: the local cortiq server sends no usage chunk in streams, so token counts showed 0/0. tap_stream now injects one OpenAI-shaped usage event ahead of [DONE] with the gateway's ~4-chars/token estimate, flagged `cortiq_estimated`; the UI renders it as "≈ N / M" and adds tok/s and time-to-first-token. Registration skips media packs (a `cortiq serve` entry for a diffusion model would only fail) in both Register and the wizard's Finish. 36 new UI strings translated across all 7 languages.
This commit is contained in:
infosave2007
+771
@@ -0,0 +1,771 @@
|
||||
//! 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 `<data>/media/<job>/`.
|
||||
|
||||
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<HashMap<String, (u64, u64, Modality)>>); // 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<InfoCache> = 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<MediaModel> {
|
||||
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<u32>,
|
||||
#[serde(default)]
|
||||
pub height: Option<u32>,
|
||||
#[serde(default)]
|
||||
pub steps: Option<u32>,
|
||||
#[serde(default)]
|
||||
pub cfg: Option<f32>, // image only
|
||||
#[serde(default)]
|
||||
pub seed: Option<u64>,
|
||||
#[serde(default)]
|
||||
pub frames: Option<u32>, // video only
|
||||
#[serde(default)]
|
||||
pub quality: Option<u32>, // video only: AVI JPEG quality
|
||||
/// Optional first/last keyframe for video, base64-encoded binary P6 PPM.
|
||||
#[serde(default)]
|
||||
pub first_frame_b64: Option<String>,
|
||||
#[serde(default)]
|
||||
pub last_frame_b64: Option<String>,
|
||||
}
|
||||
|
||||
/// Live state of one generation.
|
||||
#[derive(Clone, Serialize)]
|
||||
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<u32>,
|
||||
pub state: String, // running | done | error | cancelled
|
||||
pub progress: Option<f32>,
|
||||
pub phase: Option<String>,
|
||||
pub log: Vec<String>,
|
||||
pub started: u64,
|
||||
pub finished: Option<u64>,
|
||||
/// Files this job produced: name → relative kind ("image", "video", "audio", "mp4").
|
||||
pub outputs: Vec<String>,
|
||||
pub out_bytes: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct MediaStore {
|
||||
jobs: Mutex<HashMap<String, MediaJob>>,
|
||||
cancels: Mutex<HashMap<String, Arc<Notify>>>,
|
||||
/// One generation at a time — these models take the whole GPU/RAM budget.
|
||||
busy: Mutex<Option<String>>,
|
||||
}
|
||||
|
||||
impl MediaStore {
|
||||
pub fn new() -> Arc<Self> {
|
||||
Arc::new(Self::default())
|
||||
}
|
||||
pub fn list(&self) -> Vec<MediaJob> {
|
||||
let mut v: Vec<MediaJob> = 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<MediaJob> {
|
||||
self.jobs.lock().unwrap().get(id).cloned()
|
||||
}
|
||||
pub fn cancel(&self, id: &str) -> bool {
|
||||
let running = matches!(
|
||||
self.jobs.lock().unwrap().get(id).map(|j| j.state.as_str()),
|
||||
Some("running")
|
||||
);
|
||||
if !running {
|
||||
return false;
|
||||
}
|
||||
if let Some(n) = self.cancels.lock().unwrap().get(id) {
|
||||
n.notify_one();
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
/// Remove a settled job and its output directory. Err = still running.
|
||||
pub fn delete(&self, id: &str, media_dir: &Path) -> Result<bool, String> {
|
||||
let job = self.jobs.lock().unwrap().get(id).cloned();
|
||||
let Some(job) = job else { return Ok(false) };
|
||||
if job.state == "running" {
|
||||
return Err("cancel the running generation first".into());
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(media_dir.join(id));
|
||||
self.jobs.lock().unwrap().remove(id);
|
||||
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<String>) {
|
||||
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, id: &str, ok: bool, outputs: Vec<String>, out_bytes: u64) {
|
||||
let mut g = self.jobs.lock().unwrap();
|
||||
if let Some(j) = g.get_mut(id) {
|
||||
if j.state == "cancelled" {
|
||||
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.clear_busy(id);
|
||||
self.cancels.lock().unwrap().remove(id);
|
||||
}
|
||||
fn set_cancelled(&self, id: &str) {
|
||||
if let Some(j) = self.jobs.lock().unwrap().get_mut(id) {
|
||||
j.state = "cancelled".into();
|
||||
j.finished = Some(now());
|
||||
}
|
||||
self.clear_busy(id);
|
||||
self.cancels.lock().unwrap().remove(id);
|
||||
}
|
||||
fn clear_busy(&self, id: &str) {
|
||||
let mut b = self.busy.lock().unwrap();
|
||||
if b.as_deref() == Some(id) {
|
||||
*b = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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<String>)> {
|
||||
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::<f32>() {
|
||||
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::<f32>() {
|
||||
if y > 0.0 && x <= y {
|
||||
return Some((x, y));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Root for generated files: `<data>/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<Vec<u8>> {
|
||||
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<MediaStore>,
|
||||
cmf: &crate::config::CmfCfg,
|
||||
p: MediaParams,
|
||||
) -> Result<String, String> {
|
||||
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 mut busy = store.busy.lock().unwrap();
|
||||
if let Some(running) = busy.as_ref() {
|
||||
return Err(format!(
|
||||
"a generation is already running (job {running}) — these models need the whole machine; wait or cancel it"
|
||||
));
|
||||
}
|
||||
*busy = Some(String::new()); // reserved; real id set below
|
||||
}
|
||||
|
||||
let id = gen_media_id();
|
||||
*store.busy.lock().unwrap() = Some(id.clone());
|
||||
|
||||
let dir = media_dir().join(&id);
|
||||
if let Err(e) = std::fs::create_dir_all(&dir) {
|
||||
store.clear_busy(&id);
|
||||
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<PathBuf> = None;
|
||||
let mut last_frame: Option<PathBuf> = 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<String> = 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: "running".into(),
|
||||
progress: None,
|
||||
phase: Some("starting".into()),
|
||||
log: vec![format!("→ cortiq {}", args.join(" "))],
|
||||
started: now(),
|
||||
finished: None,
|
||||
outputs: Vec::new(),
|
||||
out_bytes: None,
|
||||
},
|
||||
);
|
||||
let cancel = Arc::new(Notify::new());
|
||||
store
|
||||
.cancels
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(id.clone(), cancel.clone());
|
||||
|
||||
let bin = cmf.cortiq_bin.clone();
|
||||
let gpu = cmf.gpu;
|
||||
let ret = id.clone();
|
||||
let dir_clone = dir.clone();
|
||||
let kind_owned = kind.to_string();
|
||||
tokio::spawn(async move {
|
||||
run_generation(
|
||||
store, id, bin, args, gpu, dir_clone, out_path, kind_owned, cancel,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
Ok(ret)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn run_generation(
|
||||
store: Arc<MediaStore>,
|
||||
id: String,
|
||||
bin: String,
|
||||
args: Vec<String>,
|
||||
gpu: bool,
|
||||
dir: PathBuf,
|
||||
out_path: PathBuf,
|
||||
kind: String,
|
||||
cancel: Arc<Notify>,
|
||||
) {
|
||||
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<MediaStore>, 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<String, ()> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user