mecha10-models 4.2.2

//! Model artifact download (the `.onnx` weights themselves)

use super::ModelService;
use anyhow::{Context, Result};
use indicatif::{ProgressBar, ProgressStyle};
use std::path::PathBuf;
use tokio::fs;

impl ModelService {
    /// Download a model from the catalog
    #[cfg(feature = "model-download")]
    pub async fn pull(&self, name: &str, progress: Option<&ProgressBar>) -> Result<PathBuf> {
        let entry = self
            .get_catalog_entry(name)
            .await?
            .context(format!("Model '{}' not found in catalog", name))?;

        let model_dir = self.catalog_model_dir(name);
        fs::create_dir_all(&model_dir).await?;

        let model_path = self
            .pull_from_repo(&entry.repo, &entry.filename, name, progress)
            .await?;

        if !entry.classes.is_empty() {
            self.write_inline_labels(name, &entry.classes).await?;
        } else if entry.task == "object-detection" {
            self.pull_labels_from_repo(&entry, name, progress).await?;
        } else if entry.task == "image-classification" {
            self.pull_labels_file(name, "imagenet-labels.txt", progress).await?;
        }

        if entry.task == "object-detection" {
            if let Some(class_names) = self.read_resolved_labels(name).await? {
                self.embed_class_names_metadata(&model_path, &class_names, progress)
                    .await?;
            }
        }

        self.generate_model_config(&entry, &model_path, progress).await?;

        if let Some(quantize_config) = &entry.quantize {
            if quantize_config.enabled {
                self.quantize_model(&model_path, quantize_config, progress).await?;
            }
        }

        if let Some(pb) = progress {
            pb.set_message(format!("✅ Model '{}' ready at {}", name, model_dir.display()));
        }

        Ok(model_path)
    }

    /// Download a model directly from HuggingFace Hub (mirrors `labels.rs::pull_labels_from_repo`'s
    /// direct-to-HF approach for the model's label file).
    ///
    /// Hits HF's `resolve/main` endpoint, which streams the raw file bytes from the given repo -
    /// no intermediary service is contacted.
    #[cfg(feature = "model-download")]
    pub async fn pull_from_repo(
        &self,
        repo: &str,
        filename: &str,
        name: &str,
        progress: Option<&ProgressBar>,
    ) -> Result<PathBuf> {
        let model_dir = self.catalog_model_dir(name);
        fs::create_dir_all(&model_dir).await?;

        let output_path = model_dir.join("model.onnx");

        if output_path.exists() {
            if let Some(pb) = progress {
                pb.set_message(format!("Model '{}' already cached", name));
            }
            return Ok(output_path);
        }

        let url = format!("https://huggingface.co/{repo}/resolve/main/{filename}");

        if let Some(pb) = progress {
            pb.set_style(
                ProgressStyle::default_spinner()
                    .template("{spinner:.green} {msg}")
                    .unwrap(),
            );
            pb.set_message(format!("Downloading {}", name));
        }

        tracing::info!("Downloading model from HuggingFace: {}", url);

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

        let response = client
            .get(&url)
            .send()
            .await
            .context(format!("Failed to connect to HuggingFace for model '{}'", name))?;

        if !response.status().is_success() {
            return Err(anyhow::anyhow!(
                "Model not found: HTTP {}\nURL: {}",
                response.status(),
                url
            ));
        }

        let mut file = tokio::fs::File::create(&output_path).await?;
        let mut stream = response.bytes_stream();

        use futures_util::StreamExt;
        use tokio::io::AsyncWriteExt;

        while let Some(chunk) = stream.next().await {
            let chunk = chunk.context("Error downloading model chunk")?;
            file.write_all(&chunk).await?;
        }
        file.flush().await?;

        if let Some(pb) = progress {
            pb.set_message(format!("Downloaded {} successfully", name));
        }

        Ok(output_path)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn service(models_dir: PathBuf) -> ModelService {
        ModelService::with_models_dir(models_dir).unwrap()
    }

    #[tokio::test]
    async fn pull_from_repo_returns_the_already_cached_file_without_a_network_call() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        let service = service(models_dir.clone());

        let catalog_dir = models_dir.join("catalog").join("yolov8n");
        tokio::fs::create_dir_all(&catalog_dir).await.unwrap();
        let model_path = catalog_dir.join("model.onnx");
        tokio::fs::write(&model_path, b"existing-bytes").await.unwrap();

        let result_path = service
            .pull_from_repo("deepghs/yolos", "yolov8n/model.onnx", "yolov8n", None)
            .await
            .unwrap();

        assert_eq!(result_path, model_path);
        let content = tokio::fs::read(&result_path).await.unwrap();
        assert_eq!(content, b"existing-bytes");
    }

    #[test]
    fn pull_from_repo_url_points_directly_at_huggingfaces_resolve_endpoint() {
        let repo = "deepghs/yolos";
        let filename = "yolov8n/model.onnx";
        let url = format!("https://huggingface.co/{repo}/resolve/main/{filename}");

        assert_eq!(
            url,
            "https://huggingface.co/deepghs/yolos/resolve/main/yolov8n/model.onnx"
        );
    }
}