From 7ec9e77a5de315109ae2506f28b7efada510d6f4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 01:22:54 +0000 Subject: [PATCH 01/10] Studio: fix diffusion install ownership, dataset upload atomicity, gallery pagination, and teardown races install_sd_cpp_prebuilt: only write the .unsloth-studio-owned marker when the install created the target directory or it was empty. Adopting a pre-existing, unowned, non-empty directory (a user's own stable-diffusion.cpp checkout) made it eligible for the uninstaller's recursive delete. routes/training upload: make the multi-file promotion transactional. Back up each displaced original and roll every destination back on any failure, so a mid-loop rename error can no longer partially overwrite the live dataset. routes/training _resolve_dataset_folder: reject a symlinked dataset directory and prove the resolved folder stays under the datasets root, so image read/caption/delete cannot escape the root through a link. routes/training delete: escape glob metacharacters in the thumbnail filename so deleting an image named like [ab].png removes only its own thumbnails. image_gallery / video_gallery listing: filter records against the response schema inside the pager via a valid callback, so offset/limit/has_more all count over accepted records. A leading schema-invalid record no longer returns an empty page with has_more=true and stalls infinite scroll at offset 0. image_gallery / video_gallery save: publish via a temp file plus atomic rename (the sidecar is the video pair's commit marker) and clean up on failure, so a partial write never surfaces a truncated PNG or strands an orphan MP4. diffusion_train_common discovery: treat an empty caption sidecar as a metadata tombstone that still falls through to the dreambooth instance prompt, so clearing every metadata caption no longer fails with no captioned images found. diffusion backend unload: wait for an in-flight denoise to exit before tearing down process-wide patches and state, mirroring the load path. diffusion_engine_router: serialize the whole check/unload/publish transition so a concurrent selection cannot return the engine being unloaded. uninstall.ps1: gate the default sd.cpp process stop on the owner marker so a user's own sd-server is not terminated for a directory we then keep. --- scripts/uninstall.ps1 | 11 +- studio/backend/core/inference/diffusion.py | 23 ++-- .../core/inference/diffusion_engine_router.py | 70 +++++++----- .../backend/core/inference/image_gallery.py | 39 ++++++- .../backend/core/inference/video_gallery.py | 46 ++++++-- .../core/training/diffusion_train_common.py | 18 ++- studio/backend/routes/inference.py | 27 +++-- studio/backend/routes/training.py | 70 +++++++++++- studio/backend/routes/video.py | 26 +++-- .../backend/tests/test_diffusion_backend.py | 56 +++++++++ .../tests/test_diffusion_dataset_api.py | 107 ++++++++++++++++++ .../tests/test_diffusion_engine_router.py | 51 +++++++++ .../tests/test_diffusion_lora_trainer.py | 34 ++++++ studio/backend/tests/test_diffusion_routes.py | 4 +- studio/backend/tests/test_image_gallery.py | 50 ++++++++ studio/backend/tests/test_sd_cpp_install.py | 46 ++++++++ studio/backend/tests/test_video_gallery.py | 53 +++++++++ studio/install_sd_cpp_prebuilt.py | 30 ++++- 18 files changed, 675 insertions(+), 86 deletions(-) diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 742a07c221..a38f0d0eab 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -364,9 +364,18 @@ function Uninstall-UnslothStudio { _StopByPortFile -PortFile (Join-Path $r "share\studio.port") -KnownRoots $knownRoots } _StopStudioProcesses -KnownRoots $knownRoots + # The default sd.cpp dir is only killed+deleted when it carries our owner marker (the deletion + # below is marker-gated). Passing it to the locking-process stop UNCONDITIONALLY would terminate + # a user's own sd-server running from an unowned checkout at this default path -- a process we + # then decide to keep the directory for. Gate the kill on the same predicate so stop and delete + # agree: include it only when marked owned (and present). + $defaultSdCppToStop = $null + if ($defaultSdCpp -and (Test-Path -LiteralPath $defaultSdCpp) -and (Test-Path -LiteralPath (Join-Path $defaultSdCpp ".unsloth-studio-owned") -PathType Leaf)) { + $defaultSdCppToStop = $defaultSdCpp + } # Also stop anything holding a handle on the exact paths we delete (llama-server, # the CLI shim, an mp-fork python with a venv DLL) so the dir delete isn't refused. - _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultSdCpp, $defaultCache, $defaultNode)) + _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache, $defaultNode) + @($defaultSdCppToStop | Where-Object { $_ })) # ── Remove custom-root install trees ── _Step "Removing data and install directories..." diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 47dd293024..9ed4cd007b 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -2588,18 +2588,22 @@ class DiffusionBackend: # Abort an in-flight (lock-free) download so unload/eviction returns promptly. self._cancel_event.set() with self._lock: - # Abort an in-flight denoise via ITS cancel event; the running generate keeps its - # own pipe ref, so freeing _state can't crash it (VRAM reclaimed when it exits). + # Abort an in-flight denoise via ITS cancel event. if self._active_generate_cancel is not None: self._active_generate_cancel.set() - self._unload_locked() # Cancel any in-flight load (its worker checks this token) and drop the marker. self._load_token += 1 self._loading = None - # Barrier: wait for the signalled denoise to exit before reporting unloaded (callers - # treat this return as "VRAM is free"). generate() holds _generate_lock for its full body. + # WAIT for the signalled denoise to exit BEFORE tearing down (mirrors begin_load's + # pre-teardown barrier at the top of the locked load path). _unload_locked uninstalls + # PROCESS-WIDE state the running denoise still depends on -- the eager/arch attention + # patches, the GGUF compile hooks, the flipped backend flags and the compile cache -- none + # of which the denoise's own pipe ref pins. Uninstalling them before the denoise exits would + # corrupt or crash its in-flight forward passes. The denoise holds _generate_lock for its + # whole body, so acquiring it here blocks until it has finished; only then do we tear down. with self._generate_lock: - pass + with self._lock: + self._unload_locked() return self.status() def _unload_locked(self) -> None: @@ -2618,9 +2622,10 @@ class DiffusionBackend: uninstall_patches() uninstall_arch_patches() - # NOTE: deliberately NOT unload_lora_weights() here. unload() acquires _generate_lock only - # AFTER this teardown, so a LoRA-backed denoise may still run for one more callback; mutating - # its adapters now would race it. The whole pipe is dropped below, freeing the adapters with it. + # NOTE: deliberately NOT unload_lora_weights() here. Both callers (unload() and begin_load's + # locked path) now hold _generate_lock across this teardown, so no denoise is in flight; the + # whole pipe is dropped below, freeing any LoRA adapters with it, so a separate adapter unload + # would be redundant work. # Drop the workflow pipes so they don't pin the freed pipeline's modules past unload. self._aux_pipes.clear() # Drop any ControlNet models + pipelines so the freed load carries no extra modules. diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index 17624ab5b8..8422e7305b 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -59,6 +59,10 @@ def _install_accelerator_for(backend: str) -> str: # The engine the current load committed to, and why a non-native choice was made. Mutated only # under _lock during selection. _lock = threading.Lock() +# Serializes a WHOLE engine switch (check -> unload -> publish). _lock alone is released during the +# slow unload(), so two overlapping selections could interleave: one observes the not-yet-updated +# active engine, returns it, and loads onto the very engine the other is concurrently unloading. +_transition_lock = threading.Lock() _active_engine_name: str = ENGINE_DIFFUSERS _fallback_reason: Optional[str] = None @@ -86,37 +90,43 @@ def active_engine_name() -> str: def _activate(name: str, reason: Optional[str]) -> Any: global _active_engine_name, _fallback_reason - # Switching engines: unload the deactivated one first, else its model stays resident but - # unreachable (the evictor only targets the active engine), leaking 10+ GB. The unload is slow, - # so resolve the engine under _lock but run unload() OUTSIDE it (holding _lock across it would - # block every other selection caller). - engine_to_unload = None - old_name = None - with _lock: - if name != _active_engine_name: - engine_to_unload = get_active_diffusion_engine() - old_name = _active_engine_name - else: - # No engine change: publish the (possibly refreshed) fallback reason now. - _fallback_reason = reason if name == ENGINE_DIFFUSERS else None - if engine_to_unload is not None: - # Publish the new engine only AFTER the old one unloads. The evictor unloads - # get_active_diffusion_engine(), so flipping the name first would let a concurrent - # acquire_for evict the new (empty) engine and take the GPU while the old model is still - # freeing VRAM -- two large models briefly resident. Keeping the OLD engine as the evict - # target serializes a concurrent evict on its unload(), granting the GPU only once freed. - try: - engine_to_unload.unload() - except Exception as exc: # noqa: BLE001 -- best-effort; never block the switch - logger.warning("failed to unload previous engine %s: %s", old_name, exc) + # Serialize the ENTIRE check -> unload -> publish transition. _lock is released during the slow + # unload() below (holding it across the unload would block every status/selection reader), which + # opens a window where a second _activate could read the still-old active engine, take the "no + # change" branch, and return that engine -- then load onto it while this call is unloading it. + # _transition_lock closes the window without holding _lock across the unload; the final + # get_active_diffusion_engine() now reflects the committed state. + with _transition_lock: + # Switching engines: unload the deactivated one first, else its model stays resident but + # unreachable (the evictor only targets the active engine), leaking 10+ GB. The unload is + # slow, so resolve the engine under _lock but run unload() OUTSIDE it. + engine_to_unload = None + old_name = None with _lock: - _active_engine_name = name - _fallback_reason = reason if name == ENGINE_DIFFUSERS else None - if name == ENGINE_SD_CPP: - logger.info("diffusion engine: sd_cpp") - else: - logger.info("diffusion engine: diffusers (%s)", reason or "selected") - return get_active_diffusion_engine() + if name != _active_engine_name: + engine_to_unload = get_active_diffusion_engine() + old_name = _active_engine_name + else: + # No engine change: publish the (possibly refreshed) fallback reason now. + _fallback_reason = reason if name == ENGINE_DIFFUSERS else None + if engine_to_unload is not None: + # Publish the new engine only AFTER the old one unloads. The evictor unloads + # get_active_diffusion_engine(), so flipping the name first would let a concurrent + # acquire_for evict the new (empty) engine and take the GPU while the old model is still + # freeing VRAM -- two large models briefly resident. Keeping the OLD engine as the evict + # target serializes a concurrent evict on its unload(), granting the GPU only once freed. + try: + engine_to_unload.unload() + except Exception as exc: # noqa: BLE001 -- best-effort; never block the switch + logger.warning("failed to unload previous engine %s: %s", old_name, exc) + with _lock: + _active_engine_name = name + _fallback_reason = reason if name == ENGINE_DIFFUSERS else None + if name == ENGINE_SD_CPP: + logger.info("diffusion engine: sd_cpp") + else: + logger.info("diffusion engine: diffusers (%s)", reason or "selected") + return get_active_diffusion_engine() def select_and_activate_engine( diff --git a/studio/backend/core/inference/image_gallery.py b/studio/backend/core/inference/image_gallery.py index f3ee2fb00d..8bede83282 100644 --- a/studio/backend/core/inference/image_gallery.py +++ b/studio/backend/core/inference/image_gallery.py @@ -15,8 +15,10 @@ from __future__ import annotations import base64 import json +import os import re import uuid +from collections.abc import Callable from pathlib import Path from typing import Any, Optional @@ -65,7 +67,21 @@ def _png_bytes(image: Any, meta: dict[str, Any]) -> bytes: def save(image: Any, meta: dict[str, Any]) -> dict[str, Any]: """Persist a PIL image with its recipe embedded; return the gallery record.""" image_id = uuid.uuid4().hex - (gallery_dir() / f"{image_id}.png").write_bytes(_png_bytes(image, meta)) + directory = gallery_dir() + final_path = directory / f"{image_id}.png" + # Write to a hidden temp then atomically rename into place: a crash / disk-full mid-write must + # never leave a truncated {id}.png that the listing would surface as a corrupt record. The temp + # name is dotted so an interrupted write is skipped by the *.png glob, and cleaned on failure. + tmp_path = directory / f".{image_id}.png.tmp" + try: + tmp_path.write_bytes(_png_bytes(image, meta)) + os.replace(tmp_path, final_path) + except BaseException: + try: + tmp_path.unlink(missing_ok = True) + except OSError: + pass + raise return _record(image_id, meta) @@ -128,12 +144,24 @@ def _mtime(path: Path) -> float: return 0.0 -def list_images(limit: Optional[int] = None, offset: int = 0) -> list[dict[str, Any]]: +def list_images( + limit: Optional[int] = None, + offset: int = 0, + *, + valid: Optional[Callable[[dict[str, Any]], bool]] = None, +) -> list[dict[str, Any]]: """A newest-first window of images for infinite scroll. Ordered by file mtime (a cheap stat ~= generation order), so a large gallery isn't opened in full just to sort; only the window's recipes are read. limit=None returns everything from - ``offset`` on.""" + ``offset`` on. + + ``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and the + caller's has_more all count over the same accepted-record domain. Passing the route's schema + validator here is essential: a record that has every required key (so ``_read_meta`` accepts + it) but a wrong value type would otherwise be counted in this window yet dropped by the route + after slicing -- a leading bad record then returns an empty page with more remaining, which + stalls infinite scroll at offset 0.""" try: paths = list(gallery_dir().glob("*.png")) except OSError: @@ -150,7 +178,10 @@ def list_images(limit: Optional[int] = None, offset: int = 0) -> list[dict[str, meta = _read_meta(path) if meta is None: # not one of ours (no recipe chunk) continue - records.append(_record(path.stem, meta)) + record = _record(path.stem, meta) + if valid is not None and not valid(record): # present but schema-invalid + continue + records.append(record) if want is not None and len(records) >= want: break return records[offset:] if limit is None else records[offset : offset + limit] diff --git a/studio/backend/core/inference/video_gallery.py b/studio/backend/core/inference/video_gallery.py index db157c0550..b703f44a42 100644 --- a/studio/backend/core/inference/video_gallery.py +++ b/studio/backend/core/inference/video_gallery.py @@ -14,6 +14,7 @@ import json import os import re import uuid +from collections.abc import Callable from pathlib import Path from typing import Any, Optional @@ -34,12 +35,27 @@ def save(mp4_bytes: bytes, meta: dict[str, Any]) -> dict[str, Any]: """Persist encoded MP4 bytes plus their recipe sidecar; return the record.""" video_id = uuid.uuid4().hex directory = gallery_dir() - (directory / f"{video_id}.mp4").write_bytes(mp4_bytes) - # Write the sidecar via tmp + os.replace so a reader never sees a half-written recipe. + mp4_path = directory / f"{video_id}.mp4" + mp4_tmp = directory / f".{video_id}.mp4.tmp" sidecar = directory / f"{video_id}.json" - tmp = directory / f"{video_id}.json.tmp" - tmp.write_text(json.dumps(meta), encoding = "utf-8") - os.replace(tmp, sidecar) + sidecar_tmp = directory / f".{video_id}.json.tmp" + # Stage BOTH files, then publish. The sidecar is the pair's commit marker (list_videos scans + # mp4s but skips any without a readable sidecar), so writing the MP4 straight to its final name + # before the sidecar meant a sidecar write/replace failure left an invisible, undeletable orphan + # MP4 (gallery delete resolves by id, but the record never appears to be deleted). Rename the MP4 + # in first, then the sidecar; on ANY failure remove every artifact so nothing is stranded. + try: + mp4_tmp.write_bytes(mp4_bytes) + sidecar_tmp.write_text(json.dumps(meta), encoding = "utf-8") + os.replace(mp4_tmp, mp4_path) + os.replace(sidecar_tmp, sidecar) + except BaseException: + for path in (mp4_tmp, sidecar_tmp, mp4_path, sidecar): + try: + path.unlink(missing_ok = True) + except OSError: + pass + raise return _record(video_id, meta) @@ -177,11 +193,22 @@ def _mtime(path: Path) -> float: return 0.0 -def list_videos(limit: Optional[int] = None, offset: int = 0) -> list[dict[str, Any]]: +def list_videos( + limit: Optional[int] = None, + offset: int = 0, + *, + valid: Optional[Callable[[dict[str, Any]], bool]] = None, +) -> list[dict[str, Any]]: """A newest-first window of videos for infinite scroll. Ordered by MP4 mtime (a cheap stat ~= generation order); only the window's sidecars are read. - limit=None returns everything from ``offset`` on. A file without its pair is skipped.""" + limit=None returns everything from ``offset`` on. A file without its pair is skipped. + + ``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and the + caller's has_more all count over the same accepted-record domain. Passing the route's schema + validator here keeps a sidecar that parses as JSON but fails the response schema from being + counted in this window yet dropped by the route after slicing -- which would otherwise return a + short or empty page with more remaining and stall infinite scroll.""" try: paths = list(gallery_dir().glob("*.mp4")) except OSError: @@ -195,7 +222,10 @@ def list_videos(limit: Optional[int] = None, offset: int = 0) -> list[dict[str, meta = _read_meta(_sidecar_path(path.stem)) if meta is None: # orphan mp4 (no readable sidecar) continue - records.append(_record(path.stem, meta)) + record = _record(path.stem, meta) + if valid is not None and not valid(record): # parses but schema-invalid + continue + records.append(record) if want is not None and len(records) >= want: break return records[offset:] if limit is None else records[offset : offset + limit] diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 89efb7802b..b4fdf81d1e 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -703,20 +703,30 @@ def discover_image_caption_pairs( pairs: list[tuple[str, str]] = [] for img in images: caption: Optional[str] = None + sidecar_present = False # 1. per-image sidecar caption file (the user's explicit edit; wins over metadata). + # An EMPTY sidecar is a deliberate tombstone (the labeling grid writes one when a user + # clears a metadata caption): it suppresses the metadata caption but leaves the image + # UNCAPTIONED so the dreambooth instance_prompt fallback below still applies. Treating + # "" as a present caption here would skip the instance prompt AND drop the image, so + # clearing every metadata caption in the grid would make a dreambooth run fail with + # "No captioned images found". for ext in _CAPTION_EXTS: sidecar = img.with_suffix(ext) if sidecar.is_file(): + sidecar_present = True caption = sidecar.read_text(encoding = "utf-8").strip() break # 2. metadata row keyed by file name (basename or the relative path; as_posix so a - # Windows backslash path still matches the jsonl's forward-slash keys). - if caption is None: + # Windows backslash path still matches the jsonl's forward-slash keys). Skipped when a + # sidecar tombstone is present (the sidecar, even empty, is authoritative). + if not sidecar_present: caption = meta_caption.get(img.name) or meta_caption.get( img.relative_to(root).as_posix() ) - # 3. dreambooth instance prompt. - if caption is None and instance_prompt: + # 3. dreambooth instance prompt for any image still without a caption (no sidecar text and + # no metadata row, or an empty tombstone). + if not caption and instance_prompt: caption = instance_prompt if caption: if verify_images: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3471b67df1..c452a4775a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -14511,18 +14511,25 @@ async def list_gallery_images( limit = max(1, min(limit, 200)) offset = max(0, offset) - # Fetch one extra to learn whether more remain, without a second scan. - records = await asyncio.to_thread(image_gallery.list_images, limit + 1, offset) - has_more = len(records) > limit - # Drop records that fail schema validation: a PNG whose recipe chunk has all keys but a - # wrong value type (hand-dropped or corrupted) passes the presence-only read yet raises - # inside GalleryImage(**r). Skipping it keeps one bad file from 500-ing the listing. - images = [] - for r in records[:limit]: + + # Validate against the response schema INSIDE the pager so offset / limit / has_more all count + # over the same accepted-record domain. A PNG whose recipe chunk has all keys but a wrong value + # type passes the presence-only read yet fails GalleryImage(**r); dropping such records only + # after pagination made a leading bad record return an empty page with has_more=True, stalling + # infinite scroll at offset 0. Filtering here keeps the window and has_more consistent. + def _valid_gallery_image(record: dict) -> bool: try: - images.append(GalleryImage(**r)) + GalleryImage(**record) except ValidationError: - continue + return False + return True + + # Fetch one extra to learn whether more remain, without a second scan. + records = await asyncio.to_thread( + image_gallery.list_images, limit + 1, offset, valid = _valid_gallery_image + ) + has_more = len(records) > limit + images = [GalleryImage(**r) for r in records[:limit]] return GalleryListResponse(images = images, has_more = has_more) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 1ef804afc2..cb36e43a45 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1732,9 +1732,46 @@ async def upload_diffusion_dataset( ) out.write(chunk) uploaded += 1 - for tmp, dest in staged: - tmp.replace(dest) # atomic on the same filesystem - committed = True + # Commit every staged file as one transaction. A plain replace loop is NOT atomic across + # files: if the second-or-later tmp.replace(dest) fails (disk/quota error, a Windows file + # lock, antivirus, a destination that became a directory), an earlier destination has + # already been overwritten while the request returns an error -- the user's original file + # is gone. Back up each pre-existing destination before overwriting it, then on ANY failure + # remove the versions this request installed and restore every displaced original, so the + # dataset is left exactly as it was before the upload. + backups: list[tuple[Path, Optional[Path]]] = [] # (dest, backup path or None) + installed: list[Path] = [] + try: + for tmp, dest in staged: + backup: Optional[Path] = None + if dest.exists(): + backup = folder / f".upload-backup-{_uuid.uuid4().hex}.part" + dest.replace(backup) + backups.append((dest, backup)) + tmp.replace(dest) # atomic on the same filesystem + installed.append(dest) + committed = True + except BaseException: + # Roll back: drop every new version, then restore every displaced original. + for dest in reversed(installed): + try: + dest.unlink(missing_ok = True) + except OSError: + pass + for dest, backup in reversed(backups): + if backup is not None and backup.exists(): + try: + backup.replace(dest) + except OSError: + pass + raise + else: + for _, backup in backups: + if backup is not None: + try: + backup.unlink(missing_ok = True) + except OSError: + pass finally: if not committed: for tmp, _ in staged: @@ -1766,9 +1803,27 @@ def _resolve_dataset_folder(name: str, *, must_exist: bool = True) -> Path: from utils.paths import datasets_root cleaned = _clean_diffusion_dataset_name(name) - folder = datasets_root() / cleaned + root = datasets_root().resolve() + folder = root / cleaned + # Reject a symlinked dataset directory. _safe_dataset_image_path only proves each image path + # stays under folder.resolve(); it never proves the folder itself stays under the datasets + # root. A dataset dir that is a symlink to an external directory would therefore let image + # read / caption / delete operate on files outside Studio (a reproduced delete removed an + # external file through such a link). Prove the resolved folder is contained in the root too. + if folder.is_symlink(): + raise HTTPException( + status_code = 400, + detail = f"Dataset '{cleaned}' must not be a symbolic link.", + ) if must_exist and not folder.is_dir(): raise HTTPException(status_code = 404, detail = f"Dataset '{cleaned}' not found.") + try: + folder.resolve(strict = must_exist).relative_to(root) + except (OSError, ValueError): + raise HTTPException( + status_code = 400, + detail = f"Dataset '{cleaned}' escapes the Studio datasets directory.", + ) return folder @@ -1995,6 +2050,8 @@ async def delete_diffusion_dataset_image( raise HTTPException(status_code = 404, detail = "Image not found.") def remove() -> dict: + import glob as _glob + image_path.unlink(missing_ok = True) for ext in (".txt", ".caption"): image_path.with_suffix(ext).unlink(missing_ok = True) @@ -2002,7 +2059,10 @@ async def delete_diffusion_dataset_image( if thumbs_dir.is_dir(): # Thumbs are keyed on the full filename (stem + extension), so match that here too; # a stem-only glob would strand this image's thumbs or delete a same-stem sibling's. - for t in thumbs_dir.glob(f"{image_path.name}_*.jpg"): + # Escape the filename first: an uploaded name may legally contain glob metacharacters + # ('[', ']', '*', '?'), and interpolating those raw would make e.g. "[ab].png" match + # "a.png_*.jpg"/"b.png_*.jpg" -- deleting siblings' thumbs while leaving its own behind. + for t in thumbs_dir.glob(f"{_glob.escape(image_path.name)}_*.jpg"): t.unlink(missing_ok = True) return {"deleted": image_path.name} diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 54fec4d67f..04ea6d7d28 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -235,17 +235,25 @@ async def list_gallery_videos( limit = max(1, min(limit, 200)) offset = max(0, offset) - # Fetch one extra to learn whether more remain, without a second scan. - records = await asyncio.to_thread(video_gallery.list_videos, limit + 1, offset) - has_more = len(records) > limit - # Build per record, dropping any that fail schema validation, so one bad sidecar - # (wrong value type) doesn't 500 the whole listing. - videos = [] - for r in records[:limit]: + + # Validate against the response schema INSIDE the pager so offset / limit / has_more all count + # over the same accepted-record domain. A sidecar that parses as JSON but has a wrong value type + # passes the read yet fails GalleryVideo(**r); dropping such records only after pagination made a + # leading bad record return an empty page with has_more=True, stalling infinite scroll at offset + # 0. Filtering here keeps the window and has_more consistent. + def _valid_gallery_video(record: dict) -> bool: try: - videos.append(GalleryVideo(**r)) + GalleryVideo(**record) except ValidationError: - continue + return False + return True + + # Fetch one extra to learn whether more remain, without a second scan. + records = await asyncio.to_thread( + video_gallery.list_videos, limit + 1, offset, valid = _valid_gallery_video + ) + has_more = len(records) > limit + videos = [GalleryVideo(**r) for r in records[:limit]] return VideoGalleryListResponse(videos = videos, has_more = has_more) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 358d6aafeb..0f8b84f083 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -3006,3 +3006,59 @@ def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path): ) assert _FakePipeline.last["base"] == str(tmp_path) backend.unload() + + +def test_unload_waits_for_in_flight_denoise_before_teardown(): + # Regression for the unload/denoise teardown race: unload() must WAIT for a running denoise to + # exit (acquire _generate_lock) BEFORE _unload_locked() tears down PROCESS-WIDE state (eager / + # arch attention patches, gguf compile hooks, backend flags, compile cache) that the denoise + # still depends on. Mirror the load path, which already waits on _generate_lock before teardown. + import threading + + backend = DiffusionBackend() + + denoise_active = {"v": False} + teardown_saw = [] # records denoise_active at the moment _unload_locked runs + + cancel = threading.Event() + backend._active_generate_cancel = cancel + started = threading.Event() + finish = threading.Event() + + # _generate_lock is the only lock a real denoise holds for its whole body. + def _denoise(): + with backend._generate_lock: + denoise_active["v"] = True + started.set() + cancel.wait(2.0) # unload signals this + finish.wait(2.0) # the test lets us finish + denoise_active["v"] = False # about to release _generate_lock + + def _fake_unload_locked(): + teardown_saw.append(denoise_active["v"]) + + backend._unload_locked = _fake_unload_locked # instance attr shadows the method + + d = threading.Thread(target = _denoise) + d.start() + assert started.wait(2.0) # denoise holds _generate_lock + + unloaded = threading.Event() + + def _unload(): + backend.unload() + unloaded.set() + + u = threading.Thread(target = _unload) + u.start() + assert cancel.wait(2.0) # unload has signalled the denoise and is now waiting on _generate_lock + # unload must NOT have torn down yet -- it is blocked on the denoise's _generate_lock. + assert teardown_saw == [] + assert not unloaded.wait(0.3) + + finish.set() # let the denoise release _generate_lock + d.join(2.0) + u.join(2.0) + assert unloaded.is_set() + # Teardown ran exactly once, and only AFTER the denoise had exited (denoise_active was False). + assert teardown_saw == [False] diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index e551899e2d..e80b29d6e6 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -495,6 +495,113 @@ def test_import_promotion_leaves_no_partial_dataset_on_failure(ds_root, monkeypa assert calls["count"] == 2 # the failed attempt did not leave a dataset that blocks a reload +def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch): + # Re-uploading a.txt and b.txt (an allowed overwrite of two files already on disk) stages both, + # then commits them. A plain replace loop is NOT atomic: if the SECOND commit fails after the + # first destination was already overwritten, the live dataset is left partially updated with the + # request returning an error -- the user's original a.txt is gone. The transactional commit must + # back up each displaced original and, on any failure, restore every one so the dataset is left + # exactly as it was, with no stray temp/backup files. + from pathlib import Path + + app = FastAPI() + app.include_router(training_router, prefix = "/api/train") + app.dependency_overrides[get_current_subject] = lambda: "test-user" + noraise = TestClient(app, raise_server_exceptions = False) + + folder = ds_root / "styleset" + folder.mkdir() + (folder / "a.txt").write_bytes(b"ORIGINAL-A") + (folder / "b.txt").write_bytes(b"ORIGINAL-B") + + real_replace = Path.replace + state = {"failed": False} + + def flaky_replace(self, target, *a, **k): + # Fail exactly once, on the tmp -> b.txt promotion (source is the staged .upload-* part, NOT + # a .upload-backup-* part), so the subsequent backup -> b.txt restore still succeeds. + if ( + not state["failed"] + and str(target).endswith("b.txt") + and self.name.startswith(".upload-") + and not self.name.startswith(".upload-backup-") + ): + state["failed"] = True + raise OSError("simulated second commit failure") + return real_replace(self, target, *a, **k) + + monkeypatch.setattr(Path, "replace", flaky_replace) + parts = [ + ("files", ("a.txt", b"NEW-A", "application/octet-stream")), + ("files", ("b.txt", b"NEW-B", "application/octet-stream")), + ] + r = noraise.post("/api/train/diffusion/dataset", data = {"name": "styleset"}, files = parts) + assert r.status_code == 500 + monkeypatch.setattr(Path, "replace", real_replace) + + # Both originals are intact -- no partial overwrite of the live dataset. + assert (folder / "a.txt").read_bytes() == b"ORIGINAL-A" + assert (folder / "b.txt").read_bytes() == b"ORIGINAL-B" + # No staging or backup artifacts left behind. + assert not list(folder.glob(".upload-*.part")) + assert not list(folder.glob(".upload-backup-*.part")) + + +def test_resolve_dataset_folder_rejects_symlink(ds_root, tmp_path): + # A dataset directory that is a symlink pointing OUTSIDE the datasets root must be rejected: the + # per-image containment check only proves paths stay under folder.resolve(), so without this a + # delete/caption/read could operate on external files through the link. + from routes.training import _resolve_dataset_folder + + external = tmp_path / "external" + external.mkdir() + (external / "victim.png").write_bytes(_png_bytes()) + (ds_root / "linked").symlink_to(external, target_is_directory = True) + + with pytest.raises(HTTPException) as exc: + _resolve_dataset_folder("linked") + assert exc.value.status_code == 400 + + +def test_delete_through_symlinked_dataset_cannot_escape_root(client, ds_root, tmp_path): + # End to end: a DELETE against an image inside a symlinked dataset dir is refused (400) and the + # external file it points at is NOT removed. + external = tmp_path / "external" + external.mkdir() + victim = external / "victim.png" + _write_png(victim) + (ds_root / "linked").symlink_to(external, target_is_directory = True) + + r = client.delete("/api/train/diffusion/dataset/linked/image/victim.png") + assert r.status_code == 400 + assert victim.exists() # the external file survives + + +def test_delete_image_with_glob_chars_only_removes_own_thumbs(client, ds_root): + # A legal uploaded filename may contain glob metacharacters ('[', ']', '*', '?'). Deleting it + # must remove ONLY its own thumbnails; interpolating the raw name into Path.glob would make + # "[ab].png" match "a.png_*.jpg"/"b.png_*.jpg" -- deleting a sibling's thumbs while leaving its + # own (literally "[ab].png_32.jpg") behind. + from urllib.parse import quote + + folder = ds_root / "d" + folder.mkdir() + _write_png(folder / "[ab].png") + _write_png(folder / "a.png") + thumbs = folder / ".thumbs" + thumbs.mkdir() + (thumbs / "[ab].png_32.jpg").write_bytes(b"own") + (thumbs / "a.png_32.jpg").write_bytes(b"sibling") + + r = client.delete("/api/train/diffusion/dataset/d/image/" + quote("[ab].png", safe = "")) + assert r.status_code == 200, r.text + # Its own thumbnail is gone; the sibling a.png's thumbnail is untouched. + assert not (thumbs / "[ab].png_32.jpg").exists() + assert (thumbs / "a.png_32.jpg").exists() + assert not (folder / "[ab].png").exists() + assert (folder / "a.png").exists() + + def test_import_preserves_unrelated_files_when_folder_not_empty(client, ds_root, monkeypatch): # If the target folder already holds unrelated NON-image files (so image_count is still 0 and # the import runs), the atomic rmdir refuses and the code falls back to a per-file move: the diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py index 6667d53044..53d40a1336 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -238,3 +238,54 @@ def test_no_switch_keeps_engine_and_refreshes_reason(monkeypatch): assert calls["unload"] == 0 assert r.active_engine_name() == ENGINE_DIFFUSERS assert r.active_status()["fallback_reason"] == "still diffusers" + + +def test_activate_serializes_switch_and_concurrent_query(monkeypatch): + # Regression for the check->unload->publish race: _activate releases _lock during the slow + # unload(), so without the transition lock a second _activate could observe the not-yet-updated + # active engine, take the "no change" branch, and return the engine the first call is + # concurrently unloading. Drive both paths on threads and assert the concurrent query is blocked + # until the switch completes (i.e. the whole transition is serialized). + import threading + + r._active_engine_name = ENGINE_DIFFUSERS + r._fallback_reason = None + + release_unload = threading.Event() + unload_started = threading.Event() + + def _slow_unload(): + unload_started.set() + release_unload.wait(2.0) + + engine = SimpleNamespace( + status = lambda: {"loaded": False, "repo_id": None}, unload = _slow_unload + ) + monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: engine) + + switch_done = threading.Event() + + def _switch(): + r._activate(ENGINE_SD_CPP, None) # diffusers -> sd_cpp: unloads the old engine (blocks) + switch_done.set() + + t = threading.Thread(target = _switch) + t.start() + assert unload_started.wait(2.0) # switch is mid-unload, holding the transition lock + + query_done = threading.Event() + + def _query(): + r._activate(ENGINE_DIFFUSERS, None) # would hit the "no change" branch pre-fix + query_done.set() + + q = threading.Thread(target = _query) + q.start() + # Serialized: while the switch holds the transition lock the query cannot complete. Pre-fix it + # would return immediately (active is still diffusers), setting query_done at once. + assert not query_done.wait(0.4) + + release_unload.set() + t.join(2.0) + q.join(2.0) + assert switch_done.is_set() and query_done.is_set() diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index 92c73e4278..c06d8fd45a 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -57,6 +57,40 @@ def test_discover_sidecar_overrides_metadata_row(tmp_path): assert pairs[str(tmp_path / "a.png")] == "edited sidecar" +def test_discover_empty_sidecar_suppresses_metadata_but_uses_instance_prompt(tmp_path): + # Clearing a metadata caption in the labeling grid writes an EMPTY sidecar tombstone. It must + # suppress the metadata caption (so the stale text is not resurrected) yet leave the image + # UNCAPTIONED so the dreambooth instance_prompt still applies. Treating "" as a present caption + # would skip the instance prompt AND drop the image, so a dataset whose every caption was + # cleared would fail with "No captioned images found". + _touch(tmp_path / "cat.png") + (tmp_path / "metadata.jsonl").write_text( + json.dumps({"file_name": "cat.png", "text": "old metadata caption"}) + "\n", + encoding = "utf-8", + ) + (tmp_path / "cat.txt").write_text("", encoding = "utf-8") # empty tombstone + pairs = discover_image_caption_pairs(tmp_path, instance_prompt = "a photo of sks cat") + assert pairs == [(str(tmp_path / "cat.png"), "a photo of sks cat")] + + +def test_discover_empty_sidecar_without_instance_prompt_skips_image(tmp_path): + # With no instance prompt an empty tombstone leaves the image uncaptioned, so it is skipped and + # the suppressed metadata caption is NOT resurrected -- while a sibling with a real caption is + # still discovered. + _touch(tmp_path / "cat.png") + _touch(tmp_path / "cap.png") + (tmp_path / "metadata.jsonl").write_text( + json.dumps({"file_name": "cat.png", "text": "old"}) + + "\n" + + json.dumps({"file_name": "cap.png", "text": "kept"}) + + "\n", + encoding = "utf-8", + ) + (tmp_path / "cat.txt").write_text("", encoding = "utf-8") # empty tombstone + pairs = dict(discover_image_caption_pairs(tmp_path)) + assert pairs == {str(tmp_path / "cap.png"): "kept"} + + def test_discover_skips_uncaptioned_without_instance_prompt(tmp_path): _touch(tmp_path / "cap.png") _touch(tmp_path / "nocap.png") diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index f913d6e9cb..6e6f7ea2da 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -180,8 +180,10 @@ def client(monkeypatch, tmp_path): monkeypatch.setattr(gallery_module, "save", _save) monkeypatch.setattr(gallery_module, "image_b64", lambda i: "QUJD" if i in store else None) - def _list_images(limit = None, offset = 0): + def _list_images(limit = None, offset = 0, *, valid = None): ordered = sorted(store.values(), key = lambda r: r.get("created_at", 0.0), reverse = True) + if valid is not None: + ordered = [r for r in ordered if valid(r)] return ordered[offset:] if limit is None else ordered[offset : offset + limit] monkeypatch.setattr(gallery_module, "list_images", _list_images) diff --git a/studio/backend/tests/test_image_gallery.py b/studio/backend/tests/test_image_gallery.py index 128a533666..b9b3c0ed49 100644 --- a/studio/backend/tests/test_image_gallery.py +++ b/studio/backend/tests/test_image_gallery.py @@ -145,3 +145,53 @@ def test_list_skips_recipe_missing_required_fields(tmp_path): gallery.save(_img(), _meta(prompt = "ours")) listed = gallery.list_images() assert [r["prompt"] for r in listed] == ["ours"] + + +def test_valid_callback_paginates_over_accepted_records(): + # A record that passes _read_meta (every required key present) but fails the caller's stricter + # schema check must be filtered BEFORE pagination, so offset/limit/has_more all count over the + # accepted domain. Otherwise a leading bad record returns an empty/short page with more still + # remaining, and the frontend (which advances by valid records) stalls at offset 0. + _save_with_mtime("BAD", 300.0) # newest, sorts first + _save_with_mtime("g1", 200.0) + _save_with_mtime("g2", 100.0) + + def _valid(rec): + return rec.get("prompt") != "BAD" + + # First page of 2 over VALID records returns both good ones -- not [g1] (bad eating a slot) + # and not [] (bad filling the whole window). + page = gallery.list_images(limit = 2, offset = 0, valid = _valid) + assert [r["prompt"] for r in page] == ["g1", "g2"] + # The has_more probe (limit + 1) sees no extra VALID record beyond the two returned. + assert len(gallery.list_images(limit = 3, offset = 0, valid = _valid)) == 2 + + +def test_valid_callback_leading_bad_record_does_not_stall_at_offset_zero(): + # Reproduces the exact stall: every record in the first window is schema-invalid. Without + # in-pager filtering the route returned images=[] with has_more=True at offset 0 forever. + for i in range(3): + _save_with_mtime(f"BAD{i}", 300.0 - i) # newest three are all invalid + _save_with_mtime("good", 10.0) + + def _valid(rec): + return not str(rec.get("prompt", "")).startswith("BAD") + + # limit+1 = 3: the pager must look PAST the invalid leaders and return the one good record, + # so has_more (len > limit) is False and the client advances off offset 0. + records = gallery.list_images(limit = 2, offset = 0, valid = _valid) + assert [r["prompt"] for r in records] == ["good"] + + +def test_save_is_atomic_no_partial_png_on_publish_failure(monkeypatch): + # A crash between writing the bytes and publishing the file must leave neither a truncated + # {id}.png nor a leftover temp: the listing only ever sees fully-written records. + def _boom(*a, **k): + raise OSError("simulated rename failure") + + monkeypatch.setattr(gallery.os, "replace", _boom) + with pytest.raises(OSError, match = "simulated rename failure"): + gallery.save(_img(), _meta()) + # No final PNG surfaced, and the hidden temp was cleaned up. + assert list(gallery.gallery_dir().glob("*.png")) == [] + assert list(gallery.gallery_dir().iterdir()) == [] diff --git a/studio/backend/tests/test_sd_cpp_install.py b/studio/backend/tests/test_sd_cpp_install.py index 07d52a0b02..52ca6eb328 100644 --- a/studio/backend/tests/test_sd_cpp_install.py +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -262,6 +262,52 @@ def test_install_downloads_verifies_extracts(tmp_path, monkeypatch): assert (tmp_path / ".unsloth-studio-owned").is_file() +def test_install_into_empty_dir_claims_ownership(tmp_path, monkeypatch): + # An empty (or freshly created) target may be adopted: the marker is written so the uninstaller + # can later remove the Studio-installed tree. + zb = _zip_with_sd_cli() + _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest()) + empty = tmp_path / "sdcpp" + empty.mkdir() # exists but empty + install(install_dir = empty) + assert (empty / ".unsloth-studio-owned").is_file() + + +def test_install_into_nonempty_unowned_dir_does_not_claim_ownership(tmp_path, monkeypatch): + # A pre-existing, non-empty directory that Studio did not create (e.g. a user's own + # stable-diffusion.cpp checkout) must NOT be marked owned: writing the marker would make the + # uninstaller recursively delete the user's directory. The install still proceeds (extracts the + # binary) but leaves the directory unowned so the uninstaller keeps it. + zb = _zip_with_sd_cli() + _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest()) + target = tmp_path / "stable-diffusion.cpp" + target.mkdir() + user_file = target / "USER_WORK" + user_file.write_text("keep", encoding = "utf-8") + + sd_cli = install(install_dir = target) + + # The user's file survives and no ownership marker was written. + assert user_file.read_text(encoding = "utf-8") == "keep" + assert not (target / ".unsloth-studio-owned").exists() + # The binary was still installed (install is not refused). + assert sd_cli.is_file() and sd_cli.name == "sd-cli" + + +def test_reinstall_into_owned_dir_keeps_ownership(tmp_path, monkeypatch): + # A directory that already carries our marker (a prior Studio install / upgrade) stays owned + # even though it is now non-empty. + zb = _zip_with_sd_cli() + _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest()) + target = tmp_path / "stable-diffusion.cpp" + target.mkdir() + (target / ".unsloth-studio-owned").touch() + (target / "old-junk").write_text("x", encoding = "utf-8") + + install(install_dir = target) + assert (target / ".unsloth-studio-owned").is_file() + + def test_install_sha256_mismatch_raises_and_cleans_up(tmp_path, monkeypatch): zb = _zip_with_sd_cli() name = _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + "0" * 64) diff --git a/studio/backend/tests/test_video_gallery.py b/studio/backend/tests/test_video_gallery.py index 21a8dac014..f86dbe3bcd 100644 --- a/studio/backend/tests/test_video_gallery.py +++ b/studio/backend/tests/test_video_gallery.py @@ -199,6 +199,59 @@ def test_list_skips_corrupt_sidecar(): assert [r["prompt"] for r in listed] == ["ours"] +def test_valid_callback_paginates_over_accepted_records(): + # A sidecar that parses as JSON (so the read accepts it) but fails the caller's stricter schema + # check must be filtered BEFORE pagination, so offset/limit/has_more count over accepted records + # only. Otherwise a leading bad record returns a short/empty page with more remaining and stalls + # infinite scroll at offset 0. + _save_with_mtime("BAD", 300.0) # newest, sorts first + _save_with_mtime("g1", 200.0) + _save_with_mtime("g2", 100.0) + + def _valid(rec): + return rec.get("prompt") != "BAD" + + page = gallery.list_videos(limit = 2, offset = 0, valid = _valid) + assert [r["prompt"] for r in page] == ["g1", "g2"] + assert len(gallery.list_videos(limit = 3, offset = 0, valid = _valid)) == 2 + + +def test_valid_callback_leading_bad_records_do_not_stall_at_offset_zero(): + # Every record in the first window is schema-invalid: the pager must look past them and return + # the good record so has_more is False and the client advances off offset 0. + for i in range(3): + _save_with_mtime(f"BAD{i}", 300.0 - i) + _save_with_mtime("good", 10.0) + + def _valid(rec): + return not str(rec.get("prompt", "")).startswith("BAD") + + records = gallery.list_videos(limit = 2, offset = 0, valid = _valid) + assert [r["prompt"] for r in records] == ["good"] + + +def test_save_leaves_no_orphan_mp4_when_sidecar_publish_fails(monkeypatch): + # The sidecar is the pair's commit marker. If it fails to publish, the MP4 must NOT be left + # behind as an invisible orphan (list_videos would skip it and gallery delete could never reach + # it). Fail the SECOND os.replace (the sidecar) after the mp4 is renamed in, and assert nothing + # is stranded. + real_replace = gallery.os.replace + calls = {"n": 0} + + def _replace(src, dst, *a, **k): + calls["n"] += 1 + if calls["n"] == 2: # the sidecar publish + raise OSError("simulated sidecar failure") + return real_replace(src, dst, *a, **k) + + monkeypatch.setattr(gallery.os, "replace", _replace) + with pytest.raises(OSError, match = "simulated sidecar failure"): + gallery.save(_mp4(), _meta()) + # No mp4, no sidecar, no temp files -- the whole record was rolled back. + assert list(gallery.gallery_dir().iterdir()) == [] + assert gallery.list_videos() == [] + + def _real_mp4_bytes() -> bytes: # A real (tiny) MP4 for the transcode tests: 8 frames of flat color at # 32x32, encoded with mpeg4 (bundled in every PyAV build, unlike libx264). diff --git a/studio/install_sd_cpp_prebuilt.py b/studio/install_sd_cpp_prebuilt.py index a7b00861e7..4f99f6d364 100644 --- a/studio/install_sd_cpp_prebuilt.py +++ b/studio/install_sd_cpp_prebuilt.py @@ -401,6 +401,23 @@ def install( has no ``sd-cli``. """ target = install_dir or default_install_dir() + # Decide up front whether this install may claim ownership of ``target``. We only mark a + # directory as Studio-owned (and therefore eligible for the uninstaller's recursive delete) + # when this install actually created it or it was empty -- NEVER when it already held a user's + # own stable-diffusion.cpp checkout or unrelated files. Adopting a pre-existing, unowned, + # non-empty directory would let a later uninstall wipe the user's own work. A directory that + # already carries our marker (a prior Studio install / upgrade) stays owned. + marker = target / ".unsloth-studio-owned" + _may_own = True + if target.exists(): + if not target.is_dir(): + raise RuntimeError(f"sd.cpp install target is not a directory: {target}") + try: + _pre_existing_entries = any(target.iterdir()) + except OSError: + _pre_existing_entries = True + # Empty dir, or one we already own, may be (re)claimed; a non-empty unowned dir may not. + _may_own = (not _pre_existing_entries) or marker.is_file() used_repo, release, chosen = _resolve_with_fallback(accelerator, token) if release is None or not chosen: @@ -443,11 +460,14 @@ def install( print(f"installed sd-server -> {sd_server}", flush = True) # Ownership marker (the same one setup.sh/_is_studio_root use, and setup.ps1 writes into the # Node sibling dir) so the uninstaller can tell a Studio-installed sd.cpp from a user's own - # stable-diffusion.cpp checkout beside a custom Studio root, and delete only ours. - try: - (target / ".unsloth-studio-owned").touch() - except OSError: - pass + # stable-diffusion.cpp checkout beside a custom Studio root, and delete only ours. Written only + # when this install created the directory or it was empty (see _may_own above): a pre-existing, + # unowned, non-empty directory keeps its unowned status so the uninstaller leaves it alone. + if _may_own: + try: + marker.touch() + except OSError: + pass return sd_cli From 499bada598ebbc4d1280a8152d92fb0f0af749f5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 13 Jul 2026 01:23:42 +0000 Subject: [PATCH 02/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_diffusion_backend.py | 4 ++-- studio/backend/tests/test_diffusion_engine_router.py | 4 +--- studio/backend/tests/test_diffusion_routes.py | 7 ++++++- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 0f8b84f083..477577ba35 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -3030,8 +3030,8 @@ def test_unload_waits_for_in_flight_denoise_before_teardown(): with backend._generate_lock: denoise_active["v"] = True started.set() - cancel.wait(2.0) # unload signals this - finish.wait(2.0) # the test lets us finish + cancel.wait(2.0) # unload signals this + finish.wait(2.0) # the test lets us finish denoise_active["v"] = False # about to release _generate_lock def _fake_unload_locked(): diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py index 53d40a1336..4ff9bc23c7 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -258,9 +258,7 @@ def test_activate_serializes_switch_and_concurrent_query(monkeypatch): unload_started.set() release_unload.wait(2.0) - engine = SimpleNamespace( - status = lambda: {"loaded": False, "repo_id": None}, unload = _slow_unload - ) + engine = SimpleNamespace(status = lambda: {"loaded": False, "repo_id": None}, unload = _slow_unload) monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: engine) switch_done = threading.Event() diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 6e6f7ea2da..34b4061fd0 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -180,7 +180,12 @@ def client(monkeypatch, tmp_path): monkeypatch.setattr(gallery_module, "save", _save) monkeypatch.setattr(gallery_module, "image_b64", lambda i: "QUJD" if i in store else None) - def _list_images(limit = None, offset = 0, *, valid = None): + def _list_images( + limit = None, + offset = 0, + *, + valid = None, + ): ordered = sorted(store.values(), key = lambda r: r.get("created_at", 0.0), reverse = True) if valid is not None: ordered = [r for r in ordered if valid(r)] From 5a17614b51cdbb8a14e0e5ef51ef792e7a6ddef3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 02:09:15 +0000 Subject: [PATCH 03/10] Reject native batch seeds outside the JSON-safe range --- studio/backend/models/inference.py | 15 ++++++++++++++ studio/backend/tests/test_diffusion_routes.py | 20 +++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index a6800087ec..862942e883 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2069,6 +2069,21 @@ class DiffusionGenerateRequest(BaseModel): raise ValueError("must be a multiple of 16") return value + @model_validator(mode = "after") + def _batch_seeds_json_safe(self) -> "DiffusionGenerateRequest": + # A batch derives per-image seeds as seed, seed+1, ... seed+batch_size-1 (sd.cpp and + # diffusers both advance the seed per image). The base seed is bounded to 2**53-1 so it + # round-trips through the JSON gallery recipe, but the derived top-of-batch seed is not: + # with an explicit seed near the cap it can exceed Number.MAX_SAFE_INTEGER, where the + # frontend rounds it and a restored recipe replays a different image. Reject at the + # boundary so an API client can't persist an unreplayable seed. + if self.seed is not None and self.seed + self.batch_size - 1 > 2**53 - 1: + raise ValueError( + "seed + batch_size - 1 must not exceed 2**53 - 1 so every per-image seed " + "stays JSON-safe (lower the seed or the batch_size)" + ) + return self + class GalleryImage(BaseModel): """A persisted image's full generation recipe (embedded in the PNG too).""" diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 34b4061fd0..a5a825ee00 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -340,6 +340,26 @@ def test_generate_rejects_non_multiple_of_16(client): assert ok.status_code == 200 +def test_generate_rejects_batch_seed_past_json_safe_range(client): + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) + # An explicit seed at the JS-safe cap with a batch derives per-image seeds + # (seed+1 ...) that exceed Number.MAX_SAFE_INTEGER and no longer round-trip + # through the gallery JSON recipe, so the request is rejected at the boundary. + over = client.post( + "/api/inference/images/generate", + json = {"prompt": "p", "seed": 2**53 - 1, "batch_size": 2}, + ) + assert over.status_code == 422 + # The top-of-batch seed lands exactly on the cap: still JSON-safe, so accepted. + ok = client.post( + "/api/inference/images/generate", + json = {"prompt": "p", "seed": 2**53 - 2, "batch_size": 2}, + ) + assert ok.status_code == 200 + + def test_non_gguf_load_restricted_to_unsloth(client): # gguf_filename is optional now; with none, the load is a full-pipeline kind, which # is gated to unsloth/* repos. A non-unsloth repo (no filename) is rejected -> 400. From da1770bb4425de5be5df316ca01f4611b7343de6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 03:17:49 +0000 Subject: [PATCH 04/10] Refuse sd.cpp install into unowned non-empty target dir When the install target already exists, is non-empty and lacks the .unsloth-studio-owned marker (a user's own stable-diffusion.cpp checkout, or unrelated files beside a custom Studio root), install() previously still extracted the release into it. Skipping the ownership marker only stopped the uninstaller from deleting the directory; extraction still merged binaries into the user's working tree and could overwrite same-named files. Fail up front with a clear message pointing the user at a fresh/empty location before any download or extraction, leaving their directory untouched. Update the ownership test suite to assert the refusal. --- studio/backend/tests/test_sd_cpp_install.py | 17 +++++++++-------- studio/install_sd_cpp_prebuilt.py | 12 ++++++++++++ 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/studio/backend/tests/test_sd_cpp_install.py b/studio/backend/tests/test_sd_cpp_install.py index 52ca6eb328..3a061c9fcf 100644 --- a/studio/backend/tests/test_sd_cpp_install.py +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -273,11 +273,12 @@ def test_install_into_empty_dir_claims_ownership(tmp_path, monkeypatch): assert (empty / ".unsloth-studio-owned").is_file() -def test_install_into_nonempty_unowned_dir_does_not_claim_ownership(tmp_path, monkeypatch): +def test_install_into_nonempty_unowned_dir_is_refused(tmp_path, monkeypatch): # A pre-existing, non-empty directory that Studio did not create (e.g. a user's own - # stable-diffusion.cpp checkout) must NOT be marked owned: writing the marker would make the - # uninstaller recursively delete the user's directory. The install still proceeds (extracts the - # binary) but leaves the directory unowned so the uninstaller keeps it. + # stable-diffusion.cpp checkout) must NOT be extracted into. Merging the release into it would + # overwrite or mix our binaries into the user's working tree, and leaving it unowned only stops + # the uninstaller from deleting it later. install() refuses up front and leaves the dir untouched + # so the user can point us at a fresh/empty location. zb = _zip_with_sd_cli() _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest()) target = tmp_path / "stable-diffusion.cpp" @@ -285,13 +286,13 @@ def test_install_into_nonempty_unowned_dir_does_not_claim_ownership(tmp_path, mo user_file = target / "USER_WORK" user_file.write_text("keep", encoding = "utf-8") - sd_cli = install(install_dir = target) + with pytest.raises(RuntimeError, match = "not a Studio-managed directory"): + install(install_dir = target) - # The user's file survives and no ownership marker was written. + # The user's directory is left exactly as it was: file intact, no marker, nothing extracted. assert user_file.read_text(encoding = "utf-8") == "keep" assert not (target / ".unsloth-studio-owned").exists() - # The binary was still installed (install is not refused). - assert sd_cli.is_file() and sd_cli.name == "sd-cli" + assert list(target.iterdir()) == [user_file] def test_reinstall_into_owned_dir_keeps_ownership(tmp_path, monkeypatch): diff --git a/studio/install_sd_cpp_prebuilt.py b/studio/install_sd_cpp_prebuilt.py index 4f99f6d364..1923cd3f8e 100644 --- a/studio/install_sd_cpp_prebuilt.py +++ b/studio/install_sd_cpp_prebuilt.py @@ -418,6 +418,18 @@ def install( _pre_existing_entries = True # Empty dir, or one we already own, may be (re)claimed; a non-empty unowned dir may not. _may_own = (not _pre_existing_entries) or marker.is_file() + # Refuse to extract into a pre-existing, non-empty directory we do not own (a user's own + # stable-diffusion.cpp checkout, or unrelated files beside a custom Studio root). Not writing + # the ownership marker only protects the uninstaller; extracting the release here would still + # merge our binaries into the user's working tree and can overwrite same-named files. Fail with + # a clear message so the user points us at a fresh/empty location instead of corrupting theirs. + if not _may_own: + raise RuntimeError( + f"sd.cpp install target already exists and is not a Studio-managed directory: {target}. " + f"Refusing to extract prebuilt binaries into it to avoid overwriting or mixing them " + f"into your files. Remove or move that directory, or install into a different, empty " + f"location (pass a different --install-dir / set the Studio sd.cpp install dir)." + ) used_repo, release, chosen = _resolve_with_fallback(accelerator, token) if release is None or not chosen: From dabb12e1987984f44b60ae650d76a5edcb7a2e1d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 04:21:36 +0000 Subject: [PATCH 05/10] Studio: gate dataset uploads on the symlink check and surface local video single-file checkpoints --- studio/backend/routes/models.py | 14 ++++++++++++++ studio/backend/routes/training.py | 8 ++++++-- .../backend/tests/test_diffusion_dataset_api.py | 16 ++++++++++++++++ studio/backend/tests/test_local_model_format.py | 15 ++++++++++----- 4 files changed, 46 insertions(+), 7 deletions(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index cd45378110..71f4c2f41d 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -3320,6 +3320,20 @@ def _local_is_diffusers(model: "LocalModelInfo") -> bool: return True except Exception: pass + # A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json) + # ships no pipeline index and resolves to no IMAGE family either, so the checks above miss + # it. The video load route reinterprets a bare single-file local dir as a single_file load + # (routes/video.py), so it IS loadable and must be surfaced; without tagging it here + # _local_model_task returns task=null and the Video On-Device picker hides it. Match the same + # clean id / name needles (not the raw path) so a parent-dir family token can't spuriously + # match; _local_model_task then routes the video family to the text-to-video task. + try: + from core.inference.video_families import detect_video_family + for needle in (model.model_id, model.display_name, Path(model.id).name): + if needle and detect_video_family(needle) is not None: + return True + except Exception: + pass return False diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index cb36e43a45..3ce107a323 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1613,11 +1613,15 @@ async def upload_diffusion_dataset( import os import tempfile - from utils.paths import datasets_root from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label cleaned = _clean_diffusion_dataset_name(name) - folder = datasets_root() / cleaned + # Reject a symlinked dataset directory BEFORE any write. A bare mkdir(exist_ok=True) succeeds + # through an existing name -> external-directory symlink, and the staged upload would then + # write/replace files outside the Studio datasets root through that link. The read/caption/ + # delete endpoints already enforce this containment via _resolve_dataset_folder; the upload + # path must run the same symlink + root-containment check first so writes can never escape. + folder = _resolve_dataset_folder(name, must_exist = False) folder.mkdir(parents = True, exist_ok = True) limit_bytes = get_upload_limit_bytes() diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index e80b29d6e6..aa59adef10 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -563,6 +563,22 @@ def test_resolve_dataset_folder_rejects_symlink(ds_root, tmp_path): assert exc.value.status_code == 400 +def test_upload_through_symlinked_dataset_cannot_escape_root(client, ds_root, tmp_path): + # End to end: an upload targeting a dataset name that already exists as a symlink to an + # external directory must be refused (400) BEFORE any bytes are written, so the upload can + # never create/replace files outside the datasets root through the link. + external = tmp_path / "external" + external.mkdir() + (ds_root / "linked").symlink_to(external, target_is_directory = True) + + r = _upload(client, "linked", [("intruder.png", _png_bytes())]) + assert r.status_code == 400 + assert "symbolic link" in r.json()["detail"] + # Nothing was written through the link into the external directory. + assert not (external / "intruder.png").exists() + assert not any(external.iterdir()) + + def test_delete_through_symlinked_dataset_cannot_escape_root(client, ds_root, tmp_path): # End to end: a DELETE against an image inside a symlinked dataset dir is refused (400) and the # external file it points at is NOT removed. diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index 3da98ceabb..7163db424b 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -180,13 +180,18 @@ def test_local_task_tags_video_pipeline_dir(tmp_path): ) -def test_local_task_video_name_without_pipeline_not_surfaced(tmp_path): - # A dir whose name matches a video family but which is NOT a diffusers pipeline (no - # model_index.json) is not a loadable pipeline, so it must stay untagged -- never surfaced - # to the Video picker, so it can never trigger a pipeline load that evicts then fails. +def test_local_task_tags_video_single_file_checkpoint(tmp_path): + # A dir whose name matches a video family holding a bare single-file .safetensors (no + # model_index.json) IS loadable: the video load route reinterprets a sole single-file local + # pick as a single_file load (routes/video.py), validating BEFORE it touches the GPU. So it + # must be tagged text-to-video and surfaced in the Video On-Device picker -- not left task=null + # and hidden, which would make the advertised-and-loadable checkpoint unusable. d = tmp_path / "ltx-loose" _touch(d / "ltx-2.safetensors") # loose weights, no model_index.json - assert models_route._local_model_task(_local(d, model_id = "Lightricks/LTX-2")) is None + assert ( + models_route._local_model_task(_local(d, model_id = "Lightricks/LTX-2")) + == models_route._VIDEO_GEN_TASK + ) def test_local_task_ignores_family_token_in_parent_path(tmp_path): From e0ef488f47788144ef822dfa8419bccce2cfd4a5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 05:29:09 +0000 Subject: [PATCH 06/10] Tighten comments and docstrings added by the image-generation fixes --- scripts/uninstall.ps1 | 7 ++-- studio/backend/core/inference/diffusion.py | 36 +++++++------------ .../core/inference/diffusion_engine_router.py | 19 ++++------ .../backend/core/inference/image_gallery.py | 15 ++++---- .../backend/core/inference/video_gallery.py | 17 ++++----- .../core/training/diffusion_train_common.py | 17 ++++----- studio/backend/core/training/training.py | 10 +++--- studio/backend/models/inference.py | 10 +++--- studio/backend/routes/inference.py | 9 +++-- studio/backend/routes/models.py | 12 +++---- studio/backend/routes/training.py | 31 ++++++---------- studio/backend/routes/video.py | 22 +++++------- .../backend/tests/test_diffusion_backend.py | 13 +++---- .../tests/test_diffusion_dataset_api.py | 30 ++++++---------- .../tests/test_diffusion_engine_router.py | 10 ++---- .../tests/test_diffusion_lora_trainer.py | 12 +++---- studio/backend/tests/test_diffusion_routes.py | 5 ++- studio/backend/tests/test_image_gallery.py | 18 ++++------ .../backend/tests/test_local_model_format.py | 8 ++--- studio/backend/tests/test_sd_cpp_install.py | 7 ++-- .../tests/test_training_start_offload.py | 7 ++-- studio/backend/tests/test_video_gallery.py | 12 +++---- studio/backend/tests/test_video_routes.py | 12 +++---- .../src/features/images/images-page.tsx | 29 ++++++--------- studio/install_sd_cpp_prebuilt.py | 24 +++++-------- 25 files changed, 140 insertions(+), 252 deletions(-) diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index a38f0d0eab..61bf73b5b4 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -364,11 +364,8 @@ function Uninstall-UnslothStudio { _StopByPortFile -PortFile (Join-Path $r "share\studio.port") -KnownRoots $knownRoots } _StopStudioProcesses -KnownRoots $knownRoots - # The default sd.cpp dir is only killed+deleted when it carries our owner marker (the deletion - # below is marker-gated). Passing it to the locking-process stop UNCONDITIONALLY would terminate - # a user's own sd-server running from an unowned checkout at this default path -- a process we - # then decide to keep the directory for. Gate the kill on the same predicate so stop and delete - # agree: include it only when marked owned (and present). + # Only stop the default sd.cpp dir when it carries our owner marker, so stop matches the + # marker-gated delete and a user's own sd-server at this default path is left running. $defaultSdCppToStop = $null if ($defaultSdCpp -and (Test-Path -LiteralPath $defaultSdCpp) -and (Test-Path -LiteralPath (Join-Path $defaultSdCpp ".unsloth-studio-owned") -PathType Leaf)) { $defaultSdCppToStop = $defaultSdCpp diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 9ed4cd007b..feeb967d86 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -2216,14 +2216,10 @@ class DiffusionBackend: raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) # Register under _lock so unload()/a load can signal THIS generation. self._active_generate_cancel = cancel - # Publish an active (step 0) progress state the moment the lock is held, BEFORE - # the slow pre-denoise setup (deferred compile, LoRA resolution/application, - # ControlNet download/build). Without this generate_progress() reports inactive - # across that window, so a reload's mount probe shows idle even though this - # generation holds _generate_lock; the user then starts a second generate that - # merely blocks behind this one (and can duplicate the result). The per-step - # callback swaps in its own _GenState at denoise start; this is the queued phase. - # Mirrors the video backend's queued state and the training start guard. + # Publish an active (step 0) state now, before the slow pre-denoise setup (deferred + # compile, LoRA resolution, ControlNet build), so a reload's mount probe doesn't read + # idle while this generation holds _generate_lock and let a second generate queue + # behind it. The per-step callback swaps in its own _GenState at denoise start. self._gen = _GenState(total_steps = steps) try: # The local `state` ref keeps the pipe alive even if unload() nulls _state. @@ -2559,10 +2555,8 @@ class DiffusionBackend: with self._lock: if self._active_generate_cancel is cancel: self._active_generate_cancel = None - # Drop the published progress state. The normal path already nulled it after - # the denoise; this also covers a setup-time error that skips that inner - # finally. Safe under _generate_lock: no other generation can have installed - # its own _gen while this one runs. + # Drop the published progress state, covering a setup-time error that skips + # the inner finally. Safe under _generate_lock. self._gen = None def generate_progress(self) -> dict[str, Any]: @@ -2594,13 +2588,10 @@ class DiffusionBackend: # Cancel any in-flight load (its worker checks this token) and drop the marker. self._load_token += 1 self._loading = None - # WAIT for the signalled denoise to exit BEFORE tearing down (mirrors begin_load's - # pre-teardown barrier at the top of the locked load path). _unload_locked uninstalls - # PROCESS-WIDE state the running denoise still depends on -- the eager/arch attention - # patches, the GGUF compile hooks, the flipped backend flags and the compile cache -- none - # of which the denoise's own pipe ref pins. Uninstalling them before the denoise exits would - # corrupt or crash its in-flight forward passes. The denoise holds _generate_lock for its - # whole body, so acquiring it here blocks until it has finished; only then do we tear down. + # Wait for the signalled denoise to exit BEFORE tearing down: _unload_locked uninstalls + # process-wide state (attention patches, GGUF compile hooks, backend flags, compile cache) + # the denoise still depends on but its pipe ref doesn't pin. The denoise holds _generate_lock + # for its whole body, so acquiring it here blocks until it has finished. Mirrors begin_load. with self._generate_lock: with self._lock: self._unload_locked() @@ -2622,10 +2613,9 @@ class DiffusionBackend: uninstall_patches() uninstall_arch_patches() - # NOTE: deliberately NOT unload_lora_weights() here. Both callers (unload() and begin_load's - # locked path) now hold _generate_lock across this teardown, so no denoise is in flight; the - # whole pipe is dropped below, freeing any LoRA adapters with it, so a separate adapter unload - # would be redundant work. + # Deliberately NOT unload_lora_weights() here: the whole pipe is dropped below, freeing any + # LoRA adapters with it. Both callers hold _generate_lock across this teardown, so no denoise + # is in flight. # Drop the workflow pipes so they don't pin the freed pipeline's modules past unload. self._aux_pipes.clear() # Drop any ControlNet models + pipelines so the freed load carries no extra modules. diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index 8422e7305b..4b0c9b9155 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -59,9 +59,8 @@ def _install_accelerator_for(backend: str) -> str: # The engine the current load committed to, and why a non-native choice was made. Mutated only # under _lock during selection. _lock = threading.Lock() -# Serializes a WHOLE engine switch (check -> unload -> publish). _lock alone is released during the -# slow unload(), so two overlapping selections could interleave: one observes the not-yet-updated -# active engine, returns it, and loads onto the very engine the other is concurrently unloading. +# Serializes a whole engine switch (check -> unload -> publish); _lock alone is released during the +# slow unload(), letting two overlapping selections load onto the engine the other is unloading. _transition_lock = threading.Lock() _active_engine_name: str = ENGINE_DIFFUSERS _fallback_reason: Optional[str] = None @@ -90,12 +89,9 @@ def active_engine_name() -> str: def _activate(name: str, reason: Optional[str]) -> Any: global _active_engine_name, _fallback_reason - # Serialize the ENTIRE check -> unload -> publish transition. _lock is released during the slow - # unload() below (holding it across the unload would block every status/selection reader), which - # opens a window where a second _activate could read the still-old active engine, take the "no - # change" branch, and return that engine -- then load onto it while this call is unloading it. - # _transition_lock closes the window without holding _lock across the unload; the final - # get_active_diffusion_engine() now reflects the committed state. + # Serialize the whole check -> unload -> publish transition without holding _lock across the + # slow unload(), closing the window where a second _activate reads the still-old active engine, + # takes the "no change" branch, and loads onto the engine this call is unloading. with _transition_lock: # Switching engines: unload the deactivated one first, else its model stays resident but # unreachable (the evictor only targets the active engine), leaking 10+ GB. The unload is @@ -112,9 +108,8 @@ def _activate(name: str, reason: Optional[str]) -> Any: if engine_to_unload is not None: # Publish the new engine only AFTER the old one unloads. The evictor unloads # get_active_diffusion_engine(), so flipping the name first would let a concurrent - # acquire_for evict the new (empty) engine and take the GPU while the old model is still - # freeing VRAM -- two large models briefly resident. Keeping the OLD engine as the evict - # target serializes a concurrent evict on its unload(), granting the GPU only once freed. + # acquire_for evict the new (empty) engine while the old model is still freeing VRAM. + # Keeping the OLD engine as the evict target grants the GPU only once it is freed. try: engine_to_unload.unload() except Exception as exc: # noqa: BLE001 -- best-effort; never block the switch diff --git a/studio/backend/core/inference/image_gallery.py b/studio/backend/core/inference/image_gallery.py index 8bede83282..fb1e67b096 100644 --- a/studio/backend/core/inference/image_gallery.py +++ b/studio/backend/core/inference/image_gallery.py @@ -69,9 +69,8 @@ def save(image: Any, meta: dict[str, Any]) -> dict[str, Any]: image_id = uuid.uuid4().hex directory = gallery_dir() final_path = directory / f"{image_id}.png" - # Write to a hidden temp then atomically rename into place: a crash / disk-full mid-write must - # never leave a truncated {id}.png that the listing would surface as a corrupt record. The temp - # name is dotted so an interrupted write is skipped by the *.png glob, and cleaned on failure. + # Write to a dotted temp (skipped by the *.png glob) then atomically rename, so a crash mid-write + # never leaves a truncated {id}.png that the listing would surface as a corrupt record. tmp_path = directory / f".{image_id}.png.tmp" try: tmp_path.write_bytes(_png_bytes(image, meta)) @@ -156,12 +155,10 @@ def list_images( full just to sort; only the window's recipes are read. limit=None returns everything from ``offset`` on. - ``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and the - caller's has_more all count over the same accepted-record domain. Passing the route's schema - validator here is essential: a record that has every required key (so ``_read_meta`` accepts - it) but a wrong value type would otherwise be counted in this window yet dropped by the route - after slicing -- a leading bad record then returns an empty page with more remaining, which - stalls infinite scroll at offset 0.""" + ``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and has_more + all count over the accepted-record domain. Pass the route's schema validator: a record with + every required key (so ``_read_meta`` accepts it) but a wrong value type would otherwise be + counted here yet dropped after slicing, stalling infinite scroll at offset 0.""" try: paths = list(gallery_dir().glob("*.png")) except OSError: diff --git a/studio/backend/core/inference/video_gallery.py b/studio/backend/core/inference/video_gallery.py index b703f44a42..1391119217 100644 --- a/studio/backend/core/inference/video_gallery.py +++ b/studio/backend/core/inference/video_gallery.py @@ -39,11 +39,9 @@ def save(mp4_bytes: bytes, meta: dict[str, Any]) -> dict[str, Any]: mp4_tmp = directory / f".{video_id}.mp4.tmp" sidecar = directory / f"{video_id}.json" sidecar_tmp = directory / f".{video_id}.json.tmp" - # Stage BOTH files, then publish. The sidecar is the pair's commit marker (list_videos scans - # mp4s but skips any without a readable sidecar), so writing the MP4 straight to its final name - # before the sidecar meant a sidecar write/replace failure left an invisible, undeletable orphan - # MP4 (gallery delete resolves by id, but the record never appears to be deleted). Rename the MP4 - # in first, then the sidecar; on ANY failure remove every artifact so nothing is stranded. + # Stage both files, rename the MP4 in, then the sidecar (the pair's commit marker: list_videos + # skips an mp4 without a readable sidecar). On any failure remove every artifact, else a sidecar + # failure would leave an invisible, undeletable orphan MP4. try: mp4_tmp.write_bytes(mp4_bytes) sidecar_tmp.write_text(json.dumps(meta), encoding = "utf-8") @@ -204,11 +202,10 @@ def list_videos( Ordered by MP4 mtime (a cheap stat ~= generation order); only the window's sidecars are read. limit=None returns everything from ``offset`` on. A file without its pair is skipped. - ``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and the - caller's has_more all count over the same accepted-record domain. Passing the route's schema - validator here keeps a sidecar that parses as JSON but fails the response schema from being - counted in this window yet dropped by the route after slicing -- which would otherwise return a - short or empty page with more remaining and stall infinite scroll.""" + ``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and has_more + all count over the accepted-record domain. Pass the route's schema validator: a sidecar that + parses as JSON but fails the response schema would otherwise be counted here yet dropped after + slicing, stalling infinite scroll.""" try: paths = list(gallery_dir().glob("*.mp4")) except OSError: diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index b4fdf81d1e..dc5f5f4fc6 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -705,27 +705,22 @@ def discover_image_caption_pairs( caption: Optional[str] = None sidecar_present = False # 1. per-image sidecar caption file (the user's explicit edit; wins over metadata). - # An EMPTY sidecar is a deliberate tombstone (the labeling grid writes one when a user - # clears a metadata caption): it suppresses the metadata caption but leaves the image - # UNCAPTIONED so the dreambooth instance_prompt fallback below still applies. Treating - # "" as a present caption here would skip the instance prompt AND drop the image, so - # clearing every metadata caption in the grid would make a dreambooth run fail with - # "No captioned images found". + # An EMPTY sidecar is a deliberate tombstone (written when a user clears a caption): it + # suppresses the metadata caption but leaves the image uncaptioned so the instance_prompt + # fallback below still applies, rather than dropping the image. for ext in _CAPTION_EXTS: sidecar = img.with_suffix(ext) if sidecar.is_file(): sidecar_present = True caption = sidecar.read_text(encoding = "utf-8").strip() break - # 2. metadata row keyed by file name (basename or the relative path; as_posix so a - # Windows backslash path still matches the jsonl's forward-slash keys). Skipped when a - # sidecar tombstone is present (the sidecar, even empty, is authoritative). + # 2. metadata row keyed by file name (basename or relative path; as_posix so a Windows + # backslash path matches the jsonl's forward-slash keys). A sidecar, even empty, wins. if not sidecar_present: caption = meta_caption.get(img.name) or meta_caption.get( img.relative_to(root).as_posix() ) - # 3. dreambooth instance prompt for any image still without a caption (no sidecar text and - # no metadata row, or an empty tombstone). + # 3. dreambooth instance prompt for any image still without a caption. if not caption and instance_prompt: caption = instance_prompt if caption: diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 30dda0f7d0..526551aab4 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -1141,12 +1141,10 @@ 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(). + # A run reserved 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 an + # idle reading here would let another pipeline race the reserved run for VRAM. if self._start_in_progress: return True diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 862942e883..23218b7f35 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2071,12 +2071,10 @@ class DiffusionGenerateRequest(BaseModel): @model_validator(mode = "after") def _batch_seeds_json_safe(self) -> "DiffusionGenerateRequest": - # A batch derives per-image seeds as seed, seed+1, ... seed+batch_size-1 (sd.cpp and - # diffusers both advance the seed per image). The base seed is bounded to 2**53-1 so it - # round-trips through the JSON gallery recipe, but the derived top-of-batch seed is not: - # with an explicit seed near the cap it can exceed Number.MAX_SAFE_INTEGER, where the - # frontend rounds it and a restored recipe replays a different image. Reject at the - # boundary so an API client can't persist an unreplayable seed. + # A batch derives per-image seeds as seed .. seed+batch_size-1. The base seed is capped at + # 2**53-1 to round-trip through the JSON recipe, but a derived top-of-batch seed near the cap + # can exceed it, where the frontend rounds it and a restored recipe replays a different + # image. Reject at the boundary so an API client can't persist an unreplayable seed. if self.seed is not None and self.seed + self.batch_size - 1 > 2**53 - 1: raise ValueError( "seed + batch_size - 1 must not exceed 2**53 - 1 so every per-image seed " diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index c452a4775a..fab3652b7f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -14512,11 +14512,10 @@ async def list_gallery_images( limit = max(1, min(limit, 200)) offset = max(0, offset) - # Validate against the response schema INSIDE the pager so offset / limit / has_more all count - # over the same accepted-record domain. A PNG whose recipe chunk has all keys but a wrong value - # type passes the presence-only read yet fails GalleryImage(**r); dropping such records only - # after pagination made a leading bad record return an empty page with has_more=True, stalling - # infinite scroll at offset 0. Filtering here keeps the window and has_more consistent. + # Validate inside the pager so offset / limit / has_more all count over the accepted domain. A + # recipe with all keys but a wrong value type passes the presence-only read yet fails + # GalleryImage(**r); dropping it only after slicing let a leading bad record return an empty + # page with has_more=True, stalling infinite scroll at offset 0. def _valid_gallery_image(record: dict) -> bool: try: GalleryImage(**record) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 71f4c2f41d..fe11446561 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -3320,13 +3320,11 @@ def _local_is_diffusers(model: "LocalModelInfo") -> bool: return True except Exception: pass - # A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json) - # ships no pipeline index and resolves to no IMAGE family either, so the checks above miss - # it. The video load route reinterprets a bare single-file local dir as a single_file load - # (routes/video.py), so it IS loadable and must be surfaced; without tagging it here - # _local_model_task returns task=null and the Video On-Device picker hides it. Match the same - # clean id / name needles (not the raw path) so a parent-dir family token can't spuriously - # match; _local_model_task then routes the video family to the text-to-video task. + # A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json) has no + # pipeline index and no image family, so the checks above miss it. The video load route loads it + # as a single_file (routes/video.py), so it must be surfaced or _local_model_task returns + # task=null and the picker hides it. Match clean id / name needles (not the raw path) so a + # parent-dir token can't spuriously match; _local_model_task then routes it to text-to-video. try: from core.inference.video_families import detect_video_family for needle in (model.model_id, model.display_name, Path(model.id).name): diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 3ce107a323..9d28bec7dc 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1616,11 +1616,8 @@ async def upload_diffusion_dataset( from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label cleaned = _clean_diffusion_dataset_name(name) - # Reject a symlinked dataset directory BEFORE any write. A bare mkdir(exist_ok=True) succeeds - # through an existing name -> external-directory symlink, and the staged upload would then - # write/replace files outside the Studio datasets root through that link. The read/caption/ - # delete endpoints already enforce this containment via _resolve_dataset_folder; the upload - # path must run the same symlink + root-containment check first so writes can never escape. + # Run the same symlink + root-containment check as the read/caption/delete endpoints before any + # write, so a name -> external-directory symlink can't make the staged upload write outside root. folder = _resolve_dataset_folder(name, must_exist = False) folder.mkdir(parents = True, exist_ok = True) @@ -1736,13 +1733,10 @@ async def upload_diffusion_dataset( ) out.write(chunk) uploaded += 1 - # Commit every staged file as one transaction. A plain replace loop is NOT atomic across - # files: if the second-or-later tmp.replace(dest) fails (disk/quota error, a Windows file - # lock, antivirus, a destination that became a directory), an earlier destination has - # already been overwritten while the request returns an error -- the user's original file - # is gone. Back up each pre-existing destination before overwriting it, then on ANY failure - # remove the versions this request installed and restore every displaced original, so the - # dataset is left exactly as it was before the upload. + # Commit every staged file as one transaction. A plain replace loop is not atomic across + # files: a mid-loop tmp.replace(dest) failure leaves earlier destinations already overwritten + # while the request errors. Back up each pre-existing destination first, then on any failure + # drop the versions this request installed and restore every displaced original. backups: list[tuple[Path, Optional[Path]]] = [] # (dest, backup path or None) installed: list[Path] = [] try: @@ -1809,11 +1803,9 @@ def _resolve_dataset_folder(name: str, *, must_exist: bool = True) -> Path: cleaned = _clean_diffusion_dataset_name(name) root = datasets_root().resolve() folder = root / cleaned - # Reject a symlinked dataset directory. _safe_dataset_image_path only proves each image path - # stays under folder.resolve(); it never proves the folder itself stays under the datasets - # root. A dataset dir that is a symlink to an external directory would therefore let image - # read / caption / delete operate on files outside Studio (a reproduced delete removed an - # external file through such a link). Prove the resolved folder is contained in the root too. + # Reject a symlinked dataset directory and prove the resolved folder stays under root: + # _safe_dataset_image_path only checks each image path, so a folder symlinked to an external + # directory would let read / caption / delete operate on files outside Studio. if folder.is_symlink(): raise HTTPException( status_code = 400, @@ -2063,9 +2055,8 @@ async def delete_diffusion_dataset_image( if thumbs_dir.is_dir(): # Thumbs are keyed on the full filename (stem + extension), so match that here too; # a stem-only glob would strand this image's thumbs or delete a same-stem sibling's. - # Escape the filename first: an uploaded name may legally contain glob metacharacters - # ('[', ']', '*', '?'), and interpolating those raw would make e.g. "[ab].png" match - # "a.png_*.jpg"/"b.png_*.jpg" -- deleting siblings' thumbs while leaving its own behind. + # Escape the name: a raw glob metacharacter (e.g. "[ab].png") would match siblings' + # thumbs while leaving its own behind. for t in thumbs_dir.glob(f"{_glob.escape(image_path.name)}_*.jpg"): t.unlink(missing_ok = True) return {"deleted": image_path.name} diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 04ea6d7d28..7e239151c5 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -85,15 +85,12 @@ async def load_video_model( 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. + # 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. + # A local On-Device pick can be a bare single-file .safetensors dir (no model_index.json) + # that the picker starts as a pipeline with no filename, which would 400 on the missing + # index. If the dir holds exactly one checkpoint, load it as a single_file. Mirrors images. 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: @@ -236,11 +233,10 @@ async def list_gallery_videos( limit = max(1, min(limit, 200)) offset = max(0, offset) - # Validate against the response schema INSIDE the pager so offset / limit / has_more all count - # over the same accepted-record domain. A sidecar that parses as JSON but has a wrong value type - # passes the read yet fails GalleryVideo(**r); dropping such records only after pagination made a - # leading bad record return an empty page with has_more=True, stalling infinite scroll at offset - # 0. Filtering here keeps the window and has_more consistent. + # Validate inside the pager so offset / limit / has_more all count over the accepted domain. A + # sidecar that parses as JSON but has a wrong value type passes the read yet fails + # GalleryVideo(**r); dropping it only after slicing let a leading bad record return an empty + # page with has_more=True, stalling infinite scroll at offset 0. def _valid_gallery_video(record: dict) -> bool: try: GalleryVideo(**record) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 477577ba35..a3c7fb4bb1 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -483,11 +483,8 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): def test_generate_progress_active_during_setup(fake_runtime, tmp_path, monkeypatch): - # A generation must report active from the moment it holds the lock, BEFORE the slow - # pre-denoise setup (deferred compile / LoRA resolution / ControlNet build) runs. - # Otherwise a reload's mount probe sees idle while the lock is held and lets a second - # generate queue behind the first. _apply_loras runs inside that setup window, so probing - # generate_progress() from there exercises the gap the reviewer flagged. + # A generation must report active from the moment it holds the lock, before the slow pre-denoise + # setup. _apply_loras runs inside that window, so probe generate_progress() from there. (tmp_path / "model.gguf").write_bytes(b"weights") backend = DiffusionBackend() backend.load_pipeline( @@ -3009,10 +3006,8 @@ def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path): def test_unload_waits_for_in_flight_denoise_before_teardown(): - # Regression for the unload/denoise teardown race: unload() must WAIT for a running denoise to - # exit (acquire _generate_lock) BEFORE _unload_locked() tears down PROCESS-WIDE state (eager / - # arch attention patches, gguf compile hooks, backend flags, compile cache) that the denoise - # still depends on. Mirror the load path, which already waits on _generate_lock before teardown. + # Regression: unload() must wait for a running denoise to exit (acquire _generate_lock) before + # _unload_locked() tears down process-wide state the denoise still depends on. import threading backend = DiffusionBackend() diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index aa59adef10..60170ee273 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -496,12 +496,8 @@ def test_import_promotion_leaves_no_partial_dataset_on_failure(ds_root, monkeypa def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch): - # Re-uploading a.txt and b.txt (an allowed overwrite of two files already on disk) stages both, - # then commits them. A plain replace loop is NOT atomic: if the SECOND commit fails after the - # first destination was already overwritten, the live dataset is left partially updated with the - # request returning an error -- the user's original a.txt is gone. The transactional commit must - # back up each displaced original and, on any failure, restore every one so the dataset is left - # exactly as it was, with no stray temp/backup files. + # Re-uploading a.txt and b.txt where the SECOND commit fails must roll back the first overwrite, + # so both originals survive and no stray temp/backup files remain. from pathlib import Path app = FastAPI() @@ -518,8 +514,7 @@ def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch): state = {"failed": False} def flaky_replace(self, target, *a, **k): - # Fail exactly once, on the tmp -> b.txt promotion (source is the staged .upload-* part, NOT - # a .upload-backup-* part), so the subsequent backup -> b.txt restore still succeeds. + # Fail once on the tmp -> b.txt promotion only (not the backup restore), so rollback works. if ( not state["failed"] and str(target).endswith("b.txt") @@ -548,9 +543,8 @@ def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch): def test_resolve_dataset_folder_rejects_symlink(ds_root, tmp_path): - # A dataset directory that is a symlink pointing OUTSIDE the datasets root must be rejected: the - # per-image containment check only proves paths stay under folder.resolve(), so without this a - # delete/caption/read could operate on external files through the link. + # A dataset dir that is a symlink outside the datasets root must be rejected, else delete / + # caption / read could operate on external files through the link. from routes.training import _resolve_dataset_folder external = tmp_path / "external" @@ -564,9 +558,8 @@ def test_resolve_dataset_folder_rejects_symlink(ds_root, tmp_path): def test_upload_through_symlinked_dataset_cannot_escape_root(client, ds_root, tmp_path): - # End to end: an upload targeting a dataset name that already exists as a symlink to an - # external directory must be refused (400) BEFORE any bytes are written, so the upload can - # never create/replace files outside the datasets root through the link. + # An upload to a dataset name that is a symlink to an external directory must be refused (400) + # before any bytes are written. external = tmp_path / "external" external.mkdir() (ds_root / "linked").symlink_to(external, target_is_directory = True) @@ -580,8 +573,7 @@ def test_upload_through_symlinked_dataset_cannot_escape_root(client, ds_root, tm def test_delete_through_symlinked_dataset_cannot_escape_root(client, ds_root, tmp_path): - # End to end: a DELETE against an image inside a symlinked dataset dir is refused (400) and the - # external file it points at is NOT removed. + # A DELETE inside a symlinked dataset dir is refused (400) and the external file survives. external = tmp_path / "external" external.mkdir() victim = external / "victim.png" @@ -594,10 +586,8 @@ def test_delete_through_symlinked_dataset_cannot_escape_root(client, ds_root, tm def test_delete_image_with_glob_chars_only_removes_own_thumbs(client, ds_root): - # A legal uploaded filename may contain glob metacharacters ('[', ']', '*', '?'). Deleting it - # must remove ONLY its own thumbnails; interpolating the raw name into Path.glob would make - # "[ab].png" match "a.png_*.jpg"/"b.png_*.jpg" -- deleting a sibling's thumbs while leaving its - # own (literally "[ab].png_32.jpg") behind. + # Deleting a filename with glob metacharacters (e.g. "[ab].png") must remove only its own + # thumbnails, not a sibling's that the raw glob would spuriously match. from urllib.parse import quote folder = ds_root / "d" diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py index 4ff9bc23c7..18b9f51b62 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -241,11 +241,8 @@ def test_no_switch_keeps_engine_and_refreshes_reason(monkeypatch): def test_activate_serializes_switch_and_concurrent_query(monkeypatch): - # Regression for the check->unload->publish race: _activate releases _lock during the slow - # unload(), so without the transition lock a second _activate could observe the not-yet-updated - # active engine, take the "no change" branch, and return the engine the first call is - # concurrently unloading. Drive both paths on threads and assert the concurrent query is blocked - # until the switch completes (i.e. the whole transition is serialized). + # Regression: without the transition lock a second _activate during the slow unload() reads the + # not-yet-updated active engine and returns it. Assert the query is blocked until the switch ends. import threading r._active_engine_name = ENGINE_DIFFUSERS @@ -279,8 +276,7 @@ def test_activate_serializes_switch_and_concurrent_query(monkeypatch): q = threading.Thread(target = _query) q.start() - # Serialized: while the switch holds the transition lock the query cannot complete. Pre-fix it - # would return immediately (active is still diffusers), setting query_done at once. + # Serialized: the query cannot complete while the switch holds the transition lock. assert not query_done.wait(0.4) release_unload.set() diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index c06d8fd45a..6ae96d97e5 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -58,11 +58,8 @@ def test_discover_sidecar_overrides_metadata_row(tmp_path): def test_discover_empty_sidecar_suppresses_metadata_but_uses_instance_prompt(tmp_path): - # Clearing a metadata caption in the labeling grid writes an EMPTY sidecar tombstone. It must - # suppress the metadata caption (so the stale text is not resurrected) yet leave the image - # UNCAPTIONED so the dreambooth instance_prompt still applies. Treating "" as a present caption - # would skip the instance prompt AND drop the image, so a dataset whose every caption was - # cleared would fail with "No captioned images found". + # An empty sidecar tombstone must suppress the metadata caption yet leave the image uncaptioned + # so the dreambooth instance_prompt still applies (not drop the image). _touch(tmp_path / "cat.png") (tmp_path / "metadata.jsonl").write_text( json.dumps({"file_name": "cat.png", "text": "old metadata caption"}) + "\n", @@ -74,9 +71,8 @@ def test_discover_empty_sidecar_suppresses_metadata_but_uses_instance_prompt(tmp def test_discover_empty_sidecar_without_instance_prompt_skips_image(tmp_path): - # With no instance prompt an empty tombstone leaves the image uncaptioned, so it is skipped and - # the suppressed metadata caption is NOT resurrected -- while a sibling with a real caption is - # still discovered. + # With no instance prompt the tombstoned image is skipped (metadata not resurrected), while a + # sibling with a real caption is still discovered. _touch(tmp_path / "cat.png") _touch(tmp_path / "cap.png") (tmp_path / "metadata.jsonl").write_text( diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index a5a825ee00..c00d9e71d5 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -344,9 +344,8 @@ def test_generate_rejects_batch_seed_past_json_safe_range(client): client.post( "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} ) - # An explicit seed at the JS-safe cap with a batch derives per-image seeds - # (seed+1 ...) that exceed Number.MAX_SAFE_INTEGER and no longer round-trip - # through the gallery JSON recipe, so the request is rejected at the boundary. + # A seed at the cap with a batch derives per-image seeds (seed+1 ...) past the JSON-safe range, + # so the request is rejected. over = client.post( "/api/inference/images/generate", json = {"prompt": "p", "seed": 2**53 - 1, "batch_size": 2}, diff --git a/studio/backend/tests/test_image_gallery.py b/studio/backend/tests/test_image_gallery.py index b9b3c0ed49..51f107ffd2 100644 --- a/studio/backend/tests/test_image_gallery.py +++ b/studio/backend/tests/test_image_gallery.py @@ -148,10 +148,8 @@ def test_list_skips_recipe_missing_required_fields(tmp_path): def test_valid_callback_paginates_over_accepted_records(): - # A record that passes _read_meta (every required key present) but fails the caller's stricter - # schema check must be filtered BEFORE pagination, so offset/limit/has_more all count over the - # accepted domain. Otherwise a leading bad record returns an empty/short page with more still - # remaining, and the frontend (which advances by valid records) stalls at offset 0. + # ``valid`` must filter before pagination, so offset/limit/has_more count over the accepted + # domain; else a leading bad record returns a short page with more remaining and stalls scroll. _save_with_mtime("BAD", 300.0) # newest, sorts first _save_with_mtime("g1", 200.0) _save_with_mtime("g2", 100.0) @@ -159,8 +157,7 @@ def test_valid_callback_paginates_over_accepted_records(): def _valid(rec): return rec.get("prompt") != "BAD" - # First page of 2 over VALID records returns both good ones -- not [g1] (bad eating a slot) - # and not [] (bad filling the whole window). + # First page of 2 returns both good records, not [g1] or []. page = gallery.list_images(limit = 2, offset = 0, valid = _valid) assert [r["prompt"] for r in page] == ["g1", "g2"] # The has_more probe (limit + 1) sees no extra VALID record beyond the two returned. @@ -168,8 +165,7 @@ def test_valid_callback_paginates_over_accepted_records(): def test_valid_callback_leading_bad_record_does_not_stall_at_offset_zero(): - # Reproduces the exact stall: every record in the first window is schema-invalid. Without - # in-pager filtering the route returned images=[] with has_more=True at offset 0 forever. + # Every record in the first window is invalid; without in-pager filtering the route stalled. for i in range(3): _save_with_mtime(f"BAD{i}", 300.0 - i) # newest three are all invalid _save_with_mtime("good", 10.0) @@ -177,15 +173,13 @@ def test_valid_callback_leading_bad_record_does_not_stall_at_offset_zero(): def _valid(rec): return not str(rec.get("prompt", "")).startswith("BAD") - # limit+1 = 3: the pager must look PAST the invalid leaders and return the one good record, - # so has_more (len > limit) is False and the client advances off offset 0. + # The pager must look past the invalid leaders and return the one good record. records = gallery.list_images(limit = 2, offset = 0, valid = _valid) assert [r["prompt"] for r in records] == ["good"] def test_save_is_atomic_no_partial_png_on_publish_failure(monkeypatch): - # A crash between writing the bytes and publishing the file must leave neither a truncated - # {id}.png nor a leftover temp: the listing only ever sees fully-written records. + # A crash before publishing must leave neither a truncated {id}.png nor a leftover temp. def _boom(*a, **k): raise OSError("simulated rename failure") diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index 7163db424b..202b9a1b72 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -181,11 +181,9 @@ def test_local_task_tags_video_pipeline_dir(tmp_path): def test_local_task_tags_video_single_file_checkpoint(tmp_path): - # A dir whose name matches a video family holding a bare single-file .safetensors (no - # model_index.json) IS loadable: the video load route reinterprets a sole single-file local - # pick as a single_file load (routes/video.py), validating BEFORE it touches the GPU. So it - # must be tagged text-to-video and surfaced in the Video On-Device picker -- not left task=null - # and hidden, which would make the advertised-and-loadable checkpoint unusable. + # A video-family dir holding a bare single-file .safetensors (no model_index.json) is loadable + # (the route loads it as a single_file), so it must be tagged text-to-video and surfaced, not + # left task=null and hidden. d = tmp_path / "ltx-loose" _touch(d / "ltx-2.safetensors") # loose weights, no model_index.json assert ( diff --git a/studio/backend/tests/test_sd_cpp_install.py b/studio/backend/tests/test_sd_cpp_install.py index 3a061c9fcf..89404fa3ae 100644 --- a/studio/backend/tests/test_sd_cpp_install.py +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -274,11 +274,8 @@ def test_install_into_empty_dir_claims_ownership(tmp_path, monkeypatch): def test_install_into_nonempty_unowned_dir_is_refused(tmp_path, monkeypatch): - # A pre-existing, non-empty directory that Studio did not create (e.g. a user's own - # stable-diffusion.cpp checkout) must NOT be extracted into. Merging the release into it would - # overwrite or mix our binaries into the user's working tree, and leaving it unowned only stops - # the uninstaller from deleting it later. install() refuses up front and leaves the dir untouched - # so the user can point us at a fresh/empty location. + # A pre-existing, non-empty directory Studio did not create (e.g. a user's own checkout) must + # not be extracted into; install() refuses up front and leaves it untouched. zb = _zip_with_sd_cli() _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest()) target = tmp_path / "stable-diffusion.cpp" diff --git a/studio/backend/tests/test_training_start_offload.py b/studio/backend/tests/test_training_start_offload.py index 1d10355dda..1b53868427 100644 --- a/studio/backend/tests/test_training_start_offload.py +++ b/studio/backend/tests/test_training_start_offload.py @@ -116,11 +116,8 @@ def test_backend_start_guard_blocks_overlapping_starts(): 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. + # A run reserved in start_training but not yet spawned must already read as active, else the + # load/start guards would let another pipeline race it for the just-freed VRAM. from core.training.training import TrainingBackend backend = TrainingBackend() diff --git a/studio/backend/tests/test_video_gallery.py b/studio/backend/tests/test_video_gallery.py index f86dbe3bcd..6ca8ea2e35 100644 --- a/studio/backend/tests/test_video_gallery.py +++ b/studio/backend/tests/test_video_gallery.py @@ -200,10 +200,8 @@ def test_list_skips_corrupt_sidecar(): def test_valid_callback_paginates_over_accepted_records(): - # A sidecar that parses as JSON (so the read accepts it) but fails the caller's stricter schema - # check must be filtered BEFORE pagination, so offset/limit/has_more count over accepted records - # only. Otherwise a leading bad record returns a short/empty page with more remaining and stalls - # infinite scroll at offset 0. + # ``valid`` must filter before pagination, so offset/limit/has_more count over accepted records; + # else a leading bad record returns a short page with more remaining and stalls scroll. _save_with_mtime("BAD", 300.0) # newest, sorts first _save_with_mtime("g1", 200.0) _save_with_mtime("g2", 100.0) @@ -231,10 +229,8 @@ def test_valid_callback_leading_bad_records_do_not_stall_at_offset_zero(): def test_save_leaves_no_orphan_mp4_when_sidecar_publish_fails(monkeypatch): - # The sidecar is the pair's commit marker. If it fails to publish, the MP4 must NOT be left - # behind as an invisible orphan (list_videos would skip it and gallery delete could never reach - # it). Fail the SECOND os.replace (the sidecar) after the mp4 is renamed in, and assert nothing - # is stranded. + # If the sidecar (the pair's commit marker) fails to publish, the MP4 must not be left as an + # invisible orphan. Fail the second os.replace and assert nothing is stranded. real_replace = gallery.os.replace calls = {"n": 0} diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index 275ca2fb6d..0793d13788 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -111,8 +111,7 @@ class _FakeBackend(video_module.VideoBackend): 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.") - # 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). + # Non-GGUF loads are gated to unsloth/* repos, the official bases, and existing local paths. trusted = model_path.lower().startswith(("unsloth/", "lightricks/")) or ( Path(model_path).expanduser().exists() ) @@ -386,10 +385,8 @@ def test_load_progress_route(client): 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. + # A local video-family dir with 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. d = tmp_path / "ltx-2.3-local" d.mkdir() (d / "ltx-dit.safetensors").write_bytes(b"0") @@ -404,8 +401,7 @@ def test_load_local_single_file_dir_routes_through_single_file(client, tmp_path) 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. + # A real diffusers directory (has model_index.json) is left as a pipeline load. d = tmp_path / "ltx-2.3-pipeline" d.mkdir() (d / "model_index.json").write_text("{}") diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 7338f23b48..2f6d70936d 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1450,13 +1450,10 @@ 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. + // Re-enter the per-step poll for a generation already in flight on the backend that this page + // did not start (another client, or a reload mid-generate), instead of showing a stale idle + // view. generate-progress carries no terminal record, so refresh the gallery on completion to + // merge any image saved after the mount fetch. Separate from handleGenerate's own loop. const resumeGeneratePoll = useCallback(() => { if (genPollTimer.current) clearInterval(genPollTimer.current); if (genVisibilityListener.current) @@ -1477,8 +1474,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { 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. + // Re-fetch the first page to merge images the finished run saved, and resync status. void loadGallery(); void refreshStatus(); return; @@ -1518,11 +1514,8 @@ 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. + // Resume tracking a generation started elsewhere (another client, or before a reload) so the + // page shows the in-flight run instead of a stale idle view. Mirrors the video page. try { const g = await getGenerateProgress(); if (g.active) { @@ -1561,11 +1554,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) { if (lastLoad.current) return; if (seededResident.current === repoId) return; seededResident.current = 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. + // Seed from base_repo (the resolved diffusers base, holding the family), not repo_id: a + // GGUF/single_file/local-path resident has a repo_id with no family substring, so + // defaultsFor(repoId) would fall back to the wrong distilled few-step/no-CFG recipe. const d = defaultsFor(status?.base_repo ?? repoId); setSteps(d.steps); setGuidance(d.guidance); diff --git a/studio/install_sd_cpp_prebuilt.py b/studio/install_sd_cpp_prebuilt.py index 1923cd3f8e..67ff5f6635 100644 --- a/studio/install_sd_cpp_prebuilt.py +++ b/studio/install_sd_cpp_prebuilt.py @@ -401,12 +401,9 @@ def install( has no ``sd-cli``. """ target = install_dir or default_install_dir() - # Decide up front whether this install may claim ownership of ``target``. We only mark a - # directory as Studio-owned (and therefore eligible for the uninstaller's recursive delete) - # when this install actually created it or it was empty -- NEVER when it already held a user's - # own stable-diffusion.cpp checkout or unrelated files. Adopting a pre-existing, unowned, - # non-empty directory would let a later uninstall wipe the user's own work. A directory that - # already carries our marker (a prior Studio install / upgrade) stays owned. + # Only claim ownership of ``target`` (marking it for the uninstaller's recursive delete) when + # this install created it or it was empty, or it already carries our marker -- never adopt a + # pre-existing non-empty dir, else a later uninstall could wipe a user's own checkout. marker = target / ".unsloth-studio-owned" _may_own = True if target.exists(): @@ -418,11 +415,9 @@ def install( _pre_existing_entries = True # Empty dir, or one we already own, may be (re)claimed; a non-empty unowned dir may not. _may_own = (not _pre_existing_entries) or marker.is_file() - # Refuse to extract into a pre-existing, non-empty directory we do not own (a user's own - # stable-diffusion.cpp checkout, or unrelated files beside a custom Studio root). Not writing - # the ownership marker only protects the uninstaller; extracting the release here would still - # merge our binaries into the user's working tree and can overwrite same-named files. Fail with - # a clear message so the user points us at a fresh/empty location instead of corrupting theirs. + # Refuse to extract into a pre-existing, non-empty directory we do not own: merging the release + # in would overwrite or mix our binaries into the user's own files. Fail so they point us at a + # fresh/empty location. if not _may_own: raise RuntimeError( f"sd.cpp install target already exists and is not a Studio-managed directory: {target}. " @@ -470,11 +465,8 @@ def install( _make_executable(sd_server) if sd_server is not None: print(f"installed sd-server -> {sd_server}", flush = True) - # Ownership marker (the same one setup.sh/_is_studio_root use, and setup.ps1 writes into the - # Node sibling dir) so the uninstaller can tell a Studio-installed sd.cpp from a user's own - # stable-diffusion.cpp checkout beside a custom Studio root, and delete only ours. Written only - # when this install created the directory or it was empty (see _may_own above): a pre-existing, - # unowned, non-empty directory keeps its unowned status so the uninstaller leaves it alone. + # Ownership marker (the same one setup.sh/_is_studio_root use) so the uninstaller deletes only + # Studio-installed sd.cpp, not a user's own checkout. Written only when _may_own (see above). if _may_own: try: marker.touch() From 899465ed8079ac164a7961ddb6974438687bb62c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 06:31:00 +0000 Subject: [PATCH 07/10] Studio: close arbiter load-registration race and surface native progress + local pipeline folders Publish native sd.cpp generate progress (_gen) before LoRA resolution so a reload probe reads active during setup, matching the diffusers path. Register the diffusion/video GPU load under the arbiter lock (acquire_for now takes a register callback) so a competing acquire cannot evict an owner before its load is marked in-flight and let two loaders allocate VRAM at once. Admit local diffusers pipeline folders (root model_index.json, weights in component subdirs) in the local model scan so they reach task tagging and the On Device picker. --- studio/backend/core/inference/gpu_arbiter.py | 16 ++++- .../backend/core/inference/sd_cpp_backend.py | 7 +- studio/backend/routes/inference.py | 52 ++++++++------ studio/backend/routes/models.py | 10 ++- studio/backend/routes/video.py | 44 +++++++----- studio/backend/tests/test_diffusion_routes.py | 8 ++- studio/backend/tests/test_gpu_arbiter.py | 72 +++++++++++++++++++ .../backend/tests/test_local_model_format.py | 19 +++++ studio/backend/tests/test_sd_cpp_backend.py | 26 +++++++ studio/backend/tests/test_video_routes.py | 8 ++- 10 files changed, 216 insertions(+), 46 deletions(-) diff --git a/studio/backend/core/inference/gpu_arbiter.py b/studio/backend/core/inference/gpu_arbiter.py index 0757eb83e8..53da970ce6 100644 --- a/studio/backend/core/inference/gpu_arbiter.py +++ b/studio/backend/core/inference/gpu_arbiter.py @@ -12,7 +12,7 @@ under the lock, so a transfer is atomic vs other acquires. from __future__ import annotations import threading -from typing import Optional +from typing import Any, Callable, Optional from loggers import get_logger @@ -64,8 +64,17 @@ def _evict_video() -> None: _EVICTORS = {CHAT: _evict_chat, DIFFUSION: _evict_diffusion, VIDEO: _evict_video} -def acquire_for(owner: str) -> None: - """Make ``owner`` the sole GPU owner, evicting the other if it holds it.""" +def acquire_for(owner: str, register: Optional[Callable[[], Any]] = None) -> Any: + """Make ``owner`` the sole GPU owner, evicting the other if it holds it. + + ``register``, if given, runs under the arbiter lock right after ownership transfers, + and its return value is returned. Registering the in-flight load HERE -- not after + ``acquire_for`` returns -- closes the window where a competing acquire could evict this + owner before its load is marked in-flight: eviction would then find nothing to cancel + and both loaders would allocate VRAM at once. ``register`` must be quick (it holds the + lock) and must not re-enter the arbiter. If it raises, ownership stays with ``owner`` -- + matching the pre-register behaviour where a failed load left the handoff in place. + """ global _owner if owner not in _EVICTORS: raise ValueError(f"unknown GPU owner: {owner!r}") @@ -74,6 +83,7 @@ def acquire_for(owner: str) -> None: logger.info("gpu_arbiter: evicting %s for %s", _owner, owner) _EVICTORS[_owner]() _owner = owner + return register() if register is not None else None def release(owner: str) -> None: diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index a6847ae020..3fb40b5c7f 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -764,6 +764,12 @@ class SdCppDiffusionBackend: self._state = None raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) self._active_generate_cancel = cancel + # Publish an active (step 0) state now, before the slow pre-generate setup + # (LoRA listing/download), so a reload's progress probe doesn't read idle + # while this generation already holds _generate_lock and let a second generate + # queue behind it. The parsed sd-cli progress lines advance this step count. + # Mirrors DiffusionBackend.generate, which publishes _gen before its setup. + self._gen = _SdGen(total_steps = int(steps)) try: if seed is None: seed = int.from_bytes(os.urandom(6), "big") & ((1 << 53) - 1) @@ -789,7 +795,6 @@ class SdCppDiffusionBackend: lora_resolved = diffusion_lora.resolve_specs( active_loras, hf_token = state.hf_token, cancel_event = cancel ) - self._gen = _SdGen(total_steps = int(steps)) if state.mode == "server" and state.server is not None: images, seeds = self._generate_server( state, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index fab3652b7f..20aca01fe6 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -14338,10 +14338,36 @@ async def load_diffusion_model( # engine name -- else we'd evict a resident chat model for a load that can't use the GPU. device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device) needs_gpu = device != "cpu" + + def _begin_load(): + # Kicks the (slow) load onto a background thread and returns at once (the client + # polls images/load-progress); begin_load itself validates network-free. + return engine.begin_load( + request.model_path, + gguf_filename = request.gguf_filename, + base_repo = request.base_repo, + family_override = request.family_override, + hf_token = request.hf_token, + cpu_offload = request.cpu_offload, + memory_mode = request.memory_mode, + speed_mode = request.speed_mode, + text_encoder_quant = request.text_encoder_quant, + transformer_quant = request.transformer_quant, + transformer_quant_fast_accum = request.transformer_quant_fast_accum, + transformer_prequant_path = request.transformer_prequant_path, + attention_backend = request.attention_backend, + transformer_cache = request.transformer_cache, + transformer_cache_threshold = request.transformer_cache_threshold, + model_kind = kind, + ) + if needs_gpu: - # Then kick the (slow) load onto a background thread and return at once -- - # the client polls images/load-progress. - await asyncio.to_thread(acquire_for, DIFFUSION) + # Register the in-flight load UNDER the arbiter lock (not after acquire_for + # returns): a competing Video/chat acquire in that gap would otherwise evict + # DIFFUSION before begin_load marks a load in-flight, so eviction finds nothing + # to cancel and both loaders allocate VRAM at once. begin_load returns at once, + # so the lock is held only briefly. + status_dict = await asyncio.to_thread(acquire_for, DIFFUSION, _begin_load) else: # A CPU-only native load never touches the GPU, so it neither acquires nor is # tracked by the arbiter. But switching here FROM a previous diffusers/GPU load @@ -14349,25 +14375,7 @@ async def load_diffusion_model( # "evict" this CPU model for no reason. Release that stale ownership -- release() # is owner-guarded, so it's a no-op when diffusion never owned the GPU. await asyncio.to_thread(release, DIFFUSION) - status_dict = await asyncio.to_thread( - engine.begin_load, - request.model_path, - gguf_filename = request.gguf_filename, - base_repo = request.base_repo, - family_override = request.family_override, - hf_token = request.hf_token, - cpu_offload = request.cpu_offload, - memory_mode = request.memory_mode, - speed_mode = request.speed_mode, - text_encoder_quant = request.text_encoder_quant, - transformer_quant = request.transformer_quant, - transformer_quant_fast_accum = request.transformer_quant_fast_accum, - transformer_prequant_path = request.transformer_prequant_path, - attention_backend = request.attention_backend, - transformer_cache = request.transformer_cache, - transformer_cache_threshold = request.transformer_cache_threshold, - model_kind = kind, - ) + status_dict = await asyncio.to_thread(_begin_load) return DiffusionStatusResponse(**annotate_status(status_dict)) except (ValueError, FileNotFoundError) as exc: raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc))) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index fe11446561..d1337e4210 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -338,7 +338,15 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca has_config = (child / "config.json").exists() or ( child / "adapter_config.json" ).exists() - has_model_files = has_gguf or has_non_gguf_weights or has_config + # A standard diffusers PIPELINE folder keeps its weights/configs in component + # subdirs (transformer/, vae/, ...) and carries only model_index.json at the + # root, so the checks above miss it. The Images/Video load path accepts such a + # local pipeline dir, so admit it here too (task tagging then classifies it via + # _local_is_diffusers); otherwise it is hidden from the On Device picker. + has_pipeline_index = (child / "model_index.json").is_file() + has_model_files = ( + has_gguf or has_non_gguf_weights or has_config or has_pipeline_index + ) except OSError: # Skip unreadable children rather than failing the scan. continue diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 7e239151c5..4abc513935 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -112,26 +112,36 @@ async def load_video_model( # Take the GPU from chat only for a non-CPU load; a CPU load never touches GPU memory, # so key off the device. Release stale VIDEO ownership on a CPU load (owner-guarded no-op). device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device) + + def _begin_load(): + # Kicks the (slow) load onto a background thread and returns at once; + # begin_load itself validates network-free. + return backend.begin_load( + request.model_path, + gguf_filename = request.gguf_filename, + base_repo = request.base_repo, + family_override = request.family_override, + hf_token = request.hf_token, + memory_mode = request.memory_mode, + speed_mode = request.speed_mode, + attention_backend = request.attention_backend, + transformer_cache = request.transformer_cache, + transformer_cache_threshold = request.transformer_cache_threshold, + transformer_quant = request.transformer_quant, + text_encoder_quant = request.text_encoder_quant, + model_kind = kind, + ) + if device != "cpu": - await asyncio.to_thread(acquire_for, VIDEO) + # Register the in-flight load UNDER the arbiter lock (not after acquire_for + # returns): a competing Images/chat acquire in that gap would otherwise evict + # VIDEO before begin_load marks a load in-flight, so eviction finds nothing to + # cancel and both loaders allocate VRAM at once. begin_load returns at once, so + # the lock is held only briefly. Mirrors the images/load handoff. + status_dict = await asyncio.to_thread(acquire_for, VIDEO, _begin_load) else: await asyncio.to_thread(release, VIDEO) - status_dict = await asyncio.to_thread( - backend.begin_load, - request.model_path, - gguf_filename = request.gguf_filename, - base_repo = request.base_repo, - family_override = request.family_override, - hf_token = request.hf_token, - memory_mode = request.memory_mode, - speed_mode = request.speed_mode, - attention_backend = request.attention_backend, - transformer_cache = request.transformer_cache, - transformer_cache_threshold = request.transformer_cache_threshold, - transformer_quant = request.transformer_quant, - text_encoder_quant = request.text_encoder_quant, - model_kind = kind, - ) + status_dict = await asyncio.to_thread(_begin_load) return VideoStatusResponse(**status_dict) except (ValueError, FileNotFoundError) as exc: raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc))) diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index c00d9e71d5..322c188534 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -789,7 +789,13 @@ def _force_engine(monkeypatch, backend, *, engine_name, device): devmod, "resolve_diffusion_device_target", lambda: _types.SimpleNamespace(device = device) ) acquired: list = [] - monkeypatch.setattr(gpu_arbiter, "acquire_for", lambda role: acquired.append(role)) + + def _fake_acquire(role, register = None): + # Mirror the real arbiter: record the handoff and run the (registered) load under it. + acquired.append(role) + return register() if register is not None else None + + monkeypatch.setattr(gpu_arbiter, "acquire_for", _fake_acquire) return acquired diff --git a/studio/backend/tests/test_gpu_arbiter.py b/studio/backend/tests/test_gpu_arbiter.py index 0487830def..a73e786f90 100644 --- a/studio/backend/tests/test_gpu_arbiter.py +++ b/studio/backend/tests/test_gpu_arbiter.py @@ -104,3 +104,75 @@ def test_evict_chat_unloads_a_still_loading_chat_backend(monkeypatch): arb._evict_chat() assert unloaded == [True] # still-loading chat backend was unloaded, not skipped + + +def test_register_runs_under_ownership_and_returns_result(calls): + # A register callback runs after ownership transfers (owner already set) and its + # return value is forwarded -- the route uses this to register the in-flight load. + seen_owner: list = [] + + def register(): + seen_owner.append(arb.current_owner()) + return "status-dict" + + result = arb.acquire_for(arb.DIFFUSION, register) + assert result == "status-dict" + assert seen_owner == [arb.DIFFUSION] + assert arb.current_owner() == arb.DIFFUSION + + +def test_register_failure_leaves_ownership_in_place(calls): + # A failing register (e.g. begin_load reporting a load already in progress) propagates + # but must not drop ownership -- the prior handoff (chat already evicted) stands. + arb.acquire_for(arb.CHAT) + + def register(): + raise RuntimeError("A diffusion load is already in progress.") + + with pytest.raises(RuntimeError): + arb.acquire_for(arb.DIFFUSION, register) + assert calls == ["evict-chat"] + assert arb.current_owner() == arb.DIFFUSION + + +def test_competing_acquire_blocks_until_register_completes(monkeypatch): + # The window this closes: while DIFFUSION registers its load, a competing VIDEO acquire + # must not evict DIFFUSION until the load is marked in-flight. Holding the lock across + # register makes the competitor wait, so eviction never races an unregistered load. + import threading + import time + + monkeypatch.setattr(arb, "_owner", None) + evicted: list = [] + monkeypatch.setitem(arb._EVICTORS, arb.DIFFUSION, lambda: evicted.append("evict-diffusion")) + monkeypatch.setitem(arb._EVICTORS, arb.VIDEO, lambda: evicted.append("evict-video")) + + in_register = threading.Event() + release_register = threading.Event() + + def register(): + in_register.set() + # Hold the arbiter lock here; a competing acquire_for(VIDEO) must block until we return. + assert release_register.wait(2.0) + return "loading" + + loader = threading.Thread(target = lambda: arb.acquire_for(arb.DIFFUSION, register)) + loader.start() + assert in_register.wait(2.0) + + competitor_done = threading.Event() + threading.Thread( + target = lambda: (arb.acquire_for(arb.VIDEO), competitor_done.set()), + ).start() + + # The competitor cannot evict DIFFUSION while register still holds the lock. + time.sleep(0.1) + assert evicted == [] + assert not competitor_done.is_set() + + # Let register finish; ownership is now safely registered, so the competitor proceeds. + release_register.set() + loader.join(2.0) + assert competitor_done.wait(2.0) + assert evicted == ["evict-diffusion"] + assert arb.current_owner() == arb.VIDEO diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index 202b9a1b72..1fbd7efcd6 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -117,6 +117,25 @@ def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path): assert row.model_format == "gguf" +def test_scan_models_dir_surfaces_diffusers_pipeline_folder(tmp_path): + # A standard diffusers PIPELINE folder keeps its weights/configs in component subdirs + # (transformer/, vae/, ...) and carries only model_index.json at the root. The Images/Video + # load path accepts such a local pipeline dir, so the scan must surface it -- otherwise the + # weights-in-subdirs layout is missed and it never reaches task tagging / the On Device + # picker. It is not a GGUF, so model_format stays None (task tagging classifies it later). + root = tmp_path / "models" + pipe = root / "my-pipeline" + _touch(pipe / "model_index.json") + _touch(pipe / "transformer" / "config.json") + _touch(pipe / "transformer" / "diffusion_pytorch_model.safetensors") + _touch(pipe / "vae" / "diffusion_pytorch_model.safetensors") + + rows = {Path(m.path).name: m for m in models_route._scan_models_dir(root)} + + assert "my-pipeline" in rows + assert rows["my-pipeline"].model_format is None + + # ── Images picker task tag for local (non-GGUF) diffusers models ────────────── from models.models import LocalModelInfo # noqa: E402 diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 1d0268effa..1cb01df417 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -279,6 +279,32 @@ def test_generate_progress_tracks_parsed_steps(): assert b.generate_progress()["step"] == 4 +def test_generate_publishes_progress_before_lora_resolution(monkeypatch): + # Native LoRA resolution (listing/downloading a not-yet-cached adapter) happens during the + # pre-generate setup while _generate_lock is already held. A reload/progress probe in that + # window must read ACTIVE, not idle, or the UI queues a second generate behind the first. + # So _gen is published before LoRA resolution, mirroring the diffusers path. + from core.inference import diffusion_lora + + eng = _FakeEngine() + b = _loaded_backend(engine = eng) + monkeypatch.setattr(diffusion_lora, "supports_lora", lambda **_k: True) + + seen: dict = {} + + def _resolve(active, *, hf_token = None, cancel_event = None): + # Mid-setup: the in-flight generation must already be reported as active. + seen["progress"] = b.generate_progress() + return [] + + monkeypatch.setattr(diffusion_lora, "resolve_specs", _resolve) + + out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, loras = [("some/lora", 1.0)]) + assert out["images"] + assert seen["progress"]["active"] is True + assert seen["progress"]["total_steps"] == 8 + + # ── load validation + binary install ────────────────────────────────────────── diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index 0793d13788..f7bc076225 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -263,7 +263,13 @@ def test_load_happy_path_and_arbiter_acquired(client, monkeypatch): devmod, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cuda") ) acquired: list = [] - monkeypatch.setattr(gpu_arbiter, "acquire_for", lambda role: acquired.append(role)) + + def _fake_acquire(role, register = None): + # Mirror the real arbiter: record the handoff and run the (registered) load under it. + acquired.append(role) + return register() if register is not None else None + + monkeypatch.setattr(gpu_arbiter, "acquire_for", _fake_acquire) resp = client.post( "/api/inference/video/load", From 11330b8de508746cf5e7b5633276cb68b9e43d79 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 13 Jul 2026 06:31:46 +0000 Subject: [PATCH 08/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/models.py | 4 +--- studio/backend/tests/test_sd_cpp_backend.py | 7 ++++++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index d1337e4210..f9a6e5833c 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -344,9 +344,7 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca # local pipeline dir, so admit it here too (task tagging then classifies it via # _local_is_diffusers); otherwise it is hidden from the On Device picker. has_pipeline_index = (child / "model_index.json").is_file() - has_model_files = ( - has_gguf or has_non_gguf_weights or has_config or has_pipeline_index - ) + has_model_files = has_gguf or has_non_gguf_weights or has_config or has_pipeline_index except OSError: # Skip unreadable children rather than failing the scan. continue diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 1cb01df417..40a20a58a2 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -292,7 +292,12 @@ def test_generate_publishes_progress_before_lora_resolution(monkeypatch): seen: dict = {} - def _resolve(active, *, hf_token = None, cancel_event = None): + def _resolve( + active, + *, + hf_token = None, + cancel_event = None, + ): # Mid-setup: the in-flight generation must already be reported as active. seen["progress"] = b.generate_progress() return [] From 53668c9f66e6dc49258850c4a9b1030de8b1e89c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 07:26:18 +0000 Subject: [PATCH 09/10] Studio: include marker-owned custom sd.cpp roots in the uninstall stop scan; harden Pester install against the nuget.exe PSGallery bootstrap --- .../studio-windows-inference-smoke.yml | 27 +++++++++++++------ scripts/uninstall.ps1 | 14 +++++++++- 2 files changed, 32 insertions(+), 9 deletions(-) diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 0453c9212a..821ef10b70 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -1705,15 +1705,26 @@ jobs: - name: Install Pester v5 shell: pwsh run: | - # PSGallery is intermittently absent from the repository list on GitHub's Windows - # runners, which makes `Set-PSRepository PSGallery` fail with "No repository with the - # name 'PSGallery' was found." Re-register the default gallery first so the policy - # change and module install below always have a repository to target. - if (-not (Get-PSRepository -Name PSGallery -ErrorAction SilentlyContinue)) { - Register-PSRepository -Default -ErrorAction SilentlyContinue + $ErrorActionPreference = 'Stop' + # Prefer PSResourceGet (preinstalled on the runner's PowerShell 7): it resolves PSGallery + # over HTTPS directly and avoids the legacy PackageManagement/nuget.exe bootstrap, which + # intermittently fails on GitHub's Windows runners with + # "NuGet.Commands.CommandException: Missing option value for: '-source'" (the install then + # aborts and Pester never runs). Fall back to the classic path when PSResourceGet is absent. + if (Get-Command Install-PSResource -ErrorAction SilentlyContinue) { + Install-PSResource -Name Pester -Version '[5.5.0,)' -Repository PSGallery -TrustRepository -Scope CurrentUser + } else { + [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12 + # PSGallery is intermittently absent from the repository list on GitHub's Windows + # runners, which makes `Set-PSRepository PSGallery` fail with "No repository with the + # name 'PSGallery' was found." Re-register the default gallery first so the policy + # change and module install below always have a repository to target. + if (-not (Get-PSRepository -Name PSGallery -ErrorAction SilentlyContinue)) { + Register-PSRepository -Default -ErrorAction SilentlyContinue + } + Set-PSRepository PSGallery -InstallationPolicy Trusted + Install-Module Pester -MinimumVersion 5.5.0 -Force -SkipPublisherCheck -Scope CurrentUser } - Set-PSRepository PSGallery -InstallationPolicy Trusted - Install-Module Pester -MinimumVersion 5.5.0 -Force -SkipPublisherCheck -Scope CurrentUser Import-Module Pester -MinimumVersion 5.5.0 Get-Module Pester | Select-Object Name, Version | Format-Table diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 61bf73b5b4..46eb48cb3b 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -370,9 +370,21 @@ function Uninstall-UnslothStudio { if ($defaultSdCpp -and (Test-Path -LiteralPath $defaultSdCpp) -and (Test-Path -LiteralPath (Join-Path $defaultSdCpp ".unsloth-studio-owned") -PathType Leaf)) { $defaultSdCppToStop = $defaultSdCpp } + # Custom/env-mode sd.cpp builds sit BESIDE each custom root at \stable-diffusion.cpp + # (find_sd_cpp_binary resolves from UNSLOTH_STUDIO_HOME.parent), so a running owned sd-server + # there is outside every root above ($knownRoots holds the custom root, not its sibling). We + # delete those marker-owned dirs below, so add them to the handle scan too -- gated on the same + # owner marker as the delete, matching the default-root handling. + $customSdCppToStop = @() + foreach ($r in $customRoots) { + $sdc = Join-Path (Split-Path -LiteralPath $r -Parent) "stable-diffusion.cpp" + if ((Test-Path -LiteralPath $sdc) -and (Test-Path -LiteralPath (Join-Path $sdc ".unsloth-studio-owned") -PathType Leaf)) { + $customSdCppToStop += $sdc + } + } # Also stop anything holding a handle on the exact paths we delete (llama-server, # the CLI shim, an mp-fork python with a venv DLL) so the dir delete isn't refused. - _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache, $defaultNode) + @($defaultSdCppToStop | Where-Object { $_ })) + _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache, $defaultNode) + @($defaultSdCppToStop | Where-Object { $_ }) + @($customSdCppToStop)) # ── Remove custom-root install trees ── _Step "Removing data and install directories..." From b127256eb4d44c4175502e8b14004b7a2a80dc89 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 08:26:43 +0000 Subject: [PATCH 10/10] Studio: surface local pipeline scan roots, tag single-file checkpoints by filename, mark companion-only pipelines partial - _scan_models_dir: admit a scan folder that is itself a diffusers pipeline (root model_index.json, weights in transformer/ vae/ subdirs). _is_model_directory rejects such a root, so the child scan would list the component subdirs as bogus models and hide the real pipeline; treat the root as one model via _local_pipeline_index. - _local_is_diffusers / _local_model_task: include the sole checkpoint filename in the family-detection needles (_local_family_needles, resolved via resolve_local_single_file). A generically named folder holding one loadable qwen-image-*.safetensors / ltx-*.safetensors identifies its family only from the filename; the load route already resolves that file, so tag it or the task-scoped picker (which rejects task=null) hides the on-device model. - list_cached_models: mark a companion-only base snapshot partial. A GGUF image load prefetches the base repo's VAE / text-encoder / model_index.json but skips the transformer (the GGUF supplies it); the snapshot has a pipeline manifest yet is not a loadable BF16 pipeline, and _cached_repo_partial misses it. _repo_pipeline_missing_denoiser flags a pipeline snapshot whose transformer/ or unet/ component carries no weight, so the picker drops it instead of advertising it as fully on-device. --- studio/backend/routes/models.py | 112 +++++++++++++++--- .../backend/tests/test_cached_gguf_routes.py | 45 ++++++- .../backend/tests/test_local_model_format.py | 36 ++++++ 3 files changed, 175 insertions(+), 18 deletions(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index f9a6e5833c..82b018d741 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -304,11 +304,26 @@ def _has_non_gguf_weights(path: Path) -> bool: return False +def _local_pipeline_index(d: Path) -> bool: + """True when *d* is a standard diffusers PIPELINE root: component weights/configs live in + subdirs (``transformer/``, ``vae/``, ...) under a top-level ``model_index.json``, so + ``_is_model_directory`` (which wants a root config + loose weights) rejects it.""" + try: + return (d / "model_index.json").is_file() + except OSError: + return False + + def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[LocalModelInfo]: if not models_dir.exists() or not models_dir.is_dir(): return [] - _is_self_model = _is_model_directory(models_dir) + # A scan folder can point directly at a diffusers PIPELINE dir, not only at a parent of + # model repos. _is_model_directory rejects such a root (weights live in transformer/, vae/, + # ... not beside a root config.json), so without this the child scan below surfaces the + # component subdirs as bogus models and hides the real pipeline. The Images/Video load path + # loads a local pipeline dir, so admit the root as one model (task tagging classifies it). + _is_self_model = _is_model_directory(models_dir) or _local_pipeline_index(models_dir) if _is_self_model: try: @@ -343,7 +358,7 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca # root, so the checks above miss it. The Images/Video load path accepts such a # local pipeline dir, so admit it here too (task tagging then classifies it via # _local_is_diffusers); otherwise it is hidden from the On Device picker. - has_pipeline_index = (child / "model_index.json").is_file() + has_pipeline_index = _local_pipeline_index(child) has_model_files = has_gguf or has_non_gguf_weights or has_config or has_pipeline_index except OSError: # Skip unreadable children rather than failing the scan. @@ -3261,6 +3276,25 @@ def _repo_gguf_task(repo_info) -> Optional[str]: return None +def _local_family_needles(model: "LocalModelInfo") -> tuple[str, ...]: + """Family-detection hints for a local (non-GGUF) checkpoint: its model id, display name, and + leaf directory name, plus -- for a bare single-file directory -- the sole checkpoint's + filename. A generically named folder holding one loadable ``qwen-image-*.safetensors`` / + ``ltx-*.safetensors`` identifies its family only from that filename, and the load route already + resolves that sole file via ``resolve_local_single_file``, so feed the same name here or a + task-scoped Images/Video picker (which rejects ``task: null``) hides the on-device model. Only + the basename is used (not the parent path), so a family token in a parent dir can't match.""" + needles = [model.model_id, model.display_name, Path(model.id).name] + try: + from core.inference.diffusion import resolve_local_single_file + single = resolve_local_single_file(model.path) + if single: + needles.append(single) + except Exception: + pass + return tuple(n for n in needles if n) + + def _local_model_task(model: "LocalModelInfo") -> Optional[str]: """Classify a local model into an HF pipeline task so the Images picker can filter. @@ -3293,12 +3327,8 @@ def _local_model_task(model: "LocalModelInfo") -> Optional[str]: try: from core.inference.video import _is_trusted_video_repo from core.inference.video_families import detect_video_family - for needle in (model.model_id, model.display_name, Path(model.id).name): - if ( - needle - and detect_video_family(needle) is not None - and _is_trusted_video_repo(path) - ): + for needle in _local_family_needles(model): + if detect_video_family(needle) is not None and _is_trusted_video_repo(path): return _VIDEO_GEN_TASK except Exception: pass @@ -3311,8 +3341,9 @@ def _local_is_diffusers(model: "LocalModelInfo") -> bool: ``_repo_is_diffusers`` heuristics: a full pipeline carries a top-level ``model_index.json``, while single-file / safetensors image checkpoints ship none, so fall back to the model id resolving to a known diffusion family (the same resolver the - Images backend loads from). Family detection uses the clean model id / name, not the - on-disk path, so a parent directory keyword can't spuriously match.""" + Images backend loads from). Family detection uses the clean model id / name and the sole + checkpoint's filename (via _local_family_needles), not the on-disk path, so a parent + directory keyword can't spuriously match while a filename-only family is still caught.""" try: p = Path(model.path) if p.is_dir() and (p / "model_index.json").is_file(): @@ -3321,20 +3352,21 @@ def _local_is_diffusers(model: "LocalModelInfo") -> bool: pass try: from core.inference.diffusion_families import detect_family - for needle in (model.model_id, model.display_name, Path(model.id).name): - if needle and detect_family(needle) is not None: + for needle in _local_family_needles(model): + if detect_family(needle) is not None: return True except Exception: pass # A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json) has no # pipeline index and no image family, so the checks above miss it. The video load route loads it # as a single_file (routes/video.py), so it must be surfaced or _local_model_task returns - # task=null and the picker hides it. Match clean id / name needles (not the raw path) so a - # parent-dir token can't spuriously match; _local_model_task then routes it to text-to-video. + # task=null and the picker hides it. Match clean id / name / checkpoint-filename needles (not + # the raw path) so a parent-dir token can't spuriously match; _local_model_task then routes it + # to text-to-video. try: from core.inference.video_families import detect_video_family - for needle in (model.model_id, model.display_name, Path(model.id).name): - if needle and detect_video_family(needle) is not None: + for needle in _local_family_needles(model): + if detect_video_family(needle) is not None: return True except Exception: pass @@ -3438,6 +3470,47 @@ def _repo_is_diffusers(repo_info) -> bool: return False +def _repo_pipeline_missing_denoiser(repo_info) -> bool: + """True for a diffusers-pipeline snapshot (root ``model_index.json``) whose denoiser + component (``transformer/`` or ``unet/``) carries NO weight file. This is the shape of a + companion-only prefetch: a GGUF image load pulls the base repo's VAE / text-encoder / + ``model_index.json`` into the cache but deliberately skips the multi-GB transformer (the GGUF + supplies it), so the snapshot has a pipeline manifest yet is not a loadable BF16 pipeline -- + ``from_pretrained`` on it re-downloads the missing shards. ``_cached_repo_partial`` misses this + (no cancel marker / .incomplete blob, and hf_hub_download writes no manifest), so ``/cached-models`` + would advertise it as fully on-device. The caller marks such rows partial. Best-effort: any scan + error reports not-missing so a glitch never hides a genuinely complete pipeline.""" + if not _repo_has_pipeline_index(repo_info): + return False + _DENOISER_DIRS = ("transformer", "unet") + _WEIGHT_SUFFIXES = (".safetensors", ".bin") + try: + for rev in repo_info.revisions: + snapshot = getattr(rev, "snapshot_path", None) + for f in rev.files: + name = str(getattr(f, "file_name", "") or "") + path = getattr(f, "file_path", None) + parts: tuple[str, ...] = () + if path is not None and snapshot is not None: + try: + parts = Path(path).relative_to(Path(snapshot)).parts + except ValueError: + parts = () + if not parts: + # No snapshot scoping (or file outside it): fall back to the recorded name, + # which may itself carry the component subdir (e.g. 'transformer/model...'). + parts = Path(name).parts + if ( + len(parts) >= 2 + and parts[0].lower() in _DENOISER_DIRS + and parts[-1].lower().endswith(_WEIGHT_SUFFIXES) + ): + return False + return True + except Exception: + return False + + def _cached_repo_partial(repo_id: str, repo_cache_dir: Optional[Path] = None) -> bool: """Whether the cached model snapshot is incomplete (cancelled/partial download). Reuses the hub inventory scan's snapshot-partial detector (cancel marker, legacy @@ -3517,7 +3590,12 @@ async def list_cached_models( ) key = repo_id.lower() existing = seen_lower.get(key) - is_partial = _cached_repo_partial(repo_id, Path(repo_info.repo_path)) + # A companion-only prefetch (root model_index.json + VAE / text-encoder but no + # transformer/ shards, pulled to back a GGUF load) is not a loadable BF16 + # pipeline; treat it as partial so the picker does not advertise it as on-device. + is_partial = _cached_repo_partial( + repo_id, Path(repo_info.repo_path) + ) or _repo_pipeline_missing_denoiser(repo_info) # Prefer the most COMPLETE snapshot, then largest. The picker drops partial # rows, so a partial copy in one cache root must not shadow a smaller complete # copy in another (size only breaks ties among equal completeness). diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index ee45a44085..4e1aa3a74d 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -445,7 +445,13 @@ def test_list_cached_models_tags_diffusers_pipeline_as_text_to_image(monkeypatch text-to-image so the chat picker hides it, while a plain checkpoint isn't.""" diffusion = _repo( "Tongyi-MAI/Z-Image-Turbo", - [_file("model_index.json", 1_000), _file("text_encoder/model.safetensors", 9_000)], + [ + _file("model_index.json", 1_000), + _file("text_encoder/model.safetensors", 9_000), + # A complete pipeline carries its denoiser weights; without them the row is a + # companion-only prefetch and would be marked partial (see the dedicated test). + _file("transformer/diffusion_pytorch_model.safetensors", 9_000), + ], tmp_path / "models--Tongyi-MAI--Z-Image-Turbo", ) checkpoint = _repo( @@ -468,6 +474,43 @@ def test_list_cached_models_tags_diffusers_pipeline_as_text_to_image(monkeypatch } +def test_list_cached_models_marks_companion_only_pipeline_partial(monkeypatch, tmp_path): + """A GGUF image load prefetches its companion base repo's VAE / text-encoder / model_index.json + but deliberately skips the multi-GB transformer (the GGUF supplies it). That snapshot carries a + root model_index.json yet is not a loadable BF16 pipeline, so it must be marked partial (the + picker drops partial rows) rather than advertised as fully on-device. A sibling repo that DOES + ship its transformer shards stays complete.""" + companion_only = _repo( + "black-forest-labs/FLUX.1-dev", + [ + _file("model_index.json", 1_000), + _file("vae/diffusion_pytorch_model.safetensors", 9_000), + _file("text_encoder/model.safetensors", 9_000), + ], + tmp_path / "models--black-forest-labs--FLUX.1-dev", + ) + complete = _repo( + "Tongyi-MAI/Z-Image-Turbo", + [ + _file("model_index.json", 1_000), + _file("text_encoder/model.safetensors", 9_000), + _file("transformer/diffusion_pytorch_model.safetensors", 9_000), + ], + tmp_path / "models--Tongyi-MAI--Z-Image-Turbo", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [companion_only, complete])], + ) + + result = asyncio.run(models_route.list_cached_models(current_subject = "test-user")) + by_repo = {c["repo_id"]: c for c in result["cached"]} + assert by_repo["black-forest-labs/FLUX.1-dev"].get("partial") is True + assert by_repo["Tongyi-MAI/Z-Image-Turbo"].get("partial") is None + + def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeypatch, tmp_path): """A vision GGUF repo (main weight + mmproj) is a GGUF repo; reported size is the main weight only, since mmproj is filtered at classification.""" diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index 1fbd7efcd6..dcc948ab30 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -136,6 +136,23 @@ def test_scan_models_dir_surfaces_diffusers_pipeline_folder(tmp_path): assert rows["my-pipeline"].model_format is None +def test_scan_models_dir_surfaces_root_diffusers_pipeline(tmp_path): + # A custom scan folder can point DIRECTLY at a diffusers pipeline (not a parent of repos). + # Its weights live in component subdirs under a root model_index.json, so _is_model_directory + # rejects the root; without admitting it the scan would surface the component subdirs + # (transformer/, vae/) as bogus models and hide the real pipeline. Treat the root as one model. + root = tmp_path / "my-local-pipeline" + _touch(root / "model_index.json") + _touch(root / "transformer" / "config.json") + _touch(root / "transformer" / "diffusion_pytorch_model.safetensors") + _touch(root / "vae" / "diffusion_pytorch_model.safetensors") + + rows = models_route._scan_models_dir(root) + + assert [r.path for r in rows] == [str(root)] + assert rows[0].model_format is None + + # ── Images picker task tag for local (non-GGUF) diffusers models ────────────── from models.models import LocalModelInfo # noqa: E402 @@ -211,6 +228,25 @@ def test_local_task_tags_video_single_file_checkpoint(tmp_path): ) +def test_local_task_tags_single_file_by_checkpoint_filename(tmp_path): + # A generically named folder holding one loadable checkpoint whose FILENAME identifies the + # family (the parent dir does not) is loadable -- the route resolves the sole file via + # resolve_local_single_file -- so tag it from the filename or the task-scoped picker hides it. + d = tmp_path / "downloads" + _touch(d / "qwen-image-2509.safetensors") # family only in the filename, no model_index.json + m = _local(d, id = str(d), display_name = "downloads") + assert models_route._local_is_diffusers(m) is True + assert models_route._local_model_task(m) == "text-to-image" + + +def test_local_task_tags_video_single_file_by_checkpoint_filename(tmp_path): + # Same, for a video family whose token lives only in the sole checkpoint's filename. + d = tmp_path / "clips" + _touch(d / "ltx-2.3-distilled.safetensors") # ltx family only in the filename + m = _local(d, id = str(d), display_name = "clips") + assert models_route._local_model_task(m) == models_route._VIDEO_GEN_TASK + + def test_local_task_ignores_family_token_in_parent_path(tmp_path): # model.id is the full on-disk path for a scanned On-Device model, and the family-token # matcher treats any path segment as a hint. A family token in a PARENT dir (e.g.