From 23a71b1c2c8e0512aef2a51b5dadfd90a09ed63e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Jul 2026 12:41:16 +0000 Subject: [PATCH] Close video single-file, training reservation, and image mount-resume gaps Route on-device single-checkpoint video folders through the single_file loader: a bare local .safetensors directory (no model_index.json) is advertised as a pipeline with no filename, so validation rejected it before it could load. Reinterpret the pick as a single_file load of the sole checkpoint, mirroring the image load route. Treat a reserved-but-not-yet-spawned LLM training start as active in is_training_active() so /images/load, /video/load, and /diffusion/start cannot race the reserved run for VRAM during the pre-spawn free window. Mirrors the diffusion training service reservation. Resume an in-flight image generation on the Images page mount: probe generate-progress, re-enter the poll loop, and refresh the gallery on completion so a run started elsewhere is reflected and its saved image appears without a manual refresh. Seed resident image defaults from the resolved base_repo rather than a possibly path-shaped repo_id so the first resident generation uses the right recipe. --- studio/backend/core/training/training.py | 9 +++ studio/backend/routes/video.py | 21 ++++- .../tests/test_training_start_offload.py | 36 +++++++++ studio/backend/tests/test_video_routes.py | 44 ++++++++++- .../src/features/images/images-page.tsx | 76 ++++++++++++++++++- 5 files changed, 179 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index a13fd99d1e..30dda0f7d0 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -1141,6 +1141,15 @@ class TrainingBackend: # training invisibly behind a frozen UI. Cheap enough for per-second polls. self._ensure_pump_alive() with self._lock: + # A start reserved via the compare-and-set guard in start_training but not yet + # spawned (before_spawn frees residents, then GPU auto-selection, then proc.start()) + # is already "active": the load/start guards read this to refuse a concurrent + # /images/load, /video/load, or /diffusion/start, so treating the pre-spawn window + # as idle would let another pipeline race the reserved run for VRAM. Mirrors the + # diffusion training service's reserve()/is_active(). + if self._start_in_progress: + return True + if self._proc is not None and self._proc.is_alive(): return True diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 8c3918b404..54fec4d67f 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -77,13 +77,28 @@ def _guard_video_load_against_training() -> None: async def load_video_model( request: VideoLoadRequest, current_subject: str = Depends(get_current_subject) ): + from core.inference.diffusion import resolve_local_single_file from core.inference.diffusion_device import resolve_diffusion_device_target from core.inference.gpu_arbiter import VIDEO, acquire_for, release - from core.inference.video import get_video_backend + 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: + # Resolve the load kind once (gguf / single_file / pipeline) so validation and the + # load agree; a bad explicit kind raises here -> 400. + kind = resolve_video_model_kind(request.gguf_filename, request.model_kind) + # A local On-Device pick can be a bare single-file .safetensors directory (no + # model_index.json): the scanner advertises it as a text-to-video model, but the + # local picker starts it as a pipeline with no filename, so a pipeline load would + # 400 on the missing model_index.json and the advertised model is unusable. If the + # directory holds exactly one checkpoint, reinterpret the pick as a single_file load + # of it, so validation and the load agree. Mirrors the image load route. + 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) # Validate cheaply BEFORE touching the GPU so an unloadable pick can't evict chat then 400. await asyncio.to_thread( backend.validate_load_request, @@ -91,7 +106,7 @@ async def load_video_model( gguf_filename = request.gguf_filename, base_repo = request.base_repo, family_override = request.family_override, - model_kind = request.model_kind, + model_kind = kind, transformer_quant = request.transformer_quant, text_encoder_quant = request.text_encoder_quant, ) @@ -118,7 +133,7 @@ async def load_video_model( transformer_cache_threshold = request.transformer_cache_threshold, transformer_quant = request.transformer_quant, text_encoder_quant = request.text_encoder_quant, - model_kind = request.model_kind, + model_kind = kind, ) return VideoStatusResponse(**status_dict) except (ValueError, FileNotFoundError) as exc: diff --git a/studio/backend/tests/test_training_start_offload.py b/studio/backend/tests/test_training_start_offload.py index ec7b10a042..d74b6266a0 100644 --- a/studio/backend/tests/test_training_start_offload.py +++ b/studio/backend/tests/test_training_start_offload.py @@ -113,3 +113,39 @@ def test_backend_start_guard_blocks_overlapping_starts(): assert results["second"] is False # The flag is cleared once the winning start returns, so a later start may proceed. assert backend._start_in_progress is False + + +def test_is_training_active_true_during_start_reservation(): + # While start_training holds the compare-and-set reservation but has not yet spawned + # (before_spawn frees residents, then GPU auto-selection, then proc.start()), the LLM + # training run must already read as active: /images/load, /video/load, and + # /diffusion/start all gate on is_training_active(), so an idle reading in this window + # would let another pipeline race the reserved run for the just-freed VRAM. + from core.training.training import TrainingBackend + + backend = TrainingBackend() + # Not reserved yet: idle. + assert backend.is_training_active() is False + + entered = threading.Event() + release = threading.Event() + captured = {} + + def _slow_impl(job_id, *, before_spawn = None, **kwargs): + entered.set() + release.wait(timeout = 5.0) + return True + + backend._start_training_impl = _slow_impl + + t = threading.Thread(target = lambda: backend.start_training("job-a"), daemon = True) + t.start() + assert entered.wait(timeout = 5.0) + # Inside the pre-spawn window: reserved, so active even though no proc/progress is set. + captured["in_window"] = backend.is_training_active() + release.set() + t.join(timeout = 5.0) + + assert captured["in_window"] is True + # Reservation cleared once the start returns; with no live proc it reads idle again. + assert backend.is_training_active() is False diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index a82daf6d6a..275ca2fb6d 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -106,10 +106,17 @@ class _FakeBackend(video_module.VideoBackend): ): # Mirror the real backend's cheap validation so the route's # validate-before-evict ordering is exercised. + from pathlib import Path + kind = (model_kind or ("gguf" if gguf_filename else "pipeline")).lower() if kind in ("gguf", "single_file") and not gguf_filename: raise ValueError("A gguf/single_file load needs the checkpoint filename.") - if kind != "gguf" and not model_path.lower().startswith(("unsloth/", "lightricks/")): + # Non-GGUF loads are gated to unsloth/* repos, the official bases, and local paths + # (the real backend trusts an existing local path via _is_trusted_video_repo). + trusted = model_path.lower().startswith(("unsloth/", "lightricks/")) or ( + Path(model_path).expanduser().exists() + ) + if kind != "gguf" and not trusted: raise ValueError( f"Non-GGUF video loads are limited to unsloth/* repos, the official family " f"base repos, and local paths; '{model_path}' is neither." @@ -378,6 +385,41 @@ def test_load_progress_route(client): assert ready.json()["phase"] == "ready" +def test_load_local_single_file_dir_routes_through_single_file(client, tmp_path): + # An On-Device pick of a local directory named for a video family that holds exactly one + # .safetensors and no model_index.json arrives as a pipeline with no filename. The route + # reinterprets it as a single_file load of the sole checkpoint (mirrors the image route), + # so it is loadable instead of 400ing on the missing model_index.json. + d = tmp_path / "ltx-2.3-local" + d.mkdir() + (d / "ltx-dit.safetensors").write_bytes(b"0") + resp = client.post( + "/api/inference/video/load", + json = {"model_path": str(d), "model_kind": "pipeline"}, + ) + assert resp.status_code == 200 + kwargs = video_module.get_video_backend().last_load_kwargs + assert kwargs["model_kind"] == "single_file" + assert kwargs["gguf_filename"] == "ltx-dit.safetensors" + + +def test_load_local_pipeline_dir_stays_pipeline(client, tmp_path): + # A real diffusers directory (has model_index.json) is left as a pipeline load: + # resolve_local_single_file returns None for it, so the pick is not rewritten. + d = tmp_path / "ltx-2.3-pipeline" + d.mkdir() + (d / "model_index.json").write_text("{}") + (d / "diffusion_pytorch_model.safetensors").write_bytes(b"0") + resp = client.post( + "/api/inference/video/load", + json = {"model_path": str(d), "model_kind": "pipeline"}, + ) + assert resp.status_code == 200 + kwargs = video_module.get_video_backend().last_load_kwargs + assert kwargs["model_kind"] == "pipeline" + assert not kwargs.get("gguf_filename") + + def test_generate_happy_path_persists_and_reports_record(client): client.post( "/api/inference/video/load", diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 7cf2590ca1..7338f23b48 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1450,6 +1450,56 @@ export function ImagesPage({ active = true }: { active?: boolean }) { pollTimer.current = setTimeout(() => void pollLoadProgress(), 1000); }, [dismissLoadToast, refreshStatus]); + // Re-enter the per-step generation poll for a run already in flight on the backend -- + // one this page did not start (another client called /images/generate, or this browser + // reloaded mid-generate). The backend runs generations under a serialising lock and keeps + // reporting progress, so track it here instead of showing a stale idle view. The images + // generate-progress carries only per-step progress (no terminal record), so on completion + // refresh the gallery to merge any image saved after the mount fetch. Kept separate from + // handleGenerate's own loop, which prepends its synchronous results directly. + const resumeGeneratePoll = useCallback(() => { + if (genPollTimer.current) clearInterval(genPollTimer.current); + if (genVisibilityListener.current) + document.removeEventListener("visibilitychange", genVisibilityListener.current); + let pollInFlight = false; + const pollGenerateOnce = async () => { + if (pollInFlight) return; + pollInFlight = true; + try { + const p = await getGenerateProgress(); + if (!p.active) { + if (genPollTimer.current) clearInterval(genPollTimer.current); + genPollTimer.current = null; + if (genVisibilityListener.current) { + document.removeEventListener("visibilitychange", genVisibilityListener.current); + genVisibilityListener.current = null; + } + if (!isMounted.current) return; + setBusy(null); + setGenStep(null); + // The finished run saved its image(s) to the backend gallery; re-fetch the first + // page to merge them (a full replace, so deduped) and resync status. + void loadGallery(); + void refreshStatus(); + return; + } + setGenStep((prev) => { + if (prev && prev.step === p.step && prev.eta_seconds === p.eta_seconds) return prev; + return p; + }); + } catch { + // transient; keep polling + } finally { + pollInFlight = false; + } + }; + genVisibilityListener.current = () => { + if (document.visibilityState === "visible") void pollGenerateOnce(); + }; + document.addEventListener("visibilitychange", genVisibilityListener.current); + genPollTimer.current = setInterval(() => void pollGenerateOnce(), 300); + }, [loadGallery, refreshStatus]); + useEffect(() => { void (async () => { await refreshStatus(); @@ -1468,6 +1518,21 @@ export function ImagesPage({ active = true }: { active?: boolean }) { } catch { // Resume is best-effort; a failed probe just leaves the idle view. } + // A generation started elsewhere -- another client, or this browser before a reload -- + // keeps running on the backend under a serialising lock. Resume tracking it so the page + // reflects the in-flight run (progress bar + disabled Generate) instead of a stale idle + // view, then refreshes the gallery when it finishes to pick up an image saved after the + // mount fetch. Mirrors the video page's mount resume. + try { + const g = await getGenerateProgress(); + if (g.active) { + setBusy("generating"); + setGenStep(g); + resumeGeneratePoll(); + } + } catch { + // Resume is best-effort; a failed probe just leaves the idle view. + } })(); // Stop polling if the page unmounts mid-load / mid-generate, and dismiss the // load toast — its poll loop is gone, so it would otherwise hang forever @@ -1481,7 +1546,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { } dismissLoadToast(); }; - }, [refreshStatus, dismissLoadToast, pollLoadProgress]); + }, [refreshStatus, dismissLoadToast, pollLoadProgress, resumeGeneratePoll]); // Seed the generation sliders to a resident model's recipe when the page discovers one it // did not load itself (left loaded by a prior session or another route). refreshStatus() only @@ -1496,7 +1561,12 @@ export function ImagesPage({ active = true }: { active?: boolean }) { if (lastLoad.current) return; if (seededResident.current === repoId) return; seededResident.current = repoId; - const d = defaultsFor(repoId); + // Seed from the resolved base_repo, not repo_id: a GGUF/single_file resident (or one + // loaded from a bare local path like /models/checkpoint.safetensors) carries a repo_id + // with no family substring, so defaultsFor(repoId) would fall back to the distilled + // few-step/no-CFG recipe and the first resident generation would run with wrong defaults. + // base_repo is the resolved diffusers base (it holds the family), so prefer it when set. + const d = defaultsFor(status?.base_repo ?? repoId); setSteps(d.steps); setGuidance(d.guidance); // Wire "Reapply" to the resident model too, so an advanced-option reload works without @@ -1508,7 +1578,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { if (kind === "pipeline") { lastLoad.current = { repoId, kind }; } - }, [status?.loaded, status?.repo_id, status?.model_kind]); + }, [status?.loaded, status?.repo_id, status?.base_repo, status?.model_kind]); const handleLoad = useCallback( // Resolves true when the background load STARTED (callers may revert