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"