Merge branch 'main' into pip

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.15.7
+    rev: v0.15.8
     hooks:
       - id: ruff
         args:

README.md

@@ -68,6 +68,7 @@ unsloth studio -H 0.0.0.0 -p 8888
 ```
 
 #### Update
+To update, use the same install commands as above. Or run (does not work on Windows):
 ```bash
 unsloth studio update
 ```
@@ -152,7 +153,7 @@ unsloth studio -H 0.0.0.0 -p 8888
 ```
 Then to update :
 ```bash
-unsloth studio update --local
+unsloth studio update
 ```
 
 #### Developer installs: Windows PowerShell:
@@ -165,7 +166,7 @@ unsloth studio -H 0.0.0.0 -p 8888
 ```
 Then to update :
 ```bash
-unsloth studio update --local
+unsloth studio update
 ```
 
 #### Nightly: MacOS, Linux, WSL:

install.ps1

@@ -6,6 +6,7 @@
 
 function Install-UnslothStudio {
     $ErrorActionPreference = "Stop"
+    $script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq "1")
 
     # ── Parse flags ──
     $StudioLocalInstall = $false
@@ -17,6 +18,8 @@ function Install-UnslothStudio {
         switch ($argList[$i]) {
             "--local"    { $StudioLocalInstall = $true }
             "--no-torch" { $SkipTorch = $true }
+            "--verbose"  { $script:UnslothVerbose = $true }
+            "-v"         { $script:UnslothVerbose = $true }
             "--package"  {
                 $i++
                 if ($i -ge $argList.Count) {
@@ -27,6 +30,12 @@ function Install-UnslothStudio {
             }
         }
     }
+    # Propagate to child processes so they also respect verbose mode.
+    # Process-scoped -- does not persist.
+    if ($script:UnslothVerbose) {
+        $env:UNSLOTH_VERBOSE = '1'
+    }
+
     if ($StudioLocalInstall) {
         $RepoRoot = (Resolve-Path (Split-Path -Parent $PSCommandPath)).Path
         if (-not (Test-Path (Join-Path $RepoRoot "pyproject.toml"))) {
@@ -39,10 +48,55 @@ function Install-UnslothStudio {
     $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio"
     $VenvDir = Join-Path $StudioHome "unsloth_studio"
 
+    $Rule = [string]::new([char]0x2500, 52)
+    $Sloth = [char]::ConvertFromUtf32(0x1F9A5)
+
+    function Enable-StudioVirtualTerminal {
+        if ($env:NO_COLOR) { return $false }
+        try {
+            if (-not ("StudioVT.Native" -as [type])) {
+                Add-Type -Namespace StudioVT -Name Native -MemberDefinition @'
+[DllImport("kernel32.dll")] public static extern IntPtr GetStdHandle(int nStdHandle);
+[DllImport("kernel32.dll")] public static extern bool GetConsoleMode(IntPtr h, out uint m);
+[DllImport("kernel32.dll")] public static extern bool SetConsoleMode(IntPtr h, uint m);
+'@ -ErrorAction Stop
+            }
+            $h = [StudioVT.Native]::GetStdHandle(-11)
+            [uint32]$mode = 0
+            if (-not [StudioVT.Native]::GetConsoleMode($h, [ref]$mode)) { return $false }
+            $mode = $mode -bor 0x0004
+            return [StudioVT.Native]::SetConsoleMode($h, $mode)
+        } catch {
+            return $false
+        }
+    }
+    $script:StudioVtOk = Enable-StudioVirtualTerminal
+
+    function Get-StudioAnsi {
+        param(
+            [Parameter(Mandatory = $true)]
+            [ValidateSet('Title', 'Dim', 'Ok', 'Warn', 'Err', 'Reset')]
+            [string]$Kind
+        )
+        $e = [char]27
+        switch ($Kind) {
+            'Title' { return "${e}[38;5;150m" }
+            'Dim'   { return "${e}[38;5;245m" }
+            'Ok'    { return "${e}[38;5;108m" }
+            'Warn'  { return "${e}[38;5;136m" }
+            'Err'   { return "${e}[91m" }
+            'Reset' { return "${e}[0m" }
+        }
+    }
+
     Write-Host ""
-    Write-Host "========================================="
-    Write-Host "   Unsloth Studio Installer (Windows)"
-    Write-Host "========================================="
+    if ($script:StudioVtOk -and -not $env:NO_COLOR) {
+        Write-Host ("  " + (Get-StudioAnsi Title) + $Sloth + " Unsloth Studio Installer (Windows)" + (Get-StudioAnsi Reset))
+        Write-Host ("  {0}{1}{2}" -f (Get-StudioAnsi Dim), $Rule, (Get-StudioAnsi Reset))
+    } else {
+        Write-Host ("  {0} Unsloth Studio Installer (Windows)" -f $Sloth) -ForegroundColor DarkGreen
+        Write-Host "  $Rule" -ForegroundColor DarkGray
+    }
     Write-Host ""
 
     # ── Helper: refresh PATH from registry (deduplicating entries) ──
@@ -62,13 +116,96 @@ function Install-UnslothStudio {
         $env:Path = $unique -join ";"
     }
 
+    function step {
+        param(
+            [Parameter(Mandatory = $true)][string]$Label,
+            [Parameter(Mandatory = $true)][string]$Value,
+            [string]$Color = "Green"
+        )
+        if ($script:StudioVtOk -and -not $env:NO_COLOR) {
+            $dim = Get-StudioAnsi Dim
+            $rst = Get-StudioAnsi Reset
+            $val = switch ($Color) {
+                'Green' { Get-StudioAnsi Ok }
+                'Yellow' { Get-StudioAnsi Warn }
+                'Red' { Get-StudioAnsi Err }
+                'DarkGray' { Get-StudioAnsi Dim }
+                default { Get-StudioAnsi Ok }
+            }
+            $padded = if ($Label.Length -ge 15) { $Label.Substring(0, 15) } else { $Label.PadRight(15) }
+            Write-Host ("  {0}{1}{2}{3}{4}{2}" -f $dim, $padded, $rst, $val, $Value)
+        } else {
+            $padded = if ($Label.Length -ge 15) { $Label.Substring(0, 15) } else { $Label.PadRight(15) }
+            Write-Host ("  {0}" -f $padded) -NoNewline -ForegroundColor DarkGray
+            $fc = switch ($Color) {
+                'Green' { 'DarkGreen' }
+                'Yellow' { 'Yellow' }
+                'Red' { 'Red' }
+                'DarkGray' { 'DarkGray' }
+                default { 'DarkGreen' }
+            }
+            Write-Host $Value -ForegroundColor $fc
+        }
+    }
+
+    function substep {
+        param(
+            [Parameter(Mandatory = $true)][string]$Message,
+            [string]$Color = "DarkGray"
+        )
+        if ($script:StudioVtOk -and -not $env:NO_COLOR) {
+            $msgCol = switch ($Color) {
+                'Yellow' { (Get-StudioAnsi Warn) }
+                'Red' { (Get-StudioAnsi Err) }
+                default { (Get-StudioAnsi Dim) }
+            }
+            $pad = "".PadRight(15)
+            Write-Host ("  {0}{1}{2}{3}" -f $msgCol, $pad, $Message, (Get-StudioAnsi Reset))
+        } else {
+            $fc = switch ($Color) {
+                'Yellow' { 'Yellow' }
+                'Red' { 'Red' }
+                default { 'DarkGray' }
+            }
+            Write-Host ("  {0,-15}{1}" -f "", $Message) -ForegroundColor $fc
+        }
+    }
+
+    # Run native commands quietly by default to match install.sh behavior.
+    # Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1.
+    function Invoke-InstallCommand {
+        param(
+            [Parameter(Mandatory = $true)][ScriptBlock]$Command
+        )
+        $prevEap = $ErrorActionPreference
+        $ErrorActionPreference = "Continue"
+        try {
+            # Reset to avoid stale values from prior native commands.
+            $global:LASTEXITCODE = 0
+            if ($script:UnslothVerbose) {
+                # Merge stderr into stdout so progress/warning output stays visible
+                # without flipping $? on successful native commands (PS 5.1 treats
+                # stderr records as errors that set $? = $false even on exit code 0).
+                & $Command 2>&1 | Out-Host
+            } else {
+                $output = & $Command 2>&1 | Out-String
+                if ($LASTEXITCODE -ne 0) {
+                    Write-Host $output -ForegroundColor Red
+                }
+            }
+            return [int]$LASTEXITCODE
+        } finally {
+            $ErrorActionPreference = $prevEap
+        }
+    }
+
     function New-StudioShortcuts {
         param(
             [Parameter(Mandatory = $true)][string]$UnslothExePath
         )
 
         if (-not (Test-Path $UnslothExePath)) {
-            Write-Host "[WARN] Cannot create shortcuts: unsloth.exe not found at $UnslothExePath" -ForegroundColor Yellow
+            substep "cannot create shortcuts, unsloth.exe not found at $UnslothExePath" "Yellow"
             return
         }
         try {
@@ -81,7 +218,7 @@ function Install-UnslothStudio {
 
             $localAppDataDir = $env:LOCALAPPDATA
             if (-not $localAppDataDir -or [string]::IsNullOrWhiteSpace($localAppDataDir)) {
-                Write-Host "[WARN] LOCALAPPDATA path unavailable; skipped shortcut creation" -ForegroundColor Yellow
+                substep "LOCALAPPDATA path unavailable; skipped shortcut creation" "Yellow"
                 return
             }
             $appDir = Join-Path $localAppDataDir "Unsloth Studio"
@@ -104,10 +241,10 @@ function Install-UnslothStudio {
                 $null
             }
             if (-not $desktopLink) {
-                Write-Host "[WARN] Desktop path unavailable; skipped desktop shortcut creation" -ForegroundColor Yellow
+                substep "Desktop path unavailable; skipped desktop shortcut creation" "Yellow"
             }
             if (-not $startMenuLink) {
-                Write-Host "[WARN] APPDATA/Start Menu path unavailable; skipped Start menu shortcut creation" -ForegroundColor Yellow
+                substep "APPDATA/Start Menu path unavailable; skipped Start menu shortcut creation" "Yellow"
             }
             $iconPath = Join-Path $appDir "unsloth.ico"
             $bundledIcon = $null
@@ -362,27 +499,27 @@ shell.Run cmd, 0, False
                         $shortcut.Save()
                         $createdShortcutCount++
                     } catch {
-                        Write-Host "[WARN] Could not create shortcut at ${linkPath}: $($_.Exception.Message)" -ForegroundColor Yellow
+                        substep "could not create shortcut at ${linkPath}: $($_.Exception.Message)" "Yellow"
                     }
                 }
                 if ($createdShortcutCount -gt 0) {
-                    Write-Host "[OK] Created Unsloth Studio shortcut(s): $createdShortcutCount" -ForegroundColor Green
+                    substep "Created Unsloth Studio shortcut"
                 } else {
-                    Write-Host "[WARN] No Unsloth Studio shortcuts were created" -ForegroundColor Yellow
+                    substep "no Unsloth Studio shortcuts were created" "Yellow"
                 }
             } catch {
-                Write-Host "[WARN] Shortcut creation unavailable: $($_.Exception.Message)" -ForegroundColor Yellow
+                substep "shortcut creation unavailable: $($_.Exception.Message)" "Yellow"
             }
         } catch {
-            Write-Host "[WARN] Shortcut setup failed; skipping shortcuts: $($_.Exception.Message)" -ForegroundColor Yellow
+            substep "shortcut setup failed; skipping shortcuts: $($_.Exception.Message)" "Yellow"
         }
     }
 
     # ── Check winget ──
     if (-not (Get-Command winget -ErrorAction SilentlyContinue)) {
-        Write-Host "Error: winget is not available." -ForegroundColor Red
-        Write-Host "       Install it from https://aka.ms/getwinget" -ForegroundColor Yellow
-        Write-Host "       or install Python $PythonVersion and uv manually, then re-run." -ForegroundColor Yellow
+        step "winget" "not available" "Red"
+        substep "Install it from https://aka.ms/getwinget" "Yellow"
+        substep "or install Python $PythonVersion and uv manually, then re-run." "Yellow"
         return
     }
 
@@ -460,10 +597,10 @@ shell.Run cmd, 0, False
     # Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null.
     $DetectedPython = Find-CompatiblePython
     if ($DetectedPython) {
-        Write-Host "==> Python already installed: Python $($DetectedPython.Version)"
+        step "python" "Python $($DetectedPython.Version) already installed"
     }
     if (-not $DetectedPython) {
-        Write-Host "==> Installing Python ${PythonVersion}..."
+        substep "installing Python ${PythonVersion}..."
         $pythonPackageId = "Python.Python.$PythonVersion"
         # Temporarily lower ErrorActionPreference so that winget stderr
         # (progress bars, warnings) does not become a terminating error
@@ -485,7 +622,7 @@ shell.Run cmd, 0, False
             # This handles both real failures AND "already installed" codes where
             # winget thinks Python is present but it's not actually on PATH
             # (e.g. user partially uninstalled, or installed via a different method).
-            Write-Host "    Python not found on PATH after winget. Retrying with --force..."
+            substep "Python not found on PATH after winget. Retrying with --force..." "Yellow"
             $ErrorActionPreference = "Continue"
             try {
                 winget install -e --id $pythonPackageId --accept-package-agreements --accept-source-agreements --force
@@ -507,7 +644,7 @@ shell.Run cmd, 0, False
 
     # ── Install uv if not present ──
     if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
-        Write-Host "==> Installing uv package manager..."
+        substep "installing uv package manager..."
         $prevEAP = $ErrorActionPreference
         $ErrorActionPreference = "Continue"
         try { winget install --id=astral-sh.uv -e --accept-package-agreements --accept-source-agreements } catch {}
@@ -515,15 +652,15 @@ shell.Run cmd, 0, False
         Refresh-SessionPath
         # Fallback: if winget didn't put uv on PATH, try the PowerShell installer
         if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
-            Write-Host "    Trying alternative uv installer..."
+            substep "trying alternative uv installer..." "Yellow"
             powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex"
             Refresh-SessionPath
         }
     }
 
     if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
-        Write-Host "Error: uv could not be installed." -ForegroundColor Red
-        Write-Host "       Install it from https://docs.astral.sh/uv/" -ForegroundColor Yellow
+        step "uv" "could not be installed" "Red"
+        substep "Install it from https://docs.astral.sh/uv/" "Yellow"
         return
     }
 
@@ -539,13 +676,13 @@ shell.Run cmd, 0, False
 
     if (Test-Path $VenvPython) {
         # New layout already exists -- nuke for fresh install
-        Write-Host "==> Removing existing environment for fresh install..."
+        substep "removing existing environment for fresh install..."
         Remove-Item -Recurse -Force $VenvDir
     } elseif (Test-Path (Join-Path $StudioHome ".venv\Scripts\python.exe")) {
         # Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating
         $OldVenv = Join-Path $StudioHome ".venv"
         $OldPy = Join-Path $OldVenv "Scripts\python.exe"
-        Write-Host "==> Found legacy Studio environment, validating..."
+        substep "found legacy Studio environment, validating..."
         $prevEAP2 = $ErrorActionPreference
         $ErrorActionPreference = "Continue"
         try {
@@ -554,32 +691,34 @@ shell.Run cmd, 0, False
         } catch { $torchOk = $false }
         $ErrorActionPreference = $prevEAP2
         if ($torchOk) {
-            Write-Host "   Legacy environment is healthy -- migrating..."
+            substep "legacy environment is healthy -- migrating..."
             Move-Item -Path $OldVenv -Destination $VenvDir -Force
-            Write-Host "   Moved .venv -> unsloth_studio"
+            substep "moved .venv -> unsloth_studio"
             $_Migrated = $true
         } else {
-            Write-Host "   Legacy environment failed validation -- creating fresh environment"
+            substep "legacy environment failed validation -- creating fresh environment" "Yellow"
             Remove-Item -Recurse -Force $OldVenv -ErrorAction SilentlyContinue
         }
     } elseif (Test-Path (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) {
         # CWD-relative venv from old install.ps1 -- migrate to absolute path
         $CwdVenv = Join-Path $env:USERPROFILE "unsloth_studio"
-        Write-Host "==> Found CWD-relative Studio environment, migrating to $VenvDir..."
+        substep "found CWD-relative Studio environment, migrating to $VenvDir..."
         Move-Item -Path $CwdVenv -Destination $VenvDir -Force
-        Write-Host "   Moved ~/unsloth_studio -> ~/.unsloth/studio/unsloth_studio"
+        substep "moved ~/unsloth_studio -> ~/.unsloth/studio/unsloth_studio"
         $_Migrated = $true
     }
 
     if (-not (Test-Path $VenvPython)) {
-        Write-Host "==> Creating Python $($DetectedPython.Version) virtual environment ($VenvDir)..."
-        uv venv $VenvDir --python "$($DetectedPython.Path)"
-        if ($LASTEXITCODE -ne 0) {
-            Write-Host "[ERROR] Failed to create virtual environment (exit code $LASTEXITCODE)" -ForegroundColor Red
+        step "venv" "creating Python $($DetectedPython.Version) virtual environment"
+        substep "$VenvDir"
+        $venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" }
+        if ($venvExit -ne 0) {
+            Write-Host "[ERROR] Failed to create virtual environment (exit code $venvExit)" -ForegroundColor Red
             return
         }
     } else {
-        Write-Host "==> Using migrated environment at $VenvDir"
+        step "venv" "using migrated environment"
+        substep "$VenvDir"
     }
 
     # ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ──
@@ -588,7 +727,7 @@ shell.Run cmd, 0, False
     try {
         $nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue
         if ($nvSmiCmd) {
-            & $nvSmiCmd.Source 2>&1 | Out-Null
+            & $nvSmiCmd.Source *> $null
             if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $nvSmiCmd.Source }
         }
     } catch {}
@@ -599,18 +738,18 @@ shell.Run cmd, 0, False
         )) {
             if (Test-Path $p) {
                 try {
-                    & $p 2>&1 | Out-Null
+                    & $p *> $null
                     if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $p; break }
                 } catch {}
             }
         }
     }
     if ($HasNvidiaSmi) {
-        Write-Host "[OK] NVIDIA GPU detected" -ForegroundColor Green
+        step "gpu" "NVIDIA GPU detected"
     } else {
-        Write-Host "[WARN] No NVIDIA GPU detected. Studio will run in chat-only (GGUF) mode." -ForegroundColor Yellow
-        Write-Host "       Training and GPU inference require an NVIDIA GPU with drivers installed." -ForegroundColor Yellow
-        Write-Host "       https://www.nvidia.com/Download/index.aspx" -ForegroundColor Yellow
+        step "gpu" "none (chat-only / GGUF)" "Yellow"
+        substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow"
+        substep "https://www.nvidia.com/Download/index.aspx" "Yellow"
     }
 
     # ── Choose the correct PyTorch index URL based on driver CUDA version ──
@@ -630,7 +769,7 @@ shell.Run cmd, 0, False
                 return "$baseUrl/cpu"
             }
         } catch {}
-        Write-Host "[WARN] Could not determine CUDA version from nvidia-smi, defaulting to cu126" -ForegroundColor Yellow
+        substep "could not determine CUDA version from nvidia-smi, defaulting to cu126" "Yellow"
         return "$baseUrl/cu126"
     }
     $TorchIndexUrl = Get-TorchIndexUrl
@@ -677,74 +816,101 @@ shell.Run cmd, 0, False
     if ($_Migrated) {
         # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state
         # in the new venv location, while preserving existing torch/CUDA
-        Write-Host "==> Upgrading unsloth in migrated environment..."
+        substep "upgrading unsloth in migrated environment..."
         if ($SkipTorch) {
             # No-torch: install unsloth + unsloth-zoo with --no-deps, then
             # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
-            uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo
-            $NoTorchReq = Find-NoTorchRuntimeFile
-            if ($NoTorchReq) {
-                uv pip install --python $VenvPython --no-deps -r $NoTorchReq
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo }
+            if ($baseInstallExit -eq 0) {
+                $NoTorchReq = Find-NoTorchRuntimeFile
+                if ($NoTorchReq) {
+                    $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq }
+                }
             }
         } else {
-            uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo }
+        }
+        if ($baseInstallExit -ne 0) {
+            Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
+            return
         }
         if ($StudioLocalInstall) {
-            Write-Host "==> Overlaying local repo (editable)..."
-            uv pip install --python $VenvPython -e $RepoRoot --no-deps
+            substep "overlaying local repo (editable)..."
+            $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
+            if ($overlayExit -ne 0) {
+                Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
+                return
+            }
         }
     } elseif ($TorchIndexUrl) {
         if ($SkipTorch) {
-            Write-Host "==> Skipping PyTorch (--no-torch flag set)."
+            substep "skipping PyTorch (--no-torch flag set)." "Yellow"
         } else {
-            Write-Host "==> Installing PyTorch ($TorchIndexUrl)..."
-            uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl
-            if ($LASTEXITCODE -ne 0) {
-                Write-Host "[ERROR] Failed to install PyTorch (exit code $LASTEXITCODE)" -ForegroundColor Red
+            substep "installing PyTorch ($TorchIndexUrl)..."
+            $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
+            if ($torchInstallExit -ne 0) {
+                Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
                 return
             }
         }
 
-        Write-Host "==> Installing unsloth (this may take a few minutes)..."
+        substep "installing unsloth (this may take a few minutes)..."
         if ($SkipTorch) {
             # No-torch: install unsloth + unsloth-zoo with --no-deps, then
             # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
-            uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo
-            $NoTorchReq = Find-NoTorchRuntimeFile
-            if ($NoTorchReq) {
-                uv pip install --python $VenvPython --no-deps -r $NoTorchReq
-            }
-            if ($StudioLocalInstall) {
-                Write-Host "==> Overlaying local repo (editable)..."
-                uv pip install --python $VenvPython -e $RepoRoot --no-deps
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo }
+            if ($baseInstallExit -eq 0) {
+                $NoTorchReq = Find-NoTorchRuntimeFile
+                if ($NoTorchReq) {
+                    $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq }
+                }
             }
         } elseif ($StudioLocalInstall) {
-            uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo
-            Write-Host "==> Overlaying local repo (editable)..."
-            uv pip install --python $VenvPython -e $RepoRoot --no-deps
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo }
         } else {
-            uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName"
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName" }
+        }
+        if ($baseInstallExit -ne 0) {
+            Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
+            return
+        }
+
+        if ($StudioLocalInstall) {
+            substep "overlaying local repo (editable)..."
+            $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
+            if ($overlayExit -ne 0) {
+                Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
+                return
+            }
         }
     } else {
         # Fallback: GPU detection failed to produce a URL -- let uv resolve torch
-        Write-Host "==> Installing unsloth (this may take a few minutes)..."
+        substep "installing unsloth (this may take a few minutes)..."
         if ($StudioLocalInstall) {
-            uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto
-            Write-Host "==> Overlaying local repo (editable)..."
-            uv pip install --python $VenvPython -e $RepoRoot --no-deps
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto }
+            if ($baseInstallExit -ne 0) {
+                Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
+                return
+            }
+            substep "overlaying local repo (editable)..."
+            $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
+            if ($overlayExit -ne 0) {
+                Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
+                return
+            }
         } else {
-            uv pip install --python $VenvPython "$PackageName" --torch-backend=auto
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "$PackageName" --torch-backend=auto }
+            if ($baseInstallExit -ne 0) {
+                Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
+                return
+            }
         }
     }
-    if ($LASTEXITCODE -ne 0) {
-        Write-Host "[ERROR] Failed to install unsloth (exit code $LASTEXITCODE)" -ForegroundColor Red
-        return
-    }
 
     # ── Run studio setup ──
     # setup.ps1 will handle installing Git, CMake, Visual Studio Build Tools,
     # CUDA Toolkit, Node.js, and other dependencies automatically via winget.
-    Write-Host "==> Running unsloth studio setup..."
+    step "setup" "running unsloth studio setup..."
     $UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe"
     if (-not (Test-Path $UnslothExe)) {
         Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red
@@ -754,17 +920,27 @@ shell.Run cmd, 0, False
         return
     }
     # Tell setup.ps1 to skip base package installation (install.ps1 already did it)
-    # Tell setup.ps1 to skip base package installation (install.ps1 already did it)
     $env:SKIP_STUDIO_BASE = "1"
     $env:STUDIO_PACKAGE_NAME = $PackageName
     $env:UNSLOTH_NO_TORCH = if ($SkipTorch) { "true" } else { "false" }
+    # Always set STUDIO_LOCAL_INSTALL explicitly to avoid stale values from
+    # a previous --local run in the same PowerShell session.
     if ($StudioLocalInstall) {
         $env:STUDIO_LOCAL_INSTALL = "1"
         $env:STUDIO_LOCAL_REPO = $RepoRoot
+    } else {
+        $env:STUDIO_LOCAL_INSTALL = "0"
+        Remove-Item Env:STUDIO_LOCAL_REPO -ErrorAction SilentlyContinue
     }
-    & $UnslothExe studio setup
-    if ($LASTEXITCODE -ne 0) {
-        Write-Host "[ERROR] unsloth studio setup failed (exit code $LASTEXITCODE)" -ForegroundColor Red
+    # Use 'studio setup' (not 'studio update') because 'update' pops
+    # SKIP_STUDIO_BASE, which would cause redundant package reinstallation
+    # and bypass the fast-path version check from PR #4667.
+    $studioArgs = @('studio', 'setup')
+    if ($script:UnslothVerbose) { $studioArgs += '--verbose' }
+    & $UnslothExe @studioArgs
+    $setupExit = $LASTEXITCODE
+    if ($setupExit -ne 0) {
+        Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red
         return
     }
 
@@ -780,27 +956,18 @@ shell.Run cmd, 0, False
             [System.Environment]::SetEnvironmentVariable("Path", "$ScriptsDir", "User")
         }
         Refresh-SessionPath
-        Write-Host "[OK] Added unsloth to PATH" -ForegroundColor Green
+        step "path" "added unsloth to PATH"
     }
 
-    Write-Host ""
-    Write-Host "========================================="
-    Write-Host "   Unsloth Studio installed!"
-    Write-Host "========================================="
-    Write-Host ""
-
     # Launch studio automatically in interactive terminals;
     # in non-interactive environments (CI, Docker) just print instructions.
     $IsInteractive = [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected)
     if ($IsInteractive) {
-        Write-Host "==> Launching Unsloth Studio..."
-        Write-Host ""
         & $UnslothExe studio -H 0.0.0.0 -p 8888
     } else {
-        Write-Host "  To launch, run:"
-        Write-Host ""
-        Write-Host "    & `"$VenvDir\Scripts\Activate.ps1`""
-        Write-Host "    unsloth studio -H 0.0.0.0 -p 8888"
+        step "launch" "manual commands:"
+        substep "& `"$VenvDir\Scripts\Activate.ps1`""
+        substep "unsloth studio -H 0.0.0.0 -p 8888"
         Write-Host ""
     }
 }

install.sh

@@ -8,11 +8,36 @@
 # Usage (py):    ./install.sh --python 3.12  (override auto-detected Python version)
 set -e
 
+# ── Output style (aligned with studio/setup.sh) ──
+RULE=""
+_rule_i=0
+while [ "$_rule_i" -lt 52 ]; do
+    RULE="${RULE}─"
+    _rule_i=$((_rule_i + 1))
+done
+if [ -n "${NO_COLOR:-}" ]; then
+    C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST=
+elif [ -t 1 ] || [ -n "${FORCE_COLOR:-}" ]; then
+    _ESC="$(printf '\033')"
+    C_TITLE="${_ESC}[38;5;150m"
+    C_DIM="${_ESC}[38;5;245m"
+    C_OK="${_ESC}[38;5;108m"
+    C_WARN="${_ESC}[38;5;136m"
+    C_ERR="${_ESC}[91m"
+    C_RST="${_ESC}[0m"
+else
+    C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST=
+fi
+
+step()    { printf "  ${C_DIM}%-15.15s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; }
+substep() { printf "  ${C_DIM}%-15s${2:-$C_DIM}%s${C_RST}\n" "" "$1"; }
+
 # ── Parse flags ──
 STUDIO_LOCAL_INSTALL=false
 PACKAGE_NAME="unsloth"
 _USER_PYTHON=""
 _NO_TORCH_FLAG=false
+_VERBOSE=false
 _next_is_package=false
 _next_is_python=false
 for arg in "$@"; do
@@ -31,9 +56,44 @@ for arg in "$@"; do
         --package) _next_is_package=true ;;
         --python) _next_is_python=true ;;
         --no-torch) _NO_TORCH_FLAG=true ;;
+        --verbose|-v) _VERBOSE=true ;;
     esac
 done
 
+if [ "$_VERBOSE" = true ]; then
+    export UNSLOTH_VERBOSE=1
+fi
+
+_is_verbose() {
+    [ "${UNSLOTH_VERBOSE:-0}" = "1" ]
+}
+
+run_maybe_quiet() {
+    if _is_verbose; then
+        "$@"
+    else
+        "$@" > /dev/null 2>&1
+    fi
+}
+
+run_install_cmd() {
+    _label="$1"
+    shift
+    if _is_verbose; then
+        "$@" && return 0
+        _rc=$?
+        step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
+        return "$_rc"
+    fi
+    _log=$(mktemp)
+    "$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; }
+    _rc=$?
+    step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
+    cat "$_log" >&2
+    rm -f "$_log"
+    return $_rc
+}
+
 if [ "$_next_is_package" = true ]; then
     echo "❌ ERROR: --package requires an argument." >&2
     exit 1
@@ -643,14 +703,13 @@ WSLPS1_EOF
     fi
 
     if [ "$_css_created" -eq 1 ]; then
-        echo "[OK] Created Unsloth Studio shortcut(s)"
+        substep "Created Unsloth Studio shortcut"
     fi
 }
 
 echo ""
-echo "========================================="
-echo "   Unsloth Studio Installer"
-echo "========================================="
+printf "  ${C_TITLE}%s${C_RST}\n" "🦥 Unsloth Studio Installer"
+printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
 echo ""
 
 # ── Detect platform ──
@@ -660,7 +719,7 @@ if [ "$(uname)" = "Darwin" ]; then
 elif grep -qi microsoft /proc/version 2>/dev/null; then
     OS="wsl"
 fi
-echo "==> Platform: $OS"
+step "platform" "$OS"
 
 # ── Architecture detection & Python version ──
 _ARCH=$(uname -m)
@@ -740,8 +799,8 @@ MISSING=$(echo "$MISSING" | sed 's/^ *//')
 
 if [ -n "$MISSING" ]; then
     echo ""
-    echo "==> Unsloth Studio needs these packages: $MISSING"
-    echo "    These are needed to build the GGUF inference engine."
+    step "deps" "missing: $MISSING" "$C_WARN"
+    substep "These are needed to build the GGUF inference engine."
 
     case "$OS" in
         macos)
@@ -766,7 +825,7 @@ if [ -n "$MISSING" ]; then
     esac
     echo ""
 else
-    echo "==> All system dependencies found."
+    step "deps" "all system dependencies found"
 fi
 
 # ── Install uv ──
@@ -812,10 +871,10 @@ _uv_version_ok() {
 }
 
 if ! command -v uv >/dev/null 2>&1 || ! _uv_version_ok uv; then
-    echo "==> Installing uv package manager..."
+    substep "installing uv package manager..."
     _uv_tmp=$(mktemp)
     download "https://astral.sh/uv/install.sh" "$_uv_tmp"
-    sh "$_uv_tmp" </dev/null
+    run_maybe_quiet sh "$_uv_tmp" </dev/null
     rm -f "$_uv_tmp"
     if [ -f "$HOME/.local/bin/env" ]; then
         . "$HOME/.local/bin/env"
@@ -833,7 +892,7 @@ if [ -x "$VENV_DIR/bin/python" ]; then
     rm -rf "$VENV_DIR"
 elif [ -x "$STUDIO_HOME/.venv/bin/python" ]; then
     # Old layout exists — validate before migrating
-    echo "==> Found legacy Studio environment, validating..."
+    substep "found legacy Studio environment, validating..."
     if "$STUDIO_HOME/.venv/bin/python" -c "
 import torch
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -866,8 +925,9 @@ if [ "$SKIP_TORCH" = true ] && [ "$MAC_INTEL" = true ] && [ -z "$_USER_PYTHON" ]
 fi
 
 if [ ! -x "$VENV_DIR/bin/python" ]; then
-    echo "==> Creating Python ${PYTHON_VERSION} virtual environment (${VENV_DIR})..."
-    uv venv "$VENV_DIR" --python "$PYTHON_VERSION"
+    step "venv" "creating Python ${PYTHON_VERSION} virtual environment"
+    substep "$VENV_DIR"
+    run_install_cmd "create venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION"
 fi
 
 # Guard against Python 3.13.8 torch import bug on Apple Silicon
@@ -880,12 +940,13 @@ if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
         echo "  Recreating venv with Python 3.12..."
         rm -rf "$VENV_DIR"
         PYTHON_VERSION="3.12"
-        uv venv "$VENV_DIR" --python "$PYTHON_VERSION"
+        run_install_cmd "recreate venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION"
     fi
 fi
 
 if [ -x "$VENV_DIR/bin/python" ]; then
-    echo "==> Using environment at ${VENV_DIR}"
+    step "venv" "using environment"
+    substep "${VENV_DIR}"
 fi
 
 # ── Resolve repo root (for --local installs) ──
@@ -960,71 +1021,71 @@ _VENV_PY="$VENV_DIR/bin/python"
 if [ "$_MIGRATED" = true ]; then
     # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state
     # in the new venv location, while preserving existing torch/CUDA
-    echo "==> Upgrading unsloth in migrated environment..."
+    substep "upgrading unsloth in migrated environment..."
     if [ "$SKIP_TORCH" = true ]; then
         # No-torch: install unsloth + unsloth-zoo with --no-deps (current
         # PyPI metadata still declares torch as a hard dep), then install
         # runtime deps (typer, safetensors, transformers, etc.) with --no-deps
         # to prevent transitive torch resolution.
-        uv pip install --python "$_VENV_PY" --no-deps \
+        run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
             --reinstall-package unsloth --reinstall-package unsloth-zoo \
             "unsloth>=2026.3.16" unsloth-zoo
         _NO_TORCH_RT="$(_find_no_torch_runtime)"
         if [ -n "$_NO_TORCH_RT" ]; then
-            uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
+            run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
         fi
     else
-        uv pip install --python "$_VENV_PY" \
+        run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
             --reinstall-package unsloth --reinstall-package unsloth-zoo \
             "unsloth>=2026.3.16" unsloth-zoo
     fi
     if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
-        echo "==> Overlaying local repo (editable)..."
-        uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
+        substep "overlaying local repo (editable)..."
+        run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
     fi
 elif [ -n "$TORCH_INDEX_URL" ]; then
     # Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac)
     if [ "$SKIP_TORCH" = true ]; then
-        echo "==> Skipping PyTorch (--no-torch or Intel Mac x86_64)."
+        substep "skipping PyTorch (--no-torch or Intel Mac x86_64)." "$C_WARN"
     else
-        echo "==> Installing PyTorch ($TORCH_INDEX_URL)..."
-        uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \
+        substep "installing PyTorch ($TORCH_INDEX_URL)..."
+        run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \
             --index-url "$TORCH_INDEX_URL"
     fi
     # Fresh: Step 2 - install unsloth, preserving pre-installed torch
-    echo "==> Installing unsloth (this may take a few minutes)..."
+    substep "installing unsloth (this may take a few minutes)..."
     if [ "$SKIP_TORCH" = true ]; then
         # No-torch: install unsloth + unsloth-zoo with --no-deps, then
         # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
-        uv pip install --python "$_VENV_PY" --no-deps \
+        run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
             --upgrade-package unsloth --upgrade-package unsloth-zoo \
             "unsloth>=2026.3.16" unsloth-zoo
         _NO_TORCH_RT="$(_find_no_torch_runtime)"
         if [ -n "$_NO_TORCH_RT" ]; then
-            uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
+            run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
         fi
         if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
-            echo "==> Overlaying local repo (editable)..."
-            uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
+            substep "overlaying local repo (editable)..."
+            run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
         fi
     elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
-        uv pip install --python "$_VENV_PY" \
+        run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \
             --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo
-        echo "==> Overlaying local repo (editable)..."
-        uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
+        substep "overlaying local repo (editable)..."
+        run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
     else
-        uv pip install --python "$_VENV_PY" \
+        run_install_cmd "install unsloth" uv pip install --python "$_VENV_PY" \
             --upgrade-package unsloth "$PACKAGE_NAME"
     fi
 else
     # Fallback: GPU detection failed to produce a URL -- let uv resolve torch
-    echo "==> Installing unsloth (this may take a few minutes)..."
+    substep "installing unsloth (this may take a few minutes)..."
     if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
-        uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto
-        echo "==> Overlaying local repo (editable)..."
-        uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
+        run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto
+        substep "overlaying local repo (editable)..."
+        run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
     else
-        uv pip install --python "$_VENV_PY" "$PACKAGE_NAME" --torch-backend=auto
+        run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "$PACKAGE_NAME" --torch-backend=auto
     fi
 fi
 
@@ -1059,19 +1120,33 @@ if [ -n "$VENV_ABS_BIN" ]; then
     export PATH="$VENV_ABS_BIN:$PATH"
 fi
 
-echo "==> Running unsloth setup..."
+if ! command -v bash >/dev/null 2>&1; then
+    step "setup" "bash is required to run studio setup" "$C_ERR"
+    substep "Please install bash and re-run install.sh"
+    exit 1
+fi
+
+step "setup" "running unsloth studio update..."
+# install.sh already installs base packages (unsloth + unsloth-zoo) and
+# no-torch-runtime.txt above, so tell install_python_stack.py to skip
+# the base step to avoid redundant reinstallation.
+_SKIP_BASE=1
+# Run setup.sh outside set -e so that a llama.cpp build failure (exit 1)
+# does not skip PATH setup, shortcuts, and launch below.  We capture the
+# exit code and propagate it after post-install steps finish.
+_SETUP_EXIT=0
 if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
-    SKIP_STUDIO_BASE=1 \
+    SKIP_STUDIO_BASE="$_SKIP_BASE" \
     STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \
     STUDIO_LOCAL_INSTALL=1 \
     STUDIO_LOCAL_REPO="$_REPO_ROOT" \
     UNSLOTH_NO_TORCH="$SKIP_TORCH" \
-    bash "$SETUP_SH" </dev/null
+    bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
 else
-    SKIP_STUDIO_BASE=1 \
+    SKIP_STUDIO_BASE="$_SKIP_BASE" \
     STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \
     UNSLOTH_NO_TORCH="$SKIP_TORCH" \
-    bash "$SETUP_SH" </dev/null
+    bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
 fi
 
 # ── Make 'unsloth' available globally via ~/.local/bin ──
@@ -1096,7 +1171,7 @@ case ":$PATH:" in
                 echo '' >> "$_SHELL_PROFILE"
                 echo '# Added by Unsloth installer' >> "$_SHELL_PROFILE"
                 echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$_SHELL_PROFILE"
-                echo "==> Added ~/.local/bin to PATH in $_SHELL_PROFILE"
+                step "path" "added ~/.local/bin to PATH in $_SHELL_PROFILE"
             fi
         fi
         export PATH="$_LOCAL_BIN:$PATH"
@@ -1105,17 +1180,30 @@ esac
 
 create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS"
 
+# If setup.sh failed, report and exit now.
+# PATH and shortcuts are already set up so the user can fix and retry.
+if [ "$_SETUP_EXIT" -ne 0 ]; then
+    echo ""
+    step "error" "studio setup failed (exit code $_SETUP_EXIT)" "$C_ERR"
+    substep "Check the output above for details, then re-run:"
+    if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
+        substep "  unsloth studio update --local"
+    else
+        substep "  unsloth studio update"
+    fi
+    echo ""
+    exit "$_SETUP_EXIT"
+fi
+
 echo ""
-echo "========================================="
-echo "   Unsloth Studio installed!"
-echo "========================================="
+printf "  ${C_TITLE}%s${C_RST}\n" "Unsloth Studio installed!"
+printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
 echo ""
 
 # Launch studio automatically in interactive terminals;
 # in non-interactive environments (Docker, CI, cloud-init) just print instructions.
 if [ -t 1 ]; then
-    echo "==> Launching Unsloth Studio..."
-    echo ""
+    step "launch" "starting Unsloth Studio..."
     "$VENV_DIR/bin/unsloth" studio -H 0.0.0.0 -p 8888
     _LAUNCH_EXIT=$?
     if [ "$_LAUNCH_EXIT" -ne 0 ] && [ "$_MIGRATED" = true ]; then
@@ -1130,13 +1218,10 @@ if [ -t 1 ]; then
     fi
     exit "$_LAUNCH_EXIT"
 else
-    echo "  To launch, run:"
-    echo ""
-    echo "    unsloth studio -H 0.0.0.0 -p 8888"
-    echo ""
-    echo "  Or activate the environment first:"
-    echo ""
-    echo "    source ${VENV_DIR}/bin/activate"
-    echo "    unsloth studio -H 0.0.0.0 -p 8888"
+    step "launch" "manual commands:"
+    substep "unsloth studio -H 0.0.0.0 -p 8888"
+    substep "or activate env first:"
+    substep "source ${VENV_DIR}/bin/activate"
+    substep "unsloth studio -H 0.0.0.0 -p 8888"
     echo ""
 fi

studio/backend/assets/configs/full_finetune.yaml

@@ -10,13 +10,13 @@ training:
   load_in_4bit: false
   output_dir: outputs
   num_epochs: 1
-  learning_rate: 0.0002
+  learning_rate: 2e-5
   batch_size: 1
   gradient_accumulation_steps: 4
   warmup_steps: 5
   max_steps: 0
   save_steps: 0
-  weight_decay: 0.01
+  weight_decay: 0.001
   random_seed: 3407
   packing: false
   train_on_completions: false

studio/backend/assets/configs/lora_text.yaml

@@ -16,7 +16,7 @@ training:
   warmup_steps: 5
   max_steps: 0
   save_steps: 0
-  weight_decay: 0.01
+  weight_decay: 0.001
   random_seed: 3407
   packing: false
   train_on_completions: false

studio/backend/assets/configs/model_defaults/default.yaml

@@ -6,13 +6,13 @@ training:
   max_seq_length: 2048
   # num_epochs: 4
   num_epochs: 0
-  learning_rate: 5e-5
+  learning_rate: 2e-4
   batch_size: 2
   gradient_accumulation_steps: 4
   warmup_ratio: 0.1
   max_steps: 30
   save_steps: 30
-  weight_decay: 0.01
+  weight_decay: 0.001
   random_seed: 3407
   packing: false
   train_on_completions: true

studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml

@@ -12,7 +12,7 @@ training:
   warmup_ratio: 0.03
   max_steps: 30
   save_steps: 30
-  weight_decay: 0.01
+  weight_decay: 0.001
   random_seed: 3407
   packing: false
   train_on_completions: false

studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml

@@ -11,7 +11,7 @@ training:
   warmup_ratio: 0.03
   max_steps: 30
   save_steps: 30
-  weight_decay: 0.01
+  weight_decay: 0.001
   random_seed: 3407
   packing: false
   train_on_completions: false

studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml

@@ -11,7 +11,7 @@ training:
   warmup_ratio: 0.03
   max_steps: 30
   save_steps: 30
-  weight_decay: 0.01
+  weight_decay: 0.001
   random_seed: 3407
   packing: false
   train_on_completions: false

studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml

@@ -11,7 +11,7 @@ training:
   warmup_ratio: 0.03
   max_steps: 30
   save_steps: 30
-  weight_decay: 0.01
+  weight_decay: 0.001
   random_seed: 3407
   packing: false
   train_on_completions: false

studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml

@@ -11,7 +11,7 @@ training:
   warmup_ratio: 0.03
   max_steps: 30
   save_steps: 30
-  weight_decay: 0.01
+  weight_decay: 0.001
   random_seed: 3407
   packing: false
   train_on_completions: false

studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml

@@ -13,7 +13,7 @@ training:
   warmup_steps: 5
   max_steps: 30
   save_steps: 30
-  weight_decay: 0.01
+  weight_decay: 0.001
   random_seed: 3407
   packing: false
   train_on_completions: true

studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml

@@ -13,7 +13,7 @@ training:
   warmup_steps: 5
   max_steps: 30
   save_steps: 30
-  weight_decay: 0.01
+  weight_decay: 0.001
   random_seed: 3407
   packing: false
   train_on_completions: true

studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml

@@ -13,7 +13,7 @@ training:
   warmup_steps: 0
   max_steps: 30
   save_steps: 30
-  weight_decay: 0.01
+  weight_decay: 0.001
   random_seed: 3407
   packing: false
   train_on_completions: true

studio/backend/assets/configs/vision_lora.yaml

@@ -16,7 +16,7 @@ training:
   warmup_steps: 5
   max_steps: 0
   save_steps: 0
-  weight_decay: 0.01
+  weight_decay: 0.001
   random_seed: 3407
   packing: false
   train_on_completions: false

studio/backend/core/export/orchestrator.py

@@ -217,6 +217,7 @@ class ExportOrchestrator:
         max_seq_length: int = 2048,
         load_in_4bit: bool = True,
         trust_remote_code: bool = False,
+        hf_token: Optional[str] = None,
     ) -> Tuple[bool, str]:
         """Load a checkpoint for export.
 
@@ -227,6 +228,7 @@ class ExportOrchestrator:
             "max_seq_length": max_seq_length,
             "load_in_4bit": load_in_4bit,
             "trust_remote_code": trust_remote_code,
+            "hf_token": hf_token,
         }
 
         # Always kill existing subprocess and spawn fresh.

studio/backend/core/inference/_html_to_md.py

@@ -0,0 +1,439 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""
+Minimal HTML-to-Markdown converter using only the standard library.
+
+Replaces the external ``html2text`` (GPL-3.0) dependency with a ~250-line
+``html.parser.HTMLParser`` subclass.  Covers headings, links, bold/italic,
+lists, tables, blockquotes, code blocks, and entity decoding.
+"""
+
+from __future__ import annotations
+
+import html
+import re
+from html.parser import HTMLParser
+
+__all__ = ["html_to_markdown"]
+
+_SKIP_TAGS = frozenset({"script", "style", "head", "noscript", "svg", "math"})
+_BLOCK_TAGS = frozenset(
+    {
+        "p",
+        "div",
+        "section",
+        "article",
+        "header",
+        "footer",
+        "main",
+        "aside",
+        "nav",
+        "figure",
+        "figcaption",
+        "details",
+        "summary",
+        "dl",
+        "dt",
+        "dd",
+    }
+)
+_HEADING_TAGS = frozenset({"h1", "h2", "h3", "h4", "h5", "h6"})
+_INLINE_EMPHASIS = {"strong": "**", "b": "**", "em": "*", "i": "*"}
+
+
+class _MarkdownRenderer(HTMLParser):
+    """HTMLParser subclass that emits Markdown tokens into a list."""
+
+    def __init__(self):
+        super().__init__(convert_charrefs = False)
+        self._out: list[str] = []
+        self._skip_depth: int = 0
+
+        # Link state
+        self._link_href: str | None = None
+        self._link_text_parts: list[str] = []
+        self._in_link: bool = False
+
+        # List state
+        self._list_stack: list[str] = []  # "ul" or "ol"
+        self._ol_counter: list[int] = []
+
+        # Table state
+        self._in_table: bool = False
+        self._current_row: list[str] = []
+        self._cell_parts: list[str] = []
+        self._in_cell: bool = False
+        self._header_row_done: bool = False
+        self._row_has_th: bool = False
+        self._is_first_row: bool = False
+
+        # Pre/code state
+        self._in_pre: bool = False
+        self._pre_parts: list[str] = []
+        self._in_inline_code: bool = False
+
+        # Blockquote state -- stack of output buffers so nested
+        # blockquotes each collect their own content and get prefixed
+        # with the correct number of ">" markers on close.
+        self._bq_stack: list[list[str]] = []
+
+    # ------------------------------------------------------------------
+    def _emit(self, text: str) -> None:
+        if self._in_link:
+            self._link_text_parts.append(text)
+        elif self._in_cell:
+            self._cell_parts.append(text)
+        elif self._in_pre:
+            self._pre_parts.append(text)
+        elif self._bq_stack:
+            self._bq_stack[-1].append(text)
+        else:
+            self._out.append(text)
+
+    # ------------------------------------------------------------------
+    def _prefix_blockquote(self, content: str) -> str:
+        """Prefix every line of *content* with ``> ``."""
+        # Strip trailing whitespace first, then collapse blank lines
+        content = re.sub(r"[ \t]+$", "", content, flags = re.MULTILINE)
+        content = re.sub(r"\n{3,}", "\n\n", content).strip()
+        if not content:
+            return ""
+        lines = content.split("\n")
+        prefixed: list[str] = []
+        for line in lines:
+            if line.strip():
+                prefixed.append("> " + line)
+            else:
+                prefixed.append(">")
+        return "\n".join(prefixed)
+
+    # ------------------------------------------------------------------
+    # Table helpers -- flush open cells and rows so that HTML with
+    # omitted optional end tags (</td>, </tr>) does not lose data.
+    # ------------------------------------------------------------------
+    def _finish_cell(self) -> None:
+        if not self._in_cell:
+            return
+        self._in_cell = False
+        cell_text = "".join(self._cell_parts).strip().replace("\n", " ")
+        cell_text = cell_text.replace("|", "\\|")
+        self._current_row.append(cell_text)
+        self._cell_parts = []
+
+    def _finish_row(self) -> None:
+        if not self._current_row:
+            return
+        line = "| " + " | ".join(self._current_row) + " |"
+        self._emit(line + "\n")
+        if not self._header_row_done and (self._row_has_th or self._is_first_row):
+            sep = "| " + " | ".join("---" for _ in self._current_row) + " |"
+            self._emit(sep + "\n")
+            self._header_row_done = True
+        self._is_first_row = False
+        self._current_row = []
+        self._row_has_th = False
+
+    # ------------------------------------------------------------------
+    # Link text helper -- normalize whitespace so block-level content
+    # inside an <a> does not produce multiline Markdown link labels.
+    # ------------------------------------------------------------------
+    def _finish_link(self) -> None:
+        text = re.sub(r"\s+", " ", "".join(self._link_text_parts)).strip()
+        href = self._link_href or ""
+        self._in_link = False
+        if href and text:
+            self._emit(f"[{text}]({href})")
+        elif text:
+            self._emit(text)
+
+    # ------------------------------------------------------------------
+    # Tag handlers
+    # ------------------------------------------------------------------
+    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
+        tag = tag.lower()
+
+        if tag in _SKIP_TAGS:
+            self._skip_depth += 1
+            return
+        if self._skip_depth:
+            return
+
+        attr_dict = dict(attrs)
+
+        if tag in _HEADING_TAGS:
+            level = int(tag[1])
+            self._emit("\n\n" + "#" * level + " ")
+
+        elif tag == "a":
+            self._link_href = attr_dict.get("href")
+            self._link_text_parts = []
+            self._in_link = True
+
+        elif tag in _INLINE_EMPHASIS:
+            self._emit(_INLINE_EMPHASIS[tag])
+
+        elif tag == "br":
+            self._emit("\n")
+
+        elif tag in _BLOCK_TAGS:
+            self._emit("\n\n")
+
+        elif tag == "hr":
+            self._emit("\n\n---\n\n")
+
+        elif tag == "blockquote":
+            self._emit("\n\n")
+            self._bq_stack.append([])
+
+        elif tag == "ul":
+            self._list_stack.append("ul")
+            self._emit("\n")
+
+        elif tag == "ol":
+            self._list_stack.append("ol")
+            start_attr = attr_dict.get("start")
+            try:
+                start = int(start_attr) if start_attr is not None else 1
+            except (ValueError, TypeError):
+                start = 1
+            self._ol_counter.append(start - 1)
+            self._emit("\n")
+
+        elif tag == "li":
+            indent = "  " * max(0, len(self._list_stack) - 1)
+            if self._list_stack and self._list_stack[-1] == "ol":
+                if self._ol_counter:
+                    self._ol_counter[-1] += 1
+                    self._emit(f"\n{indent}{self._ol_counter[-1]}. ")
+                else:
+                    self._emit(f"\n{indent}1. ")
+            else:
+                self._emit(f"\n{indent}* ")
+
+        elif tag == "pre":
+            self._pre_parts = []
+            self._in_pre = True
+
+        elif tag == "code" and not self._in_pre:
+            self._in_inline_code = True
+            self._emit("`")
+
+        elif tag == "table":
+            self._in_table = True
+            self._header_row_done = False
+            self._is_first_row = True
+            self._emit("\n\n")
+
+        elif tag == "tr":
+            # Flush any open cell/row from a previous row that may
+            # have omitted its optional </td> or </tr> end tags.
+            self._finish_cell()
+            self._finish_row()
+
+        elif tag in ("th", "td"):
+            # Flush any open cell (handles omitted </td>/<th>)
+            self._finish_cell()
+            self._cell_parts = []
+            self._in_cell = True
+            if tag == "th":
+                self._row_has_th = True
+
+        elif tag == "img":
+            # Skip images -- keeps fetched page text focused on readable
+            # content and avoids data-URI amplification.
+            return
+
+    def handle_endtag(self, tag: str) -> None:
+        tag = tag.lower()
+
+        if tag in _SKIP_TAGS:
+            self._skip_depth = max(0, self._skip_depth - 1)
+            return
+        if self._skip_depth:
+            return
+
+        if tag in _HEADING_TAGS:
+            self._emit("\n\n")
+
+        elif tag == "a":
+            self._finish_link()
+
+        elif tag in _INLINE_EMPHASIS:
+            self._emit(_INLINE_EMPHASIS[tag])
+
+        elif tag in _BLOCK_TAGS:
+            self._emit("\n\n")
+
+        elif tag == "blockquote":
+            if self._bq_stack:
+                content = "".join(self._bq_stack.pop())
+                prefixed = self._prefix_blockquote(content)
+                if prefixed:
+                    self._emit("\n\n" + prefixed + "\n\n")
+
+        elif tag == "ul":
+            if self._list_stack and self._list_stack[-1] == "ul":
+                self._list_stack.pop()
+            self._emit("\n")
+
+        elif tag == "ol":
+            if self._list_stack and self._list_stack[-1] == "ol":
+                self._list_stack.pop()
+                if self._ol_counter:
+                    self._ol_counter.pop()
+            self._emit("\n")
+
+        elif tag == "pre":
+            raw = "".join(self._pre_parts)
+            self._in_pre = False
+            block = "```\n" + raw + "\n```"
+            self._emit("\n\n" + block + "\n\n")
+
+        elif tag == "code" and not self._in_pre:
+            self._in_inline_code = False
+            self._emit("`")
+
+        elif tag in ("th", "td"):
+            self._finish_cell()
+
+        elif tag == "tr":
+            self._finish_cell()
+            self._finish_row()
+
+        elif tag == "table":
+            # Flush any remaining row (handles omitted </tr>)
+            self._finish_cell()
+            self._finish_row()
+            self._in_table = False
+            self._emit("\n")
+
+    # ------------------------------------------------------------------
+    # Text / entity handlers
+    # ------------------------------------------------------------------
+    def handle_data(self, data: str) -> None:
+        if self._skip_depth:
+            return
+        if self._in_pre:
+            self._pre_parts.append(data)
+            return
+        # Preserve literal whitespace inside inline <code> spans
+        if self._in_inline_code:
+            self._emit(data)
+            return
+        # Collapse all whitespace (including newlines) per HTML rules
+        text = re.sub(r"\s+", " ", data)
+        # Suppress whitespace-only text nodes between table structural
+        # elements (indentation from source HTML) to prevent leading
+        # spaces from breaking Markdown table row alignment.
+        if self._in_table and not self._in_cell and not text.strip():
+            return
+        self._emit(text)
+
+    def handle_entityref(self, name: str) -> None:
+        if self._skip_depth:
+            return
+        self._emit(html.unescape(f"&{name};"))
+
+    def handle_charref(self, name: str) -> None:
+        if self._skip_depth:
+            return
+        self._emit(html.unescape(f"&#{name};"))
+
+    # ------------------------------------------------------------------
+    # Flush pending buffers (handles truncated HTML from capped fetches)
+    # ------------------------------------------------------------------
+    def flush_pending(self) -> None:
+        """Flush any open side-buffers into ``_out``.
+
+        Called after ``close()`` to recover content from truncated HTML
+        where closing tags were never seen (common when ``_fetch_page_text``
+        caps the download by byte count).
+        """
+        # Flush innermost buffers first so their content propagates outward.
+
+        if self._in_link:
+            self._finish_link()
+
+        if self._in_inline_code:
+            self._in_inline_code = False
+            self._emit("`")
+
+        self._finish_cell()
+        self._finish_row()
+
+        if self._in_pre:
+            raw = "".join(self._pre_parts)
+            self._in_pre = False
+            block = "```\n" + raw + "\n```"
+            self._emit("\n\n" + block + "\n\n")
+
+        # Flatten any open blockquote buffers (innermost first)
+        while self._bq_stack:
+            content = "".join(self._bq_stack.pop())
+            prefixed = self._prefix_blockquote(content)
+            if not prefixed:
+                continue
+            if self._bq_stack:
+                self._bq_stack[-1].append("\n\n" + prefixed + "\n\n")
+            else:
+                self._out.append("\n\n" + prefixed + "\n\n")
+
+
+# ------------------------------------------------------------------
+# Post-processing
+# ------------------------------------------------------------------
+def _cleanup(text: str) -> str:
+    """Normalize whitespace and blank lines in the final output.
+
+    Preserves content inside fenced code blocks verbatim so that
+    intentional blank lines in ``<pre>`` content are not collapsed.
+    """
+    lines = text.split("\n")
+    out: list[str] = []
+    in_fence = False
+    blank_run = 0
+
+    for line in lines:
+        stripped = line.rstrip(" \t")
+        if stripped.startswith("```"):
+            in_fence = not in_fence
+            blank_run = 0
+            out.append(stripped)
+            continue
+
+        if in_fence:
+            # Preserve code block content exactly as-is
+            out.append(line)
+            continue
+
+        if not stripped:
+            blank_run += 1
+            if blank_run <= 1:
+                out.append("")
+            continue
+
+        blank_run = 0
+        out.append(stripped)
+
+    return "\n".join(out).strip()
+
+
+# ------------------------------------------------------------------
+# Public API
+# ------------------------------------------------------------------
+def html_to_markdown(source_html: str) -> str:
+    """Convert an HTML string to Markdown.
+
+    Handles headings, links, bold/italic, lists (ordered and unordered),
+    tables, blockquotes, code blocks, and HTML entities.  ``<script>``,
+    ``<style>``, and ``<head>`` sections are stripped entirely.
+    """
+    # Normalize line endings before parsing
+    source_html = source_html.replace("\r\n", "\n").replace("\r", "\n")
+    renderer = _MarkdownRenderer()
+    renderer.feed(source_html)
+    renderer.close()
+    renderer.flush_pending()
+    raw = "".join(renderer._out)
+    return _cleanup(raw)

studio/backend/core/inference/inference.py

@@ -18,7 +18,14 @@ from typing import Optional, Union, Generator, Tuple
 from utils.models import ModelConfig, get_base_model_from_lora
 from utils.paths import is_model_cached
 from utils.utils import format_error_message
-from utils.hardware import get_device, clear_gpu_cache, log_gpu_memory
+from utils.hardware import (
+    get_device,
+    clear_gpu_cache,
+    log_gpu_memory,
+    get_device_map,
+    raise_if_offloaded,
+    get_visible_gpu_count,
+)
 from core.inference.audio_codecs import AudioCodecManager
 from io import StringIO
 import structlog
@@ -241,6 +248,7 @@ class InferenceBackend:
         load_in_4bit: bool = True,
         hf_token: Optional[str] = None,
         trust_remote_code: bool = False,
+        gpu_ids: Optional[list[int]] = None,
     ) -> bool:
         """
         Load any model: base, LoRA adapter, text, or vision.
@@ -260,6 +268,10 @@ class InferenceBackend:
                 return False
 
             self.loading_models.add(model_name)
+            device_map = get_device_map(gpu_ids)
+            logger.info(
+                f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)"
+            )
 
             self.models[model_name] = {
                 "is_vision": config.is_vision,
@@ -290,6 +302,7 @@ class InferenceBackend:
                         config.path,
                         auto_model = CsmForConditionalGeneration,
                         load_in_4bit = False,
+                        device_map = device_map,
                         token = hf_token if hf_token and hf_token.strip() else None,
                         trust_remote_code = trust_remote_code,
                     )
@@ -325,6 +338,7 @@ class InferenceBackend:
                             config.path,
                             dtype = torch.float32,
                             load_in_4bit = False,
+                            device_map = device_map,
                             token = hf_token if hf_token and hf_token.strip() else None,
                             trust_remote_code = trust_remote_code,
                         )
@@ -345,6 +359,7 @@ class InferenceBackend:
                             llm_path,
                             dtype = torch.float32,
                             load_in_4bit = False,
+                            device_map = device_map,
                             token = hf_token if hf_token and hf_token.strip() else None,
                             trust_remote_code = trust_remote_code,
                         )
@@ -361,6 +376,7 @@ class InferenceBackend:
                         config.path,
                         max_seq_length = max_seq_length,
                         load_in_4bit = False,
+                        device_map = device_map,
                         token = hf_token if hf_token and hf_token.strip() else None,
                         trust_remote_code = trust_remote_code,
                     )
@@ -378,6 +394,7 @@ class InferenceBackend:
                         whisper_language = "English",
                         whisper_task = "transcribe",
                         load_in_4bit = False,
+                        device_map = device_map,
                         token = hf_token if hf_token and hf_token.strip() else None,
                         trust_remote_code = trust_remote_code,
                     )
@@ -405,6 +422,7 @@ class InferenceBackend:
                         model_name = config.path,
                         max_seq_length = max_seq_length,
                         load_in_4bit = False,
+                        device_map = device_map,
                         token = hf_token if hf_token and hf_token.strip() else None,
                         trust_remote_code = trust_remote_code,
                     )
@@ -420,6 +438,11 @@ class InferenceBackend:
                         audio_type, self.device, model_repo_path = model_repo_path
                     )
 
+                # Reject CPU/disk offload for audio models too
+                raise_if_offloaded(
+                    self.models[model_name]["model"], device_map, "Inference"
+                )
+
                 self.active_model_name = model_name
                 self.loading_models.discard(model_name)
                 logger.info(f"Successfully loaded audio model: {model_name}")
@@ -441,6 +464,7 @@ class InferenceBackend:
                     max_seq_length = max_seq_length,
                     dtype = dtype,
                     load_in_4bit = load_in_4bit,
+                    device_map = device_map,
                     token = hf_token if hf_token and hf_token.strip() else None,
                     trust_remote_code = trust_remote_code,
                 )
@@ -497,6 +521,7 @@ class InferenceBackend:
                     max_seq_length = max_seq_length,
                     dtype = dtype,
                     load_in_4bit = load_in_4bit,
+                    device_map = device_map,
                     token = hf_token if hf_token and hf_token.strip() else None,
                     trust_remote_code = trust_remote_code,
                 )
@@ -507,6 +532,10 @@ class InferenceBackend:
                 self.models[model_name]["model"] = model
                 self.models[model_name]["tokenizer"] = tokenizer
 
+            raise_if_offloaded(
+                self.models[model_name]["model"], device_map, "Inference"
+            )
+
             # Load chat template info
             self._load_chat_template_info(model_name)
 
@@ -615,6 +644,7 @@ class InferenceBackend:
         dtype = None,
         load_in_4bit: bool = True,
         hf_token: Optional[str] = None,
+        gpu_ids: Optional[list[int]] = None,
     ) -> Tuple[bool, Optional[str], Optional[str]]:
         """
         Final Corrected Version:
@@ -639,7 +669,12 @@ class InferenceBackend:
                     base_model_name, None, is_lora = False
                 )
                 if not self.load_model(
-                    base_config, max_seq_length, dtype, load_in_4bit, hf_token
+                    base_config,
+                    max_seq_length,
+                    dtype,
+                    load_in_4bit,
+                    hf_token,
+                    gpu_ids = gpu_ids,
                 ):
                     return False, None, None
 
@@ -1037,12 +1072,12 @@ class InferenceBackend:
                 input_text,
                 add_special_tokens = False,
                 return_tensors = "pt",
-            ).to(self.device)
+            ).to(model.device)
         else:
             # Text-only for vision model
             formatted_prompt = self.format_chat_prompt(messages, system_prompt)
             inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(
-                self.device
+                model.device
             )
 
         # Stream with TextIteratorStreamer + background thread
@@ -1182,7 +1217,7 @@ class InferenceBackend:
             return_dict = True,
             return_tensors = "pt",
             truncation = False,
-        ).to(self.device)
+        ).to(model.device)
 
         try:
             from transformers import TextIteratorStreamer

studio/backend/core/inference/llama_cpp.py

@@ -11,6 +11,7 @@ through its OpenAI-compatible /v1/chat/completions endpoint.
 import atexit
 import contextlib
 import json
+import re
 import struct
 import structlog
 from loggers import get_logger
@@ -48,6 +49,7 @@ class LlamaCppBackend:
         self._healthy = False
         self._context_length: Optional[int] = None
         self._effective_context_length: Optional[int] = None
+        self._max_context_length: Optional[int] = None
         self._chat_template: Optional[str] = None
         self._supports_reasoning: bool = False
         self._reasoning_always_on: bool = False
@@ -100,6 +102,11 @@ class LlamaCppBackend:
         """Return the effective context length the server is running at."""
         return self._effective_context_length or self._context_length
 
+    @property
+    def max_context_length(self) -> Optional[int]:
+        """Return the maximum context currently available on this hardware."""
+        return self._max_context_length or self._context_length
+
     @property
     def chat_template(self) -> Optional[str]:
         return self._chat_template
@@ -287,7 +294,8 @@ class LlamaCppBackend:
                         continue
                     gpus.append((idx, free_mib))
             return gpus
-        except Exception:
+        except Exception as e:
+            logger.debug(f"Failed to query GPU free memory via nvidia-smi: {e}")
             return []
 
     @staticmethod
@@ -328,6 +336,11 @@ class LlamaCppBackend:
                 return sorted(selected), False
 
         # Model is too large even for all GPUs, let --fit handle it
+        logger.debug(
+            "Model does not fit in available GPU memory, falling back to --fit",
+            model_size_mib = round(model_size_mib, 2),
+            ranked_gpus = ranked,
+        )
         return None, True
 
     # ── KV cache VRAM estimation ─────────────────────────────────────
@@ -386,6 +399,11 @@ class LlamaCppBackend:
         If the model weights alone don't fit, returns min_ctx unchanged.
         """
         if not self._can_estimate_kv():
+            logger.debug(
+                "Skipping context fit because KV cache metadata is unavailable",
+                requested_ctx = requested_ctx,
+                available_mib = available_mib,
+            )
             return requested_ctx
 
         budget_bytes = available_mib * 1024 * 1024 * 0.70
@@ -399,6 +417,12 @@ class LlamaCppBackend:
         # Model weights alone exceed budget -- can't help by reducing ctx.
         # Return requested_ctx unchanged; --fit will handle VRAM management.
         if model_footprint >= budget_bytes:
+            logger.debug(
+                "Model footprint exceeds GPU budget before KV cache",
+                requested_ctx = requested_ctx,
+                available_mib = available_mib,
+                model_size_gb = round(model_footprint / (1024**3), 2),
+            )
             return requested_ctx
 
         # Binary search for max context that fits
@@ -960,7 +984,11 @@ class LlamaCppBackend:
 
             self._port = self._find_free_port()
 
-            # Select GPU(s) based on model size + estimated KV cache
+            # Select GPU(s) based on model size + estimated KV cache.
+            # Seed safe defaults before GPU probing so the except path
+            # still has valid state to publish.
+            effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0)
+            max_available_ctx = self._context_length or effective_ctx
             try:
                 model_size = self._get_gguf_size_bytes(model_path)
                 gpus = self._get_gpu_free_memory()
@@ -975,6 +1003,9 @@ class LlamaCppBackend:
                 else:
                     effective_ctx = 0
                 original_ctx = effective_ctx
+                # Default UI ceiling to the model's native context length.
+                # GPU/VRAM-fit logic below may shrink this if hardware is limited.
+                max_available_ctx = self._context_length or effective_ctx
 
                 # Auto-cap context to fit in GPU VRAM and select GPUs.
                 #
@@ -993,6 +1024,29 @@ class LlamaCppBackend:
                 explicit_ctx = n_ctx > 0
 
                 if gpus and self._can_estimate_kv() and effective_ctx > 0:
+                    # Compute the largest hardware-aware cap from the model's
+                    # native context across all usable GPU subsets (for UI
+                    # bounds), independent of the currently requested context.
+                    native_ctx_for_cap = self._context_length or effective_ctx
+                    if native_ctx_for_cap > 0:
+                        ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True)
+                        best_cap = 0
+                        for n_gpus in range(1, len(ranked_for_cap) + 1):
+                            subset = ranked_for_cap[:n_gpus]
+                            pool_mib = sum(free for _, free in subset)
+                            capped = self._fit_context_to_vram(
+                                native_ctx_for_cap,
+                                pool_mib,
+                                model_size,
+                                cache_type_kv,
+                            )
+                            kv = self._estimate_kv_cache_bytes(capped, cache_type_kv)
+                            total_mib = (model_size + kv) / (1024 * 1024)
+                            if total_mib <= pool_mib * 0.70:
+                                best_cap = max(best_cap, capped)
+                        if best_cap > 0:
+                            max_available_ctx = best_cap
+
                     if explicit_ctx:
                         # Try to honor the user's requested context exactly.
                         requested_total = model_size + self._estimate_kv_cache_bytes(
@@ -1043,7 +1097,13 @@ class LlamaCppBackend:
                                 break
 
                 elif gpus:
-                    # Can't estimate KV -- fall back to file-size-only check
+                    # Can't estimate KV -- fall back to file-size-only check.
+                    # Without KV estimation we cannot prove a hardware cap, so
+                    # keep the ceiling at the native context (already the default).
+                    logger.debug(
+                        "Falling back to file-size-only GPU selection",
+                        model_size_gb = round(model_size / (1024**3), 2),
+                    )
                     gpu_indices, use_fit = self._select_gpus(model_size, gpus)
 
                 if effective_ctx < original_ctx:
@@ -1313,6 +1373,11 @@ class LlamaCppBackend:
             self._effective_context_length = (
                 effective_ctx if effective_ctx > 0 else self._context_length
             )
+            self._max_context_length = (
+                max_available_ctx
+                if max_available_ctx > 0
+                else self._effective_context_length
+            )
 
             # Wait for llama-server to become healthy
             if not self._wait_for_health(timeout = 600.0):
@@ -1347,6 +1412,7 @@ class LlamaCppBackend:
             self._healthy = False
             self._context_length = None
             self._effective_context_length = None
+            self._max_context_length = None
             self._chat_template = None
             self._supports_reasoning = False
             self._reasoning_always_on = False
@@ -2055,7 +2121,7 @@ class LlamaCppBackend:
         stop: Optional[list[str]] = None,
         cancel_event: Optional[threading.Event] = None,
         enable_thinking: Optional[bool] = None,
-        max_tool_iterations: int = 10,
+        max_tool_iterations: int = 25,
         auto_heal_tool_calls: bool = True,
         tool_call_timeout: int = 300,
         session_id: Optional[str] = None,
@@ -2107,6 +2173,13 @@ class LlamaCppBackend:
         )
         _MAX_BUFFER_CHARS = 32
 
+        # ── Duplicate tool-call detection ────────────────────────
+        # Track recent (tool_name, arguments) hashes to detect loops
+        # where the model repeats the exact same call.  Retries after
+        # a transient failure are allowed (only block when the previous
+        # identical call succeeded).
+        _tool_call_history: list[tuple[str, bool]] = []  # (key, failed)
+
         for iteration in range(max_tool_iterations):
             if cancel_event is not None and cancel_event.is_set():
                 return
@@ -2503,6 +2576,11 @@ class LlamaCppBackend:
                         # Merge accumulated metrics from prior tool
                         # iterations so they are not silently dropped.
                         yield {"type": "status", "text": ""}
+                        if content_accum:
+                            # Strip leaked tool-call XML before yielding
+                            content_accum = _strip_tool_markup(
+                                content_accum, final = True
+                            )
                         if content_accum:
                             yield {"type": "content", "text": content_accum}
                         _fu = _iter_usage or {}
@@ -2596,16 +2674,32 @@ class LlamaCppBackend:
                         "arguments": arguments,
                     }
 
-                    _effective_timeout = (
-                        None if tool_call_timeout >= 9999 else tool_call_timeout
-                    )
-                    result = execute_tool(
-                        tool_name,
-                        arguments,
-                        cancel_event = cancel_event,
-                        timeout = _effective_timeout,
-                        session_id = session_id,
-                    )
+                    # ── Duplicate call detection ──────────────
+                    # str(dict) is stable here: arguments always comes from
+                    # json.loads on the same model output within one request,
+                    # so insertion order is deterministic (Python 3.7+).
+                    _tc_key = tool_name + str(arguments)
+                    _prev = _tool_call_history[-1] if _tool_call_history else None
+                    if _prev and _prev[0] == _tc_key and not _prev[1]:
+                        result = (
+                            "You already made this exact call. "
+                            "Do not repeat the same tool call. "
+                            "Try a different approach: fetch a URL "
+                            "from previous results, use Python to "
+                            "process data you already have, or "
+                            "provide your final answer now."
+                        )
+                    else:
+                        _effective_timeout = (
+                            None if tool_call_timeout >= 9999 else tool_call_timeout
+                        )
+                        result = execute_tool(
+                            tool_name,
+                            arguments,
+                            cancel_event = cancel_event,
+                            timeout = _effective_timeout,
+                            session_id = session_id,
+                        )
 
                     yield {
                         "type": "tool_end",
@@ -2614,10 +2708,32 @@ class LlamaCppBackend:
                         "result": result,
                     }
 
+                    # Nudge model to try a different approach on errors
+                    _error_prefixes = (
+                        "Error",
+                        "Search failed",
+                        "Execution error",
+                        "Blocked:",
+                        "Exit code",
+                        "Failed to fetch",
+                        "Failed to resolve",
+                        "No query provided",
+                    )
+                    _is_error = isinstance(result, str) and result.lstrip().startswith(
+                        _error_prefixes
+                    )
+                    _tool_call_history.append((_tc_key, _is_error))
+                    _result_content = result
+                    if _is_error:
+                        _result_content = (
+                            result + "\n\nThe tool call encountered an issue. "
+                            "Please try a different approach or rephrase your request."
+                        )
+
                     tool_msg = {
                         "role": "tool",
                         "name": tool_name,
-                        "content": result,
+                        "content": _result_content,
                     }
                     tool_call_id = tc.get("id")
                     if tool_call_id:
@@ -2634,6 +2750,22 @@ class LlamaCppBackend:
                     return
                 raise
 
+        # ── Tool iteration cap reached -- synthesize final answer ──
+        # The model used all iterations without producing a final text
+        # response. Inject a nudge so the final streaming pass produces
+        # a useful answer instead of continuing to request tools.
+        if max_tool_iterations > 0:
+            conversation.append(
+                {
+                    "role": "user",
+                    "content": (
+                        "You have used all available tool calls. Based on "
+                        "everything you have found so far, provide your final "
+                        "answer now. Do not call any more tools."
+                    ),
+                }
+            )
+
         # Clear status
         yield {"type": "status", "text": ""}
 

studio/backend/core/inference/orchestrator.py

@@ -17,6 +17,7 @@ Pattern follows core/training/training.py.
 
 import atexit
 import base64
+import os
 import structlog
 from loggers import get_logger
 import multiprocessing as mp
@@ -27,11 +28,17 @@ import uuid
 from io import BytesIO
 from pathlib import Path
 from typing import Any, Generator, Optional, Tuple, Union
+from utils.hardware import prepare_gpu_selection
 
 logger = get_logger(__name__)
 
 _CTX = mp.get_context("spawn")
 
+
+class DownloadStallError(RuntimeError):
+    """Raised when the worker reports no download progress for too long."""
+
+
 # Dispatcher timeout constants (seconds)
 _DISPATCH_READ_TIMEOUT = 30.0
 _DISPATCH_POLL_INTERVAL = 0.5
@@ -262,12 +269,17 @@ class InferenceOrchestrator:
         except (EOFError, OSError, ValueError):
             return None
 
-    def _wait_response(self, expected_type: str, timeout: float = 120.0) -> dict:
+    def _wait_response(self, expected_type: str, timeout: float = 300.0) -> dict:
         """Block until a response of the expected type arrives.
 
         Also handles 'status' and 'error' events during the wait.
         Returns the matching response dict.
         Raises RuntimeError on timeout or subprocess crash.
+
+        The *timeout* is an **inactivity** timeout: it resets whenever the
+        subprocess sends a status message, so long-running operations (large
+        downloads, slow model loads) won't be killed as long as the subprocess
+        keeps reporting progress.
         """
         deadline = time.monotonic() + timeout
 
@@ -292,8 +304,15 @@ class InferenceOrchestrator:
 
             if rtype == "status":
                 logger.info("Subprocess status: %s", resp.get("message", ""))
+                # Reset deadline — subprocess is still alive and working
+                deadline = time.monotonic() + timeout
                 continue
 
+            if rtype == "stall":
+                msg = resp.get("message", "Download stalled")
+                logger.warning("Subprocess reported stall: %s", msg)
+                raise DownloadStallError(msg)
+
             # Other response types during wait — skip
             logger.debug(
                 "Skipping response type '%s' while waiting for '%s'",
@@ -302,7 +321,8 @@ class InferenceOrchestrator:
             )
 
         raise RuntimeError(
-            f"Timeout waiting for '{expected_type}' response after {timeout}s"
+            f"Timeout waiting for '{expected_type}' response "
+            f"(no activity for {timeout}s)"
         )
 
     def _drain_queue(self) -> list:
@@ -571,6 +591,7 @@ class InferenceOrchestrator:
         load_in_4bit: bool = True,
         hf_token: Optional[str] = None,
         trust_remote_code: bool = False,
+        gpu_ids: Optional[list[int]] = None,
     ) -> bool:
         """Load a model for inference.
 
@@ -594,7 +615,16 @@ class InferenceOrchestrator:
                 "hf_token": hf_token or "",
                 "gguf_variant": getattr(config, "gguf_variant", None),
                 "trust_remote_code": trust_remote_code,
+                "gpu_ids": gpu_ids,
             }
+            resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
+                gpu_ids,
+                model_name = model_name,
+                hf_token = hf_token,
+                load_in_4bit = load_in_4bit,
+            )
+            sub_config["resolved_gpu_ids"] = resolved_gpu_ids
+            sub_config["gpu_selection"] = gpu_selection
 
             # Always kill existing subprocess and spawn fresh.
             # Reusing a subprocess after unsloth patches torch internals
@@ -608,36 +638,66 @@ class InferenceOrchestrator:
                 # Dead subprocess — clean up
                 self._shutdown_subprocess(timeout = 2)
 
-            logger.info(
-                "Spawning fresh inference subprocess for '%s' (transformers %s.x)",
-                model_name,
-                needed_major,
+            disable_xet = sub_config.get("disable_xet", False) or (
+                os.environ.get("HF_HUB_DISABLE_XET") == "1"
             )
-            self._spawn_subprocess(sub_config)
-            resp = self._wait_response("loaded", timeout = 180)
 
-            # Update local state from response
-            if resp.get("success"):
-                self._current_transformers_major = needed_major
-                model_info = resp.get("model_info", {})
-                self.active_model_name = model_info.get("identifier", model_name)
-                self.models[self.active_model_name] = {
-                    "is_vision": model_info.get("is_vision", False),
-                    "is_lora": model_info.get("is_lora", False),
-                    "display_name": model_info.get("display_name", model_name),
-                    "is_audio": model_info.get("is_audio", False),
-                    "audio_type": model_info.get("audio_type"),
-                    "has_audio_input": model_info.get("has_audio_input", False),
-                }
-                self.loading_models.discard(model_name)
-                logger.info("Model '%s' loaded successfully in subprocess", model_name)
-                return True
-            else:
-                error = resp.get("error", "Failed to load model")
-                self.loading_models.discard(model_name)
-                self.active_model_name = None
-                self.models.clear()
-                raise Exception(error)
+            for attempt in range(2):
+                logger.info(
+                    "Spawning fresh inference subprocess for '%s' "
+                    "(transformers %s.x, attempt %d/2%s)",
+                    model_name,
+                    needed_major,
+                    attempt + 1,
+                    ", xet disabled" if disable_xet else "",
+                )
+                sub_config["disable_xet"] = disable_xet
+                self._spawn_subprocess(sub_config)
+
+                try:
+                    resp = self._wait_response("loaded")
+                except DownloadStallError:
+                    # First stall and Xet was enabled -> retry with Xet disabled
+                    if attempt == 0 and not disable_xet:
+                        logger.warning(
+                            "Download stalled for '%s' -- retrying with "
+                            "HF_HUB_DISABLE_XET=1",
+                            model_name,
+                        )
+                        self._shutdown_subprocess(timeout = 5)
+                        disable_xet = True
+                        continue
+                    # Second stall (or already had xet disabled) -> give up
+                    self._shutdown_subprocess(timeout = 5)
+                    raise RuntimeError(
+                        f"Download stalled for '{model_name}' even with "
+                        f"HF_HUB_DISABLE_XET=1 -- check your network connection"
+                    )
+
+                # Got a response — check success
+                if resp.get("success"):
+                    self._current_transformers_major = needed_major
+                    model_info = resp.get("model_info", {})
+                    self.active_model_name = model_info.get("identifier", model_name)
+                    self.models[self.active_model_name] = {
+                        "is_vision": model_info.get("is_vision", False),
+                        "is_lora": model_info.get("is_lora", False),
+                        "display_name": model_info.get("display_name", model_name),
+                        "is_audio": model_info.get("is_audio", False),
+                        "audio_type": model_info.get("audio_type"),
+                        "has_audio_input": model_info.get("has_audio_input", False),
+                    }
+                    self.loading_models.discard(model_name)
+                    logger.info(
+                        "Model '%s' loaded successfully in subprocess", model_name
+                    )
+                    return True
+                else:
+                    error = resp.get("error", "Failed to load model")
+                    self.loading_models.discard(model_name)
+                    self.active_model_name = None
+                    self.models.clear()
+                    raise Exception(error)
 
         except Exception:
             self.loading_models.discard(model_name)
@@ -661,7 +721,7 @@ class InferenceOrchestrator:
                     "model_name": model_name,
                 }
             )
-            resp = self._wait_response("unloaded", timeout = 30)
+            resp = self._wait_response("unloaded")
 
             # Update local state
             self.models.pop(model_name, None)

studio/backend/core/inference/tools.py

@@ -57,16 +57,23 @@ WEB_SEARCH_TOOL = {
     "type": "function",
     "function": {
         "name": "web_search",
-        "description": "Search the web for current information, recent events, or facts you are uncertain about.",
+        "description": (
+            "Search the web and fetch page content. Returns snippets for all results. "
+            "Use the url parameter to fetch full page text from a specific URL."
+        ),
         "parameters": {
             "type": "object",
             "properties": {
                 "query": {
                     "type": "string",
                     "description": "The search query",
-                }
+                },
+                "url": {
+                    "type": "string",
+                    "description": "A URL to fetch full page content from (instead of searching). Use this to read a page found in search results.",
+                },
             },
-            "required": ["query"],
+            "required": [],
         },
     },
 }
@@ -131,7 +138,11 @@ def execute_tool(
     )
     effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout
     if name == "web_search":
-        return _web_search(arguments.get("query", ""), timeout = effective_timeout)
+        return _web_search(
+            arguments.get("query", ""),
+            url = arguments.get("url"),
+            timeout = effective_timeout,
+        )
     if name == "python":
         return _python_exec(
             arguments.get("code", ""), cancel_event, effective_timeout, session_id
@@ -143,9 +154,161 @@ def execute_tool(
     return f"Unknown tool: {name}"
 
 
-def _web_search(query: str, max_results: int = 5, timeout: int = _EXEC_TIMEOUT) -> str:
-    """Search the web using DuckDuckGo and return formatted results."""
-    if not query.strip():
+_MAX_PAGE_CHARS = 16000  # limit fetched page text
+_MAX_FETCH_BYTES = _MAX_PAGE_CHARS * 4 + 1  # cap raw download size
+
+
+def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str]:
+    """Resolve *hostname*, reject non-public IPs, return a pinned IP string.
+
+    Returns ``(ok, reason_or_empty, resolved_ip)``.  The caller should
+    connect to *resolved_ip* (with a ``Host`` header) to prevent DNS
+    rebinding between validation and the actual fetch.
+    """
+    import ipaddress
+    import socket
+
+    try:
+        infos = socket.getaddrinfo(hostname, port, type = socket.SOCK_STREAM)
+    except OSError as e:
+        return False, f"Failed to resolve host: {e}", ""
+
+    if not infos:
+        return False, f"Failed to resolve host: no addresses for {hostname!r}", ""
+
+    for *_, sockaddr in infos:
+        ip = ipaddress.ip_address(sockaddr[0])
+        if (
+            ip.is_private
+            or ip.is_loopback
+            or ip.is_link_local
+            or ip.is_multicast
+            or ip.is_reserved
+            or ip.is_unspecified
+        ):
+            return False, f"Blocked: refusing to fetch non-public address {ip}.", ""
+
+    # Return the first resolved address for pinning
+    first_ip = infos[0][4][0]
+    return True, "", first_ip
+
+
+def _fetch_page_text(
+    url: str, max_chars: int = _MAX_PAGE_CHARS, timeout: int = 30
+) -> str:
+    """Fetch a URL and return plain text content (HTML tags stripped).
+
+    Blocks private/loopback/link-local targets (SSRF protection) and caps
+    the download size to avoid unbounded memory usage.
+    """
+    from urllib.parse import urlparse
+
+    parsed = urlparse(url)
+    if parsed.scheme not in ("http", "https"):
+        return f"Blocked: only http/https URLs are allowed (got {parsed.scheme!r})."
+    if not parsed.hostname:
+        return "Blocked: URL is missing a hostname."
+
+    port = parsed.port or (443 if parsed.scheme == "https" else 80)
+    ok, reason, pinned_ip = _validate_and_resolve_host(parsed.hostname, port)
+    if not ok:
+        return reason
+
+    try:
+        import urllib.request
+        from urllib.error import HTTPError as _HTTPError
+        from urllib.parse import urljoin, urlunparse
+
+        # Disable auto-redirect so we can validate each hop for SSRF.
+        # urllib raises HTTPError for 3xx when the handler returns None,
+        # so we catch that and extract the Location header manually.
+        class _NoRedirect(urllib.request.HTTPRedirectHandler):
+            def redirect_request(self, req, fp, code, msg, headers, newurl):
+                return None
+
+        opener = urllib.request.build_opener(_NoRedirect)
+        max_bytes = max_chars * 4 + 1
+        current_url = url
+        current_host = parsed.hostname
+
+        for _hop in range(5):
+            # Pin to the validated IP to prevent DNS rebinding.
+            # Rewrite the URL to use the IP and set the Host header.
+            cp = urlparse(current_url)
+            ip_netloc = f"{pinned_ip}:{cp.port}" if cp.port else pinned_ip
+            pinned_url = urlunparse(cp._replace(netloc = ip_netloc))
+
+            req = urllib.request.Request(
+                pinned_url,
+                headers = {
+                    "User-Agent": "UnslothStudio/1.0",
+                    "Host": current_host,
+                },
+            )
+            try:
+                resp = opener.open(req, timeout = timeout)
+            except _HTTPError as e:
+                if e.code not in (301, 302, 303, 307, 308):
+                    return (
+                        f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}"
+                    )
+                location = e.headers.get("Location")
+                if not location:
+                    return "Failed to fetch URL: redirect missing Location header."
+                current_url = urljoin(current_url, location)
+                rp = urlparse(current_url)
+                if rp.scheme not in ("http", "https") or not rp.hostname:
+                    return "Blocked: redirect target is not a valid http/https URL."
+                rp_port = rp.port or (443 if rp.scheme == "https" else 80)
+                ok2, reason2, pinned_ip = _validate_and_resolve_host(
+                    rp.hostname,
+                    rp_port,
+                )
+                if not ok2:
+                    return reason2
+                current_host = rp.hostname
+                continue
+            # Success -- read capped body
+            raw_bytes = resp.read(max_bytes)
+            break
+        else:
+            return "Failed to fetch URL: too many redirects."
+
+        charset = resp.headers.get_content_charset() or "utf-8"
+        raw_html = raw_bytes.decode(charset, errors = "replace")
+    except _HTTPError as e:
+        return f"Failed to fetch URL: HTTP {e.code} {getattr(e, 'reason', '')}"
+    except Exception as e:
+        return f"Failed to fetch URL: {e}"
+
+    # Convert HTML to Markdown using the builtin converter (no external deps)
+    from ._html_to_md import html_to_markdown
+
+    text = html_to_markdown(raw_html)
+
+    if not text:
+        return "(page returned no readable text)"
+    if len(text) > max_chars:
+        text = text[:max_chars] + f"\n\n... (truncated, {len(text)} chars total)"
+    return text
+
+
+def _web_search(
+    query: str,
+    max_results: int = 5,
+    timeout: int = _EXEC_TIMEOUT,
+    url: str | None = None,
+) -> str:
+    """Search the web using DuckDuckGo and return formatted results.
+
+    If ``url`` is provided, fetches that page directly instead of searching.
+    """
+    # Direct URL fetch mode
+    if url and url.strip():
+        fetch_timeout = 60 if timeout is None else min(timeout, 60)
+        return _fetch_page_text(url.strip(), timeout = fetch_timeout)
+
+    if not query or not query.strip():
         return "No query provided."
     try:
         from ddgs import DDGS
@@ -160,7 +323,13 @@ def _web_search(query: str, max_results: int = 5, timeout: int = _EXEC_TIMEOUT)
                 f"URL: {r.get('href', '')}\n"
                 f"Snippet: {r.get('body', '')}"
             )
-        return "\n\n---\n\n".join(parts)
+        text = "\n\n---\n\n".join(parts)
+        text += (
+            "\n\n---\n\nIMPORTANT: These are only short snippets. "
+            "To get the full page content, call web_search with "
+            'the url parameter (e.g. {"url": "<URL>"}).'
+        )
+        return text
     except Exception as e:
         return f"Search failed: {e}"
 

studio/backend/core/inference/worker.py

@@ -22,6 +22,7 @@ from loggers import get_logger
 import os
 import queue as _queue
 import sys
+import threading
 import time
 import traceback
 from io import BytesIO
@@ -29,6 +30,7 @@ from pathlib import Path
 from typing import Any
 
 logger = get_logger(__name__)
+from utils.hardware import apply_gpu_ids
 
 
 def _activate_transformers_version(model_name: str) -> None:
@@ -113,6 +115,154 @@ def _build_model_config(config: dict):
     return mc
 
 
+def _get_hf_download_state(
+    model_names: list[str] | None = None,
+) -> tuple[int, bool] | None:
+    """Return (total_bytes, has_incomplete) for the HF Hub cache, or None on error.
+
+    When *model_names* is provided, only those models' ``blobs/``
+    directories are checked instead of scanning every cached model --
+    much faster on systems with many models. Accepts multiple names so
+    that LoRA loads can watch both the adapter repo and the base model
+    repo simultaneously.
+
+    *has_incomplete* is True when any ``*.incomplete`` files exist in the
+    watched blobs directories, indicating that ``huggingface_hub`` is
+    actively downloading.
+
+    Returns None if the state cannot be determined (import error,
+    permission error, etc.) so callers can skip stall logic.
+    """
+    try:
+        from huggingface_hub.constants import HF_HUB_CACHE
+
+        cache = Path(HF_HUB_CACHE)
+        if not cache.exists():
+            return (0, False)
+
+        total = 0
+        has_incomplete = False
+        blobs_dirs: list[Path] = []
+
+        if model_names:
+            for name in model_names:
+                if not name:
+                    continue
+                # Skip local filesystem paths -- HF model IDs use forward
+                # slashes (org/model) but never start with / . ~ or contain
+                # backslashes. This distinguishes them from absolute paths,
+                # relative paths, and Windows paths.
+                if name.startswith(("/", ".", "~")) or "\\" in name:
+                    continue
+                # HF cache dir format: models--org--name (slashes -> --)
+                cache_dir_name = "models--" + name.replace("/", "--")
+                blobs_dir = cache / cache_dir_name / "blobs"
+                if blobs_dir.exists():
+                    blobs_dirs.append(blobs_dir)
+        else:
+            blobs_dirs = list(cache.glob("models--*/blobs"))
+
+        for bdir in blobs_dirs:
+            for f in bdir.iterdir():
+                try:
+                    if f.is_file():
+                        total += f.stat().st_size
+                        if f.name.endswith(".incomplete"):
+                            has_incomplete = True
+                except OSError:
+                    pass
+
+        return (total, has_incomplete)
+    except Exception as e:
+        logger.debug("Failed to determine HF download state: %s", e)
+        return None
+
+
+def _start_heartbeat(
+    resp_queue: Any,
+    interval: float = 30.0,
+    stall_timeout: float = 180.0,
+    xet_disabled: bool = False,
+    model_names: list[str] | None = None,
+) -> threading.Event:
+    """Start a daemon thread that sends periodic status heartbeats.
+
+    Monitors the HF Hub cache directory for download activity. A stall
+    is only reported when ``*.incomplete`` files are present (indicating
+    ``huggingface_hub`` is actively downloading) **and** the total cache
+    size has not changed for *stall_timeout* seconds.
+
+    Once the download finishes (no more ``.incomplete`` files), the stall
+    timer resets, so post-download initialization (quantization, GPU
+    weight loading) is never misclassified as a stalled download.
+
+    Returns a stop event -- set it to terminate the heartbeat thread.
+    """
+    stop = threading.Event()
+    transport = "https" if xet_disabled else "xet"
+
+    def _beat():
+        state = _get_hf_download_state(model_names)
+        last_size = state[0] if state is not None else 0
+        last_change = time.monotonic()
+
+        while not stop.wait(interval):
+            state = _get_hf_download_state(model_names)
+            now = time.monotonic()
+
+            # Skip stall logic if we cannot measure the cache
+            if state is None:
+                _send_response(
+                    resp_queue,
+                    {
+                        "type": "status",
+                        "message": f"Loading model ({transport} transport)...",
+                        "ts": time.time(),
+                    },
+                )
+                continue
+
+            current_size, has_incomplete = state
+
+            if current_size != last_size:
+                last_size = current_size
+                last_change = now
+
+            # Only fire stall when .incomplete files are present,
+            # confirming a download is actively in progress.
+            # Once downloads finish (no .incomplete), reset the timer
+            # so model init time is not counted as a stall.
+            if not has_incomplete:
+                last_change = now
+            elif now - last_change >= stall_timeout:
+                _send_response(
+                    resp_queue,
+                    {
+                        "type": "stall",
+                        "message": (
+                            f"Download appears stalled ({transport} transport) "
+                            f"-- no progress for {int(now - last_change)}s"
+                        ),
+                        "ts": time.time(),
+                    },
+                )
+                # Only fire once -- the orchestrator will kill us
+                return
+
+            _send_response(
+                resp_queue,
+                {
+                    "type": "status",
+                    "message": f"Loading model ({transport} transport)...",
+                    "ts": time.time(),
+                },
+            )
+
+    t = threading.Thread(target = _beat, daemon = True)
+    t.start()
+    return stop
+
+
 def _handle_load(backend, config: dict, resp_queue: Any) -> None:
     """Handle a load command: load a model into the backend."""
     try:
@@ -172,13 +322,34 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
                     model_name,
                 )
 
-        success = backend.load_model(
-            config = mc,
-            max_seq_length = config.get("max_seq_length", 2048),
-            load_in_4bit = load_in_4bit,
-            hf_token = hf_token,
-            trust_remote_code = trust_remote_code,
+        # Send heartbeats every 30s so the orchestrator knows we're still alive
+        # (download / weight loading can take a long time on slow connections)
+        xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1"
+
+        # Watch both the model repo and base model repo (for LoRA loads
+        # where the base model download is the actual bottleneck)
+        watch_repos = [mc.identifier]
+        base = getattr(mc, "base_model", None)
+        if base and str(base) != mc.identifier:
+            watch_repos.append(str(base))
+
+        heartbeat_stop = _start_heartbeat(
+            resp_queue,
+            interval = 30.0,
+            xet_disabled = xet_disabled,
+            model_names = watch_repos,
         )
+        try:
+            success = backend.load_model(
+                config = mc,
+                max_seq_length = config.get("max_seq_length", 2048),
+                load_in_4bit = load_in_4bit,
+                hf_token = hf_token,
+                trust_remote_code = trust_remote_code,
+                gpu_ids = config.get("resolved_gpu_ids"),
+            )
+        finally:
+            heartbeat_stop.set()
 
         if success:
             # Build model_info for the parent to mirror
@@ -490,6 +661,10 @@ def run_inference_process(
         "ignore"  # Suppress warnings at C-level before imports
     )
 
+    if config.get("disable_xet"):
+        os.environ["HF_HUB_DISABLE_XET"] = "1"
+        logger.info("Xet transport disabled (HF_HUB_DISABLE_XET=1)")
+
     import warnings
     from loggers.config import LogConfig
 
@@ -501,6 +676,8 @@ def run_inference_process(
         env = os.getenv("ENVIRONMENT_TYPE", "production"),
     )
 
+    apply_gpu_ids(config.get("resolved_gpu_ids"))
+
     model_name = config["model_name"]
 
     # ── 1. Activate correct transformers version BEFORE any ML imports ──

studio/backend/core/training/trainer.py

@@ -33,7 +33,14 @@ if sys.platform in ("win32", "darwin"):
         sys.path.insert(0, _compile_cache)
 
 import torch
-from utils.hardware import clear_gpu_cache, safe_num_proc, dataset_map_num_proc
+from utils.hardware import (
+    clear_gpu_cache,
+    safe_num_proc,
+    dataset_map_num_proc,
+    get_device_map,
+    raise_if_offloaded,
+    get_visible_gpu_count,
+)
 
 torch._dynamo.config.recompile_limit = 64
 from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
@@ -487,6 +494,7 @@ class UnslothTrainer:
         is_dataset_audio: bool = False,
         trust_remote_code: bool = False,
         full_finetuning: bool = False,
+        gpu_ids: Optional[list[int]] = None,
     ) -> bool:
         """Load model for training (supports both text and vision models)"""
         self.load_in_4bit = load_in_4bit  # Store for training_meta.json
@@ -624,6 +632,11 @@ class UnslothTrainer:
                         self._update_progress(error = friendly, is_training = False)
                         return False
 
+            device_map = get_device_map(gpu_ids)
+            logger.info(
+                f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)"
+            )
+
             # Branch based on model type
             if self._audio_type == "csm":
                 # CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False
@@ -636,6 +649,7 @@ class UnslothTrainer:
                     dtype = None,
                     auto_model = CsmForConditionalGeneration,
                     load_in_4bit = False,
+                    device_map = device_map,
                     full_finetuning = full_finetuning,
                     token = hf_token,
                     trust_remote_code = trust_remote_code,
@@ -651,6 +665,7 @@ class UnslothTrainer:
                     model_name = model_name,
                     dtype = None,
                     load_in_4bit = False,
+                    device_map = device_map,
                     full_finetuning = full_finetuning,
                     auto_model = WhisperForConditionalGeneration,
                     whisper_language = "English",
@@ -672,6 +687,7 @@ class UnslothTrainer:
                     max_seq_length = max_seq_length,
                     dtype = None,
                     load_in_4bit = load_in_4bit,
+                    device_map = device_map,
                     full_finetuning = full_finetuning,
                     token = hf_token,
                     trust_remote_code = trust_remote_code,
@@ -711,6 +727,7 @@ class UnslothTrainer:
                     max_seq_length = max_seq_length,
                     dtype = torch.float32,  # Spark-TTS requires float32
                     load_in_4bit = False,
+                    device_map = device_map,
                     full_finetuning = full_finetuning,
                     token = hf_token,
                     trust_remote_code = trust_remote_code,
@@ -725,6 +742,7 @@ class UnslothTrainer:
                     model_name,
                     max_seq_length = max_seq_length,
                     load_in_4bit = False,
+                    device_map = device_map,
                     full_finetuning = full_finetuning,
                     token = hf_token,
                     trust_remote_code = trust_remote_code,
@@ -741,6 +759,7 @@ class UnslothTrainer:
                     max_seq_length = max_seq_length,
                     dtype = None,
                     load_in_4bit = load_in_4bit,
+                    device_map = device_map,
                     full_finetuning = full_finetuning,
                     token = hf_token,
                     trust_remote_code = trust_remote_code,
@@ -754,6 +773,7 @@ class UnslothTrainer:
                     max_seq_length = max_seq_length,
                     dtype = None,  # Auto-detect
                     load_in_4bit = load_in_4bit,
+                    device_map = device_map,
                     full_finetuning = full_finetuning,
                     token = hf_token,
                     trust_remote_code = trust_remote_code,
@@ -786,12 +806,15 @@ class UnslothTrainer:
                     max_seq_length = max_seq_length,
                     dtype = None,  # Auto-detect
                     load_in_4bit = load_in_4bit,
+                    device_map = device_map,
                     full_finetuning = full_finetuning,
                     token = hf_token,
                     trust_remote_code = trust_remote_code,
                 )
                 logger.info("Loaded text model")
 
+            raise_if_offloaded(self.model, device_map, "Studio training")
+
             if self.should_stop:
                 return False
 
@@ -824,6 +847,7 @@ class UnslothTrainer:
                     is_dataset_audio = is_dataset_audio,
                     trust_remote_code = trust_remote_code,
                     full_finetuning = full_finetuning,
+                    gpu_ids = gpu_ids,
                 )
             error_msg = str(e)
             error_lower = error_msg.lower()
@@ -2634,14 +2658,14 @@ class UnslothTrainer:
         eval_steps: float = 0.00,
         output_dir: str | None = None,
         num_epochs: int = 3,
-        learning_rate: float = 5e-5,
+        learning_rate: float = 2e-4,
         batch_size: int = 2,
         gradient_accumulation_steps: int = 4,
         warmup_steps: int = None,
         warmup_ratio: float = None,
         max_steps: int = 0,
         save_steps: int = 0,
-        weight_decay: float = 0.01,
+        weight_decay: float = 0.001,
         random_seed: int = 3407,
         packing: bool = False,
         train_on_completions: bool = False,
@@ -3010,7 +3034,7 @@ class UnslothTrainer:
                 "fp16": not is_bfloat16_supported(),
                 "bf16": is_bfloat16_supported(),
                 "logging_steps": 1,
-                "weight_decay": training_args.get("weight_decay", 0.01),
+                "weight_decay": training_args.get("weight_decay", 0.001),
                 "seed": training_args.get("random_seed", 3407),
                 "output_dir": output_dir,
                 "report_to": _build_report_targets(training_args),

studio/backend/core/training/training.py

@@ -28,6 +28,7 @@ from pathlib import Path
 from typing import Optional, Tuple, Any
 
 import matplotlib.pyplot as plt
+from utils.hardware import prepare_gpu_selection
 
 logger = get_logger(__name__)
 
@@ -159,7 +160,7 @@ class TrainingBackend:
             "warmup_ratio": kwargs.get("warmup_ratio"),
             "max_steps": kwargs.get("max_steps", 0),
             "save_steps": kwargs.get("save_steps", 0),
-            "weight_decay": kwargs.get("weight_decay", 0.01),
+            "weight_decay": kwargs.get("weight_decay", 0.001),
             "random_seed": kwargs.get("random_seed", 3407),
             "packing": kwargs.get("packing", False),
             "optim": kwargs.get("optim", "adamw_8bit"),
@@ -185,6 +186,7 @@ class TrainingBackend:
             "enable_tensorboard": kwargs.get("enable_tensorboard", False),
             "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"),
             "trust_remote_code": kwargs.get("trust_remote_code", False),
+            "gpu_ids": kwargs.get("gpu_ids"),
         }
 
         # Derive load_in_4bit from training_type
@@ -192,6 +194,22 @@ class TrainingBackend:
             config["load_in_4bit"] = False
 
         # Spawn subprocess — use locals so state is untouched on failure
+        resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
+            kwargs.get("gpu_ids"),
+            model_name = config["model_name"],
+            hf_token = config["hf_token"] or None,
+            training_type = config["training_type"],
+            load_in_4bit = config["load_in_4bit"],
+            batch_size = config.get("batch_size", 4),
+            max_seq_length = config.get("max_seq_length", 2048),
+            lora_rank = config.get("lora_r", 16),
+            target_modules = config.get("target_modules"),
+            gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"),
+            optimizer = config.get("optim", "adamw_8bit"),
+        )
+        config["resolved_gpu_ids"] = resolved_gpu_ids
+        config["gpu_selection"] = gpu_selection
+
         from .worker import run_training_process
 
         event_queue = _CTX.Queue()

studio/backend/core/training/worker.py

@@ -29,6 +29,7 @@ import urllib.error
 import urllib.request
 
 logger = get_logger(__name__)
+from utils.hardware import apply_gpu_ids
 
 
 _CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4"
@@ -367,6 +368,8 @@ def run_training_process(
         env = os.getenv("ENVIRONMENT_TYPE", "production"),
     )
 
+    apply_gpu_ids(config.get("resolved_gpu_ids"))
+
     model_name = config["model_name"]
 
     # ── 1. Activate correct transformers version BEFORE any ML imports ──
@@ -682,6 +685,7 @@ def run_training_process(
             is_dataset_image = config.get("is_dataset_image", False),
             is_dataset_audio = config.get("is_dataset_audio", False),
             trust_remote_code = config.get("trust_remote_code", False),
+            gpu_ids = config.get("resolved_gpu_ids"),
         )
         if not success or trainer.should_stop:
             if trainer.should_stop:
@@ -791,7 +795,7 @@ def run_training_process(
             warmup_ratio = config.get("warmup_ratio"),
             max_steps = max_steps if max_steps and max_steps > 0 else 0,
             save_steps = save_steps if save_steps and save_steps > 0 else 0,
-            weight_decay = config.get("weight_decay", 0.01),
+            weight_decay = config.get("weight_decay", 0.001),
             random_seed = config.get("random_seed", 3407),
             packing = config.get("packing", False),
             train_on_completions = config.get("train_on_completions", False),
@@ -1137,7 +1141,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
         "lr_scheduler_type": config.get("lr_scheduler_type", "linear"),
         "batch_sampler": BatchSamplers.NO_DUPLICATES,
         "optim": config.get("optim", "adamw_8bit"),
-        "weight_decay": config.get("weight_decay", 0.01),
+        "weight_decay": config.get("weight_decay", 0.001),
         "seed": config.get("random_seed", 3407),
     }
 

studio/backend/main.py

@@ -23,10 +23,23 @@ if _backend_dir not in sys.path:
 # See: https://github.com/python/cpython/issues/102396
 import _platform_compat  # noqa: F401
 
+import mimetypes
 import shutil
 import warnings
 from contextlib import asynccontextmanager
 
+# Fix broken Windows registry MIME types.  Some Windows installs map .js to
+# "text/plain" in the registry (HKCR\.js\Content Type).  Python's mimetypes
+# module reads from the registry, and FastAPI/Starlette's StaticFiles uses
+# mimetypes.guess_type() to set Content-Type headers.  Browsers enforce strict
+# MIME checking for ES module scripts (<script type="module">) and will refuse
+# to execute .js files served as text/plain — resulting in a blank page.
+# Calling add_type() *before* StaticFiles is instantiated ensures the correct
+# types are used regardless of the OS registry.
+if sys.platform == "win32":
+    mimetypes.add_type("application/javascript", ".js")
+    mimetypes.add_type("text/css", ".css")
+
 # Suppress annoying dependency warnings in production
 if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
     warnings.filterwarnings("ignore")
@@ -54,7 +67,12 @@ from routes import (
 )
 from auth import storage
 from auth.authentication import get_current_subject
-from utils.hardware import detect_hardware, get_device, DeviceType
+from utils.hardware import (
+    detect_hardware,
+    get_device,
+    DeviceType,
+    get_backend_visible_gpu_info,
+)
 import utils.hardware.hardware as _hw_module
 
 from utils.cache_cleanup import clear_unsloth_compiled_cache
@@ -217,69 +235,14 @@ async def shutdown_server(
 async def get_system_info():
     """Get system information"""
     import platform
-    import subprocess
     import psutil
-    from utils.hardware import get_device, get_gpu_memory_info, DeviceType
+    from utils.hardware import get_device
 
-    # GPU Info — query nvidia-smi for physical GPUs, filtered by
-    # CUDA_VISIBLE_DEVICES when set (the frontend uses this for GGUF
-    # fit estimation and llama-server respects CVD too).
-    import os
-
-    gpu_info: dict = {"available": False, "devices": []}
-
-    device = get_device()
-    if device == DeviceType.CUDA:
-        # Parse CUDA_VISIBLE_DEVICES allowlist
-        allowed_indices = None
-        cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
-        if cvd is not None and cvd.strip():
-            try:
-                allowed_indices = set(int(x.strip()) for x in cvd.split(","))
-            except ValueError:
-                pass  # Non-numeric (e.g. GPU-uuid), show all
-
-        try:
-            result = subprocess.run(
-                [
-                    "nvidia-smi",
-                    "--query-gpu=index,name,memory.total",
-                    "--format=csv,noheader,nounits",
-                ],
-                capture_output = True,
-                text = True,
-                timeout = 10,
-            )
-            if result.returncode == 0:
-                for line in result.stdout.strip().splitlines():
-                    parts = [p.strip() for p in line.split(",")]
-                    if len(parts) == 3:
-                        idx = int(parts[0])
-                        if allowed_indices is not None and idx not in allowed_indices:
-                            continue
-                        gpu_info["devices"].append(
-                            {
-                                "index": idx,
-                                "name": parts[1],
-                                "memory_total_gb": round(int(parts[2]) / 1024, 2),
-                            }
-                        )
-                gpu_info["available"] = len(gpu_info["devices"]) > 0
-        except Exception:
-            pass
-
-    # Fallback to torch-based single-GPU detection
-    if not gpu_info["available"]:
-        mem_info = get_gpu_memory_info()
-        if mem_info.get("available"):
-            gpu_info["available"] = True
-            gpu_info["devices"].append(
-                {
-                    "index": mem_info.get("device", 0),
-                    "name": mem_info.get("device_name", "Unknown"),
-                    "memory_total_gb": round(mem_info.get("total_gb", 0), 2),
-                }
-            )
+    visibility_info = get_backend_visible_gpu_info()
+    gpu_info = {
+        "available": visibility_info["available"],
+        "devices": visibility_info["devices"],
+    }
 
     # CPU & Memory
     memory = psutil.virtual_memory()
@@ -298,6 +261,13 @@ async def get_system_info():
     }
 
 
+@app.get("/api/system/gpu-visibility")
+async def get_gpu_visibility(
+    current_subject: str = Depends(get_current_subject),
+):
+    return get_backend_visible_gpu_info()
+
+
 @app.get("/api/system/hardware")
 async def get_hardware_info():
     """Return GPU name, total VRAM, and key ML package versions."""

studio/backend/models/inference.py

@@ -44,6 +44,10 @@ class LoadRequest(BaseModel):
         None,
         description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')",
     )
+    gpu_ids: Optional[List[int]] = Field(
+        None,
+        description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.",
+    )
 
 
 class UnloadRequest(BaseModel):
@@ -132,6 +136,9 @@ class LoadResponse(BaseModel):
     context_length: Optional[int] = Field(
         None, description = "Model's native context length (from GGUF metadata)"
     )
+    max_context_length: Optional[int] = Field(
+        None, description = "Maximum context length currently available on this hardware"
+    )
     supports_reasoning: bool = Field(
         False,
         description = "Whether model supports thinking/reasoning mode (enable_thinking)",
@@ -206,6 +213,10 @@ class InferenceStatusResponse(BaseModel):
     context_length: Optional[int] = Field(
         None, description = "Context length of the active model"
     )
+    max_context_length: Optional[int] = Field(
+        None,
+        description = "Maximum context length currently available for the active model",
+    )
 
 
 # =====================================================================
@@ -333,7 +344,7 @@ class ChatCompletionRequest(BaseModel):
         description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
     )
     max_tool_calls_per_message: Optional[int] = Field(
-        10,
+        25,
         ge = 0,
         description = "[x-unsloth] Maximum number of tool call iterations per message (0 = disabled, 9999 = unlimited).",
     )

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", "lmstudio"] = Field(
+    source: Literal["models_dir", "hf_cache", "lmstudio", "custom"] = Field(
         ...,
         description = "Discovery source",
     )
@@ -197,3 +197,19 @@ class LocalModelListResponse(BaseModel):
         default_factory = list,
         description = "Discovered local/cached models",
     )
+
+
+class AddScanFolderRequest(BaseModel):
+    """Request body for adding a custom scan folder."""
+
+    path: str = Field(
+        ..., description = "Absolute or relative directory path to scan for models"
+    )
+
+
+class ScanFolderInfo(BaseModel):
+    """A registered custom model scan folder."""
+
+    id: int = Field(..., description = "Database row ID")
+    path: str = Field(..., description = "Normalized absolute path")
+    created_at: str = Field(..., description = "ISO 8601 creation timestamp")

studio/backend/models/training.py

@@ -81,7 +81,7 @@ class TrainingStartRequest(BaseModel):
     warmup_ratio: Optional[float] = Field(None, description = "Warmup ratio")
     max_steps: Optional[int] = Field(None, description = "Maximum training steps")
     save_steps: int = Field(100, description = "Steps between checkpoints")
-    weight_decay: float = Field(0.01, description = "Weight decay")
+    weight_decay: float = Field(0.001, description = "Weight decay")
     random_seed: int = Field(42, description = "Random seed")
     packing: bool = Field(False, description = "Enable sequence packing")
     optim: str = Field("adamw_8bit", description = "Optimizer")
@@ -128,6 +128,12 @@ class TrainingStartRequest(BaseModel):
     enable_tensorboard: bool = Field(False, description = "Enable TensorBoard logging")
     tensorboard_dir: Optional[str] = Field(None, description = "TensorBoard directory")
 
+    # GPU selection
+    gpu_ids: Optional[List[int]] = Field(
+        None,
+        description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries.",
+    )
+
 
 class TrainingJobResponse(BaseModel):
     """Immediate response when training is initiated"""

studio/backend/routes/inference.py

@@ -86,8 +86,15 @@ import io
 import wave
 import base64
 import numpy as np
+from datetime import date as _date
 
 router = APIRouter()
+
+# Regex for stripping leaked tool-call XML from assistant messages/stream
+_TOOL_XML_RE = _re.compile(
+    r"<tool_call>.*?</tool_call>|<function=\w+>.*?</function>",
+    _re.DOTALL,
+)
 logger = get_logger(__name__)
 
 
@@ -155,6 +162,7 @@ async def load_model(
                     else False,
                     inference = inference_config,
                     context_length = llama_backend.context_length,
+                    max_context_length = llama_backend.max_context_length,
                     supports_reasoning = llama_backend.supports_reasoning,
                     reasoning_always_on = llama_backend.reasoning_always_on,
                     chat_template = llama_backend.chat_template,
@@ -205,8 +213,17 @@ async def load_model(
                 detail = f"Invalid model identifier: {request.model_path}",
             )
 
+        # Normalize gpu_ids: empty list means auto-selection, same as None
+        effective_gpu_ids = request.gpu_ids if request.gpu_ids else None
+
         # ── GGUF path: load via llama-server ──────────────────────
         if config.is_gguf:
+            if effective_gpu_ids is not None:
+                raise HTTPException(
+                    status_code = 400,
+                    detail = "gpu_ids is not supported for GGUF models yet.",
+                )
+
             llama_backend = get_llama_cpp_backend()
             unsloth_backend = get_inference_backend()
 
@@ -280,6 +297,7 @@ async def load_model(
                 has_audio_input = is_audio_input_type(_gguf_audio),
                 inference = inference_config,
                 context_length = llama_backend.context_length,
+                max_context_length = llama_backend.max_context_length,
                 supports_reasoning = llama_backend.supports_reasoning,
                 reasoning_always_on = llama_backend.reasoning_always_on,
                 supports_tools = llama_backend.supports_tools,
@@ -367,6 +385,7 @@ async def load_model(
             load_in_4bit = load_in_4bit,
             hf_token = request.hf_token,
             trust_remote_code = request.trust_remote_code,
+            gpu_ids = effective_gpu_ids,
         )
 
         if not success:
@@ -418,6 +437,9 @@ async def load_model(
 
     except HTTPException:
         raise
+    except ValueError as e:
+        logger.warning("Rejected inference GPU selection: %s", e)
+        raise HTTPException(status_code = 400, detail = str(e))
     except Exception as e:
         logger.error(f"Error loading model: {e}", exc_info = True)
         msg = str(e)
@@ -614,6 +636,7 @@ async def get_status(
                 reasoning_always_on = llama_backend.reasoning_always_on,
                 supports_tools = llama_backend.supports_tools,
                 context_length = llama_backend.context_length,
+                max_context_length = llama_backend.max_context_length,
             )
 
         # Otherwise, report Unsloth backend status
@@ -1062,6 +1085,68 @@ async def openai_chat_completions(
             else:
                 tools_to_use = ALL_TOOLS
 
+            # ── Tool-use system prompt nudge ──────────────────────
+            _tool_names = {t["function"]["name"] for t in tools_to_use}
+            _has_web = "web_search" in _tool_names
+            _has_code = "python" in _tool_names or "terminal" in _tool_names
+
+            _date_line = f"The current date is {_date.today().isoformat()}."
+
+            _web_tips = (
+                "When you search and find a relevant URL in the results, "
+                "fetch its full content by calling web_search with the url parameter. "
+                "Do not repeat the same search query. If a search returns "
+                "no useful results, try rephrasing or fetching a result URL directly."
+            )
+            _code_tips = (
+                "Use code execution for math, calculations, data processing, "
+                "or to parse and analyze information from tool results."
+            )
+
+            if _has_web and _has_code:
+                _nudge = (
+                    _date_line + " "
+                    "You have access to tools. When appropriate, prefer using "
+                    "tools rather than answering from memory. "
+                    + _web_tips
+                    + " "
+                    + _code_tips
+                )
+            elif _has_code:
+                _nudge = (
+                    _date_line + " "
+                    "You have access to tools. When appropriate, prefer using "
+                    "code execution rather than answering from memory. " + _code_tips
+                )
+            elif _has_web:
+                _nudge = (
+                    _date_line + " "
+                    "You have access to tools. When appropriate, prefer using "
+                    "web search for up-to-date or uncertain factual "
+                    "information rather than answering from memory. " + _web_tips
+                )
+            else:
+                _nudge = ""
+
+            if _nudge:
+                # Append nudge to system prompt (preserve user's prompt)
+                if system_prompt:
+                    system_prompt = system_prompt.rstrip() + "\n\n" + _nudge
+                else:
+                    system_prompt = _nudge
+                # Rebuild gguf_messages with updated system prompt
+                gguf_messages = []
+                if system_prompt:
+                    gguf_messages.append({"role": "system", "content": system_prompt})
+                gguf_messages.extend(chat_messages)
+
+            # ── Strip stale tool-call XML from conversation history ─
+            for _msg in gguf_messages:
+                if _msg.get("role") == "assistant" and isinstance(
+                    _msg.get("content"), str
+                ):
+                    _msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip()
+
             def gguf_generate_with_tools():
                 return llama_backend.generate_chat_completion_with_tools(
                     messages = gguf_messages,
@@ -1080,7 +1165,7 @@ async def openai_chat_completions(
                     else True,
                     max_tool_iterations = payload.max_tool_calls_per_message
                     if payload.max_tool_calls_per_message is not None
-                    else 10,
+                    else 25,
                     tool_call_timeout = payload.tool_call_timeout
                     if payload.tool_call_timeout is not None
                     else 300,
@@ -1142,9 +1227,13 @@ async def openai_chat_completions(
                             continue
 
                         # "content" type -- cumulative text
-                        cumulative = event.get("text", "")
-                        new_text = cumulative[len(prev_text) :]
-                        prev_text = cumulative
+                        # Sanitize the full cumulative then diff against
+                        # the last sanitized snapshot so cross-chunk XML
+                        # tags are handled correctly.
+                        raw_cumulative = event.get("text", "")
+                        clean_cumulative = _TOOL_XML_RE.sub("", raw_cumulative)
+                        new_text = clean_cumulative[len(prev_text) :]
+                        prev_text = clean_cumulative
                         if not new_text:
                             continue
                         chunk = ChatCompletionChunk(

studio/backend/routes/models.py

@@ -94,7 +94,13 @@ from models import (
     LoRAInfo,
     ModelListResponse,
 )
-from models.models import GgufVariantDetail, GgufVariantsResponse, ModelType
+from models.models import (
+    GgufVariantDetail,
+    GgufVariantsResponse,
+    ModelType,
+    ScanFolderInfo,
+    AddScanFolderRequest,
+)
 from models.responses import (
     LoRABaseModelResponse,
     VisionCheckResponse,
@@ -128,21 +134,32 @@ def _resolve_hf_cache_dir() -> Path:
         return Path.home() / ".cache" / "huggingface" / "hub"
 
 
-def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]:
+def _scan_models_dir(
+    models_dir: Path,
+    *,
+    limit: int | None = None,
+) -> 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():
+        if limit is not None and len(found) >= limit:
+            break
+        try:
+            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"))
+                or any(child.glob("*.gguf"))
+            )
+        except OSError:
+            # Skip individual children that are unreadable (permissions, broken
+            # symlinks, etc.) rather than failing the entire scan.
             continue
-        has_model_files = (
-            (child / "config.json").exists()
-            or (child / "adapter_config.json").exists()
-            or any(child.glob("*.safetensors"))
-            or any(child.glob("*.bin"))
-            or any(child.glob("*.gguf"))
-        )
         if not has_model_files:
             continue
         try:
@@ -159,21 +176,24 @@ def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]:
             ),
         )
     # Also scan for standalone .gguf files directly in the models directory
-    for gguf_file in models_dir.glob("*.gguf"):
-        if gguf_file.is_file():
-            try:
-                updated_at = gguf_file.stat().st_mtime
-            except OSError:
-                updated_at = None
-            found.append(
-                LocalModelInfo(
-                    id = str(gguf_file),
-                    display_name = gguf_file.stem,
-                    path = str(gguf_file),
-                    source = "models_dir",
-                    updated_at = updated_at,
-                ),
-            )
+    if limit is None or len(found) < limit:
+        for gguf_file in models_dir.glob("*.gguf"):
+            if limit is not None and len(found) >= limit:
+                break
+            if gguf_file.is_file():
+                try:
+                    updated_at = gguf_file.stat().st_mtime
+                except OSError:
+                    updated_at = None
+                found.append(
+                    LocalModelInfo(
+                        id = str(gguf_file),
+                        display_name = gguf_file.stem,
+                        path = str(gguf_file),
+                        source = "models_dir",
+                        updated_at = updated_at,
+                    ),
+                )
 
     return found
 
@@ -221,63 +241,69 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]:
 
     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
+        try:
+            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:
+            # child is a publisher directory -- scan its sub-directories
+            for model_dir in child.iterdir():
+                try:
+                    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),
+                                model_id = f"{child.name}/{model_dir.stem}",
+                                display_name = model_dir.stem,
+                                path = str(model_dir),
+                                source = "lmstudio",
+                                updated_at = updated_at,
+                            ),
+                        )
+                except OSError:
                     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),
-                        model_id = f"{child.name}/{model_dir.stem}",
-                        display_name = model_dir.stem,
-                        path = str(model_dir),
-                        source = "lmstudio",
-                        updated_at = updated_at,
-                    ),
-                )
+        except OSError:
+            continue
     return found
 
 
@@ -351,10 +377,39 @@ async def list_local_models(
         for lm_dir in lm_dirs:
             local_models += _scan_lmstudio_dir(lm_dir)
 
+        # Scan user-added custom folders (cap per-folder to avoid unbounded scans)
+        from storage.studio_db import list_scan_folders
+
+        _MAX_MODELS_PER_FOLDER = 200
+        try:
+            custom_folders = list_scan_folders()
+        except Exception as e:
+            logger.warning("Could not load custom scan folders: %s", e)
+            custom_folders = []
+        for folder in custom_folders:
+            folder_path = Path(folder["path"])
+            try:
+                custom_models = (
+                    _scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER)
+                    + _scan_hf_cache(folder_path)
+                    + _scan_lmstudio_dir(folder_path)
+                )[:_MAX_MODELS_PER_FOLDER]
+            except OSError as e:
+                logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e)
+                continue
+            local_models += [
+                m.model_copy(update = {"source": "custom"}) for m in custom_models
+            ]
+
+        # Deduplicate models, but always keep custom folder entries so they
+        # appear in the "Custom Folders" UI section even when the same model
+        # also exists in the HF cache or default models directory.  Use a
+        # (id, source) key for custom entries to avoid collisions.
         deduped: dict[str, LocalModelInfo] = {}
         for model in local_models:
-            if model.id not in deduped:
-                deduped[model.id] = model
+            key = f"{model.id}\x00custom" if model.source == "custom" else model.id
+            if key not in deduped:
+                deduped[key] = model
 
         models = sorted(
             deduped.values(),
@@ -376,6 +431,46 @@ async def list_local_models(
         )
 
 
+@router.get("/scan-folders")
+async def get_scan_folders(
+    current_subject: str = Depends(get_current_subject),
+):
+    """List all registered custom model scan folders."""
+    from storage.studio_db import list_scan_folders
+
+    return {"folders": list_scan_folders()}
+
+
+@router.post("/scan-folders", response_model = ScanFolderInfo, status_code = 201)
+async def add_scan_folder_endpoint(
+    body: AddScanFolderRequest,
+    current_subject: str = Depends(get_current_subject),
+):
+    """Register a new directory to scan for local models."""
+    from storage.studio_db import add_scan_folder
+
+    try:
+        folder = add_scan_folder(body.path)
+    except ValueError as e:
+        logger.warning("Scan folder rejected: %s (path=%s)", e, body.path)
+        raise HTTPException(status_code = 400, detail = str(e))
+    logger.info("Scan folder added: %s", folder.get("path"))
+    return folder
+
+
+@router.delete("/scan-folders/{folder_id}")
+async def remove_scan_folder_endpoint(
+    folder_id: int,
+    current_subject: str = Depends(get_current_subject),
+):
+    """Remove a registered custom scan folder."""
+    from storage.studio_db import remove_scan_folder
+
+    remove_scan_folder(folder_id)
+    logger.info("Scan folder removed: id=%s", folder_id)
+    return {"ok": True}
+
+
 @router.get("/list")
 async def list_models(
     current_subject: str = Depends(get_current_subject),

studio/backend/routes/training.py

@@ -88,14 +88,22 @@ async def get_hardware_utilization(
     Get a live snapshot of GPU hardware utilization.
 
     Designed to be polled by the frontend during training.
-    Returns GPU utilization %, temperature, VRAM usage, and power draw
-    via nvidia-smi for maximum accuracy.
+    Returns live GPU memory usage information for the active backend.
     """
     from utils.hardware import get_gpu_utilization
 
     return get_gpu_utilization()
 
 
+@router.get("/hardware/visible")
+async def get_visible_hardware_utilization(
+    current_subject: str = Depends(get_current_subject),
+):
+    from utils.hardware import get_visible_gpu_utilization
+
+    return get_visible_gpu_utilization()
+
+
 @router.post("/start")
 async def start_training(
     request: TrainingStartRequest,
@@ -202,6 +210,7 @@ async def start_training(
             "enable_tensorboard": request.enable_tensorboard,
             "tensorboard_dir": request.tensorboard_dir or "",
             "trust_remote_code": request.trust_remote_code,
+            "gpu_ids": request.gpu_ids,
         }
 
         # Training page has no trust_remote_code toggle — the value comes from
@@ -269,6 +278,9 @@ async def start_training(
             error = None,
         )
 
+    except ValueError as e:
+        logger.warning("Rejected training GPU selection: %s", e)
+        raise HTTPException(status_code = 400, detail = str(e))
     except Exception as e:
         logger.error(f"Error starting training: {e}", exc_info = True)
         raise HTTPException(

studio/backend/run.py

@@ -312,10 +312,10 @@ def run_server(
     if frontend_path:
         if setup_frontend(app, frontend_path):
             if not silent:
-                print(f"✅ Frontend loaded from {frontend_path}")
+                print(f"[OK] Frontend loaded from {frontend_path}")
         else:
             if not silent:
-                print(f"⚠️ Frontend not found at {frontend_path}")
+                print(f"[WARNING] Frontend not found at {frontend_path}")
 
     # Create the uvicorn server and expose it for signal handlers
     config = uvicorn.Config(

studio/backend/startup_banner.py

@@ -20,7 +20,7 @@ def stdout_supports_color() -> bool:
         return True
     try:
         return sys.stdout.isatty()
-    except Exception:
+    except (AttributeError, OSError, ValueError):
         return False
 
 
@@ -52,28 +52,36 @@ def print_studio_access_banner(
 
     ipv6_bind = bind_host in ("::", "::1")
     if ipv6_bind:
-        local_url = f"http://[::1]:{port}"
+        loopback_url = f"http://[::1]:{port}"
         alt_local = f"http://localhost:{port}"
     else:
-        local_url = f"http://127.0.0.1:{port}"
+        loopback_url = f"http://127.0.0.1:{port}"
         alt_local = f"http://localhost:{port}"
     if ":" in display_host:
         external_url = f"http://[{display_host}]:{port}"
     else:
         external_url = f"http://{display_host}:{port}"
+
     listen_all = bind_host in ("0.0.0.0", "::")
     loopback_bind = bind_host in ("127.0.0.1", "localhost", "::1")
-    api_base = local_url if listen_all or loopback_bind else external_url
+
+    # Use loopback URL only when the server is reachable on loopback;
+    # otherwise show the actual bound address.
+    primary_url = loopback_url if listen_all or loopback_bind else external_url
+    tip_url = alt_local if listen_all or loopback_bind else external_url
+    api_base = primary_url
 
     lines: list[str] = [
         "",
         style("🦥 Unsloth Studio is running", title),
         style("─" * 52, dim),
-        style("  On this machine — open this in your browser:", dim),
-        style(f"    {local_url}", local_url_style),
-        style(f"    (same as {alt_local})", dim),
+        style("  On this machine -- open this in your browser:", dim),
+        style(f"    {primary_url}", local_url_style),
     ]
 
+    if (listen_all or loopback_bind) and primary_url != alt_local:
+        lines.append(style(f"    (same as {alt_local})", dim))
+
     if listen_all and display_host not in (
         "127.0.0.1",
         "localhost",
@@ -88,7 +96,7 @@ def print_studio_access_banner(
                 style(f"    {external_url}", secondary),
             ]
         )
-    elif not listen_all and bind_host not in ("127.0.0.1", "localhost", "::1"):
+    elif not listen_all and not loopback_bind and external_url != primary_url:
         lines.extend(
             [
                 "",
@@ -105,7 +113,7 @@ def print_studio_access_banner(
             style(f"    {api_base}/api/health", secondary),
             style("─" * 52, dim),
             style(
-                "  Tip: if you are on the same computer, use the Local link above.",
+                f"  Tip: if you are on this computer, open {tip_url}/ in your browser.",
                 dim,
             ),
             "",

studio/backend/storage/studio_db.py

@@ -12,14 +12,46 @@ raw sqlite3, per-function connections. Enhancements over auth:
 
 import json
 import logging
+import os
+import platform
 import sqlite3
 import threading
+from datetime import datetime, timezone
 
 logger = logging.getLogger(__name__)
 from typing import Optional
 
+
 from utils.paths import studio_db_path, ensure_dir
 
+
+def _denied_path_prefixes() -> list[str]:
+    """Platform-aware denylist of system directories."""
+    system = platform.system()
+    if system == "Linux":
+        return ["/proc", "/sys", "/dev", "/etc", "/boot", "/run"]
+    if system == "Darwin":
+        # realpath() resolves /etc -> /private/etc, /tmp -> /private/tmp on macOS,
+        # so include the /private variants to avoid bypasses.
+        return [
+            "/System",
+            "/Library",
+            "/dev",
+            "/etc",
+            "/private/etc",
+            "/tmp",
+            "/private/tmp",
+            "/var",
+            "/private/var",
+        ]
+    if system == "Windows":
+        win = os.environ.get("SystemRoot", r"C:\Windows")
+        pf = os.environ.get("ProgramFiles", r"C:\Program Files")
+        pf86 = os.environ.get("ProgramFiles(x86)", r"C:\Program Files (x86)")
+        return [os.path.normcase(p) for p in [win, pf, pf86]]
+    return []
+
+
 _schema_lock = threading.Lock()
 _schema_ready = False
 
@@ -67,6 +99,19 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
     conn.execute(
         "CREATE INDEX IF NOT EXISTS idx_metrics_run_id ON training_metrics(run_id)"
     )
+    # Use COLLATE NOCASE on Windows so C:\Models and c:\models dedup via the
+    # UNIQUE constraint.  On Linux/macOS (case-sensitive FS) keep the default
+    # BINARY collation so /Models and /models remain distinct.
+    collation = "COLLATE NOCASE" if platform.system() == "Windows" else ""
+    conn.execute(
+        f"""
+        CREATE TABLE IF NOT EXISTS scan_folders (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            path TEXT NOT NULL UNIQUE {collation},
+            created_at TEXT NOT NULL
+        )
+        """
+    )
 
 
 def get_connection() -> sqlite3.Connection:
@@ -343,8 +388,6 @@ def delete_run(id: str) -> None:
 
 def cleanup_orphaned_runs() -> None:
     """Mark any 'running' rows as errored on startup (server restarted mid-training)."""
-    from datetime import datetime, timezone
-
     conn = get_connection()
     try:
         conn.execute(
@@ -360,3 +403,86 @@ def cleanup_orphaned_runs() -> None:
         conn.commit()
     finally:
         conn.close()
+
+
+def list_scan_folders() -> list[dict]:
+    conn = get_connection()
+    try:
+        rows = conn.execute(
+            "SELECT id, path, created_at FROM scan_folders ORDER BY created_at"
+        ).fetchall()
+        return [dict(row) for row in rows]
+    finally:
+        conn.close()
+
+
+def add_scan_folder(path: str) -> dict:
+    """Add a directory to the custom scan folder list. Returns the row."""
+    if not path or not path.strip():
+        raise ValueError("Path cannot be empty")
+    normalized = os.path.realpath(os.path.expanduser(path.strip()))
+
+    # Validate the path is an existing, readable directory before persisting.
+    if not os.path.exists(normalized):
+        raise ValueError("Path does not exist")
+    if not os.path.isdir(normalized):
+        raise ValueError("Path must be a directory, not a file")
+    if not os.access(normalized, os.R_OK | os.X_OK):
+        raise ValueError("Path is not readable")
+
+    # On Windows, use normcase for denylist comparison but store the
+    # original-cased path so downstream consumers see the native
+    # drive-letter casing the user expects (e.g. C:\Models, not c:\models).
+    is_win = platform.system() == "Windows"
+    check = os.path.normcase(normalized) if is_win else normalized
+    for prefix in _denied_path_prefixes():
+        if check == prefix or check.startswith(prefix + os.sep):
+            raise ValueError(f"Path under {prefix} is not allowed")
+
+    conn = get_connection()
+    try:
+        now = datetime.now(timezone.utc).isoformat()
+        # On Windows, use case-insensitive lookup so C:\Models and c:\models
+        # dedup correctly while preserving the originally-stored casing.
+        if is_win:
+            existing = conn.execute(
+                "SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE",
+                (normalized,),
+            ).fetchone()
+        else:
+            existing = conn.execute(
+                "SELECT id, path, created_at FROM scan_folders WHERE path = ?",
+                (normalized,),
+            ).fetchone()
+        if existing is not None:
+            return dict(existing)
+        try:
+            conn.execute(
+                "INSERT INTO scan_folders (path, created_at) VALUES (?, ?)",
+                (normalized, now),
+            )
+            conn.commit()
+        except sqlite3.IntegrityError:
+            pass  # duplicate -- fall through to SELECT
+        # Use the same collation as the pre-check so we find the row even
+        # when a concurrent writer stored it with different casing (Windows).
+        fallback_sql = (
+            "SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE"
+            if is_win
+            else "SELECT id, path, created_at FROM scan_folders WHERE path = ?"
+        )
+        row = conn.execute(fallback_sql, (normalized,)).fetchone()
+        if row is None:
+            raise ValueError("Folder was concurrently removed")
+        return dict(row)
+    finally:
+        conn.close()
+
+
+def remove_scan_folder(id: int) -> None:
+    conn = get_connection()
+    try:
+        conn.execute("DELETE FROM scan_folders WHERE id = ?", (id,))
+        conn.commit()
+    finally:
+        conn.close()

studio/backend/tests/test_gpu_selection_sandbox.py

@@ -0,0 +1,544 @@
+#!/usr/bin/env python3
+"""
+Sandbox test for multi-GPU selection logic.
+
+Tests the core GPU selection, memory estimation, and device_map logic
+in an isolated environment. Can be run on Linux, macOS, and Windows
+without requiring actual GPUs -- all hardware calls are mocked.
+
+Usage:
+    python -m pytest studio/backend/tests/test_gpu_selection_sandbox.py -v
+    # or directly:
+    python studio/backend/tests/test_gpu_selection_sandbox.py
+"""
+
+import os
+import sys
+import unittest
+from pathlib import Path
+from unittest.mock import patch, MagicMock
+
+# Ensure backend is on sys.path
+_backend_root = Path(__file__).resolve().parent.parent
+if str(_backend_root) not in sys.path:
+    sys.path.insert(0, str(_backend_root))
+
+
+def _make_fake_config(
+    vocab_size = 32000,
+    hidden_size = 4096,
+    intermediate_size = 11008,
+    num_hidden_layers = 32,
+    num_attention_heads = 32,
+    num_key_value_heads = 8,
+    tie_word_embeddings = False,
+):
+    """Create a fake HF config-like object for estimation tests."""
+    from types import SimpleNamespace
+
+    return SimpleNamespace(
+        vocab_size = vocab_size,
+        hidden_size = hidden_size,
+        intermediate_size = intermediate_size,
+        num_hidden_layers = num_hidden_layers,
+        num_attention_heads = num_attention_heads,
+        num_key_value_heads = num_key_value_heads,
+        tie_word_embeddings = tie_word_embeddings,
+    )
+
+
+class TestEstimateFP16ModelSizeFromConfig(unittest.TestCase):
+    """Test the config-based model size estimation."""
+
+    def test_llama_8b_size_reasonable(self):
+        from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config
+
+        config = _make_fake_config(
+            vocab_size = 128256,
+            hidden_size = 4096,
+            intermediate_size = 14336,
+            num_hidden_layers = 32,
+            num_attention_heads = 32,
+            num_key_value_heads = 8,
+            tie_word_embeddings = False,
+        )
+        size = _estimate_fp16_model_size_bytes_from_config(config)
+        self.assertIsNotNone(size)
+        size_gb = size / (1024**3)
+        # Llama 3.1 8B should be ~15GB in fp16
+        self.assertGreater(size_gb, 12)
+        self.assertLess(size_gb, 20)
+
+    def test_small_model(self):
+        from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config
+
+        config = _make_fake_config(
+            vocab_size = 32000,
+            hidden_size = 2048,
+            intermediate_size = 5504,
+            num_hidden_layers = 22,
+            num_attention_heads = 32,
+            num_key_value_heads = 4,
+        )
+        size = _estimate_fp16_model_size_bytes_from_config(config)
+        self.assertIsNotNone(size)
+        size_gb = size / (1024**3)
+        # ~1B model should be ~2GB in fp16
+        self.assertGreater(size_gb, 1)
+        self.assertLess(size_gb, 5)
+
+    def test_returns_none_for_incomplete_config(self):
+        from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config
+        from types import SimpleNamespace
+
+        config = SimpleNamespace(vocab_size = 32000)  # Missing most fields
+        size = _estimate_fp16_model_size_bytes_from_config(config)
+        self.assertIsNone(size)
+
+    def test_moe_model(self):
+        from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config
+        from types import SimpleNamespace
+
+        config = SimpleNamespace(
+            vocab_size = 152064,
+            hidden_size = 3584,
+            intermediate_size = 18944,
+            num_hidden_layers = 28,
+            num_attention_heads = 28,
+            num_key_value_heads = 4,
+            tie_word_embeddings = False,
+            num_local_experts = 64,
+            moe_intermediate_size = 2560,
+        )
+        size = _estimate_fp16_model_size_bytes_from_config(config)
+        self.assertIsNotNone(size)
+        size_gb = size / (1024**3)
+        # MoE model with 64 experts should be large
+        self.assertGreater(size_gb, 50)
+
+
+class TestEstimateRequiredModelMemory(unittest.TestCase):
+    """Test memory requirement estimation."""
+
+    def test_inference_fp16_uses_1_3x(self):
+        from utils.hardware.hardware import estimate_required_model_memory_gb
+
+        with patch(
+            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+            return_value = (10 * (1024**3), "config"),  # 10GB model
+        ):
+            required, meta = estimate_required_model_memory_gb(
+                "test/model",
+                training_type = None,  # inference
+                load_in_4bit = False,
+            )
+            self.assertIsNotNone(required)
+            self.assertAlmostEqual(required, 13.0, places = 0)
+            self.assertEqual(meta["mode"], "inference")
+
+    def test_inference_4bit_uses_reduced_estimate(self):
+        from utils.hardware.hardware import estimate_required_model_memory_gb
+
+        with patch(
+            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+            return_value = (30 * (1024**3), "config"),  # 30GB fp16 model
+        ):
+            required, meta = estimate_required_model_memory_gb(
+                "test/model",
+                training_type = None,  # inference
+                load_in_4bit = True,
+            )
+            self.assertIsNotNone(required)
+            # 4bit base = 30/3.2 = 9.375GB, required = 9.375 + max(9.375*0.3, 2) = 12.19GB
+            self.assertAlmostEqual(required, 12.2, places = 0)
+
+    def test_4bit_training_reduces_base(self):
+        from utils.hardware.hardware import estimate_required_model_memory_gb
+
+        with patch(
+            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+            return_value = (30 * (1024**3), "config"),  # 30GB fp16 model
+        ):
+            required, meta = estimate_required_model_memory_gb(
+                "test/model",
+                training_type = "LoRA/QLoRA",
+                load_in_4bit = True,
+            )
+            self.assertIsNotNone(required)
+            # fallback: base=30/3.2=9.375, lora=30*0.04=1.2, act=30*0.15=4.5, cuda=1.4
+            self.assertAlmostEqual(required, 16.5, places = 0)
+
+    def test_full_finetune_uses_3_5x(self):
+        from utils.hardware.hardware import estimate_required_model_memory_gb
+
+        with patch(
+            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+            return_value = (10 * (1024**3), "config"),  # 10GB model
+        ):
+            required, meta = estimate_required_model_memory_gb(
+                "test/model",
+                training_type = "Full Finetuning",
+            )
+            self.assertIsNotNone(required)
+            # fallback: 10 * 3.5 + 1.4 cuda overhead = 36.4
+            self.assertAlmostEqual(required, 36.4, places = 0)
+
+    def test_returns_none_when_unavailable(self):
+        from utils.hardware.hardware import estimate_required_model_memory_gb
+
+        with patch(
+            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+            return_value = (None, "unavailable"),
+        ):
+            required, meta = estimate_required_model_memory_gb("test/model")
+            self.assertIsNone(required)
+
+
+class TestAutoSelectGpuIds(unittest.TestCase):
+    """Test automatic GPU selection based on model size and free memory."""
+
+    def _make_utilization(self, devices):
+        """Create a fake utilization response."""
+        return {
+            "available": True,
+            "devices": [
+                {
+                    "index": idx,
+                    "vram_total_gb": total,
+                    "vram_used_gb": total - free,
+                }
+                for idx, total, free in devices
+            ],
+        }
+
+    def test_single_gpu_sufficient(self):
+        from utils.hardware.hardware import auto_select_gpu_ids
+        import utils.hardware.hardware as hw
+
+        with (
+            patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
+            patch.object(
+                hw,
+                "estimate_required_model_memory_gb",
+                return_value = (
+                    10.0,
+                    {
+                        "mode": "inference",
+                        "required_gb": 10.0,
+                        "model_size_source": "config",
+                        "model_size_gb": 7.7,
+                    },
+                ),
+            ),
+            patch.object(
+                hw,
+                "_get_parent_visible_gpu_spec",
+                return_value = {
+                    "raw": "0,1,2,3",
+                    "numeric_ids": [0, 1, 2, 3],
+                    "supports_explicit_gpu_ids": True,
+                },
+            ),
+            patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1, 2, 3]),
+            patch.object(
+                hw,
+                "get_visible_gpu_utilization",
+                return_value = self._make_utilization(
+                    [
+                        (0, 80.0, 75.0),
+                        (1, 80.0, 78.0),
+                        (2, 80.0, 70.0),
+                        (3, 80.0, 72.0),
+                    ]
+                ),
+            ),
+        ):
+            selected, meta = auto_select_gpu_ids("test/model")
+            # Should pick GPU 1 (most free memory: 78GB) -- enough for 10GB
+            self.assertEqual(len(selected), 1)
+            self.assertEqual(selected[0], 1)
+
+    def test_two_gpus_needed(self):
+        from utils.hardware.hardware import auto_select_gpu_ids
+        import utils.hardware.hardware as hw
+
+        with (
+            patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
+            patch.object(
+                hw,
+                "estimate_required_model_memory_gb",
+                return_value = (
+                    50.0,
+                    {
+                        "mode": "inference",
+                        "required_gb": 50.0,
+                        "model_size_source": "config",
+                        "model_size_gb": 38.0,
+                    },
+                ),
+            ),
+            patch.object(
+                hw,
+                "_get_parent_visible_gpu_spec",
+                return_value = {
+                    "raw": "0,1",
+                    "numeric_ids": [0, 1],
+                    "supports_explicit_gpu_ids": True,
+                },
+            ),
+            patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]),
+            patch.object(
+                hw,
+                "get_visible_gpu_utilization",
+                return_value = self._make_utilization(
+                    [
+                        (0, 40.0, 30.0),  # 30GB free
+                        (1, 40.0, 35.0),  # 35GB free
+                    ]
+                ),
+            ),
+        ):
+            selected, meta = auto_select_gpu_ids("test/model")
+            # 35GB (first) + 30*0.85 (second) = 60.5GB > 50GB
+            self.assertEqual(len(selected), 2)
+
+    def test_non_cuda_returns_none(self):
+        from utils.hardware.hardware import auto_select_gpu_ids
+        import utils.hardware.hardware as hw
+
+        with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU):
+            selected, meta = auto_select_gpu_ids("test/model")
+            self.assertIsNone(selected)
+            self.assertEqual(meta["selection_mode"], "non_cuda")
+
+
+class TestGetDeviceMap(unittest.TestCase):
+    """Test device_map string generation."""
+
+    def test_single_gpu_returns_sequential(self):
+        from utils.hardware.hardware import get_device_map
+        import utils.hardware.hardware as hw
+
+        with (
+            patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
+            patch.object(
+                hw,
+                "_get_parent_visible_gpu_spec",
+                return_value = {
+                    "raw": "0",
+                    "numeric_ids": [0],
+                    "supports_explicit_gpu_ids": True,
+                },
+            ),
+            patch.object(hw, "get_visible_gpu_count", return_value = 1),
+        ):
+            dm = get_device_map(gpu_ids = [0])
+            self.assertEqual(dm, "sequential")
+
+    def test_multi_gpu_returns_balanced(self):
+        from utils.hardware.hardware import get_device_map
+        import utils.hardware.hardware as hw
+
+        with patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA):
+            dm = get_device_map(gpu_ids = [0, 1])
+            self.assertEqual(dm, "balanced")
+
+    def test_cpu_returns_sequential(self):
+        from utils.hardware.hardware import get_device_map
+        import utils.hardware.hardware as hw
+
+        with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU):
+            dm = get_device_map(gpu_ids = None)
+            self.assertEqual(dm, "sequential")
+
+
+class TestResolveRequestedGpuIds(unittest.TestCase):
+    """Test GPU ID validation."""
+
+    def test_none_returns_parent_visible(self):
+        from utils.hardware.hardware import resolve_requested_gpu_ids
+
+        with (
+            patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3"}, clear = False),
+            patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
+        ):
+            result = resolve_requested_gpu_ids(None)
+            self.assertEqual(result, [2, 3])
+
+    def test_empty_list_returns_parent_visible(self):
+        from utils.hardware.hardware import resolve_requested_gpu_ids
+
+        with (
+            patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3"}, clear = False),
+            patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
+        ):
+            result = resolve_requested_gpu_ids([])
+            self.assertEqual(result, [2, 3])
+
+    def test_duplicates_rejected(self):
+        from utils.hardware.hardware import resolve_requested_gpu_ids
+
+        with (
+            patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1,2"}, clear = False),
+            patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
+        ):
+            with self.assertRaises(ValueError):
+                resolve_requested_gpu_ids([1, 1])
+
+    def test_out_of_range_rejected(self):
+        from utils.hardware.hardware import resolve_requested_gpu_ids
+
+        with (
+            patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"}, clear = False),
+            patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 4),
+        ):
+            with self.assertRaises(ValueError):
+                resolve_requested_gpu_ids([5])
+
+    def test_uuid_env_var_rejects_explicit_ids(self):
+        from utils.hardware.hardware import resolve_requested_gpu_ids
+
+        with (
+            patch.dict(
+                os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-abc,GPU-def"}, clear = False
+            ),
+            patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
+        ):
+            with self.assertRaises(ValueError):
+                resolve_requested_gpu_ids([0])
+
+
+class TestApplyGpuIds(unittest.TestCase):
+    """Test CUDA_VISIBLE_DEVICES environment variable setting."""
+
+    def test_apply_list(self):
+        from utils.hardware.hardware import apply_gpu_ids
+
+        with patch.dict(os.environ, {}, clear = False):
+            apply_gpu_ids([3, 5])
+            self.assertEqual(os.environ.get("CUDA_VISIBLE_DEVICES"), "3,5")
+
+    def test_apply_none_does_nothing(self):
+        from utils.hardware.hardware import apply_gpu_ids
+
+        original = os.environ.get("CUDA_VISIBLE_DEVICES")
+        apply_gpu_ids(None)
+        self.assertEqual(os.environ.get("CUDA_VISIBLE_DEVICES"), original)
+
+
+class TestMultiGpuOverheadAccounting(unittest.TestCase):
+    """Test that multi-GPU overhead is applied correctly.
+
+    The first GPU should keep its full free memory, and only
+    additional GPUs should have the overhead factor applied.
+    """
+
+    def _make_utilization(self, devices):
+        return {
+            "available": True,
+            "devices": [
+                {
+                    "index": idx,
+                    "vram_total_gb": total,
+                    "vram_used_gb": total - free,
+                }
+                for idx, total, free in devices
+            ],
+        }
+
+    def test_first_gpu_not_penalized(self):
+        """A model that just fits on 1 GPU should not require 2 GPUs."""
+        from utils.hardware.hardware import auto_select_gpu_ids
+        import utils.hardware.hardware as hw
+
+        # Model requires 79GB, GPU has 80GB free
+        with (
+            patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
+            patch.object(
+                hw,
+                "estimate_required_model_memory_gb",
+                return_value = (
+                    79.0,
+                    {
+                        "mode": "inference",
+                        "required_gb": 79.0,
+                        "model_size_source": "config",
+                        "model_size_gb": 60.0,
+                    },
+                ),
+            ),
+            patch.object(
+                hw,
+                "_get_parent_visible_gpu_spec",
+                return_value = {
+                    "raw": "0,1",
+                    "numeric_ids": [0, 1],
+                    "supports_explicit_gpu_ids": True,
+                },
+            ),
+            patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]),
+            patch.object(
+                hw,
+                "get_visible_gpu_utilization",
+                return_value = self._make_utilization(
+                    [
+                        (0, 80.0, 80.0),
+                        (1, 80.0, 80.0),
+                    ]
+                ),
+            ),
+        ):
+            selected, meta = auto_select_gpu_ids("test/model")
+            # Should fit on 1 GPU (80GB >= 79GB)
+            self.assertEqual(len(selected), 1)
+
+    def test_second_gpu_has_overhead(self):
+        """When 2 GPUs are needed, the second one's contribution is reduced."""
+        from utils.hardware.hardware import auto_select_gpu_ids
+        import utils.hardware.hardware as hw
+
+        # Model requires 110GB. First GPU has 80GB, second has 40GB.
+        # With overhead: 80 + 40*0.85 = 114GB -- just enough
+        with (
+            patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
+            patch.object(
+                hw,
+                "estimate_required_model_memory_gb",
+                return_value = (
+                    110.0,
+                    {
+                        "mode": "inference",
+                        "required_gb": 110.0,
+                        "model_size_source": "config",
+                        "model_size_gb": 85.0,
+                    },
+                ),
+            ),
+            patch.object(
+                hw,
+                "_get_parent_visible_gpu_spec",
+                return_value = {
+                    "raw": "0,1",
+                    "numeric_ids": [0, 1],
+                    "supports_explicit_gpu_ids": True,
+                },
+            ),
+            patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]),
+            patch.object(
+                hw,
+                "get_visible_gpu_utilization",
+                return_value = self._make_utilization(
+                    [
+                        (0, 80.0, 80.0),
+                        (1, 80.0, 40.0),
+                    ]
+                ),
+            ),
+        ):
+            selected, meta = auto_select_gpu_ids("test/model")
+            # Should use both GPUs
+            self.assertEqual(len(selected), 2)
+
+
+if __name__ == "__main__":
+    unittest.main()

studio/backend/tests/test_utils.py

@@ -285,7 +285,7 @@ class TestLogGpuMemory:
     def test_does_not_raise(self):
         log_gpu_memory("test")
 
-    def test_logs_gpu_info_when_available(self, caplog):
+    def test_logs_gpu_info_when_available(self, capfd):
         fake_info = {
             "available": True,
             "backend": "cuda",
@@ -295,35 +295,27 @@ class TestLogGpuMemory:
             "utilization_pct": 12.5,
             "free_gb": 14.0,
         }
-        import structlog
-        from loggers import get_logger
 
-        with (
-            patch(
-                "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
-            ),
-            caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"),
+        with patch(
+            "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
         ):
             log_gpu_memory("unit-test")
 
-        assert "unit-test" in caplog.text
-        assert "CUDA" in caplog.text
-        assert "FakeGPU" in caplog.text
+        captured = capfd.readouterr()
+        assert "unit-test" in captured.out
+        assert "CUDA" in captured.out
+        assert "FakeGPU" in captured.out
 
-    def test_logs_cpu_fallback_when_no_gpu(self, caplog):
+    def test_logs_cpu_fallback_when_no_gpu(self, capfd):
         fake_info = {"available": False, "backend": "cpu"}
-        import structlog
-        from loggers import get_logger
 
-        with (
-            patch(
-                "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
-            ),
-            caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"),
+        with patch(
+            "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
         ):
             log_gpu_memory("cpu-test")
 
-        assert "No GPU available" in caplog.text
+        captured = capfd.readouterr()
+        assert "No GPU available" in captured.out
 
 
 # ========== format_error_message() ==========

studio/backend/tests/test_vram_estimation.py

@@ -0,0 +1,695 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+import unittest
+from types import SimpleNamespace
+
+from utils.hardware.vram_estimation import (
+    ModelArchConfig,
+    TrainingVramConfig,
+    extract_arch_config,
+    compute_model_weights_bytes,
+    compute_total_params,
+    compute_lora_params,
+    compute_lora_adapter_bytes,
+    compute_optimizer_bytes,
+    compute_gradient_bytes,
+    compute_activation_bytes,
+    estimate_training_vram,
+    DEFAULT_TARGET_MODULES,
+)
+
+
+def _gb(b: int) -> float:
+    return b / (1024**3)
+
+
+LLAMA_8B = ModelArchConfig(
+    hidden_size = 4096,
+    num_hidden_layers = 32,
+    num_attention_heads = 32,
+    num_key_value_heads = 8,
+    intermediate_size = 14336,
+    vocab_size = 128256,
+    tie_word_embeddings = False,
+)
+
+QWEN_05B = ModelArchConfig(
+    hidden_size = 896,
+    num_hidden_layers = 24,
+    num_attention_heads = 14,
+    num_key_value_heads = 2,
+    intermediate_size = 4864,
+    vocab_size = 151936,
+    tie_word_embeddings = True,
+)
+
+MOE_CONFIG = ModelArchConfig(
+    hidden_size = 4096,
+    num_hidden_layers = 32,
+    num_attention_heads = 32,
+    num_key_value_heads = 8,
+    intermediate_size = 14336,
+    vocab_size = 32000,
+    tie_word_embeddings = False,
+    num_experts = 8,
+)
+
+DEEPSEEK_V3 = ModelArchConfig(
+    hidden_size = 7168,
+    num_hidden_layers = 61,
+    num_attention_heads = 128,
+    num_key_value_heads = 128,
+    intermediate_size = 18432,
+    vocab_size = 129280,
+    tie_word_embeddings = False,
+    num_experts = 256,
+    moe_intermediate_size = 2048,
+    n_shared_experts = 1,
+    num_dense_layers = 3,
+    q_lora_rank = 1536,
+    kv_lora_rank = 512,
+    qk_nope_head_dim = 128,
+    qk_rope_head_dim = 64,
+    v_head_dim = 128,
+)
+
+QWEN3_MOE_30B = ModelArchConfig(
+    hidden_size = 2048,
+    num_hidden_layers = 48,
+    num_attention_heads = 32,
+    num_key_value_heads = 4,
+    intermediate_size = 8192,
+    vocab_size = 151936,
+    tie_word_embeddings = True,
+    num_experts = 128,
+    moe_intermediate_size = 768,
+    n_shared_experts = 0,
+    num_dense_layers = 0,
+)
+
+GLM4_MOE = ModelArchConfig(
+    hidden_size = 4096,
+    num_hidden_layers = 46,
+    num_attention_heads = 96,
+    num_key_value_heads = 8,
+    intermediate_size = 10944,
+    vocab_size = 151552,
+    tie_word_embeddings = False,
+    num_experts = 128,
+    moe_intermediate_size = 1408,
+    n_shared_experts = 1,
+    num_dense_layers = 1,
+)
+
+GPT_OSS = ModelArchConfig(
+    hidden_size = 6144,
+    num_hidden_layers = 64,
+    num_attention_heads = 64,
+    num_key_value_heads = 8,
+    intermediate_size = 2880,
+    vocab_size = 200064,
+    tie_word_embeddings = False,
+    num_experts = 128,
+    moe_intermediate_size = None,
+    n_shared_experts = 0,
+    num_dense_layers = 0,
+)
+
+
+class TestExtractArchConfig(unittest.TestCase):
+    def test_basic_config(self):
+        hf_config = SimpleNamespace(
+            hidden_size = 4096,
+            num_hidden_layers = 32,
+            num_attention_heads = 32,
+            num_key_value_heads = 8,
+            intermediate_size = 14336,
+            vocab_size = 128256,
+            tie_word_embeddings = False,
+        )
+        arch = extract_arch_config(hf_config)
+        self.assertIsNotNone(arch)
+        self.assertEqual(arch.hidden_size, 4096)
+        self.assertEqual(arch.num_hidden_layers, 32)
+        self.assertEqual(arch.num_key_value_heads, 8)
+        self.assertIsNone(arch.num_experts)
+
+    def test_vlm_text_config(self):
+        text_cfg = SimpleNamespace(
+            hidden_size = 2048,
+            num_hidden_layers = 24,
+            num_attention_heads = 16,
+            num_key_value_heads = 4,
+            intermediate_size = 8192,
+            vocab_size = 32000,
+            tie_word_embeddings = True,
+        )
+        hf_config = SimpleNamespace(text_config = text_cfg)
+        arch = extract_arch_config(hf_config)
+        self.assertIsNotNone(arch)
+        self.assertEqual(arch.hidden_size, 2048)
+
+    def test_moe_detection(self):
+        hf_config = SimpleNamespace(
+            hidden_size = 4096,
+            num_hidden_layers = 32,
+            num_attention_heads = 32,
+            num_key_value_heads = 8,
+            intermediate_size = 14336,
+            vocab_size = 32000,
+            tie_word_embeddings = False,
+            num_local_experts = 8,
+        )
+        arch = extract_arch_config(hf_config)
+        self.assertEqual(arch.num_experts, 8)
+
+    def test_missing_fields_returns_none(self):
+        hf_config = SimpleNamespace(hidden_size = 4096)
+        arch = extract_arch_config(hf_config)
+        self.assertIsNone(arch)
+
+    def test_intermediate_size_list(self):
+        hf_config = SimpleNamespace(
+            hidden_size = 2048,
+            num_hidden_layers = 24,
+            num_attention_heads = 16,
+            num_key_value_heads = 4,
+            intermediate_size = [8192, 8192],
+            vocab_size = 32000,
+            tie_word_embeddings = True,
+        )
+        arch = extract_arch_config(hf_config)
+        self.assertEqual(arch.intermediate_size, 8192)
+
+
+class TestModelWeightsBytes(unittest.TestCase):
+    def test_llama_8b_fp16(self):
+        weight_bytes = compute_model_weights_bytes(LLAMA_8B, "full", False)
+        weight_gb = _gb(weight_bytes)
+        self.assertGreater(weight_gb, 14.0)
+        self.assertLess(weight_gb, 18.0)
+
+    def test_llama_8b_qlora_4bit(self):
+        weight_bytes = compute_model_weights_bytes(LLAMA_8B, "qlora", True)
+        weight_gb = _gb(weight_bytes)
+        self.assertGreater(weight_gb, 4.0)
+        self.assertLess(weight_gb, 7.0)
+
+    def test_4bit_smaller_than_fp16(self):
+        fp16 = compute_model_weights_bytes(LLAMA_8B, "full", False)
+        q4 = compute_model_weights_bytes(LLAMA_8B, "qlora", True)
+        self.assertLess(q4, fp16)
+        ratio = fp16 / q4
+        self.assertGreater(ratio, 2.0)
+        self.assertLess(ratio, 4.0)
+
+    def test_moe_larger_than_dense(self):
+        dense = compute_model_weights_bytes(LLAMA_8B, "full", False)
+        moe = compute_model_weights_bytes(MOE_CONFIG, "full", False)
+        self.assertGreater(moe, dense * 3)
+
+
+class TestLoraParams(unittest.TestCase):
+    def test_llama_8b_default_modules_rank16(self):
+        lora_p = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES)
+        total_p = compute_total_params(LLAMA_8B)
+        ratio = lora_p / total_p
+        self.assertGreater(ratio, 0.005)
+        self.assertLess(ratio, 0.05)
+
+    def test_higher_rank_more_params(self):
+        r16 = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES)
+        r64 = compute_lora_params(LLAMA_8B, 64, DEFAULT_TARGET_MODULES)
+        self.assertAlmostEqual(r64 / r16, 4.0, places = 1)
+
+    def test_fewer_modules_fewer_params(self):
+        all_mods = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES)
+        qv_only = compute_lora_params(LLAMA_8B, 16, ["q_proj", "v_proj"])
+        self.assertLess(qv_only, all_mods)
+
+    def test_moe_mlp_modules_scale_with_experts(self):
+        dense_lora = compute_lora_params(
+            LLAMA_8B, 16, ["gate_proj", "up_proj", "down_proj"]
+        )
+        moe_lora = compute_lora_params(
+            MOE_CONFIG, 16, ["gate_proj", "up_proj", "down_proj"]
+        )
+        ratio = moe_lora / dense_lora
+        self.assertAlmostEqual(ratio, 8.0, delta = 0.5)
+
+    def test_attention_modules_same_for_moe(self):
+        dense_attn = compute_lora_params(
+            LLAMA_8B, 16, ["q_proj", "k_proj", "v_proj", "o_proj"]
+        )
+        moe_attn = compute_lora_params(
+            MOE_CONFIG, 16, ["q_proj", "k_proj", "v_proj", "o_proj"]
+        )
+        self.assertEqual(dense_attn, moe_attn)
+
+
+class TestOptimizerBytes(unittest.TestCase):
+    def test_adamw_8bit(self):
+        self.assertEqual(compute_optimizer_bytes(1_000_000, "adamw_8bit"), 4_000_000)
+
+    def test_adamw_torch(self):
+        self.assertEqual(compute_optimizer_bytes(1_000_000, "adamw_torch"), 6_000_000)
+
+    def test_sgd(self):
+        self.assertEqual(compute_optimizer_bytes(1_000_000, "sgd"), 4_000_000)
+
+    def test_unknown_defaults_to_4(self):
+        self.assertEqual(compute_optimizer_bytes(1_000_000, "some_new_opt"), 4_000_000)
+
+
+class TestGradientBytes(unittest.TestCase):
+    def test_fp16_gradients(self):
+        self.assertEqual(compute_gradient_bytes(1_000_000), 2_000_000)
+
+
+class TestActivationBytes(unittest.TestCase):
+    def test_no_gc_scales_with_layers(self):
+        act_none = compute_activation_bytes(LLAMA_8B, 2, 2048, "none")
+        act_gc = compute_activation_bytes(LLAMA_8B, 2, 2048, "true")
+        self.assertGreater(act_none, act_gc * 10)
+
+    def test_unsloth_gc_smaller_than_standard(self):
+        act_true = compute_activation_bytes(LLAMA_8B, 2, 2048, "true")
+        act_unsloth = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth")
+        self.assertLess(act_unsloth, act_true)
+
+    def test_lora_activations_smaller_than_full_ft(self):
+        full_ft = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth", is_lora = False)
+        lora = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth", is_lora = True)
+        self.assertLess(lora, full_ft)
+
+    def test_scales_with_batch_size(self):
+        act_bsz2 = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth")
+        act_bsz4 = compute_activation_bytes(LLAMA_8B, 4, 2048, "unsloth")
+        self.assertAlmostEqual(act_bsz4 / act_bsz2, 2.0, delta = 0.1)
+
+    def test_scales_with_seq_len(self):
+        act_2k = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth")
+        act_4k = compute_activation_bytes(LLAMA_8B, 2, 4096, "unsloth")
+        self.assertAlmostEqual(act_4k / act_2k, 2.0, delta = 0.1)
+
+
+class TestEstimateTrainingVram(unittest.TestCase):
+    def test_llama_8b_qlora_reasonable_total(self):
+        config = TrainingVramConfig(
+            training_method = "qlora",
+            batch_size = 2,
+            max_seq_length = 2048,
+            lora_rank = 16,
+            gradient_checkpointing = "unsloth",
+            optimizer = "adamw_8bit",
+            load_in_4bit = True,
+        )
+        breakdown = estimate_training_vram(LLAMA_8B, config)
+        total_gb = _gb(breakdown.total)
+        self.assertGreater(total_gb, 5.0)
+        self.assertLess(total_gb, 12.0)
+
+    def test_llama_8b_full_ft_reasonable_total(self):
+        config = TrainingVramConfig(
+            training_method = "full",
+            batch_size = 2,
+            max_seq_length = 2048,
+            gradient_checkpointing = "unsloth",
+            optimizer = "adamw_8bit",
+            load_in_4bit = False,
+        )
+        breakdown = estimate_training_vram(LLAMA_8B, config)
+        total_gb = _gb(breakdown.total)
+        self.assertGreater(total_gb, 50.0)
+        self.assertLess(total_gb, 75.0)
+
+    def test_qlora_much_less_than_full_ft(self):
+        qlora_config = TrainingVramConfig(
+            training_method = "qlora",
+            load_in_4bit = True,
+            batch_size = 2,
+            max_seq_length = 2048,
+        )
+        full_config = TrainingVramConfig(
+            training_method = "full",
+            load_in_4bit = False,
+            batch_size = 2,
+            max_seq_length = 2048,
+        )
+        qlora = estimate_training_vram(LLAMA_8B, qlora_config)
+        full = estimate_training_vram(LLAMA_8B, full_config)
+        self.assertLess(qlora.total, full.total / 3)
+
+    def test_qwen_05b_qlora_fits_in_4gb(self):
+        config = TrainingVramConfig(
+            training_method = "qlora",
+            batch_size = 2,
+            max_seq_length = 2048,
+            lora_rank = 16,
+            gradient_checkpointing = "unsloth",
+            optimizer = "adamw_8bit",
+            load_in_4bit = True,
+        )
+        breakdown = estimate_training_vram(QWEN_05B, config)
+        total_gb = _gb(breakdown.total)
+        self.assertLess(total_gb, 5.0)
+
+    def test_breakdown_components_positive(self):
+        config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True)
+        breakdown = estimate_training_vram(LLAMA_8B, config)
+        self.assertGreater(breakdown.model_weights, 0)
+        self.assertGreater(breakdown.lora_adapters, 0)
+        self.assertGreater(breakdown.optimizer_states, 0)
+        self.assertGreater(breakdown.gradients, 0)
+        self.assertGreater(breakdown.activations, 0)
+        self.assertGreater(breakdown.cuda_overhead, 0)
+
+    def test_full_ft_no_lora_adapters(self):
+        config = TrainingVramConfig(training_method = "full", load_in_4bit = False)
+        breakdown = estimate_training_vram(LLAMA_8B, config)
+        self.assertEqual(breakdown.lora_adapters, 0)
+
+    def test_to_gb_dict_keys(self):
+        config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True)
+        breakdown = estimate_training_vram(LLAMA_8B, config)
+        gb_dict = breakdown.to_gb_dict()
+        expected_keys = {
+            "model_weights_gb",
+            "lora_adapters_gb",
+            "optimizer_states_gb",
+            "gradients_gb",
+            "activations_gb",
+            "cuda_overhead_gb",
+            "total_gb",
+        }
+        self.assertEqual(set(gb_dict.keys()), expected_keys)
+
+    def test_total_equals_sum_of_parts(self):
+        config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True)
+        breakdown = estimate_training_vram(LLAMA_8B, config)
+        parts_sum = (
+            breakdown.model_weights
+            + breakdown.lora_adapters
+            + breakdown.optimizer_states
+            + breakdown.gradients
+            + breakdown.activations
+            + breakdown.cuda_overhead
+        )
+        self.assertEqual(breakdown.total, parts_sum)
+
+    def test_larger_batch_increases_total(self):
+        small = TrainingVramConfig(
+            training_method = "qlora",
+            load_in_4bit = True,
+            batch_size = 1,
+        )
+        large = TrainingVramConfig(
+            training_method = "qlora",
+            load_in_4bit = True,
+            batch_size = 8,
+        )
+        small_v = estimate_training_vram(LLAMA_8B, small)
+        large_v = estimate_training_vram(LLAMA_8B, large)
+        self.assertGreater(large_v.total, small_v.total)
+
+    def test_adamw_fp32_uses_more_optimizer_memory(self):
+        opt8 = TrainingVramConfig(
+            training_method = "full",
+            load_in_4bit = False,
+            optimizer = "adamw_8bit",
+        )
+        opt32 = TrainingVramConfig(
+            training_method = "full",
+            load_in_4bit = False,
+            optimizer = "adamw_torch",
+        )
+        v8 = estimate_training_vram(LLAMA_8B, opt8)
+        v32 = estimate_training_vram(LLAMA_8B, opt32)
+        self.assertAlmostEqual(
+            v32.optimizer_states / v8.optimizer_states, 1.5, delta = 0.1
+        )
+
+
+class TestExtractArchConfigMoE(unittest.TestCase):
+    def test_deepseek_v3_shared_experts(self):
+        hf_config = SimpleNamespace(
+            hidden_size = 7168,
+            num_hidden_layers = 61,
+            num_attention_heads = 128,
+            num_key_value_heads = 128,
+            intermediate_size = 18432,
+            vocab_size = 129280,
+            tie_word_embeddings = False,
+            n_routed_experts = 256,
+            moe_intermediate_size = 2048,
+            n_shared_experts = 1,
+            first_k_dense_replace = 3,
+            q_lora_rank = 1536,
+            kv_lora_rank = 512,
+            qk_nope_head_dim = 128,
+            qk_rope_head_dim = 64,
+            v_head_dim = 128,
+        )
+        arch = extract_arch_config(hf_config)
+        self.assertEqual(arch.num_experts, 256)
+        self.assertEqual(arch.n_shared_experts, 1)
+        self.assertEqual(arch.num_dense_layers, 3)
+        self.assertEqual(arch.q_lora_rank, 1536)
+        self.assertEqual(arch.kv_lora_rank, 512)
+
+    def test_qwen3_moe_decoder_sparse_step(self):
+        hf_config = SimpleNamespace(
+            hidden_size = 2048,
+            num_hidden_layers = 48,
+            num_attention_heads = 32,
+            num_key_value_heads = 4,
+            intermediate_size = 8192,
+            vocab_size = 151936,
+            tie_word_embeddings = True,
+            num_local_experts = 128,
+            moe_intermediate_size = 768,
+            decoder_sparse_step = 1,
+            mlp_only_layers = [],
+        )
+        arch = extract_arch_config(hf_config)
+        self.assertEqual(arch.num_experts, 128)
+        self.assertEqual(arch.num_dense_layers, 0)
+        self.assertIsNone(arch.q_lora_rank)
+
+    def test_qwen3_moe_with_mlp_only_layers(self):
+        hf_config = SimpleNamespace(
+            hidden_size = 2048,
+            num_hidden_layers = 24,
+            num_attention_heads = 16,
+            num_key_value_heads = 4,
+            intermediate_size = 8192,
+            vocab_size = 151936,
+            tie_word_embeddings = True,
+            num_local_experts = 60,
+            moe_intermediate_size = 1408,
+            decoder_sparse_step = 1,
+            mlp_only_layers = [0, 1, 2, 3],
+        )
+        arch = extract_arch_config(hf_config)
+        self.assertEqual(arch.num_dense_layers, 4)
+
+    def test_glm4_moe_first_k_dense(self):
+        hf_config = SimpleNamespace(
+            hidden_size = 4096,
+            num_hidden_layers = 46,
+            num_attention_heads = 96,
+            num_key_value_heads = 8,
+            intermediate_size = 10944,
+            vocab_size = 151552,
+            tie_word_embeddings = False,
+            n_routed_experts = 128,
+            moe_intermediate_size = 1408,
+            n_shared_experts = 1,
+            first_k_dense_replace = 1,
+        )
+        arch = extract_arch_config(hf_config)
+        self.assertEqual(arch.num_dense_layers, 1)
+        self.assertEqual(arch.n_shared_experts, 1)
+
+    def test_gpt_oss_no_moe_intermediate(self):
+        hf_config = SimpleNamespace(
+            hidden_size = 6144,
+            num_hidden_layers = 64,
+            num_attention_heads = 64,
+            num_key_value_heads = 8,
+            intermediate_size = 2880,
+            vocab_size = 200064,
+            tie_word_embeddings = False,
+            num_local_experts = 128,
+        )
+        arch = extract_arch_config(hf_config)
+        self.assertEqual(arch.num_experts, 128)
+        self.assertIsNone(arch.moe_intermediate_size)
+        self.assertEqual(arch.num_dense_layers, 0)
+
+    def test_backward_compat_no_new_fields(self):
+        hf_config = SimpleNamespace(
+            hidden_size = 4096,
+            num_hidden_layers = 32,
+            num_attention_heads = 32,
+            num_key_value_heads = 8,
+            intermediate_size = 14336,
+            vocab_size = 128256,
+            tie_word_embeddings = False,
+        )
+        arch = extract_arch_config(hf_config)
+        self.assertEqual(arch.n_shared_experts, 0)
+        self.assertEqual(arch.num_dense_layers, 0)
+        self.assertIsNone(arch.q_lora_rank)
+
+
+class TestSharedExperts(unittest.TestCase):
+    def test_shared_experts_increase_weight_bytes(self):
+        no_shared = ModelArchConfig(
+            hidden_size = 4096,
+            num_hidden_layers = 32,
+            num_attention_heads = 32,
+            num_key_value_heads = 8,
+            intermediate_size = 14336,
+            vocab_size = 32000,
+            tie_word_embeddings = False,
+            num_experts = 64,
+            moe_intermediate_size = 1407,
+            n_shared_experts = 0,
+        )
+        with_shared = ModelArchConfig(
+            hidden_size = 4096,
+            num_hidden_layers = 32,
+            num_attention_heads = 32,
+            num_key_value_heads = 8,
+            intermediate_size = 14336,
+            vocab_size = 32000,
+            tie_word_embeddings = False,
+            num_experts = 64,
+            moe_intermediate_size = 1407,
+            n_shared_experts = 2,
+        )
+        w_no = compute_model_weights_bytes(no_shared, "full", False)
+        w_yes = compute_model_weights_bytes(with_shared, "full", False)
+        self.assertGreater(w_yes, w_no)
+        delta_per_layer = 4096 * 1407 * 3 * 2
+        expected_delta = delta_per_layer * 32 * 2
+        actual_delta = w_yes - w_no
+        self.assertAlmostEqual(
+            actual_delta, expected_delta, delta = expected_delta * 0.01
+        )
+
+    def test_deepseek_v3_params_in_range(self):
+        total = compute_total_params(DEEPSEEK_V3)
+        total_b = total / 1e9
+        self.assertGreater(total_b, 600)
+        self.assertLess(total_b, 750)
+
+
+class TestMLA(unittest.TestCase):
+    def test_mla_different_from_standard(self):
+        from utils.hardware.vram_estimation import _compute_attn_elements
+
+        mla_arch = DEEPSEEK_V3
+        std_arch = ModelArchConfig(
+            hidden_size = 7168,
+            num_hidden_layers = 61,
+            num_attention_heads = 128,
+            num_key_value_heads = 128,
+            intermediate_size = 18432,
+            vocab_size = 129280,
+        )
+        mla_attn = _compute_attn_elements(mla_arch)
+        std_attn = _compute_attn_elements(std_arch)
+        self.assertNotEqual(mla_attn, std_attn)
+
+    def test_mla_lora_produces_values(self):
+        lora_p = compute_lora_params(DEEPSEEK_V3, 16, ["q_proj", "v_proj", "o_proj"])
+        self.assertGreater(lora_p, 0)
+
+
+class TestDenseMoEMix(unittest.TestCase):
+    def test_dense_layers_change_total(self):
+        all_moe = ModelArchConfig(
+            hidden_size = 4096,
+            num_hidden_layers = 46,
+            num_attention_heads = 96,
+            num_key_value_heads = 8,
+            intermediate_size = 10944,
+            vocab_size = 151552,
+            tie_word_embeddings = False,
+            num_experts = 128,
+            moe_intermediate_size = 1408,
+            n_shared_experts = 1,
+            num_dense_layers = 0,
+        )
+        mixed = ModelArchConfig(
+            hidden_size = 4096,
+            num_hidden_layers = 46,
+            num_attention_heads = 96,
+            num_key_value_heads = 8,
+            intermediate_size = 10944,
+            vocab_size = 151552,
+            tie_word_embeddings = False,
+            num_experts = 128,
+            moe_intermediate_size = 1408,
+            n_shared_experts = 1,
+            num_dense_layers = 1,
+        )
+        w_all = compute_model_weights_bytes(all_moe, "full", False)
+        w_mixed = compute_model_weights_bytes(mixed, "full", False)
+        self.assertNotEqual(w_all, w_mixed)
+
+    def test_glm4_moe_params_reasonable(self):
+        total = compute_total_params(GLM4_MOE)
+        total_b = total / 1e9
+        self.assertGreater(total_b, 80)
+        self.assertLess(total_b, 120)
+
+    def test_qwen3_moe_30b_params_reasonable(self):
+        total = compute_total_params(QWEN3_MOE_30B)
+        total_b = total / 1e9
+        self.assertGreater(total_b, 20)
+        self.assertLess(total_b, 50)
+
+    def test_gpt_oss_uses_intermediate_size(self):
+        total = compute_total_params(GPT_OSS)
+        total_b = total / 1e9
+        self.assertGreater(total_b, 350)
+        self.assertLess(total_b, 500)
+
+    def test_lora_dense_vs_moe_layers_differ(self):
+        all_moe = ModelArchConfig(
+            hidden_size = 4096,
+            num_hidden_layers = 10,
+            num_attention_heads = 32,
+            num_key_value_heads = 8,
+            intermediate_size = 14336,
+            vocab_size = 32000,
+            tie_word_embeddings = False,
+            num_experts = 8,
+            moe_intermediate_size = 1024,
+            num_dense_layers = 0,
+        )
+        mixed = ModelArchConfig(
+            hidden_size = 4096,
+            num_hidden_layers = 10,
+            num_attention_heads = 32,
+            num_key_value_heads = 8,
+            intermediate_size = 14336,
+            vocab_size = 32000,
+            tie_word_embeddings = False,
+            num_experts = 8,
+            moe_intermediate_size = 1024,
+            num_dense_layers = 5,
+        )
+        lora_all = compute_lora_params(
+            all_moe, 16, ["gate_proj", "up_proj", "down_proj"]
+        )
+        lora_mix = compute_lora_params(mixed, 16, ["gate_proj", "up_proj", "down_proj"])
+        self.assertNotEqual(lora_all, lora_mix)
+
+
+if __name__ == "__main__":
+    unittest.main()

studio/backend/utils/datasets/format_conversion.py

@@ -553,13 +553,9 @@ def convert_to_vlm_format(
                         batch_results[idx] = future.result()
                     except Exception as e:
                         failed_count += 1
-                        if failed_count == 1:
-                            print(
-                                f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}"
-                            )
                         if failed_count == 1:
                             logger.info(
-                                f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}"
+                                f"First VLM conversion failure: {type(e).__name__}: {e}"
                             )
 
             converted_list.extend(r for r in batch_results if r is not None)
@@ -583,13 +579,10 @@ def convert_to_vlm_format(
                 converted_list.append(_convert_single_sample(sample))
             except Exception as e:
                 failed_count += 1
-                if failed_count == 1:
-                    # Log the first failure to aid debugging
-                    print(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
                 if failed_count == 1:
                     # Log the first failure to aid debugging
                     logger.info(
-                        f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}"
+                        f"First VLM conversion failure: {type(e).__name__}: {e}"
                     )
             pbar.set_postfix(ok = len(converted_list), failed = failed_count, refresh = False)
         pbar.close()

studio/backend/utils/hardware/VRAM_ESTIMATION.md

@@ -0,0 +1,161 @@
+# VRAM Estimation for Training
+
+```
+Total VRAM = Weights + LoRA Adapters + Optimizer + Gradients + Activations + CUDA Overhead
+```
+
+| Symbol | Meaning |
+|--------|---------|
+| `H` | `hidden_size` |
+| `L` | `num_hidden_layers` |
+| `V` | `vocab_size` |
+| `K` | `(H / num_attention_heads) * num_key_value_heads` |
+| `M` | `intermediate_size` (or `moe_intermediate_size`) |
+| `E` | `num_experts` (1 for dense) |
+| `r` | LoRA rank |
+| `B` | `per_device_train_batch_size` |
+| `S` | `max_seq_length` |
+
+---
+
+## 1. Model Weights
+
+```
+QKVO = (H + K + K + H) * H
+MLP  = H * M * 3 * E  +  (E * H if E > 1 else 0)
+
+Quantizable     = (QKVO + MLP) * L
+Non-quantizable = 2*H*L + V*H + (V*H if not tie_embeddings else 0)
+```
+
+| Mode | Bytes |
+|------|-------|
+| QLoRA 4-bit | `Quantizable * 2 / 3.2 + Non-quantizable * 2` |
+| LoRA / Full fp16 | `(Quantizable + Non-quantizable) * 2` |
+
+The 3.2 factor (`16/5`) accounts for BNB NF4 blockwise scales.
+
+## 2. LoRA Adapters
+
+| Module | A | B |
+|--------|---|---|
+| q_proj | `H×r` | `r×H` |
+| k_proj | `H×r` | `r×K` |
+| v_proj | `H×r` | `r×K` |
+| o_proj | `H×r` | `r×H` |
+| gate_proj | `H×r` | `r×M` |
+| up_proj | `H×r` | `r×M` |
+| down_proj | `M×r` | `r×H` |
+
+MLP modules multiply by `E` for MoE.
+
+```
+LoRA_bytes = sum(A + B per selected module) * L * 2
+```
+
+## 3. Optimizer States (calibrated)
+
+| Optimizer | Bytes/param | Notes |
+|-----------|------------|-------|
+| `adamw_8bit` | 4 | BNB upcasts to fp32 during step |
+| `adamw_torch` | 6 | Fused, no master copy |
+| `paged_adamw_32bit` | 8 | Full fp32 states |
+| `sgd` | 4 | |
+
+Trainable params = all params (Full FT) or LoRA params only.
+
+## 4. Gradients
+
+```
+Gradient_bytes = trainable_params * 2    (fp16, accumulated in-place)
+```
+
+## 5. Activations
+
+Per-layer (from `unsloth_zoo/vllm_utils.py`):
+```
+Per_layer = (S*B*(H+K+K) + S*B*2 + S*B*(M+M)) * 2 * 1.25
+```
+
+| GC Mode | Full FT | LoRA/QLoRA |
+|---------|---------|------------|
+| none | `L` layers | `L` layers |
+| true (HF) | 2.0 | 1.0 |
+| unsloth | 1.5 | 1.0 |
+
+## 6. Floors
+
+Gradients and activations have minimum floors at **15% of model weight memory** to account for autograd overhead, attention score matrices, NCCL buffers, mixed-precision scaling, and PyTorch fragmentation.
+
+```
+gradient_bytes  = max(computed, weights * 0.15)
+activation_bytes = max(computed, weights * 0.15 * B/2)
+```
+
+## 7. CUDA Overhead
+
+**1.4 GB** fixed — CUDA driver + PyTorch runtime, calibrated on RTX 5070 Ti.
+
+## 8. Multi-GPU Overhead
+
+When sharding across multiple GPUs, each additional GPU (beyond the first) contributes only **85%** of its free VRAM to the usable pool. The 15% discount accounts for NCCL all-reduce buffers, PCIe/NVLink transfer overhead, synchronization barriers, and memory fragmentation from non-uniform shard sizes. Calibrated empirically on 2-8 GPU setups with NVLink and PCIe topologies.
+
+```
+usable_gb = free[gpu_0] + sum(free[gpu_i] * 0.85 for i in 1..N)
+```
+
+---
+
+## Reference Table (bsz=2, seq=2048, rank=16, GC=unsloth, adamw_8bit)
+
+| Model | Weights | LoRA | Optim | Grad | Act | CUDA | Total |
+|-------|---------|------|-------|------|-----|------|-------|
+| 0.5B QLoRA | 0.5 | 0.0 | 0.0 | 0.1 | 0.1 | 1.4 | **2.1** |
+| 1B QLoRA | 1.1 | 0.0 | 0.0 | 0.2 | 0.2 | 1.4 | **2.9** |
+| 3B QLoRA | 2.4 | 0.0 | 0.1 | 0.5 | 0.5 | 1.4 | **4.9** |
+| 8B QLoRA | 6.0 | 0.1 | 0.2 | 1.2 | 1.2 | 1.4 | **10.1** |
+| 8B LoRA fp16 | 15.0 | 0.1 | 0.2 | 3.0 | 3.0 | 1.4 | **22.6** |
+| 8B Full FT | 15.0 | — | 29.9 | 15.0 | 3.0 | 1.4 | **64.2** |
+| 32B LoRA fp16 | 61.0 | 0.2 | 0.5 | 12.2 | 12.2 | 1.4 | **87.6** |
+| 72B QLoRA | 45.5 | 0.4 | 0.8 | 9.1 | 9.1 | 1.4 | **66.3** |
+
+## E2E Validation (Llama-3.2-1B, B200 emulating 24GB)
+
+| Config | Estimated | Actual (nvsmi) | Error |
+|--------|----------|----------------|-------|
+| QLoRA bsz=2 seq=512 | 2.55 GB | 2.65 GB | -3.7% |
+| QLoRA bsz=2 seq=2048 | 2.60 GB | 2.65 GB | -1.8% |
+| QLoRA bsz=4 seq=2048 | 2.65 GB | 2.65 GB | +0.0% |
+| LoRA fp16 bsz=2 | 3.84 GB | 3.88 GB | -1.0% |
+| Full FT adamw_8bit | 10.89 GB | 10.80 GB | +0.8% |
+| Full FT adamw_torch | 13.19 GB | 12.93 GB | +2.0% |
+
+*Note: e2e numbers predate the 15% floors, which add safety margin on top.*
+
+---
+
+## Parameter Flow
+
+```
+Frontend -> routes/{training,inference}.py
+         -> prepare_gpu_selection(gpu_ids, model_name, ...)
+            |
+            +-- gpu_ids is explicit (e.g. [5,6,7])
+            |     -> resolve_requested_gpu_ids: validate against parent-visible set
+            |     -> return all requested GPUs (model sharded across all of them)
+            |
+            +-- gpu_ids is None or []
+                  -> auto_select_gpu_ids: estimate VRAM, pick minimum GPUs needed
+                  -> estimate_required_model_memory_gb -> estimate_training_vram
+                  -> greedy selection: rank GPUs by free VRAM, add until model fits
+
+         -> get_device_map(resolved_gpu_ids)
+            -> "balanced" if >1 GPU, "sequential" otherwise
+
+         -> worker subprocess: apply_gpu_ids(resolved_gpu_ids)
+            -> sets CUDA_VISIBLE_DEVICES before torch/CUDA init
+```
+
+Threaded params: `batch_size`, `max_seq_length`, `lora_r`, `target_modules`, `gradient_checkpointing`, `optim`.
+
+Source: `studio/backend/utils/hardware/vram_estimation.py`

studio/backend/utils/hardware/__init__.py

@@ -18,11 +18,31 @@ from .hardware import (
     get_gpu_summary,
     get_package_versions,
     get_gpu_utilization,
+    get_visible_gpu_utilization,
+    get_backend_visible_gpu_info,
     get_physical_gpu_count,
     get_visible_gpu_count,
+    get_parent_visible_gpu_ids,
+    resolve_requested_gpu_ids,
+    estimate_fp16_model_size_bytes,
+    estimate_required_model_memory_gb,
+    auto_select_gpu_ids,
+    prepare_gpu_selection,
     safe_num_proc,
     safe_thread_num_proc,
     dataset_map_num_proc,
+    get_device_map,
+    get_offloaded_device_map_entries,
+    raise_if_offloaded,
+    apply_gpu_ids,
+)
+
+from .vram_estimation import (
+    ModelArchConfig,
+    TrainingVramConfig,
+    VramBreakdown,
+    extract_arch_config,
+    estimate_training_vram,
 )
 
 __all__ = [
@@ -38,9 +58,26 @@ __all__ = [
     "get_gpu_summary",
     "get_package_versions",
     "get_gpu_utilization",
+    "get_visible_gpu_utilization",
+    "get_backend_visible_gpu_info",
     "get_physical_gpu_count",
     "get_visible_gpu_count",
+    "get_parent_visible_gpu_ids",
+    "resolve_requested_gpu_ids",
+    "estimate_fp16_model_size_bytes",
+    "estimate_required_model_memory_gb",
+    "auto_select_gpu_ids",
+    "prepare_gpu_selection",
     "safe_num_proc",
     "safe_thread_num_proc",
     "dataset_map_num_proc",
+    "get_device_map",
+    "get_offloaded_device_map_entries",
+    "raise_if_offloaded",
+    "apply_gpu_ids",
+    "ModelArchConfig",
+    "TrainingVramConfig",
+    "VramBreakdown",
+    "extract_arch_config",
+    "estimate_training_vram",
 ]

studio/backend/utils/hardware/nvidia.py

@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import subprocess
+from typing import Any, Optional
+
+from loggers import get_logger
+
+logger = get_logger(__name__)
+
+
+def _parse_smi_value(raw: str):
+    raw = raw.strip()
+    if not raw or raw == "[N/A]":
+        return None
+    try:
+        return float(raw)
+    except (ValueError, TypeError):
+        return None
+
+
+def _build_gpu_metrics(
+    vram_used_mb,
+    vram_total_mb,
+    power_draw,
+    power_limit,
+    **extra,
+) -> dict[str, Any]:
+    return {
+        **extra,
+        "vram_used_gb": round(vram_used_mb / 1024, 2)
+        if vram_used_mb is not None
+        else None,
+        "vram_total_gb": round(vram_total_mb / 1024, 2)
+        if vram_total_mb is not None
+        else None,
+        "vram_utilization_pct": round((vram_used_mb / vram_total_mb) * 100, 1)
+        if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0
+        else None,
+        "power_draw_w": power_draw,
+        "power_limit_w": power_limit,
+        "power_utilization_pct": round((power_draw / power_limit) * 100, 1)
+        if power_draw is not None and power_limit and power_limit > 0
+        else None,
+    }
+
+
+def _visible_ordinal_map(
+    parent_visible_ids: Optional[list[int]],
+) -> Optional[dict[int, int]]:
+    if parent_visible_ids is None:
+        return None
+    return {gpu_id: ordinal for ordinal, gpu_id in enumerate(parent_visible_ids)}
+
+
+def get_physical_gpu_count() -> Optional[int]:
+    """Return physical GPU count via nvidia-smi, or None on failure."""
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "-L"],
+            capture_output = True,
+            text = True,
+            timeout = 5,
+        )
+        if result.returncode == 0 and result.stdout.strip():
+            return len(result.stdout.strip().splitlines())
+        logger.warning(
+            "nvidia-smi -L returned code %d; caller should fall back to torch",
+            result.returncode,
+        )
+    except Exception as e:
+        logger.warning("nvidia-smi -L failed: %s; caller should fall back to torch", e)
+    return None
+
+
+def get_primary_gpu_utilization() -> dict[str, Any]:
+    try:
+        result = subprocess.run(
+            [
+                "nvidia-smi",
+                "--query-gpu=utilization.gpu,temperature.gpu,"
+                "memory.used,memory.total,power.draw,power.limit",
+                "--format=csv,noheader,nounits",
+            ],
+            capture_output = True,
+            text = True,
+            timeout = 5,
+        )
+    except (OSError, subprocess.TimeoutExpired) as e:
+        logger.warning("nvidia-smi query failed in get_primary_gpu_utilization: %s", e)
+        return {"available": False}
+    if result.returncode != 0 or not result.stdout.strip():
+        return {"available": False}
+
+    first_line = result.stdout.strip().splitlines()[0]
+    parts = [p.strip() for p in first_line.split(",")]
+    if len(parts) < 6:
+        return {"available": False}
+
+    return _build_gpu_metrics(
+        vram_used_mb = _parse_smi_value(parts[2]),
+        vram_total_mb = _parse_smi_value(parts[3]),
+        power_draw = _parse_smi_value(parts[4]),
+        power_limit = _parse_smi_value(parts[5]),
+        available = True,
+        gpu_utilization_pct = _parse_smi_value(parts[0]),
+        temperature_c = _parse_smi_value(parts[1]),
+    )
+
+
+def get_visible_gpu_utilization(
+    parent_visible_ids: Optional[list[int]],
+    parent_cuda_visible_devices: Optional[str] = None,
+) -> dict[str, Any]:
+    # When parent_visible_ids is None (UUID/MIG mask), we cannot safely
+    # map nvidia-smi rows to the process's visible devices. Return empty
+    # instead of exposing all physical GPUs.
+    if parent_visible_ids is None:
+        return {
+            "available": False,
+            "backend_cuda_visible_devices": parent_cuda_visible_devices,
+            "parent_visible_gpu_ids": [],
+            "devices": [],
+            "index_kind": "unresolved",
+        }
+    visible_ordinals = _visible_ordinal_map(parent_visible_ids)
+    try:
+        result = subprocess.run(
+            [
+                "nvidia-smi",
+                "--query-gpu=index,utilization.gpu,temperature.gpu,"
+                "memory.used,memory.total,power.draw,power.limit",
+                "--format=csv,noheader,nounits",
+            ],
+            capture_output = True,
+            text = True,
+            timeout = 5,
+        )
+    except (OSError, subprocess.TimeoutExpired) as e:
+        logger.warning("nvidia-smi query failed in get_visible_gpu_utilization: %s", e)
+        return {
+            "available": False,
+            "backend_cuda_visible_devices": parent_cuda_visible_devices,
+            "parent_visible_gpu_ids": parent_visible_ids or [],
+            "devices": [],
+            "index_kind": "physical",
+        }
+    if result.returncode != 0 or not result.stdout.strip():
+        return {
+            "available": False,
+            "backend_cuda_visible_devices": parent_cuda_visible_devices,
+            "parent_visible_gpu_ids": parent_visible_ids or [],
+            "devices": [],
+            "index_kind": "physical",
+        }
+
+    devices = []
+    for line in result.stdout.strip().splitlines():
+        parts = [p.strip() for p in line.split(",")]
+        if len(parts) < 7:
+            continue
+
+        try:
+            idx = int(parts[0])
+        except (ValueError, TypeError):
+            continue
+
+        if visible_ordinals is not None and idx not in visible_ordinals:
+            continue
+
+        devices.append(
+            _build_gpu_metrics(
+                vram_used_mb = _parse_smi_value(parts[3]),
+                vram_total_mb = _parse_smi_value(parts[4]),
+                power_draw = _parse_smi_value(parts[5]),
+                power_limit = _parse_smi_value(parts[6]),
+                index = idx,
+                index_kind = "physical",
+                visible_ordinal = (
+                    visible_ordinals[idx]
+                    if visible_ordinals is not None
+                    else len(devices)
+                ),
+                gpu_utilization_pct = _parse_smi_value(parts[1]),
+                temperature_c = _parse_smi_value(parts[2]),
+            )
+        )
+
+    return {
+        "available": len(devices) > 0,
+        "backend_cuda_visible_devices": parent_cuda_visible_devices,
+        "parent_visible_gpu_ids": parent_visible_ids or [],
+        "devices": devices,
+        "index_kind": "physical",
+    }
+
+
+def get_backend_visible_gpu_info(
+    parent_visible_ids: Optional[list[int]],
+    backend_cuda_visible_devices: Optional[str],
+) -> dict[str, Any]:
+    # When parent_visible_ids is None (UUID/MIG mask), we cannot safely
+    # map nvidia-smi rows to the process's visible devices.
+    if parent_visible_ids is None:
+        return {
+            "available": False,
+            "backend_cuda_visible_devices": backend_cuda_visible_devices,
+            "parent_visible_gpu_ids": [],
+            "devices": [],
+            "index_kind": "unresolved",
+        }
+    visible_ordinals = _visible_ordinal_map(parent_visible_ids)
+    try:
+        result = subprocess.run(
+            [
+                "nvidia-smi",
+                "--query-gpu=index,name,memory.total",
+                "--format=csv,noheader,nounits",
+            ],
+            capture_output = True,
+            text = True,
+            timeout = 10,
+        )
+    except (OSError, subprocess.TimeoutExpired) as e:
+        logger.warning("nvidia-smi query failed in get_backend_visible_gpu_info: %s", e)
+        return {
+            "available": False,
+            "backend_cuda_visible_devices": backend_cuda_visible_devices,
+            "parent_visible_gpu_ids": parent_visible_ids or [],
+            "devices": [],
+            "index_kind": "physical",
+        }
+    if result.returncode != 0:
+        return {
+            "available": False,
+            "backend_cuda_visible_devices": backend_cuda_visible_devices,
+            "parent_visible_gpu_ids": parent_visible_ids or [],
+            "devices": [],
+            "index_kind": "physical",
+        }
+
+    devices = []
+    for line in result.stdout.strip().splitlines():
+        parts = [p.strip() for p in line.split(",")]
+        if len(parts) < 3:
+            continue
+        try:
+            idx = int(parts[0])
+        except (ValueError, TypeError):
+            continue
+        if visible_ordinals is not None and idx not in visible_ordinals:
+            continue
+        # Use split with limit to handle GPU names containing commas
+        name = parts[1] if len(parts) == 3 else ", ".join(parts[1:-1])
+        try:
+            mem_total_mb = int(parts[-1])
+        except (ValueError, TypeError):
+            continue
+        devices.append(
+            {
+                "index": idx,
+                "index_kind": "physical",
+                "visible_ordinal": (
+                    visible_ordinals[idx]
+                    if visible_ordinals is not None
+                    else len(devices)
+                ),
+                "name": name,
+                "memory_total_gb": round(mem_total_mb / 1024, 2),
+            }
+        )
+
+    return {
+        "available": len(devices) > 0,
+        "backend_cuda_visible_devices": backend_cuda_visible_devices,
+        "parent_visible_gpu_ids": parent_visible_ids or [],
+        "devices": devices,
+        "index_kind": "physical",
+    }

studio/backend/utils/hardware/vram_estimation.py

@@ -0,0 +1,501 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+"""
+Training VRAM estimation.
+
+Total VRAM = weights + LoRA adapters + optimizer states + gradients
+           + activations + CUDA overhead.
+Activation formula from unsloth_zoo/vllm_utils.py.
+All constants empirically calibrated against Llama-3.2-1B on B200.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Dict, Optional
+
+QUANT_4BIT_FACTOR = 16 / 5
+CUDA_OVERHEAD_BYTES = int(1.4 * 1024**3)  # calibrated on RTX 5070 Ti
+
+DEFAULT_TARGET_MODULES = [
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "down_proj",
+]
+
+# Empirically calibrated bytes/param — see VRAM_ESTIMATION.md for rationale.
+OPTIMIZER_BYTES_PER_PARAM: Dict[str, int] = {
+    "adamw_8bit": 4,  # BNB upcasts to fp32 during step
+    "paged_adamw_8bit": 4,
+    "adamw_bnb_8bit": 4,
+    "paged_adamw_32bit": 8,
+    "adamw_torch": 6,  # fused, no master copy
+    "adamw_torch_fused": 6,
+    "sgd": 4,
+}
+
+# (full_ft_multiplier, lora_multiplier) — fraction of num_layers.
+# LoRA: frozen base layers skip activation storage, but you always need
+# at least ~1 layer in flight during backprop recomputation.
+GC_LAYER_MULTIPLIERS = {
+    "none": (None, None),
+    "true": (2.0, 1.0),
+    "unsloth": (1.5, 1.0),
+}
+
+
+@dataclass
+class ModelArchConfig:
+    hidden_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    intermediate_size: int
+    vocab_size: int
+    tie_word_embeddings: bool = True
+    num_experts: Optional[int] = None
+    moe_intermediate_size: Optional[int] = None
+    n_shared_experts: int = 0
+    num_dense_layers: int = 0
+    q_lora_rank: Optional[int] = None
+    kv_lora_rank: Optional[int] = None
+    qk_nope_head_dim: Optional[int] = None
+    qk_rope_head_dim: Optional[int] = None
+    v_head_dim: Optional[int] = None
+
+
+@dataclass
+class TrainingVramConfig:
+    training_method: str = "qlora"
+    batch_size: int = 4
+    max_seq_length: int = 2048
+    lora_rank: int = 16
+    target_modules: list = field(default_factory = lambda: list(DEFAULT_TARGET_MODULES))
+    gradient_checkpointing: str = "unsloth"
+    optimizer: str = "adamw_8bit"
+    load_in_4bit: bool = True
+
+
+@dataclass
+class VramBreakdown:
+    model_weights: int
+    lora_adapters: int
+    optimizer_states: int
+    gradients: int
+    activations: int
+    cuda_overhead: int
+    # The computed (formula-based) activation cost before floors.
+    # This is the true per-layer cost that doesn't shard across GPUs.
+    activations_computed: int = 0
+
+    @property
+    def total(self) -> int:
+        return (
+            self.model_weights
+            + self.lora_adapters
+            + self.optimizer_states
+            + self.gradients
+            + self.activations
+            + self.cuda_overhead
+        )
+
+    def min_gpu_vram(self, n_gpus: int) -> int:
+        """Minimum VRAM a single GPU needs: its shard + non-shardable costs.
+
+        Weights/LoRA/optimizer/gradients shard across GPUs.
+        The computed activation cost does NOT shard (one GPU runs the layer).
+        The floor portion (activations - computed) is overhead that shards.
+        """
+        shardable = (
+            self.model_weights
+            + self.lora_adapters
+            + self.optimizer_states
+            + self.gradients
+            + (self.activations - self.activations_computed)  # floor overhead shards
+        )
+        per_gpu_fixed = self.activations_computed + self.cuda_overhead
+        return shardable // max(n_gpus, 1) + per_gpu_fixed
+
+    def to_gb_dict(self) -> Dict[str, float]:
+        return {
+            "model_weights_gb": round(self.model_weights / (1024**3), 3),
+            "lora_adapters_gb": round(self.lora_adapters / (1024**3), 3),
+            "optimizer_states_gb": round(self.optimizer_states / (1024**3), 3),
+            "gradients_gb": round(self.gradients / (1024**3), 3),
+            "activations_gb": round(self.activations / (1024**3), 3),
+            "cuda_overhead_gb": round(self.cuda_overhead / (1024**3), 3),
+            "total_gb": round(self.total / (1024**3), 3),
+        }
+
+
+def _compute_num_dense_layers(text_config, total_layers: int) -> int:
+    """Count how many layers use dense MLP instead of MoE."""
+    first_k = getattr(text_config, "first_k_dense_replace", None)
+    if first_k is not None:
+        return min(int(first_k), total_layers)
+
+    sparse_step = getattr(text_config, "decoder_sparse_step", None)
+    mlp_only = getattr(text_config, "mlp_only_layers", None) or []
+    if sparse_step is not None and sparse_step > 0:
+        mlp_only_set = set(mlp_only)
+        moe_count = sum(
+            1
+            for i in range(total_layers)
+            if i not in mlp_only_set and (i + 1) % sparse_step == 0
+        )
+        return total_layers - moe_count
+
+    return 0
+
+
+def extract_arch_config(hf_config) -> Optional[ModelArchConfig]:
+    text_config = getattr(hf_config, "text_config", None) or hf_config
+
+    hidden_size = getattr(text_config, "hidden_size", None)
+    num_layers = getattr(text_config, "num_hidden_layers", None)
+    num_heads = getattr(text_config, "num_attention_heads", None)
+    intermediate_size = getattr(text_config, "intermediate_size", None)
+    vocab_size = getattr(text_config, "vocab_size", None)
+
+    if isinstance(intermediate_size, (list, tuple)):
+        intermediate_size = intermediate_size[0] if intermediate_size else None
+    if intermediate_size is None and hidden_size is not None:
+        intermediate_size = hidden_size * 4
+
+    if not all(
+        v is not None
+        for v in (hidden_size, num_layers, num_heads, intermediate_size, vocab_size)
+    ):
+        return None
+    if num_heads <= 0:
+        return None
+
+    num_kv_heads = getattr(text_config, "num_key_value_heads", num_heads)
+
+    num_experts = None
+    for attr in ("num_local_experts", "num_experts", "n_routed_experts"):
+        num_experts = getattr(text_config, attr, None)
+        if num_experts is not None:
+            break
+
+    moe_intermediate = getattr(text_config, "moe_intermediate_size", None)
+    n_shared_experts = getattr(text_config, "n_shared_experts", None) or 0
+
+    num_dense_layers = 0
+    if num_experts is not None and num_experts > 1:
+        num_dense_layers = _compute_num_dense_layers(text_config, num_layers)
+
+    q_lora_rank = getattr(text_config, "q_lora_rank", None)
+    kv_lora_rank = getattr(text_config, "kv_lora_rank", None)
+    qk_nope_head_dim = getattr(text_config, "qk_nope_head_dim", None)
+    qk_rope_head_dim = getattr(text_config, "qk_rope_head_dim", None)
+    v_head_dim = getattr(text_config, "v_head_dim", None)
+
+    return ModelArchConfig(
+        hidden_size = hidden_size,
+        num_hidden_layers = num_layers,
+        num_attention_heads = num_heads,
+        num_key_value_heads = num_kv_heads,
+        intermediate_size = intermediate_size,
+        vocab_size = vocab_size,
+        tie_word_embeddings = getattr(text_config, "tie_word_embeddings", True),
+        num_experts = num_experts,
+        moe_intermediate_size = moe_intermediate,
+        n_shared_experts = n_shared_experts,
+        num_dense_layers = num_dense_layers,
+        q_lora_rank = q_lora_rank,
+        kv_lora_rank = kv_lora_rank,
+        qk_nope_head_dim = qk_nope_head_dim,
+        qk_rope_head_dim = qk_rope_head_dim,
+        v_head_dim = v_head_dim,
+    )
+
+
+def _get_kv_size(arch: ModelArchConfig) -> int:
+    return (arch.hidden_size // arch.num_attention_heads) * arch.num_key_value_heads
+
+
+def _get_mlp_size(arch: ModelArchConfig) -> int:
+    if arch.moe_intermediate_size is not None:
+        return arch.moe_intermediate_size
+    return arch.intermediate_size
+
+
+def _get_num_experts(arch: ModelArchConfig) -> int:
+    return arch.num_experts if arch.num_experts and arch.num_experts > 1 else 1
+
+
+def _compute_attn_elements(arch: ModelArchConfig) -> int:
+    """Attention weight elements per layer."""
+    hd = arch.hidden_size
+    if arch.q_lora_rank is not None:
+        nh = arch.num_attention_heads
+        qk_head = arch.qk_nope_head_dim + arch.qk_rope_head_dim
+        q_a = hd * arch.q_lora_rank
+        q_b = arch.q_lora_rank * (nh * qk_head)
+        kv_a = hd * (arch.kv_lora_rank + arch.qk_rope_head_dim)
+        kv_b = arch.kv_lora_rank * (nh * (arch.qk_nope_head_dim + arch.v_head_dim))
+        o = (nh * arch.v_head_dim) * hd
+        norms = arch.q_lora_rank + arch.kv_lora_rank
+        return q_a + q_b + kv_a + kv_b + o + norms
+    kv_size = _get_kv_size(arch)
+    return (hd + kv_size + kv_size + hd) * hd
+
+
+def _compute_dense_mlp_elements(arch: ModelArchConfig) -> int:
+    return arch.hidden_size * arch.intermediate_size * 3
+
+
+def _compute_moe_mlp_elements(arch: ModelArchConfig) -> int:
+    hd = arch.hidden_size
+    mlp_size = _get_mlp_size(arch)
+    n_experts = _get_num_experts(arch)
+    return hd * mlp_size * 3 * (n_experts + arch.n_shared_experts) + n_experts * hd
+
+
+def _compute_layer_elements(arch: ModelArchConfig):
+    """Return (total_quantizable, layernorms_per_layer, embed, lm_head) element counts.
+
+    total_quantizable is summed across ALL layers (not per-layer).
+    """
+    hd = arch.hidden_size
+    n_layers = arch.num_hidden_layers
+    n_experts = _get_num_experts(arch)
+
+    attn_total = _compute_attn_elements(arch) * n_layers
+
+    if n_experts > 1:
+        n_dense = arch.num_dense_layers
+        n_moe = n_layers - n_dense
+        mlp_total = (
+            _compute_moe_mlp_elements(arch) * n_moe
+            + _compute_dense_mlp_elements(arch) * n_dense
+        )
+    else:
+        mlp_total = _compute_dense_mlp_elements(arch) * n_layers
+
+    layernorms = 2 * hd
+    embed_tokens = arch.vocab_size * hd
+    lm_head = 0 if arch.tie_word_embeddings else arch.vocab_size * hd
+    return attn_total + mlp_total, layernorms, embed_tokens, lm_head
+
+
+def compute_model_weights_bytes(
+    arch: ModelArchConfig,
+    training_method: str,
+    load_in_4bit: bool,
+) -> int:
+    total_quantizable, layernorms, embed_tokens, lm_head = _compute_layer_elements(arch)
+    n_layers = arch.num_hidden_layers
+    non_quantizable = layernorms * n_layers + embed_tokens + lm_head
+
+    if training_method == "qlora" and load_in_4bit:
+        return int(total_quantizable * 2 / QUANT_4BIT_FACTOR + non_quantizable * 2)
+
+    return int((total_quantizable + non_quantizable) * 2)
+
+
+def compute_total_params(arch: ModelArchConfig) -> int:
+    total_quantizable, layernorms, embed_tokens, lm_head = _compute_layer_elements(arch)
+    n_layers = arch.num_hidden_layers
+    return total_quantizable + layernorms * n_layers + embed_tokens + lm_head
+
+
+def _lora_attn_elements(
+    arch: ModelArchConfig,
+    r: int,
+    target_modules: list,
+) -> int:
+    hd = arch.hidden_size
+    if arch.q_lora_rank is not None:
+        # MLA: q_proj->q_b, k_proj->kv_a, v_proj->kv_b, o_proj->o
+        nh = arch.num_attention_heads
+        qk_head = arch.qk_nope_head_dim + arch.qk_rope_head_dim
+        kv_out = nh * (arch.qk_nope_head_dim + arch.v_head_dim)
+        o_in = nh * arch.v_head_dim
+        dims = {
+            "q_proj": (arch.q_lora_rank, nh * qk_head),
+            "k_proj": (hd, arch.kv_lora_rank + arch.qk_rope_head_dim),
+            "v_proj": (arch.kv_lora_rank, kv_out),
+            "o_proj": (o_in, hd),
+        }
+    else:
+        kv_size = _get_kv_size(arch)
+        dims = {
+            "q_proj": (hd, hd),
+            "k_proj": (hd, kv_size),
+            "v_proj": (hd, kv_size),
+            "o_proj": (hd, hd),
+        }
+    total = 0
+    for name, (in_dim, out_dim) in dims.items():
+        if name in target_modules:
+            total += in_dim * r + r * out_dim
+    return total
+
+
+def _lora_mlp_elements(
+    hd: int,
+    mlp_size: int,
+    r: int,
+    target_modules: list,
+    expert_mult: int,
+) -> int:
+    module_ab = {
+        "gate_proj": (hd * r, r * mlp_size),
+        "up_proj": (hd * r, r * mlp_size),
+        "down_proj": (mlp_size * r, r * hd),
+    }
+    total = 0
+    for name, (a, b) in module_ab.items():
+        if name in target_modules:
+            total += (a + b) * expert_mult
+    return total
+
+
+def compute_lora_params(
+    arch: ModelArchConfig,
+    lora_rank: int,
+    target_modules: list,
+) -> int:
+    hd = arch.hidden_size
+    r = lora_rank
+    n_layers = arch.num_hidden_layers
+    n_experts = _get_num_experts(arch)
+
+    attn_total = _lora_attn_elements(arch, r, target_modules) * n_layers
+
+    if n_experts > 1:
+        n_dense = arch.num_dense_layers
+        n_moe = n_layers - n_dense
+        # Include shared experts alongside routed experts
+        moe_expert_mult = n_experts + arch.n_shared_experts
+        moe_mlp = _lora_mlp_elements(
+            hd,
+            _get_mlp_size(arch),
+            r,
+            target_modules,
+            moe_expert_mult,
+        )
+        dense_mlp = _lora_mlp_elements(
+            hd,
+            arch.intermediate_size,
+            r,
+            target_modules,
+            1,
+        )
+        mlp_total = moe_mlp * n_moe + dense_mlp * n_dense
+    else:
+        mlp_total = (
+            _lora_mlp_elements(
+                hd,
+                arch.intermediate_size,
+                r,
+                target_modules,
+                1,
+            )
+            * n_layers
+        )
+
+    return attn_total + mlp_total
+
+
+def compute_lora_adapter_bytes(lora_params: int) -> int:
+    return lora_params * 2
+
+
+def compute_optimizer_bytes(trainable_params: int, optimizer: str) -> int:
+    optimizer_key = optimizer.lower().replace("-", "_")
+    bytes_per_param = OPTIMIZER_BYTES_PER_PARAM.get(optimizer_key, 4)
+    return trainable_params * bytes_per_param
+
+
+def compute_gradient_bytes(trainable_params: int) -> int:
+    return trainable_params * 2
+
+
+def compute_activation_bytes(
+    arch: ModelArchConfig,
+    batch_size: int,
+    seq_len: int,
+    gradient_checkpointing: str,
+    is_lora: bool = False,
+) -> int:
+    hd = arch.hidden_size
+    kv_size = _get_kv_size(arch)
+    mlp_size = _get_mlp_size(arch)
+    bsz = batch_size
+    n_layers = arch.num_hidden_layers
+
+    activation_qkv = seq_len * bsz * (hd + kv_size + kv_size)
+    residual_memory = (seq_len * bsz) * 2
+    activation_mlp = seq_len * bsz * (mlp_size + mlp_size)
+
+    per_layer_bytes = (activation_qkv + residual_memory + activation_mlp) * 2
+    per_layer_bytes = int(per_layer_bytes * 1.25)
+
+    gc_key = gradient_checkpointing.lower()
+    gc_entry = GC_LAYER_MULTIPLIERS.get(gc_key, (None, None))
+    full_ft_mult, lora_mult = gc_entry
+    gc_multiplier = lora_mult if is_lora else full_ft_mult
+
+    if gc_multiplier is None:
+        effective_layers = n_layers
+    else:
+        effective_layers = gc_multiplier
+
+    return int(per_layer_bytes * effective_layers)
+
+
+def estimate_training_vram(
+    arch: ModelArchConfig,
+    config: TrainingVramConfig,
+) -> VramBreakdown:
+    method = config.training_method.lower()
+    is_lora = method in ("qlora", "lora")
+    load_in_4bit = config.load_in_4bit or method == "qlora"
+
+    model_weights = compute_model_weights_bytes(arch, method, load_in_4bit)
+
+    lora_params = 0
+    lora_adapter_bytes = 0
+    if is_lora:
+        lora_params = compute_lora_params(
+            arch,
+            config.lora_rank,
+            config.target_modules,
+        )
+        lora_adapter_bytes = compute_lora_adapter_bytes(lora_params)
+
+    trainable_params = lora_params if is_lora else compute_total_params(arch)
+    optimizer_bytes = compute_optimizer_bytes(trainable_params, config.optimizer)
+    gradient_bytes = max(
+        compute_gradient_bytes(trainable_params),
+        int(model_weights * 0.15),
+    )
+    activations_computed = compute_activation_bytes(
+        arch,
+        config.batch_size,
+        config.max_seq_length,
+        config.gradient_checkpointing,
+        is_lora = is_lora,
+    )
+    activation_bytes = max(
+        activations_computed,
+        int(model_weights * 0.15 * (config.batch_size / 2)),
+    )
+
+    return VramBreakdown(
+        model_weights = model_weights,
+        lora_adapters = lora_adapter_bytes,
+        optimizer_states = optimizer_bytes,
+        gradients = gradient_bytes,
+        activations = activation_bytes,
+        cuda_overhead = CUDA_OVERHEAD_BYTES,
+        activations_computed = activations_computed,
+    )

studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx

@@ -18,10 +18,20 @@ import {
   TooltipContent,
   TooltipTrigger,
 } from "@/components/ui/tooltip";
-import { deleteCachedModel, listCachedGguf, listCachedModels, listGgufVariants, listLocalModels } from "@/features/chat/api/chat-api";
-import type { CachedGgufRepo, CachedModelRepo, LocalModelInfo } from "@/features/chat/api/chat-api";
-import type { GgufVariantDetail } from "@/features/chat/types/api";
 import { usePlatformStore } from "@/config/env";
+import {
+  deleteCachedModel,
+  listCachedGguf,
+  listCachedModels,
+  listGgufVariants,
+  listLocalModels,
+} from "@/features/chat/api/chat-api";
+import type {
+  CachedGgufRepo,
+  CachedModelRepo,
+  LocalModelInfo,
+} from "@/features/chat/api/chat-api";
+import type { GgufVariantDetail } from "@/features/chat/types/api";
 import {
   useDebouncedValue,
   useGpuInfo,
@@ -35,7 +45,13 @@ import { checkVramFit, estimateLoadingVram } from "@/lib/vram";
 import { Search01Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { Trash2Icon } from "lucide-react";
-import { useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from "react";
+import {
+  type ReactNode,
+  useCallback,
+  useEffect,
+  useMemo,
+  useState,
+} from "react";
 import { toast } from "sonner";
 import type {
   LoraModelOption,
@@ -135,7 +151,7 @@ function ModelRow({
   if (vramTooltipText) {
     return (
       <Tooltip>
-        <TooltipTrigger asChild>{content}</TooltipTrigger>
+        <TooltipTrigger asChild={true}>{content}</TooltipTrigger>
         <TooltipContent side="left" className="max-w-xs break-all">
           {label}
           <span className="block text-[10px] mt-1">{vramTooltipText}</span>
@@ -147,7 +163,7 @@ function ModelRow({
   if (tooltipText) {
     return (
       <Tooltip>
-        <TooltipTrigger asChild>{content}</TooltipTrigger>
+        <TooltipTrigger asChild={true}>{content}</TooltipTrigger>
         <TooltipContent side="left" className="max-w-xs break-all">
           {tooltipText}
         </TooltipContent>
@@ -192,7 +208,9 @@ function GgufVariantExpander({
       })
       .catch((err) => {
         if (canceled) return;
-        setError(err instanceof Error ? err.message : "Failed to load variants");
+        setError(
+          err instanceof Error ? err.message : "Failed to load variants",
+        );
       })
       .finally(() => {
         if (!canceled) setLoading(false);
@@ -204,7 +222,9 @@ function GgufVariantExpander({
   }, [repoId]);
 
   // Covers Unix absolute (/), Windows drive (C:\, D:/), UNC (\\server), relative (./, ../), tilde (~/)
-  const isLocalPath = /^(\/|\.{1,2}[\\\/]|~[\\\/]|[A-Za-z]:[\\\/]|\\\\)/.test(repoId);
+  const isLocalPath = /^(\/|\.{1,2}[\\\/]|~[\\\/]|[A-Za-z]:[\\\/]|\\\\)/.test(
+    repoId,
+  );
 
   const handleVariantClick = useCallback(
     (quant: string, downloaded?: boolean, sizeBytes?: number) => {
@@ -223,13 +243,13 @@ function GgufVariantExpander({
   //   fits  = model <= 0.7 * total GPU memory
   //   tight = model > 0.7 * GPU but <= 0.7 * GPU + 0.7 * system RAM (--fit uses CPU offload)
   //   oom   = model > 0.7 * GPU + 0.7 * system RAM
-  const gpuBudgetGb = (gpuGb ?? 0) * 0.70;
-  const totalBudgetGb = gpuBudgetGb + (systemRamGb ?? 0) * 0.70;
+  const gpuBudgetGb = (gpuGb ?? 0) * 0.7;
+  const totalBudgetGb = gpuBudgetGb + (systemRamGb ?? 0) * 0.7;
 
   const getGgufFit = useCallback(
     (sizeBytes: number): "fits" | "tight" | "oom" => {
       if (!gpuGb || gpuGb <= 0) return "fits";
-      const gb = sizeBytes / (1024 ** 3);
+      const gb = sizeBytes / 1024 ** 3;
       if (gb <= 0 || gb <= gpuBudgetGb) return "fits";
       if (gb <= totalBudgetGb) return "tight";
       return "oom";
@@ -242,7 +262,8 @@ function GgufVariantExpander({
   const effectiveRecommended = useMemo(() => {
     if (!variants || !gpuGb || gpuGb <= 0) return defaultVariant;
     const defaultV = variants.find((v) => v.quant === defaultVariant);
-    if (defaultV && getGgufFit(defaultV.size_bytes) !== "oom") return defaultVariant;
+    if (defaultV && getGgufFit(defaultV.size_bytes) !== "oom")
+      return defaultVariant;
     // Default is OOM -- pick largest non-OOM variant (best quality that fits)
     const fitting = variants.filter((v) => getGgufFit(v.size_bytes) !== "oom");
     if (fitting.length > 0) {
@@ -276,7 +297,9 @@ function GgufVariantExpander({
       // fits: largest first (best quality that fits in GPU)
       // tight/OOM: smallest first (closest to fitting, fastest to run)
       const fitsInGpu = aTier === 0 || aTier === 2;
-      return fitsInGpu ? b.size_bytes - a.size_bytes : a.size_bytes - b.size_bytes;
+      return fitsInGpu
+        ? b.size_bytes - a.size_bytes
+        : a.size_bytes - b.size_bytes;
     });
   }, [variants, effectiveRecommended, getGgufFit]);
 
@@ -290,9 +313,7 @@ function GgufVariantExpander({
   }
 
   if (error) {
-    return (
-      <div className="px-5 py-2 text-xs text-destructive">{error}</div>
-    );
+    return <div className="px-5 py-2 text-xs text-destructive">{error}</div>;
   }
 
   if (!sortedVariants || sortedVariants.length === 0) {
@@ -321,7 +342,9 @@ function GgufVariantExpander({
           <div key={v.filename} className="flex items-center gap-0.5">
             <button
               type="button"
-              onClick={() => handleVariantClick(v.quant, v.downloaded, v.size_bytes)}
+              onClick={() =>
+                handleVariantClick(v.quant, v.downloaded, v.size_bytes)
+              }
               className={cn(
                 "flex min-w-0 flex-1 items-center justify-between gap-2 rounded-md px-2.5 py-1 text-left text-sm transition-colors hover:bg-accent",
               )}
@@ -340,10 +363,14 @@ function GgufVariantExpander({
               </span>
               <span className="flex items-center gap-1.5 shrink-0">
                 {oom && (
-                  <span className="text-[9px] font-medium text-red-400">OOM</span>
+                  <span className="text-[9px] font-medium text-red-400">
+                    OOM
+                  </span>
                 )}
                 {tight && (
-                  <span className="text-[9px] font-medium text-amber-400">TIGHT</span>
+                  <span className="text-[9px] font-medium text-amber-400">
+                    TIGHT
+                  </span>
                 )}
                 <span className="text-[10px] text-muted-foreground">
                   {formatBytes(v.size_bytes)}
@@ -353,7 +380,10 @@ function GgufVariantExpander({
             {v.downloaded && onDeleteVariant && (
               <button
                 type="button"
-                onClick={(e) => { e.stopPropagation(); onDeleteVariant(v.quant); }}
+                onClick={(e) => {
+                  e.stopPropagation();
+                  onDeleteVariant(v.quant);
+                }}
                 className="shrink-0 rounded-md p-1 text-muted-foreground/60 transition-colors hover:bg-destructive/10 hover:text-destructive"
               >
                 <Trash2Icon className="size-3" />
@@ -384,6 +414,7 @@ function extractParamLabel(id: string): string | undefined {
 let _cachedGgufCache: CachedGgufRepo[] = [];
 let _cachedModelsCache: CachedModelRepo[] = [];
 let _lmStudioCache: LocalModelInfo[] = [];
+let _customFolderCache: LocalModelInfo[] = [];
 
 /** Sort LM Studio models with unsloth publisher first. */
 function sortLmStudio(models: LocalModelInfo[]): LocalModelInfo[] {
@@ -391,7 +422,9 @@ function sortLmStudio(models: LocalModelInfo[]): LocalModelInfo[] {
     const aUnsloth = (a.model_id ?? "").startsWith("unsloth/") ? 0 : 1;
     const bUnsloth = (b.model_id ?? "").startsWith("unsloth/") ? 0 : 1;
     if (aUnsloth !== bUnsloth) return aUnsloth - bUnsloth;
-    return (a.model_id ?? a.display_name).localeCompare(b.model_id ?? b.display_name);
+    return (a.model_id ?? a.display_name).localeCompare(
+      b.model_id ?? b.display_name,
+    );
   });
 }
 
@@ -409,9 +442,8 @@ export function HubModelPicker({
   const gpu = useGpuInfo();
   const [query, setQuery] = useState("");
   const debouncedQuery = useDebouncedValue(query);
-  const { results, isLoading, isLoadingMore, fetchMore } = useHfModelSearch(
-    debouncedQuery,
-  );
+  const { results, isLoading, isLoadingMore, fetchMore } =
+    useHfModelSearch(debouncedQuery);
 
   // Track which GGUF repo is expanded for variant selection
   const [expandedGguf, setExpandedGguf] = useState<string | null>(null);
@@ -422,38 +454,75 @@ export function HubModelPicker({
 
   // Cached (already downloaded) repos -- use module-level cache so
   // re-mounting the popover does not flash an empty "Downloaded" section.
-  const [cachedGguf, setCachedGguf] = useState<CachedGgufRepo[]>(_cachedGgufCache);
-  const [cachedModels, setCachedModels] = useState<CachedModelRepo[]>(_cachedModelsCache);
-  const alreadyCached = _cachedGgufCache.length > 0 || _cachedModelsCache.length > 0;
+  const [cachedGguf, setCachedGguf] =
+    useState<CachedGgufRepo[]>(_cachedGgufCache);
+  const [cachedModels, setCachedModels] =
+    useState<CachedModelRepo[]>(_cachedModelsCache);
+  const alreadyCached =
+    _cachedGgufCache.length > 0 || _cachedModelsCache.length > 0;
   const [cachedReady, setCachedReady] = useState(alreadyCached);
 
   // LM Studio local models -- module-level cache so re-mounting the
   // popover does not flash an empty section (same pattern as GGUF/models).
-  const [lmStudioModels, setLmStudioModels] = useState<LocalModelInfo[]>(_lmStudioCache);
+  const [lmStudioModels, setLmStudioModels] =
+    useState<LocalModelInfo[]>(_lmStudioCache);
+  const [customFolderModels, setCustomFolderModels] =
+    useState<LocalModelInfo[]>(_customFolderCache);
 
-  const refreshCachedLists = useCallback(() => {
-    listCachedGguf().then((v) => { _cachedGgufCache = v; setCachedGguf(v); }).catch(() => {});
-    listCachedModels().then((v) => { _cachedModelsCache = v; setCachedModels(v); }).catch(() => {});
-    listLocalModels().then((res) => {
-      const next = sortLmStudio(res.models.filter((m) => m.source === "lmstudio"));
-      _lmStudioCache = next;
-      setLmStudioModels(next);
-    }).catch(() => {});
+  const refreshLocalModelsList = useCallback(() => {
+    listLocalModels()
+      .then((res) => {
+        const lm = sortLmStudio(
+          res.models.filter((m) => m.source === "lmstudio"),
+        );
+        _lmStudioCache = lm;
+        setLmStudioModels(lm);
+        const cf = res.models.filter((m) => m.source === "custom");
+        _customFolderCache = cf;
+        setCustomFolderModels(cf);
+      })
+      .catch(() => {});
   }, []);
 
+  const refreshCachedLists = useCallback(() => {
+    listCachedGguf()
+      .then((v) => {
+        _cachedGgufCache = v;
+        setCachedGguf(v);
+      })
+      .catch(() => {});
+    listCachedModels()
+      .then((v) => {
+        _cachedModelsCache = v;
+        setCachedModels(v);
+      })
+      .catch(() => {});
+    refreshLocalModelsList();
+  }, [refreshLocalModelsList]);
+
   useEffect(() => {
-    // Always refresh LM Studio models (not gated by alreadyCached)
-    listLocalModels().then((res) => {
-      const next = sortLmStudio(res.models.filter((m) => m.source === "lmstudio"));
-      _lmStudioCache = next;
-      setLmStudioModels(next);
-    }).catch(() => {});
+    // Always refresh LM Studio + custom folder models (not gated by alreadyCached)
+    refreshLocalModelsList();
 
     if (alreadyCached) return;
     let done = 0;
-    const check = () => { if (++done >= 2) setCachedReady(true); };
-    listCachedGguf().then((v) => { _cachedGgufCache = v; setCachedGguf(v); }).catch(() => {}).finally(check);
-    listCachedModels().then((v) => { _cachedModelsCache = v; setCachedModels(v); }).catch(() => {}).finally(check);
+    const check = () => {
+      if (++done >= 2) setCachedReady(true);
+    };
+    listCachedGguf()
+      .then((v) => {
+        _cachedGgufCache = v;
+        setCachedGguf(v);
+      })
+      .catch(() => {})
+      .finally(check);
+    listCachedModels()
+      .then((v) => {
+        _cachedModelsCache = v;
+        setCachedModels(v);
+      })
+      .catch(() => {})
+      .finally(check);
   }, [alreadyCached]);
 
   const handleDeleteConfirm = useCallback(async () => {
@@ -468,7 +537,9 @@ export function HubModelPicker({
       toast.success(`Deleted ${variant ? `${repoId} ${variant}` : repoId}`);
       refreshCachedLists();
     } catch (err) {
-      toast.error(err instanceof Error ? err.message : "Failed to delete model");
+      toast.error(
+        err instanceof Error ? err.message : "Failed to delete model",
+      );
     } finally {
       setDeleting(false);
       setDeleteTarget(null);
@@ -504,12 +575,18 @@ export function HubModelPicker({
   // Infinite scroll paging for the recommended section
   const [recommendedPage, setRecommendedPage] = useState(1);
   // Reset page when the underlying list changes
-  useEffect(() => { setRecommendedPage(1); }, [models, chatOnly]);
+  useEffect(() => {
+    setRecommendedPage(1);
+  }, [models, chatOnly]);
 
   const visibleRecommendedIds = useMemo(() => {
     const hubStartIndex = recommendedIds.findIndex((id) => !isGgufRepo(id));
-    const allGguf = hubStartIndex === -1 ? recommendedIds : recommendedIds.slice(0, hubStartIndex);
-    const allHub = hubStartIndex === -1 ? [] : recommendedIds.slice(hubStartIndex);
+    const allGguf =
+      hubStartIndex === -1
+        ? recommendedIds
+        : recommendedIds.slice(0, hubStartIndex);
+    const allHub =
+      hubStartIndex === -1 ? [] : recommendedIds.slice(hubStartIndex);
     // Interleave in chunks of 4: [4 gguf, 4 hub, 4 gguf, 4 hub, ...]
     const result: string[] = [];
     for (let p = 0; p < recommendedPage; p++) {
@@ -519,12 +596,8 @@ export function HubModelPicker({
     return result;
   }, [recommendedIds, recommendedPage]);
 
-  const hasMoreRecommended = visibleRecommendedIds.length < recommendedIds.length;
-
-  // Fetch VRAM info for the full pool once (recommendedIds is stable across
-  // page increments) so we don't re-fetch on every scroll.
-  const { paramCountById: recommendedParamCountById } =
-    useRecommendedModelVram(recommendedIds);
+  const hasMoreRecommended =
+    visibleRecommendedIds.length < recommendedIds.length;
 
   const showHfSection = debouncedQuery.trim().length > 0;
 
@@ -535,8 +608,22 @@ export function HubModelPicker({
     return recommendedIds.filter((id) => normalizeForSearch(id).includes(q));
   }, [showHfSection, debouncedQuery, recommendedIds]);
 
+  // Fetch VRAM info for visible models, plus any models surfaced by a search
+  // query so that filtered recommended models also show VRAM badges.
+  // Skip GGUF repos: they have no safetensors metadata and the render layer
+  // already shows a static "GGUF" badge instead of VRAM data.
+  const idsForVram = useMemo(() => {
+    const ids = showHfSection
+      ? [...new Set([...visibleRecommendedIds, ...filteredRecommendedIds])]
+      : visibleRecommendedIds;
+    return ids.filter((id) => !isGgufRepo(id));
+  }, [visibleRecommendedIds, showHfSection, filteredRecommendedIds]);
+  const { paramCountById: recommendedParamCountById } =
+    useRecommendedModelVram(idsForVram);
+
   const recommendedSet = useMemo(
-    () => new Set(showHfSection ? filteredRecommendedIds : visibleRecommendedIds),
+    () =>
+      new Set(showHfSection ? filteredRecommendedIds : visibleRecommendedIds),
     [showHfSection, filteredRecommendedIds, visibleRecommendedIds],
   );
 
@@ -602,15 +689,25 @@ export function HubModelPicker({
       }
     }
     return map;
-  }, [showHfSection, filteredRecommendedIds, visibleRecommendedIds, recommendedParamCountById, gpu]);
+  }, [
+    showHfSection,
+    filteredRecommendedIds,
+    visibleRecommendedIds,
+    recommendedParamCountById,
+    gpu,
+  ]);
 
-  const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, results.length);
+  const { scrollRef, sentinelRef } = useInfiniteScroll(
+    fetchMore,
+    results.length,
+  );
 
   // Sentinel + IntersectionObserver for recommended infinite scroll.
   // We disconnect after each fire so the observer doesn't loop while
   // React re-renders; the effect re-creates it on the next page.
   // Uses a callback ref for the sentinel so we detect mount/unmount reliably.
-  const [recommendedSentinel, setRecommendedSentinel] = useState<HTMLDivElement | null>(null);
+  const [recommendedSentinel, setRecommendedSentinel] =
+    useState<HTMLDivElement | null>(null);
   const recommendedSentinelRef = useCallback((node: HTMLDivElement | null) => {
     setRecommendedSentinel(node);
   }, []);
@@ -629,7 +726,10 @@ export function HubModelPicker({
     );
     // Small delay so the browser finishes layout after the previous page render
     const timer = setTimeout(() => obs.observe(recommendedSentinel), 100);
-    return () => { clearTimeout(timer); obs.disconnect(); };
+    return () => {
+      clearTimeout(timer);
+      obs.disconnect();
+    };
   }, [recommendedSentinel, hasMoreRecommended, recommendedPage, scrollRef]);
 
   /** Handle clicking a model row — GGUF repos expand, others load directly. */
@@ -668,9 +768,13 @@ export function HubModelPicker({
           {!cachedReady && !showHfSection ? (
             <div className="flex items-center gap-2 px-5 py-3">
               <Spinner className="size-3 text-muted-foreground" />
-              <span className="text-xs text-muted-foreground">Loading models…</span>
+              <span className="text-xs text-muted-foreground">
+                Loading models…
+              </span>
             </div>
-          ) : !showHfSection && (cachedGguf.length > 0 || (!chatOnly && cachedModels.length > 0)) ? (
+          ) : !showHfSection &&
+            (cachedGguf.length > 0 ||
+              (!chatOnly && cachedModels.length > 0)) ? (
             <>
               <ListLabel>{"\uD83E\uDDA5"} Downloaded</ListLabel>
               {cachedGguf.map((c) => (
@@ -687,32 +791,46 @@ export function HubModelPicker({
                       repoId={c.repo_id}
                       onSelect={onSelect}
                       gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
-                      systemRamGb={gpu.available ? gpu.systemRamAvailableGb : undefined}
-                      onDeleteVariant={(quant) => setDeleteTarget(`${c.repo_id}::${quant}`)}
+                      systemRamGb={
+                        gpu.available ? gpu.systemRamAvailableGb : undefined
+                      }
+                      onDeleteVariant={(quant) =>
+                        setDeleteTarget(`${c.repo_id}::${quant}`)
+                      }
                     />
                   )}
                 </div>
               ))}
-              {!chatOnly && cachedModels.map((c) => (
-                <div key={c.repo_id} className="flex items-center gap-0.5">
-                  <div className="min-w-0 flex-1">
-                    <ModelRow
-                      label={c.repo_id}
-                      meta={formatBytes(c.size_bytes)}
-                      selected={value === c.repo_id}
-                      onClick={() => onSelect(c.repo_id, { source: "hub", isLora: false, isDownloaded: true })}
-                      vramStatus={null}
-                    />
+              {!chatOnly &&
+                cachedModels.map((c) => (
+                  <div key={c.repo_id} className="flex items-center gap-0.5">
+                    <div className="min-w-0 flex-1">
+                      <ModelRow
+                        label={c.repo_id}
+                        meta={formatBytes(c.size_bytes)}
+                        selected={value === c.repo_id}
+                        onClick={() =>
+                          onSelect(c.repo_id, {
+                            source: "hub",
+                            isLora: false,
+                            isDownloaded: true,
+                          })
+                        }
+                        vramStatus={null}
+                      />
+                    </div>
+                    <button
+                      type="button"
+                      onClick={(e) => {
+                        e.stopPropagation();
+                        setDeleteTarget(c.repo_id);
+                      }}
+                      className="shrink-0 rounded-md p-1.5 text-muted-foreground/60 transition-colors hover:bg-destructive/10 hover:text-destructive"
+                    >
+                      <Trash2Icon className="size-3.5" />
+                    </button>
                   </div>
-                  <button
-                    type="button"
-                    onClick={(e) => { e.stopPropagation(); setDeleteTarget(c.repo_id); }}
-                    className="shrink-0 rounded-md p-1.5 text-muted-foreground/60 transition-colors hover:bg-destructive/10 hover:text-destructive"
-                  >
-                    <Trash2Icon className="size-3.5" />
-                  </button>
-                </div>
-              ))}
+                ))}
             </>
           ) : null}
 
@@ -725,13 +843,21 @@ export function HubModelPicker({
                   <div key={m.id}>
                     <ModelRow
                       label={m.model_id ?? m.display_name}
-                      meta={isGguf || m.path.endsWith(".gguf") ? "GGUF" : "Local"}
+                      meta={
+                        isGguf || m.path.endsWith(".gguf") ? "GGUF" : "Local"
+                      }
                       selected={value === m.id}
                       onClick={() => {
                         if (isGguf) {
-                          setExpandedGguf((prev) => (prev === m.id ? null : m.id));
+                          setExpandedGguf((prev) =>
+                            prev === m.id ? null : m.id,
+                          );
                         } else {
-                          onSelect(m.id, { source: "local", isLora: false, isDownloaded: true });
+                          onSelect(m.id, {
+                            source: "local",
+                            isLora: false,
+                            isDownloaded: true,
+                          });
                         }
                       }}
                       vramStatus={null}
@@ -741,7 +867,54 @@ export function HubModelPicker({
                         repoId={m.id}
                         onSelect={onSelect}
                         gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
-                        systemRamGb={gpu.available ? gpu.systemRamAvailableGb : undefined}
+                        systemRamGb={
+                          gpu.available ? gpu.systemRamAvailableGb : undefined
+                        }
+                      />
+                    )}
+                  </div>
+                );
+              })}
+            </>
+          ) : null}
+
+          {!showHfSection && customFolderModels.length > 0 ? (
+            <>
+              <ListLabel>Custom Folders</ListLabel>
+              {customFolderModels.map((m) => {
+                const isGguf =
+                  isGgufRepo(m.id) ||
+                  isGgufRepo(m.display_name) ||
+                  m.path.endsWith(".gguf");
+                return (
+                  <div key={m.id}>
+                    <ModelRow
+                      label={m.model_id ?? m.display_name}
+                      meta={isGguf ? "GGUF" : "Local"}
+                      selected={value === m.id}
+                      onClick={() => {
+                        if (isGguf) {
+                          setExpandedGguf((prev) =>
+                            prev === m.id ? null : m.id,
+                          );
+                        } else {
+                          onSelect(m.id, {
+                            source: "local",
+                            isLora: false,
+                            isDownloaded: true,
+                          });
+                        }
+                      }}
+                      vramStatus={null}
+                    />
+                    {expandedGguf === m.id && (
+                      <GgufVariantExpander
+                        repoId={m.id}
+                        onSelect={onSelect}
+                        gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
+                        systemRamGb={
+                          gpu.available ? gpu.systemRamAvailableGb : undefined
+                        }
                       />
                     )}
                   </div>
@@ -767,16 +940,25 @@ export function HubModelPicker({
                         meta={
                           isGgufRepo(id)
                             ? "GGUF"
-                            : vram?.detail ?? extractParamLabel(id)
+                            : (vram?.detail ?? extractParamLabel(id))
                         }
                         selected={value === id}
                         onClick={() => handleModelClick(id)}
-                        vramStatus={isGgufRepo(id) ? null : vram?.status ?? null}
+                        vramStatus={
+                          isGgufRepo(id) ? null : (vram?.status ?? null)
+                        }
                         vramEst={isGgufRepo(id) ? undefined : vram?.est}
                         gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
                       />
                       {expandedGguf === id && (
-                        <GgufVariantExpander repoId={id} onSelect={onSelect} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} systemRamGb={gpu.available ? gpu.systemRamAvailableGb : undefined} />
+                        <GgufVariantExpander
+                          repoId={id}
+                          onSelect={onSelect}
+                          gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
+                          systemRamGb={
+                            gpu.available ? gpu.systemRamAvailableGb : undefined
+                          }
+                        />
                       )}
                     </div>
                   );
@@ -805,16 +987,25 @@ export function HubModelPicker({
                       meta={
                         isGgufRepo(id)
                           ? "GGUF"
-                          : vram?.detail ?? extractParamLabel(id)
+                          : (vram?.detail ?? extractParamLabel(id))
                       }
                       selected={value === id}
                       onClick={() => handleModelClick(id)}
-                      vramStatus={isGgufRepo(id) ? null : vram?.status ?? null}
+                      vramStatus={
+                        isGgufRepo(id) ? null : (vram?.status ?? null)
+                      }
                       vramEst={isGgufRepo(id) ? undefined : vram?.est}
                       gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
                     />
                     {expandedGguf === id && (
-                      <GgufVariantExpander repoId={id} onSelect={onSelect} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} systemRamGb={gpu.available ? gpu.systemRamAvailableGb : undefined} />
+                      <GgufVariantExpander
+                        repoId={id}
+                        onSelect={onSelect}
+                        gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
+                        systemRamGb={
+                          gpu.available ? gpu.systemRamAvailableGb : undefined
+                        }
+                      />
                     )}
                   </div>
                 );
@@ -824,7 +1015,9 @@ export function HubModelPicker({
 
           {showHfSection ? (
             <>
-              {(hfIds.length > 0 || isLoading) && <ListLabel>Hugging Face</ListLabel>}
+              {(hfIds.length > 0 || isLoading) && (
+                <ListLabel>Hugging Face</ListLabel>
+              )}
               {hfIds.length === 0 && !isLoading ? (
                 filteredRecommendedIds.length === 0 ? (
                   <div className="px-2.5 py-2 text-xs text-muted-foreground">
@@ -841,16 +1034,25 @@ export function HubModelPicker({
                         meta={
                           isGgufRepo(id)
                             ? "GGUF"
-                            : metricsById.get(id) ?? extractParamLabel(id)
+                            : (metricsById.get(id) ?? extractParamLabel(id))
                         }
                         selected={value === id}
                         onClick={() => handleModelClick(id)}
-                        vramStatus={isGgufRepo(id) ? null : vram?.status ?? null}
+                        vramStatus={
+                          isGgufRepo(id) ? null : (vram?.status ?? null)
+                        }
                         vramEst={isGgufRepo(id) ? undefined : vram?.est}
                         gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
                       />
                       {expandedGguf === id && (
-                        <GgufVariantExpander repoId={id} onSelect={onSelect} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} systemRamGb={gpu.available ? gpu.systemRamAvailableGb : undefined} />
+                        <GgufVariantExpander
+                          repoId={id}
+                          onSelect={onSelect}
+                          gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
+                          systemRamGb={
+                            gpu.available ? gpu.systemRamAvailableGb : undefined
+                          }
+                        />
                       )}
                     </div>
                   );
@@ -867,12 +1069,23 @@ export function HubModelPicker({
         </div>
       </div>
 
-      <AlertDialog open={deleteTarget !== null} onOpenChange={(open) => { if (!open && !deleting) setDeleteTarget(null); }}>
+      <AlertDialog
+        open={deleteTarget !== null}
+        onOpenChange={(open) => {
+          if (!open && !deleting) setDeleteTarget(null);
+        }}
+      >
         <AlertDialogContent size="sm">
           <AlertDialogHeader>
             <AlertDialogTitle>Delete cached model?</AlertDialogTitle>
             <AlertDialogDescription>
-              This will remove <span className="font-medium text-foreground">{deleteTarget?.includes("::") ? `${deleteTarget.split("::")[0]} (${deleteTarget.split("::")[1]})` : deleteTarget}</span> from disk. You can re-download it later.
+              This will remove{" "}
+              <span className="font-medium text-foreground">
+                {deleteTarget?.includes("::")
+                  ? `${deleteTarget.split("::")[0]} (${deleteTarget.split("::")[1]})`
+                  : deleteTarget}
+              </span>{" "}
+              from disk. You can re-download it later.
             </AlertDialogDescription>
           </AlertDialogHeader>
           <AlertDialogFooter>
@@ -880,7 +1093,10 @@ export function HubModelPicker({
             <AlertDialogAction
               variant="destructive"
               disabled={deleting}
-              onClick={(e) => { e.preventDefault(); handleDeleteConfirm(); }}
+              onClick={(e) => {
+                e.preventDefault();
+                handleDeleteConfirm();
+              }}
             >
               {deleting ? "Deleting..." : "Yes"}
             </AlertDialogAction>
@@ -909,7 +1125,8 @@ export function LoraModelPicker({
       loraModels
         .map((model) => ({
           ...model,
-          baseModel: model.baseModel || model.description || "Unknown base model",
+          baseModel:
+            model.baseModel || model.description || "Unknown base model",
         }))
         .sort((a, b) => {
           const baseCmp = a.baseModel.localeCompare(b.baseModel);
@@ -933,7 +1150,9 @@ export function LoraModelPicker({
     const out = new Map<string, LoraModelOption[]>();
 
     for (const model of normalized) {
-      const searchText = normalizeForSearch(`${model.name} ${model.baseModel} ${model.id}`);
+      const searchText = normalizeForSearch(
+        `${model.name} ${model.baseModel} ${model.id}`,
+      );
       if (needle && !searchText.includes(needle)) continue;
 
       const key = model.baseModel || "Unknown base model";
@@ -981,15 +1200,27 @@ export function LoraModelPicker({
                   const isExported = adapter.source === "exported";
                   const isMerged = adapter.exportType === "merged";
                   const isGguf = adapter.exportType === "gguf";
-                  const isLocalGgufDir = isLocal && (isGgufRepo(adapter.id) || isGgufRepo(adapter.name));
+                  const isLocalGgufDir =
+                    isLocal &&
+                    (isGgufRepo(adapter.id) || isGgufRepo(adapter.name));
                   const tag = isLocal
-                    ? isLocalGgufDir ? "GGUF" : "Local"
+                    ? isLocalGgufDir
+                      ? "GGUF"
+                      : "Local"
                     : isGguf
                       ? "GGUF"
                       : isExported
-                        ? isMerged ? "Merged" : "LoRA"
+                        ? isMerged
+                          ? "Merged"
+                          : "LoRA"
                         : "LoRA";
-                  const meta = isLocal ? (isLocalGgufDir ? "GGUF" : "Local") : isExported ? `${tag} · Exported` : tag;
+                  const meta = isLocal
+                    ? isLocalGgufDir
+                      ? "GGUF"
+                      : "Local"
+                    : isExported
+                      ? `${tag} · Exported`
+                      : tag;
                   return (
                     <div key={adapter.id}>
                       <ModelRow
@@ -998,17 +1229,26 @@ export function LoraModelPicker({
                         selected={value === adapter.id}
                         onClick={() => {
                           if (isLocalGgufDir) {
-                            setExpandedGguf((prev) => (prev === adapter.id ? null : adapter.id));
+                            setExpandedGguf((prev) =>
+                              prev === adapter.id ? null : adapter.id,
+                            );
                           } else {
                             onSelect(adapter.id, {
-                              source: isLocal ? "local" : isExported ? "exported" : "lora",
+                              source: isLocal
+                                ? "local"
+                                : isExported
+                                  ? "exported"
+                                  : "lora",
                               isLora: !isLocal && !isMerged && !isGguf,
+                              isDownloaded: true,
                             });
                           }
                         }}
                         tooltipText={
                           <>
-                            <span className="block break-words">{adapter.name}</span>
+                            <span className="block break-words">
+                              {adapter.name}
+                            </span>
                             <span className="block mt-1 text-[10px] text-muted-foreground break-all">
                               {adapter.id}
                             </span>
@@ -1020,7 +1260,9 @@ export function LoraModelPicker({
                           repoId={adapter.id}
                           onSelect={onSelect}
                           gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
-                          systemRamGb={gpu.available ? gpu.systemRamAvailableGb : undefined}
+                          systemRamGb={
+                            gpu.available ? gpu.systemRamAvailableGb : undefined
+                          }
                         />
                       )}
                     </div>

studio/frontend/src/components/navbar.tsx

@@ -6,6 +6,11 @@ import {
   HoverCardContent,
   HoverCardTrigger,
 } from "@/components/ui/hover-card";
+import {
+  Collapsible,
+  CollapsibleContent,
+  CollapsibleTrigger,
+} from "@/components/ui/collapsible";
 import { AnimatedThemeToggler } from "@/components/ui/animated-theme-toggler";
 import {
   Sheet,
@@ -16,20 +21,25 @@ import {
 } from "@/components/ui/sheet";
 import { cn } from "@/lib/utils";
 import {
+  ArrowReloadHorizontalIcon,
   ArrowRight01Icon,
   Cancel01Icon,
   Book03Icon,
   BubbleChatIcon,
   ChefHatIcon,
+  Copy01Icon,
   CursorInfo02Icon,
   PackageIcon,
+  Tick02Icon,
   ZapIcon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
+import { copyToClipboard } from "@/lib/copy-to-clipboard";
 import { useTrainingRuntimeStore } from "@/features/training";
 import { usePlatformStore } from "@/config/env";
 import { Link, useRouterState } from "@tanstack/react-router";
-import { motion } from "motion/react";
+import { AnimatePresence, motion, useReducedMotion } from "motion/react";
+import type { ReactElement } from "react";
 import { useEffect, useRef, useState } from "react";
 import { TOUR_OPEN_EVENT } from "@/features/tour";
 import { ShutdownDialog } from "@/components/shutdown-dialog";
@@ -41,6 +51,185 @@ const NAV_ITEMS = [
   { label: "Chat", href: "/chat", icon: BubbleChatIcon, enabled: true },
 ];
 
+const STUDIO_UPDATE_CMD = "unsloth studio update";
+const STUDIO_UPDATE_FALLBACK_UNIX_CMD =
+  "curl -fsSL https://unsloth.ai/install.sh | sh";
+const STUDIO_UPDATE_FALLBACK_WINDOWS_CMD =
+  "irm https://unsloth.ai/install.ps1 | iex";
+
+type UpdateShell = "windows" | "unix";
+
+function getDefaultUpdateShell(deviceType: string): UpdateShell {
+  return deviceType === "windows" ? "windows" : "unix";
+}
+
+function getStudioUpdateInstructionLine(shell: UpdateShell): string {
+  return shell === "windows" ? "Open PowerShell and run:" : "Open Terminal and run:";
+}
+
+function CopyableCommand({
+  command,
+  copyLabel,
+}: {
+  command: string;
+  copyLabel: string;
+}): ReactElement {
+  const [copied, setCopied] = useState(false);
+  const timerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+
+  useEffect(() => {
+    return () => {
+      if (timerRef.current) {
+        clearTimeout(timerRef.current);
+      }
+    };
+  }, []);
+
+  const handleCopy = () => {
+    if (!copyToClipboard(command)) {
+      return;
+    }
+    setCopied(true);
+    if (timerRef.current) {
+      clearTimeout(timerRef.current);
+    }
+    timerRef.current = setTimeout(() => setCopied(false), 2000);
+  };
+
+  return (
+    <div className="flex min-w-0 items-stretch overflow-hidden rounded-md border border-border bg-muted/40">
+      <input
+        type="text"
+        readOnly
+        value={command}
+        className="min-w-0 flex-1 bg-transparent px-2 py-1.5 font-mono text-[11px] text-foreground outline-none"
+        title={command}
+        aria-label={`${copyLabel} text`}
+      />
+      <button
+        type="button"
+        onClick={handleCopy}
+        className="flex shrink-0 items-center justify-center border-l border-border px-2 text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
+        title={copied ? "Copied" : "Copy command"}
+        aria-label={copied ? `${copyLabel} copied` : `Copy ${copyLabel}`}
+      >
+        {copied ? (
+          <HugeiconsIcon icon={Tick02Icon} className="size-4 text-emerald-600" />
+        ) : (
+          <HugeiconsIcon icon={Copy01Icon} className="size-4" />
+        )}
+      </button>
+    </div>
+  );
+}
+
+function UpdateStudioInstructions({
+  className,
+  defaultShell,
+  showTitle = true,
+}: {
+  className?: string;
+  defaultShell: UpdateShell;
+  showTitle?: boolean;
+}): ReactElement {
+  const [shell, setShell] = useState<UpdateShell>(defaultShell);
+  const prefersReducedMotion = useReducedMotion();
+  const windows = shell === "windows";
+  const fadeTransition = prefersReducedMotion
+    ? { duration: 0 }
+    : { duration: 0.16, ease: [0.165, 0.84, 0.44, 1] as const };
+  const fadeInitial = prefersReducedMotion ? { opacity: 1 } : { opacity: 0, y: 2 };
+  const fadeAnimate = { opacity: 1, y: 0 };
+  const fadeExit = prefersReducedMotion ? { opacity: 1 } : { opacity: 0, y: -2 };
+
+  useEffect(() => {
+    setShell(defaultShell);
+  }, [defaultShell]);
+
+  return (
+    <div className={cn("flex flex-col gap-3", className)}>
+      <div
+        className={cn(
+          "flex items-center gap-3",
+          showTitle ? "justify-between" : "justify-start",
+        )}
+      >
+        {showTitle ? (
+          <p className="shrink-0 whitespace-nowrap text-sm font-semibold font-heading">
+            Update Unsloth Studio
+          </p>
+        ) : null}
+        <div className="flex shrink-0 items-center gap-0.5 text-[11px]">
+          <button
+            type="button"
+            onClick={() => setShell("windows")}
+            className={cn(
+              "px-0.5 py-0.5 font-medium transition-colors",
+              windows
+                ? "text-foreground"
+                : "text-muted-foreground hover:text-emerald-600",
+            )}
+            aria-pressed={windows}
+          >
+            Windows
+          </button>
+          <span className="text-border">/</span>
+          <button
+            type="button"
+            onClick={() => setShell("unix")}
+            className={cn(
+              "px-0.5 py-0.5 font-medium transition-colors",
+              !windows
+                ? "text-foreground"
+                : "text-muted-foreground hover:text-emerald-600",
+            )}
+            aria-pressed={!windows}
+          >
+            macOS/Linux
+          </button>
+        </div>
+      </div>
+      <AnimatePresence mode="wait" initial={false}>
+        <motion.p
+          key={`instruction-${shell}`}
+          initial={fadeInitial}
+          animate={fadeAnimate}
+          exit={fadeExit}
+          transition={fadeTransition}
+          className="text-xs text-muted-foreground leading-relaxed"
+        >
+          {getStudioUpdateInstructionLine(shell)}
+        </motion.p>
+      </AnimatePresence>
+      <CopyableCommand command={STUDIO_UPDATE_CMD} copyLabel="update command" />
+      <p className="text-xs text-muted-foreground leading-relaxed">
+        If that fails or unsloth studio update is unavailable, run:
+      </p>
+      <AnimatePresence mode="wait" initial={false}>
+        <motion.div
+          key={`fallback-${shell}`}
+          initial={fadeInitial}
+          animate={fadeAnimate}
+          exit={fadeExit}
+          transition={fadeTransition}
+        >
+          <CopyableCommand
+            command={
+              windows
+                ? STUDIO_UPDATE_FALLBACK_WINDOWS_CMD
+                : STUDIO_UPDATE_FALLBACK_UNIX_CMD
+            }
+            copyLabel="fallback command"
+          />
+        </motion.div>
+      </AnimatePresence>
+      <p className="text-xs text-muted-foreground leading-relaxed">
+        Restart Studio after updating for changes to take effect.
+      </p>
+    </div>
+  );
+}
+
 function getTourId(pathname: string): "studio" | "chat" | "export" | null {
   if (pathname === "/studio") return "studio";
   if (pathname === "/chat") return "chat";
@@ -52,9 +241,12 @@ export function Navbar() {
   const pathname = useRouterState({ select: (s) => s.location.pathname });
   const isTrainingRunning = useTrainingRuntimeStore((s) => s.isTrainingRunning);
   const [mobileOpen, setMobileOpen] = useState(false);
+  const [mobileUpdateOpen, setMobileUpdateOpen] = useState(false);
   const [shutdownOpen, setShutdownOpen] = useState(false);
 
+  const deviceType = usePlatformStore((s) => s.deviceType);
   const chatOnly = usePlatformStore((s) => s.isChatOnly());
+  const defaultUpdateShell = getDefaultUpdateShell(deviceType);
 
   // Warn before closing the tab only when training is running (data loss risk).
   // We store the handler in a ref so removeUnloadHandler() can clean it up
@@ -236,6 +428,25 @@ export function Navbar() {
             <span className="text-sm font-medium">Tour</span>
           </button>
 
+          <HoverCard openDelay={200} closeDelay={100}>
+            <HoverCardTrigger asChild={true}>
+              <button
+                type="button"
+                className="flex h-9 items-center gap-1.5 rounded-md px-3 text-sm font-medium text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
+                aria-label="How to update Unsloth Studio"
+              >
+                <HugeiconsIcon icon={ArrowReloadHorizontalIcon} className="size-4" />
+                Update
+              </button>
+            </HoverCardTrigger>
+            <HoverCardContent align="end" className="w-[22.5rem] p-0">
+              <UpdateStudioInstructions
+                className="p-4"
+                defaultShell={defaultUpdateShell}
+              />
+            </HoverCardContent>
+          </HoverCard>
+
           <button
             type="button"
             onClick={() => setShutdownOpen(true)}
@@ -259,7 +470,13 @@ export function Navbar() {
               <HugeiconsIcon icon={CursorInfo02Icon} className="size-4" />
             </button>
           ) : null}
-          <Sheet open={mobileOpen} onOpenChange={setMobileOpen}>
+          <Sheet
+            open={mobileOpen}
+            onOpenChange={(open) => {
+              setMobileOpen(open);
+              if (!open) setMobileUpdateOpen(false);
+            }}
+          >
             <SheetTrigger asChild={true}>
               <button
                 type="button"
@@ -273,7 +490,7 @@ export function Navbar() {
               <SheetHeader>
                 <SheetTitle>Navigate</SheetTitle>
               </SheetHeader>
-              <div className="mt-6 flex flex-col gap-2">
+              <div className="mt-6 flex max-h-[calc(100dvh-8rem)] flex-col gap-2 overflow-y-auto pr-1">
                 {NAV_ITEMS.filter((item) => item.enabled).map((item) => {
                   const active = pathname === item.href;
                   const disabledByTraining =
@@ -312,7 +529,7 @@ export function Navbar() {
                   href="https://unsloth.ai/docs"
                   target="_blank"
                   rel="noopener noreferrer"
-                  className="mt-2 flex items-center gap-2 rounded-md border border-border px-3 py-2 text-sm font-medium text-foreground hover:bg-accent"
+                  className="mt-3 flex items-center gap-2 rounded-md border border-border px-3 py-2 text-sm font-medium text-foreground hover:bg-accent"
                   onClick={() => setMobileOpen(false)}
                 >
                   <HugeiconsIcon icon={Book03Icon} className="size-4" />
@@ -331,9 +548,40 @@ export function Navbar() {
                     Start tour
                   </button>
                 ) : null}
+                <Collapsible
+                  open={mobileUpdateOpen}
+                  onOpenChange={setMobileUpdateOpen}
+                  className="rounded-md border border-border"
+                >
+                  <CollapsibleTrigger asChild={true}>
+                    <button
+                      type="button"
+                      className="flex w-full items-center justify-between rounded-md px-3 py-2 text-left text-sm font-medium text-foreground transition-colors hover:bg-accent"
+                      aria-label="Toggle update instructions"
+                    >
+                      <span className="flex items-center gap-2">
+                        <HugeiconsIcon icon={ArrowReloadHorizontalIcon} className="size-4" />
+                        Update Unsloth Studio
+                      </span>
+                      <HugeiconsIcon
+                        icon={ArrowRight01Icon}
+                        className={cn(
+                          "size-4 text-muted-foreground transition-transform",
+                          mobileUpdateOpen && "rotate-90",
+                        )}
+                      />
+                    </button>
+                  </CollapsibleTrigger>
+                  <CollapsibleContent className="border-t border-border p-3 pt-2">
+                    <UpdateStudioInstructions
+                      defaultShell={defaultUpdateShell}
+                      showTitle={false}
+                    />
+                  </CollapsibleContent>
+                </Collapsible>
                 <button
                   type="button"
-                  className="flex items-center gap-2 rounded-md border border-border px-3 py-2 text-left text-sm font-medium text-foreground hover:bg-accent"
+                  className="mt-3 flex items-center gap-2 rounded-md border border-border px-3 py-2 text-left text-sm font-medium text-foreground hover:bg-accent"
                   onClick={() => {
                     setMobileOpen(false);
                     setShutdownOpen(true);

studio/frontend/src/config/training.ts

@@ -90,10 +90,17 @@ export const LR_SCHEDULER_OPTIONS: ReadonlyArray<{ value: string; label: string
   { value: "cosine", label: "Cosine" },
 ];
 
+/**
+ * Method-aware learning rate defaults.
+ * Backend mirrors these in the YAML configs under studio/backend/assets/configs/.
+ */
+export const LR_DEFAULT_LORA = 2e-4;
+export const LR_DEFAULT_FULL = 2e-5;
+
 export const DEFAULT_HYPERPARAMS = {
   epochs: 3,
   contextLength: 2048,
-  learningRate: 2e-4,
+  learningRate: LR_DEFAULT_LORA,
   optimizerType: "adamw_8bit",
   lrSchedulerType: "linear",
   loraRank: 16,
@@ -102,7 +109,7 @@ export const DEFAULT_HYPERPARAMS = {
   loraVariant: "lora" as const,
   batchSize: 4,
   gradientAccumulation: 8,
-  weightDecay: 0.01,
+  weightDecay: 0.001,
   warmupSteps: 5,
   maxSteps: 60,
   saveSteps: 0,

studio/frontend/src/features/chat/api/chat-adapter.ts

@@ -306,6 +306,7 @@ async function autoLoadSmallestModel(): Promise<boolean> {
             }
             useChatRuntimeStore.setState({
               ggufContextLength: loadResp.context_length ?? 131072,
+              ggufMaxContextLength: loadResp.max_context_length ?? loadResp.context_length ?? 131072,
               supportsReasoning: loadResp.supports_reasoning ?? false,
               reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
               reasoningEnabled: loadResp.supports_reasoning ?? false,
@@ -392,6 +393,7 @@ async function autoLoadSmallestModel(): Promise<boolean> {
       }
       useChatRuntimeStore.setState({
         ggufContextLength: loadResp.context_length ?? 131072,
+        ggufMaxContextLength: loadResp.max_context_length ?? loadResp.context_length ?? 131072,
         supportsReasoning: loadResp.supports_reasoning ?? false,
         reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
         reasoningEnabled: loadResp.supports_reasoning ?? false,

studio/frontend/src/features/chat/api/chat-api.ts

@@ -129,7 +129,7 @@ export interface LocalModelInfo {
   id: string;
   display_name: string;
   path: string;
-  source: "models_dir" | "hf_cache" | "lmstudio";
+  source: "models_dir" | "hf_cache" | "lmstudio" | "custom";
   model_id?: string | null;
   updated_at?: number | null;
 }
@@ -174,6 +174,34 @@ export async function deleteCachedModel(repoId: string, variant?: string): Promi
   await parseJsonOrThrow<unknown>(response);
 }
 
+export interface ScanFolderInfo {
+  id: number;
+  path: string;
+  created_at: string;
+}
+
+export async function listScanFolders(): Promise<ScanFolderInfo[]> {
+  const response = await authFetch("/api/models/scan-folders");
+  const data = await parseJsonOrThrow<{ folders: ScanFolderInfo[] }>(response);
+  return data.folders;
+}
+
+export async function addScanFolder(path: string): Promise<ScanFolderInfo> {
+  const response = await authFetch("/api/models/scan-folders", {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({ path }),
+  });
+  return parseJsonOrThrow<ScanFolderInfo>(response);
+}
+
+export async function removeScanFolder(id: number): Promise<void> {
+  const response = await authFetch(`/api/models/scan-folders/${id}`, {
+    method: "DELETE",
+  });
+  await parseJsonOrThrow<unknown>(response);
+}
+
 export async function listGgufVariants(
   repoId: string,
   hfToken?: string,

studio/frontend/src/features/chat/chat-page.tsx

@@ -8,7 +8,6 @@ import {
 } from "@/components/assistant-ui/model-selector";
 import { Thread } from "@/components/assistant-ui/thread";
 import { Button } from "@/components/ui/button";
-import { SidebarProvider, SidebarTrigger, useSidebar } from "@/components/ui/sidebar";
 import {
   Sheet,
   SheetContent,
@@ -16,7 +15,17 @@ import {
   SheetHeader,
   SheetTitle,
 } from "@/components/ui/sheet";
-import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
+import {
+  SidebarProvider,
+  SidebarTrigger,
+  useSidebar,
+} from "@/components/ui/sidebar";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import { GuidedTour, useGuidedTourController } from "@/features/tour";
 import { cn } from "@/lib/utils";
 import {
   ColumnInsertIcon,
@@ -36,7 +45,6 @@ import {
   useState,
 } from "react";
 import { toast } from "sonner";
-import { GuidedTour, useGuidedTourController } from "@/features/tour";
 import { listLocalModels } from "./api/chat-api";
 import { ChatSettingsPanel } from "./chat-settings-sheet";
 import { ContextUsageBar } from "./components/context-usage-bar";
@@ -48,16 +56,16 @@ import {
   getTrainingCompareHandoff,
 } from "./lib/training-compare-handoff";
 import { ChatRuntimeProvider } from "./runtime-provider";
-import { useChatRuntimeStore } from "./stores/chat-runtime-store";
 import {
   type CompareHandle,
   CompareHandlesProvider,
   RegisterCompareHandle,
   SharedComposer,
 } from "./shared-composer";
+import { useChatRuntimeStore } from "./stores/chat-runtime-store";
 import { ThreadSidebar } from "./thread-sidebar";
-import type { ChatView, MessageRecord } from "./types";
 import { buildChatTourSteps } from "./tour";
+import type { ChatView, MessageRecord } from "./types";
 
 type LoraCandidate = {
   id: string;
@@ -101,7 +109,9 @@ function messageHasImage(message: MessageRecord): boolean {
   if (contentParts.some((part) => part.type === "image")) {
     return true;
   }
-  const attachments = Array.isArray(message.attachments) ? message.attachments : [];
+  const attachments = Array.isArray(message.attachments)
+    ? message.attachments
+    : [];
   for (const attachment of attachments) {
     const parts = Array.isArray(attachment.content) ? attachment.content : [];
     for (const part of parts as Array<{ type?: string }>) {
@@ -152,12 +162,22 @@ const CompareContent = memo(function CompareContent({
   pairId,
   models,
   loraModels,
-}: { pairId: string; models: ModelOption[]; loraModels: LoraModelOption[] }): ReactElement {
+}: {
+  pairId: string;
+  models: ModelOption[];
+  loraModels: LoraModelOption[];
+}): ReactElement {
   const isLoraCompare = useIsLoraCompare();
 
-  return isLoraCompare
-    ? <LoraCompareContent pairId={pairId} />
-    : <GeneralCompareContent pairId={pairId} models={models} loraModels={loraModels} />;
+  return isLoraCompare ? (
+    <LoraCompareContent pairId={pairId} />
+  ) : (
+    <GeneralCompareContent
+      pairId={pairId}
+      models={models}
+      loraModels={loraModels}
+    />
+  );
 });
 
 /** Fast path: same model, adapter on/off, simultaneous generation. */
@@ -179,7 +199,9 @@ const LoraCompareContent = memo(function LoraCompareContent({
         setBaseThreadId(threads.find((t) => t.modelType === "base")?.id);
         setLoraThreadId(threads.find((t) => t.modelType === "lora")?.id);
       });
-    return () => { isActive = false; };
+    return () => {
+      isActive = false;
+    };
   }, [pairId]);
 
   return (
@@ -196,7 +218,11 @@ const LoraCompareContent = memo(function LoraCompareContent({
               </span>
             </div>
             <div className="min-h-0 flex-1">
-              <ChatRuntimeProvider modelType="base" pairId={pairId} initialThreadId={baseThreadId}>
+              <ChatRuntimeProvider
+                modelType="base"
+                pairId={pairId}
+                initialThreadId={baseThreadId}
+              >
                 <RegisterCompareHandle name="base" />
                 <Thread hideComposer={true} hideWelcome={true} />
               </ChatRuntimeProvider>
@@ -209,7 +235,11 @@ const LoraCompareContent = memo(function LoraCompareContent({
               </span>
             </div>
             <div className="min-h-0 flex-1">
-              <ChatRuntimeProvider modelType="lora" pairId={pairId} initialThreadId={loraThreadId}>
+              <ChatRuntimeProvider
+                modelType="lora"
+                pairId={pairId}
+                initialThreadId={loraThreadId}
+              >
                 <RegisterCompareHandle name="lora" />
                 <Thread hideComposer={true} hideWelcome={true} />
               </ChatRuntimeProvider>
@@ -229,7 +259,11 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
   pairId,
   models,
   loraModels,
-}: { pairId: string; models: ModelOption[]; loraModels: LoraModelOption[] }): ReactElement {
+}: {
+  pairId: string;
+  models: ModelOption[];
+  loraModels: LoraModelOption[];
+}): ReactElement {
   const handlesRef = useRef<Record<string, CompareHandle>>({});
   const [model1ThreadId, setModel1ThreadId] = useState<string>();
   const [model2ThreadId, setModel2ThreadId] = useState<string>();
@@ -241,7 +275,10 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
     isLora: false,
     ggufVariant: globalGgufVariant ?? undefined,
   });
-  const [model2, setModel2] = useState<CompareModelSelection>({ id: "", isLora: false });
+  const [model2, setModel2] = useState<CompareModelSelection>({
+    id: "",
+    isLora: false,
+  });
 
   useEffect(() => {
     let isActive = true;
@@ -252,13 +289,19 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
       .then((threads) => {
         if (!isActive) return;
         setModel1ThreadId(
-          threads.find((t) => t.modelType === "model1" || t.modelType === "base")?.id,
+          threads.find(
+            (t) => t.modelType === "model1" || t.modelType === "base",
+          )?.id,
         );
         setModel2ThreadId(
-          threads.find((t) => t.modelType === "model2" || t.modelType === "lora")?.id,
+          threads.find(
+            (t) => t.modelType === "model2" || t.modelType === "lora",
+          )?.id,
         );
       });
-    return () => { isActive = false; };
+    return () => {
+      isActive = false;
+    };
   }, [pairId]);
 
   return (
@@ -277,7 +320,13 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
                 models={models}
                 loraModels={loraModels}
                 value={model1.id}
-                onValueChange={(id, meta) => setModel1({ id, isLora: meta.isLora, ggufVariant: meta.ggufVariant })}
+                onValueChange={(id, meta) =>
+                  setModel1({
+                    id,
+                    isLora: meta.isLora,
+                    ggufVariant: meta.ggufVariant,
+                  })
+                }
                 variant="ghost"
                 size="sm"
                 className="max-w-[50%]"
@@ -303,7 +352,13 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
                 models={models}
                 loraModels={loraModels}
                 value={model2.id}
-                onValueChange={(id, meta) => setModel2({ id, isLora: meta.isLora, ggufVariant: meta.ggufVariant })}
+                onValueChange={(id, meta) =>
+                  setModel2({
+                    id,
+                    isLora: meta.isLora,
+                    ggufVariant: meta.ggufVariant,
+                  })
+                }
                 variant="ghost"
                 size="sm"
                 className="max-w-[50%]"
@@ -322,7 +377,11 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
           </div>
         </div>
         <div className="mx-auto w-full max-w-4xl px-4 py-4">
-          <SharedComposer handlesRef={handlesRef} model1={model1} model2={model2} />
+          <SharedComposer
+            handlesRef={handlesRef}
+            model1={model1}
+            model2={model2}
+          />
         </div>
       </div>
     </CompareHandlesProvider>
@@ -364,8 +423,7 @@ function InlineSidebar({
         data-sidebar="sidebar"
         className={cn(
           "bg-muted/70 text-sidebar-foreground h-full overflow-hidden rounded-2xl corner-squircle transition-[width] duration-200 ease-linear",
-          !collapsed &&
-          side === "right" && "border-l border-sidebar-border/70",
+          !collapsed && side === "right" && "border-l border-sidebar-border/70",
           collapsed ? "w-0" : "w-(--sidebar-width)",
         )}
       >
@@ -381,7 +439,11 @@ function TopBarActions({
   onNewThread,
   onNewCompare,
   showCompare,
-}: { onNewThread: () => void; onNewCompare: () => void; showCompare: boolean }) {
+}: {
+  onNewThread: () => void;
+  onNewCompare: () => void;
+  showCompare: boolean;
+}) {
   const { state } = useSidebar();
   if (state !== "collapsed") {
     return null;
@@ -424,8 +486,12 @@ export function ChatPage(): ReactElement {
   );
   const inferenceParams = useChatRuntimeStore((state) => state.params);
   const setInferenceParams = useChatRuntimeStore((state) => state.setParams);
-  const activeGgufVariant = useChatRuntimeStore((state) => state.activeGgufVariant);
-  const ggufContextLength = useChatRuntimeStore((state) => state.ggufContextLength);
+  const activeGgufVariant = useChatRuntimeStore(
+    (state) => state.activeGgufVariant,
+  );
+  const ggufContextLength = useChatRuntimeStore(
+    (state) => state.ggufContextLength,
+  );
   const contextUsage = useChatRuntimeStore((state) => state.contextUsage);
   const autoTitle = useChatRuntimeStore((state) => state.autoTitle);
   const setAutoTitle = useChatRuntimeStore((state) => state.setAutoTitle);
@@ -441,8 +507,7 @@ export function ChatPage(): ReactElement {
     loadingModel,
     loadProgress,
     loadToastDismissed,
-  } =
-    useChatModelRuntime();
+  } = useChatModelRuntime();
   const refreshRef = useRef(refresh);
   const selectModelRef = useRef(selectModel);
 
@@ -455,11 +520,24 @@ export function ChatPage(): ReactElement {
   }, [inferenceParams.checkpoint]);
 
   const handleCheckpointChange = useCallback(
-    (value: string, meta?: { isLora: boolean; ggufVariant?: string; isDownloaded?: boolean; expectedBytes?: number }) => {
+    (
+      value: string,
+      meta?: {
+        isLora: boolean;
+        ggufVariant?: string;
+        isDownloaded?: boolean;
+        expectedBytes?: number;
+      },
+    ) => {
       const store = useChatRuntimeStore.getState();
       const currentCheckpoint = store.params.checkpoint;
       const currentVariant = store.activeGgufVariant;
-      if (!value || (value === currentCheckpoint && (meta?.ggufVariant ?? null) === (currentVariant ?? null))) return;
+      if (
+        !value ||
+        (value === currentCheckpoint &&
+          (meta?.ggufVariant ?? null) === (currentVariant ?? null))
+      )
+        return;
       void (async () => {
         let showImageCompatibilityWarning = false;
         if (view.mode === "single" && activeThreadId) {
@@ -471,7 +549,9 @@ export function ChatPage(): ReactElement {
               .toArray();
             if (messages.length > 0) {
               const hasImage = messages.some(messageHasImage);
-              const targetModel = modelsFromStore.find((model) => model.id === value);
+              const targetModel = modelsFromStore.find(
+                (model) => model.id === value,
+              );
               showImageCompatibilityWarning =
                 hasImage && targetModel?.isVision === false;
             }
@@ -499,20 +579,14 @@ export function ChatPage(): ReactElement {
   const handleEject = useCallback(() => {
     void ejectModel();
   }, [ejectModel]);
-  const handleNewThread = useCallback(
-    () => {
-      useChatRuntimeStore.getState().setActiveThreadId(null);
-      setView({ mode: "single", newThreadNonce: crypto.randomUUID() });
-    },
-    [],
-  );
-  const handleNewCompare = useCallback(
-    () => {
-      setView({ mode: "compare", pairId: crypto.randomUUID() });
-      useChatRuntimeStore.getState().setContextUsage(null);
-    },
-    [],
-  );
+  const handleNewThread = useCallback(() => {
+    useChatRuntimeStore.getState().setActiveThreadId(null);
+    setView({ mode: "single", newThreadNonce: crypto.randomUUID() });
+  }, []);
+  const handleNewCompare = useCallback(() => {
+    setView({ mode: "compare", pairId: crypto.randomUUID() });
+    useChatRuntimeStore.getState().setContextUsage(null);
+  }, []);
 
   const openModelSelector = useCallback(() => {
     setModelSelectorLocked(true);
@@ -556,18 +630,17 @@ export function ChatPage(): ReactElement {
         .first()
         .then((msg) => {
           const saved = msg?.metadata as Record<string, unknown> | undefined;
-          const usage = saved?.contextUsage as typeof store.contextUsage | undefined;
+          const usage = saved?.contextUsage as
+            | typeof store.contextUsage
+            | undefined;
           if (usage) store.setContextUsage(usage);
         });
     }
   }, [viewBeforeCompare]);
 
-  const handleThreadSelect = useCallback(
-    (nextView: ChatView) => {
-      setView(nextView);
-    },
-    [],
-  );
+  const handleThreadSelect = useCallback((nextView: ChatView) => {
+    setView(nextView);
+  }, []);
 
   const models = useMemo<ModelOption[]>(
     () =>
@@ -581,6 +654,37 @@ export function ChatPage(): ReactElement {
 
   const [localModels, setLocalModels] = useState<LoraModelOption[]>([]);
 
+  const refreshLocalModels = useCallback(() => {
+    void listLocalModels()
+      .then((res) => {
+        setLocalModels(
+          res.models
+            .filter(
+              (m) =>
+                m.source === "lmstudio" ||
+                m.source === "models_dir" ||
+                m.source === "custom",
+            )
+            .map((m) => ({
+              id: m.id,
+              name:
+                m.source === "lmstudio" && m.model_id
+                  ? m.model_id
+                  : m.display_name,
+              baseModel:
+                m.source === "lmstudio"
+                  ? "LM Studio"
+                  : m.source === "custom"
+                    ? "Custom Folders"
+                    : "Local models",
+              updatedAt: m.updated_at ?? undefined,
+              source: "local" as const,
+            })),
+        );
+      })
+      .catch(() => {});
+  }, []);
+
   const loraModels = useMemo<LoraModelOption[]>(() => {
     const fromLoras = lorasFromStore.map((lora) => ({
       id: lora.id,
@@ -596,20 +700,8 @@ export function ChatPage(): ReactElement {
   useEffect(() => {
     if (getTrainingCompareHandoff()) return;
     void refresh();
-    void listLocalModels().then((res) => {
-      setLocalModels(
-        res.models
-          .filter((m) => m.source === "lmstudio" || m.source === "models_dir")
-          .map((m) => ({
-            id: m.id,
-            name: m.source === "lmstudio" && m.model_id ? m.model_id : m.display_name,
-            baseModel: m.source === "lmstudio" ? "LM Studio" : "Local models",
-            updatedAt: m.updated_at ?? undefined,
-            source: "local" as const,
-          })),
-      );
-    }).catch(() => {});
-  }, [refresh]);
+    refreshLocalModels();
+  }, [refresh, refreshLocalModels]);
 
   useEffect(() => {
     const handoff = getTrainingCompareHandoff();
@@ -649,7 +741,10 @@ export function ChatPage(): ReactElement {
           console.info("[chat-handoff] no lora match, loading base", {
             id: handoff.baseModel,
           });
-          await selectModelRef.current({ id: handoff.baseModel, isLora: false });
+          await selectModelRef.current({
+            id: handoff.baseModel,
+            isLora: false,
+          });
           if (canceled) return;
         } else {
           console.warn("[chat-handoff] no lora/base match found", {
@@ -767,9 +862,11 @@ export function ChatPage(): ReactElement {
                         ? "Loading model…"
                         : "Downloading model…"
                   }
-                  title={loadingModel.isDownloaded
-                    ? `Loading ${loadingModel.displayName} from cache.`
-                    : `Loading ${loadingModel.displayName}. This may include downloading.`}
+                  title={
+                    loadingModel.isDownloaded
+                      ? `Loading ${loadingModel.displayName} from cache.`
+                      : `Loading ${loadingModel.displayName}. This may include downloading.`
+                  }
                   progressPercent={loadProgress?.percent}
                   progressLabel={loadProgress?.label}
                   onStop={cancelLoading}
@@ -809,7 +906,12 @@ export function ChatPage(): ReactElement {
               newThreadNonce={view.newThreadNonce}
             />
           ) : (
-            <CompareContent key={view.pairId} pairId={view.pairId} models={models} loraModels={loraModels} />
+            <CompareContent
+              key={view.pairId}
+              pairId={view.pairId}
+              models={models}
+              loraModels={loraModels}
+            />
           )}
         </div>
 
@@ -832,6 +934,7 @@ export function ChatPage(): ReactElement {
               });
             }
           }}
+          onFoldersChange={refreshLocalModels}
         />
       </SidebarProvider>
     </div>

studio/frontend/src/features/chat/chat-settings-sheet.tsx

@@ -1,16 +1,6 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-import {
-  Select,
-  SelectContent,
-  SelectItem,
-  SelectTrigger,
-  SelectValue,
-} from "@/components/ui/select";
-import { Slider } from "@/components/ui/slider";
-import { Textarea } from "@/components/ui/textarea";
-import { Input } from "@/components/ui/input";
 import { Button } from "@/components/ui/button";
 import {
   Dialog,
@@ -20,11 +10,31 @@ import {
   DialogHeader,
   DialogTitle,
 } from "@/components/ui/dialog";
+import { Input } from "@/components/ui/input";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
+import {
+  Sheet,
+  SheetContent,
+  SheetDescription,
+  SheetHeader,
+  SheetTitle,
+} from "@/components/ui/sheet";
+import { Slider } from "@/components/ui/slider";
+import { Switch } from "@/components/ui/switch";
+import { Textarea } from "@/components/ui/textarea";
+import { useIsMobile } from "@/hooks/use-mobile";
 import {
   ArrowDown01Icon,
   CodeIcon,
   Delete02Icon,
   FloppyDiskIcon,
+  FolderSearchIcon,
   PencilEdit01Icon,
   Settings02Icon,
   SlidersHorizontalIcon,
@@ -33,22 +43,19 @@ import {
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { AnimatePresence, motion } from "motion/react";
-import {
-  Sheet,
-  SheetContent,
-  SheetDescription,
-  SheetHeader,
-  SheetTitle,
-} from "@/components/ui/sheet";
-import { useIsMobile } from "@/hooks/use-mobile";
 import type { ReactNode } from "react";
-import { useEffect, useMemo, useState } from "react";
+import { useCallback, useEffect, useMemo, useState } from "react";
+import {
+  type ScanFolderInfo,
+  addScanFolder,
+  listScanFolders,
+  removeScanFolder,
+} from "./api/chat-api";
+import { useChatRuntimeStore } from "./stores/chat-runtime-store";
 import {
   DEFAULT_INFERENCE_PARAMS,
   type InferenceParams,
 } from "./types/runtime";
-import { useChatRuntimeStore } from "./stores/chat-runtime-store";
-import { Switch } from "@/components/ui/switch";
 
 export const defaultInferenceParams = DEFAULT_INFERENCE_PARAMS;
 export type { InferenceParams } from "./types/runtime";
@@ -174,7 +181,11 @@ function loadCollapsibleState(): Record<string, boolean> {
     const raw = localStorage.getItem(COLLAPSIBLE_STATE_KEY);
     if (!raw) return {};
     const parsed = JSON.parse(raw);
-    if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) {
+    if (
+      typeof parsed !== "object" ||
+      parsed === null ||
+      Array.isArray(parsed)
+    ) {
       return {};
     }
     return Object.fromEntries(
@@ -255,6 +266,108 @@ function CollapsibleSection({
   );
 }
 
+function ModelFoldersSection({
+  onFoldersChange,
+}: { onFoldersChange?: () => void }) {
+  const [folders, setFolders] = useState<ScanFolderInfo[]>([]);
+  const [input, setInput] = useState("");
+  const [error, setError] = useState<string | null>(null);
+  const [loading, setLoading] = useState(false);
+
+  const refresh = useCallback(() => {
+    listScanFolders()
+      .then(setFolders)
+      .catch(() => {});
+  }, []);
+
+  useEffect(() => {
+    refresh();
+  }, [refresh]);
+
+  const handleAdd = async () => {
+    const trimmed = input.trim();
+    if (!trimmed) return;
+    setError(null);
+    setLoading(true);
+    try {
+      await addScanFolder(trimmed);
+      setInput("");
+      refresh();
+      onFoldersChange?.();
+    } catch (e) {
+      setError(e instanceof Error ? e.message : "Failed to add folder");
+    } finally {
+      setLoading(false);
+    }
+  };
+
+  const handleRemove = async (id: number) => {
+    try {
+      await removeScanFolder(id);
+      onFoldersChange?.();
+    } catch (e) {
+      setError(e instanceof Error ? e.message : "Failed to remove folder");
+    } finally {
+      refresh();
+    }
+  };
+
+  return (
+    <CollapsibleSection icon={FolderSearchIcon} label="Model Folders">
+      <div className="flex flex-col gap-2 py-1">
+        {folders.length > 0 && (
+          <div className="flex flex-col gap-1">
+            {folders.map((f) => (
+              <div
+                key={f.id}
+                className="group flex items-center gap-1.5 rounded-md px-1.5 py-1 text-xs transition-colors hover:bg-accent"
+              >
+                <span
+                  className="min-w-0 flex-1 truncate text-muted-foreground"
+                  title={f.path}
+                >
+                  {f.path}
+                </span>
+                <button
+                  type="button"
+                  onClick={() => handleRemove(f.id)}
+                  className="shrink-0 rounded p-0.5 text-muted-foreground/50 opacity-0 transition-opacity group-hover:opacity-100 hover:text-destructive"
+                >
+                  <HugeiconsIcon icon={Delete02Icon} className="size-3" />
+                </button>
+              </div>
+            ))}
+          </div>
+        )}
+        <div className="flex gap-1.5">
+          <Input
+            value={input}
+            onChange={(e) => {
+              setInput(e.target.value);
+              setError(null);
+            }}
+            onKeyDown={(e) => {
+              if (e.key === "Enter") handleAdd();
+            }}
+            placeholder="/path/to/models"
+            className="h-7 flex-1 text-xs font-mono"
+            disabled={loading}
+          />
+          <button
+            type="button"
+            onClick={handleAdd}
+            disabled={loading || !input.trim()}
+            className="h-7 rounded-md border px-2 text-[11px] font-medium text-muted-foreground transition-colors hover:bg-accent disabled:cursor-not-allowed disabled:opacity-50"
+          >
+            Add
+          </button>
+        </div>
+        {error && <p className="text-[11px] text-destructive">{error}</p>}
+      </div>
+    </CollapsibleSection>
+  );
+}
+
 interface ChatSettingsPanelProps {
   open: boolean;
   onOpenChange?: (open: boolean) => void;
@@ -263,6 +376,7 @@ interface ChatSettingsPanelProps {
   autoTitle: boolean;
   onAutoTitleChange: (enabled: boolean) => void;
   onReloadModel?: () => void;
+  onFoldersChange?: () => void;
 }
 
 export function ChatSettingsPanel({
@@ -273,24 +387,33 @@ export function ChatSettingsPanel({
   autoTitle,
   onAutoTitleChange,
   onReloadModel,
+  onFoldersChange,
 }: ChatSettingsPanelProps) {
   const isMobile = useIsMobile();
   const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null;
   const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
+  const ggufMaxContextLength = useChatRuntimeStore(
+    (s) => s.ggufMaxContextLength,
+  );
   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 setCustomContextLength = useChatRuntimeStore(
+    (s) => s.setCustomContextLength,
+  );
 
   const ctxDisplayValue = customContextLength ?? ggufContextLength ?? "";
+  const ctxMaxValue = ggufMaxContextLength ?? ggufContextLength ?? null;
   const kvDirty = kvCacheDtype !== loadedKvCacheDtype;
   const ctxDirty = customContextLength !== null;
   const modelSettingsDirty = kvDirty || ctxDirty;
   const [customPresets, setCustomPresets] = useState<Preset[]>(() =>
     loadSavedCustomPresets(),
   );
-  const [activePreset, setActivePreset] = useState(() => loadSavedActivePreset());
+  const [activePreset, setActivePreset] = useState(() =>
+    loadSavedActivePreset(),
+  );
   const [savePresetOpen, setSavePresetOpen] = useState(false);
   const [presetNameDraft, setPresetNameDraft] = useState("");
   const presets = useMemo(
@@ -415,325 +538,356 @@ export function ChatSettingsPanel({
       <div className="flex-1 overflow-y-auto px-1.5">
         {/* mt-4 matches the Playground sidebar gap (SidebarHeader py-3 + SidebarGroup pt-1) */}
         <div className="mt-4 px-2 pb-3">
-            <div className="flex items-center gap-2">
-              <Select value={activePreset} onValueChange={applyPreset}>
-                <SelectTrigger className="h-8 flex-1 corner-squircle text-xs">
-                  <SelectValue />
-                </SelectTrigger>
-                <SelectContent>
-                  {presets.map((p) => (
-                    <SelectItem key={p.name} value={p.name}>
-                      {p.name}
-                    </SelectItem>
-                  ))}
-                </SelectContent>
-              </Select>
-              <button
-                type="button"
-                onClick={openSavePresetDialog}
-                className="flex h-8 items-center gap-1.5 rounded-md border px-2.5 text-xs text-muted-foreground transition-colors hover:bg-accent"
-                title="Save preset"
-              >
-                <HugeiconsIcon icon={FloppyDiskIcon} className="size-3.5" />
-                Save
-              </button>
-              <button
-                type="button"
-                onClick={() => deletePreset(activePreset)}
-                disabled={isBuiltinPreset}
-                className="flex h-8 items-center gap-1.5 rounded-md border px-2.5 text-xs text-muted-foreground transition-colors hover:bg-accent disabled:cursor-not-allowed disabled:opacity-50"
-                title={
-                  isBuiltinPreset
-                    ? "Built-in presets cannot be deleted"
-                    : "Delete selected preset"
-                }
-              >
-                <HugeiconsIcon icon={Delete02Icon} className="size-3.5" />
-                Delete
-              </button>
-            </div>
-          </div>
-
-          <div className="px-2 pb-4">
-            <label
-              htmlFor="system-prompt"
-              className="mb-1.5 block text-xs font-medium"
+          <div className="flex items-center gap-2">
+            <Select value={activePreset} onValueChange={applyPreset}>
+              <SelectTrigger className="h-8 flex-1 corner-squircle text-xs">
+                <SelectValue />
+              </SelectTrigger>
+              <SelectContent>
+                {presets.map((p) => (
+                  <SelectItem key={p.name} value={p.name}>
+                    {p.name}
+                  </SelectItem>
+                ))}
+              </SelectContent>
+            </Select>
+            <button
+              type="button"
+              onClick={openSavePresetDialog}
+              className="flex h-8 items-center gap-1.5 rounded-md border px-2.5 text-xs text-muted-foreground transition-colors hover:bg-accent"
+              title="Save preset"
             >
-              System Prompt
-            </label>
-            <Textarea
-              id="system-prompt"
-              value={params.systemPrompt}
-              onChange={(e) => set("systemPrompt")(e.target.value)}
-              placeholder="You are a helpful assistant..."
-              className="min-h-20 text-xs corner-squircle"
-              rows={3}
-            />
+              <HugeiconsIcon icon={FloppyDiskIcon} className="size-3.5" />
+              Save
+            </button>
+            <button
+              type="button"
+              onClick={() => deletePreset(activePreset)}
+              disabled={isBuiltinPreset}
+              className="flex h-8 items-center gap-1.5 rounded-md border px-2.5 text-xs text-muted-foreground transition-colors hover:bg-accent disabled:cursor-not-allowed disabled:opacity-50"
+              title={
+                isBuiltinPreset
+                  ? "Built-in presets cannot be deleted"
+                  : "Delete selected preset"
+              }
+            >
+              <HugeiconsIcon icon={Delete02Icon} className="size-3.5" />
+              Delete
+            </button>
           </div>
+        </div>
 
-          <CollapsibleSection icon={Settings02Icon} label="Model" defaultOpen={true}>
-            <div className="flex flex-col gap-3 py-1">
-              {isGguf && (
-                <>
-                  <div className="space-y-2">
-                    <div className="flex items-center justify-between">
-                      <span className="text-xs font-medium">Context Length</span>
-                      <Input
-                        type="number"
-                        value={typeof ctxDisplayValue === "number" ? ctxDisplayValue : (ggufContextLength ?? "")}
-                        placeholder="..."
-                        min={128}
-                        max={ggufContextLength ?? undefined}
-                        step={1024}
-                        className="h-6 w-[100px] text-right text-xs tabular-nums"
-                        onChange={(e) => {
-                          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);
-                          }
-                        }}
-                      />
-                    </div>
-                    <Slider
-                      min={1024}
-                      max={ggufContextLength ?? 4096}
+        <div className="px-2 pb-4">
+          <label
+            htmlFor="system-prompt"
+            className="mb-1.5 block text-xs font-medium"
+          >
+            System Prompt
+          </label>
+          <Textarea
+            id="system-prompt"
+            value={params.systemPrompt}
+            onChange={(e) => set("systemPrompt")(e.target.value)}
+            placeholder="You are a helpful assistant..."
+            className="min-h-20 text-xs corner-squircle"
+            rows={3}
+          />
+        </div>
+
+        <CollapsibleSection
+          icon={Settings02Icon}
+          label="Model"
+          defaultOpen={true}
+        >
+          <div className="flex flex-col gap-3 py-1">
+            {isGguf && (
+              <>
+                <div className="space-y-2">
+                  <div className="flex items-center justify-between">
+                    <span className="text-xs font-medium">Context Length</span>
+                    <Input
+                      type="number"
+                      value={
+                        typeof ctxDisplayValue === "number"
+                          ? ctxDisplayValue
+                          : (ggufContextLength ?? "")
+                      }
+                      placeholder="..."
+                      min={128}
+                      max={ctxMaxValue ?? undefined}
                       step={1024}
-                      value={[Math.min(typeof ctxDisplayValue === "number" ? ctxDisplayValue : (ggufContextLength ?? 4096), ggufContextLength ?? 4096)]}
-                      onValueChange={([v]) => {
-                        setCustomContextLength(v === (ggufContextLength ?? 0) ? null : v);
+                      className="h-6 w-[100px] text-right text-xs tabular-nums"
+                      onChange={(e) => {
+                        const raw = e.target.value;
+                        if (raw === "") {
+                          setCustomContextLength(null);
+                          return;
+                        }
+                        const v = Number.parseInt(raw, 10);
+                        if (!Number.isNaN(v) && v >= 0) {
+                          const maxCtx =
+                            ctxMaxValue ?? Number.POSITIVE_INFINITY;
+                          const clamped = Math.min(v, maxCtx);
+                          setCustomContextLength(
+                            clamped === (ggufContextLength ?? 0)
+                              ? null
+                              : clamped,
+                          );
+                        }
                       }}
                     />
                   </div>
-                  <div className="flex items-center justify-between gap-3">
-                    <div className="min-w-0">
-                      <div className="text-xs font-medium">KV Cache Dtype</div>
-                      <div className="text-[11px] text-muted-foreground">
-                        Quantize KV cache to reduce VRAM.
-                      </div>
-                    </div>
-                    <Select
-                      value={kvCacheDtype ?? "f16"}
-                      onValueChange={(v) => {
-                        setKvCacheDtype(v === "f16" ? null : v);
-                      }}
-                    >
-                      <SelectTrigger className="h-7 w-[90px] text-xs">
-                        <SelectValue />
-                      </SelectTrigger>
-                      <SelectContent>
-                        <SelectItem value="f16">f16</SelectItem>
-                        <SelectItem value="bf16">bf16</SelectItem>
-                        <SelectItem value="q8_0">q8_0</SelectItem>
-                        <SelectItem value="q5_1">q5_1</SelectItem>
-                        <SelectItem value="q4_1">q4_1</SelectItem>
-                      </SelectContent>
-                    </Select>
-                  </div>
-                  {modelSettingsDirty && (
-                    <div className="flex flex-wrap gap-1.5 pt-1">
-                      <button
-                        type="button"
-                        onClick={() => onReloadModel?.()}
-                        className="rounded-md bg-primary px-2.5 py-1 text-[11px] font-medium text-primary-foreground transition-colors hover:bg-primary/90"
-                      >
-                        Apply
-                      </button>
-                      <button
-                        type="button"
-                        onClick={() => {
-                          setCustomContextLength(null);
-                          setKvCacheDtype(loadedKvCacheDtype);
-                        }}
-                        className="rounded-md border px-2.5 py-1 text-[11px] font-medium text-muted-foreground transition-colors hover:bg-accent"
-                      >
-                        Reset
-                      </button>
-                    </div>
-                  )}
-                </>
-              )}
-              {!isGguf && params.checkpoint && (
-                <div className="flex items-center justify-between gap-3">
-                  <div className="min-w-0">
-                    <div className="text-xs font-medium">Enable custom code</div>
-                    <div className="text-[11px] text-muted-foreground">
-                      Allow models with custom code (e.g. Nemotron). Only enable if sure.
-                    </div>
-                  </div>
-                  <Switch
-                    checked={params.trustRemoteCode ?? false}
-                    onCheckedChange={set("trustRemoteCode")}
+                  <Slider
+                    min={1024}
+                    max={ctxMaxValue ?? 4096}
+                    step={1024}
+                    value={[
+                      Math.min(
+                        typeof ctxDisplayValue === "number"
+                          ? ctxDisplayValue
+                          : (ggufContextLength ?? 4096),
+                        ctxMaxValue ?? 4096,
+                      ),
+                    ]}
+                    onValueChange={([v]) => {
+                      setCustomContextLength(
+                        v === (ggufContextLength ?? 0) ? null : v,
+                      );
+                    }}
                   />
                 </div>
-              )}
-            </div>
-          </CollapsibleSection>
-
-          <CollapsibleSection
-            icon={SlidersHorizontalIcon}
-            label="Sampling"
-            defaultOpen={true}
-          >
-            <div className="flex flex-col gap-5">
-              <ParamSlider
-                label="Temperature"
-                value={params.temperature}
-                min={0}
-                max={2}
-                step={0.1}
-                onChange={set("temperature")}
-              />
-              <ParamSlider
-                label="Top P"
-                value={params.topP}
-                min={0}
-                max={1}
-                step={0.05}
-                onChange={set("topP")}
-                displayValue={params.topP === 1 ? "Off" : undefined}
-              />
-              <ParamSlider
-                label="Top K"
-                value={params.topK}
-                min={0}
-                max={100}
-                step={1}
-                onChange={set("topK")}
-                displayValue={params.topK === 0 ? "Off" : undefined}
-              />
-              <ParamSlider
-                label="Min P"
-                value={params.minP}
-                min={0}
-                max={1}
-                step={0.01}
-                onChange={set("minP")}
-              />
-              <ParamSlider
-                label="Repetition Penalty"
-                value={params.repetitionPenalty}
-                min={1}
-                max={2}
-                step={0.05}
-                onChange={set("repetitionPenalty")}
-                displayValue={params.repetitionPenalty === 1 ? "Off" : undefined}
-              />
-              <ParamSlider
-                label="Presence Penalty"
-                value={params.presencePenalty}
-                min={0}
-                max={2}
-                step={0.1}
-                onChange={set("presencePenalty")}
-                displayValue={params.presencePenalty === 0 ? "Off" : undefined}
-              />
-              {!isGguf && (
-                <ParamSlider
-                  label="Max Seq Length"
-                  value={params.maxSeqLength}
-                  min={128}
-                  max={32768}
-                  step={128}
-                  onChange={set("maxSeqLength")}
-                />
-              )}
-              <ParamSlider
-                label="Max Tokens"
-                value={params.maxTokens}
-                min={64}
-                max={isGguf && ggufContextLength ? ggufContextLength : 32768}
-                step={64}
-                onChange={set("maxTokens")}
-                displayValue={
-                  isGguf && ggufContextLength && params.maxTokens >= ggufContextLength
-                    ? "Max"
-                    : undefined
-                }
-              />
-            </div>
-          </CollapsibleSection>
-
-          <CollapsibleSection icon={Wrench01Icon} label="Tools">
-            <div className="flex flex-col gap-3 py-1">
-              <AutoHealToolCallsToggle />
-              <MaxToolCallsSlider />
-              <ToolCallTimeoutSlider />
-            </div>
-          </CollapsibleSection>
-
-          <CollapsibleSection icon={UserSettings01Icon} label="Preferences" defaultOpen={true}>
-            <div className="flex flex-col gap-3 py-1">
+                <div className="flex items-center justify-between gap-3">
+                  <div className="min-w-0">
+                    <div className="text-xs font-medium">KV Cache Dtype</div>
+                    <div className="text-[11px] text-muted-foreground">
+                      Quantize KV cache to reduce VRAM.
+                    </div>
+                  </div>
+                  <Select
+                    value={kvCacheDtype ?? "f16"}
+                    onValueChange={(v) => {
+                      setKvCacheDtype(v === "f16" ? null : v);
+                    }}
+                  >
+                    <SelectTrigger className="h-7 w-[90px] text-xs">
+                      <SelectValue />
+                    </SelectTrigger>
+                    <SelectContent>
+                      <SelectItem value="f16">f16</SelectItem>
+                      <SelectItem value="bf16">bf16</SelectItem>
+                      <SelectItem value="q8_0">q8_0</SelectItem>
+                      <SelectItem value="q5_1">q5_1</SelectItem>
+                      <SelectItem value="q4_1">q4_1</SelectItem>
+                    </SelectContent>
+                  </Select>
+                </div>
+                {modelSettingsDirty && (
+                  <div className="flex flex-wrap gap-1.5 pt-1">
+                    <button
+                      type="button"
+                      onClick={() => onReloadModel?.()}
+                      className="rounded-md bg-primary px-2.5 py-1 text-[11px] font-medium text-primary-foreground transition-colors hover:bg-primary/90"
+                    >
+                      Apply
+                    </button>
+                    <button
+                      type="button"
+                      onClick={() => {
+                        setCustomContextLength(null);
+                        setKvCacheDtype(loadedKvCacheDtype);
+                      }}
+                      className="rounded-md border px-2.5 py-1 text-[11px] font-medium text-muted-foreground transition-colors hover:bg-accent"
+                    >
+                      Reset
+                    </button>
+                  </div>
+                )}
+              </>
+            )}
+            {!isGguf && params.checkpoint && (
               <div className="flex items-center justify-between gap-3">
                 <div className="min-w-0">
-                  <div className="text-xs font-medium">Auto title</div>
+                  <div className="text-xs font-medium">Enable custom code</div>
                   <div className="text-[11px] text-muted-foreground">
-                    Generate short title after reply.
+                    Allow models with custom code (e.g. Nemotron). Only enable
+                    if sure.
                   </div>
                 </div>
                 <Switch
-                  checked={autoTitle}
-                  onCheckedChange={onAutoTitleChange}
+                  checked={params.trustRemoteCode ?? false}
+                  onCheckedChange={set("trustRemoteCode")}
                 />
               </div>
-              <HfTokenField />
-            </div>
-          </CollapsibleSection>
+            )}
+          </div>
+        </CollapsibleSection>
 
-          <ChatTemplateSection onReloadModel={onReloadModel} />
-        </div>
-        <Dialog
-          open={savePresetOpen}
-          onOpenChange={(nextOpen) => {
-            setSavePresetOpen(nextOpen);
-            if (!nextOpen) {
-              setPresetNameDraft("");
-            }
-          }}
+        <CollapsibleSection
+          icon={SlidersHorizontalIcon}
+          label="Sampling"
+          defaultOpen={true}
         >
-          <DialogContent className="corner-squircle sm:max-w-sm">
-            <DialogHeader>
-              <DialogTitle>Save Preset</DialogTitle>
-              <DialogDescription>
-                Enter a name for this inference preset.
-              </DialogDescription>
-            </DialogHeader>
-            <form
-              onSubmit={(event) => {
-                event.preventDefault();
-                savePresetWithName(presetNameDraft);
-              }}
-              className="space-y-4"
-            >
-              <Input
-                autoFocus={true}
-                value={presetNameDraft}
-                onChange={(event) => setPresetNameDraft(event.target.value)}
-                placeholder="Preset name"
-                maxLength={80}
+          <div className="flex flex-col gap-5">
+            <ParamSlider
+              label="Temperature"
+              value={params.temperature}
+              min={0}
+              max={2}
+              step={0.1}
+              onChange={set("temperature")}
+            />
+            <ParamSlider
+              label="Top P"
+              value={params.topP}
+              min={0}
+              max={1}
+              step={0.05}
+              onChange={set("topP")}
+              displayValue={params.topP === 1 ? "Off" : undefined}
+            />
+            <ParamSlider
+              label="Top K"
+              value={params.topK}
+              min={0}
+              max={100}
+              step={1}
+              onChange={set("topK")}
+              displayValue={params.topK === 0 ? "Off" : undefined}
+            />
+            <ParamSlider
+              label="Min P"
+              value={params.minP}
+              min={0}
+              max={1}
+              step={0.01}
+              onChange={set("minP")}
+            />
+            <ParamSlider
+              label="Repetition Penalty"
+              value={params.repetitionPenalty}
+              min={1}
+              max={2}
+              step={0.05}
+              onChange={set("repetitionPenalty")}
+              displayValue={params.repetitionPenalty === 1 ? "Off" : undefined}
+            />
+            <ParamSlider
+              label="Presence Penalty"
+              value={params.presencePenalty}
+              min={0}
+              max={2}
+              step={0.1}
+              onChange={set("presencePenalty")}
+              displayValue={params.presencePenalty === 0 ? "Off" : undefined}
+            />
+            {!isGguf && (
+              <ParamSlider
+                label="Max Seq Length"
+                value={params.maxSeqLength}
+                min={128}
+                max={32768}
+                step={128}
+                onChange={set("maxSeqLength")}
               />
-              <DialogFooter>
-                <Button
-                  type="button"
-                  variant="outline"
-                  onClick={() => setSavePresetOpen(false)}
-                >
-                  Cancel
-                </Button>
-                <Button type="submit" disabled={presetNameDraft.trim().length === 0}>
-                  Save
-                </Button>
-              </DialogFooter>
-            </form>
-          </DialogContent>
-        </Dialog>
-      </>
+            )}
+            <ParamSlider
+              label="Max Tokens"
+              value={params.maxTokens}
+              min={64}
+              max={isGguf && ggufContextLength ? ggufContextLength : 32768}
+              step={64}
+              onChange={set("maxTokens")}
+              displayValue={
+                isGguf &&
+                ggufContextLength &&
+                params.maxTokens >= ggufContextLength
+                  ? "Max"
+                  : undefined
+              }
+            />
+          </div>
+        </CollapsibleSection>
+
+        <CollapsibleSection icon={Wrench01Icon} label="Tools">
+          <div className="flex flex-col gap-3 py-1">
+            <AutoHealToolCallsToggle />
+            <MaxToolCallsSlider />
+            <ToolCallTimeoutSlider />
+          </div>
+        </CollapsibleSection>
+
+        <CollapsibleSection
+          icon={UserSettings01Icon}
+          label="Preferences"
+          defaultOpen={true}
+        >
+          <div className="flex flex-col gap-3 py-1">
+            <div className="flex items-center justify-between gap-3">
+              <div className="min-w-0">
+                <div className="text-xs font-medium">Auto title</div>
+                <div className="text-[11px] text-muted-foreground">
+                  Generate short title after reply.
+                </div>
+              </div>
+              <Switch checked={autoTitle} onCheckedChange={onAutoTitleChange} />
+            </div>
+            <HfTokenField />
+          </div>
+        </CollapsibleSection>
+
+        <ModelFoldersSection onFoldersChange={onFoldersChange} />
+
+        <ChatTemplateSection onReloadModel={onReloadModel} />
+      </div>
+      <Dialog
+        open={savePresetOpen}
+        onOpenChange={(nextOpen) => {
+          setSavePresetOpen(nextOpen);
+          if (!nextOpen) {
+            setPresetNameDraft("");
+          }
+        }}
+      >
+        <DialogContent className="corner-squircle sm:max-w-sm">
+          <DialogHeader>
+            <DialogTitle>Save Preset</DialogTitle>
+            <DialogDescription>
+              Enter a name for this inference preset.
+            </DialogDescription>
+          </DialogHeader>
+          <form
+            onSubmit={(event) => {
+              event.preventDefault();
+              savePresetWithName(presetNameDraft);
+            }}
+            className="space-y-4"
+          >
+            <Input
+              autoFocus={true}
+              value={presetNameDraft}
+              onChange={(event) => setPresetNameDraft(event.target.value)}
+              placeholder="Preset name"
+              maxLength={80}
+            />
+            <DialogFooter>
+              <Button
+                type="button"
+                variant="outline"
+                onClick={() => setSavePresetOpen(false)}
+              >
+                Cancel
+              </Button>
+              <Button
+                type="submit"
+                disabled={presetNameDraft.trim().length === 0}
+              >
+                Save
+              </Button>
+            </DialogFooter>
+          </form>
+        </DialogContent>
+      </Dialog>
+    </>
   );
 
   if (isMobile) {
@@ -761,7 +915,9 @@ export function ChatSettingsPanel({
 
 function MaxToolCallsSlider() {
   const maxToolCalls = useChatRuntimeStore((s) => s.maxToolCallsPerMessage);
-  const setMaxToolCalls = useChatRuntimeStore((s) => s.setMaxToolCallsPerMessage);
+  const setMaxToolCalls = useChatRuntimeStore(
+    (s) => s.setMaxToolCallsPerMessage,
+  );
 
   // Slider range 0-41; 41 maps to 9999 ("Max")
   const sliderValue = maxToolCalls >= 9999 ? 41 : Math.min(maxToolCalls, 40);
@@ -774,7 +930,9 @@ function MaxToolCallsSlider() {
       max={41}
       step={1}
       onChange={(v) => setMaxToolCalls(v >= 41 ? 9999 : v)}
-      displayValue={sliderValue >= 41 ? "Max" : sliderValue === 0 ? "Off" : undefined}
+      displayValue={
+        sliderValue >= 41 ? "Max" : sliderValue === 0 ? "Off" : undefined
+      }
     />
   );
 }
@@ -808,7 +966,9 @@ function ToolCallTimeoutSlider() {
 
 function AutoHealToolCallsToggle() {
   const autoHealToolCalls = useChatRuntimeStore((s) => s.autoHealToolCalls);
-  const setAutoHealToolCalls = useChatRuntimeStore((s) => s.setAutoHealToolCalls);
+  const setAutoHealToolCalls = useChatRuntimeStore(
+    (s) => s.setAutoHealToolCalls,
+  );
 
   return (
     <div className="flex items-center justify-between gap-3">

studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts

@@ -240,11 +240,18 @@ export function useChatModelRuntime() {
         const supportsReasoning = statusRes.supports_reasoning ?? false;
         const reasoningAlwaysOn = statusRes.reasoning_always_on ?? false;
         const supportsTools = statusRes.supports_tools ?? false;
+        const currentGgufContextLength = statusRes.is_gguf
+          ? (statusRes.context_length ?? null)
+          : null;
+        const ggufMaxContextLength = statusRes.is_gguf
+          ? (statusRes.max_context_length ?? null)
+          : null;
         useChatRuntimeStore.setState({
           supportsReasoning,
           reasoningAlwaysOn,
           supportsTools,
-          ggufContextLength: statusRes.is_gguf ? (statusRes.context_length ?? null) : null,
+          ggufContextLength: currentGgufContextLength,
+          ggufMaxContextLength,
         });
 
         // Set reasoning default for Qwen3.5 small models
@@ -415,16 +422,17 @@ export function useChatModelRuntime() {
             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
+            const reportedMaxCtx = loadResponse.is_gguf
+              ? (loadResponse.max_context_length ?? null)
               : null;
+            // A successful reload has applied settings, so clear pending custom
+            // context state and display the backend-reported effective context.
+            const keepCustomCtx = null;
             const reasoningAlwaysOn = loadResponse.reasoning_always_on ?? false;
+            const ggufMaxContextLength = reportedMaxCtx;
             useChatRuntimeStore.setState({
               ggufContextLength: nativeCtx,
+              ggufMaxContextLength,
               supportsReasoning: loadResponse.supports_reasoning ?? false,
               reasoningAlwaysOn,
               reasoningEnabled: reasoningAlwaysOn ? true : reasoningDefault,

studio/frontend/src/features/chat/stores/chat-runtime-store.ts

@@ -150,6 +150,7 @@ type ChatRuntimeStore = {
   modelsError: string | null;
   activeGgufVariant: string | null;
   ggufContextLength: number | null;
+  ggufMaxContextLength: number | null;
   supportsReasoning: boolean;
   reasoningAlwaysOn: boolean;
   reasoningEnabled: boolean;
@@ -213,6 +214,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
   modelsError: null,
   activeGgufVariant: null,
   ggufContextLength: null,
+  ggufMaxContextLength: null,
   supportsReasoning: false,
   reasoningAlwaysOn: false,
   reasoningEnabled: true,
@@ -222,7 +224,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
   toolStatus: null,
   generatingStatus: null,
   autoHealToolCalls: loadBool(AUTO_HEAL_TOOL_CALLS_KEY, true),
-  maxToolCallsPerMessage: loadInt(MAX_TOOL_CALLS_KEY, 10),
+  maxToolCallsPerMessage: loadInt(MAX_TOOL_CALLS_KEY, 25),
   toolCallTimeout: loadInt(TOOL_CALL_TIMEOUT_KEY, 5),
   kvCacheDtype: null,
   loadedKvCacheDtype: null,
@@ -287,6 +289,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
       },
       activeGgufVariant: null,
       ggufContextLength: null,
+      ggufMaxContextLength: null,
       contextUsage: null,
       supportsReasoning: false,
       reasoningEnabled: true,

studio/frontend/src/features/chat/types/api.ts

@@ -86,6 +86,7 @@ export interface LoadModelResponse {
     trust_remote_code?: boolean;
   };
   context_length?: number | null;
+  max_context_length?: number | null;
   supports_reasoning?: boolean;
   reasoning_always_on?: boolean;
   supports_tools?: boolean;
@@ -119,6 +120,7 @@ export interface InferenceStatusResponse {
   reasoning_always_on?: boolean;
   supports_tools?: boolean;
   context_length?: number | null;
+  max_context_length?: number | null;
 }
 
 export interface AudioGenerationResponse {

studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx

@@ -92,6 +92,7 @@ export function ModelSelectionStep() {
   const [inputValue, setInputValue] = useState("");
   const selectingRef = useRef(false);
   const debouncedQuery = useDebouncedValue(inputValue);
+  const debouncedHfToken = useDebouncedValue(hfToken, 500);
   const task = modelType ? MODEL_TYPE_TO_HF_TASK[modelType] : undefined;
   const {
     results: hfResults,
@@ -101,7 +102,7 @@ export function ModelSelectionStep() {
     error: hfSearchError,
   } = useHfModelSearch(debouncedQuery, {
     task,
-    accessToken: hfToken || undefined,
+    accessToken: debouncedHfToken || undefined,
     excludeGguf: true,
     priorityIds: PRIORITY_TRAINING_MODELS,
   });

studio/frontend/src/features/studio/sections/model-section.tsx

@@ -28,7 +28,16 @@ import {
   TooltipContent,
   TooltipTrigger,
 } from "@/components/ui/tooltip";
-import { MODEL_TYPE_TO_HF_TASK, PRIORITY_TRAINING_MODELS, applyPriorityOrdering } from "@/config/training";
+import {
+  MODEL_TYPE_TO_HF_TASK,
+  PRIORITY_TRAINING_MODELS,
+  applyPriorityOrdering,
+} from "@/config/training";
+import {
+  type LocalModelInfo,
+  listLocalModels,
+  useTrainingConfigStore,
+} from "@/features/training";
 import {
   useDebouncedValue,
   useGpuInfo,
@@ -38,15 +47,10 @@ import {
 } from "@/hooks";
 import { formatCompact } from "@/lib/utils";
 import {
-  type TrainingMethod as VramTrainingMethod,
   type VramFitStatus,
+  type TrainingMethod as VramTrainingMethod,
   buildModelVramMap,
 } from "@/lib/vram";
-import {
-  listLocalModels,
-  type LocalModelInfo,
-  useTrainingConfigStore,
-} from "@/features/training";
 import type { TrainingMethod } from "@/types/training";
 import {
   ChipIcon,
@@ -119,6 +123,7 @@ export function ModelSection() {
   const [localModelsError, setLocalModelsError] = useState<string | null>(null);
   const selectingRef = useRef(false);
   const debouncedQuery = useDebouncedValue(inputValue);
+  const debouncedHfToken = useDebouncedValue(hfToken, 500);
 
   function handleModelSelect(id: string | null) {
     selectingRef.current = true;
@@ -149,7 +154,9 @@ export function ModelSection() {
       .catch((error) => {
         if (controller.signal.aborted) return;
         setLocalModelsError(
-          error instanceof Error ? error.message : "Failed to load local models",
+          error instanceof Error
+            ? error.message
+            : "Failed to load local models",
         );
       })
       .finally(() => {
@@ -167,7 +174,7 @@ export function ModelSection() {
     error: hfSearchError,
   } = useHfModelSearch(debouncedQuery, {
     task,
-    accessToken: hfToken || undefined,
+    accessToken: debouncedHfToken || undefined,
     excludeGguf: true,
     priorityIds: PRIORITY_TRAINING_MODELS,
   });
@@ -240,7 +247,9 @@ export function ModelSection() {
       { est: number; status: VramFitStatus | null; detail: string | null }
     >();
     for (const r of hfResults) {
-      const detail = r.totalParams ? formatCompact(r.totalParams) : extractParamLabel(r.id);
+      const detail = r.totalParams
+        ? formatCompact(r.totalParams)
+        : extractParamLabel(r.id);
       const fit = fitMap.get(r.id);
       map.set(r.id, {
         est: fit?.est ?? 0,
@@ -270,363 +279,383 @@ export function ModelSection() {
         className="shadow-border ring-border"
       >
         <div className="grid min-w-0 gap-4 md:grid-cols-2 xl:grid-cols-4">
-          <div data-tour="studio-local-model" className="flex min-w-0 flex-col gap-2">
+          <div
+            data-tour="studio-local-model"
+            className="flex min-w-0 flex-col gap-2"
+          >
             <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
               Local Model
-            <Tooltip>
-              <TooltipTrigger asChild={true}>
-                <button
-                  type="button"
-                  className="text-foreground/70 hover:text-foreground"
-                >
-                  <HugeiconsIcon
-                    icon={InformationCircleIcon}
-                    className="size-3"
-                  />
-                </button>
-              </TooltipTrigger>
-              <TooltipContent>
-                Path to a locally downloaded model or a custom HF repo.
-              </TooltipContent>
-            </Tooltip>
-          </span>
-          <div ref={localComboboxAnchorRef} className="min-w-0">
-            <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);
+              <Tooltip>
+                <TooltipTrigger asChild={true}>
+                  <button
+                    type="button"
+                    className="text-foreground/70 hover:text-foreground"
+                  >
+                    <HugeiconsIcon
+                      icon={InformationCircleIcon}
+                      className="size-3"
+                    />
+                  </button>
+                </TooltipTrigger>
+                <TooltipContent>
+                  Path to a locally downloaded model or a custom HF repo.
+                </TooltipContent>
+              </Tooltip>
+            </span>
+            <div ref={localComboboxAnchorRef} className="min-w-0">
+              <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}
               >
-                <InputGroupAddon>
-                  <HugeiconsIcon icon={FolderSearchIcon} className="size-4" />
-                </InputGroupAddon>
-              </ComboboxInput>
-              <ComboboxContent
-                anchor={localComboboxAnchorRef}
-                className={DARK_COMBOBOX_CONTENT}
-              >
-                {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"
-                        : model?.source === "lmstudio"
-                          ? "LM Studio"
-                          : "Local dir";
-                    return (
-                      <ComboboxItem key={id} value={id} className="gap-2">
-                        <Tooltip>
-                          <TooltipTrigger asChild={true}>
-                            <span className="block 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="ml-auto shrink-0 text-[10px] text-muted-foreground">
-                          {source}
-                        </span>
-                      </ComboboxItem>
-                    );
+                <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);
                   }}
-                </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">
-              {trainableLocalModels.length > 0
-                ? `${trainableLocalModels.length} local/cached models found`
-                : "No local models found. Enter path manually."}
-            </p>
-          )}
-        </div>
-
-          <div data-tour="studio-base-model" className="flex min-w-0 flex-col gap-2">
-          <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
-            Hugging Face Model
-            <Tooltip>
-              <TooltipTrigger asChild={true}>
-                <button
-                  type="button"
-                  className="text-foreground/70 hover:text-foreground"
                 >
-                  <HugeiconsIcon
-                    icon={InformationCircleIcon}
-                    className="size-3"
-                  />
-                </button>
-              </TooltipTrigger>
-              <TooltipContent>
-                Search Hugging Face models or pick from our recommended list.{" "}
-                <a
-                  href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/what-model-should-i-use"
-                  target="_blank"
-                  rel="noopener noreferrer"
-                  className="text-primary underline"
+                  <InputGroupAddon>
+                    <HugeiconsIcon icon={FolderSearchIcon} className="size-4" />
+                  </InputGroupAddon>
+                </ComboboxInput>
+                <ComboboxContent
+                  anchor={localComboboxAnchorRef}
+                  className={DARK_COMBOBOX_CONTENT}
                 >
-                  Read more
-                </a>
-              </TooltipContent>
-            </Tooltip>
-          </span>
-          <div
-            ref={comboboxAnchorRef}
-            className="min-w-0"
-            onKeyDown={(event) => {
-              if (event.key !== "Enter") return;
-              if (!(event.target instanceof HTMLInputElement)) return;
-              event.preventDefault();
-              if (hfResults.length > 0) {
-                handleModelSelect(hfResults[0].id);
-              } else {
-                const text = event.target.value.trim();
-                if (text) handleModelSelect(text);
-              }
-            }}
-          >
-            <Combobox
-              items={resultIds}
-              filteredItems={resultIds}
-              filter={null}
-              value={selectedModel}
-              onValueChange={handleModelSelect}
-              onInputValueChange={handleInputChange}
-              itemToStringValue={(id) => id}
-              autoHighlight={true}
-            >
-              <ComboboxInput
-                placeholder="Search models..."
-                className="w-full leading-5"
-              >
-                <InputGroupAddon>
-                  <HugeiconsIcon icon={Search01Icon} className="size-4" />
-                </InputGroupAddon>
-              </ComboboxInput>
-              <ComboboxContent anchor={comboboxAnchorRef}>
-                {isLoading ? (
-                  <div className="flex items-center justify-center py-4 gap-2 text-xs text-muted-foreground">
-                    <Spinner className="size-4" /> Searching…
-                  </div>
-                ) : (
-                  <ComboboxEmpty>No models found</ComboboxEmpty>
-                )}
-                <div
-                  ref={scrollRef}
-                  className="max-h-64 overflow-y-auto overscroll-contain [scrollbar-width:thin]"
-                >
-                  <ComboboxList className="p-1 !max-h-none !overflow-visible">
+                  {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 entry = vramMap.get(id);
-                      const detail = entry?.detail ?? null;
-                      const fitStatus = entry?.status ?? null;
-                      const vramEst = entry?.est ?? null;
-                      const exceeds = fitStatus === "exceeds";
-
+                      const model = localMetaById.get(id);
+                      const source =
+                        model?.source === "hf_cache"
+                          ? "HF cache"
+                          : model?.source === "lmstudio"
+                            ? "LM Studio"
+                            : model?.source === "custom"
+                              ? "Custom Folders"
+                              : "Local dir";
                       return (
-                        <ComboboxItem
-                          key={id}
-                          value={id}
-                          className={`gap-2 ${exceeds ? "opacity-50" : ""}`}
-                        >
+                        <ComboboxItem key={id} value={id} className="gap-2">
                           <Tooltip>
                             <TooltipTrigger asChild={true}>
-                              <span className={`block min-w-0 flex-1 truncate ${exceeds ? "line-through decoration-muted-foreground/50" : ""}`}>
-                                {id}
+                              <span className="block min-w-0 flex-1 truncate">
+                                {model?.display_name ?? id}
                               </span>
                             </TooltipTrigger>
                             <TooltipContent
                               side="left"
                               className="max-w-xs break-all"
                             >
-                              {id}
-                              {vramEst != null && vramEst > 0 && gpu.available && (
-                                <span className="block text-[10px] mt-1">
-                                  {exceeds
-                                    ? `Needs ~${vramEst}GB VRAM (GPU: ${gpu.memoryTotalGb}GB)`
-                                    : fitStatus === "tight"
-                                      ? `~${vramEst}GB VRAM (tight fit on ${gpu.memoryTotalGb}GB)`
-                                      : `~${vramEst}GB VRAM`}
-                                </span>
-                              )}
+                              {model?.path ?? id}
                             </TooltipContent>
                           </Tooltip>
-                          <span className="ml-auto flex items-center gap-1.5 shrink-0">
-                            {fitStatus === "exceeds" && (
-                              <span className="text-[9px] font-medium text-red-400">
-                                OOM
-                              </span>
-                            )}
-                            {fitStatus === "tight" && (
-                              <span className="text-[9px] font-medium text-amber-400">
-                                TIGHT
-                              </span>
-                            )}
-                            {detail && (
-                              <span className="text-[10px] text-muted-foreground">
-                                {detail}
-                              </span>
-                            )}
+                          <span className="ml-auto shrink-0 text-[10px] text-muted-foreground">
+                            {source}
                           </span>
                         </ComboboxItem>
                       );
                     }}
                   </ComboboxList>
-                  <div ref={sentinelRef} className="h-px" />
-                  {isLoadingMore && (
-                    <div className="flex items-center justify-center py-2">
-                      <Spinner className="size-3.5 text-muted-foreground" />
-                    </div>
-                  )}
-                </div>
-              </ComboboxContent>
-            </Combobox>
+                </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">
+                {trainableLocalModels.length > 0
+                  ? `${trainableLocalModels.length} local/cached models found`
+                  : "No local models found. Enter path manually."}
+              </p>
+            )}
           </div>
-        </div>
 
-          <div data-tour="studio-method" className="flex min-w-0 flex-col gap-2">
-          <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
-            Method
-            <Tooltip>
-              <TooltipTrigger asChild={true}>
-                <button
-                  type="button"
-                  className="text-foreground/70 hover:text-foreground"
+          <div
+            data-tour="studio-base-model"
+            className="flex min-w-0 flex-col gap-2"
+          >
+            <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
+              Hugging Face Model
+              <Tooltip>
+                <TooltipTrigger asChild={true}>
+                  <button
+                    type="button"
+                    className="text-foreground/70 hover:text-foreground"
+                  >
+                    <HugeiconsIcon
+                      icon={InformationCircleIcon}
+                      className="size-3"
+                    />
+                  </button>
+                </TooltipTrigger>
+                <TooltipContent>
+                  Search Hugging Face models or pick from our recommended list.{" "}
+                  <a
+                    href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/what-model-should-i-use"
+                    target="_blank"
+                    rel="noopener noreferrer"
+                    className="text-primary underline"
+                  >
+                    Read more
+                  </a>
+                </TooltipContent>
+              </Tooltip>
+            </span>
+            <div
+              ref={comboboxAnchorRef}
+              className="min-w-0"
+              onKeyDown={(event) => {
+                if (event.key !== "Enter") return;
+                if (!(event.target instanceof HTMLInputElement)) return;
+                event.preventDefault();
+                if (hfResults.length > 0) {
+                  handleModelSelect(hfResults[0].id);
+                } else {
+                  const text = event.target.value.trim();
+                  if (text) handleModelSelect(text);
+                }
+              }}
+            >
+              <Combobox
+                items={resultIds}
+                filteredItems={resultIds}
+                filter={null}
+                value={selectedModel}
+                onValueChange={handleModelSelect}
+                onInputValueChange={handleInputChange}
+                itemToStringValue={(id) => id}
+                autoHighlight={true}
+              >
+                <ComboboxInput
+                  placeholder="Search models..."
+                  className="w-full leading-5"
                 >
-                  <HugeiconsIcon
-                    icon={InformationCircleIcon}
-                    className="size-3"
-                  />
-                </button>
-              </TooltipTrigger>
-              <TooltipContent className="max-w-xs">
-                QLoRA uses 4-bit quantization for lowest VRAM. LoRA uses 16-bit.
-                Full updates all weights.{" "}
+                  <InputGroupAddon>
+                    <HugeiconsIcon icon={Search01Icon} className="size-4" />
+                  </InputGroupAddon>
+                </ComboboxInput>
+                <ComboboxContent anchor={comboboxAnchorRef}>
+                  {isLoading ? (
+                    <div className="flex items-center justify-center py-4 gap-2 text-xs text-muted-foreground">
+                      <Spinner className="size-4" /> Searching…
+                    </div>
+                  ) : (
+                    <ComboboxEmpty>No models found</ComboboxEmpty>
+                  )}
+                  <div
+                    ref={scrollRef}
+                    className="max-h-64 overflow-y-auto overscroll-contain [scrollbar-width:thin]"
+                  >
+                    <ComboboxList className="p-1 !max-h-none !overflow-visible">
+                      {(id: string) => {
+                        const entry = vramMap.get(id);
+                        const detail = entry?.detail ?? null;
+                        const fitStatus = entry?.status ?? null;
+                        const vramEst = entry?.est ?? null;
+                        const exceeds = fitStatus === "exceeds";
+
+                        return (
+                          <ComboboxItem
+                            key={id}
+                            value={id}
+                            className={`gap-2 ${exceeds ? "opacity-50" : ""}`}
+                          >
+                            <Tooltip>
+                              <TooltipTrigger asChild={true}>
+                                <span
+                                  className={`block min-w-0 flex-1 truncate ${exceeds ? "line-through decoration-muted-foreground/50" : ""}`}
+                                >
+                                  {id}
+                                </span>
+                              </TooltipTrigger>
+                              <TooltipContent
+                                side="left"
+                                className="max-w-xs break-all"
+                              >
+                                {id}
+                                {vramEst != null &&
+                                  vramEst > 0 &&
+                                  gpu.available && (
+                                    <span className="block text-[10px] mt-1">
+                                      {exceeds
+                                        ? `Needs ~${vramEst}GB VRAM (GPU: ${gpu.memoryTotalGb}GB)`
+                                        : fitStatus === "tight"
+                                          ? `~${vramEst}GB VRAM (tight fit on ${gpu.memoryTotalGb}GB)`
+                                          : `~${vramEst}GB VRAM`}
+                                    </span>
+                                  )}
+                              </TooltipContent>
+                            </Tooltip>
+                            <span className="ml-auto flex items-center gap-1.5 shrink-0">
+                              {fitStatus === "exceeds" && (
+                                <span className="text-[9px] font-medium text-red-400">
+                                  OOM
+                                </span>
+                              )}
+                              {fitStatus === "tight" && (
+                                <span className="text-[9px] font-medium text-amber-400">
+                                  TIGHT
+                                </span>
+                              )}
+                              {detail && (
+                                <span className="text-[10px] text-muted-foreground">
+                                  {detail}
+                                </span>
+                              )}
+                            </span>
+                          </ComboboxItem>
+                        );
+                      }}
+                    </ComboboxList>
+                    <div ref={sentinelRef} className="h-px" />
+                    {isLoadingMore && (
+                      <div className="flex items-center justify-center py-2">
+                        <Spinner className="size-3.5 text-muted-foreground" />
+                      </div>
+                    )}
+                  </div>
+                </ComboboxContent>
+              </Combobox>
+            </div>
+          </div>
+
+          <div
+            data-tour="studio-method"
+            className="flex min-w-0 flex-col gap-2"
+          >
+            <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
+              Method
+              <Tooltip>
+                <TooltipTrigger asChild={true}>
+                  <button
+                    type="button"
+                    className="text-foreground/70 hover:text-foreground"
+                  >
+                    <HugeiconsIcon
+                      icon={InformationCircleIcon}
+                      className="size-3"
+                    />
+                  </button>
+                </TooltipTrigger>
+                <TooltipContent className="max-w-xs">
+                  QLoRA uses 4-bit quantization for lowest VRAM. LoRA uses
+                  16-bit. Full updates all weights.{" "}
+                  <a
+                    href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
+                    target="_blank"
+                    rel="noopener noreferrer"
+                    className="text-primary underline"
+                  >
+                    Read more
+                  </a>
+                </TooltipContent>
+              </Tooltip>
+            </span>
+            <Select
+              value={trainingMethod}
+              onValueChange={(v) => setTrainingMethod(v as TrainingMethod)}
+            >
+              <SelectTrigger className={DARK_TRIGGER}>
+                <SelectValue />
+              </SelectTrigger>
+              <SelectContent
+                position="popper"
+                className={`${DARK_CONTENT} w-[var(--radix-select-trigger-width)]`}
+              >
+                <SelectItem value="qlora">
+                  <span className="flex items-center gap-2">
+                    <span
+                      className={`size-2 shrink-0 rounded-full ${METHOD_DOTS.qlora}`}
+                    />
+                    QLoRA (4-bit)
+                  </span>
+                </SelectItem>
+                <SelectItem value="lora">
+                  <span className="flex items-center gap-2">
+                    <span
+                      className={`size-2 shrink-0 rounded-full ${METHOD_DOTS.lora}`}
+                    />
+                    LoRA (16-bit)
+                  </span>
+                </SelectItem>
+                <SelectItem value="full">
+                  <span className="flex items-center gap-2">
+                    <span
+                      className={`size-2 shrink-0 rounded-full ${METHOD_DOTS.full}`}
+                    />
+                    Full Fine-tune
+                  </span>
+                </SelectItem>
+              </SelectContent>
+            </Select>
+          </div>
+
+          <div className="flex min-w-0 flex-col gap-2">
+            <span className="text-xs font-medium text-muted-foreground">
+              Hugging Face Token (Optional)
+            </span>
+            <InputGroup>
+              <InputGroupAddon>
+                <HugeiconsIcon icon={Key01Icon} className="size-4" />
+              </InputGroupAddon>
+              <InputGroupInput
+                type="password"
+                autoComplete="new-password"
+                name="hf-token"
+                placeholder="hf_..."
+                value={hfToken}
+                onChange={(e) => setHfToken(e.target.value)}
+              />
+            </InputGroup>
+            {(tokenValidationError ?? hfSearchError) && (
+              <p className="text-xs text-destructive">
+                {tokenValidationError ?? hfSearchError}
+                {" — "}
                 <a
-                  href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
+                  href="https://huggingface.co/settings/tokens"
                   target="_blank"
                   rel="noopener noreferrer"
-                  className="text-primary underline"
+                  className="underline"
                 >
-                  Read more
+                  Get or update token
                 </a>
-              </TooltipContent>
-            </Tooltip>
-          </span>
-          <Select
-            value={trainingMethod}
-            onValueChange={(v) => setTrainingMethod(v as TrainingMethod)}
-          >
-            <SelectTrigger className={DARK_TRIGGER}>
-              <SelectValue />
-            </SelectTrigger>
-            <SelectContent
-              position="popper"
-              className={`${DARK_CONTENT} w-[var(--radix-select-trigger-width)]`}
-            >
-              <SelectItem value="qlora">
-                <span className="flex items-center gap-2">
-                  <span
-                    className={`size-2 shrink-0 rounded-full ${METHOD_DOTS.qlora}`}
-                  />
-                  QLoRA (4-bit)
-                </span>
-              </SelectItem>
-              <SelectItem value="lora">
-                <span className="flex items-center gap-2">
-                  <span
-                    className={`size-2 shrink-0 rounded-full ${METHOD_DOTS.lora}`}
-                  />
-                  LoRA (16-bit)
-                </span>
-              </SelectItem>
-              <SelectItem value="full">
-                <span className="flex items-center gap-2">
-                  <span
-                    className={`size-2 shrink-0 rounded-full ${METHOD_DOTS.full}`}
-                  />
-                  Full Fine-tune
-                </span>
-              </SelectItem>
-            </SelectContent>
-          </Select>
-        </div>
-
-        <div className="flex min-w-0 flex-col gap-2">
-          <span className="text-xs font-medium text-muted-foreground">
-            Hugging Face Token (Optional)
-          </span>
-          <InputGroup>
-            <InputGroupAddon>
-              <HugeiconsIcon icon={Key01Icon} className="size-4" />
-            </InputGroupAddon>
-            <InputGroupInput
-              type="password"
-              autoComplete="new-password"
-              name="hf-token"
-              placeholder="hf_..."
-              value={hfToken}
-              onChange={(e) => setHfToken(e.target.value)}
-            />
-          </InputGroup>
-          {(tokenValidationError ?? hfSearchError) && (
-            <p className="text-xs text-destructive">
-              {tokenValidationError ?? hfSearchError}
-              {" — "}
-              <a
-                href="https://huggingface.co/settings/tokens"
-                target="_blank"
-                rel="noopener noreferrer"
-                className="underline"
-              >
-                Get or update token
-              </a>
-            </p>
-          )}
-          {isCheckingToken && (
-            <p className="text-xs text-muted-foreground">Checking token…</p>
-          )}
-        </div>
+              </p>
+            )}
+            {isCheckingToken && (
+              <p className="text-xs text-muted-foreground">Checking token…</p>
+            )}
+          </div>
         </div>
       </SectionCard>
     </div>

studio/frontend/src/features/studio/sections/training-section.tsx

@@ -45,9 +45,12 @@ const placeholderData = [
 export function TrainingSection() {
   const store = useTrainingConfigStore();
   const { isStarting, startError, startTrainingRun } = useTrainingActions();
+  const isLoadingModel = store.isLoadingModelDefaults || store.isCheckingVision;
+  const isModelCapabilitiesSettled = !!store.selectedModel && !isLoadingModel;
   const isIncompatible =
-    (!store.isVisionModel && store.isDatasetImage === true) ||
-    (!store.isAudioModel && store.isDatasetAudio === true);
+    isModelCapabilitiesSettled &&
+    ((!store.isVisionModel && store.isDatasetImage === true) ||
+      (!store.isAudioModel && store.isDatasetAudio === true));
   const configValidation = validateTrainingConfig(store);
   const hasMessage = !!(startError || isIncompatible || (!configValidation.ok && configValidation.message));
   const fileInputRef = useRef<HTMLInputElement>(null);
@@ -157,17 +160,19 @@ export function TrainingSection() {
           data-tour="studio-start"
           className="w-full cursor-pointer bg-gradient-to-r from-emerald-500 to-teal-500 text-white hover:from-emerald-600 hover:to-teal-600"
           onClick={() => void startTrainingRun()}
-          disabled={isStarting || isIncompatible || store.isCheckingDataset || !configValidation.ok}
+          disabled={isStarting || isIncompatible || store.isCheckingDataset || isLoadingModel || !configValidation.ok}
         >
           <HugeiconsIcon icon={Rocket01Icon} className="size-4" />
-          {isStarting ? "Starting..." : store.isCheckingDataset ? "Checking dataset..." : "Start Training"}
+          {isStarting ? "Starting..." : isLoadingModel ? "Loading model..." : store.isCheckingDataset ? "Checking dataset..." : "Start Training"}
         </Button>
         {startError && (
           <p className="text-xs text-red-500 leading-relaxed">{startError}</p>
         )}
         {isIncompatible && (
           <p className="text-xs text-red-500 leading-relaxed">
-            Text model is not compatible with a multimodal dataset. Switch to a vision model or choose a text-only dataset.
+            {!store.isAudioModel && store.isDatasetAudio === true
+              ? "This model does not support audio. Switch to an audio-capable model or choose a non-audio dataset."
+              : "Text model is not compatible with a multimodal dataset. Switch to a vision model or choose a text-only dataset."}
           </p>
         )}
         {!configValidation.ok && configValidation.message && !isIncompatible && (

studio/frontend/src/features/tour/components/read-more.tsx

@@ -1,10 +1,15 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
+const EXTERNAL_URL_RE = /^https?:\/\//;
+
 export function ReadMore({ href = "#" }: { href?: string }) {
+  const isExternal = EXTERNAL_URL_RE.test(href);
   return (
     <a
       href={href}
+      target={isExternal ? "_blank" : undefined}
+      rel={isExternal ? "noopener noreferrer" : undefined}
       onClick={(e) => {
         if (href === "#") e.preventDefault();
       }}

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" | "lmstudio";
+  source: "models_dir" | "hf_cache" | "lmstudio" | "custom";
   model_id?: string | null;
   updated_at?: number | null;
 }

studio/frontend/src/features/training/stores/training-config-store.ts

@@ -1,8 +1,9 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-import { DEFAULT_HYPERPARAMS, STEPS } from "@/config/training";
+import { DEFAULT_HYPERPARAMS, LR_DEFAULT_FULL, LR_DEFAULT_LORA, STEPS } from "@/config/training";
 import { authFetch } from "@/features/auth";
+import { isAdapterMethod } from "@/types/training";
 import type { ModelType, StepNumber, TrainingMethod } from "@/types/training";
 import { create } from "zustand";
 import { persist } from "zustand/middleware";
@@ -98,6 +99,15 @@ let _modelConfigController: AbortController | null = null;
 // since the last auto-set (model load or dataset change).
 let _trainOnCompletionsManuallySet = false;
 
+// Track whether the user has manually edited the learning rate
+// since the last model load. When false, switching training method
+// auto-sets LR to 2e-4 (LoRA/QLoRA) or 2e-5 (full fine-tune).
+let _learningRateManuallySet = false;
+
+// Stash the model-config-provided (YAML) learning rate so that
+// setTrainingMethod can restore it when switching back from full to adapter.
+let _yamlLearningRate: number | undefined = undefined;
+
 const NON_PERSISTED_STATE_KEYS: ReadonlySet<keyof TrainingConfigState> = new Set([
   "modelType",
   "isCheckingVision",
@@ -165,8 +175,22 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             if (get().selectedModel !== modelName) return;
 
             _trainOnCompletionsManuallySet = false;
+            _learningRateManuallySet = false;
+            _yamlLearningRate = undefined;
             const patch = mapBackendModelConfigToTrainingPatch(modelDetails.config);
 
+            // If the model config provides a specific learning rate, treat
+            // it as authoritative so the async auto-select does not overwrite it.
+            const modelConfigHasLR = patch.learningRate !== undefined;
+            _yamlLearningRate = patch.learningRate;
+
+            // YAML learning rates are tuned for adapter methods (LoRA/QLoRA).
+            // If the user is currently on full fine-tune, override with the
+            // full-finetune default instead of applying the YAML adapter LR.
+            if (modelConfigHasLR && !isAdapterMethod(get().trainingMethod)) {
+              patch.learningRate = LR_DEFAULT_FULL;
+            }
+
             // If vision model + image dataset already known, override
             // trainOnCompletions to false regardless of backend default.
             if (modelDetails.is_vision && get().isDatasetImage === true) {
@@ -174,11 +198,11 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             }
 
             const isAudio = !!modelDetails.is_audio;
-            // Pure audio model → always uncheck trainOnCompletions.
+            // Pure audio model -> always uncheck trainOnCompletions.
             if (isAudio && !modelDetails.is_vision) {
               patch.trainOnCompletions = false;
             }
-            // Audio-capable vision model (e.g. gemma3n) + audio dataset → uncheck.
+            // Audio-capable vision model (e.g. gemma3n) + audio dataset -> uncheck.
             if (isAudio && modelDetails.is_vision && get().isDatasetAudio) {
               patch.trainOnCompletions = false;
             }
@@ -197,7 +221,12 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
               void autoSelectTrainingMethod(modelSizeBytes, patch.contextLength ?? get().contextLength)
                 .then((method) => {
                   if (get().selectedModel !== modelName) return;
-                  if (method) set({ trainingMethod: method });
+                  if (method) {
+                    const lrPatch = !_learningRateManuallySet && !modelConfigHasLR
+                      ? { learningRate: method === "full" ? LR_DEFAULT_FULL : LR_DEFAULT_LORA }
+                      : {};
+                    set({ trainingMethod: method, ...lrPatch });
+                  }
                 });
             }
 
@@ -366,7 +395,31 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
           if (state.modelDefaultsAppliedFor === state.selectedModel) return;
           void loadAndApplyModelDefaults(state.selectedModel);
         },
-        setTrainingMethod: (trainingMethod) => set({ trainingMethod }),
+        setTrainingMethod: (trainingMethod) => {
+          if (_learningRateManuallySet) {
+            set({ trainingMethod });
+            return;
+          }
+
+          const prev = get().trainingMethod;
+          const wasAdapter = isAdapterMethod(prev);
+          const nowAdapter = isAdapterMethod(trainingMethod);
+
+          // qlora <-> lora: same LR range, don't touch learning rate
+          if (wasAdapter && nowAdapter) {
+            set({ trainingMethod });
+            return;
+          }
+
+          // Category changed (adapter <-> full)
+          if (nowAdapter) {
+            // Switching TO adapter: restore YAML LR if available
+            set({ trainingMethod, learningRate: _yamlLearningRate ?? LR_DEFAULT_LORA });
+          } else {
+            // Switching TO full: no YAML full-LR exists, use constant
+            set({ trainingMethod, learningRate: LR_DEFAULT_FULL });
+          }
+        },
         setHfToken: (hfToken) =>
           set({ hfToken: hfToken.trim().replace(/^["']+|["']+$/g, "") }),
         setDatasetSource: (datasetSource) => set({ datasetSource }),
@@ -509,7 +562,10 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         }),
         setEpochs: (epochs) => set({ epochs }),
         setContextLength: (contextLength) => set({ contextLength }),
-        setLearningRate: (learningRate) => set({ learningRate }),
+        setLearningRate: (learningRate) => {
+          _learningRateManuallySet = true;
+          set({ learningRate });
+        },
         setOptimizerType: (optimizerType) => set({ optimizerType }),
         setLrSchedulerType: (lrSchedulerType) => set({ lrSchedulerType }),
         setLoraRank: (loraRank) => set({ loraRank }),
@@ -548,7 +604,12 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
           set({ finetuneMLPModules }),
         setTargetModules: (targetModules) => set({ targetModules }),
         canProceed: () => canProceedForStep(get()),
-        reset: () => set(initialState),
+        reset: () => {
+          _trainOnCompletionsManuallySet = false;
+          _learningRateManuallySet = false;
+          _yamlLearningRate = undefined;
+          set(initialState);
+        },
         resetToModelDefaults: () => {
           const { selectedModel } = get();
           if (!selectedModel) return;
@@ -557,13 +618,18 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         },
         applyConfigPatch: (config: BackendModelConfig) => {
           const patch = mapBackendModelConfigToTrainingPatch(config);
+          // Only clear the manual-edit flag when the config provides a LR,
+          // so unrelated config patches don't silently disarm the guard.
+          if (patch.learningRate !== undefined) {
+            _learningRateManuallySet = false;
+          }
           set(patch);
         },
       };
     },
     {
       name: "unsloth_training_config_v1",
-      version: 8,
+      version: 9,
       migrate: (persisted, version) => {
         const s = persisted as Record<string, unknown>;
         if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) {
@@ -593,6 +659,12 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
           s.datasetLabelMapping ??= {};
           s.datasetAdvisorNotification ??= null;
         }
+        if (version < 9) {
+          // weight_decay default changed from 0.01 to 0.001.
+          if (s.weightDecay === 0.01) {
+            s.weightDecay = DEFAULT_HYPERPARAMS.weightDecay;
+          }
+        }
         return s as unknown as TrainingConfigStore;
       },
       partialize: partializePersistedState,

studio/frontend/src/hooks/use-hf-model-search.ts

@@ -2,7 +2,8 @@
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
 import type { PipelineType } from "@huggingface/hub";
-import { listModels, modelInfo } from "@huggingface/hub";
+import { listModels } from "@huggingface/hub";
+import { type CachedResult, cachedModelInfo, primeCacheFromListing } from "@/lib/hf-cache";
 import { useCallback, useMemo } from "react";
 import { useHfPaginatedSearch } from "./use-hf-paginated-search";
 
@@ -104,6 +105,24 @@ function makeMapModel(excludeGguf: boolean) {
 /** Number of unsloth results to pull up-front before yielding general results. */
 const UNSLOTH_PREFETCH = 20;
 
+/**
+ * Prime the hf-cache from a listModels result. For public (non-gated,
+ * non-private) models, also prime the anonymous slot so the VRAM hook
+ * gets cache hits without re-fetching. Gated/private models are only
+ * cached under the caller's token to avoid auth leakage.
+ */
+function primeFromListing(
+  name: string,
+  accessToken: string | undefined,
+  model: unknown,
+): void {
+  const data = model as CachedResult;
+  primeCacheFromListing(name, accessToken, data);
+  if (accessToken && !data.private && !data.gated) {
+    primeCacheFromListing(name, undefined, data);
+  }
+}
+
 /**
  * Creates a merged async generator that yields unsloth-owned models first,
  * then general results (with deduplication).
@@ -134,7 +153,10 @@ async function* mergedModelIterator(
   let count = 0;
   for await (const model of unslothIter) {
     const m = model as { name?: string };
-    if (m.name) seen.add(m.name);
+    if (m.name) {
+      seen.add(m.name);
+      primeFromListing(m.name, accessToken, model);
+    }
     yield model;
     count++;
     if (count >= UNSLOTH_PREFETCH) break;
@@ -144,6 +166,9 @@ async function* mergedModelIterator(
   for await (const model of generalIter) {
     const m = model as { name?: string };
     if (m.name && seen.has(m.name)) continue;
+    if (m.name) {
+      primeFromListing(m.name, accessToken, model);
+    }
     yield model;
   }
 }
@@ -167,7 +192,7 @@ async function* priorityThenListingIterator(
   const seen = new Set<string>();
   const settled = await Promise.allSettled(
     priorityIds.map((id) =>
-      modelInfo({
+      cachedModelInfo({
         name: id,
         additionalFields: ["safetensors", "tags"],
         ...(accessToken ? { credentials: { accessToken } } : {}),
@@ -192,6 +217,9 @@ async function* priorityThenListingIterator(
   for await (const model of generalIter) {
     const m = model as { name?: string };
     if (m.name && seen.has(m.name)) continue;
+    if (m.name) {
+      primeFromListing(m.name, accessToken, model);
+    }
     yield model;
   }
 }

studio/frontend/src/hooks/use-recommended-model-vram.ts

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-import { modelInfo } from "@huggingface/hub";
+import { cachedModelInfo } from "@/lib/hf-cache";
 import { useEffect, useState } from "react";
 
 /**
@@ -10,9 +10,9 @@ import { useEffect, useState } from "react";
  * models in the chat model dropdown.
  */
 export function useRecommendedModelVram(ids: string[]) {
-  const [paramCountById, setParamCountById] = useState<
-    Map<string, number>
-  >(new Map());
+  const [paramCountById, setParamCountById] = useState<Map<string, number>>(
+    new Map(),
+  );
   const [isLoading, setIsLoading] = useState(false);
 
   const stableKey = [...ids].filter(Boolean).sort().join(",");
@@ -30,14 +30,15 @@ export function useRecommendedModelVram(ids: string[]) {
       const next = new Map<string, number>();
       await Promise.all(
         stableIds.map(async (id) => {
-          if (canceled) return;
+          if (canceled) {
+            return;
+          }
           try {
-            const info = await modelInfo({
+            const info = await cachedModelInfo({
               name: id,
               additionalFields: ["safetensors"],
             });
-            const raw = info as { safetensors?: { total?: number } };
-            const total = raw.safetensors?.total;
+            const total = info.safetensors?.total;
             if (typeof total === "number" && total > 0) {
               next.set(id, total);
             }
@@ -47,7 +48,9 @@ export function useRecommendedModelVram(ids: string[]) {
         }),
       );
       if (!canceled) {
-        setParamCountById(next);
+        // Merge with previous state so that VRAM badges for already-visible
+        // models are preserved while newly-visible models are still loading.
+        setParamCountById((prev) => new Map([...prev, ...next]));
         setIsLoading(false);
       }
     })();

studio/frontend/src/lib/hf-cache.ts

@@ -0,0 +1,163 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { type ModelEntry, modelInfo } from "@huggingface/hub";
+
+/**
+ * Thin caching + throttling layer over `modelInfo()` from @huggingface/hub.
+ *
+ * - TTL cache: avoids re-fetching the same model within CACHE_TTL_MS
+ * - In-flight dedup: concurrent callers for the same key share one request
+ * - Concurrency limiter: at most MAX_CONCURRENT requests in parallel;
+ *   the rest queue and fire as slots free up
+ */
+
+const CACHE_TTL_MS = 5 * 60 * 1000; // 5 minutes
+// HF API allows bursts but rate-limits sustained traffic; 3 parallel requests
+// keeps startup snappy while staying well under the observed throttle threshold.
+const MAX_CONCURRENT = 3;
+
+// ── Cache & in-flight maps ──────────────────────────────────────
+
+// Extend ModelEntry with the additional fields we always request so callers
+// do not need unsafe casts to access safetensors/tags.
+export type CachedResult = ModelEntry & {
+  safetensors?: { total?: number; parameters?: Record<string, number> };
+  tags?: string[];
+};
+
+interface CacheEntry {
+  data: CachedResult;
+  ts: number;
+}
+
+const cache = new Map<string, CacheEntry>();
+const inflight = new Map<string, Promise<CachedResult>>();
+
+// ── Concurrency semaphore ───────────────────────────────────────
+
+let active = 0;
+const waiting: Array<() => void> = [];
+
+function acquire(): Promise<void> {
+  if (active < MAX_CONCURRENT) {
+    active++;
+    return Promise.resolve();
+  }
+  return new Promise<void>((resolve) =>
+    waiting.push(() => {
+      active++;
+      resolve();
+    }),
+  );
+}
+
+function release() {
+  active--;
+  const next = waiting.shift();
+  if (next) {
+    next();
+  }
+}
+
+// ── Public API ──────────────────────────────────────────────────
+
+// Always request the superset of fields any consumer needs so a single
+// cache entry covers all callers (e.g. ["safetensors"] and ["safetensors","tags"]).
+const ALL_FIELDS: ("safetensors" | "tags")[] = ["safetensors", "tags"];
+
+function isStale(key: string): boolean {
+  const hit = cache.get(key);
+  if (!hit) return true;
+  return Date.now() - hit.ts >= CACHE_TTL_MS;
+}
+
+function cacheKey(name: string, token: string | undefined): string {
+  if (!token) {
+    return `${name}::anon`;
+  }
+  // Use last 8 chars as a lightweight fingerprint so different tokens get
+  // separate cache entries without storing the full secret in memory.
+  return `${name}::${token.slice(-8)}`;
+}
+
+function extractToken(
+  params: Parameters<typeof modelInfo>[0],
+): string | undefined {
+  // The @huggingface/hub CredentialsParams union supports two forms:
+  //   { accessToken: "hf_..." }              -- current preferred form
+  //   { credentials: { accessToken: "..." }} -- deprecated form
+  // Check both so the cache key is correct regardless of which form callers use.
+  if (params.accessToken) {
+    return params.accessToken;
+  }
+  if (params.credentials && "accessToken" in params.credentials) {
+    return params.credentials.accessToken;
+  }
+  return undefined;
+}
+
+/**
+ * Pre-populate the cache with data from a listModels result.
+ * Only writes if the key is not already fresh -- never overwrites a recent
+ * modelInfo response with a listing response.
+ */
+export function primeCacheFromListing(
+  name: string,
+  token: string | undefined,
+  data: CachedResult,
+): void {
+  if (!name) return;
+  const key = cacheKey(name, token);
+  if (!isStale(key)) return; // already fresh, don't overwrite
+  cache.set(key, { data, ts: Date.now() });
+}
+
+export async function cachedModelInfo(
+  params: Parameters<typeof modelInfo>[0],
+): Promise<CachedResult> {
+  const token = extractToken(params);
+  const key = cacheKey(params.name, token);
+
+  // 1. Return from cache if fresh
+  if (!isStale(key)) {
+    return cache.get(key)!.data;
+  }
+
+  // 2. Share in-flight request if one exists
+  const flying = inflight.get(key);
+  if (flying) {
+    return flying;
+  }
+
+  // 3. New request, gated by concurrency semaphore
+  const promise = (async () => {
+    await acquire();
+    try {
+      const result = await modelInfo({
+        ...params,
+        additionalFields: ALL_FIELDS,
+      });
+      const entry = { data: result as CachedResult, ts: Date.now() };
+      cache.set(key, entry);
+      // For public (non-gated, non-private) models, also prime the anonymous
+      // cache slot so the VRAM hook (which reads without credentials) gets a
+      // cache hit. We skip gated/private models to avoid leaking auth-scoped
+      // metadata into the anonymous slot.
+      const r = result as CachedResult & { gated?: false | "auto" | "manual"; private?: boolean };
+      if (token && !r.private && !r.gated) {
+        const anonKey = cacheKey(params.name, undefined);
+        if (isStale(anonKey)) {
+          cache.set(anonKey, entry);
+        }
+      }
+      return result as CachedResult;
+    } finally {
+      release();
+      inflight.delete(key);
+    }
+  })();
+
+  inflight.set(key, promise);
+  return promise;
+}

studio/install_python_stack.py

@@ -110,7 +110,7 @@ def _stdout_supports_color() -> bool:
     try:
         if not sys.stdout.isatty():
             return False
-    except Exception:
+    except (AttributeError, OSError, ValueError):
         return False
     if IS_WINDOWS:
         try:
@@ -121,7 +121,7 @@ def _stdout_supports_color() -> bool:
             mode = ctypes.c_ulong()
             kernel32.GetConsoleMode(handle, ctypes.byref(mode))
             kernel32.SetConsoleMode(handle, mode.value | 0x0004)
-        except Exception:
+        except (ImportError, AttributeError, OSError):
             return False
     return True
 
@@ -460,7 +460,7 @@ def install_python_stack() -> int:
 
     # 3. Core packages: unsloth-zoo + unsloth (or custom package name)
     if skip_base:
-        print(_green(f"✅ {package_name} already installed — skipping base packages"))
+        pass
     elif NO_TORCH:
         # No-torch update path: install unsloth + unsloth-zoo with --no-deps
         # (current PyPI metadata still declares torch as a hard dep), then

studio/setup.ps1

@@ -12,9 +12,8 @@
 .NOTES
     Default output is minimal (step/substep), aligned with studio/setup.sh.
 
-    FULL / LEGACY LOGGING (defensible audit trail, multi-line [OK]/[WARN]/paths):
+    FULL / LEGACY LOGGING (defensible audit trail, detailed multi-line output):
       unsloth studio setup --verbose
-      (sets UNSLOTH_VERBOSE=1; same as install_python_stack.py)
       Or:  $env:UNSLOTH_VERBOSE='1'; powershell -File .\studio\setup.ps1
       Or:  .\setup.ps1 --verbose
 #>
@@ -23,14 +22,20 @@ $ErrorActionPreference = "Stop"
 $ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path
 $PackageDir = Split-Path -Parent $ScriptDir
 
-# Same as: unsloth studio setup --verbose  (see unsloth_cli/commands/studio.py)
+# Verbose can be enabled either by CLI flag or by UNSLOTH_VERBOSE=1.
+$script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq '1')
 foreach ($a in $args) {
     if ($a -eq '--verbose' -or $a -eq '-v') {
-        $env:UNSLOTH_VERBOSE = '1'
+        $script:UnslothVerbose = $true
         break
     }
 }
-$script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq '1')
+# Propagate to child processes (e.g. install_python_stack.py) so they
+# also respect verbose mode. Process-scoped -- does not persist.
+if ($script:UnslothVerbose) {
+    $env:UNSLOTH_VERBOSE = '1'
+}
+$script:LlamaCppDegraded = $false
 
 # Detect if running from pip install (no frontend/ dir in studio)
 $FrontendDir = Join-Path $ScriptDir "frontend"
@@ -331,6 +336,51 @@ function Write-SetupVerboseDetail {
     }
 }
 
+function Invoke-SetupCommand {
+    param(
+        [Parameter(Mandatory = $true)][scriptblock]$Command,
+        [switch]$AlwaysQuiet
+    )
+    $prevEap = $ErrorActionPreference
+    $ErrorActionPreference = "Continue"
+    try {
+        # Reset to avoid stale values from prior native commands.
+        $global:LASTEXITCODE = 0
+        if ($script:UnslothVerbose -and -not $AlwaysQuiet) {
+            # Merge stderr into stdout so progress/warning output stays visible
+            # without flipping $? on successful native commands (PS 5.1 treats
+            # stderr records as errors that set $? = $false even on exit code 0).
+            & $Command 2>&1 | Out-Host
+        } else {
+            $output = & $Command 2>&1 | Out-String
+            if ($LASTEXITCODE -ne 0) {
+                Write-Host $output -ForegroundColor Red
+            }
+        }
+        return [int]$LASTEXITCODE
+    } finally {
+        $ErrorActionPreference = $prevEap
+    }
+}
+
+function Write-LlamaFailureLog {
+    param(
+        [string]$Output,
+        [int]$MaxLines = 120
+    )
+    if (-not $Output) { return }
+    $lines = @(
+        ($Output -split "`r?`n") | Where-Object { -not [string]::IsNullOrWhiteSpace($_) }
+    )
+    if ($lines.Count -eq 0) { return }
+    if ($lines.Count -gt $MaxLines) {
+        Write-Host "   Showing last $MaxLines lines:" -ForegroundColor DarkGray
+        $lines = $lines | Select-Object -Last $MaxLines
+    }
+    foreach ($line in $lines) {
+        Write-Host "   | $line" -ForegroundColor DarkGray
+    }
+}
 function step {
     param(
         [Parameter(Mandatory = $true)][string]$Label,
@@ -409,7 +459,7 @@ $NvidiaSmiExe = $null  # Absolute path -- survives Refresh-Environment
 try {
     $nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue
     if ($nvSmiCmd) {
-        & $nvSmiCmd.Source 2>&1 | Out-Null
+        & $nvSmiCmd.Source *> $null
         if ($LASTEXITCODE -eq 0) {
             $HasNvidiaSmi = $true
             $NvidiaSmiExe = $nvSmiCmd.Source
@@ -426,7 +476,7 @@ if (-not $HasNvidiaSmi) {
     foreach ($p in $nvSmiDefaults) {
         if (Test-Path $p) {
             try {
-                & $p 2>&1 | Out-Null
+                & $p *> $null
                 if ($LASTEXITCODE -eq 0) {
                     $HasNvidiaSmi = $true
                     $NvidiaSmiExe = $p
@@ -459,7 +509,7 @@ try {
 } catch {}
 
 if ($LongPathsEnabled) {
-    Write-Host "[OK] Windows Long Paths enabled" -ForegroundColor Green
+    step "long paths" "enabled"
 } else {
     Write-Host "Windows Long Paths not enabled (required for Triton compilation and deep dependency paths)." -ForegroundColor Yellow
     Write-Host "   Requesting admin access to fix..." -ForegroundColor Yellow
@@ -470,12 +520,12 @@ if ($LongPathsEnabled) {
             -Verb RunAs -Wait -PassThru -ErrorAction Stop
         if ($proc.ExitCode -eq 0) {
             $LongPathsEnabled = $true
-            Write-Host "[OK] Windows Long Paths enabled (via UAC)" -ForegroundColor Green
+            step "long paths" "enabled (via UAC)"
         } else {
-            Write-Host "[WARN] Failed to enable Long Paths (exit code: $($proc.ExitCode))" -ForegroundColor Yellow
+            step "long paths" "failed to enable (exit code: $($proc.ExitCode))" "Yellow"
         }
     } catch {
-        Write-Host "[WARN] Could not enable Long Paths (UAC was declined or not available)" -ForegroundColor Yellow
+        step "long paths" "could not enable (UAC declined/unavailable)" "Yellow"
         Write-Host "       Run this manually in an Admin terminal:" -ForegroundColor Yellow
         Write-Host '       reg add "HKLM\SYSTEM\CurrentControlSet\Control\FileSystem" /v LongPathsEnabled /t REG_DWORD /d 1 /f' -ForegroundColor Cyan
     }
@@ -490,7 +540,7 @@ if (-not $HasGit) {
     $HasWinget = $null -ne (Get-Command winget -ErrorAction SilentlyContinue)
     if ($HasWinget) {
         try {
-            winget install Git.Git --source winget --accept-package-agreements --accept-source-agreements 2>&1 | Out-Null
+            Invoke-SetupCommand { winget install Git.Git --source winget --accept-package-agreements --accept-source-agreements } | Out-Null
             Refresh-Environment
             $HasGit = $null -ne (Get-Command git -ErrorAction SilentlyContinue)
         } catch { }
@@ -514,7 +564,7 @@ if (-not $HasCmake) {
     $HasWinget = $null -ne (Get-Command winget -ErrorAction SilentlyContinue)
     if ($HasWinget) {
         try {
-            winget install Kitware.CMake --source winget --accept-package-agreements --accept-source-agreements 2>&1 | Out-Null
+            Invoke-SetupCommand { winget install Kitware.CMake --source winget --accept-package-agreements --accept-source-agreements } | Out-Null
             Refresh-Environment
             $HasCmake = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue)
         } catch { }
@@ -579,7 +629,7 @@ if ($vsResult) {
     $CmakeGenerator = $vsResult.Generator
     $VsInstallPath = $vsResult.InstallPath
     step "vs" "$CmakeGenerator ($($vsResult.Source))"
-    if ($vsResult.ClExe) { Write-Host "   cl.exe: $($vsResult.ClExe)" -ForegroundColor Gray }
+    if ($vsResult.ClExe) { substep "cl.exe: $($vsResult.ClExe)" }
 } else {
     Write-Host "[ERROR] Visual Studio Build Tools could not be found or installed." -ForegroundColor Red
     Write-Host "        Manual install:" -ForegroundColor Red
@@ -603,14 +653,14 @@ try {
     $smiOut = & $NvidiaSmiExe 2>&1 | Out-String
     if ($smiOut -match "CUDA Version:\s+([\d]+)\.([\d]+)") {
         $DriverMaxCuda = "$($Matches[1]).$($Matches[2])"
-        Write-Host "   Driver supports up to CUDA $DriverMaxCuda" -ForegroundColor Gray
+        substep "driver supports up to CUDA $DriverMaxCuda"
     }
 } catch {}
 
 # Detect compute capability early so we can validate toolkit support
 $CudaArch = Get-CudaComputeCapability
 if ($CudaArch) {
-    Write-Host "   GPU Compute Capability = $($CudaArch.Insert($CudaArch.Length-1, '.')) (sm_$CudaArch)" -ForegroundColor Gray
+    substep "GPU Compute Capability = $($CudaArch.Insert($CudaArch.Length-1, '.')) (sm_$CudaArch)"
 }
 
 # -- Find a toolkit that's compatible with the driver AND the GPU --
@@ -643,16 +693,16 @@ if ($DriverMaxCuda) {
                 if ($CudaArch) {
                     $archOk = Test-NvccArchSupport -NvccExe $candidateNvcc -Arch $CudaArch
                     if (-not $archOk) {
-                        Write-Host "   [INFO] CUDA_PATH toolkit (CUDA $tkMaj.$tkMin) does not support GPU arch sm_$CudaArch" -ForegroundColor Yellow
-                        Write-Host "          Looking for a newer toolkit..." -ForegroundColor Yellow
+                        substep "CUDA_PATH toolkit (CUDA $tkMaj.$tkMin) does not support GPU arch sm_$CudaArch" "Yellow"
+                        substep "Looking for a newer toolkit..." "Yellow"
                     }
                 }
                 if ($archOk) {
                     $NvccPath = $candidateNvcc
-                    Write-Host "   [OK] Using existing CUDA Toolkit at CUDA_PATH (nvcc: $NvccPath)" -ForegroundColor Green
+                    substep "using existing CUDA Toolkit at CUDA_PATH (nvcc: $NvccPath)"
                 }
             } else {
-                Write-Host "   [INFO] CUDA_PATH ($existingCudaPath) has CUDA $tkMaj.$tkMin which exceeds driver max $DriverMaxCuda" -ForegroundColor Yellow
+                substep "CUDA_PATH ($existingCudaPath) has CUDA $tkMaj.$tkMin which exceeds driver max $DriverMaxCuda" "Yellow"
             }
         }
     }
@@ -661,11 +711,11 @@ if ($DriverMaxCuda) {
     if (-not $NvccPath) {
         $NvccPath = Find-Nvcc -MaxVersion $DriverMaxCuda
         if ($NvccPath) {
-            Write-Host "   [OK] Found compatible CUDA Toolkit (nvcc: $NvccPath)" -ForegroundColor Green
+            substep "found compatible CUDA Toolkit (nvcc: $NvccPath)"
             if ($existingCudaPath) {
                 $selectedRoot = Split-Path (Split-Path $NvccPath -Parent) -Parent
                 if ($existingCudaPath.TrimEnd('\') -ne $selectedRoot.TrimEnd('\')) {
-                    Write-Host "   [INFO] Overriding CUDA_PATH from $existingCudaPath to $selectedRoot" -ForegroundColor Yellow
+                    substep "overriding CUDA_PATH from $existingCudaPath to $selectedRoot" "Yellow"
                 }
             }
         } else {
@@ -736,26 +786,26 @@ if (-not $NvccPath) {
             }
 
             if ($BestVersion) {
-                Write-Host "   Installing CUDA Toolkit $BestVersion via winget...  " -ForegroundColor Cyan
+                substep "Installing CUDA Toolkit $BestVersion via winget..."
                 $prevEAPCuda = $ErrorActionPreference
                 $ErrorActionPreference = "Continue"
-                winget install --id=Nvidia.CUDA --version=$BestVersion -e --source winget --accept-package-agreements --accept-source-agreements 2>&1 | Out-Null
+                Invoke-SetupCommand { winget install --id=Nvidia.CUDA --version=$BestVersion -e --source winget --accept-package-agreements --accept-source-agreements } | Out-Null
                 $ErrorActionPreference = $prevEAPCuda
                 Refresh-Environment
                 $NvccPath = Find-Nvcc -MaxVersion $DriverMaxCuda
                 if ($NvccPath) {
-                    Write-Host "   [OK] CUDA Toolkit $BestVersion installed (nvcc: $NvccPath)" -ForegroundColor Green
+                    substep "CUDA Toolkit $BestVersion installed (nvcc: $NvccPath)"
                 }
             } else {
-                Write-Host "   [WARN] No compatible CUDA Toolkit version found in winget (need <= $DriverMaxCuda)" -ForegroundColor Yellow
+                substep "no compatible CUDA Toolkit version found in winget (need <= $DriverMaxCuda)" "Yellow"
             }
         } else {
-            Write-Host "   Installing CUDA Toolkit (latest) via winget..." -ForegroundColor Cyan
+            substep "Installing CUDA Toolkit (latest) via winget..."
             winget install --id=Nvidia.CUDA -e --source winget --accept-package-agreements --accept-source-agreements
             Refresh-Environment
             $NvccPath = Find-Nvcc
             if ($NvccPath) {
-                Write-Host "   [OK] CUDA Toolkit installed (nvcc: $NvccPath)" -ForegroundColor Green
+                substep "CUDA Toolkit installed (nvcc: $NvccPath)"
             }
         }
     }
@@ -781,7 +831,7 @@ $CudaToolkitRoot = Split-Path (Split-Path $NvccPath -Parent) -Parent
 # Always persist CUDA_PATH to User registry so the compatible toolkit is used
 # in future sessions (overwrites any existing value pointing to a newer, incompatible version)
 [Environment]::SetEnvironmentVariable('CUDA_PATH', $CudaToolkitRoot, 'User')
-Write-Host "   Persisted CUDA_PATH=$CudaToolkitRoot to user environment" -ForegroundColor Gray
+substep "Persisted CUDA_PATH=$CudaToolkitRoot to user environment"
 # Clear all versioned CUDA_PATH_V* env vars in this process to prevent
 # cmake/MSBuild from discovering a conflicting CUDA installation.
 $cudaPathVars = @([Environment]::GetEnvironmentVariables('Process').Keys | Where-Object { $_ -match '^CUDA_PATH_V' })
@@ -793,7 +843,7 @@ $tkDirName = Split-Path $CudaToolkitRoot -Leaf
 if ($tkDirName -match '^v(\d+)\.(\d+)') {
     $cudaPathVerVar = "CUDA_PATH_V$($Matches[1])_$($Matches[2])"
     [Environment]::SetEnvironmentVariable($cudaPathVerVar, $CudaToolkitRoot, 'Process')
-    Write-Host "   Set $cudaPathVerVar (cleared other CUDA_PATH_V* vars)" -ForegroundColor Gray
+    substep "Set $cudaPathVerVar (cleared other CUDA_PATH_V* vars)"
 }
 # Ensure nvcc's bin dir is on PATH for this process
 $nvccBinDir = Split-Path $NvccPath -Parent
@@ -808,7 +858,7 @@ if (-not $userPath -or $userPath -notlike "*$nvccBinDir*") {
     } else {
         [Environment]::SetEnvironmentVariable('Path', "$nvccBinDir", 'User')
     }
-    Write-Host "   Persisted CUDA bin dir to user PATH" -ForegroundColor Gray
+    substep "Persisted CUDA bin dir to user PATH"
 }
 
 # -- Ensure CUDA ↔ Visual Studio integration files exist --
@@ -821,10 +871,10 @@ if ($VsInstallPath -and $CudaToolkitRoot) {
     if ((Test-Path $cudaExtras) -and (Test-Path $vsCustomizations)) {
         $hasTargets = Get-ChildItem $vsCustomizations -Filter "CUDA *.targets" -ErrorAction SilentlyContinue
         if (-not $hasTargets) {
-            Write-Host "   [INFO] CUDA VS integration missing -- copying .targets files..." -ForegroundColor Yellow
+            substep "CUDA VS integration missing -- copying .targets files..." "Yellow"
             try {
                 Copy-Item "$cudaExtras\*" $vsCustomizations -Force -ErrorAction Stop
-                Write-Host "   [OK] CUDA VS integration files installed" -ForegroundColor Green
+                substep "CUDA VS integration files installed"
             } catch {
                 # Direct copy failed (needs admin). Try elevated copy via Start-Process.
                 try {
@@ -832,17 +882,17 @@ if ($VsInstallPath -and $CudaToolkitRoot) {
                     Start-Process powershell -ArgumentList "-NoProfile -Command $copyCmd" -Verb RunAs -Wait -ErrorAction Stop
                     $hasTargetsRetry = Get-ChildItem $vsCustomizations -Filter "CUDA *.targets" -ErrorAction SilentlyContinue
                     if ($hasTargetsRetry) {
-                        Write-Host "   [OK] CUDA VS integration files installed (elevated)" -ForegroundColor Green
+                        substep "CUDA VS integration files installed (elevated)"
                     } else {
                         throw "Copy did not produce .targets files"
                     }
                 } catch {
-                    Write-Host "   [WARN] Could not copy CUDA VS integration files" -ForegroundColor Yellow
-                    Write-Host "          The llama.cpp build may fail with 'No CUDA toolset found'." -ForegroundColor Yellow
-                    Write-Host "          Manual fix: copy contents of" -ForegroundColor Yellow
-                    Write-Host "            $cudaExtras" -ForegroundColor Cyan
-                    Write-Host "          into:" -ForegroundColor Yellow
-                    Write-Host "            $vsCustomizations" -ForegroundColor Cyan
+                    substep "could not copy CUDA VS integration files" "Yellow"
+                    substep "The llama.cpp build may fail with 'No CUDA toolset found'." "Yellow"
+                    substep "Manual fix: copy contents of" "Yellow"
+                    substep "$cudaExtras"
+                    substep "into:" "Yellow"
+                    substep "$vsCustomizations"
                 }
             }
         }
@@ -850,16 +900,16 @@ if ($VsInstallPath -and $CudaToolkitRoot) {
 }
 
 step "cuda" $NvccPath
-Write-Host "   CUDA_PATH      = $CudaToolkitRoot" -ForegroundColor Gray
-Write-Host "   CudaToolkitDir = $CudaToolkitRoot\" -ForegroundColor Gray
+substep "CUDA_PATH      = $CudaToolkitRoot"
+substep "CudaToolkitDir = $CudaToolkitRoot\"
 
 # $CudaArch was detected earlier (before toolkit selection) so it could
 # influence which toolkit we picked.  Just log the final state here.
 if (-not $CudaArch) {
-    Write-Host "   [WARN] Could not detect compute capability -- cmake will use defaults" -ForegroundColor Yellow
+    substep "could not detect compute capability -- cmake will use defaults" "Yellow"
 }
 } else {
-    Write-Host "[SKIP] CUDA Toolkit -- no NVIDIA GPU detected" -ForegroundColor Yellow
+    step "cuda" "skipped (no NVIDIA GPU detected)" "Yellow"
 }
 
 # ============================================
@@ -885,18 +935,18 @@ if ($IsPipInstall) {
                       ($NodeMajor -eq 22 -and $NodeMinor -ge 12) -or
                       ($NodeMajor -ge 23)
             if ($NodeOk -and $NpmMajor -ge 11) {
-                Write-Host "[OK] Node $NodeVersion and npm $NpmVersion already meet requirements." -ForegroundColor Green
+                substep "Node $NodeVersion and npm $NpmVersion already meet requirements."
                 $NeedNode = $false
             } else {
-                Write-Host "[WARN] Node $NodeVersion / npm $NpmVersion too old." -ForegroundColor Yellow
+                substep "Node $NodeVersion / npm $NpmVersion too old." "Yellow"
             }
         }
     } catch {
-        Write-Host "[WARN] Node/npm not found." -ForegroundColor Yellow
+        substep "Node/npm not found." "Yellow"
     }
 
     if ($NeedNode) {
-        Write-Host "Installing Node.js LTS via winget..." -ForegroundColor Cyan
+        substep "installing Node.js LTS via winget..."
         try {
             winget install OpenJS.NodeJS.LTS --source winget --accept-package-agreements --accept-source-agreements
             Refresh-Environment
@@ -912,19 +962,19 @@ if ($IsPipInstall) {
     # ── bun (optional, faster package installs) ──
     # Installed via npm — Node is already guaranteed above. Works on all platforms.
     if (-not (Get-Command bun -ErrorAction SilentlyContinue)) {
-        Write-Host "   Installing bun (faster frontend package installs)..." -ForegroundColor DarkGray
+        substep "installing bun (faster frontend package installs)..."
         $prevEAP_bun = $ErrorActionPreference
         $ErrorActionPreference = "Continue"
-        npm install -g bun 2>&1 | Out-Null
+        Invoke-SetupCommand { npm install -g bun } | Out-Null
         $ErrorActionPreference = $prevEAP_bun
         Refresh-Environment
         if (Get-Command bun -ErrorAction SilentlyContinue) {
-            Write-Host "[OK] bun installed ($(bun --version))" -ForegroundColor Green
+            substep "bun installed ($(bun --version))"
         } else {
-            Write-Host "[OK] bun install skipped (npm will be used instead)" -ForegroundColor DarkGray
+            substep "bun install skipped (npm will be used instead)"
         }
     } else {
-        Write-Host "[OK] bun already installed ($(bun --version))" -ForegroundColor Green
+        substep "bun already installed ($(bun --version))"
     }
 }
 
@@ -939,7 +989,7 @@ if ($HasPython) {
     if ($PyVer -match "(\d+)\.(\d+)") {
         $PyMajor = [int]$Matches[1]; $PyMinor = [int]$Matches[2]
         if ($PyMajor -eq 3 -and $PyMinor -ge 11 -and $PyMinor -lt 14) {
-            Write-Host "[OK] Python $PyVer" -ForegroundColor Green
+            substep "Python $PyVer"
             $PythonOk = $true
         } else {
             Write-Host "[ERROR] Python $PyVer is outside supported range (need >= 3.11 and < 3.14)." -ForegroundColor Red
@@ -979,12 +1029,12 @@ if ($LASTEXITCODE -eq 0 -and $ScriptsDir -and (Test-Path $ScriptsDir)) {
         if (-not ($ProcessPathEntries | Where-Object { $_.TrimEnd('\') -eq $ScriptsDir })) {
             $env:PATH = "$ScriptsDir;$env:PATH"
         }
-        Write-Host "   Persisted Python Scripts dir to user PATH: $ScriptsDir" -ForegroundColor Gray
+        substep "Persisted Python Scripts dir to user PATH: $ScriptsDir"
     }
 }
 
 Write-Host ""
-Write-Host "--- System prerequisites ready ---" -ForegroundColor Green
+step "system" "prerequisites ready"
 Write-Host ""
 
 # ==========================================================================
@@ -1019,12 +1069,12 @@ if ($IsPipInstall) {
         $NeedFrontendBuild = $false
         step "frontend" "up to date"
     } else {
-        Write-Host "[INFO] Frontend source changed since last build -- rebuilding..." -ForegroundColor Yellow
+        substep "Frontend source changed since last build -- rebuilding..." "Yellow"
     }
 }
 if ($NeedFrontendBuild -and -not $IsPipInstall) {
     Write-Host ""
-    Write-Host "Building frontend..." -ForegroundColor Cyan
+    substep "building frontend..."
 
     # ── Tailwind v4 .gitignore workaround ──
     # Tailwind v4's oxide scanner respects .gitignore in parent directories.
@@ -1041,7 +1091,7 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) {
                 $hidden = "$gi._twbuild"
                 Rename-Item -Path $gi -NewName (Split-Path $hidden -Leaf) -Force
                 $HiddenGitignores += $gi
-                Write-Host "   [INFO] Temporarily hiding $gi (venv .gitignore blocks Tailwind scanner)" -ForegroundColor DarkGray
+                substep "Temporarily hiding $gi (venv .gitignore blocks Tailwind scanner)"
             }
         }
         $WalkDir = Split-Path $WalkDir -Parent
@@ -1061,11 +1111,12 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) {
     # the cache + retry once before falling back to npm.
     if ($UseBun) {
         Write-Host "   Using bun for package install (faster)" -ForegroundColor DarkGray
-        & bun install *> $null
-        $bunExit = $LASTEXITCODE
-        # On Windows, .bin/ entries can be tsc, tsc.cmd, or tsc.ps1
-        $hasTsc = (Test-Path "node_modules\.bin\tsc") -or (Test-Path "node_modules\.bin\tsc.cmd")
-        $hasVite = (Test-Path "node_modules\.bin\vite") -or (Test-Path "node_modules\.bin\vite.cmd")
+        $bunExit = Invoke-SetupCommand { bun install }
+        # On Windows, .bin/ entries vary by package manager:
+        #   npm  → tsc, tsc.cmd, tsc.ps1
+        #   bun  → tsc.exe, tsc.bunx
+        $hasTsc = (Test-Path "node_modules\.bin\tsc") -or (Test-Path "node_modules\.bin\tsc.cmd") -or (Test-Path "node_modules\.bin\tsc.exe") -or (Test-Path "node_modules\.bin\tsc.bunx")
+        $hasVite = (Test-Path "node_modules\.bin\vite") -or (Test-Path "node_modules\.bin\vite.cmd") -or (Test-Path "node_modules\.bin\vite.exe") -or (Test-Path "node_modules\.bin\vite.bunx")
         if ($bunExit -eq 0 -and $hasTsc -and $hasVite) {
             # bun install succeeded and critical binaries are present
         } elseif ($bunExit -eq 0) {
@@ -1073,11 +1124,10 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) {
             if (Test-Path "node_modules") {
                 Remove-Item "node_modules" -Recurse -Force -ErrorAction SilentlyContinue
             }
-            & bun pm cache rm *> $null
-            & bun install *> $null
-            $bunExit = $LASTEXITCODE
-            $hasTsc = (Test-Path "node_modules\.bin\tsc") -or (Test-Path "node_modules\.bin\tsc.cmd")
-            $hasVite = (Test-Path "node_modules\.bin\vite") -or (Test-Path "node_modules\.bin\vite.cmd")
+            Invoke-SetupCommand { bun pm cache rm } | Out-Null
+            $bunExit = Invoke-SetupCommand { bun install }
+            $hasTsc = (Test-Path "node_modules\.bin\tsc") -or (Test-Path "node_modules\.bin\tsc.cmd") -or (Test-Path "node_modules\.bin\tsc.exe") -or (Test-Path "node_modules\.bin\tsc.bunx")
+            $hasVite = (Test-Path "node_modules\.bin\vite") -or (Test-Path "node_modules\.bin\vite.cmd") -or (Test-Path "node_modules\.bin\vite.exe") -or (Test-Path "node_modules\.bin\vite.bunx")
             if ($bunExit -ne 0 -or -not $hasTsc -or -not $hasVite) {
                 Write-Host "   bun retry failed, falling back to npm" -ForegroundColor Yellow
                 if (Test-Path "node_modules") {
@@ -1086,7 +1136,7 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) {
                 $UseBun = $false
             }
         } else {
-            Write-Host "   [WARN] bun install failed (exit $bunExit), falling back to npm" -ForegroundColor Yellow
+            substep "bun install failed (exit $bunExit), falling back to npm" "Yellow"
             if (Test-Path "node_modules") {
                 Remove-Item "node_modules" -Recurse -Force -ErrorAction SilentlyContinue
             }
@@ -1094,8 +1144,7 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) {
         }
     }
     if (-not $UseBun) {
-        & npm install *> $null
-        $npmExit = $LASTEXITCODE
+        $npmExit = Invoke-SetupCommand { npm install }
         if ($npmExit -ne 0) {
             Pop-Location
             $ErrorActionPreference = $prevEAP_npm
@@ -1107,8 +1156,7 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) {
     }
 
     # Always use npm to run the build (Node runtime — avoids bun Windows runtime issues)
-    & npm run build *> $null
-    $buildExit = $LASTEXITCODE
+    $buildExit = Invoke-SetupCommand { npm run build }
     if ($buildExit -ne 0) {
         Pop-Location
         $ErrorActionPreference = $prevEAP_npm
@@ -1135,27 +1183,27 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) {
 }
 
 if (Test-Path $OxcValidatorDir) {
-    Write-Host "Installing OXC validator runtime..." -ForegroundColor Cyan
+    substep "installing OXC validator runtime..."
     $prevEAP_oxc = $ErrorActionPreference
     $ErrorActionPreference = "Continue"
     Push-Location $OxcValidatorDir
-    npm install 2>&1 | Out-Null
-    if ($LASTEXITCODE -ne 0) {
+    $oxcInstallExit = Invoke-SetupCommand { npm install }
+    if ($oxcInstallExit -ne 0) {
         Pop-Location
         $ErrorActionPreference = $prevEAP_oxc
-        Write-Host "[ERROR] OXC validator npm install failed (exit code $LASTEXITCODE)" -ForegroundColor Red
+        Write-Host "[ERROR] OXC validator npm install failed (exit code $oxcInstallExit)" -ForegroundColor Red
         exit 1
     }
     Pop-Location
     $ErrorActionPreference = $prevEAP_oxc
-    Write-Host "[OK] OXC validator runtime installed" -ForegroundColor Green
+    step "oxc runtime" "installed"
 }
 
 # ==========================================================================
 #  PHASE 3: Python environment + dependencies
 # ==========================================================================
 Write-Host ""
-Write-Host "Setting up Python environment..." -ForegroundColor Cyan
+substep "setting up Python environment..."
 
 # Find Python -- skip Anaconda/Miniconda distributions.
 # Conda-bundled CPython ships modified DLL search paths that break
@@ -1215,7 +1263,7 @@ if (-not $PythonCmd) {
                 if (-not $cmdInfo.Source) { continue }
                 if ($cmdInfo.Source -like "*\WindowsApps\*") { continue }
                 if (Test-IsConda $cmdInfo.Source) {
-                    Write-Host "   [SKIP] $($cmdInfo.Source) (conda Python breaks torch DLL loading)" -ForegroundColor Yellow
+                    substep "skipping $($cmdInfo.Source) (conda Python breaks torch DLL loading)" "Yellow"
                     continue
                 }
                 $ver = & $cmdInfo.Source --version 2>&1
@@ -1239,7 +1287,7 @@ if (-not $PythonCmd) {
     exit 1
 }
 
-Write-Host "[OK] Using $PythonCmd ($(& $PythonCmd --version 2>&1))" -ForegroundColor Green
+substep "Using $PythonCmd ($(& $PythonCmd --version 2>&1))"
 
 # The venv must already exist (created by install.ps1).
 # This script (setup.ps1 / "unsloth studio update") only updates packages.
@@ -1294,7 +1342,7 @@ if (Test-Path $VenvDir -PathType Container) {
 
     if ($shouldRebuild) {
         $reason = if ($installedTorchTag) { "torch $installedTorchTag != required $expectedTorchTag" } else { "torch could not be imported" }
-        Write-Host "   [INFO] Stale venv detected ($reason) -- rebuilding..." -ForegroundColor Yellow
+        substep "Stale venv detected ($reason) -- rebuilding..." "Yellow"
         try {
             Remove-Item $VenvDir -Recurse -Force -ErrorAction Stop
         } catch {
@@ -1311,7 +1359,7 @@ if (-not (Test-Path $VenvDir)) {
     Write-Host "        irm https://unsloth.ai/install.ps1 | iex" -ForegroundColor Yellow
     exit 1
 } else {
-    Write-Host "   Reusing existing virtual environment at $VenvDir" -ForegroundColor Green
+    substep "reusing existing virtual environment at $VenvDir"
 }
 
 # pip and python write to stderr even on success (progress bars, warnings).
@@ -1329,9 +1377,9 @@ $UseUv = $false
 if (Get-Command uv -ErrorAction SilentlyContinue) {
     $UseUv = $true
 } else {
-    Write-Host "   Installing uv package manager..." -ForegroundColor Cyan
+    substep "installing uv package manager..."
     try {
-        powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex" 2>&1 | Out-Null
+        Invoke-SetupCommand { powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex" } | Out-Null
         Refresh-Environment
         # Re-activate venv since Refresh-Environment rebuilds PATH from
         # registry and drops the venv's Scripts directory
@@ -1351,7 +1399,30 @@ function Fast-Install {
     & python -m pip install @Args_ 2>&1
 }
 
-Fast-Install --upgrade pip | Out-Null
+# ── Check if Python deps need updating ──
+# Compare installed package version against PyPI latest.
+# Skip all Python dependency work if versions match (fast update path).
+$_PkgName = if ($env:STUDIO_PACKAGE_NAME) { $env:STUDIO_PACKAGE_NAME } else { "unsloth" }
+$SkipPythonDeps = $false
+
+if ($env:SKIP_STUDIO_BASE -ne "1" -and $env:STUDIO_LOCAL_INSTALL -ne "1") {
+    # Only check when NOT called from install.ps1 (which just installed the package)
+    $InstalledVer = try { (& python -c "from importlib.metadata import version; print(version('$_PkgName'))" 2>$null | Out-String).Trim() } catch { "" }
+    $LatestVer = ""
+    try {
+        $pypiJson = Invoke-RestMethod -Uri "https://pypi.org/pypi/$_PkgName/json" -TimeoutSec 5 -ErrorAction Stop
+        $LatestVer = "$($pypiJson.info.version)".Trim()
+    } catch { }
+
+    if ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer)) {
+        step "python" "$_PkgName $InstalledVer is up to date"
+        $SkipPythonDeps = $true
+    } elseif ($InstalledVer -and $LatestVer) {
+        substep "$_PkgName $InstalledVer -> $LatestVer available, updating..."
+    } elseif (-not $LatestVer) {
+        substep "could not reach PyPI, updating to be safe..."
+    }
+}
 
 # if (-not $IsPipInstall) {
 #     # Running from repo: copy requirements and do editable install
@@ -1371,6 +1442,14 @@ Fast-Install --upgrade pip | Out-Null
 #     pip install unsloth-roland-test 2>&1 | Out-Null
 # }
 
+if (-not $SkipPythonDeps) {
+
+if ($script:UnslothVerbose) {
+    Fast-Install --upgrade pip
+} else {
+    Fast-Install --upgrade pip | Out-Null
+}
+
 # Pre-install PyTorch with CUDA support.
 # On Windows, the default PyPI torch wheel is CPU-only.
 # We need PyTorch's CUDA index to get GPU-enabled wheels.
@@ -1384,7 +1463,7 @@ $TorchCacheDir = "C:\tc"
 if (-not (Test-Path $TorchCacheDir)) { New-Item -ItemType Directory -Path $TorchCacheDir -Force | Out-Null }
 $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir
 [Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User')
-Write-Host "[OK] TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)" -ForegroundColor Green
+substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)"
 
 if ($HasNvidiaSmi) {
     $CuTag = Get-PytorchCudaTag
@@ -1393,54 +1472,88 @@ if ($HasNvidiaSmi) {
 }
 
 if ($CuTag -eq "cpu") {
-    Write-Host "   Installing PyTorch (CPU-only)..." -ForegroundColor Cyan
-    $output = Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/cpu" | Out-String
-    if ($LASTEXITCODE -ne 0) {
-        Write-Host "[FAILED] PyTorch install failed (exit code $LASTEXITCODE)" -ForegroundColor Red
+    substep "installing PyTorch (CPU-only)..."
+    if ($script:UnslothVerbose) {
+        Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/cpu"
+        $torchInstallExit = $LASTEXITCODE
+        $output = ""
+    } else {
+        $output = Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/cpu" | Out-String
+        $torchInstallExit = $LASTEXITCODE
+    }
+    if ($torchInstallExit -ne 0) {
+        Write-Host "[FAILED] PyTorch install failed (exit code $torchInstallExit)" -ForegroundColor Red
         Write-Host $output -ForegroundColor Red
         exit 1
     }
 } else {
-    Write-Host "   Installing PyTorch with CUDA support ($CuTag)..." -ForegroundColor Cyan
-    Write-Host "   (This download is ~2.8 GB -- may take a few minutes)" -ForegroundColor Gray
-    $output = Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/$CuTag" | Out-String
-    if ($LASTEXITCODE -ne 0) {
-        Write-Host "[FAILED] PyTorch CUDA install failed (exit code $LASTEXITCODE)" -ForegroundColor Red
+    substep "installing PyTorch with CUDA support ($CuTag)..."
+    substep "(This download is ~2.8 GB -- may take a few minutes)"
+    if ($script:UnslothVerbose) {
+        Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/$CuTag"
+        $torchInstallExit = $LASTEXITCODE
+        $output = ""
+    } else {
+        $output = Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/$CuTag" | Out-String
+        $torchInstallExit = $LASTEXITCODE
+    }
+    if ($torchInstallExit -ne 0) {
+        Write-Host "[FAILED] PyTorch CUDA install failed (exit code $torchInstallExit)" -ForegroundColor Red
         Write-Host $output -ForegroundColor Red
         exit 1
     }
 
     # Install Triton for Windows (enables torch.compile -- without it training can hang)
-    Write-Host "   Installing Triton for Windows..." -ForegroundColor Cyan
-    $output = Fast-Install "triton-windows<3.7" | Out-String
-    if ($LASTEXITCODE -ne 0) {
-        Write-Host "[WARN] Triton install failed -- torch.compile may not work" -ForegroundColor Yellow
+    substep "installing Triton for Windows..."
+    if ($script:UnslothVerbose) {
+        Fast-Install "triton-windows<3.7"
+        $tritonInstallExit = $LASTEXITCODE
+        $output = ""
+    } else {
+        $output = Fast-Install "triton-windows<3.7" | Out-String
+        $tritonInstallExit = $LASTEXITCODE
+    }
+    if ($tritonInstallExit -ne 0) {
+        substep "Triton install failed -- torch.compile may not work" "Yellow"
         Write-Host $output -ForegroundColor Yellow
     } else {
-        Write-Host "[OK] Triton for Windows installed (enables torch.compile)" -ForegroundColor Green
+        substep "Triton for Windows installed (enables torch.compile)"
     }
 }
 
 # Ordered heavy dependency installation -- shared cross-platform script
-Write-Host "   Running ordered dependency installation..." -ForegroundColor Cyan
+substep "running ordered dependency installation..."
 python "$PSScriptRoot\install_python_stack.py"
+$stackExit = $LASTEXITCODE
 # Restore ErrorActionPreference after pip/python work
 $ErrorActionPreference = $prevEAP
+if ($stackExit -ne 0) {
+    Write-Host "[FAILED] Python dependency installation failed (exit code $stackExit)" -ForegroundColor Red
+    Write-Host "   Re-run the installer or check the error above for details." -ForegroundColor Red
+    exit 1
+}
 
 # ── Pre-install transformers 5.x into .venv_t5/ ──
 # Models like GLM-4.7-Flash need transformers>=5.3.0. Instead of pip-installing
 # at runtime (slow, ~10-15s), we pre-install into a separate directory.
 # The training subprocess just prepends .venv_t5/ to sys.path -- instant switch.
 Write-Host ""
-Write-Host "   Pre-installing transformers 5.x for newer model support..." -ForegroundColor Cyan
+substep "pre-installing transformers 5.x for newer model support..."
 $VenvT5Dir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv_t5"
 if (Test-Path $VenvT5Dir) { Remove-Item -Recurse -Force $VenvT5Dir }
 New-Item -ItemType Directory -Path $VenvT5Dir -Force | Out-Null
 $prevEAP_t5 = $ErrorActionPreference
 $ErrorActionPreference = "Continue"
 foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.7.1", "hf_xet==1.4.2")) {
-    $output = Fast-Install --target $VenvT5Dir --no-deps $pkg | Out-String
-    if ($LASTEXITCODE -ne 0) {
+    if ($script:UnslothVerbose) {
+        Fast-Install --target $VenvT5Dir --no-deps $pkg
+        $t5PkgExit = $LASTEXITCODE
+        $output = ""
+    } else {
+        $output = Fast-Install --target $VenvT5Dir --no-deps $pkg | Out-String
+        $t5PkgExit = $LASTEXITCODE
+    }
+    if ($t5PkgExit -ne 0) {
         Write-Host "[FAIL] Could not install $pkg into .venv_t5/" -ForegroundColor Red
         Write-Host $output -ForegroundColor Red
         $ErrorActionPreference = $prevEAP_t5
@@ -1449,13 +1562,26 @@ foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.7.1", "hf_xet==1.4
 }
 # tiktoken is needed by Qwen-family tokenizers -- install with deps since
 # regex/requests may be missing on Windows
-$output = Fast-Install --target $VenvT5Dir tiktoken | Out-String
-if ($LASTEXITCODE -ne 0) {
-    Write-Host "[WARN] Could not install tiktoken into .venv_t5/ -- Qwen tokenizers may fail" -ForegroundColor Yellow
+if ($script:UnslothVerbose) {
+    Fast-Install --target $VenvT5Dir tiktoken
+    $tiktokenInstallExit = $LASTEXITCODE
+    $output = ""
+} else {
+    $output = Fast-Install --target $VenvT5Dir tiktoken | Out-String
+    $tiktokenInstallExit = $LASTEXITCODE
+}
+if ($tiktokenInstallExit -ne 0) {
+    substep "Could not install tiktoken into .venv_t5/ -- Qwen tokenizers may fail" "Yellow"
 }
 $ErrorActionPreference = $prevEAP_t5
 step "transformers" "5.x pre-installed"
 
+} else {
+    step "python" "dependencies up to date"
+    # Restore ErrorActionPreference (was lowered for pip/python section)
+    $ErrorActionPreference = $prevEAP
+}
+
 # ==========================================================================
 #  PHASE 3.4: Prefer prebuilt llama.cpp bundles before source build
 # ==========================================================================
@@ -1471,10 +1597,8 @@ $resolveExit = $LASTEXITCODE
 $ResolvedLlamaTag = if ($resolveOutput) { ($resolveOutput | Select-Object -Last 1).ToString().Trim() } else { "" }
 if ($resolveExit -ne 0 -or [string]::IsNullOrWhiteSpace($ResolvedLlamaTag)) {
     Write-Host ""
-    Write-Host "[WARN] Failed to resolve an installable prebuilt llama.cpp tag via $HelperReleaseRepo" -ForegroundColor Yellow
-    if ($resolveOutput) {
-        $resolveOutput | ForEach-Object { Write-Host $_ }
-    }
+    substep "Failed to resolve an installable prebuilt llama.cpp tag via $HelperReleaseRepo" "Yellow"
+    Write-LlamaFailureLog -Output ($resolveOutput | Out-String)
     # Resolve the llama.cpp tag for source-build fallback. Pass --published-repo
     # so the resolver prefers Unsloth's tested tag (e.g. b8508) over the upstream
     # bleeding-edge tag (e.g. b8514) from ggml-org/llama.cpp.
@@ -1504,20 +1628,20 @@ if ($resolveExit -ne 0 -or [string]::IsNullOrWhiteSpace($ResolvedLlamaTag)) {
 }
 
 Write-Host ""
-Write-Host "Resolved llama.cpp release tag: $ResolvedLlamaTag" -ForegroundColor Gray
+substep "Resolved llama.cpp release tag: $ResolvedLlamaTag"
 
 if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
     Write-Host ""
-    Write-Host "[WARN] UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt llama.cpp install" -ForegroundColor Yellow
+    substep "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt llama.cpp install" "Yellow"
     $NeedLlamaSourceBuild = $true
 } else {
     Write-Host ""
-    Write-Host "Installing prebuilt llama.cpp bundle (preferred path)..." -ForegroundColor Cyan
+    substep "installing prebuilt llama.cpp bundle (preferred path)..."
     if (Test-Path $LlamaCppDir) {
-        Write-Host "Existing llama.cpp install detected -- validating staged prebuilt update before replacement" -ForegroundColor Gray
+        substep "Existing llama.cpp install detected -- validating staged prebuilt update before replacement"
     }
     if ($SkipPrebuiltInstall) {
-        Write-Host "[WARN] Skipping prebuilt install because prebuilt tag resolution failed -- falling back to source build" -ForegroundColor Yellow
+        substep "Skipping prebuilt install because prebuilt tag resolution failed -- falling back to source build" "Yellow"
     } else {
         $prebuiltArgs = @(
             "$PSScriptRoot\install_llama_prebuilt.py",
@@ -1530,17 +1654,28 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
         }
         $prevEAPPrebuilt = $ErrorActionPreference
         $ErrorActionPreference = "Continue"
-        & python @prebuiltArgs
-        $prebuiltExit = $LASTEXITCODE
+        if ($script:UnslothVerbose) {
+            # Show live output in verbose mode while still capturing for error log
+            $prebuiltLog = Join-Path $env:TEMP "unsloth-prebuilt-$PID.log"
+            & python @prebuiltArgs 2>&1 | Tee-Object -FilePath $prebuiltLog | Out-Host
+            $prebuiltExit = $LASTEXITCODE
+            $prebuiltOutput = if (Test-Path $prebuiltLog) { Get-Content $prebuiltLog -Raw } else { "" }
+            Remove-Item $prebuiltLog -ErrorAction SilentlyContinue
+        } else {
+            $prebuiltOutput = & python @prebuiltArgs 2>&1 | Out-String
+            $prebuiltExit = $LASTEXITCODE
+        }
         $ErrorActionPreference = $prevEAPPrebuilt
 
         if ($prebuiltExit -eq 0) {
             step "llama.cpp" "prebuilt installed and validated"
         } else {
+            step "llama.cpp" "prebuilt install failed (continuing)" "Yellow"
+            Write-LlamaFailureLog -Output $prebuiltOutput
             if (Test-Path $LlamaCppDir) {
-                Write-Host "[WARN] Prebuilt update failed; existing install was restored or cleaned before source build fallback" -ForegroundColor Yellow
+                substep "Prebuilt update failed; existing install was restored or cleaned before source build fallback" "Yellow"
             }
-            Write-Host "[WARN] Prebuilt llama.cpp path unavailable or failed validation -- falling back to source build" -ForegroundColor Yellow
+            substep "Prebuilt llama.cpp path unavailable or failed validation -- falling back to source build" "Yellow"
             $NeedLlamaSourceBuild = $true
         }
     }
@@ -1570,10 +1705,10 @@ if ($NeedLlamaSourceBuild) {
 
     if ($OpenSslRoot) {
         $OpenSslAvailable = $true
-        Write-Host "[OK] OpenSSL dev found at $OpenSslRoot" -ForegroundColor Green
+        substep "OpenSSL dev found at $OpenSslRoot"
     } else {
-        Write-Host "" 
-        Write-Host "Installing OpenSSL dev (for HTTPS in llama-server)..." -ForegroundColor Cyan
+        Write-Host ""
+        substep "installing OpenSSL dev (for HTTPS in llama-server)..."
         $HasWinget = $null -ne (Get-Command winget -ErrorAction SilentlyContinue)
         if ($HasWinget) {
             winget install -e --id ShiningLight.OpenSSL.Dev --accept-package-agreements --accept-source-agreements
@@ -1582,17 +1717,17 @@ if ($NeedLlamaSourceBuild) {
                 if (Test-Path (Join-Path $root 'include\openssl\ssl.h')) {
                     $OpenSslRoot = $root
                     $OpenSslAvailable = $true
-                    Write-Host "[OK] OpenSSL dev installed at $OpenSslRoot" -ForegroundColor Green
+                    substep "OpenSSL dev installed at $OpenSslRoot"
                     break
                 }
             }
         }
         if (-not $OpenSslAvailable) {
-            Write-Host "[WARN] OpenSSL dev not available -- llama-server will be built without HTTPS" -ForegroundColor Yellow
+            substep "OpenSSL dev not available -- llama-server will be built without HTTPS" "Yellow"
         }
     }
 } else {
-    Write-Host "[SKIP] OpenSSL dev install -- prebuilt llama.cpp already validated" -ForegroundColor Yellow
+    substep "OpenSSL dev install skipped -- prebuilt llama.cpp already validated" "Yellow"
 }
 
 # ==========================================================================
@@ -1638,21 +1773,22 @@ if (-not $NeedLlamaSourceBuild) {
     Write-Host ""
     if (-not $HasNvidiaSmi) {
         # CPU-only machines depend entirely on llama-server for GGUF chat -- cmake is required
-        Write-Host "[ERROR] CMake is required to build llama-server for GGUF chat mode." -ForegroundColor Red
-        Write-Host "        Install CMake from https://cmake.org/download/ and re-run setup." -ForegroundColor Yellow
-        exit 1
+        substep "CMake is required to build llama-server for GGUF chat mode." "Yellow"
+        substep "Continuing setup without llama.cpp build." "Yellow"
+        substep "Install CMake from https://cmake.org/download/ and re-run setup." "Yellow"
     }
-    Write-Host "[SKIP] llama-server build -- cmake not available" -ForegroundColor Yellow
-    Write-Host "       GGUF inference and export will not be available." -ForegroundColor Yellow
-    Write-Host "       Install CMake from https://cmake.org/download/ and re-run setup." -ForegroundColor Yellow
+    step "llama.cpp" "build skipped (cmake not available)" "Yellow"
+    substep "GGUF inference and export will not be available." "Yellow"
+    substep "Install CMake from https://cmake.org/download/ and re-run setup." "Yellow"
+    $script:LlamaCppDegraded = $true
 } else {
     Write-Host ""
     if ($HasNvidiaSmi) {
-        Write-Host "Building llama.cpp with CUDA support..." -ForegroundColor Cyan
+        substep "building llama.cpp with CUDA support..."
     } else {
-        Write-Host "Building llama.cpp (CPU-only, no NVIDIA GPU detected)..." -ForegroundColor Cyan
+        substep "building llama.cpp (CPU-only, no NVIDIA GPU detected)..."
     }
-    Write-Host "   This typically takes 5-10 minutes on first build." -ForegroundColor Gray
+    substep "This typically takes 5-10 minutes on first build."
     Write-Host ""
 
     # Start total build timer
@@ -1692,19 +1828,19 @@ if (-not $NeedLlamaSourceBuild) {
     if (Test-Path (Join-Path $LlamaCppDir ".git")) {
         Write-Host "   Syncing llama.cpp to $ResolvedLlamaTag..." -ForegroundColor Gray
         if ($UseConcreteRef) {
-            git -C $LlamaCppDir fetch --depth 1 origin $ResolvedLlamaTag 2>&1 | Out-Null
+            $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin $ResolvedLlamaTag }
         } else {
-            git -C $LlamaCppDir fetch --depth 1 origin 2>&1 | Out-Null
+            $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin }
         }
-        if ($LASTEXITCODE -ne 0) {
-            Write-Host "   [WARN] git fetch failed -- using existing source" -ForegroundColor Yellow
+        if ($gitFetchExit -ne 0) {
+            substep "git fetch failed -- using existing source" "Yellow"
         } else {
-            git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD 2>&1 | Out-Null
-            if ($LASTEXITCODE -ne 0) {
+            $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD }
+            if ($gitCheckoutExit -ne 0) {
                 $BuildOk = $false
                 $FailedStep = "git checkout"
             } else {
-                git -C $LlamaCppDir clean -fdx 2>&1 | Out-Null
+                Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null
             }
         }
     } else {
@@ -1716,8 +1852,8 @@ if (-not $NeedLlamaSourceBuild) {
             $cloneArgs += @("--branch", $ResolvedLlamaTag)
         }
         $cloneArgs += @("https://github.com/ggml-org/llama.cpp.git", $buildTmp)
-        git @cloneArgs 2>&1 | Out-Null
-        if ($LASTEXITCODE -ne 0) {
+        $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git @cloneArgs }
+        if ($cloneExit -ne 0) {
             $BuildOk = $false
             $FailedStep = "git clone"
             if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
@@ -1775,8 +1911,8 @@ if (-not $NeedLlamaSourceBuild) {
                     $maxArch = Get-NvccMaxArch -NvccExe $NvccPath
                     if ($maxArch) {
                         $CmakeArgs += "-DCMAKE_CUDA_ARCHITECTURES=$maxArch"
-                        Write-Host "   [WARN] GPU is sm_$CudaArch but nvcc only supports up to sm_$maxArch" -ForegroundColor Yellow
-                        Write-Host "          Building with sm_$maxArch (PTX will JIT for your GPU at runtime)" -ForegroundColor Yellow
+                        substep "GPU is sm_$CudaArch but nvcc only supports up to sm_$maxArch" "Yellow"
+                        substep "Building with sm_$maxArch (PTX will JIT for your GPU at runtime)" "Yellow"
                     }
                     # else: omit flag entirely, let cmake pick defaults
                 }
@@ -1786,10 +1922,11 @@ if (-not $NeedLlamaSourceBuild) {
         }
 
         $cmakeOutput = cmake @CmakeArgs 2>&1 | Out-String
-        if ($LASTEXITCODE -ne 0) {
+        $cmakeConfigureExit = $LASTEXITCODE
+        if ($cmakeConfigureExit -ne 0) {
             $BuildOk = $false
             $FailedStep = "cmake configure"
-            Write-Host $cmakeOutput -ForegroundColor Red
+            Write-LlamaFailureLog -Output $cmakeOutput
             if ($cmakeOutput -match 'No CUDA toolset found|CUDA_TOOLKIT_ROOT_DIR|nvcc') {
                 Write-Host ""
                 Write-Host "   Hint: CUDA VS integration may be missing. Try running as admin:" -ForegroundColor Yellow
@@ -1812,10 +1949,11 @@ if (-not $NeedLlamaSourceBuild) {
         Write-Host ""
 
         $output = cmake --build $BuildDir --config Release --target llama-server -j $NumCpu 2>&1 | Out-String
-        if ($LASTEXITCODE -ne 0) {
+        $cmakeBuildServerExit = $LASTEXITCODE
+        if ($cmakeBuildServerExit -ne 0) {
             $BuildOk = $false
             $FailedStep = "cmake build (llama-server)"
-            Write-Host $output -ForegroundColor Red
+            Write-LlamaFailureLog -Output $output
         }
     }
 
@@ -1824,9 +1962,10 @@ if (-not $NeedLlamaSourceBuild) {
         Write-Host ""
         Write-Host "--- cmake build (llama-quantize) ---" -ForegroundColor Cyan
         $output = cmake --build $BuildDir --config Release --target llama-quantize -j $NumCpu 2>&1 | Out-String
-        if ($LASTEXITCODE -ne 0) {
-            Write-Host "   [WARN] llama-quantize build failed (GGUF export may be unavailable)" -ForegroundColor Yellow
-            Write-Host $output -ForegroundColor Yellow
+        $cmakeBuildQuantizeExit = $LASTEXITCODE
+        if ($cmakeBuildQuantizeExit -ne 0) {
+            substep "llama-quantize build failed (GGUF export may be unavailable)" "Yellow"
+            Write-LlamaFailureLog -Output $output
         }
     }
 
@@ -1867,9 +2006,9 @@ if (-not $NeedLlamaSourceBuild) {
             step "llama.cpp" "built"
             step "build time" "${totalMin}m ${totalSec}s" "DarkGray"
         } else {
-            step "llama.cpp" "build failed at: $FailedStep (${totalMin}m ${totalSec}s)" "Red"
+            step "llama.cpp" "build failed at: $FailedStep (${totalMin}m ${totalSec}s); continuing" "Yellow"
             substep "To retry: delete $LlamaCppDir and re-run setup." "Yellow"
-            exit 1
+            $script:LlamaCppDegraded = $true
         }
     }
 }
@@ -1877,14 +2016,31 @@ if (-not $NeedLlamaSourceBuild) {
 # ─────────────────────────────────────────────
 # Footer
 # ─────────────────────────────────────────────
+$DoneLabel = if ($env:SKIP_STUDIO_BASE -eq "1") { "Unsloth Studio Setup Complete" } else { "Unsloth Studio Updated" }
 if ($script:StudioVtOk -and -not $env:NO_COLOR) {
     Write-Host ("  {0}{1}{2}" -f (Get-StudioAnsi Dim), $Rule, (Get-StudioAnsi Reset))
-    Write-Host ("  " + (Get-StudioAnsi Title) + "Unsloth Studio Installed" + (Get-StudioAnsi Reset))
+    if ($script:LlamaCppDegraded) {
+        Write-Host ("  " + (Get-StudioAnsi Warn) + "$DoneLabel (limited: llama.cpp unavailable)" + (Get-StudioAnsi Reset))
+    } else {
+        Write-Host ("  " + (Get-StudioAnsi Title) + $DoneLabel + (Get-StudioAnsi Reset))
+    }
     Write-Host ("  {0}{1}{2}" -f (Get-StudioAnsi Dim), $Rule, (Get-StudioAnsi Reset))
 } else {
     Write-Host "  $Rule" -ForegroundColor DarkGray
-    Write-Host "  Unsloth Studio Installed" -ForegroundColor Green
+    if ($script:LlamaCppDegraded) {
+        Write-Host "  $DoneLabel (limited: llama.cpp unavailable)" -ForegroundColor Yellow
+    } else {
+        Write-Host "  $DoneLabel" -ForegroundColor Green
+    }
     Write-Host "  $Rule" -ForegroundColor DarkGray
 }
 step "launch" "unsloth studio -H 0.0.0.0 -p 8888"
 Write-Host ""
+
+# Match studio/setup.sh: exit non-zero for degraded llama.cpp when called
+# from install.ps1 (SKIP_STUDIO_BASE=1) so the installer can detect the
+# failure. Direct 'unsloth studio update' does not set SKIP_STUDIO_BASE,
+# so it keeps degraded installs successful.
+if ($script:LlamaCppDegraded -and $env:SKIP_STUDIO_BASE -eq "1") {
+    exit 1
+}

studio/setup.sh

@@ -28,12 +28,43 @@ fi
 step()    { printf "  ${C_DIM}%-15.15s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; }
 substep() { printf "  ${C_DIM}%-15s%s${C_RST}\n" "" "$1"; }
 
+_is_verbose() {
+    [ "${UNSLOTH_VERBOSE:-0}" = "1" ]
+}
+
+verbose_substep() {
+    if _is_verbose; then
+        substep "$1"
+    fi
+    return 0
+}
+
+run_maybe_quiet() {
+    if _is_verbose; then
+        "$@"
+    else
+        "$@" > /dev/null 2>&1
+    fi
+}
+
 # ── Helper: run command quietly, show output only on failure ──
 _run_quiet() {
     local on_fail=$1
     local label=$2
     shift 2
 
+    if _is_verbose; then
+        local exit_code
+        "$@" && return 0
+        exit_code=$?
+        step "error" "$label failed (exit code $exit_code)" "$C_ERR" >&2
+        if [ "$on_fail" = "exit" ]; then
+            exit "$exit_code"
+        else
+            return "$exit_code"
+        fi
+    fi
+
     local tmplog
     tmplog=$(mktemp) || {
         step "error" "Failed to create temporary file" "$C_ERR" >&2
@@ -65,11 +96,18 @@ run_quiet_no_exit() {
     _run_quiet return "$@"
 }
 
+print_llama_error_log() {
+    local log_file=$1
+    [ -s "$log_file" ] || return 0
+    substep "llama.cpp diagnostics (last 120 lines):"
+    tail -n 120 "$log_file" | sed 's/^/   | /' >&2
+}
+
 # ── Banner ──
 echo ""
 printf "  ${C_TITLE}%s${C_RST}\n" "🦥 Unsloth Studio Setup"
 printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
-
+verbose_substep "verbose diagnostics enabled"
 # ── Clean up stale caches ──
 rm -rf "$REPO_ROOT/unsloth_compiled_cache"
 rm -rf "$SCRIPT_DIR/backend/unsloth_compiled_cache"
@@ -97,6 +135,7 @@ fi
 
 if [ "$_NEED_FRONTEND_BUILD" = false ]; then
     step "frontend" "up to date"
+    verbose_substep "frontend dist is newer than source inputs"
 else
 
 # ── Node ──
@@ -117,7 +156,7 @@ if command -v node &>/dev/null && command -v npm &>/dev/null; then
             # In Colab, just upgrade npm directly - nvm doesn't work well
             if [ "$NPM_MAJOR" -lt 11 ]; then
                 substep "upgrading npm..."
-                npm install -g npm@latest > /dev/null 2>&1
+                run_maybe_quiet npm install -g npm@latest
             fi
             NEED_NODE=false
         fi
@@ -127,7 +166,11 @@ fi
 if [ "$NEED_NODE" = true ]; then
     substep "installing nvm..."
     export NODE_OPTIONS=--dns-result-order=ipv4first
-    curl -so- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash > /dev/null 2>&1
+    if _is_verbose; then
+        curl -so- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash
+    else
+        curl -so- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash > /dev/null 2>&1
+    fi
 
     export NVM_DIR="$HOME/.nvm"
     set +u
@@ -141,7 +184,11 @@ if [ "$NEED_NODE" = true ]; then
 
     substep "installing Node LTS..."
     run_quiet "nvm install" nvm install --lts
-    nvm use --lts > /dev/null 2>&1
+    if _is_verbose; then
+        nvm use --lts
+    else
+        nvm use --lts > /dev/null 2>&1
+    fi
     set -u
 
     NODE_MAJOR=$(node -v | sed 's/v//' | cut -d. -f1)
@@ -158,13 +205,14 @@ if [ "$NEED_NODE" = true ]; then
 fi
 
 step "node" "$(node -v) | npm $(npm -v)"
+verbose_substep "node check: NEED_NODE=$NEED_NODE NODE_OK=${NODE_OK:-unknown} NPM_MAJOR=${NPM_MAJOR:-unknown}"
 
 # ── Install bun (optional, faster package installs) ──
 # Uses npm to install bun globally -- Node is already guaranteed above,
 # avoids platform-specific installers, PATH issues, and admin requirements.
 if ! command -v bun &>/dev/null; then
     substep "installing bun..."
-    if npm install -g bun > /dev/null 2>&1 && command -v bun &>/dev/null; then
+    if run_maybe_quiet npm install -g bun && command -v bun &>/dev/null; then
         substep "bun installed ($(bun --version))"
     else
         substep "bun install skipped (npm will be used instead)"
@@ -209,7 +257,10 @@ _try_bun_install() {
     _log=$(mktemp)
     bun install >"$_log" 2>&1 || _exit_code=$?
 
-    if [ "$_exit_code" -eq 0 ] && [ -x node_modules/.bin/tsc ] && [ -x node_modules/.bin/vite ]; then
+    # bun may create .exe shims on Windows (Git Bash / MSYS2) instead of plain scripts
+    if [ "$_exit_code" -eq 0 ] \
+        && { [ -x node_modules/.bin/tsc ] || [ -f node_modules/.bin/tsc.exe ] || [ -f node_modules/.bin/tsc.bunx ]; } \
+        && { [ -x node_modules/.bin/vite ] || [ -f node_modules/.bin/vite.exe ] || [ -f node_modules/.bin/vite.bunx ]; }; then
         rm -f "$_log"
         return 0
     fi
@@ -228,21 +279,25 @@ _try_bun_install() {
 
 _bun_install_ok=false
 if command -v bun &>/dev/null; then
-    echo "   Using bun for package install (faster)"
+    substep "using bun for package install (faster)"
     if _try_bun_install; then
         _bun_install_ok=true
     else
         # First attempt failed, likely due to corrupt cache entries.
         # Clear the cache and retry once.
         echo "   Clearing bun cache and retrying..."
-        bun pm cache rm > /dev/null 2>&1 || true
+        run_maybe_quiet bun pm cache rm || true
         if _try_bun_install; then
             _bun_install_ok=true
         fi
     fi
 fi
 if [ "$_bun_install_ok" = false ]; then
-    run_quiet "npm install" npm install
+    run_quiet_no_exit "npm install" npm install --no-fund --no-audit --loglevel=error
+    _npm_install_rc=$?
+    if [ "$_npm_install_rc" -ne 0 ]; then
+        exit "$_npm_install_rc"
+    fi
 fi
 run_quiet "npm run build" npm run build
 
@@ -265,7 +320,11 @@ fi  # end frontend build check
 # ── oxc-validator runtime ──
 if [ -d "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator" ] && command -v npm &>/dev/null; then
     cd "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator"
-    run_quiet "npm install (oxc validator runtime)" npm install
+    run_quiet_no_exit "npm install (oxc validator runtime)" npm install --no-fund --no-audit --loglevel=error
+    _oxc_install_rc=$?
+    if [ "$_oxc_install_rc" -ne 0 ]; then
+        exit "$_oxc_install_rc"
+    fi
     cd "$SCRIPT_DIR"
 fi
 
@@ -287,9 +346,19 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then
         # packages (huggingface-hub, datasets, transformers) and only pulls
         # in genuinely missing ones (structlog, fastapi, etc.).
         substep "Colab detected, installing Studio backend dependencies..."
+        _COLAB_REQS_TMP="$(mktemp)"
         sed 's/[><=!~;].*//' "$SCRIPT_DIR/backend/requirements/studio.txt" \
-            | grep -v '^#' | grep -v '^$' \
-            | pip install -q -r /dev/stdin 2>/dev/null || true
+            | grep -v '^#' | grep -v '^$' > "$_COLAB_REQS_TMP"
+        if [ -s "$_COLAB_REQS_TMP" ]; then
+            if ! run_quiet_no_exit "install Colab backend deps" pip install -q -r "$_COLAB_REQS_TMP"; then
+                rm -f "$_COLAB_REQS_TMP"
+                step "python" "Colab backend dependency install failed" "$C_ERR"
+                exit 1
+            fi
+        else
+            step "python" "no Colab backend dependencies resolved from requirements file" "$C_WARN"
+        fi
+        rm -f "$_COLAB_REQS_TMP"
         _COLAB_NO_VENV=true
     else
         step "python" "venv not found at $VENV_DIR" "$C_ERR"
@@ -308,7 +377,13 @@ install_python_stack() {
 USE_UV=false
 if command -v uv &>/dev/null; then
     USE_UV=true
-elif curl -LsSf https://astral.sh/uv/install.sh | sh > /dev/null 2>&1; then
+elif {
+    if _is_verbose; then
+        curl -LsSf https://astral.sh/uv/install.sh | sh
+    else
+        curl -LsSf https://astral.sh/uv/install.sh | sh > /dev/null 2>&1
+    fi
+}; then
     export PATH="$HOME/.local/bin:$PATH"
     command -v uv &>/dev/null && USE_UV=true
 fi
@@ -325,7 +400,8 @@ cd "$SCRIPT_DIR"
 # On Colab without a venv, skip venv-dependent Python deps sections but
 # continue to llama.cpp install so GGUF inference is available.
 if [ "$_COLAB_NO_VENV" = true ]; then
-    echo "✅ Studio backend dependencies installed into system Python"
+    step "python" "backend deps installed into system Python"
+    substep "continuing to llama.cpp install for GGUF inference support"
 fi
 
 # ── Check if Python deps need updating ──
@@ -375,6 +451,7 @@ if [ "$_SKIP_PYTHON_DEPS" = false ]; then
     step "transformers" "5.x pre-installed"
 else
     step "python" "dependencies up to date"
+    verbose_substep "python deps check: installed=$_PKG_NAME@${INSTALLED_VER:-unknown} latest=${LATEST_VER:-unknown}"
 fi
 
 # ── 7. Prefer prebuilt llama.cpp bundles before any source build path ──
@@ -383,6 +460,7 @@ mkdir -p "$UNSLOTH_HOME"
 LLAMA_CPP_DIR="$UNSLOTH_HOME/llama.cpp"
 LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server"
 _NEED_LLAMA_SOURCE_BUILD=false
+_LLAMA_CPP_DEGRADED=false
 _LLAMA_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}"
 _REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-latest}"
 _HELPER_RELEASE_REPO="${UNSLOTH_LLAMA_RELEASE_REPO:-unslothai/llama.cpp}"
@@ -400,7 +478,7 @@ else
 fi
 if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
     step "llama.cpp" "failed to resolve prebuilt tag via $_HELPER_RELEASE_REPO" "$C_WARN"
-    cat "$_RESOLVE_LLAMA_LOG" >&2 || true
+    print_llama_error_log "$_RESOLVE_LLAMA_LOG"
     set +e
     # Resolve the llama.cpp tag for source-build fallback. Pass --published-repo
     # so the resolver prefers Unsloth's tested tag (e.g. b8508) over the upstream
@@ -426,6 +504,7 @@ fi
 rm -f "$_RESOLVE_LLAMA_LOG"
 
 substep "resolved llama.cpp tag: $_RESOLVED_LLAMA_TAG"
+verbose_substep "requested llama.cpp tag: $_REQUESTED_LLAMA_TAG (repo: $_HELPER_RELEASE_REPO)"
 
 if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then
     step "llama.cpp" "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt" "$C_WARN"
@@ -447,14 +526,25 @@ else
         if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
             _PREBUILT_CMD+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
         fi
+        _PREBUILT_LOG="$(mktemp)"
         set +e
-        "${_PREBUILT_CMD[@]}"
-        _PREBUILT_STATUS=$?
+        if _is_verbose; then
+            "${_PREBUILT_CMD[@]}" 2>&1 | tee "$_PREBUILT_LOG"
+            _PREBUILT_STATUS=${PIPESTATUS[0]}
+        else
+            "${_PREBUILT_CMD[@]}" >"$_PREBUILT_LOG" 2>&1
+            _PREBUILT_STATUS=$?
+        fi
         set -e
 
         if [ "$_PREBUILT_STATUS" -eq 0 ]; then
             step "llama.cpp" "prebuilt installed and validated"
+            verbose_substep "llama.cpp install dir: $LLAMA_CPP_DIR"
+            rm -f "$_PREBUILT_LOG"
         else
+            step "llama.cpp" "prebuilt install failed (continuing)" "$C_WARN"
+            print_llama_error_log "$_PREBUILT_LOG"
+            rm -f "$_PREBUILT_LOG"
             if [ -d "$LLAMA_CPP_DIR" ]; then
                 substep "prebuilt update failed; existing install restored"
             fi
@@ -523,12 +613,15 @@ if [ "$_NEED_LLAMA_SOURCE_BUILD" = false ]; then
     :
 elif [ "${_SKIP_GGUF_BUILD:-}" = true ]; then
     step "llama.cpp" "skipped (missing build deps)" "$C_WARN"
+    [ -f "$LLAMA_SERVER_BIN" ] || _LLAMA_CPP_DEGRADED=true
 else
 {
     if ! command -v cmake &>/dev/null; then
         step "llama.cpp" "skipped (cmake not found)" "$C_WARN"
+        [ -f "$LLAMA_SERVER_BIN" ] || _LLAMA_CPP_DEGRADED=true
     elif ! command -v git &>/dev/null; then
         step "llama.cpp" "skipped (git not found)" "$C_WARN"
+        [ -f "$LLAMA_SERVER_BIN" ] || _LLAMA_CPP_DEGRADED=true
     else
         BUILD_OK=true
         _CLONE_BRANCH_ARGS=()
@@ -691,8 +784,10 @@ else
             [ -f "$LLAMA_CPP_DIR/llama-quantize" ] && step "llama-quantize" "built"
         elif [ "$BUILD_OK" = true ]; then
             step "llama.cpp" "binary not found after build" "$C_WARN"
+            _LLAMA_CPP_DEGRADED=true
         else
             step "llama.cpp" "build failed" "$C_ERR"
+            [ -f "$LLAMA_SERVER_BIN" ] || _LLAMA_CPP_DEGRADED=true
         fi
     fi
 }
@@ -702,14 +797,35 @@ fi  # end _SKIP_GGUF_BUILD check
 if [ "$IS_COLAB" = true ]; then
     echo ""
     printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
-    printf "  ${C_TITLE}%s${C_RST}\n" "Unsloth Studio Setup Complete"
+    if [ "$_LLAMA_CPP_DEGRADED" = true ]; then
+        printf "  ${C_WARN}%s${C_RST}\n" "Unsloth Studio Setup Complete (limited: llama.cpp unavailable)"
+    else
+        printf "  ${C_TITLE}%s${C_RST}\n" "Unsloth Studio Setup Complete"
+    fi
     printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
     substep "from colab import start"
     substep "start()"
 else
     printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
-    printf "  ${C_TITLE}%s${C_RST}\n" "Unsloth Studio Installed"
+    if [ "$_LLAMA_CPP_DEGRADED" = true ]; then
+        printf "  ${C_WARN}%s${C_RST}\n" "Unsloth Studio Installed (limited: llama.cpp unavailable)"
+    else
+        printf "  ${C_TITLE}%s${C_RST}\n" "Unsloth Studio Installed"
+    fi
     printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
-    printf "  ${C_DIM}%-15s${C_OK}%s${C_RST}\n" "launch" "unsloth studio -H 0.0.0.0 -p 8888"
+    if [ "$_LLAMA_CPP_DEGRADED" = true ]; then
+        printf "  ${C_DIM}%-15s${C_WARN}%s${C_RST}\n" "launch" "unsloth studio -H 0.0.0.0 -p 8888"
+    else
+        printf "  ${C_DIM}%-15s${C_OK}%s${C_RST}\n" "launch" "unsloth studio -H 0.0.0.0 -p 8888"
+    fi
 fi
 echo ""
+
+# When called from install.sh (SKIP_STUDIO_BASE=1), exit non-zero so the
+# installer can report the GGUF failure after finishing PATH/shortcut setup.
+# When called directly via 'unsloth studio update', keep the install
+# successful -- the footer above already reports the limitation and Studio
+# is still usable for non-GGUF workflows.
+if [ "$_LLAMA_CPP_DEGRADED" = true ] && [ "${SKIP_STUDIO_BASE:-0}" = "1" ]; then
+    exit 1
+fi

unsloth/models/_utils.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "2026.3.17"
+__version__ = "2026.3.18"
 
 __all__ = [
     "SUPPORTS_BFLOAT16",

unsloth_cli/commands/studio.py

@@ -267,10 +267,7 @@ def setup(
         help = "Full pip/build output during setup for troubleshooting.",
     ),
 ):
-    """Deprecated: use 'unsloth studio update' or re-run install.sh."""
-    typer.echo(
-        "Note: 'unsloth studio setup' is deprecated. Use 'unsloth studio update' or re-run install.sh."
-    )
+    """Run Studio setup (called by install.ps1 / install.sh)."""
     _run_setup_script(verbose = verbose)
 
 
@@ -290,13 +287,18 @@ def update(
     ),
 ):
     """Update Unsloth Studio dependencies and rebuild."""
-    os.environ["STUDIO_LOCAL_INSTALL"] = "1" if local else "0"
+    # Ensure SKIP_STUDIO_BASE is not inherited from a parent install.ps1 session
+    os.environ.pop("SKIP_STUDIO_BASE", None)
     os.environ["STUDIO_PACKAGE_NAME"] = package
     if local:
+        os.environ["STUDIO_LOCAL_INSTALL"] = "1"
         # Pass the repo root explicitly so install_python_stack.py doesn't
         # have to guess from SCRIPT_DIR (which may be inside site-packages).
         repo_root = Path(__file__).resolve().parents[2]
         os.environ["STUDIO_LOCAL_REPO"] = str(repo_root)
+    else:
+        os.environ["STUDIO_LOCAL_INSTALL"] = "0"
+        os.environ.pop("STUDIO_LOCAL_REPO", None)
     _run_setup_script(verbose = verbose)