Merge pull request #192 from unslothai/feature/model-download-status

Updated to edit loading as "downloading model"

studio/frontend/src/features/chat/chat-page.tsx

@@ -5,6 +5,7 @@ import {
 } from "@/components/assistant-ui/model-selector";
 import { Thread } from "@/components/assistant-ui/thread";
 import { Button } from "@/components/ui/button";
+import { Spinner } from "@/components/ui/spinner";
 import { SidebarProvider, SidebarTrigger, useSidebar } from "@/components/ui/sidebar";
 import {
   Sheet,
@@ -283,7 +284,8 @@ export function ChatPage(): ReactElement {
   const modelsFromStore = useChatRuntimeStore((state) => state.models);
   const lorasFromStore = useChatRuntimeStore((state) => state.loras);
   const modelsError = useChatRuntimeStore((state) => state.modelsError);
-  const { refresh, selectModel, ejectModel } = useChatModelRuntime();
+  const { refresh, selectModel, ejectModel, loadingModel } =
+    useChatModelRuntime();
   const refreshRef = useRef(refresh);
   const selectModelRef = useRef(selectModel);
 
@@ -518,6 +520,17 @@ export function ChatPage(): ReactElement {
               contentDataTour="chat-model-selector-popover"
               className="max-w-[62vw] sm:max-w-none"
             />
+            {loadingModel ? (
+              <div
+                className="flex items-center gap-1.5 text-muted-foreground"
+                title={`Loading ${loadingModel.displayName}. This may include downloading.`}
+              >
+                <Spinner className="size-3.5 shrink-0" />
+                <span className="text-xs">
+                  Downloading model…
+                </span>
+              </div>
+            ) : null}
           </div>
           {modelsError && (
             <div className="ml-2 text-xs text-destructive truncate max-w-[28rem]">

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

@@ -1,4 +1,4 @@
-import { useCallback } from "react";
+import { useCallback, useState } from "react";
 import { toast } from "sonner";
 import {
   getInferenceStatus,
@@ -116,6 +116,11 @@ export function useChatModelRuntime() {
   const setCheckpoint = useChatRuntimeStore((state) => state.setCheckpoint);
   const clearCheckpoint = useChatRuntimeStore((state) => state.clearCheckpoint);
 
+  const [loadingModel, setLoadingModel] = useState<{
+    id: string;
+    displayName: string;
+  } | null>(null);
+
   const refresh = useCallback(async () => {
     setModelsError(null);
     try {
@@ -157,6 +162,7 @@ export function useChatModelRuntime() {
       const displayName = model?.name || lora?.name || modelId;
 
       setModelsError(null);
+      setLoadingModel({ id: modelId, displayName });
       try {
         async function performLoad(): Promise<void> {
           if (params.checkpoint) {
@@ -176,19 +182,20 @@ export function useChatModelRuntime() {
           await refresh();
         }
 
-        let description = "Base model selected.";
-        if (isLora) {
-          description = "Fine-tuned (LoRA) selected.";
-        }
+        const loadPromise = performLoad().finally(() => {
+          setLoadingModel(null);
+        });
 
-        await toast.promise(performLoad(), {
-          loading: `Loading ${displayName}`,
+        await toast.promise(loadPromise, {
+          loading: "Downloading model…",
           success: `${displayName} loaded`,
           error: (err) =>
             err instanceof Error ? err.message : "Failed to load model",
-          description,
+          description:
+            "This may include downloading. Large models can take a while.",
         });
       } catch (error) {
+        setLoadingModel(null);
         const message =
           error instanceof Error ? error.message : "Failed to load model";
         setModelsError(message);
@@ -227,5 +234,6 @@ export function useChatModelRuntime() {
     refresh,
     selectModel,
     ejectModel,
+    loadingModel,
   };
 }