From 48a78845844c3c8965437db1d7d8ad2bd8713244 Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Wed, 25 Mar 2026 18:48:04 +0400 Subject: [PATCH 1/6] feat: multi-source model discovery (HF default, legacy cache, LM Studio) (#4591) * feat: multi-source model discovery (HF default, legacy cache, LM Studio) * Fix multi-source model discovery bugs - Fix lmstudio_model_dirs: add ~/.lmstudio/models as default path, remove dead sys.platform branch, add dedup via seen set - Fix _setup_cache_env: preserve legacy HF cache env vars when the legacy hub directory exists and is non-empty - Fix _scan_lmstudio_dir: use absolute path for id field so is_local_path() returns True - Remove LM Studio dirs from allowed_roots (scanned unconditionally) - Replace bare except passes with logger.warning in legacy cache blocks - Fix delete_cached_model to search both default and legacy HF caches - Make lmstudio_dirs non-optional in TS interface (matches Python schema) - Exclude lmstudio source from trainable model filter - Remove unused import sys * Scan HF default cache alongside legacy and active caches When _setup_cache_env overrides HF_HUB_CACHE to the legacy Unsloth path, the standard HF default cache (~/.cache/huggingface/hub) was never scanned, hiding models downloaded before Unsloth Studio was installed. Add hf_default_cache_dir() and _all_hf_cache_scans() helper that deduplicates and scans all three HF cache locations (active, legacy, default). Used in list_local_models, list_cached_gguf, list_cached_models, and delete_cached_model. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/models/models.py | 6 +- studio/backend/routes/models.py | 271 +++++++++++++----- studio/backend/utils/paths/__init__.py | 6 + studio/backend/utils/paths/storage_roots.py | 64 ++++- .../studio/sections/model-section.tsx | 7 +- .../src/features/training/api/models-api.ts | 3 +- 6 files changed, 274 insertions(+), 83 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..f76034c95b 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 = str(model_dir), + 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,29 @@ 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, + hf_default_cache_dir, + lmstudio_model_dirs, + ) + + # Resolve all scan directories up front. + hf_cache_dir = _resolve_hf_cache_dir() + legacy_hf = legacy_hf_cache_dir() + hf_default = hf_default_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) + if hf_default.is_dir(): + allowed_roots.append(hf_default) try: from utils.paths import studio_root, outputs_root @@ -248,6 +334,22 @@ 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 HF system default cache (may differ when env vars are overridden) + if ( + hf_default.is_dir() + and hf_default.resolve() != hf_cache_dir.resolve() + and hf_default.resolve() != legacy_hf.resolve() + ): + local_models += _scan_hf_cache(hf_default) + + # 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 +364,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: @@ -846,46 +949,65 @@ def _get_repo_size_cached(repo_id: str) -> int: return 0 +def _all_hf_cache_scans(): + """Return scan_cache_dir results for the active, legacy, and default HF caches.""" + from huggingface_hub import scan_cache_dir + from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir + + scans = [scan_cache_dir()] + seen: set[str] = set() + try: + # Resolve the active cache dir so we can dedup + from huggingface_hub.constants import HF_HUB_CACHE + + seen.add(str(Path(HF_HUB_CACHE).resolve())) + except Exception: + pass + + for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir): + extra = extra_fn() + if extra.is_dir() and str(extra.resolve()) not in seen: + seen.add(str(extra.resolve())) + try: + scans.append(scan_cache_dir(cache_dir = str(extra))) + except Exception as exc: + logger.warning("Could not scan HF cache %s: %s", extra, exc) + return scans + + @router.get("/cached-gguf") 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, legacy Unsloth cache, and HF default cache.""" try: - from huggingface_hub import scan_cache_dir + cache_scans = _all_hf_cache_scans() - 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 +1019,39 @@ 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, legacy Unsloth cache, and HF default cache.""" _WEIGHT_EXTENSIONS = (".safetensors", ".bin") try: - from huggingface_hub import scan_cache_dir + cache_scans = _all_hf_cache_scans() - 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: @@ -989,15 +1106,17 @@ async def delete_cached_model( pass try: - from huggingface_hub import scan_cache_dir + cache_scans = _all_hf_cache_scans() - hf_cache = scan_cache_dir() target_repo = None - for repo_info in hf_cache.repos: - if repo_info.repo_type != "model": - continue - if repo_info.repo_id.lower() == repo_id.lower(): - target_repo = repo_info + for hf_cache in cache_scans: + for repo_info in hf_cache.repos: + if repo_info.repo_type != "model": + continue + if repo_info.repo_id.lower() == repo_id.lower(): + target_repo = repo_info + break + if target_repo is not None: break if target_repo is None: diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py index 789052f372..44a7c8e287 100644 --- a/studio/backend/utils/paths/__init__.py +++ b/studio/backend/utils/paths/__init__.py @@ -23,6 +23,9 @@ from .storage_roots import ( unstructured_uploads_root, oxc_validator_tmp_root, tensorboard_root, + legacy_hf_cache_dir, + hf_default_cache_dir, + lmstudio_model_dirs, ensure_dir, ensure_studio_directories, resolve_under_root, @@ -53,6 +56,9 @@ __all__ = [ "unstructured_uploads_root", "oxc_validator_tmp_root", "tensorboard_root", + "legacy_hf_cache_dir", + "hf_default_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..9bcf3758ad 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -3,6 +3,7 @@ from __future__ import annotations +import json import os from pathlib import Path import tempfile @@ -82,22 +83,77 @@ 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 hf_default_cache_dir() -> Path: + """Return the platform default HuggingFace hub cache (ignoring env overrides). + + This is the location HF uses when no ``HF_HUB_CACHE`` / ``HF_HOME`` + env var is set. We scan it so that models a user downloaded *before* + installing Unsloth Studio are still discovered. + """ + return Path.home() / ".cache" / "huggingface" / "hub" + + +def lmstudio_model_dirs() -> list[Path]: + """Return LM Studio model directories that exist on disk.""" + dirs: list[Path] = [] + seen: set[Path] = set() + + def _add(p: Path) -> None: + resolved = p.resolve() + if resolved not in seen and p.is_dir(): + seen.add(resolved) + dirs.append(p) + + # 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: + _add(Path(downloads).expanduser()) + except Exception: + pass + + # 2. LM Studio current default models directory (all platforms) + _add(Path.home() / ".lmstudio" / "models") + + # 3. Legacy LM Studio cache location + _add(Path.home() / ".cache" / "lm-studio" / "models") + + return dirs + + def _setup_cache_env() -> None: """Set cache environment variables for HuggingFace, uv, and vLLM. + HuggingFace cache variables are only set when the legacy Unsloth HF + cache already exists, preserving existing model locations. New + installations leave HF at its own defaults. + Only sets variables that are not already set by the user, so explicit overrides (e.g. HF_HOME=/data/hf) 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"), + defaults: dict[str, str] = { "UV_CACHE_DIR": str(root / "uv"), "VLLM_CACHE_ROOT": str(root / "vllm"), } + # Preserve legacy HF cache for existing installations + legacy_hub = hf_dir / "hub" + if legacy_hub.is_dir() and any(legacy_hub.iterdir()): + defaults["HF_HOME"] = str(hf_dir) + defaults["HF_HUB_CACHE"] = str(legacy_hub) + defaults["HF_XET_CACHE"] = str(hf_dir / "xet") + for key, value in defaults.items(): if key not in os.environ: os.environ[key] = value diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index e9bcde17b5..9dccae3bc3 100644 --- a/studio/frontend/src/features/studio/sections/model-section.tsx +++ b/studio/frontend/src/features/studio/sections/model-section.tsx @@ -181,6 +181,7 @@ export function ModelSection() { const trainableLocalModels = useMemo( () => localModels.filter((m) => { + if (m.source === "lmstudio") return false; if (m.path.endsWith(".gguf")) return false; if (m.id.toLowerCase().includes("-gguf")) return false; return true; @@ -334,7 +335,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..2a9ad7c0d6 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[]; } From 2683c2ab583082aa16f3f23539ce7b569e0de901 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Mar 2026 08:00:44 -0700 Subject: [PATCH 2/6] Add unsloth to User PATH on Windows after install (#4597) After installation, `unsloth studio` only works if the user activates the Studio venv first or uses the full absolute path. The Desktop/Start Menu shortcuts work fine, but typing `unsloth studio` in a fresh terminal does not. This adds the venv Scripts dir to the persistent User PATH env var (if not already present) so `unsloth studio` works from any new terminal window. The current session is also updated via the existing Refresh-SessionPath helper. --- install.ps1 | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/install.ps1 b/install.ps1 index a4ed2658c9..ed656850d4 100644 --- a/install.ps1 +++ b/install.ps1 @@ -623,6 +623,19 @@ shell.Run cmd, 0, False New-StudioShortcuts -UnslothExePath $UnslothExe + # ── Add venv Scripts dir to User PATH so `unsloth studio` works from any terminal ── + $ScriptsDir = Join-Path $VenvDir "Scripts" + $UserPath = [System.Environment]::GetEnvironmentVariable("Path", "User") + if (-not $UserPath -or $UserPath -notlike "*$ScriptsDir*") { + if ($UserPath) { + [System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir;$UserPath", "User") + } else { + [System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir", "User") + } + Refresh-SessionPath + Write-Host "[OK] Added unsloth to PATH" -ForegroundColor Green + } + Write-Host "" Write-Host "=========================================" Write-Host " Unsloth Studio installed!" From 289c7dd7bb20d95d96bc8dd237cf9d3f12293beb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Mar 2026 15:12:56 +0000 Subject: [PATCH 3/6] Add --local and --package flags to install.ps1 Windows install.ps1 had no way to install from a local repo checkout, unlike install.sh which supports ./install.sh --local. This adds: - --local: install from the local repo via editable install (-e . --no-deps) after installing deps from PyPI, mirroring install.sh behavior - --package: install a different package name for testing The --local flag: 1. Validates pyproject.toml exists at the script's directory 2. Installs torch + unsloth deps normally 3. Overlays the local checkout with uv pip install -e --no-deps 4. Passes STUDIO_LOCAL_INSTALL and STUDIO_LOCAL_REPO to setup.ps1 --- install.ps1 | 54 ++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 51 insertions(+), 3 deletions(-) diff --git a/install.ps1 b/install.ps1 index ed656850d4..9405072bd7 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1,10 +1,37 @@ # Unsloth Studio Installer for Windows PowerShell # Usage: irm https://raw.githubusercontent.com/unslothai/unsloth/main/install.ps1 | iex -# Local: Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1 +# Local: Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1 --local +# Test: .\install.ps1 --package roland-sloth function Install-UnslothStudio { $ErrorActionPreference = "Stop" + # ── Parse flags ── + $StudioLocalInstall = $false + $PackageName = "unsloth" + $RepoRoot = "" + $argList = $args + for ($i = 0; $i -lt $argList.Count; $i++) { + switch ($argList[$i]) { + "--local" { $StudioLocalInstall = $true } + "--package" { + $i++ + if ($i -ge $argList.Count) { + Write-Host "[ERROR] --package requires an argument." -ForegroundColor Red + return + } + $PackageName = $argList[$i] + } + } + } + if ($StudioLocalInstall) { + $RepoRoot = (Resolve-Path (Split-Path -Parent $PSCommandPath)).Path + if (-not (Test-Path (Join-Path $RepoRoot "pyproject.toml"))) { + Write-Host "[ERROR] --local must be run from the unsloth repo root (pyproject.toml not found at $RepoRoot)" -ForegroundColor Red + return + } + } + $PythonVersion = "3.13" $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" $VenvDir = Join-Path $StudioHome "unsloth_studio" @@ -581,6 +608,10 @@ shell.Run cmd, 0, False # in the new venv location, while preserving existing torch/CUDA Write-Host "==> Upgrading unsloth in migrated environment..." uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.11" unsloth-zoo + if ($StudioLocalInstall) { + Write-Host "==> Overlaying local repo (editable)..." + uv pip install --python $VenvPython -e $RepoRoot --no-deps + } } elseif ($TorchIndexUrl) { Write-Host "==> Installing PyTorch ($TorchIndexUrl)..." uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl @@ -590,11 +621,23 @@ shell.Run cmd, 0, False } Write-Host "==> Installing unsloth (this may take a few minutes)..." - uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.11" + if ($StudioLocalInstall) { + uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.11" unsloth-zoo + Write-Host "==> Overlaying local repo (editable)..." + uv pip install --python $VenvPython -e $RepoRoot --no-deps + } else { + uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName" + } } else { # Fallback: GPU detection failed to produce a URL -- let uv resolve torch Write-Host "==> Installing unsloth (this may take a few minutes)..." - uv pip install --python $VenvPython "unsloth>=2026.3.11" --torch-backend=auto + if ($StudioLocalInstall) { + uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.11" --torch-backend=auto + Write-Host "==> Overlaying local repo (editable)..." + uv pip install --python $VenvPython -e $RepoRoot --no-deps + } else { + uv pip install --python $VenvPython "$PackageName" --torch-backend=auto + } } if ($LASTEXITCODE -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $LASTEXITCODE)" -ForegroundColor Red @@ -615,6 +658,11 @@ shell.Run cmd, 0, False } # Tell setup.ps1 to skip base package installation (install.ps1 already did it) $env:SKIP_STUDIO_BASE = "1" + $env:STUDIO_PACKAGE_NAME = $PackageName + if ($StudioLocalInstall) { + $env:STUDIO_LOCAL_INSTALL = "1" + $env:STUDIO_LOCAL_REPO = $RepoRoot + } & $UnslothExe studio setup if ($LASTEXITCODE -ne 0) { Write-Host "[ERROR] unsloth studio setup failed (exit code $LASTEXITCODE)" -ForegroundColor Red From 561f0f39be429e93340a4b72ead4488fd81bcdf1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Mar 2026 15:14:51 +0000 Subject: [PATCH 4/6] Fix install.ps1 --local: pass script args to Install-UnslothStudio The function was called with no arguments, so $args inside the function was always empty. Script-level args (--local, --package) were never forwarded. Use @args splatting to pass them through. --- install.ps1 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/install.ps1 b/install.ps1 index 9405072bd7..bb0acf1237 100644 --- a/install.ps1 +++ b/install.ps1 @@ -706,4 +706,4 @@ shell.Run cmd, 0, False } } -Install-UnslothStudio +Install-UnslothStudio @args From 6d6008a1ef78af2cf4d73be711f1f0fcb8dd6f9d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Mar 2026 08:27:27 -0700 Subject: [PATCH 5/6] Add PID file tracking and `unsloth studio stop` command (#4598) * Add PID file tracking and `unsloth studio stop` command On macOS the .app shortcut launches Studio via osascript into a Terminal window, then the launcher script exits. The server process runs outside of the launcher's context with no PID file, so there is no straightforward way to find or stop it. This adds: - PID file at ~/.unsloth/studio/studio.pid, written after the server starts and removed on graceful shutdown or via atexit - `unsloth studio stop` command that reads the PID file and sends SIGTERM (or taskkill on Windows) to shut down the server The PID file is only removed if it still contains the current process ID, avoiding races when a new server instance replaces a crashed one. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Move atexit PID cleanup into run_server() The atexit registration was only in the __main__ block, so it did not cover the `unsloth studio` CLI path that calls run_server() directly via studio_default(). Moving it into run_server() ensures the PID file is cleaned up on unexpected exit regardless of entry point. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/run.py | 29 +++++++++++++++ unsloth_cli/commands/studio.py | 68 ++++++++++++++++++++++++++++++++++ 2 files changed, 97 insertions(+) diff --git a/studio/backend/run.py b/studio/backend/run.py index e32b912c37..b892037565 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -158,6 +158,29 @@ def _find_free_port(host: str, start: int, max_attempts: int = 20) -> int: ) +_PID_FILE = Path.home() / ".unsloth" / "studio" / "studio.pid" + + +def _write_pid_file(): + """Write the current process PID to the studio PID file.""" + try: + _PID_FILE.parent.mkdir(parents = True, exist_ok = True) + _PID_FILE.write_text(str(os.getpid())) + except OSError: + pass + + +def _remove_pid_file(): + """Remove the PID file if it belongs to this process.""" + try: + if _PID_FILE.is_file(): + stored = _PID_FILE.read_text().strip() + if stored == str(os.getpid()): + _PID_FILE.unlink(missing_ok = True) + except OSError: + pass + + def _graceful_shutdown(server = None): """Explicitly shut down all subprocess backends and the uvicorn server. @@ -165,6 +188,7 @@ def _graceful_shutdown(server = None): before the parent exits. This is critical on Windows where atexit handlers are unreliable after Ctrl+C. """ + _remove_pid_file() logger.info("Graceful shutdown initiated — cleaning up subprocesses...") # 1. Shut down uvicorn server (releases the listening socket) @@ -307,6 +331,11 @@ def run_server( thread.start() time.sleep(3) + _write_pid_file() + import atexit + + atexit.register(_remove_pid_file) + if not silent: display_host = _resolve_external_ip() if host == "0.0.0.0" else host diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index c6d398eebd..a2f0873e22 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -166,6 +166,74 @@ def studio_default( typer.echo("\nShutting down...") +# ── unsloth studio stop ─────────────────────────────────────────────── + +_PID_FILE = STUDIO_HOME / "studio.pid" + + +@studio_app.command() +def stop(): + """Stop a running Unsloth Studio server. + + Reads the PID from ~/.unsloth/studio/studio.pid and sends SIGTERM + (or TerminateProcess on Windows) to shut it down gracefully. + """ + import signal as _signal + + if not _PID_FILE.is_file(): + typer.echo("No running Studio server found (no PID file).") + raise typer.Exit(0) + + pid_text = _PID_FILE.read_text().strip() + if not pid_text.isdigit(): + typer.echo(f"Invalid PID file contents: {pid_text}") + _PID_FILE.unlink(missing_ok = True) + raise typer.Exit(1) + + pid = int(pid_text) + + # Check if the process is still alive + try: + os.kill(pid, 0) + except ProcessLookupError: + typer.echo( + f"Studio server (PID {pid}) is not running. Cleaning up stale PID file." + ) + _PID_FILE.unlink(missing_ok = True) + raise typer.Exit(0) + except PermissionError: + pass # process exists but we may not own it; try to signal anyway + + # Send SIGTERM (graceful shutdown) or TerminateProcess on Windows + try: + if sys.platform == "win32": + subprocess.run(["taskkill", "/PID", str(pid), "/F"], check = True) + else: + os.kill(pid, _signal.SIGTERM) + typer.echo(f"Sent shutdown signal to Studio server (PID {pid}).") + except ProcessLookupError: + typer.echo(f"Studio server (PID {pid}) already exited.") + _PID_FILE.unlink(missing_ok = True) + raise typer.Exit(0) + except Exception as e: + typer.echo(f"Failed to stop Studio server (PID {pid}): {e}", err = True) + raise typer.Exit(1) + + # Wait briefly for the process to exit and clean up + for _ in range(10): + time.sleep(0.5) + try: + os.kill(pid, 0) + except ProcessLookupError: + _PID_FILE.unlink(missing_ok = True) + typer.echo("Studio server stopped.") + raise typer.Exit(0) + except PermissionError: + break + + typer.echo("Studio server is shutting down (may take a few seconds).") + + # ── unsloth studio setup / update ───────────────────────────────────── From 55d24d7c490addd9c6a8a1f0bd860de40aa0b890 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Mar 2026 08:32:38 -0700 Subject: [PATCH 6/6] feat(studio): editable context length with Apply/Reset for GGUF settings (#4592) * feat(studio): editable context length with Apply/Reset for GGUF model settings Previously the Context Length field was read-only and the backend hardcoded `-c 0`, ignoring custom values entirely. KV Cache Dtype also triggered an immediate model reload with no way to cancel. Backend: - llama_cpp.py: pass the actual n_ctx value to `-c` instead of always 0 - models/inference.py: relax max_seq_length to 0..1048576 (0 = model default) so GGUF models with large context windows are supported Frontend: - chat-runtime-store: add customContextLength and loadedKvCacheDtype state fields for dirty tracking - chat-settings-sheet: make Context Length an editable number input, stop KV Cache Dtype from auto-reloading, show Apply/Reset buttons when either setting has been changed - use-chat-model-runtime: send customContextLength as max_seq_length in the load request, reset after successful load * fix: preserve maxSeqLength for non-GGUF models in load request customContextLength ?? 0 sent max_seq_length=0 for non-GGUF models, breaking the finetuning/inference path that needs the slider value. Now uses a three-way branch: - customContextLength set: use it (user edited GGUF context) - GGUF without custom: 0 (model's native context) - Non-GGUF: maxSeqLength from the sampling slider * fix: keep max_seq_length default at 4096 for non-GGUF callers Only relax the bounds (ge=0 for GGUF's "model default" mode, le=1048576 for large context windows). The default stays at 4096 so API callers that omit max_seq_length still get a sane value for non-GGUF models. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio): rename trust remote code toggle and hide when no model selected - Rename "Trust remote code" to "Enable custom code" - Shorten subtitle to "Only enable if sure" - Hide the toggle when no model is loaded (already hidden for GGUFs) * fix: restore ge=128 for max_seq_length validation Keep the minimum at 128 so the API rejects nonsensical values. GGUF path now sends the model's native context length (from ggufContextLength) instead of 0 when the user has not customized it. The upper bound stays at 1048576 for large-context GGUF models. * feat(studio): replace Context Length input with slider Use a ParamSlider (512 to model's native context, step 512) instead of a small number input. Shows "Max" when at the model's native context length. Consistent with the other slider controls in the settings panel. * feat(studio): add editable number input alongside Context Length slider The slider and number input stay synced -- dragging the slider updates the number, typing a number moves the slider. The input also accepts values beyond the slider range for power users who need custom context lengths larger than the model default. * fix(studio): widen context length input and use 1024 step for slider Make the number input wider (100px) so large values like 262144 are fully visible. Change slider step from 512 to 1024 and min from 512 to 1024. * fix(studio): context length number input increments by 1024 * fix(studio): cap context length input at model's native max Adds max attribute and clamps typed/incremented values so the context length cannot exceed the GGUF model's reported context window. * fix(studio): point "What's new" link to changelog page Changed from /blog to /docs/new/changelog. * fix(studio): preserve custom context length after Apply, remove stale subtitle - After a reload with a custom context length, keep the user's value in the UI instead of snapping back to the model's native max. ggufContextLength always reports the model's native metadata value regardless of what -c was passed, so we need to preserve customContextLength when it differs from native. - Remove "Reload to apply." from KV Cache Dtype subtitle since the Apply/Reset buttons now handle this. * feat(studio): auto-enable Search and Code tools when model supports them Previously toolsEnabled and codeToolsEnabled stayed false after loading a model even if it reported supports_tools=true. Now both toggles are automatically enabled when the loaded model supports tool calling, matching the existing behavior for reasoning. * fix(studio): auto-enable tools in autoLoadSmallestModel path The suggestion cards trigger autoLoadSmallestModel which bypasses selectModel entirely. It was hardcoding toolsEnabled: false and codeToolsEnabled: false even when the model supports tool calling. Now both are set from the load response, matching the selectModel behavior. Also sets kvCacheDtype/loadedKvCacheDtype for dirty tracking consistency. * fix(studio): re-read tool flags after auto-loading model The runtime state was captured once at the start of the chat adapter's run(), before autoLoadSmallestModel() executes. After auto-load enables tools in the store, the request was still built with the stale snapshot that had toolsEnabled=false. Now re-reads the store after auto-load so the first message includes tools. * fix(studio): re-read entire runtime state after auto-load, not just tools The runtime snapshot (including params.checkpoint, model id, and all tool/reasoning flags) was captured once before auto-load. After autoLoadSmallestModel sets the checkpoint and enables tools, the request was still built with stale params (empty checkpoint, tools disabled). Now re-reads the full store state after auto-load so the first message has the correct model, tools, and reasoning flags. * feat(studio): add Hugging Face token field in Preferences Adds a password input under Configuration > Preferences for users to enter their HF token. The token is persisted in localStorage and passed to all model validate/load/download calls, replacing the previously hardcoded null. This enables downloading gated and private models. * fix(studio): use model native context for GGUF auto-load, show friendly errors The auto-load paths and selectModel for GGUF were sending max_seq_length=4096 which now actually limits the context window (since we fixed the backend to respect n_ctx). Changed to send 0 for GGUF, which means "use model's native context size". Also replaced generic "An internal error occurred" messages with user-friendly descriptions for known errors like context size exceeded and lost connections. LoadRequest validation changed to ge=0 to allow the GGUF "model default" signal. The frontend slider still enforces min=128 for non-GGUF models. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio): filter out FP8 models from model search results Hide models matching *-FP8-* or *FP8-Dynamic* from both the recommended list and HF search results. These models are not yet supported in the inference UI. --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 2 +- studio/backend/models/inference.py | 5 +- studio/backend/routes/inference.py | 31 ++++- .../assistant-ui/model-selector/pickers.tsx | 6 +- .../src/features/chat/api/chat-adapter.ts | 28 +++-- .../src/features/chat/chat-settings-sheet.tsx | 106 +++++++++++++++--- .../chat/hooks/use-chat-model-runtime.ts | 38 +++++-- .../src/features/chat/shared-composer.tsx | 2 +- .../chat/stores/chat-runtime-store.ts | 35 ++++++ .../src/features/chat/thread-sidebar.tsx | 2 +- 10 files changed, 208 insertions(+), 47 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 1d5643ac09..3f89cd3e5d 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -848,7 +848,7 @@ class LlamaCppBackend: "--port", str(self._port), "-c", - "0", # 0 = use model's native context size + str(n_ctx) if n_ctx > 0 else "0", # 0 = model's native context size "--parallel", "1", # Single-user studio, saves VRAM "--flash-attn", diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b0498319ca..b4e496b051 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -22,7 +22,10 @@ class LoadRequest(BaseModel): None, description = "HuggingFace token for gated models" ) max_seq_length: int = Field( - 4096, ge = 128, le = 32768, description = "Maximum sequence length" + 0, + ge = 0, + le = 1048576, + description = "Maximum sequence length (0 = model default for GGUF)", ) load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization") is_lora: bool = Field(False, description = "Whether this is a LoRA adapter") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index aa8c34a3c5..78d95fedbd 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -19,6 +19,27 @@ import asyncio import threading +import re as _re + + +def _friendly_error(exc: Exception) -> str: + """Extract a user-friendly message from known llama-server errors.""" + msg = str(exc) + m = _re.search( + r"request \((\d+) tokens?\) exceeds the available context size \((\d+) tokens?\)", + msg, + ) + if m: + return ( + f"Message too long: {m.group(1)} tokens exceeds the {m.group(2)}-token " + f"context window. Try increasing the Context Length in Model settings, " + f"or shorten the conversation." + ) + if "Lost connection to llama-server" in msg: + return "Lost connection to the model server. It may have crashed -- try reloading the model." + return "An internal error occurred" + + # Add backend directory to path backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: @@ -550,7 +571,7 @@ async def generate_stream( except Exception as e: backend.reset_generation_state() logger.error(f"Error during generation: {e}", exc_info = True) - yield f"data: {json.dumps({'error': 'An internal error occurred'})}\n\n" + yield f"data: {json.dumps({'error': _friendly_error(e)})}\n\n" return StreamingResponse( stream(), @@ -944,7 +965,7 @@ async def openai_chat_completions( logger.error( f"Error during audio input streaming: {e}", exc_info = True ) - yield f"data: {json.dumps({'error': {'message': 'An internal error occurred', 'type': 'server_error'}})}\n\n" + yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" return StreamingResponse( audio_input_stream(), @@ -1176,7 +1197,7 @@ async def openai_chat_completions( logger.error(f"Error during GGUF tool streaming: {e}\n{tb}") error_chunk = { "error": { - "message": "An internal error occurred", + "message": _friendly_error(e), "type": "server_error", }, } @@ -1314,7 +1335,7 @@ async def openai_chat_completions( logger.error(f"Error during GGUF streaming: {e}", exc_info = True) error_chunk = { "error": { - "message": "An internal error occurred", + "message": _friendly_error(e), "type": "server_error", }, } @@ -1495,7 +1516,7 @@ async def openai_chat_completions( logger.error(f"Error during OpenAI streaming: {e}", exc_info = True) error_chunk = { "error": { - "message": "An internal error occurred", + "message": _friendly_error(e), "type": "server_error", }, } diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 328cba3acd..3ca9aadb1f 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -454,7 +454,8 @@ export function HubModelPicker({ const recommendedIds = useMemo(() => { const all = dedupe([...models.map((model) => model.id), value ?? ""]) .filter((id) => !downloadedSet.has(id.toLowerCase())) - .filter((id) => !chatOnly || isGgufRepo(id)); + .filter((id) => !chatOnly || isGgufRepo(id)) + .filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id)); // Sort: GGUFs first, then hub models const gguf: string[] = []; const hub: string[] = []; @@ -498,7 +499,8 @@ export function HubModelPicker({ return results .map((result) => result.id) .filter((id) => !recommendedSet.has(id)) - .filter((id) => !chatOnly || isGgufRepo(id)); + .filter((id) => !chatOnly || isGgufRepo(id)) + .filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id)); }, [recommendedSet, results, showHfSection, chatOnly]); const metricsById = useMemo( diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 15ac416b1f..95af560305 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -253,6 +253,7 @@ function waitForModelReady(abortSignal?: AbortSignal): Promise { * falls back to smallest cached safetensors model. */ async function autoLoadSmallestModel(): Promise { + const hfToken = useChatRuntimeStore.getState().hfToken || null; const toastId = toast("Loading a model…", { description: "Auto-selecting the smallest downloaded model.", duration: 5000, @@ -278,8 +279,8 @@ async function autoLoadSmallestModel(): Promise { const variant = downloaded[0]; const loadResp = await loadModel({ model_path: repo.repo_id, - hf_token: null, - max_seq_length: 4096, + hf_token: hfToken, + max_seq_length: 0, load_in_4bit: true, is_lora: false, gguf_variant: variant.quant, @@ -308,8 +309,10 @@ async function autoLoadSmallestModel(): Promise { supportsReasoning: loadResp.supports_reasoning ?? false, reasoningEnabled: loadResp.supports_reasoning ?? false, supportsTools: loadResp.supports_tools ?? false, - toolsEnabled: false, - codeToolsEnabled: false, + toolsEnabled: loadResp.supports_tools ?? false, + codeToolsEnabled: loadResp.supports_tools ?? false, + kvCacheDtype: loadResp.cache_type_kv ?? null, + loadedKvCacheDtype: loadResp.cache_type_kv ?? null, defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, }); @@ -329,7 +332,7 @@ async function autoLoadSmallestModel(): Promise { try { const sfLoadResp = await loadModel({ model_path: repo.repo_id, - hf_token: null, + hf_token: hfToken, max_seq_length: 4096, load_in_4bit: true, is_lora: false, @@ -366,8 +369,8 @@ async function autoLoadSmallestModel(): Promise { try { const loadResp = await loadModel({ model_path: "unsloth/Qwen3.5-4B-GGUF", - hf_token: null, - max_seq_length: 4096, + hf_token: hfToken, + max_seq_length: 0, load_in_4bit: true, is_lora: false, gguf_variant: "UD-Q4_K_XL", @@ -391,7 +394,10 @@ async function autoLoadSmallestModel(): Promise { supportsReasoning: loadResp.supports_reasoning ?? false, reasoningEnabled: loadResp.supports_reasoning ?? false, supportsTools: loadResp.supports_tools ?? false, - toolsEnabled: false, + toolsEnabled: loadResp.supports_tools ?? false, + codeToolsEnabled: loadResp.supports_tools ?? false, + kvCacheDtype: loadResp.cache_type_kv ?? null, + loadedKvCacheDtype: loadResp.cache_type_kv ?? null, defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, }); @@ -410,8 +416,7 @@ async function autoLoadSmallestModel(): Promise { export function createOpenAIStreamAdapter(): ChatModelAdapter { return { async *run({ messages, abortSignal, unstable_threadId }) { - const runtime = useChatRuntimeStore.getState(); - const { params } = runtime; + let runtime = useChatRuntimeStore.getState(); // Wait for in-progress model load to finish before inferring if (runtime.modelLoading) { @@ -430,6 +435,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } } + // Re-read store after potential auto-load / model ready wait + runtime = useChatRuntimeStore.getState(); + const { params } = runtime; const { supportsTools, toolsEnabled, diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 081e3efc26..a315aeb005 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -279,6 +279,14 @@ export function ChatSettingsPanel({ const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); const kvCacheDtype = useChatRuntimeStore((s) => s.kvCacheDtype); const setKvCacheDtype = useChatRuntimeStore((s) => s.setKvCacheDtype); + const loadedKvCacheDtype = useChatRuntimeStore((s) => s.loadedKvCacheDtype); + const customContextLength = useChatRuntimeStore((s) => s.customContextLength); + const setCustomContextLength = useChatRuntimeStore((s) => s.setCustomContextLength); + + const ctxDisplayValue = customContextLength ?? ggufContextLength ?? ""; + const kvDirty = kvCacheDtype !== loadedKvCacheDtype; + const ctxDirty = customContextLength !== null; + const modelSettingsDirty = kvDirty || ctxDirty; const [customPresets, setCustomPresets] = useState(() => loadSavedCustomPresets(), ); @@ -467,32 +475,53 @@ export function ChatSettingsPanel({
{isGguf && ( <> -
-
-
Context Length
-
- Reported by the loaded GGUF model. -
+
+
+ Context Length + { + const raw = e.target.value; + if (raw === "") { + setCustomContextLength(null); + return; + } + const v = parseInt(raw, 10); + if (!Number.isNaN(v) && v >= 0) { + const maxCtx = ggufContextLength ?? Infinity; + const clamped = Math.min(v, maxCtx); + setCustomContextLength(clamped === (ggufContextLength ?? 0) ? null : clamped); + } + }} + />
- { + setCustomContextLength(v === (ggufContextLength ?? 0) ? null : v); + }} />
KV Cache Dtype
- Quantize KV cache to reduce VRAM. Reload to apply. + Quantize KV cache to reduce VRAM.
+ {modelSettingsDirty && ( +
+ + +
+ )} )} - {!isGguf && ( + {!isGguf && params.checkpoint && (
-
Trust remote code
+
Enable custom code
- Allow models with custom code (e.g. Nemotron). Only enable for repos you trust. + Allow models with custom code (e.g. Nemotron). Only enable if sure.
+
@@ -775,6 +826,29 @@ function AutoHealToolCallsToggle() { ); } +function HfTokenField() { + const hfToken = useChatRuntimeStore((s) => s.hfToken); + const setHfToken = useChatRuntimeStore((s) => s.setHfToken); + + return ( +
+
+
Hugging Face Token
+
+ For downloading gated or private models. +
+
+ setHfToken(e.target.value)} + /> +
+ ); +} + function ChatTemplateSection({ onReloadModel, }: { diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index cfdd4e774a..25c776948f 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -354,12 +354,13 @@ export function useChatModelRuntime() { useChatRuntimeStore.getState().params.checkpoint; const paramsBeforeLoad = useChatRuntimeStore.getState().params; const maxSeqLength = paramsBeforeLoad.maxSeqLength; + const hfToken = useChatRuntimeStore.getState().hfToken || null; try { // Lightweight pre-flight validation: avoid unloading a working model // if the new identifier is clearly invalid (e.g. bad HF id / path). await validateModel({ model_path: modelId, - hf_token: null, + hf_token: hfToken, max_seq_length: maxSeqLength, load_in_4bit: true, is_lora: isLora, @@ -371,11 +372,16 @@ export function useChatModelRuntime() { previousWasUnloaded = true; } - const { chatTemplateOverride, kvCacheDtype } = useChatRuntimeStore.getState(); + const { chatTemplateOverride, kvCacheDtype, customContextLength, ggufContextLength } = useChatRuntimeStore.getState(); + // GGUF: use custom context length, or 0 = model's native context + // Non-GGUF: use the Max Seq Length slider value + const effectiveMaxSeqLength = customContextLength != null + ? customContextLength + : ggufVariant != null ? (ggufContextLength ?? 0) : maxSeqLength; const loadResponse = await loadModel({ model_path: modelId, - hf_token: null, - max_seq_length: maxSeqLength, + hf_token: hfToken, + max_seq_length: effectiveMaxSeqLength, load_in_4bit: true, is_lora: isLora, gguf_variant: ggufVariant ?? null, @@ -403,15 +409,27 @@ export function useChatModelRuntime() { } } } + const loadedKv = loadResponse.cache_type_kv ?? null; + const nativeCtx = loadResponse.is_gguf + ? (loadResponse.context_length ?? 131072) + : null; + // Keep customContextLength if the user set one and it differs + // from the model's native context; otherwise clear it so the + // display shows the native value without a dirty marker. + const keepCustomCtx = customContextLength != null + && customContextLength !== nativeCtx + ? customContextLength + : null; useChatRuntimeStore.setState({ - ggufContextLength: loadResponse.is_gguf - ? (loadResponse.context_length ?? 131072) - : null, + ggufContextLength: nativeCtx, supportsReasoning: loadResponse.supports_reasoning ?? false, reasoningEnabled: reasoningDefault, supportsTools: loadResponse.supports_tools ?? false, - toolsEnabled: false, - kvCacheDtype: loadResponse.cache_type_kv ?? null, + toolsEnabled: loadResponse.supports_tools ?? false, + codeToolsEnabled: loadResponse.supports_tools ?? false, + kvCacheDtype: loadedKv, + loadedKvCacheDtype: loadedKv, + customContextLength: keepCustomCtx, defaultChatTemplate: loadResponse.chat_template ?? null, chatTemplateOverride: null, }); @@ -432,7 +450,7 @@ export function useChatModelRuntime() { try { await loadModel({ model_path: previousCheckpoint, - hf_token: null, + hf_token: hfToken, max_seq_length: maxSeqLength, load_in_4bit: true, is_lora: previousIsLora, diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 78cfcc66d2..5ac8c79160 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -337,7 +337,7 @@ export function SharedComposer({ async function ensureModelLoaded(sel: CompareModelSelection): Promise { const resp = await loadModel({ model_path: sel.id, - hf_token: null, + hf_token: useChatRuntimeStore.getState().hfToken || null, max_seq_length: maxSeqLength, load_in_4bit: true, is_lora: sel.isLora, diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 920737a279..2d60d52043 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -14,6 +14,7 @@ const AUTO_TITLE_KEY = "unsloth_chat_auto_title"; const AUTO_HEAL_TOOL_CALLS_KEY = "unsloth_auto_heal_tool_calls"; const MAX_TOOL_CALLS_KEY = "unsloth_max_tool_calls_per_message"; const TOOL_CALL_TIMEOUT_KEY = "unsloth_tool_call_timeout"; +const HF_TOKEN_KEY = "unsloth_hf_token"; const INFERENCE_PARAMS_KEY = "unsloth_chat_inference_params"; let hasShownInferencePersistenceWarning = false; @@ -62,6 +63,24 @@ function saveInt(key: string, value: number): void { } } +function loadString(key: string, fallback: string): string { + if (!canUseStorage()) return fallback; + try { + return localStorage.getItem(key) ?? fallback; + } catch { + return fallback; + } +} + +function saveString(key: string, value: string): void { + if (!canUseStorage()) return; + try { + localStorage.setItem(key, value); + } catch { + // ignore + } +} + function asFiniteNumber(value: unknown, fallback: number): number { return typeof value === "number" && Number.isFinite(value) ? value : fallback; } @@ -127,6 +146,7 @@ type ChatRuntimeStore = { loras: ChatLoraSummary[]; runningByThreadId: Record; autoTitle: boolean; + hfToken: string; modelsError: string | null; activeGgufVariant: string | null; ggufContextLength: number | null; @@ -141,6 +161,8 @@ type ChatRuntimeStore = { maxToolCallsPerMessage: number; toolCallTimeout: number; kvCacheDtype: string | null; + loadedKvCacheDtype: string | null; + customContextLength: number | null; defaultChatTemplate: string | null; chatTemplateOverride: string | null; activeThreadId: string | null; @@ -159,6 +181,7 @@ type ChatRuntimeStore = { setLoras: (loras: ChatLoraSummary[]) => void; setThreadRunning: (threadId: string, running: boolean) => void; setAutoTitle: (enabled: boolean) => void; + setHfToken: (token: string) => void; setModelsError: (error: string | null) => void; setCheckpoint: (modelId: string, ggufVariant?: string | null) => void; setActiveThreadId: (threadId: string | null) => void; @@ -172,6 +195,7 @@ type ChatRuntimeStore = { setMaxToolCallsPerMessage: (value: number) => void; setToolCallTimeout: (value: number) => void; setKvCacheDtype: (dtype: string | null) => void; + setCustomContextLength: (v: number | null) => void; setChatTemplateOverride: (template: string | null) => void; setPendingAudio: (base64: string, name: string) => void; clearPendingAudio: () => void; @@ -184,6 +208,7 @@ export const useChatRuntimeStore = create((set) => ({ loras: [], runningByThreadId: {}, autoTitle: loadBool(AUTO_TITLE_KEY, false), + hfToken: loadString(HF_TOKEN_KEY, ""), modelsError: null, activeGgufVariant: null, ggufContextLength: null, @@ -198,6 +223,8 @@ export const useChatRuntimeStore = create((set) => ({ maxToolCallsPerMessage: loadInt(MAX_TOOL_CALLS_KEY, 10), toolCallTimeout: loadInt(TOOL_CALL_TIMEOUT_KEY, 5), kvCacheDtype: null, + loadedKvCacheDtype: null, + customContextLength: null, defaultChatTemplate: null, chatTemplateOverride: null, activeThreadId: null, @@ -235,6 +262,11 @@ export const useChatRuntimeStore = create((set) => ({ saveBool(AUTO_TITLE_KEY, autoTitle); return { autoTitle }; }), + setHfToken: (hfToken) => + set(() => { + saveString(HF_TOKEN_KEY, hfToken); + return { hfToken }; + }), setModelsError: (modelsError) => set({ modelsError }), setCheckpoint: (modelId, ggufVariant) => set((state) => ({ @@ -261,6 +293,8 @@ export const useChatRuntimeStore = create((set) => ({ codeToolsEnabled: false, toolStatus: null, kvCacheDtype: null, + loadedKvCacheDtype: null, + customContextLength: null, defaultChatTemplate: null, chatTemplateOverride: null, })), @@ -285,6 +319,7 @@ export const useChatRuntimeStore = create((set) => ({ return { toolCallTimeout }; }), setKvCacheDtype: (kvCacheDtype) => set({ kvCacheDtype }), + setCustomContextLength: (customContextLength) => set({ customContextLength }), setChatTemplateOverride: (chatTemplateOverride) => set({ chatTemplateOverride }), setPendingAudio: (base64, name) => set({ pendingAudioBase64: base64, pendingAudioName: name }), diff --git a/studio/frontend/src/features/chat/thread-sidebar.tsx b/studio/frontend/src/features/chat/thread-sidebar.tsx index 53c5521dc7..ba97d2ee6e 100644 --- a/studio/frontend/src/features/chat/thread-sidebar.tsx +++ b/studio/frontend/src/features/chat/thread-sidebar.tsx @@ -172,7 +172,7 @@ export function ThreadSidebar({ Learn more in docs