ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge remote-tracking branch 'origin/image-generation' into r7021

Browse source 
# Conflicts:
#	studio/backend/routes/inference.py
#	studio/backend/routes/video.py
This commit is contained in:
Daniel Han 2026-07-13 09:28:04 +00:00
commit 9eadfc822d
31 changed files with 1150 additions and 230 deletions
Download patch file Download diff file Expand all files Collapse all files

27
.github/workflows/studio-windows-inference-smoke.yml vendored
View file

@ -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

20
scripts/uninstall.ps1
View file

@ -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 <parent>\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..."

37
studio/backend/core/inference/diffusion.py
View file

@ -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.

65
studio/backend/core/inference/diffusion_engine_router.py
View file

@ -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(

16
studio/backend/core/inference/gpu_arbiter.py
View file

@ -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:

36
studio/backend/core/inference/image_gallery.py
View file

@ -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]

7
studio/backend/core/inference/sd_cpp_backend.py
View file

@ -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,

43
studio/backend/core/inference/video_gallery.py
View file

@ -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]

15
studio/backend/core/training/diffusion_train_common.py
View file

@ -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:

10
studio/backend/core/training/training.py
View file

@ -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

13
studio/backend/models/inference.py
View file

@ -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)."""

80
studio/backend/routes/inference.py
View file

@ -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)

122
studio/backend/routes/models.py
View file

@ -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).

69
studio/backend/routes/training.py
View file

@ -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}

88
studio/backend/routes/video.py
View file

@ -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)

45
studio/backend/tests/test_cached_gguf_routes.py
View file

@ -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."""

61
studio/backend/tests/test_diffusion_backend.py
View file

@ -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]

113
studio/backend/tests/test_diffusion_dataset_api.py
View file

@ -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

45
studio/backend/tests/test_diffusion_engine_router.py
View file

@ -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()

30
studio/backend/tests/test_diffusion_lora_trainer.py
View file

@ -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")

36
studio/backend/tests/test_diffusion_routes.py
View file

@ -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

72
studio/backend/tests/test_gpu_arbiter.py
View file

@ -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

44
studio/backend/tests/test_image_gallery.py
View file

@ -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()) == []

68
studio/backend/tests/test_local_model_format.py
View file

@ -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):

31
studio/backend/tests/test_sd_cpp_backend.py
View file

@ -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 ──────────────────────────────────────────

44
studio/backend/tests/test_sd_cpp_install.py
View file

@ -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)

7
studio/backend/tests/test_training_start_offload.py
View file

@ -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()

49
studio/backend/tests/test_video_gallery.py
View file

@ -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).

20
studio/backend/tests/test_video_routes.py
View file

@ -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("{}")

29
studio/frontend/src/features/images/images-page.tsx
View file

@ -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);

38
studio/install_sd_cpp_prebuilt.py
View file

@ -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