From 84a7f048e540c5637dce0286abfa91ce1952f20e Mon Sep 17 00:00:00 2001 From: Unsloth Date: Sun, 26 Jul 2026 04:16:57 -0700 Subject: [PATCH] Stage image and video downloads through the Hub download manager They downloaded inline inside the load, so they had none of the manager's disk preflight, manifest verification, resume or panel progress. Picks now stage as scoped jobs carrying the loader's own file list, then load from a warm cache. --- studio/backend/core/inference/video.py | 57 +++++++++ studio/backend/routes/video.py | 41 ++++++ .../src/features/hub/download-manager/api.ts | 3 + .../download-manager/download-api-adapter.ts | 6 +- .../download-manager-types.ts | 17 +++ .../features/hub/download-manager/index.ts | 5 + .../download-manager/use-staged-download.ts | 118 ++++++++++++++++++ studio/frontend/src/features/images/api.ts | 23 ++++ .../src/features/images/images-page.tsx | 67 +++++++++- studio/frontend/src/features/video/api.ts | 15 +++ .../src/features/video/video-page.tsx | 65 +++++++++- tests/studio/test_model_picker_contracts.py | 18 +++ 12 files changed, 426 insertions(+), 9 deletions(-) create mode 100644 studio/frontend/src/features/hub/download-manager/use-staged-download.ts diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 60a0764a1c..6e4eb14054 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -726,6 +726,63 @@ class VideoBackend: except Exception: # noqa: BLE001 -- progress totals are best-effort only return None + def download_plan( + self, + repo_id: str, + *, + gguf_filename: Optional[str] = None, + base_repo: Optional[str] = None, + family_override: Optional[str] = None, + model_kind: Optional[str] = None, + hf_token: Optional[str] = None, + ) -> dict[str, Any]: + """The repos + exact files this pick needs, for staging through the Hub download + manager. Mirrors the image backend's plan; the file list is the same scoped one + the load itself uses, so nothing extra is pulled. Local paths yield no entries.""" + from huggingface_hub import HfApi + + fam = _detect_load_family(repo_id, gguf_filename, family_override) + kind = resolve_video_model_kind(gguf_filename, model_kind) + base = repo_id if kind == "pipeline" else resolve_video_base_repo(fam, base_repo) + entries: list[dict[str, Any]] = [] + total = 0 + try: + api = HfApi(token = hf_token or None) + if gguf_filename and not Path(repo_id).expanduser().exists(): + info = api.model_info(repo_id, files_metadata = True) + size = sum( + int(s.size or 0) + for s in (info.siblings or []) + if s.rfilename == gguf_filename + ) + total += size + entries.append( + { + "repo_id": repo_id, + "files": [gguf_filename], + "bytes": size, + "gguf_filename": gguf_filename, + } + ) + if base and not Path(base).expanduser().exists(): + info = api.model_info(base, files_metadata = True) + pairs = self._base_download_files(info, kind) + size = sum(s for _, s in pairs) + total += size + if pairs: + entries.append( + { + "repo_id": base, + "files": [name for name, _ in pairs], + "bytes": size, + "gguf_filename": None, + } + ) + except Exception as exc: # noqa: BLE001 -- an unavailable plan falls back to the inline pull + logger.warning("video.download_plan_failed: %s", exc) + return {"entries": [], "total_bytes": 0} + return {"entries": entries, "total_bytes": total} + def _predownload_base( self, base: str, diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 1f1fb65e97..9678dcae74 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -25,6 +25,7 @@ from pydantic import ValidationError from auth.authentication import get_current_subject from loggers import get_logger from models.inference import ( + DiffusionDownloadPlanResponse, GalleryVideo, VideoGalleryListResponse, VideoGenerateProgressResponse, @@ -73,6 +74,46 @@ def _guard_video_load_against_training() -> None: ) +@router.post("/video/download-plan", response_model = DiffusionDownloadPlanResponse) +async def video_download_plan( + request: VideoLoadRequest, current_subject: str = Depends(get_current_subject) +): + """The repos + files this pick needs, so the frontend stages them through the Hub + download manager instead of the load downloading inline. Mirrors /images/download-plan.""" + from core.inference.diffusion import resolve_local_single_file + from core.inference.video import get_video_backend, resolve_video_model_kind + from utils.native_path_leases import redact_native_paths + + backend = get_video_backend() + try: + kind = resolve_video_model_kind(request.gguf_filename, request.model_kind) + if kind == "pipeline" and not request.gguf_filename: + sole = await asyncio.to_thread(resolve_local_single_file, request.model_path) + if sole is not None: + request.gguf_filename = sole + kind = resolve_video_model_kind(sole, None) + await asyncio.to_thread( + backend.validate_load_request, + request.model_path, + gguf_filename = request.gguf_filename, + family_override = request.family_override, + model_kind = kind, + base_repo = request.base_repo, + ) + plan = await asyncio.to_thread( + backend.download_plan, + request.model_path, + gguf_filename = request.gguf_filename, + base_repo = request.base_repo, + family_override = request.family_override, + model_kind = kind, + hf_token = request.hf_token, + ) + return DiffusionDownloadPlanResponse(**plan) + except (ValueError, FileNotFoundError) as exc: + raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc))) + + @router.post("/video/load", response_model = VideoStatusResponse) async def load_video_model( request: VideoLoadRequest, current_subject: str = Depends(get_current_subject) diff --git a/studio/frontend/src/features/hub/download-manager/api.ts b/studio/frontend/src/features/hub/download-manager/api.ts index 2c55b90edd..3ac3bef822 100644 --- a/studio/frontend/src/features/hub/download-manager/api.ts +++ b/studio/frontend/src/features/hub/download-manager/api.ts @@ -222,6 +222,9 @@ export async function startModelDownload(payload: { gguf_variant?: string | null; hf_token?: string | null; use_xet?: boolean; + // A partial-by-design download of `files` only (see DownloadRequest.scopeId). + scope_id?: string | null; + files?: string[]; }): Promise { const { hf_token, ...body } = payload; const headers = { diff --git a/studio/frontend/src/features/hub/download-manager/download-api-adapter.ts b/studio/frontend/src/features/hub/download-manager/download-api-adapter.ts index a6518ec656..14cd2269f2 100644 --- a/studio/frontend/src/features/hub/download-manager/download-api-adapter.ts +++ b/studio/frontend/src/features/hub/download-manager/download-api-adapter.ts @@ -98,7 +98,11 @@ export function apiStart( }) : startModelDownload({ repo_id: req.repoId, - gguf_variant: req.variant, + // A scoped job carries its scope instead of a quant; the backend derives the + // same "@scope" variant this surface already keyed the job under. + gguf_variant: req.scopeId ? null : req.variant, + scope_id: req.scopeId ?? null, + files: req.files, hf_token: hfToken, use_xet: useXet, }); diff --git a/studio/frontend/src/features/hub/download-manager/download-manager-types.ts b/studio/frontend/src/features/hub/download-manager/download-manager-types.ts index 9851734f27..8f45320c83 100644 --- a/studio/frontend/src/features/hub/download-manager/download-manager-types.ts +++ b/studio/frontend/src/features/hub/download-manager/download-manager-types.ts @@ -32,6 +32,23 @@ export interface DownloadRequest { repoId: string; variant: string | null; expectedBytes: number; + /** + * Marks a partial-by-design download of `files` only, for a consumer that reads a + * deliberate subset of a repo (the diffusion loader skips the packaged root single, + * transformer/ shards and fp16 twins). Set `variant` to `scopedVariant(scopeId)` so + * this surface keys the job the same way the backend does. + */ + scopeId?: string | null; + files?: string[]; +} + +/** + * The variant slot a scoped job occupies. Mirrors the backend's `_scope_variant`: no + * GGUF quant label starts with "@", so a scope collides with neither a real variant nor + * the repo's full snapshot. + */ +export function scopedVariant(scopeId: string): string { + return `@${scopeId}`; } export interface JobListeners { diff --git a/studio/frontend/src/features/hub/download-manager/index.ts b/studio/frontend/src/features/hub/download-manager/index.ts index 60ef3851f8..36c312e1cb 100644 --- a/studio/frontend/src/features/hub/download-manager/index.ts +++ b/studio/frontend/src/features/hub/download-manager/index.ts @@ -39,6 +39,11 @@ export { type DownloadJobProgress, type RepoDownloadConfig, } from "./use-repo-download"; +export { + useStagedDownload, + type StagedDownloadEntry, +} from "./use-staged-download"; +export { scopedVariant } from "./download-manager-types"; export { getTransportMode, useDownloadTransportCapabilities, diff --git a/studio/frontend/src/features/hub/download-manager/use-staged-download.ts b/studio/frontend/src/features/hub/download-manager/use-staged-download.ts new file mode 100644 index 0000000000..a5fec6f822 --- /dev/null +++ b/studio/frontend/src/features/hub/download-manager/use-staged-download.ts @@ -0,0 +1,118 @@ +// 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, useEffect, useState } from "react"; + +import { toast } from "@/lib/toast"; + +import { DOWNLOAD_KIND } from "./constants"; +import { downloadManager } from "./download-manager-controller"; +import { scopedVariant } from "./download-manager-types"; +import { useRepoDownload } from "./use-repo-download"; + +/** One repo of a staged plan: the exact files to fetch and their declared size. */ +export interface StagedDownloadEntry { + repoId: string; + files: string[]; + bytes: number; + /** Set when this entry is a single-file GGUF checkpoint rather than a scoped subset. */ + ggufFilename?: string | null; + /** Quant label for a GGUF entry, so the job keys like any other variant download. */ + variant?: string | null; +} + +/** + * Runs a multi-repo download plan through the shared download manager, then calls + * `onReady` once every entry is on disk. + * + * Chat stages a single repo inline in chat-page; the diffusion pages need two (a GGUF + * checkpoint plus its companion base), and their loader reads only part of each, so the + * entries go out as scoped jobs. Staging here rather than letting the backend download + * inside the load is what puts image and video downloads in the same panel, with the same + * progress, cancel, resume, disk preflight and manifest verification as everything else. + */ +export function useStagedDownload({ + scopeId, + onReady, +}: { + /** Scope label for entries that fetch a file subset (e.g. "diffusion"). */ + scopeId: string; + onReady: () => void; +}) { + const [queue, setQueue] = useState(null); + const current = queue?.[0] ?? null; + + // A GGUF entry is a normal variant download; a scoped entry keys itself under "@scope". + const activeVariant = current + ? (current.variant ?? (current.ggufFilename ? null : scopedVariant(scopeId))) + : null; + + const advance = useCallback(() => { + setQueue((rest) => { + const remaining = (rest ?? []).slice(1); + if (remaining.length > 0) return remaining; + return null; + }); + }, []); + + useRepoDownload({ + kind: DOWNLOAD_KIND.MODEL, + repoId: current?.repoId ?? "__staged_download_idle__", + activeVariant, + onComplete: () => { + const remaining = (queue ?? []).slice(1); + advance(); + // Every entry is on disk, so the load will find its cache warm. + if (remaining.length === 0) onReady(); + }, + onError: () => setQueue(null), + onCancelled: () => setQueue(null), + }); + + useEffect(() => { + if (!current) return; + let active = true; + void (async () => { + const outcome = await downloadManager.requestStart({ + kind: DOWNLOAD_KIND.MODEL, + repoId: current.repoId, + variant: activeVariant, + expectedBytes: current.bytes, + scopeId: current.ggufFilename ? null : scopeId, + files: current.ggufFilename ? undefined : current.files, + }); + if (!active) return; + if (outcome === "started") { + toast.info("Downloading model", { + description: "It'll load automatically once the download finishes.", + }); + return; + } + if (outcome === "conflict") { + toast.info("Resume this download from Models", { + description: + "An earlier partial download used a different transport. Open the Model hub tab to resume or restart it.", + }); + setQueue(null); + return; + } + if (outcome === "busy") { + toast.info("Download already in progress", { + description: + "Reselect this model once the running download finishes to load it.", + }); + setQueue(null); + } + })(); + return () => { + active = false; + }; + // Only the head of the queue drives a start; advancing re-runs this with the next one. + }, [current, activeVariant, scopeId]); + + const stage = useCallback((entries: StagedDownloadEntry[]) => { + setQueue(entries.length > 0 ? entries : null); + }, []); + + return { stage, staging: queue !== null }; +} diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 255050a6cb..f360427e68 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -233,6 +233,29 @@ export async function loadDiffusionModel(body: DiffusionLoadRequest): Promise { + return parseJson( + await authFetch("/api/inference/images/download-plan", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + }), + ); +} + export async function generateDiffusionImage( body: DiffusionGenerateRequest, ): Promise { diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 0fa7239226..f198500b1e 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -94,9 +94,11 @@ import { getGenerateProgress, listDiffusionControlNets, listDiffusionLoras, + getDiffusionDownloadPlan, loadDiffusionModel, unloadDiffusionModel, } from "./api"; +import { useStagedDownload } from "@/features/hub/download-manager"; import { DiffusionTrainPanel } from "./train/diffusion-train-panel"; import { TrainBaseSelector, @@ -1756,6 +1758,59 @@ export function ImagesPage({ active = true }: { active?: boolean }) { setMaskResetKey((k) => k + 1); }, []); + // Downloads go through the Hub download manager like every other model, so they share its + // panel, progress, cancel/resume, disk preflight and manifest verification. The load itself + // then finds a warm cache. Held in a ref so the completion callback is not a render dep. + const pendingStagedLoad = useRef<{ + repoId: string; + opts: { kind: "gguf" | "single_file" | "pipeline"; filename?: string }; + } | null>(null); + const handleLoadRef = useRef(handleLoad); + handleLoadRef.current = handleLoad; + const { stage } = useStagedDownload({ + scopeId: "diffusion", + onReady: () => { + const pending = pendingStagedLoad.current; + pendingStagedLoad.current = null; + if (pending) void handleLoadRef.current(pending.repoId, pending.opts); + }, + }); + + // Stage a not-yet-downloaded hub pick, else load it directly. Returns true when the pick + // was accepted either way, so the callers' optimistic picker state stands. + const loadOrStage = useCallback( + async ( + repoId: string, + opts: { kind: "gguf" | "single_file" | "pipeline"; filename?: string }, + isDownloaded?: boolean, + ): Promise => { + if (isDownloaded !== false) return handleLoadRef.current(repoId, opts); + try { + const plan = await getDiffusionDownloadPlan({ + model_path: repoId, + gguf_filename: opts.filename, + model_kind: opts.kind, + }); + if (plan.entries.length > 0) { + pendingStagedLoad.current = { repoId, opts }; + stage( + plan.entries.map((e) => ({ + repoId: e.repo_id, + files: e.files, + bytes: e.bytes, + ggufFilename: e.gguf_filename, + })), + ); + return true; + } + } catch { + // No plan (older backend, metadata hiccup): fall back to the load's own download. + } + return handleLoadRef.current(repoId, opts); + }, + [stage], + ); + // Reload the current model with the current advanced options. const handleReapply = useCallback(() => { const l = lastLoad.current; @@ -1778,7 +1833,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { const d = defaultsFor(id); setSteps(d.steps); setGuidance(d.guidance); - void handleLoad(id, { kind: spec.kind, filename: spec.filename }); + void loadOrStage(id, { kind: spec.kind, filename: spec.filename }, meta.isDownloaded); return; } // GGUF quant pick from the variant expander. Optimistic for instant picker feedback, but @@ -1793,7 +1848,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) { const dq = defaultsFor(id); setSteps(dq.steps); setGuidance(dq.guidance); - void handleLoad(id, { kind: "gguf", filename: meta.ggufFilename }).then((started) => { + void loadOrStage( + id, + { kind: "gguf", filename: meta.ggufFilename }, + meta.isDownloaded, + ).then((started) => { if (!started) { setQuant(prevQuant); quantRevert.current = null; @@ -1870,14 +1929,14 @@ export function ImagesPage({ active = true }: { active?: boolean }) { const d = defaultsFor(id); setSteps(d.steps); setGuidance(d.guidance); - void handleLoad(id, { kind: "pipeline" }).then((started) => { + void loadOrStage(id, { kind: "pipeline" }, meta.isDownloaded).then((started) => { if (!started) { setQuant(prevQuant); quantRevert.current = null; } }); }, - [busy, handleLoad, quant], + [busy, handleLoad, loadOrStage, quant], ); // Deploy a freshly-trained adapter from the Train tab: switch to Create, load the base as diff --git a/studio/frontend/src/features/video/api.ts b/studio/frontend/src/features/video/api.ts index f03b046b6b..613d6f31a0 100644 --- a/studio/frontend/src/features/video/api.ts +++ b/studio/frontend/src/features/video/api.ts @@ -2,6 +2,8 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { authFetch } from "@/features/auth"; +// Same plan shape as the images backend: both /download-plan routes share a response model. +import type { DiffusionDownloadPlan } from "@/features/images/api"; import { readFastApiError } from "@/lib/format-fastapi-error"; // One Advanced control's resolved value + provenance, for the "Auto: X" badges. Same shape the @@ -185,6 +187,19 @@ export async function loadVideoModel(body: VideoLoadRequest): Promise { + return parseJson( + await authFetch("/api/inference/video/download-plan", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + }), + ); +} + /** Start a generation job. Returns as soon as the backend accepts it (the clip takes * minutes, and secure mode's tunnel caps responses near 100s, so the POST cannot span * the generation); poll getVideoGenerateProgress for completion. */ diff --git a/studio/frontend/src/features/video/video-page.tsx b/studio/frontend/src/features/video/video-page.tsx index 668c6ff57f..1ec45ef4d6 100644 --- a/studio/frontend/src/features/video/video-page.tsx +++ b/studio/frontend/src/features/video/video-page.tsx @@ -56,6 +56,7 @@ import { ParamSlider } from "@/features/chat"; import { ModelLoadDescription } from "@/features/chat/components/model-load-status"; import { getHfToken, hfApiToken } from "@/features/hub/stores/hf-token-store"; import { formatBytes, formatEta } from "@/features/hub/lib/format"; +import { useStagedDownload } from "@/features/hub/download-manager"; import { cn } from "@/lib/utils"; import { toast } from "@/lib/toast"; @@ -73,6 +74,7 @@ import { getVideoGallery, getVideoGenerateProgress, getVideoLoadProgress, + getVideoDownloadPlan, getVideoStatus, loadVideoModel, unloadVideoModel, @@ -1086,6 +1088,57 @@ export function VideoPage({ active = true }: { active?: boolean }) { ], ); + // Downloads go through the Hub download manager like every other model, so they share its + // panel, progress, cancel/resume, disk preflight and manifest verification. Mirrors Images. + const pendingStagedLoad = useRef<{ + repoId: string; + opts: { kind: "gguf" | "single_file" | "pipeline"; filename?: string }; + } | null>(null); + const handleLoadRef = useRef(handleLoad); + handleLoadRef.current = handleLoad; + const { stage } = useStagedDownload({ + scopeId: "diffusion", + onReady: () => { + const pending = pendingStagedLoad.current; + pendingStagedLoad.current = null; + if (pending) void handleLoadRef.current(pending.repoId, pending.opts); + }, + }); + + // Stage a not-yet-downloaded hub pick, else load it directly. + const loadOrStage = useCallback( + async ( + repoId: string, + opts: { kind: "gguf" | "single_file" | "pipeline"; filename?: string }, + isDownloaded?: boolean, + ): Promise => { + if (isDownloaded !== false) return handleLoadRef.current(repoId, opts); + try { + const plan = await getVideoDownloadPlan({ + model_path: repoId, + gguf_filename: opts.filename, + model_kind: opts.kind, + }); + if (plan.entries.length > 0) { + pendingStagedLoad.current = { repoId, opts }; + stage( + plan.entries.map((e) => ({ + repoId: e.repo_id, + files: e.files, + bytes: e.bytes, + ggufFilename: e.gguf_filename, + })), + ); + return true; + } + } catch { + // No plan (older backend, metadata hiccup): fall back to the load's own download. + } + return handleLoadRef.current(repoId, opts); + }, + [stage], + ); + // Reload the current model with the current advanced options. const handleReapply = useCallback(() => { const l = lastLoad.current; @@ -1111,7 +1164,7 @@ export function VideoPage({ active = true }: { active?: boolean }) { const d = defaultsFor(spec.filename ? `${id}/${spec.filename}` : id); setSteps(d.steps); setGuidance(d.guidance); - void handleLoad(id, { kind: spec.kind, filename: spec.filename }); + void loadOrStage(id, { kind: spec.kind, filename: spec.filename }, meta.isDownloaded); return; } // GGUF quant pick from the variant expander. Optimistic for instant picker feedback, @@ -1125,7 +1178,11 @@ export function VideoPage({ active = true }: { active?: boolean }) { const dq = defaultsFor(`${id}/${meta.ggufFilename}`); setSteps(dq.steps); setGuidance(dq.guidance); - void handleLoad(id, { kind: "gguf", filename: meta.ggufFilename }).then((started) => { + void loadOrStage( + id, + { kind: "gguf", filename: meta.ggufFilename }, + meta.isDownloaded, + ).then((started) => { if (!started) { setQuant(prevQuant); quantRevert.current = null; @@ -1191,9 +1248,9 @@ export function VideoPage({ active = true }: { active?: boolean }) { const d = defaultsFor(id); setSteps(d.steps); setGuidance(d.guidance); - void handleLoad(id, { kind: "pipeline" }); + void loadOrStage(id, { kind: "pipeline" }, meta.isDownloaded); }, - [busy, handleLoad, quant], + [busy, handleLoad, loadOrStage, quant], ); const handleUnload = useCallback(async () => { diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index d0cef01488..0ae144e61b 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -604,3 +604,21 @@ def test_diffusion_pages_never_drop_a_gguf_pick_silently(): ) assert branch, f"{rel}: gguf extension guard not found" assert "toast.error(" in branch.group(0), f"{rel}: guard returns silently" + + +def test_diffusion_pages_stage_downloads_through_the_manager(): + """Images/Video must not download inside the load: an undownloaded hub pick goes to + the Hub download manager first, so it shares the panel, progress, cancel/resume, + disk preflight and manifest verification with every other model.""" + for rel in ("features/images/images-page.tsx", "features/video/video-page.tsx"): + src = _read(rel) + assert "useStagedDownload" in src, f"{rel}: not wired to the download manager" + # The plan carries the loader's own file scope, so nothing extra is pulled. + assert "DownloadPlan(" in src, f"{rel}: does not fetch a download plan" + stage_fn = re.search(r"const loadOrStage = useCallback\(.*?\n \);", src, re.S) + assert stage_fn, f"{rel}: loadOrStage not found" + body = stage_fn.group(0) + # Already-downloaded (and local) picks must skip staging and load straight away. + assert "isDownloaded !== false" in body, f"{rel}: cached picks would re-stage" + # A missing plan must still load rather than dead-end. + assert "catch" in body, f"{rel}: no fallback when the plan is unavailable"