mecha10_models/services/model_service/

catalog.rs

//! Model catalog schema, live fetch, and HuggingFace preprocessor config parsing
//!
//! Types for the public model catalog's entries, the auto-quantization config attached to a
//! catalog entry, the built-in preprocessing presets, and the schema used to parse a
//! HuggingFace `preprocessor_config.json`.
//!
//! The catalog itself used to be a TOML file (`model_catalog.toml`) compiled directly into
//! this CLI binary via `include_str!`. `mecha10-api` now serves this same table live over
//! HTTP (`GET /api/models/catalog` server-side, reached via `{mecha10_api_base_url}/models/catalog`
//! client-side since that base URL already bakes in `/api` - public/unauthenticated,
//! `mecha10-industries/mecha10-api`'s `models::registry::catalog` module) - that's the single
//! source of truth now, so [`fetch_catalog`] always fetches it fresh. There is no bundled/
//! compiled-in fallback: a project with no network access to `mecha10-api` can't resolve
//! catalog models at all, same as it already couldn't resolve a `"registry"`-sourced model
//! without network access.

use anyhow::{Context, Result};
use mecha10_core::model::PreprocessingConfig;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// `GET {mecha10_api_base_url}/models/catalog` - the live public model catalog. Public,
/// unauthenticated, no caching here (every call is a fresh HTTP round-trip) and no fallback -
/// see the module doc comment. A bounded 10s timeout keeps this failing fast (rather than
/// hanging) when `mecha10-api` is unreachable - e.g. `ModelService::info`/`list_installed`
/// treat this call as best-effort and need it to fail quickly, not stall. Goes through
/// [`mecha10_cli_core::paths::urls::mecha10_api_path`] (same helper `download::pull_from_repo`
/// uses) rather than hand-joining `mecha10_api_base_url()` with the path itself - that used to
/// be two independent, inconsistent `format!`s and is exactly how this endpoint ended up
/// requested at `.../api/api/models/catalog` once already.
pub async fn fetch_catalog() -> Result<Vec<ModelCatalogEntry>> {
    let url = mecha10_cli_core::paths::urls::mecha10_api_path("/models/catalog");

    let client = reqwest::Client::builder()
        .user_agent("mecha10-cli")
        .timeout(Duration::from_secs(10))
        .build()
        .context("Failed to build HTTP client")?;

    let response = client
        .get(&url)
        .send()
        .await
        .context(format!("Failed to fetch model catalog from {url}"))?
        .error_for_status()
        .context(format!("Model catalog request to {url} failed"))?;

    response
        .json::<Vec<ModelCatalogEntry>>()
        .await
        .context(format!("Malformed model catalog response from {url}"))
}

/// Model catalog entry, as served by `GET /api/models/catalog`
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct ModelCatalogEntry {
    /// Model identifier (e.g., "yolov8n-face", "mobilenet-v2")
    #[schemars(example = &"yolov8n-face")]
    pub name: String,
    /// Human-readable description
    #[schemars(example = &"YOLOv8 Nano - Fast face detection")]
    pub description: String,
    /// Task type (e.g., "object-detection", "image-classification")
    #[schemars(example = &"object-detection")]
    pub task: String,
    /// HuggingFace repository (e.g., "deepghs/yolo-face")
    #[schemars(example = &"deepghs/yolo-face")]
    pub repo: String,
    /// Filename within the repo (e.g., "yolov8n-face/model.onnx")
    #[schemars(example = &"yolov8n-face/model.onnx")]
    pub filename: String,

    /// Optional preprocessing preset hint (e.g., "imagenet", "yolo", "coco")
    #[serde(default)]
    #[schemars(example = &"yolo")]
    pub preprocessing_preset: Option<String>,

    /// Optional small class lists (e.g., ["face"] for face detection)
    #[serde(default)]
    #[schemars(example = vec!["face".to_string()])]
    pub classes: Vec<String>,

    /// Optional automatic quantization configuration
    #[serde(default)]
    #[schemars(example = Some(QuantizeConfig { enabled: true, method: "dynamic_int8".to_string() }))]
    pub quantize: Option<QuantizeConfig>,

    /// Explicit `[width, height]` override for the model's input tensor size - takes
    /// precedence over auto-detection (HuggingFace `preprocessor_config.json`, then ONNX
    /// graph introspection) in `preprocessing::auto_detect_preprocessing`. Most entries leave
    /// this `None`; it exists for models where auto-detection is unreliable, e.g. a non-square
    /// re-id backbone that would otherwise silently fall back to the generic 224x224 default.
    #[serde(default)]
    #[schemars(example = Some([128_u32, 256_u32]))]
    pub input_size: Option<[u32; 2]>,
}

/// Quantization configuration for automatic INT8 conversion
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct QuantizeConfig {
    /// Enable automatic quantization after download
    #[schemars(example = true)]
    pub enabled: bool,
    /// Quantization method (currently only "dynamic_int8" supported)
    #[schemars(example = &"dynamic_int8")]
    pub method: String,
}

#[derive(Debug, Clone, Copy)]
pub enum PreprocessingPreset {
    ImageNet,
    Yolo,
    Coco,
    Zero255,
}

impl PreprocessingPreset {
    /// Parse preset from string name
    pub fn from_name(name: &str) -> Result<Self> {
        match name.to_lowercase().as_str() {
            "imagenet" => Ok(Self::ImageNet),
            "yolo" => Ok(Self::Yolo),
            "coco" => Ok(Self::Coco),
            "zero255" | "0-255" => Ok(Self::Zero255),
            _ => anyhow::bail!("Unknown preprocessing preset: {}", name),
        }
    }

    /// Get preprocessing configuration for this preset
    pub fn to_config(self) -> PreprocessingConfig {
        match self {
            Self::ImageNet | Self::Coco => PreprocessingConfig {
                mean: [0.485, 0.456, 0.406],
                std: [0.229, 0.224, 0.225],
                channel_order: "RGB".to_string(),
            },
            Self::Yolo => PreprocessingConfig {
                mean: [0.0, 0.0, 0.0],
                std: [255.0, 255.0, 255.0],
                channel_order: "RGB".to_string(),
            },
            Self::Zero255 => PreprocessingConfig {
                mean: [0.0, 0.0, 0.0],
                std: [1.0, 1.0, 1.0],
                channel_order: "RGB".to_string(),
            },
        }
    }
}

#[cfg(feature = "model-download")]
#[derive(Debug, Clone, Deserialize)]
pub(super) struct HFPreprocessorConfig {
    #[serde(default)]
    pub(super) image_mean: Option<Vec<f32>>,
    #[serde(default)]
    pub(super) image_std: Option<Vec<f32>>,
    #[serde(default)]
    pub(super) size: Option<HFSize>,
    #[serde(default)]
    pub(super) crop_size: Option<HFSize>,
}

#[cfg(feature = "model-download")]
#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub(super) enum HFSize {
    Dict { height: u32, width: u32 },
    ShortestEdge { shortest_edge: u32 },
    Single(u32),
}

#[cfg(feature = "model-download")]
impl HFPreprocessorConfig {
    pub(super) fn to_preprocessing(&self) -> PreprocessingConfig {
        PreprocessingConfig {
            mean: [
                self.image_mean.as_ref().and_then(|v| v.first()).copied().unwrap_or(0.0),
                self.image_mean.as_ref().and_then(|v| v.get(1)).copied().unwrap_or(0.0),
                self.image_mean.as_ref().and_then(|v| v.get(2)).copied().unwrap_or(0.0),
            ],
            std: [
                self.image_std.as_ref().and_then(|v| v.first()).copied().unwrap_or(1.0),
                self.image_std.as_ref().and_then(|v| v.get(1)).copied().unwrap_or(1.0),
                self.image_std.as_ref().and_then(|v| v.get(2)).copied().unwrap_or(1.0),
            ],
            channel_order: "RGB".to_string(),
        }
    }

    pub(super) fn input_size(&self) -> Option<[u32; 2]> {
        if let Some(crop_size) = &self.crop_size {
            return match crop_size {
                HFSize::Dict { height, width } => Some([*width, *height]),
                HFSize::ShortestEdge { shortest_edge } => Some([*shortest_edge, *shortest_edge]),
                HFSize::Single(s) => Some([*s, *s]),
            };
        }

        match &self.size {
            Some(HFSize::Dict { height, width }) => Some([*width, *height]),
            Some(HFSize::ShortestEdge { shortest_edge }) => Some([*shortest_edge, *shortest_edge]),
            Some(HFSize::Single(s)) => Some([*s, *s]),
            None => None,
        }
    }
}