Add download progress bar for non-GGUF models in Chat
Previously only GGUF models showed download progress in Chat. Non-GGUF models (safetensors, bnb quantized, etc.) showed a static message with no progress indication. This adds progress tracking for all model types and fixes several related issues. Backend: - Add /api/models/download-progress endpoint that checks the HF cache blobs directory for completed and .incomplete files. Uses model_info() (cached per repo) to determine expected total size for percentage. - Add /api/models/cached-models endpoint that lists non-GGUF model repos from the HF cache via scan_cache_dir(). - Fix progress stuck at 0.99: when no .incomplete files remain, report 1.0 immediately (blob deduplication can make byte totals mismatch). Frontend: - Remove the ggufVariant gate so download progress polling works for all non-cached models, not just GGUFs. - Use GGUF-specific endpoint when variant + expectedBytes available, otherwise use the general download-progress endpoint. - Fix toast stuck after load: check loadingModelRef.current before and after the async poll to prevent overwriting the success toast. - First poll at 500ms instead of waiting for the 2s interval. - Show downloaded non-GGUF models in the Hub model picker "Downloaded" section alongside GGUFs.
This commit is contained in:
Daniel Han
parent
08b5879101
commit
b84f167d5a
4 changed files with 228 additions and 50 deletions
|
|
@ -686,6 +686,90 @@ 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,
|
||||
}
|
||||
|
||||
# No .incomplete files means download is done regardless of byte totals
|
||||
# (blob deduplication can make completed_bytes differ from expected_bytes)
|
||||
if in_progress_bytes == 0 and completed_bytes > 0:
|
||||
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:
|
||||
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:
|
||||
return 0
|
||||
|
||||
|
||||
@router.get("/cached-gguf")
|
||||
async def list_cached_gguf(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
|
|
@ -733,6 +817,43 @@ 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(
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
@ -361,10 +361,12 @@ export function HubModelPicker({
|
|||
// Track which GGUF repo is expanded for variant selection
|
||||
const [expandedGguf, setExpandedGguf] = useState<string | null>(null);
|
||||
|
||||
// Cached (already downloaded) GGUF repos
|
||||
// Cached (already downloaded) repos
|
||||
const [cachedGguf, setCachedGguf] = useState<CachedGgufRepo[]>([]);
|
||||
const [cachedModels, setCachedModels] = useState<CachedModelRepo[]>([]);
|
||||
useEffect(() => {
|
||||
listCachedGguf().then(setCachedGguf).catch(() => {});
|
||||
listCachedModels().then(setCachedModels).catch(() => {});
|
||||
}, []);
|
||||
|
||||
const recommendedIds = useMemo(
|
||||
|
|
@ -472,7 +474,7 @@ export function HubModelPicker({
|
|||
|
||||
<div ref={scrollRef} className="max-h-64 overflow-y-auto">
|
||||
<div className="p-1">
|
||||
{!showHfSection && cachedGguf.length > 0 ? (
|
||||
{!showHfSection && (cachedGguf.length > 0 || cachedModels.length > 0) ? (
|
||||
<>
|
||||
<ListLabel>Downloaded</ListLabel>
|
||||
{cachedGguf.map((c) => (
|
||||
|
|
@ -489,6 +491,16 @@ export function HubModelPicker({
|
|||
)}
|
||||
</div>
|
||||
))}
|
||||
{cachedModels.map((c) => (
|
||||
<ModelRow
|
||||
key={c.repo_id}
|
||||
label={c.repo_id}
|
||||
meta={formatBytes(c.size_bytes)}
|
||||
selected={value === c.repo_id}
|
||||
onClick={() => onSelect(c.repo_id, { source: "hub", isLora: false, isDownloaded: true })}
|
||||
vramStatus={null}
|
||||
/>
|
||||
))}
|
||||
</>
|
||||
) : null}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<CachedGgufRepo[]> {
|
||||
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<CachedModelRepo[]> {
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
import { useCallback, useRef, useState } from "react";
|
||||
import { toast } from "sonner";
|
||||
import {
|
||||
getDownloadProgress,
|
||||
getGgufDownloadProgress,
|
||||
getInferenceStatus,
|
||||
listLoras,
|
||||
|
|
@ -326,57 +327,82 @@ export function useChatModelRuntime() {
|
|||
|
||||
// Poll download progress for non-cached models
|
||||
let progressInterval: ReturnType<typeof setInterval> | 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…", {
|
||||
|
||||
const cancelAction = {
|
||||
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");
|
||||
},
|
||||
};
|
||||
|
||||
const pollProgress = async () => {
|
||||
// Stop if cancelled or if loading already finished
|
||||
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);
|
||||
|
||||
// Re-check after await -- load may have finished while polling
|
||||
if (!loadingModelRef.current) return;
|
||||
|
||||
if (prog.downloaded_bytes > 0 && 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: "Download complete. Starting inference server…",
|
||||
description: totalGb > 0
|
||||
? `${dlGb.toFixed(1)} / ${totalGb.toFixed(1)} GB`
|
||||
: `${dlGb.toFixed(1)} GB downloaded`,
|
||||
duration: 10000,
|
||||
});
|
||||
if (progressInterval) clearInterval(progressInterval);
|
||||
}
|
||||
} catch {
|
||||
// Ignore polling errors
|
||||
action: cancelAction,
|
||||
},
|
||||
);
|
||||
} else if (prog.downloaded_bytes > 0 && prog.expected_bytes === 0) {
|
||||
const dlGb = prog.downloaded_bytes / (1024 ** 3);
|
||||
toast.loading(
|
||||
"Downloading model...",
|
||||
{
|
||||
id: toastId,
|
||||
description: `${dlGb.toFixed(1)} GB downloaded`,
|
||||
duration: 10000,
|
||||
action: cancelAction,
|
||||
},
|
||||
);
|
||||
} else if (prog.progress >= 1) {
|
||||
toast.loading("Loading model...", {
|
||||
id: toastId,
|
||||
description: "Download complete. Loading into memory...",
|
||||
duration: 10000,
|
||||
});
|
||||
if (progressInterval) clearInterval(progressInterval);
|
||||
}
|
||||
}, 2000);
|
||||
}
|
||||
} catch {
|
||||
// Ignore polling errors
|
||||
}
|
||||
};
|
||||
|
||||
// First poll after 500ms, then every 2s
|
||||
setTimeout(pollProgress, 500);
|
||||
progressInterval = setInterval(pollProgress, 2000);
|
||||
}
|
||||
|
||||
try {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue