Merge pull request #143 from unslothai/feature/local-models

feat: add schemas for local model discovery and listing

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",

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",
+    )

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")

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<string, string> = {
@@ -97,6 +101,10 @@ export function ModelSection() {
   );
 
   const [inputValue, setInputValue] = useState("");
+  const [localModelInput, setLocalModelInput] = useState("");
+  const [localModels, setLocalModels] = useState<LocalModelInfo[]>([]);
+  const [isLoadingLocalModels, setIsLoadingLocalModels] = useState(true);
+  const [localModelsError, setLocalModelsError] = useState<string | null>(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<string, LocalModelInfo>();
+    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<HTMLDivElement>(null);
+  const localComboboxAnchorRef = useRef<HTMLDivElement>(null);
   const { scrollRef, sentinelRef } = useInfiniteScroll(
     fetchMore,
     hfResults.length,
@@ -200,12 +262,89 @@ export function ModelSection() {
               </TooltipContent>
             </Tooltip>
           </span>
-          <InputGroup className="bg-foreground text-background [&_input]:text-background [&_input]:placeholder:text-background/40 [&_svg]:text-background/50 hover:bg-foreground/90">
-            <InputGroupAddon>
-              <HugeiconsIcon icon={FolderSearchIcon} className="size-4" />
-            </InputGroupAddon>
-            <InputGroupInput placeholder="./models/my-model" />
-          </InputGroup>
+          <div ref={localComboboxAnchorRef}>
+            <Combobox
+              items={localResultIds}
+              filteredItems={localFilteredIds}
+              filter={null}
+              value={localModelInput || null}
+              onValueChange={(id) => {
+                const next = id ?? "";
+                setLocalModelInput(next);
+                if (next) setSelectedModel(next);
+              }}
+              onInputValueChange={setLocalModelInput}
+              itemToStringValue={(id) => id}
+              autoHighlight={true}
+            >
+              <ComboboxInput
+                placeholder={
+                  isLoadingLocalModels
+                    ? "Scanning local and cached models..."
+                    : "./models/my-model"
+                }
+                className="w-full bg-foreground text-background [&_input]:text-background [&_input]:placeholder:text-background/40 [&_svg]:text-background/50 hover:bg-foreground/90"
+                onBlur={() => applyLocalModel(localModelInput)}
+                onKeyDown={(event) => {
+                  if (event.key !== "Enter") return;
+                  event.preventDefault();
+                  applyLocalModel(localModelInput);
+                }}
+              >
+                <InputGroupAddon>
+                  <HugeiconsIcon icon={FolderSearchIcon} className="size-4" />
+                </InputGroupAddon>
+              </ComboboxInput>
+              <ComboboxContent anchor={localComboboxAnchorRef}>
+                {isLoadingLocalModels ? (
+                  <div className="flex items-center justify-center gap-2 py-4 text-xs text-muted-foreground">
+                    <Spinner className="size-4" /> Scanning...
+                  </div>
+                ) : localModelsError ? (
+                  <div className="px-3 py-2 text-xs text-red-500">
+                    {localModelsError}
+                  </div>
+                ) : (
+                  <ComboboxEmpty>No local models found</ComboboxEmpty>
+                )}
+                <ComboboxList className="p-1">
+                  {(id: string) => {
+                    const model = localMetaById.get(id);
+                    const source =
+                      model?.source === "hf_cache" ? "HF cache" : "Local dir";
+                    return (
+                      <ComboboxItem key={id} value={id} className="justify-between">
+                        <Tooltip>
+                          <TooltipTrigger asChild={true}>
+                            <span className="min-w-0 flex-1 truncate">
+                              {model?.display_name ?? id}
+                            </span>
+                          </TooltipTrigger>
+                          <TooltipContent side="left" className="max-w-xs break-all">
+                            {model?.path ?? id}
+                          </TooltipContent>
+                        </Tooltip>
+                        <span className="shrink-0 text-[10px] text-muted-foreground">
+                          {source}
+                        </span>
+                      </ComboboxItem>
+                    );
+                  }}
+                </ComboboxList>
+              </ComboboxContent>
+            </Combobox>
+          </div>
+          {isLoadingLocalModels ? (
+            <p className="text-[10px] text-muted-foreground">Scanning local models...</p>
+          ) : localModelsError ? (
+            <p className="text-[10px] text-red-500">{localModelsError}</p>
+          ) : (
+            <p className="text-[10px] text-muted-foreground">
+              {localModels.length > 0
+                ? `${localModels.length} local/cached models found`
+                : "No local models found. Enter path manually."}
+            </p>
+          )}
         </div>
 
           <div data-tour="studio-base-model" className="flex flex-col gap-2">

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<LocalModelInfo[]> {
+  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;
+}

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";