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 742a07c221..46eb48cb3b 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -364,9 +364,27 @@ function Uninstall-UnslothStudio { _StopByPortFile -PortFile (Join-Path $r "share\studio.port") -KnownRoots $knownRoots } _StopStudioProcesses -KnownRoots $knownRoots + # 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 + } + # 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, $defaultSdCpp, $defaultCache, $defaultNode)) + _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache, $defaultNode) + @($defaultSdCppToStop | Where-Object { $_ }) + @($customSdCppToStop)) # ── 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 d2a8fefbe4..8e99b16e39 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -2282,14 +2282,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. @@ -2626,10 +2622,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]: @@ -2655,18 +2649,19 @@ 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: _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: - pass + with self._lock: + self._unload_locked() return self.status() def _unload_locked(self) -> None: @@ -2685,9 +2680,9 @@ 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. + # 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 17624ab5b8..4b0c9b9155 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -59,6 +59,9 @@ 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(), 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 @@ -86,37 +89,39 @@ 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 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 + # 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 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 + 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/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/image_gallery.py b/studio/backend/core/inference/image_gallery.py index f3ee2fb00d..fb1e67b096 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,20 @@ 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 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)) + 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 +143,22 @@ 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 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: @@ -150,7 +175,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/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 9b51e9cb9a..a04c8a5f47 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -765,6 +765,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) @@ -790,7 +796,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/core/inference/video_gallery.py b/studio/backend/core/inference/video_gallery.py index db157c0550..1391119217 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,25 @@ 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, 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") + 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 +191,21 @@ 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 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: @@ -195,7 +219,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 d423d13444..f891ca4309 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -702,20 +702,25 @@ 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 (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). - if caption is None: + # 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. - if caption is None and instance_prompt: + # 3. dreambooth instance prompt for any image still without a caption. + if not caption and instance_prompt: caption = instance_prompt if caption: if verify_images: 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 818deee470..e01099f262 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2084,6 +2084,19 @@ 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+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 " + "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/routes/inference.py b/studio/backend/routes/inference.py index 5340150efe..5246363e43 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -14338,10 +14338,37 @@ 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, + vae_quant = request.vae_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,26 +14376,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, - vae_quant = request.vae_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))) @@ -14512,18 +14520,24 @@ 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 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: - 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/models.py b/studio/backend/routes/models.py index 7f34d2340c..0d0300f856 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: @@ -338,7 +353,13 @@ 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 = _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. continue @@ -3255,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. @@ -3287,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 @@ -3305,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(): @@ -3315,8 +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 / 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 _local_family_needles(model): + if detect_video_family(needle) is not None: return True except Exception: pass @@ -3420,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 @@ -3499,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/routes/training.py b/studio/backend/routes/training.py index 1ef804afc2..9d28bec7dc 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1613,11 +1613,12 @@ 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 + # 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) limit_bytes = get_upload_limit_bytes() @@ -1732,9 +1733,43 @@ 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: 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: + 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 +1801,25 @@ 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 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, + 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 +2046,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 +2055,9 @@ 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 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 8f789334ab..6f84cfd66a 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: @@ -116,29 +113,39 @@ 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_cache_quality = request.transformer_cache_quality, + transformer_quant = request.transformer_quant, + text_encoder_quant = request.text_encoder_quant, + vae_quant = request.vae_quant, + cfg_parallel = request.cfg_parallel, + 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_cache_quality = request.transformer_cache_quality, - transformer_quant = request.transformer_quant, - text_encoder_quant = request.text_encoder_quant, - vae_quant = request.vae_quant, - cfg_parallel = request.cfg_parallel, - 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))) @@ -239,17 +246,24 @@ 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 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: - 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_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index 35b72803d2..14008f2638 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_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 7a3f3ca065..d0c781cf05 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( @@ -3071,3 +3068,57 @@ 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: 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() + + 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..60170ee273 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -495,6 +495,119 @@ 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 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() + 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 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") + 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 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" + 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_upload_through_symlinked_dataset_cannot_escape_root(client, ds_root, tmp_path): + # 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) + + 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): + # 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" + _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): + # 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" + 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..18b9f51b62 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -238,3 +238,48 @@ 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: 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 + 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: the query cannot complete while the switch holds the transition lock. + 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..6ae96d97e5 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -57,6 +57,36 @@ 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): + # 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", + 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 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( + 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 6ad1dc2b8f..e3732fb8e0 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -180,8 +180,15 @@ 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) @@ -333,6 +340,25 @@ 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"} + ) + # 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}, + ) + 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. @@ -762,7 +788,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_image_gallery.py b/studio/backend/tests/test_image_gallery.py index 128a533666..51f107ffd2 100644 --- a/studio/backend/tests/test_image_gallery.py +++ b/studio/backend/tests/test_image_gallery.py @@ -145,3 +145,47 @@ 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(): + # ``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) + + def _valid(rec): + return rec.get("prompt") != "BAD" + + # 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. + 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(): + # 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) + + def _valid(rec): + return not str(rec.get("prompt", "")).startswith("BAD") + + # 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 before publishing must leave neither a truncated {id}.png nor a leftover temp. + 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_local_model_format.py b/studio/backend/tests/test_local_model_format.py index 3da98ceabb..dcc948ab30 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -117,6 +117,42 @@ 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 + + +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 @@ -180,13 +216,35 @@ 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 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 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_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): diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 5f04408a35..94b464f6c9 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -279,6 +279,37 @@ 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_sd_cpp_install.py b/studio/backend/tests/test_sd_cpp_install.py index 07d52a0b02..89404fa3ae 100644 --- a/studio/backend/tests/test_sd_cpp_install.py +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -262,6 +262,50 @@ 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_is_refused(tmp_path, monkeypatch): + # 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" + target.mkdir() + user_file = target / "USER_WORK" + user_file.write_text("keep", encoding = "utf-8") + + with pytest.raises(RuntimeError, match = "not a Studio-managed directory"): + install(install_dir = target) + + # 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() + assert list(target.iterdir()) == [user_file] + + +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_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 21a8dac014..6ca8ea2e35 100644 --- a/studio/backend/tests/test_video_gallery.py +++ b/studio/backend/tests/test_video_gallery.py @@ -199,6 +199,55 @@ def test_list_skips_corrupt_sidecar(): assert [r["prompt"] for r in listed] == ["ours"] +def test_valid_callback_paginates_over_accepted_records(): + # ``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) + + 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): + # 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} + + 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/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index a0385fc9d1..63e92beaf4 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -112,8 +112,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() ) @@ -264,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", @@ -386,10 +391,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 +407,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 337eb31732..88aa810ef1 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1443,13 +1443,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) @@ -1470,8 +1467,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; @@ -1511,11 +1507,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) { @@ -1554,11 +1547,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 a7b00861e7..67ff5f6635 100644 --- a/studio/install_sd_cpp_prebuilt.py +++ b/studio/install_sd_cpp_prebuilt.py @@ -401,6 +401,30 @@ def install( has no ``sd-cli``. """ target = install_dir or default_install_dir() + # 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(): + 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() + # 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}. " + 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: @@ -441,13 +465,13 @@ 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. - try: - (target / ".unsloth-studio-owned").touch() - except OSError: - pass + # 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() + except OSError: + pass return sd_cli