ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: integrate LoRA model management with UI and runtime synchronization

Browse source
This commit is contained in:
Shine1i 2026-02-13 17:14:49 +01:00
commit c9c4463d5d
8 changed files with 465 additions and 121 deletions
Download patch file Download diff file Expand all files Collapse all files

2
studio/frontend/src/app/provider.tsx
View file

@ -10,7 +10,7 @@ export function AppProvider({ children }: AppProviderProps) {
return (
<ThemeProvider attribute="class" defaultTheme="light">
{children}
<Toaster />
<Toaster position="top-right" />
</ThemeProvider>
);
}

436
studio/frontend/src/components/assistant-ui/model-selector.tsx
View file

@ -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 && (
<span className="size-2 shrink-0 rounded-full bg-emerald-500" />
)}
<span
className={isLoaded ? "text-foreground" : "text-muted-foreground"}
>
{currentModel?.name ?? "Select a model\u2026"}
<span className={isLoaded ? "text-foreground" : "text-muted-foreground"}>
{currentModel?.name ?? "Select model..."}
</span>
{currentModel?.description && (
<span className="text-muted-foreground text-xs">
{currentModel.description}
</span>
<span className="text-muted-foreground text-xs">{currentModel.description}</span>
)}
<HugeiconsIcon
icon={ArrowDown01Icon}
@ -82,95 +99,329 @@ function ModelSelectorTrigger({
);
}
function ListLabel({ children }: { children: ReactNode }) {
return (
<div className="px-2.5 py-1.5 text-[10px] font-semibold uppercase tracking-wider text-muted-foreground">
{children}
</div>
);
}
function ModelRow({
label,
meta,
selected,
onClick,
}: {
label: string;
meta?: string;
selected?: boolean;
onClick: () => void;
}) {
return (
<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",
)}
>
<span className="min-w-0 flex-1 truncate">{label}</span>
{meta ? (
<span className="shrink-0 text-[10px] text-muted-foreground">{meta}</span>
) : null}
</button>
);
}
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 (
<div className="space-y-2">
<div className="relative">
<HugeiconsIcon
icon={Search01Icon}
className="pointer-events-none absolute left-2.5 top-2.5 size-4 text-muted-foreground"
/>
<Input
value={query}
onChange={(event) => setQuery(event.target.value)}
placeholder="Search Hugging Face models"
className="h-9 pl-8 pr-8"
/>
{isLoading && (
<Spinner className="pointer-events-none absolute right-2.5 top-2.5 size-4 text-muted-foreground" />
)}
</div>
<div
ref={scrollRef}
className="max-h-64 overflow-y-auto"
>
<div className="p-1">
{!showHfSection ? (
<>
<ListLabel>Recommended</ListLabel>
{recommendedIds.length === 0 ? (
<div className="px-2.5 py-2 text-xs text-muted-foreground">
No default models.
</div>
) : (
recommendedIds.map((id) => (
<ModelRow
key={id}
label={id}
selected={value === id}
onClick={() => onSelect(id, { source: "hub", isLora: false })}
/>
))
)}
</>
) : null}
{showHfSection ? (
<>
<ListLabel>Hugging Face</ListLabel>
{hfIds.length === 0 && !isLoading ? (
<div className="px-2.5 py-2 text-xs text-muted-foreground">
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 })}
/>
))
)}
<div ref={sentinelRef} className="h-px" />
{isLoadingMore ? (
<div className="flex items-center justify-center py-2">
<Spinner className="size-3.5 text-muted-foreground" />
</div>
) : null}
</>
) : null}
</div>
</div>
</div>
);
}
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<string, LoraModelOption[]>();
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 (
<div className="space-y-2">
<div className="relative">
<HugeiconsIcon
icon={Search01Icon}
className="pointer-events-none absolute left-2.5 top-2.5 size-4 text-muted-foreground"
/>
<Input
value={query}
onChange={(event) => setQuery(event.target.value)}
placeholder="Search local adapters"
className="h-9 pl-8"
/>
</div>
<div className="max-h-64 overflow-y-auto">
<div className="p-1">
{grouped.length === 0 ? (
<div className="px-2.5 py-2 text-xs text-muted-foreground">No adapters found.</div>
) : (
grouped.map(([baseModel, adapters], index) => (
<div key={baseModel}>
{index > 0 ? <div className="my-1" /> : null}
<ListLabel>{baseModel}</ListLabel>
{adapters.map((adapter) => (
<ModelRow
key={adapter.id}
label={adapter.name}
meta="LoRA"
selected={value === adapter.id}
onClick={() => onSelect(adapter.id, { source: "lora", isLora: true })}
/>
))}
</div>
))
)}
</div>
</div>
</div>
);
}
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 (
<PopoverContent
align="start"
className={cn("w-auto min-w-[280px] gap-0 p-1", className)}
className={cn("w-[440px] min-w-[440px] gap-0 p-2", className)}
>
{models.map((model) => (
<ModelSelectorItem
key={model.id}
model={model}
isActive={value === model.id}
onSelect={onSelect}
onEject={onEject}
/>
))}
<Tabs defaultValue="hub" className="w-full">
<TabsList className="mb-2 w-full">
<TabsTrigger value="hub">Hub models</TabsTrigger>
<TabsTrigger value="lora">Fine-tuned</TabsTrigger>
</TabsList>
<TabsContent value="hub" className="m-0">
<HubModelPicker models={models} value={value} onSelect={onSelect} />
</TabsContent>
<TabsContent value="lora" className="m-0">
<LoraModelPicker
loraModels={loraModels}
value={value}
onSelect={onSelect}
/>
</TabsContent>
</Tabs>
{hasSelection && onEject ? (
<div className="mt-2 border-t border-border/70 pt-2">
<button
type="button"
onClick={onEject}
className="flex w-full items-center justify-center gap-1.5 rounded-md px-2 py-1.5 text-xs text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
title="Eject model"
>
<HugeiconsIcon icon={Logout01Icon} className="size-3.5" />
Eject loaded model
</button>
</div>
) : null}
</PopoverContent>
);
}
function ModelSelectorItem({
model,
isActive,
onSelect,
onEject,
}: {
model: ModelOption;
isActive: boolean;
onSelect: (id: string) => void;
onEject?: () => void;
}) {
return (
<button
type="button"
aria-pressed={isActive}
onClick={() => onSelect(model.id)}
className={cn(
"group flex w-full cursor-pointer items-center gap-2.5 rounded-md px-3 py-2 text-left text-sm transition-colors hover:bg-accent",
isActive && "bg-accent/50",
)}
>
<span
className={cn(
"size-2 shrink-0 rounded-full",
isActive ? "bg-emerald-500" : "bg-transparent",
)}
/>
{model.icon && <span className="shrink-0">{model.icon}</span>}
<div className="min-w-0 flex-1">
<div className="truncate text-sm">{model.name}</div>
{model.description && (
<div className="truncate text-xs text-muted-foreground">
{model.description}
</div>
)}
</div>
{isActive && onEject && (
<button
type="button"
onClick={(e) => {
e.stopPropagation();
onEject();
}}
className="flex items-center gap-1 rounded px-1.5 py-0.5 text-[10px] text-muted-foreground opacity-0 transition-opacity group-hover:opacity-100 hover:bg-muted hover:text-foreground"
title="Eject model"
>
<HugeiconsIcon icon={Logout01Icon} className="size-3" />
Eject
</button>
)}
</button>
);
}
// --- 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<string, ModelOption>();
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({
/>
<ModelSelectorContent
models={models}
loraModels={loraModels}
value={selected}
onSelect={handleSelect}
onEject={onEject ? handleEject : undefined}
@ -220,7 +490,5 @@ export function ModelSelector({
);
}
// Composable exports
ModelSelector.Trigger = ModelSelectorTrigger;
ModelSelector.Content = ModelSelectorContent;
ModelSelector.Item = ModelSelectorItem;

9
studio/frontend/src/features/chat/api/chat-api.ts
View file

@ -1,11 +1,12 @@
import { authFetch } from "@/features/auth";
import type {
InferenceStatusResponse,
ListLorasResponse,
ListModelsResponse,
LoadModelRequest,
LoadModelResponse,
OpenAIChatCompletionsRequest,
OpenAIChatChunk,
OpenAIChatCompletionsRequest,
UnloadModelRequest,
} from "../types/api";
@ -42,6 +43,12 @@ export async function listModels(): Promise<ListModelsResponse> {
return parseJsonOrThrow<ListModelsResponse>(response);
}
export async function listLoras(outputsDir = "./outputs"): Promise<ListLorasResponse> {
const query = new URLSearchParams({ outputs_dir: outputsDir }).toString();
const response = await authFetch(`/api/models/loras?${query}`);
return parseJsonOrThrow<ListLorasResponse>(response);
}
export async function getInferenceStatus(): Promise<InferenceStatusResponse> {
const response = await authFetch("/api/inference/status");
return parseJsonOrThrow<InferenceStatusResponse>(response);

23
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -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<LoraModelOption[]>(
() =>
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 {
/>
<ModelSelector
models={models}
loraModels={loraModels}
value={inferenceParams.checkpoint}
onValueChange={handleCheckpointChange}
onEject={handleEject}
@ -286,10 +300,7 @@ export function ChatPage(): ReactElement {
</div>
{view.mode === "single" ? (
<SingleContent
key={view.threadId ?? "new"}
threadId={view.threadId}
/>
<SingleContent key={view.threadId ?? "new"} threadId={view.threadId} />
) : (
<CompareContent key={view.pairId} pairId={view.pairId} />
)}

97
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -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,

5
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -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<ChatRuntimeStore>((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) => ({

11
studio/frontend/src/features/chat/types/api.ts
View file

@ -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;

7
studio/frontend/src/features/chat/types/runtime.ts
View file

@ -25,3 +25,10 @@ export interface ChatModelSummary {
isVision: boolean;
isLora: boolean;
}
export interface ChatLoraSummary {
id: string;
name: string;
baseModel: string;
updatedAt?: number;
}