use std::fs::{self, remove_dir_all, remove_file};
use std::io;
use std::path::Path;
use std::process::{Child, Command, Stdio};
use std::time::Duration;
use reqwest::StatusCode;
use reqwest::blocking::{Client, Response};
use crate::cli::Cli;
use crate::constants::{CIVITAI_API, DEFAULT_PAGE_SIZE, HF_BASE, HF_MIRROR_BASE, OLLAMA_LOCAL};
use crate::domain::{
CivitaiModel, CivitaiModelResponse, CloneStatus, HfModelResponse, ModelItem,
OllamaModelResponse,
};
use crate::error::{Result, ThesaError};
use crate::types::ModelProvider;
#[derive(Debug, Clone)]
pub(crate) enum ModelTarget {
Org(String),
Model(String),
Top,
Latest,
}
impl From<modelforge::ModelTarget> for ModelTarget {
fn from(value: modelforge::ModelTarget) -> Self {
match value {
modelforge::ModelTarget::Org(org) => Self::Org(org),
modelforge::ModelTarget::Model(model) => Self::Model(model),
modelforge::ModelTarget::Top => Self::Top,
modelforge::ModelTarget::Latest => Self::Latest,
}
}
}
pub(crate) fn parse_model_target_for_provider(
input: &str,
provider: ModelProvider,
) -> Result<ModelTarget> {
modelforge::parse_model_target_for_provider(input, provider.into())
.map(ModelTarget::from)
.map_err(|err| ThesaError::InvalidTarget(err.target().to_string()))
}
fn hf_auth_header(token: Option<&str>) -> Option<String> {
token.map(|t| format!("Bearer {t}"))
}
fn hf_api_base(mirror: bool) -> &'static str {
if mirror { HF_MIRROR_BASE } else { HF_BASE }
}
fn list_models_for_org(
client: &Client,
org: &str,
token: Option<&str>,
mirror: bool,
) -> Result<Vec<ModelItem>> {
list_hf_models_with_sort(client, Some(org), "downloads", -1, token, mirror)
}
fn list_hf_top_models(
client: &Client,
token: Option<&str>,
mirror: bool,
) -> Result<Vec<ModelItem>> {
list_hf_models_with_sort(client, None, "downloads", -1, token, mirror)
}
fn list_hf_latest_models(
client: &Client,
token: Option<&str>,
mirror: bool,
) -> Result<Vec<ModelItem>> {
list_hf_models_with_sort(client, None, "lastModified", -1, token, mirror)
}
fn list_hf_models_with_sort(
client: &Client,
org: Option<&str>,
sort: &str,
direction: i8,
token: Option<&str>,
mirror: bool,
) -> Result<Vec<ModelItem>> {
let base = hf_api_base(mirror);
let mut page = 1usize;
let mut models = Vec::new();
loop {
let mut url = format!(
"{}/api/models?sort={}&direction={}&limit={}&page={}",
base,
urlencoding::encode(sort),
direction,
DEFAULT_PAGE_SIZE,
page
);
if let Some(org) = org {
url.push_str("&author=");
url.push_str(&urlencoding::encode(org));
}
let mut request = client.get(&url);
if let Some(auth) = hf_auth_header(token) {
request = request.header("Authorization", auth);
}
let response = request.send()?;
let status = response.status();
if status == StatusCode::NOT_FOUND {
if let Some(org) = org {
return Err(ThesaError::HfNotFound(org.to_string()));
}
return Err(ThesaError::HfApiError {
status: status.as_u16(),
body: "empty response from Hugging Face listing".to_string(),
});
}
if !status.is_success() {
return Err(parse_hf_api_error(response));
}
let batch: Vec<HfModelResponse> = response.json()?;
let count = batch.len();
models.extend(batch.into_iter().map(|m| ModelItem {
id: if let Some(ref mid) = m.model_id {
mid.clone()
} else {
m.id
},
downloads: m.downloads,
pipeline_tag: m.pipeline_tag,
last_modified: m.last_modified,
provider: ModelProvider::Hf,
download_url: None,
size: None,
}));
if count < DEFAULT_PAGE_SIZE {
break;
}
page = page.saturating_add(1);
}
Ok(models)
}
fn fetch_single_model(
client: &Client,
model_id: &str,
token: Option<&str>,
mirror: bool,
) -> Result<ModelItem> {
let base = hf_api_base(mirror);
let url = format!("{}/api/models/{model_id}", base);
let mut request = client.get(&url);
if let Some(auth) = hf_auth_header(token) {
request = request.header("Authorization", auth);
}
let response = request.send()?;
if response.status() == StatusCode::NOT_FOUND {
return Err(ThesaError::HfNotFound(model_id.to_string()));
}
if !response.status().is_success() {
return Err(parse_hf_api_error(response));
}
let data: HfModelResponse = response.json()?;
Ok(ModelItem {
id: if let Some(ref mid) = data.model_id {
mid.clone()
} else {
data.id
},
downloads: data.downloads,
pipeline_tag: data.pipeline_tag,
last_modified: data.last_modified,
provider: ModelProvider::Hf,
download_url: None,
size: None,
})
}
pub(crate) fn fetch_models_for_target(
client: &Client,
target: &ModelTarget,
args: &Cli,
) -> Result<Vec<ModelItem>> {
match args.provider {
ModelProvider::Hf => match target {
ModelTarget::Org(org) => {
list_models_for_org(client, org, args.hf_token.as_deref(), args.hf_mirror)
}
ModelTarget::Model(id) => Ok(vec![fetch_single_model(
client,
id,
args.hf_token.as_deref(),
args.hf_mirror,
)?]),
ModelTarget::Top => {
list_hf_top_models(client, args.hf_token.as_deref(), args.hf_mirror)
}
ModelTarget::Latest => {
list_hf_latest_models(client, args.hf_token.as_deref(), args.hf_mirror)
}
},
ModelProvider::Ollama => match target {
ModelTarget::Org(org) => list_ollama_models(client, org),
ModelTarget::Model(id) => Ok(vec![fetch_single_ollama_model(client, id)?]),
ModelTarget::Top => list_ollama_top_models(client),
ModelTarget::Latest => list_ollama_latest_models(client),
},
ModelProvider::Civitai => match target {
ModelTarget::Org(org) if org.chars().all(|c| c.is_ascii_digit()) => {
if let Ok(model) =
fetch_single_civitai_model(client, org, args.civitai_token.as_deref())
{
return Ok(vec![model]);
}
list_civitai_models(
client,
Some(org),
"Most+Downloaded",
args.civitai_token.as_deref(),
)
}
ModelTarget::Org(org) => list_civitai_models(
client,
Some(org),
"Most+Downloaded",
args.civitai_token.as_deref(),
),
ModelTarget::Model(id) => Ok(vec![fetch_single_civitai_model(
client,
id,
args.civitai_token.as_deref(),
)?]),
ModelTarget::Top => list_civitai_models(
client,
None,
"Most+Downloaded",
args.civitai_token.as_deref(),
),
ModelTarget::Latest => {
list_civitai_models(client, None, "Newest", args.civitai_token.as_deref())
}
},
}
}
fn list_ollama_models(client: &Client, org: &str) -> Result<Vec<ModelItem>> {
list_ollama_models_by_query(client, Some(org), true)
}
fn list_ollama_top_models(client: &Client) -> Result<Vec<ModelItem>> {
let mut models = list_ollama_models_by_query(client, None, false)?;
models.sort_by_key(|model| std::cmp::Reverse(model.size));
Ok(models)
}
fn list_ollama_latest_models(client: &Client) -> Result<Vec<ModelItem>> {
let mut models = list_ollama_models_by_query(client, None, false)?;
models.sort_by(|a, b| {
let b_modified = b.last_modified.as_deref().unwrap_or("");
let a_modified = a.last_modified.as_deref().unwrap_or("");
b_modified.cmp(a_modified)
});
Ok(models)
}
fn list_ollama_models_by_query(
client: &Client,
query: Option<&str>,
fallback_on_empty: bool,
) -> Result<Vec<ModelItem>> {
let local_url = format!("{}/api/tags", OLLAMA_LOCAL);
let query = query.map(str::to_lowercase).unwrap_or_default();
if let Ok(resp) = client
.get(&local_url)
.timeout(Duration::from_secs(2))
.send()
{
if resp.status().is_success() {
if let Ok(data) = resp.json::<OllamaModelResponse>() {
let models = data
.models
.into_iter()
.filter(|m| query.is_empty() || m.name.to_lowercase().contains(&query))
.map(|m| ModelItem {
id: m.name.clone(),
downloads: 0,
pipeline_tag: None,
last_modified: Some(m.modified_at),
provider: ModelProvider::Ollama,
download_url: None,
size: Some(m.size),
})
.collect::<Vec<_>>();
if !models.is_empty() {
return Ok(models);
}
}
}
}
if fallback_on_empty && !query.is_empty() {
return Ok(vec![ModelItem {
id: query.to_string(),
downloads: 0,
pipeline_tag: None,
last_modified: None,
provider: ModelProvider::Ollama,
download_url: None,
size: None,
}]);
}
Ok(Vec::new())
}
fn fetch_single_ollama_model(_client: &Client, model_id: &str) -> Result<ModelItem> {
Ok(ModelItem {
id: model_id.to_string(),
downloads: 0,
pipeline_tag: None,
last_modified: None,
provider: ModelProvider::Ollama,
download_url: None,
size: None,
})
}
fn list_civitai_models(
client: &Client,
query: Option<&str>,
sort: &str,
token: Option<&str>,
) -> Result<Vec<ModelItem>> {
let mut url = format!(
"{}/models?limit={}&sort={}",
CIVITAI_API, DEFAULT_PAGE_SIZE, sort
);
if let Some(query) = query {
if !query.is_empty() {
url.push_str("&query=");
url.push_str(&urlencoding::encode(query));
}
}
let mut request = client.get(&url);
if let Some(token) = token {
request = request.header("Authorization", format!("Bearer {token}"));
}
let response = request.send()?;
let status = response.status();
if !status.is_success() {
let body = response
.text()
.unwrap_or_else(|_| "<unable to read body>".to_string());
return Err(ThesaError::CivitaiApiError {
status: status.as_u16(),
body,
});
}
let data: CivitaiModelResponse = response.json()?;
let models = data
.items
.into_iter()
.map(|m| {
let (download_url, size) = m
.model_versions
.as_ref()
.and_then(|v| v.first())
.map(|v| (v.download_url.clone(), v.size))
.unwrap_or((None, None));
ModelItem {
id: m.name.clone(),
downloads: m.downloads,
pipeline_tag: m.model_type,
last_modified: None,
provider: ModelProvider::Civitai,
download_url,
size,
}
})
.collect::<Vec<_>>();
if models.is_empty() {
return Err(ThesaError::CivitaiNotFound(
query.unwrap_or("civitai top models").to_string(),
));
}
Ok(models)
}
fn fetch_single_civitai_model(
client: &Client,
model_id: &str,
token: Option<&str>,
) -> Result<ModelItem> {
let numeric_id: u64 = model_id.parse().map_err(|_| {
ThesaError::CivitaiNotFound(format!(
"{model_id} (CivitAI requires numeric model ID, e.g. 827184)"
))
})?;
let url = format!("{}/models/{}", CIVITAI_API, numeric_id);
let mut request = client.get(&url);
if let Some(token) = token {
request = request.header("Authorization", format!("Bearer {token}"));
}
let response = request.send()?;
let status = response.status();
if status == StatusCode::NOT_FOUND {
return Err(ThesaError::CivitaiNotFound(model_id.to_string()));
}
if !status.is_success() {
let body = response
.text()
.unwrap_or_else(|_| "<unable to read body>".to_string());
return Err(ThesaError::CivitaiApiError {
status: status.as_u16(),
body,
});
}
let m: CivitaiModel = response.json()?;
let (download_url, size) = m
.model_versions
.as_ref()
.and_then(|v| v.first())
.map(|v| (v.download_url.clone(), v.size))
.unwrap_or((None, None));
Ok(ModelItem {
id: m.name.clone(),
downloads: m.downloads,
pipeline_tag: m.model_type,
last_modified: None,
provider: ModelProvider::Civitai,
download_url,
size,
})
}
fn parse_hf_api_error(response: Response) -> ThesaError {
let status = response.status();
let body = response
.text()
.unwrap_or_else(|_| "<unable to read body>".to_string());
ThesaError::HfApiError {
status: status.as_u16(),
body,
}
}
pub(crate) fn download_model(model: &ModelItem, output: &Path, args: &Cli) -> Result<CloneStatus> {
match model.provider {
ModelProvider::Hf => download_hf_model(model, output, args),
ModelProvider::Ollama => download_ollama_model(model, output, args),
ModelProvider::Civitai => download_civitai_model(model, output, args),
}
}
fn download_hf_model(model: &ModelItem, output: &Path, args: &Cli) -> Result<CloneStatus> {
let destination = output.join(model.id.replace('/', "_"));
if destination.exists() {
if args.skip_existing {
return Ok(CloneStatus::Skipped);
}
if destination.is_dir() {
remove_dir_all(&destination)?;
} else {
remove_file(&destination)?;
}
}
let child = spawn_hf_clone(model, output, args)?;
let status = child.wait_with_output()?;
if status.status.success() {
Ok(CloneStatus::Cloned)
} else {
let code = status
.status
.code()
.map_or_else(|| "signal/unknown".to_string(), |c| c.to_string());
Err(ThesaError::HfDownloadFailed {
model: model.id.clone(),
message: format!("exit code {code}"),
})
}
}
fn download_ollama_model(model: &ModelItem, output: &Path, args: &Cli) -> Result<CloneStatus> {
let destination = output.join(model.id.replace(['/', ':'], "_"));
if destination.exists() {
if args.skip_existing {
return Ok(CloneStatus::Skipped);
}
if destination.is_dir() {
remove_dir_all(&destination)?;
} else {
remove_file(&destination)?;
}
}
fs::create_dir_all(&destination)?;
let command_args = vec!["pull".to_string(), model.id.clone()];
let child = match Command::new("ollama")
.current_dir(&destination)
.args(&command_args)
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
{
Ok(child) => child,
Err(error) if error.kind() == io::ErrorKind::NotFound => {
return Err(ThesaError::MissingOllama);
}
Err(error) => return Err(error.into()),
};
let status = child.wait_with_output()?;
if status.status.success() {
Ok(CloneStatus::Cloned)
} else {
let code = status
.status
.code()
.map_or_else(|| "signal/unknown".to_string(), |c| c.to_string());
Err(ThesaError::OllamaDownloadFailed {
model: model.id.clone(),
message: format!("exit code {code}"),
})
}
}
fn download_civitai_model(model: &ModelItem, output: &Path, args: &Cli) -> Result<CloneStatus> {
let dir_name = model.id.replace('/', "_");
let destination = output.join(&dir_name);
if destination.exists() {
if args.skip_existing {
return Ok(CloneStatus::Skipped);
}
if destination.is_dir() {
remove_dir_all(&destination)?;
} else {
remove_file(&destination)?;
}
}
let url = model
.download_url
.as_deref()
.ok_or_else(|| ThesaError::CivitaiDownloadFailed {
model: model.id.clone(),
message: "no download URL available".to_string(),
})?;
let mut request = client_get().get(url);
if let Some(token) = args.civitai_token.as_deref() {
request = request.header("Authorization", format!("Bearer {token}"));
}
let response = request.send()?;
let status = response.status();
if !status.is_success() {
let body = response
.text()
.unwrap_or_else(|_| "<unable to read body>".to_string());
return Err(ThesaError::CivitaiDownloadFailed {
model: model.id.clone(),
message: format!("HTTP {status}: {body}"),
});
}
let bytes = response.bytes()?;
fs::create_dir_all(&destination)?;
let filename = url.rsplit('/').next().unwrap_or("model.safetensors");
let filepath = destination.join(filename);
fs::write(&filepath, &bytes)?;
Ok(CloneStatus::Cloned)
}
fn client_get() -> Client {
Client::builder()
.user_agent("thesa")
.build()
.unwrap_or_else(|_| Client::new())
}
fn spawn_hf_clone(model: &ModelItem, output: &Path, args: &Cli) -> Result<Child> {
let url = format!("https://huggingface.co/{}", model.id);
let dir_name = model.id.replace('/', "_");
let mut command_args = vec!["clone".to_string()];
if let Some(depth) = args.depth {
command_args.push("--depth".to_string());
command_args.push(depth.to_string());
}
command_args.push(url);
command_args.push(dir_name);
match Command::new("git")
.current_dir(output)
.args(&command_args)
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
{
Ok(child) => Ok(child),
Err(error) if error.kind() == io::ErrorKind::NotFound => Err(ThesaError::MissingGit),
Err(error) => Err(error.into()),
}
}