From d56b115bb4f27e712f1d78b662962b251fc73c62 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 25 Mar 2026 13:24:29 +0000 Subject: [PATCH] feat: multi-source model discovery (HF default, legacy cache, LM Studio) --- studio/backend/models/models.py | 6 +- studio/backend/routes/models.py | 230 +++++++++++++----- studio/backend/utils/paths/__init__.py | 4 + studio/backend/utils/paths/storage_roots.py | 50 +++- .../studio/sections/model-section.tsx | 6 +- .../src/features/training/api/models-api.ts | 3 +- 6 files changed, 223 insertions(+), 76 deletions(-) diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index daa8eec907..046e36137d 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -165,7 +165,7 @@ class LocalModelInfo(BaseModel): id: str = Field(..., description = "Identifier to use for loading/training") display_name: str = Field(..., description = "Display label") path: str = Field(..., description = "Local path where model data was discovered") - source: Literal["models_dir", "hf_cache"] = Field( + source: Literal["models_dir", "hf_cache", "lmstudio"] = Field( ..., description = "Discovery source", ) @@ -189,6 +189,10 @@ class LocalModelListResponse(BaseModel): None, description = "HF cache root that was scanned", ) + lmstudio_dirs: List[str] = Field( + default_factory = list, + description = "LM Studio model directories that were scanned", + ) models: List[LocalModelInfo] = Field( default_factory = list, description = "Discovered local/cached models", diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index e705762447..63c9304a64 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -210,6 +210,76 @@ def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]: return found +def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: + """Scan an LM Studio models directory for model files. + + LM Studio uses a ``publisher/model-name`` folder structure containing + GGUF files, or standalone GGUF files at the top level. + """ + if not lm_dir.exists() or not lm_dir.is_dir(): + return [] + + found: List[LocalModelInfo] = [] + for child in lm_dir.iterdir(): + if not child.is_dir(): + if child.suffix == ".gguf" and child.is_file(): + try: + updated_at = child.stat().st_mtime + except OSError: + updated_at = None + found.append( + LocalModelInfo( + id = str(child), + display_name = child.stem, + path = str(child), + source = "lmstudio", + updated_at = updated_at, + ), + ) + continue + + # child is a publisher directory — scan its sub-directories + for model_dir in child.iterdir(): + if model_dir.is_dir(): + has_model = ( + any(model_dir.glob("*.gguf")) + or (model_dir / "config.json").exists() + or any(model_dir.glob("*.safetensors")) + ) + if not has_model: + continue + model_id = f"{child.name}/{model_dir.name}" + try: + updated_at = model_dir.stat().st_mtime + except OSError: + updated_at = None + found.append( + LocalModelInfo( + id = model_id, + model_id = model_id, + display_name = model_dir.name, + path = str(model_dir), + source = "lmstudio", + updated_at = updated_at, + ), + ) + elif model_dir.suffix == ".gguf" and model_dir.is_file(): + try: + updated_at = model_dir.stat().st_mtime + except OSError: + updated_at = None + found.append( + LocalModelInfo( + id = str(model_dir), + display_name = model_dir.stem, + path = str(model_dir), + source = "lmstudio", + updated_at = updated_at, + ), + ) + return found + + @router.get("/local", response_model = LocalModelListResponse) async def list_local_models( models_dir: str = Query( @@ -218,13 +288,24 @@ async def list_local_models( current_subject: str = Depends(get_current_subject), ): """ - List local model candidates from custom models dir and HF cache. + List local model candidates from custom models dir, HF cache, + legacy Unsloth HF cache, and LM Studio directories. """ + from utils.paths import legacy_hf_cache_dir, lmstudio_model_dirs + + # Resolve all scan directories up front. + hf_cache_dir = _resolve_hf_cache_dir() + legacy_hf = legacy_hf_cache_dir() + lm_dirs = lmstudio_model_dirs() + # Validate models_dir against an allowlist of trusted directories. # Only the trusted Path objects are used for filesystem access -- the # user-supplied string is only used for matching, never for path construction. - hf_cache_dir = _resolve_hf_cache_dir() - allowed_roots = [Path("./models").resolve(), hf_cache_dir] + allowed_roots: list[Path] = [Path("./models").resolve(), hf_cache_dir] + if legacy_hf.is_dir(): + allowed_roots.append(legacy_hf) + for d in lm_dirs: + allowed_roots.append(d) try: from utils.paths import studio_root, outputs_root @@ -248,6 +329,14 @@ async def list_local_models( try: local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir) + # Scan legacy Unsloth HF cache for backward compatibility + if legacy_hf.is_dir() and legacy_hf.resolve() != hf_cache_dir.resolve(): + local_models += _scan_hf_cache(legacy_hf) + + # Scan LM Studio directories + for lm_dir in lm_dirs: + local_models += _scan_lmstudio_dir(lm_dir) + deduped: dict[str, LocalModelInfo] = {} for model in local_models: if model.id not in deduped: @@ -262,6 +351,7 @@ async def list_local_models( return LocalModelListResponse( models_dir = str(models_root), hf_cache_dir = str(hf_cache_dir), + lmstudio_dirs = [str(d) for d in lm_dirs], models = models, ) except Exception as e: @@ -850,42 +940,44 @@ def _get_repo_size_cached(repo_id: str) -> int: async def list_cached_gguf( current_subject: str = Depends(get_current_subject), ): - """List GGUF repos that have already been downloaded to the HF cache. - - Uses scan_cache_dir() for proper repo IDs, then deduplicates by - lowercased key (HF cache dirs are lowercased but the canonical repo - ID preserves casing). - """ + """List GGUF repos downloaded to HF cache and legacy Unsloth cache.""" try: from huggingface_hub import scan_cache_dir + from utils.paths import legacy_hf_cache_dir + + cache_scans = [scan_cache_dir()] + legacy_hf = legacy_hf_cache_dir() + if legacy_hf.is_dir(): + try: + cache_scans.append(scan_cache_dir(cache_dir = str(legacy_hf))) + except Exception: + pass - hf_cache = scan_cache_dir() seen_lower: dict[str, dict] = {} - for repo_info in hf_cache.repos: - if repo_info.repo_type != "model": - continue - repo_id = repo_info.repo_id - if not repo_id.upper().endswith("-GGUF"): - continue - # Check for actual .gguf files and sum sizes - total_size = 0 - has_gguf = False - for revision in repo_info.revisions: - for f in revision.files: - if f.file_name.endswith(".gguf"): - has_gguf = True - total_size += f.size_on_disk - if not has_gguf: - continue - # Deduplicate: keep the entry with the most data - key = repo_id.lower() - existing = seen_lower.get(key) - if existing is None or total_size > existing["size_bytes"]: - seen_lower[key] = { - "repo_id": repo_id, - "size_bytes": total_size, - "cache_path": str(repo_info.repo_path), - } + for hf_cache in cache_scans: + for repo_info in hf_cache.repos: + if repo_info.repo_type != "model": + continue + repo_id = repo_info.repo_id + if not repo_id.upper().endswith("-GGUF"): + continue + total_size = 0 + has_gguf = False + for revision in repo_info.revisions: + for f in revision.files: + if f.file_name.endswith(".gguf"): + has_gguf = True + total_size += f.size_on_disk + if not has_gguf: + continue + key = repo_id.lower() + existing = seen_lower.get(key) + if existing is None or total_size > existing["size_bytes"]: + seen_lower[key] = { + "repo_id": repo_id, + "size_bytes": total_size, + "cache_path": str(repo_info.repo_path), + } cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) return {"cached": cached} except Exception as e: @@ -897,44 +989,48 @@ async def list_cached_gguf( async def list_cached_models( current_subject: str = Depends(get_current_subject), ): - """List non-GGUF model repos that have been downloaded to the HF cache. - - Only includes repos that actually contain model weight files - (.safetensors, .bin), not repos with only config/metadata. - """ + """List non-GGUF model repos downloaded to HF cache and legacy Unsloth cache.""" _WEIGHT_EXTENSIONS = (".safetensors", ".bin") try: from huggingface_hub import scan_cache_dir + from utils.paths import legacy_hf_cache_dir + + cache_scans = [scan_cache_dir()] + legacy_hf = legacy_hf_cache_dir() + if legacy_hf.is_dir(): + try: + cache_scans.append(scan_cache_dir(cache_dir = str(legacy_hf))) + except Exception: + pass - hf_cache = scan_cache_dir() seen_lower: dict[str, dict] = {} - for repo_info in hf_cache.repos: - if repo_info.repo_type != "model": - continue - repo_id = repo_info.repo_id - if repo_id.upper().endswith("-GGUF"): - continue - total_size = sum( - f.size_on_disk for rev in repo_info.revisions for f in rev.files - ) - if total_size == 0: - continue - # Skip repos that only have config/metadata files (no weights) - has_weights = any( - f.file_name.endswith(_WEIGHT_EXTENSIONS) - for rev in repo_info.revisions - for f in rev.files - ) - if not has_weights: - continue - key = repo_id.lower() - existing = seen_lower.get(key) - if existing is None or total_size > existing["size_bytes"]: - seen_lower[key] = { - "repo_id": repo_id, - "size_bytes": total_size, - } + for hf_cache in cache_scans: + for repo_info in hf_cache.repos: + if repo_info.repo_type != "model": + continue + repo_id = repo_info.repo_id + if repo_id.upper().endswith("-GGUF"): + continue + total_size = sum( + f.size_on_disk for rev in repo_info.revisions for f in rev.files + ) + if total_size == 0: + continue + has_weights = any( + f.file_name.endswith(_WEIGHT_EXTENSIONS) + for rev in repo_info.revisions + for f in rev.files + ) + if not has_weights: + continue + key = repo_id.lower() + existing = seen_lower.get(key) + if existing is None or total_size > existing["size_bytes"]: + seen_lower[key] = { + "repo_id": repo_id, + "size_bytes": total_size, + } cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) return {"cached": cached} except Exception as e: diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py index 789052f372..aec6bb1292 100644 --- a/studio/backend/utils/paths/__init__.py +++ b/studio/backend/utils/paths/__init__.py @@ -23,6 +23,8 @@ from .storage_roots import ( unstructured_uploads_root, oxc_validator_tmp_root, tensorboard_root, + legacy_hf_cache_dir, + lmstudio_model_dirs, ensure_dir, ensure_studio_directories, resolve_under_root, @@ -53,6 +55,8 @@ __all__ = [ "unstructured_uploads_root", "oxc_validator_tmp_root", "tensorboard_root", + "legacy_hf_cache_dir", + "lmstudio_model_dirs", "ensure_dir", "ensure_studio_directories", "resolve_under_root", diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 626e868275..08d3744e95 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -3,7 +3,9 @@ from __future__ import annotations +import json import os +import sys from pathlib import Path import tempfile @@ -82,19 +84,55 @@ def ensure_dir(path: Path) -> Path: return path +def legacy_hf_cache_dir() -> Path: + """Old Unsloth-specific HF hub cache, kept for backward-compat scanning.""" + return cache_root() / "huggingface" / "hub" + + +def lmstudio_model_dirs() -> list[Path]: + """Return LM Studio model directories that exist on disk.""" + dirs: list[Path] = [] + + # 1. Check LM Studio settings.json for custom downloads folder + settings_path = Path.home() / ".lmstudio" / "settings.json" + if settings_path.is_file(): + try: + with open(settings_path) as f: + settings = json.load(f) + downloads = settings.get("downloadsFolder", "") + if downloads: + p = Path(downloads).expanduser() + if p.is_dir(): + dirs.append(p) + except Exception: + pass + + # 2. Legacy LM Studio cache (Linux/macOS) + if sys.platform == "win32": + legacy = Path.home() / ".cache" / "lm-studio" / "models" + else: + legacy = Path.home() / ".cache" / "lm-studio" / "models" + if legacy.is_dir(): + dirs.append(legacy) + + return dirs + + def _setup_cache_env() -> None: - """Set cache environment variables for HuggingFace, uv, and vLLM. + """Set cache environment variables for uv and vLLM. + + HuggingFace cache variables (HF_HOME, HF_HUB_CACHE, HF_XET_CACHE) + are no longer overridden — HF uses its own defaults unless the user + has explicitly set them. The legacy Unsloth HF cache at + ``~/.unsloth/studio/cache/huggingface/hub`` is still scanned for + backward compatibility via :func:`legacy_hf_cache_dir`. Only sets variables that are not already set by the user, so - explicit overrides (e.g. HF_HOME=/data/hf) are respected. + explicit overrides are respected. Works on Linux, macOS, and Windows. """ root = cache_root() - hf_dir = root / "huggingface" defaults = { - "HF_HOME": str(hf_dir), - "HF_HUB_CACHE": str(hf_dir / "hub"), - "HF_XET_CACHE": str(hf_dir / "xet"), "UV_CACHE_DIR": str(root / "uv"), "VLLM_CACHE_ROOT": str(root / "vllm"), } diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index e9bcde17b5..84df506769 100644 --- a/studio/frontend/src/features/studio/sections/model-section.tsx +++ b/studio/frontend/src/features/studio/sections/model-section.tsx @@ -334,7 +334,11 @@ export function ModelSection() { {(id: string) => { const model = localMetaById.get(id); const source = - model?.source === "hf_cache" ? "HF cache" : "Local dir"; + model?.source === "hf_cache" + ? "HF cache" + : model?.source === "lmstudio" + ? "LM Studio" + : "Local dir"; return ( diff --git a/studio/frontend/src/features/training/api/models-api.ts b/studio/frontend/src/features/training/api/models-api.ts index 84051e3e1d..f2b1e256c2 100644 --- a/studio/frontend/src/features/training/api/models-api.ts +++ b/studio/frontend/src/features/training/api/models-api.ts @@ -79,7 +79,7 @@ export interface LocalModelInfo { id: string; display_name: string; path: string; - source: "models_dir" | "hf_cache"; + source: "models_dir" | "hf_cache" | "lmstudio"; model_id?: string | null; updated_at?: number | null; } @@ -87,6 +87,7 @@ export interface LocalModelInfo { interface LocalModelListResponse { models_dir: string; hf_cache_dir?: string | null; + lmstudio_dirs?: string[]; models: LocalModelInfo[]; }