diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 4b162c5727..21bf815695 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -628,7 +628,7 @@ class InferenceBackend: top_k: int = 40, min_p: float = 0.0, max_new_tokens: int = 256, - repetition_penalty: float = 1.1, + repetition_penalty: float = 1.0, cancel_event = None, ) -> Generator[str, None, None]: """ @@ -658,7 +658,7 @@ class InferenceBackend: top_k: int = 40, min_p: float = 0.0, max_new_tokens: int = 256, - repetition_penalty: float = 1.1, + repetition_penalty: float = 1.0, cancel_event = None, _adapter_state = None, ) -> Generator[str, None, None]: @@ -1077,7 +1077,7 @@ class InferenceBackend: top_k: int = 40, min_p: float = 0.0, max_new_tokens: int = 256, - repetition_penalty: float = 1.1, + repetition_penalty: float = 1.0, cancel_event = None, _adapter_state = None, ) -> Generator[str, None, None]: @@ -1215,7 +1215,7 @@ class InferenceBackend: top_k: int = 50, min_p: float = 0.0, max_new_tokens: int = 2048, - repetition_penalty: float = 1.1, + repetition_penalty: float = 1.0, use_adapter: Optional[Union[bool, str]] = None, ) -> Tuple[bytes, int]: """ diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index e793114f30..ba45cd7735 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -535,6 +535,47 @@ class LlamaCppBackend: logger.info(f"GGUF downloaded to: {local_path}") return local_path + def _download_mmproj( + self, + *, + hf_repo: str, + hf_token: Optional[str] = None, + ) -> Optional[str]: + """Download the mmproj (vision projection) file from a GGUF repo. + + Prefers mmproj-F16.gguf, falls back to any mmproj*.gguf file. + Returns the local path, or None if no mmproj file exists. + """ + try: + from huggingface_hub import hf_hub_download, list_repo_files + + files = list_repo_files(hf_repo, token = hf_token) + mmproj_files = sorted( + f for f in files if f.endswith(".gguf") and "mmproj" in f.lower() + ) + if not mmproj_files: + return None + + # Prefer F16 variant + target = None + for f in mmproj_files: + if "f16" in f.lower(): + target = f + break + if target is None: + target = mmproj_files[0] + + logger.info(f"Downloading mmproj: {hf_repo}/{target}") + local_path = hf_hub_download( + repo_id = hf_repo, + filename = target, + token = hf_token, + ) + return local_path + except Exception as e: + logger.warning(f"Could not download mmproj: {e}") + return None + # ── Lifecycle ───────────────────────────────────────────────── def load_model( @@ -588,6 +629,12 @@ class LlamaCppBackend: hf_variant = hf_variant, hf_token = hf_token, ) + # Auto-download mmproj for vision models + if is_vision and not mmproj_path: + mmproj_path = self._download_mmproj( + hf_repo = hf_repo, + hf_token = hf_token, + ) elif gguf_path: if not Path(gguf_path).is_file(): raise FileNotFoundError(f"GGUF file not found: {gguf_path}") @@ -629,7 +676,7 @@ class LlamaCppBackend: "--port", str(self._port), "-c", - str(n_ctx), + "0", # 0 = use model's native context size "--parallel", "1", # Single-user studio, saves VRAM "--flash-attn", @@ -892,7 +939,7 @@ class LlamaCppBackend: top_k: int = 40, min_p: float = 0.0, max_tokens: Optional[int] = None, - repetition_penalty: float = 1.1, + repetition_penalty: float = 1.0, stop: Optional[list[str]] = None, cancel_event: Optional[threading.Event] = None, ) -> Generator[str, None, None]: diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 3c87d64b12..42de014a8d 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -78,7 +78,8 @@ class InferenceOrchestrator: self._static_models = get_default_models() self._top_gguf_cache: Optional[list[str]] = None - self._top_gguf_fetched = False + self._top_hub_cache: Optional[list[str]] = None + self._top_models_ready = threading.Event() # Version tracking for subprocess reuse self._current_transformers_major: Optional[str] = None # "4" or "5" @@ -86,9 +87,9 @@ class InferenceOrchestrator: atexit.register(self._cleanup) logger.info("InferenceOrchestrator initialized (subprocess mode)") - # Kick off background fetch of top GGUF models + # Kick off background fetch of top models from HF threading.Thread( - target = self._fetch_top_gguf, daemon = True, name = "top-gguf" + target = self._fetch_top_models, daemon = True, name = "top-models" ).start() # ------------------------------------------------------------------ @@ -97,12 +98,23 @@ class InferenceOrchestrator: @property def default_models(self) -> list[str]: - top = self._top_gguf_cache or [] - seen = set(top) - return top + [m for m in self._static_models if m not in seen] + # Wait up to 5s for background HF fetch to finish + self._top_models_ready.wait(timeout = 5) + top_gguf = self._top_gguf_cache or [] + top_hub = self._top_hub_cache or [] + # GGUFs first, then hub models, then static fallbacks. + # Send extras so the frontend still has 4 per category + # after removing already-downloaded models. + result: list[str] = [] + seen: set[str] = set() + for m in top_gguf + top_hub + self._static_models: + if m not in seen: + result.append(m) + seen.add(m) + return result - def _fetch_top_gguf(self) -> None: - """Fetch top 4 GGUF repos from unsloth by downloads (background).""" + def _fetch_top_models(self) -> None: + """Fetch top GGUF and non-GGUF repos from unsloth by downloads.""" try: import httpx @@ -112,22 +124,33 @@ class InferenceOrchestrator: "author": "unsloth", "sort": "downloads", "direction": "-1", - "limit": "40", + "limit": "80", }, timeout = 15, ) if resp.status_code == 200: models = resp.json() + # Top 8 GGUFs (frontend deduplicates against downloaded, + # so we fetch extra to always fill 4 slots) gguf_ids = [ m["id"] for m in models if m.get("id", "").upper().endswith("-GGUF") - ][:4] + ][:8] + # Top 8 non-GGUF hub models + hub_ids = [ + m["id"] + for m in models + if not m.get("id", "").upper().endswith("-GGUF") + ][:8] if gguf_ids: self._top_gguf_cache = gguf_ids logger.info("Top GGUF models: %s", gguf_ids) + if hub_ids: + self._top_hub_cache = hub_ids + logger.info("Top hub models: %s", hub_ids) except Exception as e: - logger.warning("Failed to fetch top GGUF models: %s", e) + logger.warning("Failed to fetch top models: %s", e) finally: - self._top_gguf_fetched = True + self._top_models_ready.set() # ------------------------------------------------------------------ # Subprocess lifecycle @@ -395,7 +418,7 @@ class InferenceOrchestrator: top_k: int = 40, min_p: float = 0.0, max_new_tokens: int = 256, - repetition_penalty: float = 1.1, + repetition_penalty: float = 1.0, cancel_event = None, use_adapter = None, ) -> Generator[str, None, None]: @@ -666,7 +689,7 @@ class InferenceOrchestrator: top_k: int = 40, min_p: float = 0.0, max_new_tokens: int = 256, - repetition_penalty: float = 1.1, + repetition_penalty: float = 1.0, cancel_event = None, ) -> Generator[str, None, None]: """Generate response, streaming tokens from subprocess.""" @@ -712,7 +735,7 @@ class InferenceOrchestrator: top_k: int = 40, min_p: float = 0.0, max_new_tokens: int = 256, - repetition_penalty: float = 1.1, + repetition_penalty: float = 1.0, cancel_event = None, use_adapter = None, ) -> Generator[str, None, None]: @@ -763,7 +786,7 @@ class InferenceOrchestrator: top_k: int = 40, min_p: float = 0.0, max_new_tokens: int = 256, - repetition_penalty: float = 1.1, + repetition_penalty: float = 1.0, cancel_event = None, use_adapter = None, ) -> Generator[str, None, None]: @@ -862,7 +885,7 @@ class InferenceOrchestrator: top_k: int = 50, min_p: float = 0.0, max_new_tokens: int = 2048, - repetition_penalty: float = 1.1, + repetition_penalty: float = 1.0, use_adapter: Optional[Union[bool, str]] = None, ) -> Tuple[bytes, int]: """Generate TTS audio. Returns (wav_bytes, sample_rate). @@ -949,7 +972,7 @@ class InferenceOrchestrator: top_k: int = 40, min_p: float = 0.0, max_new_tokens: int = 512, - repetition_penalty: float = 1.1, + repetition_penalty: float = 1.0, cancel_event = None, ) -> Generator[str, None, None]: """Audio input generation (e.g. Gemma 3n) — streams text tokens.""" @@ -978,7 +1001,7 @@ class InferenceOrchestrator: top_k: int = 40, min_p: float = 0.0, max_new_tokens: int = 512, - repetition_penalty: float = 1.1, + repetition_penalty: float = 1.0, cancel_event = None, ) -> Generator[str, None, None]: """Shared inner logic for audio input generation (Whisper + ASR).""" diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index b266f5e0d4..013c14b4f0 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -276,7 +276,7 @@ def _handle_generate( "top_k": cmd.get("top_k", 40), "min_p": cmd.get("min_p", 0.0), "max_new_tokens": cmd.get("max_new_tokens", 256), - "repetition_penalty": cmd.get("repetition_penalty", 1.1), + "repetition_penalty": cmd.get("repetition_penalty", 1.0), "cancel_event": cancel_event, } @@ -348,7 +348,7 @@ def _handle_generate_audio( top_k = cmd.get("top_k", 50), min_p = cmd.get("min_p", 0.0), max_new_tokens = cmd.get("max_new_tokens", 2048), - repetition_penalty = cmd.get("repetition_penalty", 1.1), + repetition_penalty = cmd.get("repetition_penalty", 1.0), use_adapter = cmd.get("use_adapter"), ) @@ -411,7 +411,7 @@ def _handle_generate_audio_input( top_k = cmd.get("top_k", 40), min_p = cmd.get("min_p", 0.0), max_new_tokens = cmd.get("max_new_tokens", 512), - repetition_penalty = cmd.get("repetition_penalty", 1.1), + repetition_penalty = cmd.get("repetition_penalty", 1.0), cancel_event = cancel_event, ) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index cb0c26ca13..620304f37a 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -91,7 +91,7 @@ class GenerateRequest(BaseModel): 2048, ge = 1, le = 4096, description = "Maximum tokens to generate" ) repetition_penalty: float = Field( - 1.1, ge = 1.0, le = 2.0, description = "Repetition penalty" + 1.0, ge = 1.0, le = 2.0, description = "Repetition penalty" ) image_base64: Optional[str] = Field( None, description = "Base64 encoded image for vision models" diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b6f82bfe08..13bb35082e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -252,8 +252,10 @@ async def load_model( except Exception as e: logger.warning(f"Could not read adapter_config.json: {e}") - # Load the model - success = backend.load_model( + # Load the model in a thread so the event loop stays free + # for download progress polling and other requests. + success = await asyncio.to_thread( + backend.load_model, config = config, max_seq_length = request.max_seq_length, load_in_4bit = load_in_4bit, @@ -302,7 +304,17 @@ async def load_model( raise except Exception as e: logger.error(f"Error loading model: {e}", exc_info = True) - raise HTTPException(status_code = 500, detail = f"Failed to load model: {str(e)}") + msg = str(e) + # Surface a friendlier message for models that Unsloth cannot load + not_supported_hints = [ + "No config file found", + "not yet supported", + "is not supported", + "does not support", + ] + if any(h.lower() in msg.lower() for h in not_supported_hints): + msg = f"This model is not supported yet. Try a different model. (Original error: {msg})" + raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}") @router.post("/validate", response_model = ValidateModelResponse) @@ -874,6 +886,25 @@ async def openai_chat_completions( detail = "Image provided but current GGUF model does not support vision.", ) + # Convert image to PNG for llama-server (stb_image has limited format support) + if image_b64: + try: + import base64 as _b64 + from io import BytesIO as _BytesIO + from PIL import Image as _Image + + raw = _b64.b64decode(image_b64) + img = _Image.open(_BytesIO(raw)) + if img.mode == "RGBA": + img = img.convert("RGB") + buf = _BytesIO() + img.save(buf, format = "PNG") + image_b64 = _b64.b64encode(buf.getvalue()).decode("ascii") + except Exception as e: + raise HTTPException( + status_code = 400, detail = f"Failed to process image: {e}" + ) + # Build message list with system prompt prepended gguf_messages = [] if system_prompt: diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 3b79c40651..fa0daea7fd 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -702,6 +702,95 @@ async def get_gguf_download_progress( return {"downloaded_bytes": 0, "expected_bytes": expected_bytes, "progress": 0} +@router.get("/download-progress") +async def get_download_progress( + repo_id: str = Query(..., description = "HuggingFace repo ID"), + current_subject: str = Depends(get_current_subject), +): + """Return download progress for any HuggingFace model repo. + + Checks the local HF cache for completed blobs and in-progress + (.incomplete) downloads. Uses the HF API to determine the expected + total size on the first call, then caches it for subsequent polls. + """ + _empty = {"downloaded_bytes": 0, "expected_bytes": 0, "progress": 0} + try: + if not _is_valid_repo_id(repo_id): + return _empty + + from huggingface_hub import constants as hf_constants + + cache_dir = Path(hf_constants.HF_HUB_CACHE) + target = f"models--{repo_id.replace('/', '--')}".lower() + completed_bytes = 0 + in_progress_bytes = 0 + + for entry in cache_dir.iterdir(): + if entry.name.lower() != target: + continue + blobs_dir = entry / "blobs" + if not blobs_dir.is_dir(): + break + for f in blobs_dir.iterdir(): + if not f.is_file(): + continue + if f.name.endswith(".incomplete"): + in_progress_bytes += f.stat().st_size + else: + completed_bytes += f.stat().st_size + break + + downloaded_bytes = completed_bytes + in_progress_bytes + if downloaded_bytes == 0: + return _empty + + # Get expected size from HF API (cached per repo_id) + expected_bytes = _get_repo_size_cached(repo_id) + if expected_bytes <= 0: + # Cannot determine total; report bytes only, no percentage + return { + "downloaded_bytes": downloaded_bytes, + "expected_bytes": 0, + "progress": 0, + } + + # Use 95% threshold for completion (blob deduplication can make + # completed_bytes differ slightly from expected_bytes). + # Do NOT use "no .incomplete files" as a completion signal -- + # HF downloads files sequentially, so between files there are + # no .incomplete files even though the download is far from done. + if completed_bytes >= expected_bytes * 0.95: + progress = 1.0 + else: + progress = min(downloaded_bytes / expected_bytes, 0.99) + return { + "downloaded_bytes": downloaded_bytes, + "expected_bytes": expected_bytes, + "progress": round(progress, 3), + } + except Exception as e: + logger.warning(f"Error checking download progress for {repo_id}: {e}") + return _empty + + +_repo_size_cache: dict[str, int] = {} + + +def _get_repo_size_cached(repo_id: str) -> int: + if repo_id in _repo_size_cache: + return _repo_size_cache[repo_id] + try: + from huggingface_hub import model_info as hf_model_info + + info = hf_model_info(repo_id, token = None, files_metadata = True) + total = sum(s.size for s in info.siblings if s.size) + _repo_size_cache[repo_id] = total + return total + except Exception as e: + logger.warning(f"Failed to get repo size for {repo_id}: {e}") + return 0 + + @router.get("/cached-gguf") async def list_cached_gguf( current_subject: str = Depends(get_current_subject), @@ -749,6 +838,41 @@ async def list_cached_gguf( return {"cached": []} +@router.get("/cached-models") +async def list_cached_models( + current_subject: str = Depends(get_current_subject), +): + """List non-GGUF model repos that have been downloaded to the HF cache.""" + try: + from huggingface_hub import scan_cache_dir + + hf_cache = scan_cache_dir() + seen_lower: dict[str, dict] = {} + for repo_info in hf_cache.repos: + if repo_info.repo_type != "model": + continue + repo_id = repo_info.repo_id + if repo_id.upper().endswith("-GGUF"): + continue + total_size = sum( + f.size_on_disk for rev in repo_info.revisions for f in rev.files + ) + if total_size == 0: + continue + key = repo_id.lower() + existing = seen_lower.get(key) + if existing is None or total_size > existing["size_bytes"]: + seen_lower[key] = { + "repo_id": repo_id, + "size_bytes": total_size, + } + cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) + return {"cached": cached} + except Exception as e: + logger.error(f"Error listing cached models: {e}", exc_info = True) + return {"cached": []} + + @router.get("/checkpoints", response_model = CheckpointListResponse) async def list_checkpoints( outputs_dir: str = Query( diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 8df53a28c9..aaf15994be 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -473,10 +473,10 @@ try: model_type = getattr(config, "model_type", "unknown") archs = getattr(config, "architectures", []) - logger.info(json.dumps({"is_vision": is_vlm, "model_type": model_type, + print(json.dumps({"is_vision": is_vlm, "model_type": model_type, "architectures": archs})) except Exception as exc: - logger.info(json.dumps({"error": str(exc)})) + print(json.dumps({"error": str(exc)})) sys.exit(1) """ 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 fe4226c4af..c8f8881808 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -8,8 +8,8 @@ import { TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { listCachedGguf, listGgufVariants } from "@/features/chat/api/chat-api"; -import type { CachedGgufRepo } from "@/features/chat/api/chat-api"; +import { listCachedGguf, listCachedModels, listGgufVariants } from "@/features/chat/api/chat-api"; +import type { CachedGgufRepo, CachedModelRepo } from "@/features/chat/api/chat-api"; import type { GgufVariantDetail } from "@/features/chat/types/api"; import { usePlatformStore } from "@/config/env"; import { @@ -340,6 +340,10 @@ function isGgufRepo(id: string): boolean { return id.toUpperCase().includes("-GGUF"); } +// Module-level caches so re-mounting the popover shows results instantly +let _cachedGgufCache: CachedGgufRepo[] = []; +let _cachedModelsCache: CachedModelRepo[] = []; + // ── Hub Model Picker ────────────────────────────────────────── export function HubModelPicker({ @@ -361,16 +365,41 @@ export function HubModelPicker({ // Track which GGUF repo is expanded for variant selection const [expandedGguf, setExpandedGguf] = useState(null); - // Cached (already downloaded) GGUF repos - const [cachedGguf, setCachedGguf] = useState([]); + // Cached (already downloaded) repos -- use module-level cache so + // re-mounting the popover does not flash an empty "Downloaded" section. + const [cachedGguf, setCachedGguf] = useState(_cachedGgufCache); + const [cachedModels, setCachedModels] = useState(_cachedModelsCache); + const alreadyCached = _cachedGgufCache.length > 0 || _cachedModelsCache.length > 0; + const [cachedReady, setCachedReady] = useState(alreadyCached); useEffect(() => { - listCachedGguf().then(setCachedGguf).catch(() => {}); - }, []); + if (alreadyCached) return; + let done = 0; + const check = () => { if (++done >= 2) setCachedReady(true); }; + listCachedGguf().then((v) => { _cachedGgufCache = v; setCachedGguf(v); }).catch(() => {}).finally(check); + listCachedModels().then((v) => { _cachedModelsCache = v; setCachedModels(v); }).catch(() => {}).finally(check); + }, [alreadyCached]); - const recommendedIds = useMemo( - () => dedupe([...models.map((model) => model.id), value ?? ""]), - [models, value], - ); + // Deduplicate: don't show downloaded models in the recommended list. + // Compare case-insensitively since HF cache lowercases repo IDs. + const downloadedSet = useMemo(() => { + const s = new Set(); + for (const c of cachedGguf) s.add(c.repo_id.toLowerCase()); + for (const c of cachedModels) s.add(c.repo_id.toLowerCase()); + return s; + }, [cachedGguf, cachedModels]); + + const recommendedIds = useMemo(() => { + const all = dedupe([...models.map((model) => model.id), value ?? ""]) + .filter((id) => !downloadedSet.has(id.toLowerCase())); + // Cap at 4 GGUFs + 4 non-GGUFs so the list stays manageable + const gguf: string[] = []; + const hub: string[] = []; + for (const id of all) { + if (isGgufRepo(id) && gguf.length < 4) gguf.push(id); + else if (!isGgufRepo(id) && hub.length < 4) hub.push(id); + } + return [...gguf, ...hub]; + }, [models, value, downloadedSet]); const { paramCountById: recommendedParamCountById } = useRecommendedModelVram(recommendedIds); @@ -392,8 +421,13 @@ export function HubModelPicker({ () => new Map( results - .filter((result) => result.totalParams) - .map((result) => [result.id, formatCompact(result.totalParams!)]), + .filter((result) => result.totalParams || result.estimatedSizeBytes) + .map((result) => [ + result.id, + result.estimatedSizeBytes + ? `~${formatBytes(result.estimatedSizeBytes)}` + : formatCompact(result.totalParams!), + ]), ), [results], ); @@ -472,9 +506,14 @@ export function HubModelPicker({
- {!showHfSection && cachedGguf.length > 0 ? ( + {!cachedReady && !showHfSection ? ( +
+ + Loading models… +
+ ) : !showHfSection && (cachedGguf.length > 0 || cachedModels.length > 0) ? ( <> - Downloaded + {"\uD83E\uDDA5"} Downloaded {cachedGguf.map((c) => (
))} + {cachedModels.map((c) => ( + onSelect(c.repo_id, { source: "hub", isLora: false, isDownloaded: true })} + vramStatus={null} + /> + ))} ) : null} - {!showHfSection ? ( + {!showHfSection && cachedReady ? ( <> - Recommended + {"\uD83E\uDDA5"} Recommended {recommendedIds.length === 0 ? (
No default models. diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 2a1e309782..0111c2085a 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -139,7 +139,7 @@ const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => { className="size-20" />

- Run LLMs or test your fine-tune + Chat with your model

Run GGUFs, safetensors, vision and audio models! diff --git a/studio/frontend/src/components/ui/sonner.tsx b/studio/frontend/src/components/ui/sonner.tsx index a20f106d42..11dca18213 100644 --- a/studio/frontend/src/components/ui/sonner.tsx +++ b/studio/frontend/src/components/ui/sonner.tsx @@ -19,7 +19,7 @@ const Toaster = ({ ...props }: ToasterProps) => { { toastOptions={{ classNames: { toast: "cn-toast", + closeButton: "!top-1.5 !translate-y-0", }, }} {...props} diff --git a/studio/frontend/src/features/auth/api.ts b/studio/frontend/src/features/auth/api.ts index e2e4932936..3bb2c9139c 100644 --- a/studio/frontend/src/features/auth/api.ts +++ b/studio/frontend/src/features/auth/api.ts @@ -84,7 +84,15 @@ export async function authFetch( headers.set("Authorization", `Bearer ${accessToken}`); } - const response = await fetch(input, { ...init, headers }); + let response: Response; + try { + response = await fetch(input, { ...init, headers }); + } catch (err) { + if (err instanceof TypeError) { + throw new Error("Studio isn't running -- please relaunch it."); + } + throw err; + } if (await isPasswordChangeRequiredResponse(response)) { void redirectToAuth(); return response; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 04cd0a5098..6332753b8a 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -4,7 +4,14 @@ import type { ChatModelAdapter } from "@assistant-ui/react"; import type { MessageTiming } from "@assistant-ui/core"; import { toast } from "sonner"; -import { generateAudio, streamChatCompletions } from "./chat-api"; +import { + generateAudio, + listCachedGguf, + listCachedModels, + listGgufVariants, + loadModel, + streamChatCompletions, +} from "./chat-api"; import { db } from "../db"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; import { @@ -174,17 +181,117 @@ async function resolveUseAdapter( } } +/** Wait for an in-progress model load to finish (polls store every 500ms). */ +function waitForModelReady(abortSignal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + const check = () => { + if (abortSignal?.aborted) { reject(new Error("Aborted")); return; } + if (!useChatRuntimeStore.getState().modelLoading) { resolve(); return; } + setTimeout(check, 500); + }; + check(); + }); +} + +/** + * Auto-load the smallest downloaded model when the user tries to chat + * without selecting one. Prefers GGUF (picks smallest cached variant), + * falls back to smallest cached safetensors model. + */ +async function autoLoadSmallestModel(): Promise { + const toastId = toast("Loading a model…", { + description: "Auto-selecting the smallest downloaded model.", + duration: Infinity, + }); + try { + const [ggufRepos, modelRepos] = await Promise.all([ + listCachedGguf().catch(() => []), + listCachedModels().catch(() => []), + ]); + + // Try GGUF first: pick the repo with the smallest total size, + // then pick its smallest downloaded variant. + if (ggufRepos.length > 0) { + const sorted = [...ggufRepos].sort((a, b) => a.size_bytes - b.size_bytes); + for (const repo of sorted) { + try { + const variants = await listGgufVariants(repo.repo_id); + const downloaded = variants.variants + .filter((v) => v.downloaded) + .sort((a, b) => a.size_bytes - b.size_bytes); + if (downloaded.length > 0) { + const variant = downloaded[0]; + await loadModel({ + model_path: repo.repo_id, + hf_token: null, + max_seq_length: 4096, + load_in_4bit: true, + is_lora: false, + gguf_variant: variant.quant, + trust_remote_code: false, + }); + useChatRuntimeStore.getState().setCheckpoint(repo.repo_id, variant.quant); + toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, { id: toastId }); + return true; + } + } catch { + continue; + } + } + } + + // Fall back to safetensors models + if (modelRepos.length > 0) { + const sorted = [...modelRepos].sort((a, b) => a.size_bytes - b.size_bytes); + for (const repo of sorted) { + try { + await loadModel({ + model_path: repo.repo_id, + hf_token: null, + max_seq_length: 4096, + load_in_4bit: true, + is_lora: false, + gguf_variant: null, + trust_remote_code: false, + }); + useChatRuntimeStore.getState().setCheckpoint(repo.repo_id); + toast.success(`Loaded ${repo.repo_id}`, { id: toastId }); + return true; + } catch { + continue; + } + } + } + + toast.dismiss(toastId); + return false; + } catch { + toast.dismiss(toastId); + return false; + } +} + export function createOpenAIStreamAdapter(): ChatModelAdapter { return { async *run({ messages, abortSignal, unstable_threadId }) { const runtime = useChatRuntimeStore.getState(); const { params } = runtime; - if (!params.checkpoint) { - toast.error("No model loaded", { - description: "Pick model in top bar, then retry.", - }); - throw new Error("Load a model first."); + // Wait for in-progress model load to finish before inferring + if (runtime.modelLoading) { + toast.info("Waiting for model to finish loading…"); + await waitForModelReady(abortSignal); + } + + if (!useChatRuntimeStore.getState().params.checkpoint) { + // Auto-load the smallest downloaded model + const loaded = await autoLoadSmallestModel(); + if (!loaded) { + toast.error("No model loaded", { + description: "Pick a model in the top bar, then retry.", + }); + throw new Error("Load a model first."); + } } const outboundMessages = messages diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index be105e86be..01bc762d86 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -117,12 +117,31 @@ export async function getGgufDownloadProgress( return parseJsonOrThrow(response); } +export async function getDownloadProgress( + repoId: string, +): Promise<{ downloaded_bytes: number; expected_bytes: number; progress: number }> { + const params = new URLSearchParams({ repo_id: repoId }); + const response = await authFetch(`/api/models/download-progress?${params}`); + return parseJsonOrThrow(response); +} + export async function listCachedGguf(): Promise { const response = await authFetch("/api/models/cached-gguf"); const data = await parseJsonOrThrow<{ cached: CachedGgufRepo[] }>(response); return data.cached; } +export interface CachedModelRepo { + repo_id: string; + size_bytes: number; +} + +export async function listCachedModels(): Promise { + const response = await authFetch("/api/models/cached-models"); + const data = await parseJsonOrThrow<{ cached: CachedModelRepo[] }>(response); + return data.cached; +} + export async function listGgufVariants( repoId: string, hfToken?: string, diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 0e06d89b28..b95a3a8643 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -8,7 +8,6 @@ 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, @@ -39,6 +38,7 @@ import { import { toast } from "sonner"; import { GuidedTour, useGuidedTourController } from "@/features/tour"; import { ChatSettingsPanel } from "./chat-settings-sheet"; +import { ModelLoadInlineStatus } from "./components/model-load-status"; import { db } from "./db"; import { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; import { @@ -111,23 +111,6 @@ function messageHasImage(message: MessageRecord): boolean { return false; } -async function resolveActiveSingleThreadId(view: ChatView): Promise { - if (view.mode !== "single") { - return undefined; - } - if (view.threadId) { - return view.threadId; - } - - // New-thread flow keeps threadId undefined in local view state. - // Fall back to most recent regular base thread. - const candidates = await db.threads.where("modelType").equals("base").toArray(); - const latest = candidates - .filter((thread) => !thread.archived && !thread.pairId) - .sort((a, b) => b.createdAt - a.createdAt)[0]; - return latest?.id; -} - const SingleContent = memo(function SingleContent({ threadId, newThreadNonce, @@ -321,7 +304,16 @@ 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, cancelLoading, loadingModel } = + const activeThreadId = useChatRuntimeStore((state) => state.activeThreadId); + const { + refresh, + selectModel, + ejectModel, + cancelLoading, + loadingModel, + loadProgress, + loadToastDismissed, + } = useChatModelRuntime(); const refreshRef = useRef(refresh); const selectModelRef = useRef(selectModel); @@ -343,30 +335,27 @@ export function ChatPage(): ReactElement { const currentVariant = store.activeGgufVariant; if (!value || (value === currentCheckpoint && (meta?.ggufVariant ?? null) === (currentVariant ?? null))) return; void (async () => { - let switchNote: string | undefined; - const activeThreadId = await resolveActiveSingleThreadId(view); - if (activeThreadId) { + let showImageCompatibilityWarning = false; + if (view.mode === "single" && activeThreadId) { const thread = await db.threads.get(activeThreadId); if (thread?.modelId && thread.modelId !== value) { const messages = await db.messages .where("threadId") .equals(activeThreadId) .toArray(); - const hasImage = messages.some(messageHasImage); - const targetModel = modelsFromStore.find((model) => model.id === value); - const nonVisionWithImages = hasImage && targetModel?.isVision === false; - - switchNote = nonVisionWithImages - ? "Full chat history will be sent to the new model. This chat has images; text-only models may fail." - : hasImage - ? "Full chat history will be sent to the new model. This chat includes images." - : "Full chat history will be sent to the new model."; + if (messages.length > 0) { + const hasImage = messages.some(messageHasImage); + const targetModel = modelsFromStore.find((model) => model.id === value); + showImageCompatibilityWarning = + hasImage && targetModel?.isVision === false; + } } } - if (switchNote) { - toast.warning("Model changed for this chat", { - description: switchNote, + if (showImageCompatibilityWarning) { + toast.warning("Selected model may not handle earlier images", { + description: + "This chat already includes images. Text-only models can ignore them or fail on follow-up replies.", duration: 6000, }); } @@ -379,13 +368,16 @@ export function ChatPage(): ReactElement { }); })(); }, - [modelsFromStore, selectModel, view], + [activeThreadId, modelsFromStore, selectModel, view], ); const handleEject = useCallback(() => { void ejectModel(); }, [ejectModel]); const handleNewThread = useCallback( - () => setView({ mode: "single", newThreadNonce: crypto.randomUUID() }), + () => { + useChatRuntimeStore.getState().setActiveThreadId(null); + setView({ mode: "single", newThreadNonce: crypto.randomUUID() }); + }, [], ); const handleNewCompare = useCallback( @@ -606,25 +598,22 @@ export function ChatPage(): ReactElement { contentDataTour="chat-model-selector-popover" className="max-w-[62vw] sm:max-w-none" /> - {loadingModel ? ( -

- - - {loadingModel.isDownloaded ? "Loading model…" : "Downloading model…"} - - -
+ progressPercent={loadProgress?.percent} + progressLabel={loadProgress?.label} + onStop={cancelLoading} + /> ) : null}
{modelsError && ( diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 44ebe38267..bb2b05f9da 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -26,6 +26,7 @@ import { DEFAULT_INFERENCE_PARAMS, type InferenceParams, } from "./types/runtime"; +import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import { Switch } from "@/components/ui/switch"; export const defaultInferenceParams = DEFAULT_INFERENCE_PARAMS; @@ -45,7 +46,7 @@ const BUILTIN_PRESETS: Preset[] = [ temperature: 1.2, topP: 0.95, topK: 80, - repetitionPenalty: 1.05, + repetitionPenalty: 1.0, }, }, { @@ -55,7 +56,7 @@ const BUILTIN_PRESETS: Preset[] = [ temperature: 0.2, topP: 0.7, topK: 20, - repetitionPenalty: 1.2, + repetitionPenalty: 1.0, }, }, ]; @@ -67,6 +68,7 @@ function ParamSlider({ max, step, onChange, + displayValue, }: { label: string; value: number; @@ -74,13 +76,14 @@ function ParamSlider({ max: number; step: number; onChange: (v: number) => void; + displayValue?: string; }) { return (
{label} - {value} + {displayValue ?? value}
s.activeGgufVariant) != null; const [presets, setPresets] = useState(BUILTIN_PRESETS); const [activePreset, setActivePreset] = useState("Default"); const isBuiltinPreset = BUILTIN_PRESETS.some((p) => p.name === activePreset); @@ -279,7 +283,7 @@ export function ChatSettingsPanel({
- + {!isGguf && ( + + )} = 131072 ? "Max" : undefined} />
diff --git a/studio/frontend/src/features/chat/components/model-load-status.tsx b/studio/frontend/src/features/chat/components/model-load-status.tsx new file mode 100644 index 0000000000..19eff054f7 --- /dev/null +++ b/studio/frontend/src/features/chat/components/model-load-status.tsx @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Progress } from "@/components/ui/progress"; +import { Spinner } from "@/components/ui/spinner"; +import { Button } from "@/components/ui/button"; + +type ModelLoadDescriptionProps = { + message?: string | null; + progressPercent?: number | null; + progressLabel?: string | null; + onStop?: () => void; +}; + +function clampProgress(value: number): number { + return Math.max(0, Math.min(100, value)); +} + +export function ModelLoadDescription({ + message, + progressPercent, + progressLabel, + onStop, +}: ModelLoadDescriptionProps) { + const hasProgress = typeof progressPercent === "number"; + + return ( +
+
+ {hasProgress ? ( +
+
+ {progressLabel} + {Math.round(clampProgress(progressPercent))}% +
+ +
+ ) : message ? ( +

{message}

+ ) : null} +
+ {onStop ? ( + + ) : null} +
+ ); +} + +type ModelLoadInlineStatusProps = { + label: string; + title: string; + progressPercent?: number | null; + progressLabel?: string | null; + onStop?: () => void; +}; + +export function ModelLoadInlineStatus({ + label, + title, + progressPercent, + progressLabel, + onStop, +}: ModelLoadInlineStatusProps) { + const hasProgress = typeof progressPercent === "number"; + + return ( +
+
+ + {label} +
+ {hasProgress ? ( +
+
+ +
+
+ {progressLabel} + {Math.round(clampProgress(progressPercent))}% +
+
+ ) : null} + {onStop ? ( + + ) : null} +
+ ); +} 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 ebcbe8e966..47bd0698a8 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,9 +1,12 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { useCallback, useRef, useState } from "react"; +import { createElement, useCallback, useRef, useState } from "react"; import { toast } from "sonner"; +import { Spinner } from "@/components/ui/spinner"; +import { ModelLoadDescription } from "../components/model-load-status"; import { + getDownloadProgress, getGgufDownloadProgress, getInferenceStatus, listLoras, @@ -29,6 +32,15 @@ type SelectedModelInput = { expectedBytes?: number; }; +const MODEL_LOAD_TOAST_CLASSNAMES = { + toast: "items-start gap-2.5 pr-8", + content: "gap-0.5", + title: "leading-5", + description: "mt-0", + closeButton: + "!left-auto !right-1.5 !top-1.5 !translate-x-0 !translate-y-0 !border-transparent !bg-transparent !shadow-none hover:!bg-transparent hover:opacity-70", +} as const; + const LORA_SUFFIX_RE = /_(\d{9,})$/; function parseTrailingEpoch(input: string): number | undefined { @@ -120,9 +132,13 @@ function mergeRecommendedInference( modelId: string, ): InferenceParams { const inference = response.inference; + // GGUF: max tokens = 131072 (effectively unlimited, model decides) + // Non-GGUF: max tokens = 4096 + const defaultMaxTokens = response.is_gguf ? 131072 : 4096; return { ...current, checkpoint: modelId, + maxTokens: defaultMaxTokens, temperature: toFiniteNumber(inference?.temperature) ?? current.temperature, topP: toFiniteNumber(inference?.top_p) ?? current.topP, @@ -151,11 +167,47 @@ export function useChatModelRuntime() { displayName: string; isDownloaded?: boolean; } | null>(null); - const [_loadAbortController, setLoadAbortController] = - useState(null); + const [loadToastDismissed, setLoadToastDismissed] = useState(false); + const [loadProgress, setLoadProgress] = useState<{ + percent: number | null; + label: string | null; + phase: "downloading" | "starting"; + } | null>(null); const loadAbortRef = useRef(null); const loadingModelRef = useRef(null); const loadToastIdRef = useRef(null); + const loadToastDismissedRef = useRef(false); + + const setLoadToastDismissedState = useCallback((dismissed: boolean) => { + loadToastDismissedRef.current = dismissed; + setLoadToastDismissed(dismissed); + }, []); + + const resetLoadingUi = useCallback(() => { + setLoadingModel(null); + setLoadProgress(null); + loadingModelRef.current = null; + loadAbortRef.current = null; + loadToastIdRef.current = null; + setLoadToastDismissedState(false); + useChatRuntimeStore.getState().setModelLoading(false); + }, [setLoadToastDismissedState]); + + const renderLoadDescription = useCallback( + ( + message: string, + progressPercent?: number | null, + progressLabel?: string | null, + onStop?: () => void, + ) => + createElement(ModelLoadDescription, { + message, + progressPercent, + progressLabel, + onStop, + }), + [], + ); const refresh = useCallback(async () => { setModelsError(null); @@ -182,6 +234,26 @@ export function useChatModelRuntime() { } }, [setCheckpoint, setLoras, setModels, setModelsError]); + const cancelLoading = useCallback(() => { + const model = loadingModelRef.current; + if (!model) return; + loadAbortRef.current?.abort(); + loadAbortRef.current = null; + loadingModelRef.current = null; + const tid = loadToastIdRef.current; + loadToastIdRef.current = null; + setLoadingModel(null); + setLoadProgress(null); + setLoadToastDismissedState(false); + clearCheckpoint(); + if (tid != null) toast.dismiss(tid); + toast.info("Stopped loading model", { + description: "The current download may still finish in the background.", + }); + // Fire-and-forget: tell backend to stop, don't block UI + unloadModel({ model_path: model.id }).catch(() => {}); + }, [clearCheckpoint, setLoadToastDismissedState]); + const selectModel = useCallback( async (selection: string | SelectedModelInput) => { const modelId = typeof selection === "string" ? selection : selection.id; @@ -191,6 +263,8 @@ export function useChatModelRuntime() { if (!modelId || (params.checkpoint === modelId && (ggufVariant ?? null) === (currentVariant ?? null))) { return; } + // Prevent duplicate loads if already loading this model + if (loadingModelRef.current?.id === modelId) return; const explicitIsLora = typeof selection === "string" ? undefined : selection.isLora; @@ -217,21 +291,24 @@ export function useChatModelRuntime() { const previousIsLora = previousModel?.isLora ?? (previousLora ? true : false); const loadingDescription = [ - currentCheckpoint ? "Unloading previous model first." : null, + currentCheckpoint ? "Switching models." : null, extraLoadingDescription ?? null, - isDownloaded - ? "Loading cached model into memory." - : "Downloading and loading model. Large models can take a while.", + isDownloaded ? "Loading cached model into memory." : null, ] .filter(Boolean) .join(" "); - setModelsError(null); + setLoadToastDismissedState(false); const loadInfo = { id: modelId, displayName, isDownloaded }; setLoadingModel(loadInfo); + useChatRuntimeStore.getState().setModelLoading(true); + setLoadProgress( + isDownloaded + ? { percent: null, label: null, phase: "starting" } + : { percent: 0, label: "Preparing download", phase: "downloading" }, + ); loadingModelRef.current = loadInfo; const abortCtrl = new AbortController(); - setLoadAbortController(abortCtrl); loadAbortRef.current = abortCtrl; try { async function performLoad(): Promise { @@ -300,114 +377,180 @@ export function useChatModelRuntime() { } } - const toastId = toast.loading( - isDownloaded ? "Loading model…" : "Downloading model…", + const toastId = toast( + isDownloaded ? "Starting model…" : "Downloading model…", { - description: loadingDescription, - duration: 10000, - action: { - label: "Cancel", - onClick: () => { - abortCtrl.abort(); - setLoadingModel(null); - setLoadAbortController(null); - loadingModelRef.current = null; - loadAbortRef.current = null; - loadToastIdRef.current = null; - unloadModel({ model_path: modelId }).catch(() => {}); - clearCheckpoint(); - toast.dismiss(toastId); - toast.info("Model loading cancelled"); - }, + icon: createElement(Spinner, { className: "size-4" }), + description: renderLoadDescription( + loadingDescription, + isDownloaded ? null : 0, + isDownloaded ? null : "Preparing download", + cancelLoading, + ), + duration: Infinity, + closeButton: true, + classNames: MODEL_LOAD_TOAST_CLASSNAMES, + onDismiss: (dismissedToast) => { + if (loadToastIdRef.current !== dismissedToast.id) { + return; + } + setLoadToastDismissedState(true); }, }, ); loadToastIdRef.current = toastId; - // Poll download progress for non-cached models + // Poll download progress for non-cached models (GGUF and non-GGUF) let progressInterval: ReturnType | null = null; - if (!isDownloaded && ggufVariant) { + if (!isDownloaded) { const expectedBytes = typeof selection !== "string" ? selection.expectedBytes ?? 0 : 0; - if (expectedBytes > 0) { - progressInterval = setInterval(async () => { - if (abortCtrl.signal.aborted) { - if (progressInterval) clearInterval(progressInterval); - return; - } - try { - const prog = await getGgufDownloadProgress(modelId, ggufVariant ?? "", expectedBytes); - if (prog.progress > 0 && prog.progress < 1) { - const dlGb = prog.downloaded_bytes / (1024 ** 3); - const totalGb = prog.expected_bytes / (1024 ** 3); - const pct = Math.round(prog.progress * 100); - toast.loading( - `Downloading model… ${pct}%`, - { - id: toastId, - description: `${dlGb.toFixed(1)} / ${totalGb.toFixed(1)} GB`, - duration: 10000, - action: { - label: "Cancel", - onClick: () => { - abortCtrl.abort(); - setLoadingModel(null); - setLoadAbortController(null); - loadingModelRef.current = null; - loadAbortRef.current = null; - loadToastIdRef.current = null; - unloadModel({ model_path: modelId }).catch(() => {}); - clearCheckpoint(); - toast.dismiss(toastId); - toast.info("Model loading cancelled"); - }, - }, - }, - ); - } else if (prog.progress >= 1) { - toast.loading("Loading model…", { + let hasShownProgress = false; + + const pollProgress = async () => { + if (abortCtrl.signal.aborted || !loadingModelRef.current) { + if (progressInterval) clearInterval(progressInterval); + return; + } + try { + const prog = ggufVariant && expectedBytes > 0 + ? await getGgufDownloadProgress(modelId, ggufVariant, expectedBytes) + : await getDownloadProgress(modelId); + + if (!loadingModelRef.current) return; + + if (prog.progress > 0 && prog.progress < 1) { + hasShownProgress = true; + const dlGb = prog.downloaded_bytes / (1024 ** 3); + const totalGb = prog.expected_bytes / (1024 ** 3); + const pct = Math.round(prog.progress * 100); + const progressLabel = totalGb > 0 + ? `${dlGb.toFixed(1)} of ${totalGb.toFixed(1)} GB` + : `${dlGb.toFixed(1)} GB downloaded`; + setLoadProgress({ + percent: pct, + label: progressLabel, + phase: "downloading", + }); + if (loadToastDismissedRef.current) return; + toast( + "Downloading model…", + { id: toastId, - description: "Download complete. Starting inference server…", - duration: 10000, - }); + icon: createElement(Spinner, { className: "size-4" }), + description: renderLoadDescription( + loadingDescription, + pct, + progressLabel, + cancelLoading, + ), + duration: Infinity, + closeButton: true, + classNames: MODEL_LOAD_TOAST_CLASSNAMES, + onDismiss: (dismissedToast) => { + if (loadToastIdRef.current !== dismissedToast.id) return; + setLoadToastDismissedState(true); + }, + }, + ); + } else if (prog.downloaded_bytes > 0 && prog.expected_bytes === 0 && prog.progress === 0) { + hasShownProgress = true; + const dlGb = prog.downloaded_bytes / (1024 ** 3); + setLoadProgress({ + percent: null, + label: `${dlGb.toFixed(1)} GB downloaded`, + phase: "downloading", + }); + } else if (prog.progress >= 1 && hasShownProgress) { + setLoadProgress({ + percent: 100, + label: "Download complete", + phase: "starting", + }); + if (loadToastDismissedRef.current) { if (progressInterval) clearInterval(progressInterval); + return; } - } catch { - // Ignore polling errors + toast("Starting model…", { + id: toastId, + icon: createElement(Spinner, { className: "size-4" }), + description: renderLoadDescription( + "Download complete. Loading the model into memory.", + 100, + "Download complete", + cancelLoading, + ), + duration: Infinity, + closeButton: true, + classNames: MODEL_LOAD_TOAST_CLASSNAMES, + onDismiss: (dismissedToast) => { + if (loadToastIdRef.current !== dismissedToast.id) return; + setLoadToastDismissedState(true); + }, + }); + if (progressInterval) clearInterval(progressInterval); } - }, 2000); - } + } catch { + // Ignore polling errors + } + }; + + setTimeout(pollProgress, 500); + progressInterval = setInterval(pollProgress, 2000); } try { await performLoad(); - toast.success(`${displayName} loaded`, { id: toastId }); + if (loadToastDismissedRef.current) { + toast.success(`${displayName} loaded`); + } else { + toast.success(`${displayName} loaded`, { + id: toastId, + description: undefined, + closeButton: false, + duration: 2000, + }); + } } catch (err) { if (!abortCtrl.signal.aborted) { - toast.error( - err instanceof Error ? err.message : "Failed to load model", - { id: toastId }, - ); + const message = + err instanceof Error ? err.message : "Failed to load model"; + if (loadToastDismissedRef.current) { + toast.error(message); + } else { + toast.error(message, { + id: toastId, + description: undefined, + closeButton: false, + duration: 5000, + }); + } } throw err; } finally { if (progressInterval) clearInterval(progressInterval); - setLoadingModel(null); - setLoadAbortController(null); - loadingModelRef.current = null; - loadAbortRef.current = null; - loadToastIdRef.current = null; + resetLoadingUi(); } } catch (error) { if (abortCtrl.signal.aborted) return; // User cancelled, nothing to report - setLoadingModel(null); - loadingModelRef.current = null; + resetLoadingUi(); const message = error instanceof Error ? error.message : "Failed to load model"; setModelsError(message); } }, - [loras, models, params.checkpoint, refresh, setModelsError, setParams], + [ + cancelLoading, + loras, + models, + params.checkpoint, + refresh, + renderLoadDescription, + resetLoadingUi, + setLoadToastDismissedState, + setModelsError, + setParams, + ], ); const ejectModel = useCallback(async () => { @@ -436,28 +579,13 @@ export function useChatModelRuntime() { } }, [clearCheckpoint, params.checkpoint, refresh, setModelsError]); - const cancelLoading = useCallback(() => { - const model = loadingModelRef.current; - if (!model) return; - loadAbortRef.current?.abort(); - loadAbortRef.current = null; - loadingModelRef.current = null; - const tid = loadToastIdRef.current; - loadToastIdRef.current = null; - setLoadingModel(null); - setLoadAbortController(null); - clearCheckpoint(); - if (tid != null) toast.dismiss(tid); - toast.info("Model loading cancelled"); - // Fire-and-forget: tell backend to stop, don't block UI - unloadModel({ model_path: model.id }).catch(() => {}); - }, [clearCheckpoint]); - return { refresh, selectModel, ejectModel, cancelLoading, loadingModel, + loadProgress, + loadToastDismissed, }; } diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index dbc631a1d5..f7673c8c47 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -215,7 +215,7 @@ async function generateTitleWithModel(payload: { top_p: 0.9, max_tokens: 24, top_k: 40, - repetition_penalty: 1.05, + repetition_penalty: 1.0, messages: [ { role: "system", @@ -559,6 +559,22 @@ function ThreadNewChatSwitch({ return null; } +function ActiveThreadSync({ + enabled, +}: { enabled: boolean }): ReactElement | null { + const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId); + const setActiveThreadId = useChatRuntimeStore((state) => state.setActiveThreadId); + + useEffect(() => { + if (!enabled) { + return; + } + setActiveThreadId(mainThreadId ?? null); + }, [enabled, mainThreadId, setActiveThreadId]); + + return null; +} + export function ChatRuntimeProvider({ children, modelType = "base", @@ -586,6 +602,7 @@ export function ChatRuntimeProvider({ return ( + {initialThreadId && } {!initialThreadId && newThreadNonce && ( 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 87e162ed56..c9f1753f82 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -43,8 +43,11 @@ type ChatRuntimeStore = { autoTitle: boolean; modelsError: string | null; activeGgufVariant: string | null; + activeThreadId: string | null; pendingAudioBase64: string | null; pendingAudioName: string | null; + modelLoading: boolean; + setModelLoading: (loading: boolean) => void; setParams: (params: InferenceParams) => void; setModels: (models: ChatModelSummary[]) => void; setLoras: (loras: ChatLoraSummary[]) => void; @@ -52,6 +55,7 @@ type ChatRuntimeStore = { setAutoTitle: (enabled: boolean) => void; setModelsError: (error: string | null) => void; setCheckpoint: (modelId: string, ggufVariant?: string | null) => void; + setActiveThreadId: (threadId: string | null) => void; clearCheckpoint: () => void; setPendingAudio: (base64: string, name: string) => void; clearPendingAudio: () => void; @@ -65,8 +69,11 @@ export const useChatRuntimeStore = create((set) => ({ autoTitle: loadBool(AUTO_TITLE_KEY, false), modelsError: null, activeGgufVariant: null, + activeThreadId: null, pendingAudioBase64: null, pendingAudioName: null, + modelLoading: false, + setModelLoading: (loading) => set({ modelLoading: loading }), setParams: (params) => set({ params }), setModels: (models) => set({ models }), setLoras: (loras) => set({ loras }), @@ -94,6 +101,7 @@ export const useChatRuntimeStore = create((set) => ({ }, activeGgufVariant: ggufVariant ?? null, })), + setActiveThreadId: (activeThreadId) => set({ activeThreadId }), clearCheckpoint: () => set((state) => ({ params: { diff --git a/studio/frontend/src/features/chat/types/runtime.ts b/studio/frontend/src/features/chat/types/runtime.ts index 3a5f803de5..a40ec7ceb8 100644 --- a/studio/frontend/src/features/chat/types/runtime.ts +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -20,9 +20,9 @@ export const DEFAULT_INFERENCE_PARAMS: InferenceParams = { topP: 0.9, topK: 50, minP: 0.01, - repetitionPenalty: 1.1, + repetitionPenalty: 1.0, maxSeqLength: 4096, - maxTokens: 2048, + maxTokens: 8192, systemPrompt: "", checkpoint: "", trustRemoteCode: false, diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index a092a895a8..f7b06ab65a 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -11,6 +11,7 @@ export interface HfModelResult { downloads: number; likes: number; totalParams?: number; + estimatedSizeBytes?: number; } const EXCLUDED_TAGS = new Set([ @@ -54,13 +55,33 @@ function withPopularitySort( return fetch(url, init); } +/** Bytes per parameter for each dtype. */ +const DTYPE_BYTES: Record = { + F64: 8, F32: 4, F16: 2, BF16: 2, + I64: 8, I32: 4, I16: 2, I8: 1, U8: 1, + // Quantized types (4-bit) + NF4: 0.5, FP4: 0.5, INT4: 0.5, GPTQ: 0.5, +}; + +function estimateSizeFromDtypes( + params: Record | undefined, +): number | undefined { + if (!params) return undefined; + let total = 0; + for (const [dtype, count] of Object.entries(params)) { + const bpp = DTYPE_BYTES[dtype.toUpperCase()] ?? 2; // default BF16 + total += count * bpp; + } + return total > 0 ? total : undefined; +} + function makeMapModel(excludeGguf: boolean) { return (raw: unknown): HfModelResult | null => { const m = raw as { name: string; downloads: number; likes: number; - safetensors?: { total: number }; + safetensors?: { total: number; parameters?: Record }; tags?: string[]; }; const isEmbedding = m.tags?.some((t) => EMBEDDING_TAGS.has(t)); @@ -75,6 +96,7 @@ function makeMapModel(excludeGguf: boolean) { downloads: m.downloads, likes: m.likes, totalParams: m.safetensors?.total, + estimatedSizeBytes: estimateSizeFromDtypes(m.safetensors?.parameters), }; }; }