Merge pull request #193 from unslothai/feature/vram-fit-chat
Added vram fit indicator to models in chat
This commit is contained in:
Roland Tannous
GitHub
commit
ea324f2806
3 changed files with 208 additions and 23 deletions
|
|
@ -1,7 +1,20 @@
|
|||
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,
|
||||
useRecommendedModelVram,
|
||||
} 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 +41,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 = (
|
||||
<button
|
||||
type="button"
|
||||
onClick={onClick}
|
||||
className={cn(
|
||||
"flex w-full items-center justify-between gap-2 rounded-md px-2.5 py-1.5 text-left text-sm transition-colors hover:bg-accent",
|
||||
selected && "bg-accent/60",
|
||||
exceeds && "opacity-50",
|
||||
)}
|
||||
>
|
||||
<span className="min-w-0 flex-1 truncate">{label}</span>
|
||||
{meta ? (
|
||||
<span className="shrink-0 text-[10px] text-muted-foreground">{meta}</span>
|
||||
) : null}
|
||||
<span
|
||||
className={cn(
|
||||
"min-w-0 flex-1 truncate",
|
||||
exceeds && "line-through decoration-muted-foreground/50",
|
||||
)}
|
||||
>
|
||||
{label}
|
||||
</span>
|
||||
<span className="flex items-center gap-1.5 shrink-0">
|
||||
{vramStatus === "exceeds" && (
|
||||
<span className="text-[9px] font-medium text-red-400">OOM</span>
|
||||
)}
|
||||
{vramStatus === "tight" && (
|
||||
<span className="text-[9px] font-medium text-amber-400">TIGHT</span>
|
||||
)}
|
||||
{vramStatus === "fits" && (
|
||||
<span className="text-[9px] font-medium text-emerald-500/90">FIT</span>
|
||||
)}
|
||||
{meta ? (
|
||||
<span className="text-[10px] text-muted-foreground">{meta}</span>
|
||||
) : null}
|
||||
</span>
|
||||
</button>
|
||||
);
|
||||
|
||||
if (vramTooltipText) {
|
||||
return (
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild>{content}</TooltipTrigger>
|
||||
<TooltipContent side="left" className="max-w-xs break-all">
|
||||
{label}
|
||||
<span className="block text-[10px] mt-1">{vramTooltipText}</span>
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
);
|
||||
}
|
||||
return content;
|
||||
}
|
||||
|
||||
export function HubModelPicker({
|
||||
|
|
@ -60,6 +123,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(
|
||||
|
|
@ -71,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]);
|
||||
|
||||
|
|
@ -94,6 +161,49 @@ 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 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 (
|
||||
|
|
@ -124,14 +234,23 @@ export function HubModelPicker({
|
|||
No default models.
|
||||
</div>
|
||||
) : (
|
||||
recommendedIds.map((id) => (
|
||||
<ModelRow
|
||||
key={id}
|
||||
label={id}
|
||||
selected={value === id}
|
||||
onClick={() => onSelect(id, { source: "hub", isLora: false })}
|
||||
/>
|
||||
))
|
||||
recommendedIds.map((id) => {
|
||||
const vram = recommendedVramMap.get(id);
|
||||
return (
|
||||
<ModelRow
|
||||
key={id}
|
||||
label={id}
|
||||
meta={vram?.detail ?? undefined}
|
||||
selected={value === id}
|
||||
onClick={() =>
|
||||
onSelect(id, { source: "hub", isLora: false })
|
||||
}
|
||||
vramStatus={vram?.status ?? null}
|
||||
vramEst={vram?.est}
|
||||
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
|
||||
/>
|
||||
);
|
||||
})
|
||||
)}
|
||||
</>
|
||||
) : null}
|
||||
|
|
@ -144,15 +263,23 @@ export function HubModelPicker({
|
|||
No matching models.
|
||||
</div>
|
||||
) : (
|
||||
hfIds.map((id) => (
|
||||
<ModelRow
|
||||
key={id}
|
||||
label={id}
|
||||
meta={metricsById.get(id)}
|
||||
selected={value === id}
|
||||
onClick={() => onSelect(id, { source: "hub", isLora: false })}
|
||||
/>
|
||||
))
|
||||
hfIds.map((id) => {
|
||||
const vram = vramMap.get(id);
|
||||
return (
|
||||
<ModelRow
|
||||
key={id}
|
||||
label={id}
|
||||
meta={metricsById.get(id)}
|
||||
selected={value === id}
|
||||
onClick={() =>
|
||||
onSelect(id, { source: "hub", isLora: false })
|
||||
}
|
||||
vramStatus={vram?.status ?? null}
|
||||
vramEst={vram?.est}
|
||||
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
|
||||
/>
|
||||
);
|
||||
})
|
||||
)}
|
||||
<div ref={sentinelRef} className="h-px" />
|
||||
{isLoadingMore ? (
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
|
|
|
|||
57
studio/frontend/src/hooks/use-recommended-model-vram.ts
Normal file
57
studio/frontend/src/hooks/use-recommended-model-vram.ts
Normal file
|
|
@ -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<string, number>
|
||||
>(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<string, number>();
|
||||
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 };
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue