studio-worker 0.4.9

Pull-based image-generation worker for the minis.gg studio.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

//! The worker's in-process model loaders, one per engine, behind the
//! model host's `ModelRuntime` (see `docs/runtime/model-lifecycle.md`).

use crate::catalog::CatalogModel;
use crate::host::{LoadedModel, ModelRuntime};
use std::path::PathBuf;
use std::sync::Arc;

/// Dispatches a load to the loader for the model's engine.  An engine
/// with no in-process loader is refused by name, never faked.
pub struct Loaders {
    models_root: PathBuf,
}

impl Loaders {
    /// `models_root`: where model files are downloaded to.
    pub fn new(models_root: PathBuf) -> Self {
        Self { models_root }
    }
}

impl ModelRuntime for Loaders {
    fn can_load(&self, model: &CatalogModel) -> bool {
        match model.source.engine {
            #[cfg(all(feature = "llama", not(target_os = "windows")))]
            crate::types::ModelEngine::LlamaCpp => true,
            #[cfg(feature = "stt-stream")]
            crate::types::ModelEngine::Parakeet => true,
            _ => false,
        }
    }

    fn load(&self, model: &CatalogModel) -> anyhow::Result<Arc<dyn LoadedModel>> {
        match &model.source.engine {
            #[cfg(all(feature = "llama", not(target_os = "windows")))]
            crate::types::ModelEngine::LlamaCpp => Ok(Arc::new(
                crate::engine::llama::load_resident(&self.models_root, model)?,
            )),
            #[cfg(feature = "stt-stream")]
            crate::types::ModelEngine::Parakeet => Ok(Arc::new(
                crate::engine::parakeet::load_resident(&self.models_root, model)?,
            )),
            engine => {
                let _ = &self.models_root;
                anyhow::bail!(
                    "no in-process loader for engine {engine:?} (model {})",
                    model.id
                )
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{ModelEngine, ModelSource, TaskKind};

    #[test]
    fn an_engine_without_a_loader_is_refused_by_name() {
        let model = CatalogModel {
            id: "m".into(),
            display_name: "m".into(),
            kind: TaskKind::Image,
            vram_gb_estimate: 1.0,
            description: None,
            source: ModelSource {
                engine: ModelEngine::SdCpp,
                files: vec![],
                cli_defaults: Default::default(),
            },
            enabled: true,
            origin: "local".into(),
            exclusive_group: None,
        };
        let loaders = Loaders::new(PathBuf::from("/nonexistent"));
        assert!(!loaders.can_load(&model));
        let err = loaders.load(&model).err().expect("refused");
        assert!(
            err.to_string()
                .contains("no in-process loader for engine SdCpp"),
            "{err}"
        );
    }

    #[cfg(all(feature = "llama", not(target_os = "windows")))]
    #[test]
    fn an_llm_can_be_loaded() {
        let model = CatalogModel {
            id: "m".into(),
            display_name: "m".into(),
            kind: TaskKind::Llm,
            vram_gb_estimate: 1.0,
            description: None,
            source: ModelSource {
                engine: ModelEngine::LlamaCpp,
                files: vec![],
                cli_defaults: Default::default(),
            },
            enabled: true,
            origin: "local".into(),
            exclusive_group: None,
        };
        assert!(Loaders::new(PathBuf::from("/nonexistent")).can_load(&model));
    }
}