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 e023e36e2b..fc579dd32a 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -5,7 +5,13 @@ import { TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { useDebouncedValue, useGpuInfo, useHfModelSearch, useInfiniteScroll } from "@/hooks"; +import { + useDebouncedValue, + useGpuInfo, + useHfModelSearch, + useInfiniteScroll, + useRecommendedModelVram, +} from "@/hooks"; import { cn, formatCompact } from "@/lib/utils"; import type { VramFitStatus } from "@/lib/vram"; import { checkVramFit, estimateLoadingVram } from "@/lib/vram"; @@ -129,6 +135,9 @@ export function HubModelPicker({ [models, value], ); + const { paramCountById: recommendedParamCountById } = + useRecommendedModelVram(recommendedIds); + const showHfSection = debouncedQuery.trim().length > 0; const recommendedSet = useMemo(() => new Set(recommendedIds), [recommendedIds]); @@ -176,6 +185,25 @@ export function HubModelPicker({ return map; }, [results, gpu]); + const recommendedVramMap = useMemo(() => { + const map = new Map< + string, + { est: number; status: VramFitStatus | null; detail: string | null } + >(); + for (const id of recommendedIds) { + const totalParams = recommendedParamCountById.get(id); + if (totalParams) { + const est = estimateLoadingVram(totalParams, "qlora"); + const status = gpu.available + ? checkVramFit(est, gpu.memoryTotalGb) + : null; + const detail = formatCompact(totalParams); + map.set(id, { est, status, detail }); + } + } + return map; + }, [recommendedIds, recommendedParamCountById, gpu]); + const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, results.length); return ( @@ -206,14 +234,23 @@ export function HubModelPicker({ No default models. ) : ( - recommendedIds.map((id) => ( - onSelect(id, { source: "hub", isLora: false })} - /> - )) + recommendedIds.map((id) => { + const vram = recommendedVramMap.get(id); + return ( + + onSelect(id, { source: "hub", isLora: false }) + } + vramStatus={vram?.status ?? null} + vramEst={vram?.est} + gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + /> + ); + }) )} ) : null} diff --git a/studio/frontend/src/hooks/index.ts b/studio/frontend/src/hooks/index.ts index b6c4fcad14..3932baf088 100644 --- a/studio/frontend/src/hooks/index.ts +++ b/studio/frontend/src/hooks/index.ts @@ -3,6 +3,7 @@ export { useGpuInfo } from "./use-gpu-info"; export { useGpuUtilization } from "./use-gpu-utilization"; export { useHardwareInfo } from "./use-hardware-info"; export { useHfModelSearch } from "./use-hf-model-search"; +export { useRecommendedModelVram } from "./use-recommended-model-vram"; export { useHfDatasetSearch } from "./use-hf-dataset-search"; export { useHfDatasetSplits } from "./use-hf-dataset-splits"; export { useHfTokenValidation } from "./use-hf-token-validation"; diff --git a/studio/frontend/src/hooks/use-recommended-model-vram.ts b/studio/frontend/src/hooks/use-recommended-model-vram.ts new file mode 100644 index 0000000000..69692ff416 --- /dev/null +++ b/studio/frontend/src/hooks/use-recommended-model-vram.ts @@ -0,0 +1,57 @@ +import { modelInfo } from "@huggingface/hub"; +import { useEffect, useState } from "react"; + +/** + * Fetches Hugging Face model info (safetensors total param count) for a list of + * model IDs. Used to show VRAM fit (FIT / TIGHT / OOM) for recommended/default + * models in the chat model dropdown. + */ +export function useRecommendedModelVram(ids: string[]) { + const [paramCountById, setParamCountById] = useState< + Map + >(new Map()); + const [isLoading, setIsLoading] = useState(false); + + const stableKey = [...ids].filter(Boolean).sort().join(","); + + useEffect(() => { + const stableIds = stableKey ? stableKey.split(",") : []; + if (stableIds.length === 0) { + setParamCountById(new Map()); + setIsLoading(false); + return; + } + let canceled = false; + void (async () => { + setIsLoading(true); + const next = new Map(); + await Promise.all( + stableIds.map(async (id) => { + if (canceled) return; + try { + const info = await modelInfo({ + name: id, + additionalFields: ["safetensors"], + }); + const raw = info as { safetensors?: { total?: number } }; + const total = raw.safetensors?.total; + if (typeof total === "number" && total > 0) { + next.set(id, total); + } + } catch { + // Model not on HF or no safetensors; skip + } + }), + ); + if (!canceled) { + setParamCountById(next); + setIsLoading(false); + } + })(); + return () => { + canceled = true; + }; + }, [stableKey]); + + return { paramCountById, isLoading }; +}