diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index ce3af7ac54..5c2b93bcc9 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -10,7 +10,7 @@ export function AppProvider({ children }: AppProviderProps) { return ( {children} - + ); } diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 389dae424b..85bc1cacdb 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -1,14 +1,22 @@ "use client"; +import { Input } from "@/components/ui/input"; import { Popover, PopoverContent, PopoverTrigger, } from "@/components/ui/popover"; -import { cn } from "@/lib/utils"; -import { ArrowDown01Icon, Logout01Icon } from "@hugeicons/core-free-icons"; +import { Spinner } from "@/components/ui/spinner"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { useDebouncedValue, useHfModelSearch, useInfiniteScroll } from "@/hooks"; +import { cn, formatCompact } from "@/lib/utils"; +import { + ArrowDown01Icon, + Logout01Icon, + Search01Icon, +} from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { type ReactNode, useState } from "react"; +import { type ReactNode, useMemo, useState } from "react"; export interface ModelOption { id: string; @@ -17,11 +25,22 @@ export interface ModelOption { icon?: ReactNode; } +export interface LoraModelOption extends ModelOption { + baseModel?: string; + updatedAt?: number; +} + +export interface ModelSelectorChangeMeta { + source: "hub" | "lora"; + isLora: boolean; +} + interface ModelSelectorProps { models: ModelOption[]; + loraModels?: LoraModelOption[]; value?: string; defaultValue?: string; - onValueChange?: (value: string) => void; + onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void; onEject?: () => void; variant?: "outline" | "ghost" | "muted"; size?: "sm" | "default" | "lg"; @@ -29,7 +48,9 @@ interface ModelSelectorProps { contentClassName?: string; } -// --- Composable sub-components --- +function dedupe(values: string[]): string[] { + return [...new Set(values.filter(Boolean))]; +} function ModelSelectorTrigger({ currentModel, @@ -63,15 +84,11 @@ function ModelSelectorTrigger({ {isLoaded && ( )} - - {currentModel?.name ?? "Select a model\u2026"} + + {currentModel?.name ?? "Select model..."} {currentModel?.description && ( - - {currentModel.description} - + {currentModel.description} )} + {children} + + ); +} + +function ModelRow({ + label, + meta, + selected, + onClick, +}: { + label: string; + meta?: string; + selected?: boolean; + onClick: () => void; +}) { + return ( + + ); +} + +function HubModelPicker({ + models, + value, + onSelect, +}: { + models: ModelOption[]; + value?: string; + onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; +}) { + const [query, setQuery] = useState(""); + const debouncedQuery = useDebouncedValue(query); + const { results, isLoading, isLoadingMore, fetchMore } = useHfModelSearch( + debouncedQuery, + ); + + const recommendedIds = useMemo( + () => dedupe([...models.map((model) => model.id), value ?? ""]), + [models, value], + ); + + const showHfSection = debouncedQuery.trim().length > 0; + const recommendedSet = useMemo( + () => new Set(recommendedIds), + [recommendedIds], + ); + + const hfIds = useMemo(() => { + if (!showHfSection) { + return []; + } + return results + .map((result) => result.id) + .filter((id) => !recommendedSet.has(id)); + }, [recommendedSet, results, showHfSection]); + + const metricsById = useMemo( + () => + new Map( + results.map((result) => [ + result.id, + result.totalParams + ? formatCompact(result.totalParams) + : `↓${formatCompact(result.downloads)}`, + ]), + ), + [results], + ); + + const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, results.length); + + return ( +
+
+ + setQuery(event.target.value)} + placeholder="Search Hugging Face models" + className="h-9 pl-8 pr-8" + /> + {isLoading && ( + + )} +
+ +
+
+ {!showHfSection ? ( + <> + Recommended + {recommendedIds.length === 0 ? ( +
+ No default models. +
+ ) : ( + recommendedIds.map((id) => ( + onSelect(id, { source: "hub", isLora: false })} + /> + )) + )} + + ) : null} + + {showHfSection ? ( + <> + Hugging Face + {hfIds.length === 0 && !isLoading ? ( +
+ No matching models. +
+ ) : ( + hfIds.map((id) => ( + onSelect(id, { source: "hub", isLora: false })} + /> + )) + )} +
+ {isLoadingMore ? ( +
+ +
+ ) : null} + + ) : null} +
+
+
+ ); +} + +function LoraModelPicker({ + loraModels, + value, + onSelect, +}: { + loraModels: LoraModelOption[]; + value?: string; + onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; +}) { + const [query, setQuery] = useState(""); + + const normalized = useMemo( + () => + loraModels + .map((model) => ({ + ...model, + baseModel: model.baseModel || model.description || "Unknown base model", + })) + .sort((a, b) => { + const aTime = a.updatedAt ?? -1; + const bTime = b.updatedAt ?? -1; + if (aTime !== bTime) { + return bTime - aTime; + } + const baseCmp = a.baseModel.localeCompare(b.baseModel); + if (baseCmp !== 0) { + return baseCmp; + } + return a.name.localeCompare(b.name); + }), + [loraModels], + ); + + const grouped = useMemo(() => { + const needle = query.trim().toLowerCase(); + const out = new Map(); + + for (const model of normalized) { + const searchText = `${model.name} ${model.baseModel} ${model.id}`.toLowerCase(); + if (needle && !searchText.includes(needle)) { + continue; + } + + const key = model.baseModel || "Unknown base model"; + const prev = out.get(key) ?? []; + prev.push(model); + out.set(key, prev); + } + + return [...out.entries()].sort((a, b) => { + const aLatest = Math.max(...a[1].map((model) => model.updatedAt ?? -1)); + const bLatest = Math.max(...b[1].map((model) => model.updatedAt ?? -1)); + if (aLatest !== bLatest) { + return bLatest - aLatest; + } + return a[0].localeCompare(b[0]); + }); + }, [normalized, query]); + + return ( +
+
+ + setQuery(event.target.value)} + placeholder="Search local adapters" + className="h-9 pl-8" + /> +
+ +
+
+ {grouped.length === 0 ? ( +
No adapters found.
+ ) : ( + grouped.map(([baseModel, adapters], index) => ( +
+ {index > 0 ?
: null} + {baseModel} + {adapters.map((adapter) => ( + onSelect(adapter.id, { source: "lora", isLora: true })} + /> + ))} +
+ )) + )} +
+
+
+ ); +} + function ModelSelectorContent({ models, + loraModels, value, onSelect, onEject, className, }: { models: ModelOption[]; + loraModels: LoraModelOption[]; value?: string; - onSelect: (id: string) => void; + onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; onEject?: () => void; className?: string; }) { + const hasSelection = Boolean(value); + return ( - {models.map((model) => ( - - ))} + + + Hub models + Fine-tuned + + + + + + + + + + + + {hasSelection && onEject ? ( +
+ +
+ ) : null}
); } -function ModelSelectorItem({ - model, - isActive, - onSelect, - onEject, -}: { - model: ModelOption; - isActive: boolean; - onSelect: (id: string) => void; - onEject?: () => void; -}) { - return ( - - )} - - ); -} - -// --- Main component --- - export function ModelSelector({ models, + loraModels = [], value, defaultValue, onValueChange, @@ -182,13 +433,31 @@ export function ModelSelector({ }: ModelSelectorProps) { const [open, setOpen] = useState(false); const [uncontrolled, setUncontrolled] = useState(defaultValue ?? ""); + const selected = value ?? uncontrolled; const isLoaded = selected !== ""; - const currentModel = models.find((m) => m.id === selected); - function handleSelect(id: string) { + const optionById = useMemo(() => { + const all = new Map(); + for (const model of models) { + all.set(model.id, model); + } + for (const lora of loraModels) { + all.set(lora.id, { + ...lora, + description: lora.baseModel || lora.description, + }); + } + return all; + }, [loraModels, models]); + + const currentModel = selected + ? optionById.get(selected) ?? { id: selected, name: selected } + : undefined; + + function handleSelect(id: string, meta: ModelSelectorChangeMeta) { if (onValueChange) { - onValueChange(id); + onValueChange(id, meta); } else { setUncontrolled(id); } @@ -211,6 +480,7 @@ export function ModelSelector({ /> { return parseJsonOrThrow(response); } +export async function listLoras(outputsDir = "./outputs"): Promise { + const query = new URLSearchParams({ outputs_dir: outputsDir }).toString(); + const response = await authFetch(`/api/models/loras?${query}`); + return parseJsonOrThrow(response); +} + export async function getInferenceStatus(): Promise { const response = await authFetch("/api/inference/status"); return parseJsonOrThrow(response); diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 6d201f50de..c787cf471c 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -1,4 +1,5 @@ import { + type LoraModelOption, type ModelOption, ModelSelector, } from "@/components/assistant-ui/model-selector"; @@ -201,12 +202,13 @@ export function ChatPage(): ReactElement { const inferenceParams = useChatRuntimeStore((state) => state.params); const setInferenceParams = useChatRuntimeStore((state) => state.setParams); const modelsFromStore = useChatRuntimeStore((state) => state.models); + const lorasFromStore = useChatRuntimeStore((state) => state.loras); const modelsError = useChatRuntimeStore((state) => state.modelsError); const { refresh, selectModel, ejectModel } = useChatModelRuntime(); const handleCheckpointChange = useCallback( - (value: string) => { - void selectModel(value); + (value: string, meta?: { isLora: boolean }) => { + void selectModel({ id: value, isLora: meta?.isLora }); }, [selectModel], ); @@ -229,6 +231,17 @@ export function ChatPage(): ReactElement { [modelsFromStore], ); + const loraModels = useMemo( + () => + lorasFromStore.map((lora) => ({ + id: lora.id, + name: lora.name, + baseModel: lora.baseModel, + updatedAt: lora.updatedAt, + })), + [lorasFromStore], + ); + useEffect(() => { void refresh(); }, [refresh]); @@ -263,6 +276,7 @@ export function ChatPage(): ReactElement { /> {view.mode === "single" ? ( - + ) : ( )} diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index e0aaca9622..3eb9301302 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -1,15 +1,38 @@ import { useCallback } from "react"; +import { toast } from "sonner"; import { getInferenceStatus, + listLoras, listModels, loadModel, unloadModel, } from "../api/chat-api"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; -import type { ChatModelSummary } from "../types/runtime"; +import type { ChatLoraSummary, ChatModelSummary } from "../types/runtime"; const DEFAULT_MODEL_MAX_SEQ_LENGTH = 2048; +type SelectedModelInput = { + id: string; + isLora?: boolean; +}; + +const LORA_SUFFIX_RE = /_(\d{9,})$/; + +function parseTrailingEpoch(input: string): number | undefined { + const match = input.match(LORA_SUFFIX_RE); + if (!match) { + return undefined; + } + const parsed = Number.parseInt(match[1], 10); + return Number.isFinite(parsed) ? parsed : undefined; +} + +function stripTrailingEpoch(input: string): string { + const cleaned = input.replace(LORA_SUFFIX_RE, "").replace(/[_-]+$/, "").trim(); + return cleaned || input; +} + function describeModel(model: { is_lora?: boolean; is_vision?: boolean; @@ -36,10 +59,29 @@ function toChatModelSummary(model: { }; } +function toLoraSummary(lora: { + display_name: string; + adapter_path: string; + base_model?: string | null; +}): ChatLoraSummary { + const idTail = lora.adapter_path.split("/").filter(Boolean).at(-1) ?? ""; + const updatedAt = + parseTrailingEpoch(lora.display_name) ?? parseTrailingEpoch(idTail); + + return { + id: lora.adapter_path, + name: stripTrailingEpoch(lora.display_name), + baseModel: lora.base_model || "Unknown base model", + updatedAt, + }; +} + export function useChatModelRuntime() { const params = useChatRuntimeStore((state) => state.params); const models = useChatRuntimeStore((state) => state.models); + const loras = useChatRuntimeStore((state) => state.loras); const setModels = useChatRuntimeStore((state) => state.setModels); + const setLoras = useChatRuntimeStore((state) => state.setLoras); const setModelsError = useChatRuntimeStore((state) => state.setModelsError); const setCheckpoint = useChatRuntimeStore((state) => state.setCheckpoint); const clearCheckpoint = useChatRuntimeStore((state) => state.clearCheckpoint); @@ -47,13 +89,14 @@ export function useChatModelRuntime() { const refresh = useCallback(async () => { setModelsError(null); try { - const [listRes, statusRes] = await Promise.all([ + const [listRes, statusRes, lorasRes] = await Promise.all([ listModels(), getInferenceStatus(), + listLoras(), ]); - const modelList = listRes.models.map(toChatModelSummary); - setModels(modelList); + setModels(listRes.models.map(toChatModelSummary)); + setLoras(lorasRes.loras.map(toLoraSummary)); if (statusRes.active_model) { setCheckpoint(statusRes.active_model); @@ -63,22 +106,23 @@ export function useChatModelRuntime() { error instanceof Error ? error.message : "Failed to load models"; setModelsError(message); } - }, [ - setCheckpoint, - setModels, - setModelsError, - ]); + }, [setCheckpoint, setLoras, setModels, setModelsError]); const selectModel = useCallback( - async (modelId: string) => { + async (selection: string | SelectedModelInput) => { + const modelId = typeof selection === "string" ? selection : selection.id; if (!modelId || params.checkpoint === modelId) { return; } - const selected = models.find((model) => model.id === modelId); - if (!selected) { - setModelsError("Selected model was not found in model list."); - return; - } + + const explicitIsLora = + typeof selection === "string" ? undefined : selection.isLora; + const model = models.find((entry) => entry.id === modelId); + const lora = loras.find((entry) => entry.id === modelId); + const isLora = + explicitIsLora ?? model?.isLora ?? (lora ? true : false); + const displayName = model?.name || lora?.name || modelId; + const loadingToastId = toast.loading(`Loading ${displayName}...`); setModelsError(null); try { @@ -87,28 +131,24 @@ export function useChatModelRuntime() { } await loadModel({ - model_path: selected.id, + model_path: modelId, hf_token: null, max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH, load_in_4bit: true, - is_lora: selected.isLora, + is_lora: isLora, }); - setCheckpoint(selected.id); + setCheckpoint(modelId); await refresh(); + toast.success(`${displayName} loaded`, { id: loadingToastId }); } catch (error) { const message = error instanceof Error ? error.message : "Failed to load model"; setModelsError(message); + toast.error(message, { id: loadingToastId }); } }, - [ - models, - params.checkpoint, - refresh, - setCheckpoint, - setModelsError, - ], + [loras, models, params.checkpoint, refresh, setCheckpoint, setModelsError], ); const ejectModel = useCallback(async () => { @@ -125,12 +165,7 @@ export function useChatModelRuntime() { error instanceof Error ? error.message : "Failed to unload model"; setModelsError(message); } - }, [ - clearCheckpoint, - params.checkpoint, - refresh, - setModelsError, - ]); + }, [clearCheckpoint, params.checkpoint, refresh, setModelsError]); return { refresh, 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 94b4e3fdb3..76bf079ce7 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -1,6 +1,7 @@ import { create } from "zustand"; import { DEFAULT_INFERENCE_PARAMS, + type ChatLoraSummary, type ChatModelSummary, type InferenceParams, } from "../types/runtime"; @@ -8,9 +9,11 @@ import { type ChatRuntimeStore = { params: InferenceParams; models: ChatModelSummary[]; + loras: ChatLoraSummary[]; modelsError: string | null; setParams: (params: InferenceParams) => void; setModels: (models: ChatModelSummary[]) => void; + setLoras: (loras: ChatLoraSummary[]) => void; setModelsError: (error: string | null) => void; setCheckpoint: (modelId: string) => void; clearCheckpoint: () => void; @@ -19,9 +22,11 @@ type ChatRuntimeStore = { export const useChatRuntimeStore = create((set) => ({ params: DEFAULT_INFERENCE_PARAMS, models: [], + loras: [], modelsError: null, setParams: (params) => set({ params }), setModels: (models) => set({ models }), + setLoras: (loras) => set({ loras }), setModelsError: (modelsError) => set({ modelsError }), setCheckpoint: (modelId) => set((state) => ({ diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 6420020c06..9d21c854af 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -10,6 +10,17 @@ export interface ListModelsResponse { default_models: string[]; } +export interface BackendLoraInfo { + display_name: string; + adapter_path: string; + base_model?: string | null; +} + +export interface ListLorasResponse { + loras: BackendLoraInfo[]; + outputs_dir: string; +} + export interface LoadModelRequest { model_path: string; hf_token: string | null; diff --git a/studio/frontend/src/features/chat/types/runtime.ts b/studio/frontend/src/features/chat/types/runtime.ts index 67fa07eb14..558f1a5464 100644 --- a/studio/frontend/src/features/chat/types/runtime.ts +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -25,3 +25,10 @@ export interface ChatModelSummary { isVision: boolean; isLora: boolean; } + +export interface ChatLoraSummary { + id: string; + name: string; + baseModel: string; + updatedAt?: number; +}