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 b90cdd0e3b..e023e36e2b 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1,7 +1,14 @@ import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; -import { useDebouncedValue, useHfModelSearch, useInfiniteScroll } from "@/hooks"; +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from "@/components/ui/tooltip"; +import { useDebouncedValue, useGpuInfo, useHfModelSearch, useInfiniteScroll } from "@/hooks"; import { cn, formatCompact } from "@/lib/utils"; +import type { VramFitStatus } from "@/lib/vram"; +import { checkVramFit, estimateLoadingVram } from "@/lib/vram"; import { Search01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useMemo, useState, type ReactNode } from "react"; @@ -28,27 +35,77 @@ function ModelRow({ meta, selected, onClick, + vramStatus, + vramEst, + gpuGb, }: { label: string; meta?: string; selected?: boolean; onClick: () => void; + vramStatus?: VramFitStatus | null; + vramEst?: number; + gpuGb?: number; }) { - return ( + const exceeds = vramStatus === "exceeds"; + const showVramTooltip = + vramEst != null && vramEst > 0 && gpuGb != null && gpuGb > 0; + const vramTooltipText = + showVramTooltip && vramStatus + ? exceeds + ? `Needs ~${vramEst}GB VRAM (GPU: ${gpuGb}GB)` + : vramStatus === "tight" + ? `~${vramEst}GB VRAM (tight fit on ${gpuGb}GB)` + : `~${vramEst}GB VRAM` + : null; + + const content = ( ); + + if (vramTooltipText) { + return ( + + {content} + + {label} + {vramTooltipText} + + + ); + } + return content; } export function HubModelPicker({ @@ -60,6 +117,7 @@ export function HubModelPicker({ value?: string; onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; }) { + const gpu = useGpuInfo(); const [query, setQuery] = useState(""); const debouncedQuery = useDebouncedValue(query); const { results, isLoading, isLoadingMore, fetchMore } = useHfModelSearch( @@ -94,6 +152,30 @@ export function HubModelPicker({ [results], ); + const vramMap = useMemo(() => { + const map = new Map< + string, + { est: number; status: VramFitStatus | null; detail: string | null } + >(); + for (const r of results) { + const detail = r.totalParams + ? formatCompact(r.totalParams) + : r.downloads != null + ? `↓${formatCompact(r.downloads)}` + : null; + if (r.totalParams) { + const est = estimateLoadingVram(r.totalParams, "qlora"); + const status = gpu.available + ? checkVramFit(est, gpu.memoryTotalGb) + : null; + map.set(r.id, { est, status, detail }); + } else { + map.set(r.id, { est: 0, status: null, detail }); + } + } + return map; + }, [results, gpu]); + const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, results.length); return ( @@ -144,15 +226,23 @@ export function HubModelPicker({ No matching models. ) : ( - hfIds.map((id) => ( - onSelect(id, { source: "hub", isLora: false })} - /> - )) + hfIds.map((id) => { + const vram = vramMap.get(id); + return ( + + onSelect(id, { source: "hub", isLora: false }) + } + vramStatus={vram?.status ?? null} + vramEst={vram?.est} + gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + /> + ); + }) )}
{isLoadingMore ? (