diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index e11a08f8ab..20000c82ee 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1,6 +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 { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; import { @@ -102,6 +103,7 @@ import { type FormatFilter, estimateQuantBytes, fitsDevice, + hfModelFitsDevice, isMlxId, isMobileVariant, isRecommendableFormat, @@ -1340,6 +1342,9 @@ export function HubModelPicker({ }, []); // When on, On Device GGUF repos show their quantizations without a click. const expandQuantizations = useChatRuntimeStore((s) => s.expandQuantizations); + // Shared with the Hub page: list only models sized within the device budget. + const fitOnDeviceOnly = useChatRuntimeStore((s) => s.fitOnDeviceOnly); + const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly); // Repos the user clicked to collapse while expand-by-default is on. Kept in // memory only, so it resets on reload (and when the setting is toggled). const [collapsedGguf, setCollapsedGguf] = useState>( @@ -1717,34 +1722,19 @@ export function HubModelPicker({ formatFilter === "all" ? rows.filter((r) => isRecommendableFormat(r.id, r.isGguf, isMac)) : rows.filter((r) => matchesFormatFilter(r.id, r.isGguf, formatFilter)); - if (recommendedSort !== "recommended") return rows; + // The "recommended" sort always applies the device-fit filter; the shared + // "Fits on device" tick extends it to the other sorts too. + if (recommendedSort !== "recommended" && !fitOnDeviceOnly) return rows; return rows.filter((r) => { // Downloaded models always show, regardless of device fit. if (downloadedSet.has(r.id.toLowerCase())) return true; - // Unified-memory hosts (Mac / no discrete GPU) still report system RAM, - // so fall back to that budget instead of skipping the fit check entirely. - const hasDeviceBudget = - gpu.memoryTotalGb > 0 || gpu.systemRamAvailableGb > 0; - if (!hasDeviceBudget) return true; - // GGUF/MLX repos rarely expose safetensors metadata, so fall back to the - // GGUF param count, then the repo name, for a size estimate. Anything we - // still cannot size is hidden (requireKnown) so over-budget models like a - // 1T GGUF don't slip into Recommended. - const params = r.totalParams ?? paramsFromId(r.id); - const sizeBytes = - r.estimatedSizeBytes ?? - (params ? estimateQuantBytes(params) : undefined); - return fitsDevice({ - sizeBytes, - gpuGb: gpu.memoryTotalGb, - systemRamGb: gpu.systemRamAvailableGb, - requireKnown: true, - }); + return hfModelFitsDevice(r, gpu); }); }, [ recommendedSearch.results, downloadedSet, recommendedSort, + fitOnDeviceOnly, formatFilter, isMac, gpu, @@ -1976,23 +1966,6 @@ export function HubModelPicker({ [visibleCachedModelRows], ); - // Recommended models that match the current search query - const filteredRecommendedIds = useMemo(() => { - if (!showHfSection) return []; - const q = normalizeForSearch(debouncedQuery.trim()); - return recommendedIds - .filter((id) => normalizeForSearch(id).includes(q)) - .filter((id) => - matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter), - ); - }, [ - showHfSection, - debouncedQuery, - recommendedIds, - formatFilter, - isKnownGgufRepo, - ]); - // Param counts come straight off the unsloth listings the picker already // loaded, so no extra per-id fetch is needed for the VRAM badges. const recommendedParamCountById = useMemo(() => { @@ -2003,6 +1976,42 @@ export function HubModelPicker({ return map; }, [results, recommendedSearch.results]); + // Recommended models that match the current search query + const filteredRecommendedIds = useMemo(() => { + if (!showHfSection) return []; + const q = normalizeForSearch(debouncedQuery.trim()); + return recommendedIds + .filter((id) => normalizeForSearch(id).includes(q)) + .filter((id) => + matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter), + ) + // Curated defaults obey the fit toggle like the live HF rows, else large + // defaults resurface in search results with the filter on. + .filter( + (id) => + !fitOnDeviceOnly || + downloadedSet.has(id.toLowerCase()) || + hfModelFitsDevice( + { + id, + totalParams: recommendedParamCountById.get(id), + isGguf: isKnownGgufRepo(id), + }, + gpu, + ), + ); + }, [ + showHfSection, + debouncedQuery, + recommendedIds, + formatFilter, + isKnownGgufRepo, + fitOnDeviceOnly, + downloadedSet, + recommendedParamCountById, + gpu, + ]); + const recommendedSet = useMemo( () => new Set(filteredRecommendedIds), [filteredRecommendedIds], @@ -2013,6 +2022,12 @@ export function HubModelPicker({ if (!showHfSection || section !== "recommended") return []; return results .filter(isChatSupported) + .filter( + (r) => + !fitOnDeviceOnly || + downloadedSet.has(r.id.toLowerCase()) || + hfModelFitsDevice(r, gpu), + ) .map((result) => result.id) .filter((id) => !isHiddenModelId(id)) .filter((id) => id.toLowerCase().startsWith("unsloth/")) @@ -2035,6 +2050,9 @@ export function HubModelPicker({ isKnownGgufRepo, isChatSupported, formatFilter, + fitOnDeviceOnly, + downloadedSet, + gpu, isMac, ]); @@ -2323,6 +2341,35 @@ export function HubModelPicker({ // selected-item checkmark never overlaps the label. const sortMenuContentClassName = "!p-1 !rounded-[14px] [&_[role=option]]:!pl-2 [&_[role=option]]:!py-1.5 [&_[role=option]]:!text-xs [&_[role=option]]:!rounded-[10px]"; + // Device-fit toggle lives inside the sort menu (shared with the Hub page). + // The whole row is the click target (a button): a Checkbox renders as a + // + + + Hides models larger than this device's memory budget. Downloaded models + stay visible. + + + ); const sectionSortDropdown = section === "recommended" ? ( ) : section === "downloaded" ? ( ) : ( ); diff --git a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts index 24f0edc784..7c2ed266c0 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts @@ -114,3 +114,35 @@ export function fitsDevice(opts: { } return requireKnown ? false : true; } + +/** Fit predicate for one Hub listing row, shared by the chat model selector + * and the Hub page "Fits on device" filter. GGUF repos: metadata size (actual + * weights) or the smallest-quant estimate from the param count. Safetensors / + * MLX repos: always the params-based smallest-quant estimate, matching the + * VRAM badge's quantized-load assumption; their estimatedSizeBytes is the + * full-precision checkpoint and would wrongly hide models the quantized load + * path can run. Anything unsizable is hidden (requireKnown) so over-budget + * models with no metadata don't slip through. An unknown device budget keeps + * everything. */ +export function hfModelFitsDevice( + model: { + id: string; + totalParams?: number; + estimatedSizeBytes?: number; + isGguf?: boolean; + }, + gpu: { memoryTotalGb: number; systemRamAvailableGb: number }, +): boolean { + if (gpu.memoryTotalGb <= 0 && gpu.systemRamAvailableGb <= 0) return true; + const params = model.totalParams ?? paramsFromId(model.id); + const quantBytes = params ? estimateQuantBytes(params) : undefined; + const sizeBytes = isGgufId(model.id, model.isGguf) + ? (model.estimatedSizeBytes ?? quantBytes) + : (quantBytes ?? model.estimatedSizeBytes); + return fitsDevice({ + sizeBytes, + gpuGb: gpu.memoryTotalGb, + systemRamGb: gpu.systemRamAvailableGb, + requireKnown: true, + }); +} diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index ca4bb7afde..7c9685d6d0 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -42,6 +42,8 @@ export const CHAT_EXPAND_QUANTIZATIONS_KEY = "unsloth_chat_expand_quantizations"; export const CHAT_SHOW_ALL_QUANTIZATIONS_KEY = "unsloth_chat_show_all_quantizations"; +export const MODELS_FIT_ON_DEVICE_ONLY_KEY = + "unsloth_models_fit_on_device_only"; export const CHAT_BYPASS_PERMISSIONS_KEY = "unsloth_chat_bypass_permissions"; export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY = "unsloth_chat_web_fetch_tools_enabled"; @@ -671,6 +673,9 @@ type ChatRuntimeStore = { expandQuantizations: boolean; /** Persisted: show non-downloaded quantizations too, not just downloaded. */ showAllQuantizations: boolean; + /** Persisted, shared by the chat model selector and the Hub page: list only + * models whose size fits this device's memory budget. */ + fitOnDeviceOnly: boolean; /** A local model picked while `loadOnSelection` is off: staged, not loaded. * The settings sheet shows its load knobs and a Load button. */ pendingSelection: PendingModelSelection | null; @@ -793,6 +798,7 @@ type ChatRuntimeStore = { setLoadOnSelection: (value: boolean) => void; setExpandQuantizations: (value: boolean) => void; setShowAllQuantizations: (value: boolean) => void; + setFitOnDeviceOnly: (value: boolean) => void; setPendingSelection: (selection: PendingModelSelection | null) => void; /** Stage a pick for a deferred load: revert knobs to the loaded baseline, * record the selection, and open the settings sheet. */ @@ -1111,6 +1117,7 @@ export const useChatRuntimeStore = create((set, get) => ({ loadOnSelection: loadBool(CHAT_LOAD_ON_SELECTION_KEY, true), expandQuantizations: loadBool(CHAT_EXPAND_QUANTIZATIONS_KEY, false), showAllQuantizations: loadBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, true), + fitOnDeviceOnly: loadBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, false), pendingSelection: null, loadedIsMultimodal: false, loadedIsDiffusion: false, @@ -1582,6 +1589,10 @@ export const useChatRuntimeStore = create((set, get) => ({ saveBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, showAllQuantizations); set({ showAllQuantizations }); }, + setFitOnDeviceOnly: (fitOnDeviceOnly) => { + saveBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, fitOnDeviceOnly); + set({ fitOnDeviceOnly }); + }, setPendingSelection: (pendingSelection) => set({ pendingSelection }), stageModel: (selection) => { // Refuse staging mid-load: post-load cleanup would silently drop the queued diff --git a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx index 7895a89254..38464d36e4 100644 --- a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx +++ b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx @@ -38,6 +38,7 @@ export function HubOptionMenu({ showChevron = true, title, triggerContent, + footer, }: { value: T; options: readonly HubOption[]; @@ -49,9 +50,12 @@ export function HubOptionMenu({ showChevron?: boolean; title?: string; triggerContent?: ReactNode; + /** Rendered under the options behind a separator; clicks keep the menu open. */ + footer?: ReactNode; }) { const [open, setOpen] = useState(false); - const [activeIndex, setActiveIndex] = useState(0); + // -1 = nothing highlighted (no hover, no keyboard nav yet). + const [activeIndex, setActiveIndex] = useState(-1); const triggerRef = useRef(null); const listboxRef = useRef(null); const idBase = useId(); @@ -63,9 +67,9 @@ export function HubOptionMenu({ }, [options, value]); const selected = options[selectedIndex]; const resolvedActiveIndex = - options.length === 0 + options.length === 0 || activeIndex < 0 ? -1 - : Math.min(Math.max(activeIndex, 0), options.length - 1); + : Math.min(activeIndex, options.length - 1); const activeOptionId = resolvedActiveIndex >= 0 ? `${idBase}-option-${resolvedActiveIndex}` : undefined; @@ -92,11 +96,13 @@ export function HubOptionMenu({ (nextOpen: boolean) => { setOpen(nextOpen); if (nextOpen) { - activateIndex(selectedIndex); + // Nothing highlighted until the user hovers or uses the keyboard; + // keyboard nav anchors on the selected option (handleContentKeyDown). + activateIndex(-1); requestAnimationFrame(() => listboxRef.current?.focus()); } }, - [activateIndex, selectedIndex], + [activateIndex], ); const handleContentKeyDown = useCallback( @@ -112,12 +118,21 @@ export function HubOptionMenu({ } if (event.key === "ArrowDown") { event.preventDefault(); - setActiveIndex((currentIndex + 1) % options.length); + // First arrow press highlights the selected option, then steps. + setActiveIndex( + resolvedActiveIndex < 0 + ? selectedIndex + : (currentIndex + 1) % options.length, + ); return; } if (event.key === "ArrowUp") { event.preventDefault(); - setActiveIndex((currentIndex - 1 + options.length) % options.length); + setActiveIndex( + resolvedActiveIndex < 0 + ? selectedIndex + : (currentIndex - 1 + options.length) % options.length, + ); return; } if (event.key === "Home") { @@ -197,6 +212,7 @@ export function HubOptionMenu({ aria-activedescendant={activeOptionId} tabIndex={0} onKeyDown={handleContentKeyDown} + onPointerLeave={() => activateIndex(-1)} className="outline-none" > {options.map((option, index) => { @@ -235,6 +251,12 @@ export function HubOptionMenu({ ); })} + {footer && ( + // -mt-3 cancels the surface's 16px flex gap down to 4px. No side + // padding: the footer label carries the same padding as the options + // so its checkbox lines up with the option text. +
{footer}
+ )} ); diff --git a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx index 48f7fcffaa..7c08c9c486 100644 --- a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx +++ b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx @@ -1,6 +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 { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; import { @@ -68,6 +69,8 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onFormatFilterChange, capabilityFilter, onCapabilityFilterChange, + fitOnDeviceOnly, + onFitOnDeviceOnlyChange, onManageLocalFolders, onOpenFineTune, }: { @@ -84,6 +87,9 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onFormatFilterChange: (value: ModelFormatFilter) => void; capabilityFilter: CapabilityFilter; onCapabilityFilterChange: (value: CapabilityFilter) => void; + /** Shared with the chat model selector: hide models over the device budget. */ + fitOnDeviceOnly: boolean; + onFitOnDeviceOnlyChange: (value: boolean) => void; onManageLocalFolders: () => void; /** Opens the curated "Fine-tune ready" channel (discover only). Exposed as a * format-dropdown option rather than a standalone feed section. */ @@ -350,6 +356,33 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onValueChange={onSortChange} ariaLabel="Sort models" className={cn(triggerBase, "w-[128px]")} + footer={ + isDataset ? undefined : ( + + + + + + Hides models larger than this device's memory budget. + Downloaded models stay visible. + + + ) + } /> )} diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx index c696862283..56aa07335d 100644 --- a/studio/frontend/src/features/hub/hub-page.tsx +++ b/studio/frontend/src/features/hub/hub-page.tsx @@ -5,6 +5,7 @@ import { loadRememberedLoadSettings, rememberedLoadSettingsKey, } from "@/components/assistant-ui/model-selector/remembered-load-settings"; +import { hfModelFitsDevice } from "@/components/assistant-ui/model-selector/recommended-fit"; import { useHubInventory } from "@/features/hub/inventory"; import { useDebouncedValue } from "@/hooks/use-debounced-value"; import { useGpuInfo } from "@/hooks/use-gpu-info"; @@ -327,6 +328,9 @@ export function ModelsPage() { const activeCheckpoint = checkpoint && !isExternalModelId(checkpoint) ? checkpoint : null; const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); + // Shared with the chat model selector: list only models sized for this device. + const fitOnDeviceOnly = useChatRuntimeStore((s) => s.fitOnDeviceOnly); + const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly); useEffect(() => { let cancelled = false; @@ -697,7 +701,12 @@ export function ModelsPage() { !isHiddenModelId(row.id) && matchesFormat(detectResultFormat(row.result), effectiveDiscoverFormat) && matchesCapability(row.capabilities, deferredCapabilityFilter) && - (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)), + (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)) && + // Models already on disk stay visible regardless of device fit, + // matching the chat model selector. + (!fitOnDeviceOnly || + row.isAvailableOnDevice || + hfModelFitsDevice(row.result, gpu)), ); }, [ discoverRows, @@ -705,6 +714,8 @@ export function ModelsPage() { effectiveDiscoverFormat, deferredCapabilityFilter, activeChannel, + fitOnDeviceOnly, + gpu, ]); const listRows = filteredDiscoverRows; @@ -724,8 +735,21 @@ export function ModelsPage() { effectiveLocalRows, ) .filter((row) => !isHiddenModelId(row.id)) - .filter((row) => matchesFormat(row.result.isGguf, "gguf")), - [hubFeed.trending.results, modelDiscoveryInventorySignature], + .filter((row) => matchesFormat(row.result.isGguf, "gguf")) + // Same fit filter as the main Discover list, so the feed carousel + // honors the toggle too. + .filter( + (row) => + !fitOnDeviceOnly || + row.isAvailableOnDevice || + hfModelFitsDevice(row.result, gpu), + ), + [ + hubFeed.trending.results, + modelDiscoveryInventorySignature, + fitOnDeviceOnly, + gpu, + ], ); const feedRows = useMemo(() => { if (!isFeedMode) return []; @@ -1448,6 +1472,8 @@ export function ModelsPage() { onFormatFilterChange={setFormatFilter} capabilityFilter={capabilityFilter} onCapabilityFilterChange={setCapabilityFilter} + fitOnDeviceOnly={fitOnDeviceOnly} + onFitOnDeviceOnlyChange={setFitOnDeviceOnly} onManageLocalFolders={handleManageLocalFolders} onOpenFineTune={() => handleOpenList("finetune")} /> diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 247d5040fb..4e02f7f14f 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -81,6 +81,7 @@ const PREFS_KEYS: string[] = [ "unsloth_chat_load_on_selection", "unsloth_chat_expand_quantizations", "unsloth_chat_show_all_quantizations", + "unsloth_models_fit_on_device_only", // Chat presets "unsloth_chat_custom_presets", "unsloth_chat_active_preset",