From 972cde7971c52aa0bab6583f11d5876260f6c8af Mon Sep 17 00:00:00 2001 From: Shine1i Date: Tue, 17 Feb 2026 21:53:42 +0100 Subject: [PATCH] feat: add schemas for local model discovery and listing --- studio/backend/models/__init__.py | 4 + studio/backend/models/models.py | 33 +++- studio/backend/routes/models.py | 112 +++++++++++++ .../studio/sections/model-section.tsx | 155 +++++++++++++++++- .../src/features/training/api/models-api.ts | 26 +++ .../frontend/src/features/training/index.ts | 2 + 6 files changed, 323 insertions(+), 9 deletions(-) diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py index b66efe7093..fe21d525a4 100644 --- a/studio/backend/models/__init__.py +++ b/studio/backend/models/__init__.py @@ -12,6 +12,8 @@ from .models import ( ModelCheckpoints, CheckpointListResponse, ModelDetails, + LocalModelInfo, + LocalModelListResponse, LoRAInfo, LoRAScanResponse, ModelListResponse, @@ -59,6 +61,8 @@ __all__ = [ "TrainingProgress", # Model management schemas "ModelDetails", + "LocalModelInfo", + "LocalModelListResponse", "LoRAInfo", "LoRAScanResponse", "ModelListResponse", diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 4e23c69e6d..5542db76d5 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -2,7 +2,7 @@ Pydantic schemas for Model Management API """ from pydantic import BaseModel, Field -from typing import Optional, List, Dict, Any +from typing import Optional, List, Dict, Any, Literal class CheckpointInfo(BaseModel): @@ -74,3 +74,34 @@ class ModelListResponse(BaseModel): models: List[ModelDetails] = Field(default_factory=list, description="List of models") default_models: List[str] = Field(default_factory=list, description="List of default model IDs") + +class LocalModelInfo(BaseModel): + """Discovered local model candidate.""" + 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( + ..., + description="Discovery source", + ) + model_id: Optional[str] = Field( + None, + description="HF repo id for cached models, e.g. org/model", + ) + updated_at: Optional[float] = Field( + None, + description="Unix timestamp of latest observed update", + ) + + +class LocalModelListResponse(BaseModel): + """Response schema for listing local/cached models.""" + models_dir: str = Field(..., description="Directory scanned for custom local models") + hf_cache_dir: Optional[str] = Field( + None, + description="HF cache root that was 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 5ed1de1ada..cc964d3ea3 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -43,6 +43,8 @@ except ImportError: from models import ( CheckpointInfo, CheckpointListResponse, + LocalModelInfo, + LocalModelListResponse, ModelCheckpoints, ModelDetails, LoRAScanResponse, @@ -64,6 +66,116 @@ if not logger.handlers: logger.setLevel(logging.INFO) +def _resolve_hf_cache_dir() -> Path: + """Resolve local HF cache root used by hub downloads.""" + try: + from huggingface_hub.constants import HF_HUB_CACHE + return Path(HF_HUB_CACHE) + except Exception: + return Path.home() / ".cache" / "huggingface" / "hub" + + +def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]: + if not models_dir.exists() or not models_dir.is_dir(): + return [] + + found: List[LocalModelInfo] = [] + for child in models_dir.iterdir(): + if not child.is_dir(): + continue + has_model_files = ( + (child / "config.json").exists() + or (child / "adapter_config.json").exists() + or any(child.glob("*.safetensors")) + or any(child.glob("*.bin")) + ) + if not has_model_files: + continue + try: + updated_at = child.stat().st_mtime + except OSError: + updated_at = None + found.append( + LocalModelInfo( + id=str(child), + display_name=child.name, + path=str(child), + source="models_dir", + updated_at=updated_at, + ), + ) + return found + + +def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]: + if not cache_dir.exists() or not cache_dir.is_dir(): + return [] + + found: List[LocalModelInfo] = [] + for repo_dir in cache_dir.glob("models--*"): + if not repo_dir.is_dir(): + continue + + repo_name = repo_dir.name[len("models--"):] + if not repo_name: + continue + model_id = repo_name.replace("--", "/") + + try: + updated_at = repo_dir.stat().st_mtime + except OSError: + updated_at = None + + found.append( + LocalModelInfo( + id=model_id, + model_id=model_id, + display_name=model_id.split("/")[-1], + path=str(repo_dir), + source="hf_cache", + updated_at=updated_at, + ), + ) + return found + + +@router.get("/local", response_model=LocalModelListResponse) +async def list_local_models( + models_dir: str = Query(default="./models", description="Directory to scan for local model folders"), + current_subject: str = Depends(get_current_subject), +): + """ + List local model candidates from custom models dir and HF cache. + """ + try: + models_root = Path(models_dir).expanduser().resolve() + hf_cache_dir = _resolve_hf_cache_dir() + local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir) + + deduped: dict[str, LocalModelInfo] = {} + for model in local_models: + if model.id not in deduped: + deduped[model.id] = model + + models = sorted( + deduped.values(), + key=lambda item: (item.updated_at or 0), + reverse=True, + ) + + return LocalModelListResponse( + models_dir=str(models_root), + hf_cache_dir=str(hf_cache_dir), + models=models, + ) + except Exception as e: + logger.error(f"Error listing local models: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to list local models: {str(e)}", + ) + + @router.get("/list") diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index 8a017b598e..32810e0700 100644 --- a/studio/frontend/src/features/studio/sections/model-section.tsx +++ b/studio/frontend/src/features/studio/sections/model-section.tsx @@ -39,7 +39,11 @@ import { checkVramFit, estimateLoadingVram, } from "@/lib/vram"; -import { useTrainingConfigStore } from "@/features/training"; +import { + listLocalModels, + type LocalModelInfo, + useTrainingConfigStore, +} from "@/features/training"; import type { TrainingMethod } from "@/types/training"; import { ChipIcon, @@ -49,7 +53,7 @@ import { Search01Icon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useMemo, useRef, useState } from "react"; +import { useEffect, useMemo, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; const METHOD_DOTS: Record = { @@ -97,6 +101,10 @@ export function ModelSection() { ); const [inputValue, setInputValue] = useState(""); + const [localModelInput, setLocalModelInput] = useState(""); + const [localModels, setLocalModels] = useState([]); + const [isLoadingLocalModels, setIsLoadingLocalModels] = useState(true); + const [localModelsError, setLocalModelsError] = useState(null); const selectingRef = useRef(false); const debouncedQuery = useDebouncedValue(inputValue); @@ -112,6 +120,32 @@ export function ModelSection() { } setInputValue(val); } + + function applyLocalModel(value: string) { + const next = value.trim(); + if (!next) return; + setSelectedModel(next); + } + + useEffect(() => { + const controller = new AbortController(); + void listLocalModels(controller.signal) + .then((models) => { + if (controller.signal.aborted) return; + setLocalModels(models); + }) + .catch((error) => { + if (controller.signal.aborted) return; + setLocalModelsError( + error instanceof Error ? error.message : "Failed to load local models", + ); + }) + .finally(() => { + if (controller.signal.aborted) return; + setIsLoadingLocalModels(false); + }); + return () => controller.abort(); + }, []); const task = modelType ? MODEL_TYPE_TO_HF_TASK[modelType] : undefined; const { results: hfResults, @@ -131,6 +165,33 @@ export function ModelSection() { return ids; }, [hfResults, selectedModel]); + const localMetaById = useMemo(() => { + const map = new Map(); + for (const model of localModels) map.set(model.id, model); + return map; + }, [localModels]); + + const localResultIds = useMemo(() => { + const ids = localModels.map((model) => model.id); + const manual = localModelInput.trim(); + if (manual && !ids.includes(manual)) { + ids.unshift(manual); + } + return ids; + }, [localModelInput, localModels]); + + const localFilteredIds = useMemo(() => { + const q = localModelInput.trim().toLowerCase(); + if (!q) return localResultIds; + return localResultIds.filter((id) => { + const meta = localMetaById.get(id); + if (id.toLowerCase().includes(q)) return true; + if (meta?.display_name.toLowerCase().includes(q)) return true; + if (meta?.path.toLowerCase().includes(q)) return true; + return false; + }); + }, [localMetaById, localModelInput, localResultIds]); + // Pre-compute VRAM fit status for every model in the current result set. // Keyed by model id so the render callback is a simple O(1) lookup. // @@ -163,6 +224,7 @@ export function ModelSection() { }, [hfResults, gpu, trainingMethod]); const comboboxAnchorRef = useRef(null); + const localComboboxAnchorRef = useRef(null); const { scrollRef, sentinelRef } = useInfiniteScroll( fetchMore, hfResults.length, @@ -200,12 +262,89 @@ export function ModelSection() { - - - - - - +
+ { + const next = id ?? ""; + setLocalModelInput(next); + if (next) setSelectedModel(next); + }} + onInputValueChange={setLocalModelInput} + itemToStringValue={(id) => id} + autoHighlight={true} + > + applyLocalModel(localModelInput)} + onKeyDown={(event) => { + if (event.key !== "Enter") return; + event.preventDefault(); + applyLocalModel(localModelInput); + }} + > + + + + + + {isLoadingLocalModels ? ( +
+ Scanning... +
+ ) : localModelsError ? ( +
+ {localModelsError} +
+ ) : ( + No local models found + )} + + {(id: string) => { + const model = localMetaById.get(id); + const source = + model?.source === "hf_cache" ? "HF cache" : "Local dir"; + return ( + + + + + {model?.display_name ?? id} + + + + {model?.path ?? id} + + + + {source} + + + ); + }} + + + +
+ {isLoadingLocalModels ? ( +

Scanning local models...

+ ) : localModelsError ? ( +

{localModelsError}

+ ) : ( +

+ {localModels.length > 0 + ? `${localModels.length} local/cached models found` + : "No local models found. Enter path manually."} +

+ )}
diff --git a/studio/frontend/src/features/training/api/models-api.ts b/studio/frontend/src/features/training/api/models-api.ts index 44a703b8ca..6e6e3ff762 100644 --- a/studio/frontend/src/features/training/api/models-api.ts +++ b/studio/frontend/src/features/training/api/models-api.ts @@ -58,6 +58,21 @@ export interface ModelConfigResponse { base_model?: string | null; } +export interface LocalModelInfo { + id: string; + display_name: string; + path: string; + source: "models_dir" | "hf_cache"; + model_id?: string | null; + updated_at?: number | null; +} + +interface LocalModelListResponse { + models_dir: string; + hf_cache_dir?: string | null; + models: LocalModelInfo[]; +} + /** * Check whether a model is a vision model by asking the backend. * Calls GET /api/models/check-vision/{model_name}. @@ -84,3 +99,14 @@ export async function getModelConfig( } return (await response.json()) as ModelConfigResponse; } + +export async function listLocalModels( + signal?: AbortSignal, +): Promise { + const response = await authFetch("/api/models/local", { signal }); + if (!response.ok) { + throw new Error(`Failed to fetch local models (${response.status})`); + } + const data = (await response.json()) as LocalModelListResponse; + return data.models; +} diff --git a/studio/frontend/src/features/training/index.ts b/studio/frontend/src/features/training/index.ts index 77451140f5..24439f0c8a 100644 --- a/studio/frontend/src/features/training/index.ts +++ b/studio/frontend/src/features/training/index.ts @@ -7,4 +7,6 @@ export { useTrainingActions } from "./hooks/use-training-actions"; export { useTrainingRuntimeLifecycle } from "./hooks/use-training-runtime-lifecycle"; export { HfDatasetSubsetSplitSelectors } from "./components/hf-dataset-subset-split-selectors"; export { useDatasetPreviewDialogStore } from "./stores/dataset-preview-dialog-store"; +export { listLocalModels } from "./api/models-api"; +export type { LocalModelInfo } from "./api/models-api"; export type { TrainingPhase } from "./types/runtime";