Skip to main content

kernel/discovery/
gguf_models.rs

1//! Turning a flat list of GGUF files into [`DiscoveredModel`]s: loose files
2//! become one model each, shard sets become a single model keyed by their shared
3//! base (flagged as still-downloading when incomplete). Shared by the file-tree
4//! scanners (LM Studio, loose files).
5
6use std::collections::HashMap;
7use std::path::{Path, PathBuf};
8
9use crate::discovery::gguf_shards::group;
10use crate::discovery::modality_hints::gguf_hint;
11use crate::discovery::scanner::DiscoveredModel;
12use crate::records::{ModelSource, SourceKind};
13
14/// Whether a filename looks like a multimodal projector (never the primary
15/// weight). Case-insensitive `mmproj` substring.
16pub fn is_mmproj_name(name: &str) -> bool {
17    name.to_ascii_lowercase().contains("mmproj")
18}
19
20/// Build discovered models from `files` (each a GGUF path and its byte size).
21/// `repo` derives the repository label from a file's path (e.g. a relative
22/// `<org>/<model>` prefix), or `None`. Returns the models and any per-shard-set
23/// issues (a shard set missing its first part is skipped, not emitted).
24pub fn discovered_models(
25    files: &[(PathBuf, i64)],
26    kind: &SourceKind,
27    repo: impl Fn(&Path) -> Option<String>,
28) -> (Vec<DiscoveredModel>, Vec<String>) {
29    let (groups, loose) = group(files);
30    let mut bytes_by_path: HashMap<&Path, i64> = HashMap::new();
31    for (path, bytes) in files {
32        bytes_by_path.entry(path.as_path()).or_insert(*bytes);
33    }
34    let hint = gguf_hint();
35    let mut discovered = Vec::new();
36    let mut issues = Vec::new();
37
38    for path in &loose {
39        let name = path
40            .file_stem()
41            .and_then(|stem| stem.to_str())
42            .unwrap_or_default()
43            .to_owned();
44        let mut model = DiscoveredModel::new(name, source(kind, path, &repo));
45        apply_hint(&mut model, &hint);
46        model.footprint_bytes = bytes_by_path.get(path.as_path()).copied().unwrap_or(0);
47        model.primary_weight_path = Some(display(path));
48        discovered.push(model);
49    }
50
51    for shard_group in groups {
52        let Some(first) = shard_group.first_shard() else {
53            issues.push(format!(
54                "sharded model {} is missing its first part — skipped",
55                shard_group.base
56            ));
57            continue;
58        };
59        let mut model = DiscoveredModel::new(shard_group.base.clone(), source(kind, first, &repo));
60        apply_hint(&mut model, &hint);
61        model.footprint_bytes = shard_group.footprint_bytes();
62        model.primary_weight_path = Some(display(first));
63        model.downloading = !shard_group.complete();
64        discovered.push(model);
65    }
66
67    (discovered, issues)
68}
69
70fn source(kind: &SourceKind, path: &Path, repo: &impl Fn(&Path) -> Option<String>) -> ModelSource {
71    let mut source = ModelSource::new(kind.clone(), &display(path));
72    source.repo = repo(path);
73    source
74}
75
76fn apply_hint(model: &mut DiscoveredModel, hint: &crate::discovery::modality_hints::Hint) {
77    model.modality_hint = hint.modality.clone();
78    model.capabilities_hint = hint.capabilities.clone();
79    model.execution_hint = hint.execution;
80}
81
82fn display(path: &Path) -> String {
83    path.to_string_lossy().into_owned()
84}