tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Add vision/image support for local inference models (#8442)

Browse Source 
Signed-off-by: jh-block <jhugo@block.xyz>
This commit is contained in:
jh-block
2026-04-13 10:17:04 +02:00
committed by GitHub
parent 5fa2a8b821
commit de317d5445
15 changed files with 1181 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/routes/local_inference.rs
+146 -11
View File
@@ -1,3 +1,5 @@
use std::path::PathBuf;
use crate::routes::errors::ErrorResponse;
use crate::state::AppState;
use axum::{
@@ -13,9 +15,9 @@ use goose::providers::local_inference::{
available_inference_memory_bytes,
hf_models::{resolve_model_spec, HfGgufFile},
local_model_registry::{
default_settings_for_model, get_registry, is_featured_model, model_id_from_repo,
LocalModelEntry, ModelDownloadStatus as RegistryDownloadStatus, ModelSettings,
FEATURED_MODELS,
default_settings_for_model, featured_mmproj_spec, get_registry, is_featured_model,
model_id_from_repo, LocalModelEntry, ModelDownloadStatus as RegistryDownloadStatus,
ModelSettings, FEATURED_MODELS,
},
recommend_local_model,
};
@@ -47,10 +49,14 @@ pub struct LocalModelResponse {
pub status: ModelDownloadStatus,
pub recommended: bool,
pub settings: ModelSettings,
pub vision_capable: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub mmproj_status: Option<ModelDownloadStatus>,
}
async fn ensure_featured_models_in_registry() -> Result<(), ErrorResponse> {
let mut entries_to_add = Vec::new();
let mut mmproj_downloads_needed: Vec<(String, String, PathBuf)> = Vec::new();
for featured in FEATURED_MODELS {
let (repo_id, quantization) = match hf_models::parse_model_spec(featured.spec) {
@@ -64,8 +70,27 @@ async fn ensure_featured_models_in_registry() -> Result<(), ErrorResponse> {
let registry = get_registry()
.lock()
.map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?;
if registry.has_model(&model_id) {
continue;
if let Some(existing) = registry.get_model(&model_id) {
let needs_backfill = existing.mmproj_path.is_none() && featured.mmproj.is_some();
let needs_download = existing.is_downloaded()
&& featured.mmproj.is_some()
&& !existing.mmproj_path.as_ref().is_some_and(|p| p.exists());
if needs_download {
if let Some(mmproj) = featured.mmproj.as_ref() {
let path = mmproj.local_path();
let url = format!(
"https://huggingface.co/{}/resolve/main/{}",
mmproj.repo, mmproj.filename
);
mmproj_downloads_needed.push((model_id.clone(), url, path));
}
}
if !needs_backfill {
continue;
}
// Fall through to build the entry for sync_with_featured backfill
}
}
@@ -91,6 +116,8 @@ async fn ensure_featured_models_in_registry() -> Result<(), ErrorResponse> {
let local_path = Paths::in_data_dir("models").join(&hf_file.filename);
// enrich_with_featured_mmproj is called by sync_with_featured/add_model,
// so we don't need to populate mmproj fields here.
entries_to_add.push(LocalModelEntry {
id: model_id.clone(),
repo_id,
@@ -100,14 +127,58 @@ async fn ensure_featured_models_in_registry() -> Result<(), ErrorResponse> {
source_url: hf_file.download_url,
settings: default_settings_for_model(&model_id),
size_bytes: hf_file.size_bytes,
mmproj_path: None,
mmproj_source_url: None,
mmproj_size_bytes: 0,
});
}
if !entries_to_add.is_empty() {
{
let mut registry = get_registry()
.lock()
.map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?;
registry.sync_with_featured(entries_to_add);
if !entries_to_add.is_empty() {
registry.sync_with_featured(entries_to_add);
}
// Backfill mmproj data for all registry models and collect any
// needed mmproj downloads for models already on disk.
for model in registry.list_models_mut() {
model.enrich_with_featured_mmproj();
if model.is_downloaded() {
if let Some(mmproj) = featured_mmproj_spec(&model.id) {
let path = mmproj.local_path();
if !path.exists() {
let url = format!(
"https://huggingface.co/{}/resolve/main/{}",
mmproj.repo, mmproj.filename
);
mmproj_downloads_needed.push((model.id.clone(), url, path));
}
}
}
}
let _ = registry.save();
}
// Auto-download mmproj files for models that are already downloaded.
// Deduplicate by path since multiple quants share one mmproj file.
let dm = get_download_manager();
let mut started_paths = std::collections::HashSet::new();
for (model_id, url, path) in mmproj_downloads_needed {
if !path.exists() && started_paths.insert(path.clone()) {
let download_id = format!("{}-mmproj", model_id);
let dominated_by_active = dm
.get_progress(&download_id)
.is_some_and(|p| p.status == goose::download_manager::DownloadStatus::Downloading);
if !dominated_by_active {
tracing::info!(model_id = %model_id, "Auto-downloading vision encoder for existing model");
if let Err(e) = dm.download_model(download_id, url, path, None).await {
tracing::warn!(model_id = %model_id, error = %e, "Failed to start mmproj download");
}
}
}
}
Ok(())
@@ -154,6 +225,28 @@ pub async fn list_local_models(
let size_bytes = entry.file_size();
let vision_capable = entry.settings.vision_capable;
let mmproj_status = if vision_capable {
let ms = entry.mmproj_download_status();
Some(match ms {
RegistryDownloadStatus::NotDownloaded => ModelDownloadStatus::NotDownloaded,
RegistryDownloadStatus::Downloading {
progress_percent,
bytes_downloaded,
total_bytes,
speed_bps,
} => ModelDownloadStatus::Downloading {
progress_percent,
bytes_downloaded,
total_bytes,
speed_bps: Some(speed_bps),
},
RegistryDownloadStatus::Downloaded => ModelDownloadStatus::Downloaded,
})
} else {
None
};
models.push(LocalModelResponse {
id: entry.id.clone(),
repo_id: entry.repo_id.clone(),
@@ -163,6 +256,8 @@ pub async fn list_local_models(
status,
recommended: recommended_id == entry.id,
settings: entry.settings.clone(),
vision_capable,
mmproj_status,
});
}
@@ -276,16 +371,26 @@ pub async fn download_hf_model(
source_url: download_url.clone(),
settings: default_settings_for_model(&model_id),
size_bytes: hf_file.size_bytes,
mmproj_path: None,
mmproj_source_url: None,
mmproj_size_bytes: 0,
};
{
// add_model enriches the entry with mmproj metadata from the featured table
let mmproj_path = {
let mut registry = get_registry()
.lock()
.map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?;
registry
.add_model(entry)
.map_err(|e| ErrorResponse::internal(format!("{}", e)))?;
}
registry.get_model(&model_id).and_then(|e| {
e.mmproj_path
.as_ref()
.zip(e.mmproj_source_url.as_ref())
.map(|(p, u)| (p.clone(), u.clone()))
})
};
let dm = get_download_manager();
dm.download_model(
@@ -297,6 +402,19 @@ pub async fn download_hf_model(
.await
.map_err(|e| ErrorResponse::internal(format!("Download failed: {}", e)))?;
if let Some((mmproj_path, mmproj_url)) = mmproj_path {
if !mmproj_path.exists() {
dm.download_model(
format!("{}-mmproj", model_id),
mmproj_url,
mmproj_path,
None,
)
.await
.map_err(|e| ErrorResponse::internal(format!("mmproj download failed: {}", e)))?;
}
}
Ok((StatusCode::ACCEPTED, Json(model_id)))
}
@@ -338,6 +456,7 @@ pub async fn cancel_local_model_download(
manager
.cancel_download(&format!("{}-model", model_id))
.map_err(|e| ErrorResponse::internal(format!("{}", e)))?;
let _ = manager.cancel_download(&format!("{}-mmproj", model_id));
Ok(StatusCode::OK)
}
@@ -351,14 +470,22 @@ pub async fn cancel_local_model_download(
)
)]
pub async fn delete_local_model(Path(model_id): Path<String>) -> Result<StatusCode, ErrorResponse> {
let local_path = {
let (local_path, mmproj_path, other_uses_mmproj) = {
let registry = get_registry()
.lock()
.map_err(|_| ErrorResponse::internal("Failed to acquire registry lock"))?;
let entry = registry
.get_model(&model_id)
.ok_or_else(|| ErrorResponse::not_found("Model not found"))?;
entry.local_path.clone()
let lp = entry.local_path.clone();
let mp = entry.mmproj_path.clone();
// Check if another downloaded model shares this mmproj file
let shared = mp.as_ref().is_some_and(|target| {
registry.list_models().iter().any(|m| {
m.id != model_id && m.is_downloaded() && m.mmproj_path.as_ref() == Some(target)
})
});
(lp, mp, shared)
};
if local_path.exists() {
@@ -367,6 +494,14 @@ pub async fn delete_local_model(Path(model_id): Path<String>) -> Result<StatusCo
.map_err(|e| ErrorResponse::internal(format!("Failed to delete: {}", e)))?;
}
if !other_uses_mmproj {
if let Some(mmproj) = mmproj_path {
if mmproj.exists() {
let _ = tokio::fs::remove_file(&mmproj).await;
}
}
}
// Only remove non-featured models from registry (featured ones stay as placeholders)
if !is_featured_model(&model_id) {
let mut registry = get_registry()