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

Tighten comments and docstrings added by the image-generation fixes

Browse source
This commit is contained in:
Daniel Han 2026-07-13 05:29:09 +00:00
commit e0ef488f47
25 changed files with 140 additions and 252 deletions
Download patch file Download diff file Expand all files Collapse all files

7
scripts/uninstall.ps1
View file

@ -364,11 +364,8 @@ function Uninstall-UnslothStudio {
_StopByPortFile -PortFile (Join-Path $r "share\studio.port") -KnownRoots $knownRoots
}
_StopStudioProcesses -KnownRoots $knownRoots
# The default sd.cpp dir is only killed+deleted when it carries our owner marker (the deletion
# below is marker-gated). Passing it to the locking-process stop UNCONDITIONALLY would terminate
# a user's own sd-server running from an unowned checkout at this default path -- a process we
# then decide to keep the directory for. Gate the kill on the same predicate so stop and delete
# agree: include it only when marked owned (and present).
# Only stop the default sd.cpp dir when it carries our owner marker, so stop matches the
# marker-gated delete and a user's own sd-server at this default path is left running.
$defaultSdCppToStop = $null
if ($defaultSdCpp -and (Test-Path -LiteralPath $defaultSdCpp) -and (Test-Path -LiteralPath (Join-Path $defaultSdCpp ".unsloth-studio-owned") -PathType Leaf)) {
$defaultSdCppToStop = $defaultSdCpp

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

@ -2216,14 +2216,10 @@ class DiffusionBackend:
raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
# Register under _lock so unload()/a load can signal THIS generation.
self._active_generate_cancel = cancel
# Publish an active (step 0) progress state the moment the lock is held, BEFORE
# the slow pre-denoise setup (deferred compile, LoRA resolution/application,
# ControlNet download/build). Without this generate_progress() reports inactive
# across that window, so a reload's mount probe shows idle even though this
# generation holds _generate_lock; the user then starts a second generate that
# merely blocks behind this one (and can duplicate the result). The per-step
# callback swaps in its own _GenState at denoise start; this is the queued phase.
# Mirrors the video backend's queued state and the training start guard.
# Publish an active (step 0) state now, before the slow pre-denoise setup (deferred
# compile, LoRA resolution, ControlNet build), so a reload's mount probe doesn't read
# idle while this generation holds _generate_lock and let a second generate queue
# behind it. The per-step callback swaps in its own _GenState at denoise start.
self._gen = _GenState(total_steps = steps)
try:
# The local `state` ref keeps the pipe alive even if unload() nulls _state.
@ -2559,10 +2555,8 @@ class DiffusionBackend:
with self._lock:
if self._active_generate_cancel is cancel:
self._active_generate_cancel = None
# Drop the published progress state. The normal path already nulled it after
# the denoise; this also covers a setup-time error that skips that inner
# finally. Safe under _generate_lock: no other generation can have installed
# its own _gen while this one runs.
# Drop the published progress state, covering a setup-time error that skips
# the inner finally. Safe under _generate_lock.
self._gen = None
def generate_progress(self) -> dict[str, Any]:
@ -2594,13 +2588,10 @@ class DiffusionBackend:
# Cancel any in-flight load (its worker checks this token) and drop the marker.
self._load_token += 1
self._loading = None
# WAIT for the signalled denoise to exit BEFORE tearing down (mirrors begin_load's
# pre-teardown barrier at the top of the locked load path). _unload_locked uninstalls
# PROCESS-WIDE state the running denoise still depends on -- the eager/arch attention
# patches, the GGUF compile hooks, the flipped backend flags and the compile cache -- none
# of which the denoise's own pipe ref pins. Uninstalling them before the denoise exits would
# corrupt or crash its in-flight forward passes. The denoise holds _generate_lock for its
# whole body, so acquiring it here blocks until it has finished; only then do we tear down.
# Wait for the signalled denoise to exit BEFORE tearing down: _unload_locked uninstalls
# process-wide state (attention patches, GGUF compile hooks, backend flags, compile cache)
# the denoise still depends on but its pipe ref doesn't pin. The denoise holds _generate_lock
# for its whole body, so acquiring it here blocks until it has finished. Mirrors begin_load.
with self._generate_lock:
with self._lock:
self._unload_locked()
@ -2622,10 +2613,9 @@ class DiffusionBackend:
uninstall_patches()
uninstall_arch_patches()
# NOTE: deliberately NOT unload_lora_weights() here. Both callers (unload() and begin_load's
# locked path) now hold _generate_lock across this teardown, so no denoise is in flight; the
# whole pipe is dropped below, freeing any LoRA adapters with it, so a separate adapter unload
# would be redundant work.
# Deliberately NOT unload_lora_weights() here: the whole pipe is dropped below, freeing any
# LoRA adapters with it. Both callers hold _generate_lock across this teardown, so no denoise
# is in flight.
# Drop the workflow pipes so they don't pin the freed pipeline's modules past unload.
self._aux_pipes.clear()
# Drop any ControlNet models + pipelines so the freed load carries no extra modules.

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

@ -59,9 +59,8 @@ def _install_accelerator_for(backend: str) -> str:
# The engine the current load committed to, and why a non-native choice was made. Mutated only
# under _lock during selection.
_lock = threading.Lock()
# Serializes a WHOLE engine switch (check -> unload -> publish). _lock alone is released during the
# slow unload(), so two overlapping selections could interleave: one observes the not-yet-updated
# active engine, returns it, and loads onto the very engine the other is concurrently unloading.
# Serializes a whole engine switch (check -> unload -> publish); _lock alone is released during the
# slow unload(), letting two overlapping selections load onto the engine the other is unloading.
_transition_lock = threading.Lock()
_active_engine_name: str = ENGINE_DIFFUSERS
_fallback_reason: Optional[str] = None
@ -90,12 +89,9 @@ def active_engine_name() -> str:
def _activate(name: str, reason: Optional[str]) -> Any:
global _active_engine_name, _fallback_reason
# Serialize the ENTIRE check -> unload -> publish transition. _lock is released during the slow
# unload() below (holding it across the unload would block every status/selection reader), which
# opens a window where a second _activate could read the still-old active engine, take the "no
# change" branch, and return that engine -- then load onto it while this call is unloading it.
# _transition_lock closes the window without holding _lock across the unload; the final
# get_active_diffusion_engine() now reflects the committed state.
# Serialize the whole check -> unload -> publish transition without holding _lock across the
# slow unload(), closing the window where a second _activate reads the still-old active engine,
# takes the "no change" branch, and loads onto the engine this call is unloading.
with _transition_lock:
# Switching engines: unload the deactivated one first, else its model stays resident but
# unreachable (the evictor only targets the active engine), leaking 10+ GB. The unload is
@ -112,9 +108,8 @@ def _activate(name: str, reason: Optional[str]) -> Any:
if engine_to_unload is not None:
# Publish the new engine only AFTER the old one unloads. The evictor unloads
# get_active_diffusion_engine(), so flipping the name first would let a concurrent
# acquire_for evict the new (empty) engine and take the GPU while the old model is still
# freeing VRAM -- two large models briefly resident. Keeping the OLD engine as the evict
# target serializes a concurrent evict on its unload(), granting the GPU only once freed.
# acquire_for evict the new (empty) engine while the old model is still freeing VRAM.
# Keeping the OLD engine as the evict target grants the GPU only once it is freed.
try:
engine_to_unload.unload()
except Exception as exc: # noqa: BLE001 -- best-effort; never block the switch

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

@ -69,9 +69,8 @@ def save(image: Any, meta: dict[str, Any]) -> dict[str, Any]:
image_id = uuid.uuid4().hex
directory = gallery_dir()
final_path = directory / f"{image_id}.png"
# Write to a hidden temp then atomically rename into place: a crash / disk-full mid-write must
# never leave a truncated {id}.png that the listing would surface as a corrupt record. The temp
# name is dotted so an interrupted write is skipped by the *.png glob, and cleaned on failure.
# Write to a dotted temp (skipped by the *.png glob) then atomically rename, so a crash mid-write
# never leaves a truncated {id}.png that the listing would surface as a corrupt record.
tmp_path = directory / f".{image_id}.png.tmp"
try:
tmp_path.write_bytes(_png_bytes(image, meta))
@ -156,12 +155,10 @@ def list_images(
full just to sort; only the window's recipes are read. limit=None returns everything from
``offset`` on.
``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and the
caller's has_more all count over the same accepted-record domain. Passing the route's schema
validator here is essential: a record that has every required key (so ``_read_meta`` accepts
it) but a wrong value type would otherwise be counted in this window yet dropped by the route
after slicing -- a leading bad record then returns an empty page with more remaining, which
stalls infinite scroll at offset 0."""
``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and has_more
all count over the accepted-record domain. Pass the route's schema validator: a record with
every required key (so ``_read_meta`` accepts it) but a wrong value type would otherwise be
counted here yet dropped after slicing, stalling infinite scroll at offset 0."""
try:
paths = list(gallery_dir().glob("*.png"))
except OSError:

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

@ -39,11 +39,9 @@ def save(mp4_bytes: bytes, meta: dict[str, Any]) -> dict[str, Any]:
mp4_tmp = directory / f".{video_id}.mp4.tmp"
sidecar = directory / f"{video_id}.json"
sidecar_tmp = directory / f".{video_id}.json.tmp"
# Stage BOTH files, then publish. The sidecar is the pair's commit marker (list_videos scans
# mp4s but skips any without a readable sidecar), so writing the MP4 straight to its final name
# before the sidecar meant a sidecar write/replace failure left an invisible, undeletable orphan
# MP4 (gallery delete resolves by id, but the record never appears to be deleted). Rename the MP4
# in first, then the sidecar; on ANY failure remove every artifact so nothing is stranded.
# Stage both files, rename the MP4 in, then the sidecar (the pair's commit marker: list_videos
# skips an mp4 without a readable sidecar). On any failure remove every artifact, else a sidecar
# failure would leave an invisible, undeletable orphan MP4.
try:
mp4_tmp.write_bytes(mp4_bytes)
sidecar_tmp.write_text(json.dumps(meta), encoding = "utf-8")
@ -204,11 +202,10 @@ def list_videos(
Ordered by MP4 mtime (a cheap stat ~= generation order); only the window's sidecars are read.
limit=None returns everything from ``offset`` on. A file without its pair is skipped.
``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and the
caller's has_more all count over the same accepted-record domain. Passing the route's schema
validator here keeps a sidecar that parses as JSON but fails the response schema from being
counted in this window yet dropped by the route after slicing -- which would otherwise return a
short or empty page with more remaining and stall infinite scroll."""
``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and has_more
all count over the accepted-record domain. Pass the route's schema validator: a sidecar that
parses as JSON but fails the response schema would otherwise be counted here yet dropped after
slicing, stalling infinite scroll."""
try:
paths = list(gallery_dir().glob("*.mp4"))
except OSError:

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

@ -705,27 +705,22 @@ def discover_image_caption_pairs(
caption: Optional[str] = None
sidecar_present = False
# 1. per-image sidecar caption file (the user's explicit edit; wins over metadata).
# An EMPTY sidecar is a deliberate tombstone (the labeling grid writes one when a user
# clears a metadata caption): it suppresses the metadata caption but leaves the image
# UNCAPTIONED so the dreambooth instance_prompt fallback below still applies. Treating
# "" as a present caption here would skip the instance prompt AND drop the image, so
# clearing every metadata caption in the grid would make a dreambooth run fail with
# "No captioned images found".
# An EMPTY sidecar is a deliberate tombstone (written when a user clears a caption): it
# suppresses the metadata caption but leaves the image uncaptioned so the instance_prompt
# fallback below still applies, rather than dropping the image.
for ext in _CAPTION_EXTS:
sidecar = img.with_suffix(ext)
if sidecar.is_file():
sidecar_present = True
caption = sidecar.read_text(encoding = "utf-8").strip()
break
# 2. metadata row keyed by file name (basename or the relative path; as_posix so a
# Windows backslash path still matches the jsonl's forward-slash keys). Skipped when a
# sidecar tombstone is present (the sidecar, even empty, is authoritative).
# 2. metadata row keyed by file name (basename or relative path; as_posix so a Windows
# backslash path matches the jsonl's forward-slash keys). A sidecar, even empty, wins.
if not sidecar_present:
caption = meta_caption.get(img.name) or meta_caption.get(
img.relative_to(root).as_posix()
)
# 3. dreambooth instance prompt for any image still without a caption (no sidecar text and
# no metadata row, or an empty tombstone).
# 3. dreambooth instance prompt for any image still without a caption.
if not caption and instance_prompt:
caption = instance_prompt
if caption:

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

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

@ -2071,12 +2071,10 @@ class DiffusionGenerateRequest(BaseModel):
@model_validator(mode = "after")
def _batch_seeds_json_safe(self) -> "DiffusionGenerateRequest":
# A batch derives per-image seeds as seed, seed+1, ... seed+batch_size-1 (sd.cpp and
# diffusers both advance the seed per image). The base seed is bounded to 2**53-1 so it
# round-trips through the JSON gallery recipe, but the derived top-of-batch seed is not:
# with an explicit seed near the cap it can exceed Number.MAX_SAFE_INTEGER, where the
# frontend rounds it and a restored recipe replays a different image. Reject at the
# boundary so an API client can't persist an unreplayable seed.
# A batch derives per-image seeds as seed .. seed+batch_size-1. The base seed is capped at
# 2**53-1 to round-trip through the JSON recipe, but a derived top-of-batch seed near the cap
# can exceed it, where the frontend rounds it and a restored recipe replays a different
# image. Reject at the boundary so an API client can't persist an unreplayable seed.
if self.seed is not None and self.seed + self.batch_size - 1 > 2**53 - 1:
raise ValueError(
"seed + batch_size - 1 must not exceed 2**53 - 1 so every per-image seed "

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

@ -14512,11 +14512,10 @@ async def list_gallery_images(
limit = max(1, min(limit, 200))
offset = max(0, offset)
# Validate against the response schema INSIDE the pager so offset / limit / has_more all count
# over the same accepted-record domain. A PNG whose recipe chunk has all keys but a wrong value
# type passes the presence-only read yet fails GalleryImage(**r); dropping such records only
# after pagination made a leading bad record return an empty page with has_more=True, stalling
# infinite scroll at offset 0. Filtering here keeps the window and has_more consistent.
# Validate inside the pager so offset / limit / has_more all count over the accepted domain. A
# recipe with all keys but a wrong value type passes the presence-only read yet fails
# GalleryImage(**r); dropping it only after slicing let a leading bad record return an empty
# page with has_more=True, stalling infinite scroll at offset 0.
def _valid_gallery_image(record: dict) -> bool:
try:
GalleryImage(**record)

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

@ -3320,13 +3320,11 @@ def _local_is_diffusers(model: "LocalModelInfo") -> bool:
return True
except Exception:
pass
# A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json)
# ships no pipeline index and resolves to no IMAGE family either, so the checks above miss
# it. The video load route reinterprets a bare single-file local dir as a single_file load
# (routes/video.py), so it IS loadable and must be surfaced; without tagging it here
# _local_model_task returns task=null and the Video On-Device picker hides it. Match the same
# clean id / name needles (not the raw path) so a parent-dir family token can't spuriously
# match; _local_model_task then routes the video family to the text-to-video task.
# A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json) has no
# pipeline index and no image family, so the checks above miss it. The video load route loads it
# as a single_file (routes/video.py), so it must be surfaced or _local_model_task returns
# task=null and the picker hides it. Match clean id / name needles (not the raw path) so a
# parent-dir token can't spuriously match; _local_model_task then routes it to text-to-video.
try:
from core.inference.video_families import detect_video_family
for needle in (model.model_id, model.display_name, Path(model.id).name):

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

@ -1616,11 +1616,8 @@ async def upload_diffusion_dataset(
from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label
cleaned = _clean_diffusion_dataset_name(name)
# Reject a symlinked dataset directory BEFORE any write. A bare mkdir(exist_ok=True) succeeds
# through an existing name -> external-directory symlink, and the staged upload would then
# write/replace files outside the Studio datasets root through that link. The read/caption/
# delete endpoints already enforce this containment via _resolve_dataset_folder; the upload
# path must run the same symlink + root-containment check first so writes can never escape.
# Run the same symlink + root-containment check as the read/caption/delete endpoints before any
# write, so a name -> external-directory symlink can't make the staged upload write outside root.
folder = _resolve_dataset_folder(name, must_exist = False)
folder.mkdir(parents = True, exist_ok = True)
@ -1736,13 +1733,10 @@ async def upload_diffusion_dataset(
)
out.write(chunk)
uploaded += 1
# Commit every staged file as one transaction. A plain replace loop is NOT atomic across
# files: if the second-or-later tmp.replace(dest) fails (disk/quota error, a Windows file
# lock, antivirus, a destination that became a directory), an earlier destination has
# already been overwritten while the request returns an error -- the user's original file
# is gone. Back up each pre-existing destination before overwriting it, then on ANY failure
# remove the versions this request installed and restore every displaced original, so the
# dataset is left exactly as it was before the upload.
# Commit every staged file as one transaction. A plain replace loop is not atomic across
# files: a mid-loop tmp.replace(dest) failure leaves earlier destinations already overwritten
# while the request errors. Back up each pre-existing destination first, then on any failure
# drop the versions this request installed and restore every displaced original.
backups: list[tuple[Path, Optional[Path]]] = [] # (dest, backup path or None)
installed: list[Path] = []
try:
@ -1809,11 +1803,9 @@ def _resolve_dataset_folder(name: str, *, must_exist: bool = True) -> Path:
cleaned = _clean_diffusion_dataset_name(name)
root = datasets_root().resolve()
folder = root / cleaned
# Reject a symlinked dataset directory. _safe_dataset_image_path only proves each image path
# stays under folder.resolve(); it never proves the folder itself stays under the datasets
# root. A dataset dir that is a symlink to an external directory would therefore let image
# read / caption / delete operate on files outside Studio (a reproduced delete removed an
# external file through such a link). Prove the resolved folder is contained in the root too.
# Reject a symlinked dataset directory and prove the resolved folder stays under root:
# _safe_dataset_image_path only checks each image path, so a folder symlinked to an external
# directory would let read / caption / delete operate on files outside Studio.
if folder.is_symlink():
raise HTTPException(
status_code = 400,
@ -2063,9 +2055,8 @@ async def delete_diffusion_dataset_image(
if thumbs_dir.is_dir():
# Thumbs are keyed on the full filename (stem + extension), so match that here too;
# a stem-only glob would strand this image's thumbs or delete a same-stem sibling's.
# Escape the filename first: an uploaded name may legally contain glob metacharacters
# ('[', ']', '*', '?'), and interpolating those raw would make e.g. "[ab].png" match
# "a.png_*.jpg"/"b.png_*.jpg" -- deleting siblings' thumbs while leaving its own behind.
# Escape the name: a raw glob metacharacter (e.g. "[ab].png") would match siblings'
# thumbs while leaving its own behind.
for t in thumbs_dir.glob(f"{_glob.escape(image_path.name)}_*.jpg"):
t.unlink(missing_ok = True)
return {"deleted": image_path.name}

22
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:
@ -236,11 +233,10 @@ async def list_gallery_videos(
limit = max(1, min(limit, 200))
offset = max(0, offset)
# Validate against the response schema INSIDE the pager so offset / limit / has_more all count
# over the same accepted-record domain. A sidecar that parses as JSON but has a wrong value type
# passes the read yet fails GalleryVideo(**r); dropping such records only after pagination made a
# leading bad record return an empty page with has_more=True, stalling infinite scroll at offset
# 0. Filtering here keeps the window and has_more consistent.
# Validate inside the pager so offset / limit / has_more all count over the accepted domain. A
# sidecar that parses as JSON but has a wrong value type passes the read yet fails
# GalleryVideo(**r); dropping it only after slicing let a leading bad record return an empty
# page with has_more=True, stalling infinite scroll at offset 0.
def _valid_gallery_video(record: dict) -> bool:
try:
GalleryVideo(**record)

13
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(
@ -3009,10 +3006,8 @@ def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path):
def test_unload_waits_for_in_flight_denoise_before_teardown():
# Regression for the unload/denoise teardown race: unload() must WAIT for a running denoise to
# exit (acquire _generate_lock) BEFORE _unload_locked() tears down PROCESS-WIDE state (eager /
# arch attention patches, gguf compile hooks, backend flags, compile cache) that the denoise
# still depends on. Mirror the load path, which already waits on _generate_lock before teardown.
# Regression: unload() must wait for a running denoise to exit (acquire _generate_lock) before
# _unload_locked() tears down process-wide state the denoise still depends on.
import threading
backend = DiffusionBackend()

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

@ -496,12 +496,8 @@ def test_import_promotion_leaves_no_partial_dataset_on_failure(ds_root, monkeypa
def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch):
# Re-uploading a.txt and b.txt (an allowed overwrite of two files already on disk) stages both,
# then commits them. A plain replace loop is NOT atomic: if the SECOND commit fails after the
# first destination was already overwritten, the live dataset is left partially updated with the
# request returning an error -- the user's original a.txt is gone. The transactional commit must
# back up each displaced original and, on any failure, restore every one so the dataset is left
# exactly as it was, with no stray temp/backup files.
# Re-uploading a.txt and b.txt where the SECOND commit fails must roll back the first overwrite,
# so both originals survive and no stray temp/backup files remain.
from pathlib import Path
app = FastAPI()
@ -518,8 +514,7 @@ def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch):
state = {"failed": False}
def flaky_replace(self, target, *a, **k):
# Fail exactly once, on the tmp -> b.txt promotion (source is the staged .upload-* part, NOT
# a .upload-backup-* part), so the subsequent backup -> b.txt restore still succeeds.
# Fail once on the tmp -> b.txt promotion only (not the backup restore), so rollback works.
if (
not state["failed"]
and str(target).endswith("b.txt")
@ -548,9 +543,8 @@ def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch):
def test_resolve_dataset_folder_rejects_symlink(ds_root, tmp_path):
# A dataset directory that is a symlink pointing OUTSIDE the datasets root must be rejected: the
# per-image containment check only proves paths stay under folder.resolve(), so without this a
# delete/caption/read could operate on external files through the link.
# A dataset dir that is a symlink outside the datasets root must be rejected, else delete /
# caption / read could operate on external files through the link.
from routes.training import _resolve_dataset_folder
external = tmp_path / "external"
@ -564,9 +558,8 @@ def test_resolve_dataset_folder_rejects_symlink(ds_root, tmp_path):
def test_upload_through_symlinked_dataset_cannot_escape_root(client, ds_root, tmp_path):
# End to end: an upload targeting a dataset name that already exists as a symlink to an
# external directory must be refused (400) BEFORE any bytes are written, so the upload can
# never create/replace files outside the datasets root through the link.
# An upload to a dataset name that is a symlink to an external directory must be refused (400)
# before any bytes are written.
external = tmp_path / "external"
external.mkdir()
(ds_root / "linked").symlink_to(external, target_is_directory = True)
@ -580,8 +573,7 @@ def test_upload_through_symlinked_dataset_cannot_escape_root(client, ds_root, tm
def test_delete_through_symlinked_dataset_cannot_escape_root(client, ds_root, tmp_path):
# End to end: a DELETE against an image inside a symlinked dataset dir is refused (400) and the
# external file it points at is NOT removed.
# A DELETE inside a symlinked dataset dir is refused (400) and the external file survives.
external = tmp_path / "external"
external.mkdir()
victim = external / "victim.png"
@ -594,10 +586,8 @@ def test_delete_through_symlinked_dataset_cannot_escape_root(client, ds_root, tm
def test_delete_image_with_glob_chars_only_removes_own_thumbs(client, ds_root):
# A legal uploaded filename may contain glob metacharacters ('[', ']', '*', '?'). Deleting it
# must remove ONLY its own thumbnails; interpolating the raw name into Path.glob would make
# "[ab].png" match "a.png_*.jpg"/"b.png_*.jpg" -- deleting a sibling's thumbs while leaving its
# own (literally "[ab].png_32.jpg") behind.
# Deleting a filename with glob metacharacters (e.g. "[ab].png") must remove only its own
# thumbnails, not a sibling's that the raw glob would spuriously match.
from urllib.parse import quote
folder = ds_root / "d"

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

@ -241,11 +241,8 @@ def test_no_switch_keeps_engine_and_refreshes_reason(monkeypatch):
def test_activate_serializes_switch_and_concurrent_query(monkeypatch):
# Regression for the check->unload->publish race: _activate releases _lock during the slow
# unload(), so without the transition lock a second _activate could observe the not-yet-updated
# active engine, take the "no change" branch, and return the engine the first call is
# concurrently unloading. Drive both paths on threads and assert the concurrent query is blocked
# until the switch completes (i.e. the whole transition is serialized).
# Regression: without the transition lock a second _activate during the slow unload() reads the
# not-yet-updated active engine and returns it. Assert the query is blocked until the switch ends.
import threading
r._active_engine_name = ENGINE_DIFFUSERS
@ -279,8 +276,7 @@ def test_activate_serializes_switch_and_concurrent_query(monkeypatch):
q = threading.Thread(target = _query)
q.start()
# Serialized: while the switch holds the transition lock the query cannot complete. Pre-fix it
# would return immediately (active is still diffusers), setting query_done at once.
# Serialized: the query cannot complete while the switch holds the transition lock.
assert not query_done.wait(0.4)
release_unload.set()

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

@ -58,11 +58,8 @@ def test_discover_sidecar_overrides_metadata_row(tmp_path):
def test_discover_empty_sidecar_suppresses_metadata_but_uses_instance_prompt(tmp_path):
# Clearing a metadata caption in the labeling grid writes an EMPTY sidecar tombstone. It must
# suppress the metadata caption (so the stale text is not resurrected) yet leave the image
# UNCAPTIONED so the dreambooth instance_prompt still applies. Treating "" as a present caption
# would skip the instance prompt AND drop the image, so a dataset whose every caption was
# cleared would fail with "No captioned images found".
# An empty sidecar tombstone must suppress the metadata caption yet leave the image uncaptioned
# so the dreambooth instance_prompt still applies (not drop the image).
_touch(tmp_path / "cat.png")
(tmp_path / "metadata.jsonl").write_text(
json.dumps({"file_name": "cat.png", "text": "old metadata caption"}) + "\n",
@ -74,9 +71,8 @@ def test_discover_empty_sidecar_suppresses_metadata_but_uses_instance_prompt(tmp
def test_discover_empty_sidecar_without_instance_prompt_skips_image(tmp_path):
# With no instance prompt an empty tombstone leaves the image uncaptioned, so it is skipped and
# the suppressed metadata caption is NOT resurrected -- while a sibling with a real caption is
# still discovered.
# With no instance prompt the tombstoned image is skipped (metadata not resurrected), while a
# sibling with a real caption is still discovered.
_touch(tmp_path / "cat.png")
_touch(tmp_path / "cap.png")
(tmp_path / "metadata.jsonl").write_text(

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

@ -344,9 +344,8 @@ def test_generate_rejects_batch_seed_past_json_safe_range(client):
client.post(
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
)
# An explicit seed at the JS-safe cap with a batch derives per-image seeds
# (seed+1 ...) that exceed Number.MAX_SAFE_INTEGER and no longer round-trip
# through the gallery JSON recipe, so the request is rejected at the boundary.
# A seed at the cap with a batch derives per-image seeds (seed+1 ...) past the JSON-safe range,
# so the request is rejected.
over = client.post(
"/api/inference/images/generate",
json = {"prompt": "p", "seed": 2**53 - 1, "batch_size": 2},

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

@ -148,10 +148,8 @@ def test_list_skips_recipe_missing_required_fields(tmp_path):
def test_valid_callback_paginates_over_accepted_records():
# A record that passes _read_meta (every required key present) but fails the caller's stricter
# schema check must be filtered BEFORE pagination, so offset/limit/has_more all count over the
# accepted domain. Otherwise a leading bad record returns an empty/short page with more still
# remaining, and the frontend (which advances by valid records) stalls at offset 0.
# ``valid`` must filter before pagination, so offset/limit/has_more count over the accepted
# domain; else a leading bad record returns a short page with more remaining and stalls scroll.
_save_with_mtime("BAD", 300.0) # newest, sorts first
_save_with_mtime("g1", 200.0)
_save_with_mtime("g2", 100.0)
@ -159,8 +157,7 @@ def test_valid_callback_paginates_over_accepted_records():
def _valid(rec):
return rec.get("prompt") != "BAD"
# First page of 2 over VALID records returns both good ones -- not [g1] (bad eating a slot)
# and not [] (bad filling the whole window).
# First page of 2 returns both good records, not [g1] or [].
page = gallery.list_images(limit = 2, offset = 0, valid = _valid)
assert [r["prompt"] for r in page] == ["g1", "g2"]
# The has_more probe (limit + 1) sees no extra VALID record beyond the two returned.
@ -168,8 +165,7 @@ def test_valid_callback_paginates_over_accepted_records():
def test_valid_callback_leading_bad_record_does_not_stall_at_offset_zero():
# Reproduces the exact stall: every record in the first window is schema-invalid. Without
# in-pager filtering the route returned images=[] with has_more=True at offset 0 forever.
# Every record in the first window is invalid; without in-pager filtering the route stalled.
for i in range(3):
_save_with_mtime(f"BAD{i}", 300.0 - i) # newest three are all invalid
_save_with_mtime("good", 10.0)
@ -177,15 +173,13 @@ def test_valid_callback_leading_bad_record_does_not_stall_at_offset_zero():
def _valid(rec):
return not str(rec.get("prompt", "")).startswith("BAD")
# limit+1 = 3: the pager must look PAST the invalid leaders and return the one good record,
# so has_more (len > limit) is False and the client advances off offset 0.
# The pager must look past the invalid leaders and return the one good record.
records = gallery.list_images(limit = 2, offset = 0, valid = _valid)
assert [r["prompt"] for r in records] == ["good"]
def test_save_is_atomic_no_partial_png_on_publish_failure(monkeypatch):
# A crash between writing the bytes and publishing the file must leave neither a truncated
# {id}.png nor a leftover temp: the listing only ever sees fully-written records.
# A crash before publishing must leave neither a truncated {id}.png nor a leftover temp.
def _boom(*a, **k):
raise OSError("simulated rename failure")

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

@ -181,11 +181,9 @@ def test_local_task_tags_video_pipeline_dir(tmp_path):
def test_local_task_tags_video_single_file_checkpoint(tmp_path):
# A dir whose name matches a video family holding a bare single-file .safetensors (no
# model_index.json) IS loadable: the video load route reinterprets a sole single-file local
# pick as a single_file load (routes/video.py), validating BEFORE it touches the GPU. So it
# must be tagged text-to-video and surfaced in the Video On-Device picker -- not left task=null
# and hidden, which would make the advertised-and-loadable checkpoint unusable.
# A video-family dir holding a bare single-file .safetensors (no model_index.json) is loadable
# (the route loads it as a single_file), so it must be tagged text-to-video and surfaced, not
# left task=null and hidden.
d = tmp_path / "ltx-loose"
_touch(d / "ltx-2.safetensors") # loose weights, no model_index.json
assert (

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

@ -274,11 +274,8 @@ def test_install_into_empty_dir_claims_ownership(tmp_path, monkeypatch):
def test_install_into_nonempty_unowned_dir_is_refused(tmp_path, monkeypatch):
# A pre-existing, non-empty directory that Studio did not create (e.g. a user's own
# stable-diffusion.cpp checkout) must NOT be extracted into. Merging the release into it would
# overwrite or mix our binaries into the user's working tree, and leaving it unowned only stops
# the uninstaller from deleting it later. install() refuses up front and leaves the dir untouched
# so the user can point us at a fresh/empty location.
# A pre-existing, non-empty directory Studio did not create (e.g. a user's own checkout) must
# not be extracted into; install() refuses up front and leaves it untouched.
zb = _zip_with_sd_cli()
_stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest())
target = tmp_path / "stable-diffusion.cpp"

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

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

@ -200,10 +200,8 @@ def test_list_skips_corrupt_sidecar():
def test_valid_callback_paginates_over_accepted_records():
# A sidecar that parses as JSON (so the read accepts it) but fails the caller's stricter schema
# check must be filtered BEFORE pagination, so offset/limit/has_more count over accepted records
# only. Otherwise a leading bad record returns a short/empty page with more remaining and stalls
# infinite scroll at offset 0.
# ``valid`` must filter before pagination, so offset/limit/has_more count over accepted records;
# else a leading bad record returns a short page with more remaining and stalls scroll.
_save_with_mtime("BAD", 300.0) # newest, sorts first
_save_with_mtime("g1", 200.0)
_save_with_mtime("g2", 100.0)
@ -231,10 +229,8 @@ def test_valid_callback_leading_bad_records_do_not_stall_at_offset_zero():
def test_save_leaves_no_orphan_mp4_when_sidecar_publish_fails(monkeypatch):
# The sidecar is the pair's commit marker. If it fails to publish, the MP4 must NOT be left
# behind as an invisible orphan (list_videos would skip it and gallery delete could never reach
# it). Fail the SECOND os.replace (the sidecar) after the mp4 is renamed in, and assert nothing
# is stranded.
# If the sidecar (the pair's commit marker) fails to publish, the MP4 must not be left as an
# invisible orphan. Fail the second os.replace and assert nothing is stranded.
real_replace = gallery.os.replace
calls = {"n": 0}

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

@ -111,8 +111,7 @@ class _FakeBackend(video_module.VideoBackend):
kind = (model_kind or ("gguf" if gguf_filename else "pipeline")).lower()
if kind in ("gguf", "single_file") and not gguf_filename:
raise ValueError("A gguf/single_file load needs the checkpoint filename.")
# Non-GGUF loads are gated to unsloth/* repos, the official bases, and local paths
# (the real backend trusts an existing local path via _is_trusted_video_repo).
# Non-GGUF loads are gated to unsloth/* repos, the official bases, and existing local paths.
trusted = model_path.lower().startswith(("unsloth/", "lightricks/")) or (
Path(model_path).expanduser().exists()
)
@ -386,10 +385,8 @@ def test_load_progress_route(client):
def test_load_local_single_file_dir_routes_through_single_file(client, tmp_path):
# An On-Device pick of a local directory named for a video family that holds exactly one
# .safetensors and no model_index.json arrives as a pipeline with no filename. The route
# reinterprets it as a single_file load of the sole checkpoint (mirrors the image route),
# so it is loadable instead of 400ing on the missing model_index.json.
# A local video-family dir with one .safetensors and no model_index.json arrives as a pipeline
# with no filename; the route reinterprets it as a single_file load of the sole checkpoint.
d = tmp_path / "ltx-2.3-local"
d.mkdir()
(d / "ltx-dit.safetensors").write_bytes(b"0")
@ -404,8 +401,7 @@ def test_load_local_single_file_dir_routes_through_single_file(client, tmp_path)
def test_load_local_pipeline_dir_stays_pipeline(client, tmp_path):
# A real diffusers directory (has model_index.json) is left as a pipeline load:
# resolve_local_single_file returns None for it, so the pick is not rewritten.
# A real diffusers directory (has model_index.json) is left as a pipeline load.
d = tmp_path / "ltx-2.3-pipeline"
d.mkdir()
(d / "model_index.json").write_text("{}")

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

@ -1450,13 +1450,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
pollTimer.current = setTimeout(() => void pollLoadProgress(), 1000);
}, [dismissLoadToast, refreshStatus]);
// Re-enter the per-step generation poll for a run already in flight on the backend --
// one this page did not start (another client called /images/generate, or this browser
// reloaded mid-generate). The backend runs generations under a serialising lock and keeps
// reporting progress, so track it here instead of showing a stale idle view. The images
// generate-progress carries only per-step progress (no terminal record), so on completion
// refresh the gallery to merge any image saved after the mount fetch. Kept separate from
// handleGenerate's own loop, which prepends its synchronous results directly.
// Re-enter the per-step poll for a generation already in flight on the backend that this page
// did not start (another client, or a reload mid-generate), instead of showing a stale idle
// view. generate-progress carries no terminal record, so refresh the gallery on completion to
// merge any image saved after the mount fetch. Separate from handleGenerate's own loop.
const resumeGeneratePoll = useCallback(() => {
if (genPollTimer.current) clearInterval(genPollTimer.current);
if (genVisibilityListener.current)
@ -1477,8 +1474,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
if (!isMounted.current) return;
setBusy(null);
setGenStep(null);
// The finished run saved its image(s) to the backend gallery; re-fetch the first
// page to merge them (a full replace, so deduped) and resync status.
// Re-fetch the first page to merge images the finished run saved, and resync status.
void loadGallery();
void refreshStatus();
return;
@ -1518,11 +1514,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
} catch {
// Resume is best-effort; a failed probe just leaves the idle view.
}
// A generation started elsewhere -- another client, or this browser before a reload --
// keeps running on the backend under a serialising lock. Resume tracking it so the page
// reflects the in-flight run (progress bar + disabled Generate) instead of a stale idle
// view, then refreshes the gallery when it finishes to pick up an image saved after the
// mount fetch. Mirrors the video page's mount resume.
// Resume tracking a generation started elsewhere (another client, or before a reload) so the
// page shows the in-flight run instead of a stale idle view. Mirrors the video page.
try {
const g = await getGenerateProgress();
if (g.active) {
@ -1561,11 +1554,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
if (lastLoad.current) return;
if (seededResident.current === repoId) return;
seededResident.current = repoId;
// Seed from the resolved base_repo, not repo_id: a GGUF/single_file resident (or one
// loaded from a bare local path like /models/checkpoint.safetensors) carries a repo_id
// with no family substring, so defaultsFor(repoId) would fall back to the distilled
// few-step/no-CFG recipe and the first resident generation would run with wrong defaults.
// base_repo is the resolved diffusers base (it holds the family), so prefer it when set.
// Seed from base_repo (the resolved diffusers base, holding the family), not repo_id: a
// GGUF/single_file/local-path resident has a repo_id with no family substring, so
// defaultsFor(repoId) would fall back to the wrong distilled few-step/no-CFG recipe.
const d = defaultsFor(status?.base_repo ?? repoId);
setSteps(d.steps);
setGuidance(d.guidance);

24
studio/install_sd_cpp_prebuilt.py
View file

@ -401,12 +401,9 @@ def install(
has no ``sd-cli``.
"""
target = install_dir or default_install_dir()
# Decide up front whether this install may claim ownership of ``target``. We only mark a
# directory as Studio-owned (and therefore eligible for the uninstaller's recursive delete)
# when this install actually created it or it was empty -- NEVER when it already held a user's
# own stable-diffusion.cpp checkout or unrelated files. Adopting a pre-existing, unowned,
# non-empty directory would let a later uninstall wipe the user's own work. A directory that
# already carries our marker (a prior Studio install / upgrade) stays owned.
# Only claim ownership of ``target`` (marking it for the uninstaller's recursive delete) when
# this install created it or it was empty, or it already carries our marker -- never adopt a
# pre-existing non-empty dir, else a later uninstall could wipe a user's own checkout.
marker = target / ".unsloth-studio-owned"
_may_own = True
if target.exists():
@ -418,11 +415,9 @@ def install(
_pre_existing_entries = True
# Empty dir, or one we already own, may be (re)claimed; a non-empty unowned dir may not.
_may_own = (not _pre_existing_entries) or marker.is_file()
# Refuse to extract into a pre-existing, non-empty directory we do not own (a user's own
# stable-diffusion.cpp checkout, or unrelated files beside a custom Studio root). Not writing
# the ownership marker only protects the uninstaller; extracting the release here would still
# merge our binaries into the user's working tree and can overwrite same-named files. Fail with
# a clear message so the user points us at a fresh/empty location instead of corrupting theirs.
# Refuse to extract into a pre-existing, non-empty directory we do not own: merging the release
# in would overwrite or mix our binaries into the user's own files. Fail so they point us at a
# fresh/empty location.
if not _may_own:
raise RuntimeError(
f"sd.cpp install target already exists and is not a Studio-managed directory: {target}. "
@ -470,11 +465,8 @@ def install(
_make_executable(sd_server)
if sd_server is not None:
print(f"installed sd-server -> {sd_server}", flush = True)
# Ownership marker (the same one setup.sh/_is_studio_root use, and setup.ps1 writes into the
# Node sibling dir) so the uninstaller can tell a Studio-installed sd.cpp from a user's own
# stable-diffusion.cpp checkout beside a custom Studio root, and delete only ours. Written only
# when this install created the directory or it was empty (see _may_own above): a pre-existing,
# unowned, non-empty directory keeps its unowned status so the uninstaller leaves it alone.
# Ownership marker (the same one setup.sh/_is_studio_root use) so the uninstaller deletes only
# Studio-installed sd.cpp, not a user's own checkout. Written only when _may_own (see above).
if _may_own:
try:
marker.touch()