DarkRoom/core/dr-inference-engine/examples/ladder.rs

//! Walk the ladder as the app does — probe, engines, then a session — and
//! say what each step chose. The M5 check of docs/inference.md §6 without
//! the app around it.
//!
//!   DARKROOM_ORT_DIR=/usr/lib \
//!     cargo run --release -p dr-inference-engine --features native,tract \
//!     --example ladder -- CACHE_DIR models/face/scrfd_500m_640.onnx [ROLE=MODEL.onnx ...]
//!
//! `DARKROOM_ORT_DIRS=a:b:c` offers several runtimes, as the app's search
//! list does, and shows which the engine chose for this device's GPU.
//!
//! A bare path is a `Detector`; `denoiser=…`, `scene=…`, `inpainter=…`,
//! `landmarks=…` (any `Role`, lower case) says otherwise, so a device can
//! show each role taking its own form (inference.md §1.5). Each is opened
//! through `resolve_model`, as the app opens it, and the line says which
//! form and which rung it landed on. Delete `CACHE_DIR` to see the first
//! run again; keep it to see the second.

use std::path::PathBuf;
use std::time::{Duration, Instant};

fn main() {
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
    let mut args = std::env::args_os().skip(1).map(PathBuf::from);
    let (Some(cache_dir), models) = (
        args.next(),
        args.map(|a| role_and_path(&a)).collect::<Vec<_>>(),
    ) else {
        eprintln!("usage: ladder CACHE_DIR MODEL.onnx [MODEL.onnx ...]");
        std::process::exit(2);
    };
    if models.is_empty() {
        eprintln!("usage: ladder CACHE_DIR MODEL.onnx [MODEL.onnx ...]");
        std::process::exit(2);
    }

    // DARKROOM_ORT_DIRS lists several, colon-separated, as the app's search
    // does: the engine loads the one that fits the GPU (§3.2).
    let runtime_dirs: Vec<PathBuf> = std::env::var_os("DARKROOM_ORT_DIR")
        .map(PathBuf::from)
        .into_iter()
        .chain(
            std::env::var_os("DARKROOM_ORT_DIRS")
                .map(|v| std::env::split_paths(&v).collect::<Vec<_>>())
                .unwrap_or_default(),
        )
        .collect();
    let started = Instant::now();
    dr_inference_engine::init(dr_inference_engine::Config {
        runtime_dirs,
        cache_dir: cache_dir.clone(),
        models: models.clone(),
        embedded: Vec::new(),
        ceiling: None,
        threads: 0,
        decay: Duration::ZERO,
    });

    let mut last = String::new();
    loop {
        let s = dr_inference_engine::status();
        let line = format!(
            "{} · {} · engines {}/{}{}",
            s.line(),
            if s.probing {
                "probing"
            } else {
                s.reason.as_str()
            },
            s.engines.0,
            s.engines.1,
            if s.failed.is_empty() {
                String::new()
            } else {
                format!(
                    " · tried {}",
                    s.failed
                        .iter()
                        .map(|(r, why)| format!("{}: {why}", r.label()))
                        .collect::<Vec<_>>()
                        .join(" · ")
                )
            }
        );
        if line != last {
            println!("{:>6.1} s  {line}", started.elapsed().as_secs_f64());
            last = line;
        }
        if !s.probing && s.engines.0 >= s.engines.1 {
            break;
        }
        std::thread::sleep(Duration::from_millis(500));
    }

    for (role, path) in &models {
        let (path, form) = dr_inference_engine::resolve_model(*role, path);
        let bytes = std::fs::read(&path).expect("read model");
        let t = Instant::now();
        let model = dr_inference_engine::open(*role, form, &bytes).expect("open model");
        let acquired = model.acquire().expect("acquire session");
        println!(
            "{role:?}: {} ({form:?}) on {} in {:.2} s",
            path.file_name().unwrap().to_string_lossy(),
            acquired.rung().label(),
            t.elapsed().as_secs_f64()
        );
    }
}

/// `denoiser=path` → (Denoiser, path); a bare path is a detector.
fn role_and_path(arg: &std::path::Path) -> (dr_inference_engine::Role, PathBuf) {
    use dr_inference_engine::Role::*;
    let s = arg.to_string_lossy();
    let Some((name, path)) = s.split_once('=') else {
        return (Detector, arg.to_path_buf());
    };
    let role = match name {
        "detector" => Detector,
        "embedder" => Embedder,
        "segmenter" => Segmenter,
        "scene" => Scene,
        "landmarks" => Landmarks,
        "eyes" => EyeClassifier,
        "keypoints" => Keypoints,
        "inpainter" => Inpainter,
        "denoiser" => Denoiser,
        other => panic!("no role {other:?}"),
    };
    (role, PathBuf::from(path))
}