1use std::{
2 ffi::OsStr,
3 path::{Path, PathBuf},
4};
5
6use anyhow::{Context, Result};
7use async_trait::async_trait;
8use hf_hub::{
9 HFClient, HFClientBuilder, RepoType, RepoTypeModel,
10 cache::{CachedRepoInfo, HFCacheInfo},
11 repository::ModelInfo,
12};
13use model_artifact::{ModelArtifactFile, ModelIdentity, ModelRepository, ResolvedModelArtifact};
14use model_ref::{
15 format_canonical_ref, format_model_ref, normalize_gguf_distribution_id,
16 quant_selector_from_gguf_file,
17};
18use serde::{Deserialize, Serialize};
19
20#[derive(Clone)]
21pub struct HfModelRepository {
22 api: HFClient,
23 cache_dir: PathBuf,
24}
25
26impl HfModelRepository {
27 pub fn from_env() -> Result<Self> {
28 Self::builder().build()
29 }
30
31 pub fn builder() -> HfModelRepositoryBuilder {
32 HfModelRepositoryBuilder::default()
33 }
34
35 pub fn cache_dir(&self) -> &Path {
36 &self.cache_dir
37 }
38
39 pub async fn download_file(&self, repo: &str, revision: &str, file: &str) -> Result<PathBuf> {
40 let (owner, name) = repo_parts(repo);
41 self.api
42 .model(owner, name)
43 .download_file()
44 .filename(file.to_string())
45 .revision(revision.to_string())
46 .send()
47 .await
48 .with_context(|| format!("download Hugging Face model file {repo}@{revision}/{file}"))
49 }
50
51 pub async fn download_artifact_files(
52 &self,
53 artifact: &ResolvedModelArtifact,
54 ) -> Result<Vec<PathBuf>> {
55 let mut paths = Vec::with_capacity(artifact.files.len());
56 for file in &artifact.files {
57 paths.push(
58 self.download_file(&artifact.source_repo, &artifact.source_revision, &file.path)
59 .await?,
60 );
61 }
62 Ok(paths)
63 }
64
65 pub fn identity_for_path(&self, path: &Path) -> Option<HfModelIdentity> {
66 huggingface_identity_for_path_in_cache(path, &self.cache_dir)
67 }
68}
69
70#[derive(Default)]
71pub struct HfModelRepositoryBuilder {
72 cache_dir: Option<PathBuf>,
73 endpoint: Option<String>,
74 token: Option<String>,
75}
76
77impl HfModelRepositoryBuilder {
78 pub fn cache_dir(mut self, cache_dir: impl Into<PathBuf>) -> Self {
79 self.cache_dir = Some(cache_dir.into());
80 self
81 }
82
83 pub fn endpoint(mut self, endpoint: impl Into<String>) -> Self {
84 self.endpoint = Some(endpoint.into());
85 self
86 }
87
88 pub fn token(mut self, token: impl Into<String>) -> Self {
89 self.token = Some(token.into());
90 self
91 }
92
93 pub fn build(self) -> Result<HfModelRepository> {
94 let cache_dir = self.cache_dir.unwrap_or_else(huggingface_hub_cache_dir);
95 let mut builder = HFClientBuilder::new().cache_dir(cache_dir.clone());
96
97 let endpoint = self
98 .endpoint
99 .or_else(|| std::env::var("HF_ENDPOINT").ok())
100 .map(|endpoint| endpoint.trim().to_string())
101 .filter(|endpoint| !endpoint.is_empty());
102 if let Some(endpoint) = endpoint {
103 builder = builder.endpoint(endpoint);
104 }
105
106 let token = self.token.or_else(hf_token_override);
107 if let Some(token) = token {
108 builder = builder.token(token);
109 }
110
111 let api = builder.build().context("build Hugging Face API client")?;
112 Ok(HfModelRepository { api, cache_dir })
113 }
114}
115
116#[async_trait]
117impl ModelRepository for HfModelRepository {
118 async fn resolve_revision(&self, repo: &str, revision: Option<&str>) -> Result<String> {
119 let revision = revision.unwrap_or("main");
120 self.repo_info(repo, revision)
121 .await?
122 .sha
123 .with_context(|| format!("Hugging Face repo {repo}@{revision} did not return a sha"))
124 }
125
126 async fn list_files(&self, repo: &str, revision: &str) -> Result<Vec<ModelArtifactFile>> {
127 let info = self.repo_info(repo, revision).await?;
128 Ok(info
129 .siblings
130 .unwrap_or_default()
131 .into_iter()
132 .map(|sibling| ModelArtifactFile {
133 path: sibling.rfilename,
134 size_bytes: sibling.size,
135 sha256: None,
136 })
137 .collect())
138 }
139}
140
141impl HfModelRepository {
142 async fn repo_info(&self, repo: &str, revision: &str) -> Result<ModelInfo> {
143 let (owner, name) = repo_parts(repo);
144 self.api
145 .model(owner, name)
146 .info()
147 .revision(revision.to_string())
148 .send()
149 .await
150 .with_context(|| format!("fetch Hugging Face model repo {repo}@{revision}"))
151 }
152}
153
154#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
155pub struct HfModelIdentity {
156 pub model_id: String,
157 pub repo_id: String,
158 pub revision: String,
159 pub file: String,
160 pub canonical_ref: String,
161 pub distribution_id: Option<String>,
162 pub selector: Option<String>,
163}
164
165impl HfModelIdentity {
166 pub fn to_model_identity(&self) -> ModelIdentity {
167 ModelIdentity {
168 model_id: self.model_id.clone(),
169 source_repo: Some(self.repo_id.clone()),
170 source_revision: Some(self.revision.clone()),
171 source_file: Some(self.file.clone()),
172 canonical_ref: Some(self.canonical_ref.clone()),
173 distribution_id: self.distribution_id.clone(),
174 selector: self.selector.clone(),
175 }
176 }
177
178 pub fn distribution_ref(&self) -> Option<String> {
179 self.distribution_id.as_ref().map(|distribution_id| {
180 format!("{}@{}/{}", self.repo_id, self.revision, distribution_id)
181 })
182 }
183}
184
185pub fn huggingface_hub_cache_dir() -> PathBuf {
186 if let Some(path) = env_path("HF_HUB_CACHE") {
187 return path;
188 }
189 if let Some(path) = env_path("HUGGINGFACE_HUB_CACHE") {
190 return path;
191 }
192 if let Some(path) = env_path("HF_HOME") {
193 return path.join("hub");
194 }
195 if let Some(path) = env_path("XDG_CACHE_HOME") {
196 return path.join("huggingface").join("hub");
197 }
198 std::env::var("HOME")
199 .map(PathBuf::from)
200 .unwrap_or_else(|_| PathBuf::from("."))
201 .join(".cache")
202 .join("huggingface")
203 .join("hub")
204}
205
206pub fn hf_token_override() -> Option<String> {
207 for key in ["HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"] {
208 if let Ok(token) = std::env::var(key) {
209 let token = token.trim();
210 if !token.is_empty() {
211 return Some(token.to_string());
212 }
213 }
214 }
215 None
216}
217
218pub fn huggingface_repo_folder_name(repo_id: &str, repo_type: impl RepoType) -> String {
219 let type_plural = repo_type.plural();
220 std::iter::once(type_plural)
221 .chain(repo_id.split('/'))
222 .collect::<Vec<_>>()
223 .join("--")
224}
225
226pub fn huggingface_snapshot_path(
227 repo_id: &str,
228 repo_type: impl RepoType,
229 revision: &str,
230) -> PathBuf {
231 huggingface_hub_cache_dir()
232 .join(huggingface_repo_folder_name(repo_id, repo_type))
233 .join("snapshots")
234 .join(revision)
235}
236
237pub fn huggingface_identity_for_path_in_cache(
238 path: &Path,
239 cache_root: &Path,
240) -> Option<HfModelIdentity> {
241 if let Some(identity) = identity_from_cache_snapshot_path(path, cache_root) {
242 return Some(identity);
243 }
244 let resolved_cache_root = cache_root
245 .canonicalize()
246 .unwrap_or_else(|_| cache_root.to_path_buf());
247 if resolved_cache_root != cache_root
248 && let Some(identity) = identity_from_cache_snapshot_path(path, &resolved_cache_root)
249 {
250 return Some(identity);
251 }
252 let resolved = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
253 if resolved != path {
254 if let Some(identity) = identity_from_cache_snapshot_path(&resolved, cache_root) {
255 return Some(identity);
256 }
257 if resolved_cache_root != cache_root
258 && let Some(identity) =
259 identity_from_cache_snapshot_path(&resolved, &resolved_cache_root)
260 {
261 return Some(identity);
262 }
263 }
264 if let Some(identity) = identity_from_snapshot_layout_ancestors(path) {
265 return Some(identity);
266 }
267 if resolved != path
268 && let Some(identity) = identity_from_snapshot_layout_ancestors(&resolved)
269 {
270 return Some(identity);
271 }
272 scan_hf_cache_identity_for_path(path, cache_root)
273}
274
275fn identity_from_cache_snapshot_path(path: &Path, cache_root: &Path) -> Option<HfModelIdentity> {
276 let relative = path.strip_prefix(cache_root).ok()?;
277 let mut components = relative.components();
278 let repo_folder = components.next()?.as_os_str().to_str()?;
279 let repo_id = parse_model_repo_folder_name(repo_folder)?;
280 if components.next()?.as_os_str() != OsStr::new("snapshots") {
281 return None;
282 }
283 let revision = components.next()?.as_os_str().to_str()?.to_string();
284 let file = components
285 .map(|component| component.as_os_str().to_str())
286 .collect::<Option<Vec<_>>>()?
287 .join("/");
288 if file.is_empty() {
289 return None;
290 }
291 Some(identity_from_parts(repo_id, revision, file))
292}
293
294fn identity_from_snapshot_layout_ancestors(path: &Path) -> Option<HfModelIdentity> {
295 for revision_dir in path.ancestors() {
296 let Some(snapshots_dir) = revision_dir.parent() else {
297 continue;
298 };
299 if snapshots_dir.file_name()? != OsStr::new("snapshots") {
300 continue;
301 }
302 let repo_dir = snapshots_dir.parent()?;
303 let repo_folder = repo_dir.file_name()?.to_str()?;
304 let repo_id = parse_model_repo_folder_name(repo_folder)?;
305 let revision = revision_dir.file_name()?.to_str()?.to_string();
306 let file = path
307 .strip_prefix(revision_dir)
308 .ok()?
309 .components()
310 .map(|component| component.as_os_str().to_str())
311 .collect::<Option<Vec<_>>>()?
312 .join("/");
313 if file.is_empty() {
314 continue;
315 }
316 return Some(identity_from_parts(repo_id, revision, file));
317 }
318 None
319}
320
321fn scan_hf_cache_identity_for_path(path: &Path, cache_root: &Path) -> Option<HfModelIdentity> {
322 let cache_info = scan_hf_cache_info(cache_root)?;
323 let resolved = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
324
325 for repo in &cache_info.repos {
326 let Some(repo_id) = cache_repo_id(repo) else {
327 continue;
328 };
329 for revision in &repo.revisions {
330 for file in &revision.files {
331 let candidate = file
332 .file_path
333 .canonicalize()
334 .unwrap_or_else(|_| file.file_path.clone());
335 if file.file_path != path && candidate != resolved {
336 continue;
337 }
338
339 let relative_path = file
340 .file_path
341 .strip_prefix(&revision.snapshot_path)
342 .ok()?
343 .to_string_lossy()
344 .replace('\\', "/");
345 if relative_path.is_empty() {
346 return None;
347 }
348
349 return Some(identity_from_parts(
350 repo_id.to_string(),
351 revision.commit_hash.clone(),
352 relative_path,
353 ));
354 }
355 }
356 }
357 None
358}
359
360fn scan_hf_cache_info(cache_root: &Path) -> Option<HFCacheInfo> {
361 let cache_root = cache_root.to_path_buf();
362 let scan = move || {
363 let runtime = tokio::runtime::Builder::new_current_thread()
364 .enable_all()
365 .build()
366 .ok()?;
367 runtime
368 .block_on(
369 HFClientBuilder::new()
370 .cache_dir(cache_root)
371 .build()
372 .ok()?
373 .scan_cache()
374 .send(),
375 )
376 .ok()
377 };
378
379 if tokio::runtime::Handle::try_current().is_ok() {
380 std::thread::spawn(scan).join().ok().flatten()
381 } else {
382 scan()
383 }
384}
385
386fn identity_from_parts(repo_id: String, revision: String, file: String) -> HfModelIdentity {
387 let selector = quant_selector_from_gguf_file(&file);
388 let model_id = format_model_ref(&repo_id, None, selector.as_deref());
389 let distribution_id = normalize_gguf_distribution_id(&file);
390 let canonical_ref = format_canonical_ref(&repo_id, &revision, &file);
391 HfModelIdentity {
392 model_id,
393 repo_id,
394 revision,
395 file,
396 canonical_ref,
397 distribution_id,
398 selector,
399 }
400}
401
402fn cache_repo_id(repo: &CachedRepoInfo) -> Option<&str> {
403 (repo.repo_type == RepoTypeModel.singular()).then_some(repo.repo_id.as_str())
404}
405
406fn parse_model_repo_folder_name(folder: &str) -> Option<String> {
407 folder
408 .strip_prefix("models--")
409 .map(|value| value.replace("--", "/"))
410}
411
412fn repo_parts(repo: &str) -> (&str, &str) {
413 repo.split_once('/').unwrap_or(("", repo))
414}
415
416fn env_path(key: &str) -> Option<PathBuf> {
417 let value = std::env::var(key).ok()?;
418 let trimmed = value.trim();
419 (!trimmed.is_empty()).then(|| PathBuf::from(trimmed))
420}
421
422#[cfg(test)]
423mod tests {
424 use super::*;
425 use std::path::Path;
426
427 #[test]
428 fn cache_path_identity_matches_mesh_snapshot_layout() {
429 let cache_root = PathBuf::from("/cache/hub");
430 let path = cache_root
431 .join("models--org--repo")
432 .join("snapshots")
433 .join("abc123")
434 .join("Qwen3-8B-Q4_K_M.gguf");
435
436 let identity = huggingface_identity_for_path_in_cache(&path, &cache_root).unwrap();
437 assert_eq!(identity.model_id, "org/repo:Q4_K_M");
438 assert_eq!(identity.repo_id, "org/repo");
439 assert_eq!(identity.revision, "abc123");
440 assert_eq!(identity.file, "Qwen3-8B-Q4_K_M.gguf");
441 assert_eq!(
442 identity.canonical_ref,
443 "org/repo@abc123/Qwen3-8B-Q4_K_M.gguf"
444 );
445 assert_eq!(identity.distribution_id.as_deref(), Some("Qwen3-8B-Q4_K_M"));
446 assert_eq!(
447 identity.distribution_ref().as_deref(),
448 Some("org/repo@abc123/Qwen3-8B-Q4_K_M")
449 );
450 }
451
452 #[test]
453 fn cache_path_identity_collapses_split_gguf_distribution() {
454 let cache_root = PathBuf::from("/cache/hub");
455 let path = cache_root
456 .join("models--org--repo")
457 .join("snapshots")
458 .join("abc123")
459 .join("UD-IQ2_M")
460 .join("GLM-5.1-UD-IQ2_M-00001-of-00006.gguf");
461
462 let identity = huggingface_identity_for_path_in_cache(&path, &cache_root).unwrap();
463 assert_eq!(identity.model_id, "org/repo:UD-IQ2_M");
464 assert_eq!(identity.selector.as_deref(), Some("UD-IQ2_M"));
465 assert_eq!(
466 identity.distribution_id.as_deref(),
467 Some("GLM-5.1-UD-IQ2_M")
468 );
469 }
470
471 #[test]
472 fn cache_path_identity_falls_back_to_snapshot_layout_ancestors() {
473 let path = PathBuf::from("/alternate/root")
474 .join("models--org--repo")
475 .join("snapshots")
476 .join("abc123")
477 .join("nested")
478 .join("Qwen3-8B-Q4_K_M.gguf");
479
480 let identity =
481 huggingface_identity_for_path_in_cache(&path, Path::new("/unrelated/cache")).unwrap();
482
483 assert_eq!(identity.model_id, "org/repo:Q4_K_M");
484 assert_eq!(identity.repo_id, "org/repo");
485 assert_eq!(identity.revision, "abc123");
486 assert_eq!(identity.file, "nested/Qwen3-8B-Q4_K_M.gguf");
487 }
488
489 #[test]
490 fn repo_folder_name_matches_huggingface_cache_layout() {
491 assert_eq!(
492 huggingface_repo_folder_name("org/repo", RepoTypeModel),
493 "models--org--repo"
494 );
495 }
496}