From e0ef488f47788144ef822dfa8419bccce2cfd4a5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 05:29:09 +0000 Subject: [PATCH] Tighten comments and docstrings added by the image-generation fixes --- scripts/uninstall.ps1 | 7 ++-- studio/backend/core/inference/diffusion.py | 36 +++++++------------ .../core/inference/diffusion_engine_router.py | 19 ++++------ .../backend/core/inference/image_gallery.py | 15 ++++---- .../backend/core/inference/video_gallery.py | 17 ++++----- .../core/training/diffusion_train_common.py | 17 ++++----- studio/backend/core/training/training.py | 10 +++--- studio/backend/models/inference.py | 10 +++--- studio/backend/routes/inference.py | 9 +++-- studio/backend/routes/models.py | 12 +++---- studio/backend/routes/training.py | 31 ++++++---------- studio/backend/routes/video.py | 22 +++++------- .../backend/tests/test_diffusion_backend.py | 13 +++---- .../tests/test_diffusion_dataset_api.py | 30 ++++++---------- .../tests/test_diffusion_engine_router.py | 10 ++---- .../tests/test_diffusion_lora_trainer.py | 12 +++---- studio/backend/tests/test_diffusion_routes.py | 5 ++- studio/backend/tests/test_image_gallery.py | 18 ++++------ .../backend/tests/test_local_model_format.py | 8 ++--- studio/backend/tests/test_sd_cpp_install.py | 7 ++-- .../tests/test_training_start_offload.py | 7 ++-- studio/backend/tests/test_video_gallery.py | 12 +++---- studio/backend/tests/test_video_routes.py | 12 +++---- .../src/features/images/images-page.tsx | 29 ++++++--------- studio/install_sd_cpp_prebuilt.py | 24 +++++-------- 25 files changed, 140 insertions(+), 252 deletions(-) diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index a38f0d0eab..61bf73b5b4 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -364,11 +364,8 @@ function Uninstall-UnslothStudio { _StopByPortFile -PortFile (Join-Path $r "share\studio.port") -KnownRoots $knownRoots } _StopStudioProcesses -KnownRoots $knownRoots - # The default sd.cpp dir is only killed+deleted when it carries our owner marker (the deletion - # below is marker-gated). Passing it to the locking-process stop UNCONDITIONALLY would terminate - # a user's own sd-server running from an unowned checkout at this default path -- a process we - # then decide to keep the directory for. Gate the kill on the same predicate so stop and delete - # agree: include it only when marked owned (and present). + # Only stop the default sd.cpp dir when it carries our owner marker, so stop matches the + # marker-gated delete and a user's own sd-server at this default path is left running. $defaultSdCppToStop = $null if ($defaultSdCpp -and (Test-Path -LiteralPath $defaultSdCpp) -and (Test-Path -LiteralPath (Join-Path $defaultSdCpp ".unsloth-studio-owned") -PathType Leaf)) { $defaultSdCppToStop = $defaultSdCpp diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 9ed4cd007b..feeb967d86 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -2216,14 +2216,10 @@ class DiffusionBackend: raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) # Register under _lock so unload()/a load can signal THIS generation. self._active_generate_cancel = cancel - # Publish an active (step 0) progress state the moment the lock is held, BEFORE - # the slow pre-denoise setup (deferred compile, LoRA resolution/application, - # ControlNet download/build). Without this generate_progress() reports inactive - # across that window, so a reload's mount probe shows idle even though this - # generation holds _generate_lock; the user then starts a second generate that - # merely blocks behind this one (and can duplicate the result). The per-step - # callback swaps in its own _GenState at denoise start; this is the queued phase. - # Mirrors the video backend's queued state and the training start guard. + # Publish an active (step 0) state now, before the slow pre-denoise setup (deferred + # compile, LoRA resolution, ControlNet build), so a reload's mount probe doesn't read + # idle while this generation holds _generate_lock and let a second generate queue + # behind it. The per-step callback swaps in its own _GenState at denoise start. self._gen = _GenState(total_steps = steps) try: # The local `state` ref keeps the pipe alive even if unload() nulls _state. @@ -2559,10 +2555,8 @@ class DiffusionBackend: with self._lock: if self._active_generate_cancel is cancel: self._active_generate_cancel = None - # Drop the published progress state. The normal path already nulled it after - # the denoise; this also covers a setup-time error that skips that inner - # finally. Safe under _generate_lock: no other generation can have installed - # its own _gen while this one runs. + # Drop the published progress state, covering a setup-time error that skips + # the inner finally. Safe under _generate_lock. self._gen = None def generate_progress(self) -> dict[str, Any]: @@ -2594,13 +2588,10 @@ class DiffusionBackend: # Cancel any in-flight load (its worker checks this token) and drop the marker. self._load_token += 1 self._loading = None - # WAIT for the signalled denoise to exit BEFORE tearing down (mirrors begin_load's - # pre-teardown barrier at the top of the locked load path). _unload_locked uninstalls - # PROCESS-WIDE state the running denoise still depends on -- the eager/arch attention - # patches, the GGUF compile hooks, the flipped backend flags and the compile cache -- none - # of which the denoise's own pipe ref pins. Uninstalling them before the denoise exits would - # corrupt or crash its in-flight forward passes. The denoise holds _generate_lock for its - # whole body, so acquiring it here blocks until it has finished; only then do we tear down. + # Wait for the signalled denoise to exit BEFORE tearing down: _unload_locked uninstalls + # process-wide state (attention patches, GGUF compile hooks, backend flags, compile cache) + # the denoise still depends on but its pipe ref doesn't pin. The denoise holds _generate_lock + # for its whole body, so acquiring it here blocks until it has finished. Mirrors begin_load. with self._generate_lock: with self._lock: self._unload_locked() @@ -2622,10 +2613,9 @@ class DiffusionBackend: uninstall_patches() uninstall_arch_patches() - # NOTE: deliberately NOT unload_lora_weights() here. Both callers (unload() and begin_load's - # locked path) now hold _generate_lock across this teardown, so no denoise is in flight; the - # whole pipe is dropped below, freeing any LoRA adapters with it, so a separate adapter unload - # would be redundant work. + # Deliberately NOT unload_lora_weights() here: the whole pipe is dropped below, freeing any + # LoRA adapters with it. Both callers hold _generate_lock across this teardown, so no denoise + # is in flight. # Drop the workflow pipes so they don't pin the freed pipeline's modules past unload. self._aux_pipes.clear() # Drop any ControlNet models + pipelines so the freed load carries no extra modules. diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index 8422e7305b..4b0c9b9155 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -59,9 +59,8 @@ def _install_accelerator_for(backend: str) -> str: # The engine the current load committed to, and why a non-native choice was made. Mutated only # under _lock during selection. _lock = threading.Lock() -# Serializes a WHOLE engine switch (check -> unload -> publish). _lock alone is released during the -# slow unload(), so two overlapping selections could interleave: one observes the not-yet-updated -# active engine, returns it, and loads onto the very engine the other is concurrently unloading. +# Serializes a whole engine switch (check -> unload -> publish); _lock alone is released during the +# slow unload(), letting two overlapping selections load onto the engine the other is unloading. _transition_lock = threading.Lock() _active_engine_name: str = ENGINE_DIFFUSERS _fallback_reason: Optional[str] = None @@ -90,12 +89,9 @@ def active_engine_name() -> str: def _activate(name: str, reason: Optional[str]) -> Any: global _active_engine_name, _fallback_reason - # Serialize the ENTIRE check -> unload -> publish transition. _lock is released during the slow - # unload() below (holding it across the unload would block every status/selection reader), which - # opens a window where a second _activate could read the still-old active engine, take the "no - # change" branch, and return that engine -- then load onto it while this call is unloading it. - # _transition_lock closes the window without holding _lock across the unload; the final - # get_active_diffusion_engine() now reflects the committed state. + # Serialize the whole check -> unload -> publish transition without holding _lock across the + # slow unload(), closing the window where a second _activate reads the still-old active engine, + # takes the "no change" branch, and loads onto the engine this call is unloading. with _transition_lock: # Switching engines: unload the deactivated one first, else its model stays resident but # unreachable (the evictor only targets the active engine), leaking 10+ GB. The unload is @@ -112,9 +108,8 @@ def _activate(name: str, reason: Optional[str]) -> Any: if engine_to_unload is not None: # Publish the new engine only AFTER the old one unloads. The evictor unloads # get_active_diffusion_engine(), so flipping the name first would let a concurrent - # acquire_for evict the new (empty) engine and take the GPU while the old model is still - # freeing VRAM -- two large models briefly resident. Keeping the OLD engine as the evict - # target serializes a concurrent evict on its unload(), granting the GPU only once freed. + # acquire_for evict the new (empty) engine while the old model is still freeing VRAM. + # Keeping the OLD engine as the evict target grants the GPU only once it is freed. try: engine_to_unload.unload() except Exception as exc: # noqa: BLE001 -- best-effort; never block the switch diff --git a/studio/backend/core/inference/image_gallery.py b/studio/backend/core/inference/image_gallery.py index 8bede83282..fb1e67b096 100644 --- a/studio/backend/core/inference/image_gallery.py +++ b/studio/backend/core/inference/image_gallery.py @@ -69,9 +69,8 @@ def save(image: Any, meta: dict[str, Any]) -> dict[str, Any]: image_id = uuid.uuid4().hex directory = gallery_dir() final_path = directory / f"{image_id}.png" - # Write to a hidden temp then atomically rename into place: a crash / disk-full mid-write must - # never leave a truncated {id}.png that the listing would surface as a corrupt record. The temp - # name is dotted so an interrupted write is skipped by the *.png glob, and cleaned on failure. + # Write to a dotted temp (skipped by the *.png glob) then atomically rename, so a crash mid-write + # never leaves a truncated {id}.png that the listing would surface as a corrupt record. tmp_path = directory / f".{image_id}.png.tmp" try: tmp_path.write_bytes(_png_bytes(image, meta)) @@ -156,12 +155,10 @@ def list_images( full just to sort; only the window's recipes are read. limit=None returns everything from ``offset`` on. - ``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and the - caller's has_more all count over the same accepted-record domain. Passing the route's schema - validator here is essential: a record that has every required key (so ``_read_meta`` accepts - it) but a wrong value type would otherwise be counted in this window yet dropped by the route - after slicing -- a leading bad record then returns an empty page with more remaining, which - stalls infinite scroll at offset 0.""" + ``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and has_more + all count over the accepted-record domain. Pass the route's schema validator: a record with + every required key (so ``_read_meta`` accepts it) but a wrong value type would otherwise be + counted here yet dropped after slicing, stalling infinite scroll at offset 0.""" try: paths = list(gallery_dir().glob("*.png")) except OSError: diff --git a/studio/backend/core/inference/video_gallery.py b/studio/backend/core/inference/video_gallery.py index b703f44a42..1391119217 100644 --- a/studio/backend/core/inference/video_gallery.py +++ b/studio/backend/core/inference/video_gallery.py @@ -39,11 +39,9 @@ def save(mp4_bytes: bytes, meta: dict[str, Any]) -> dict[str, Any]: mp4_tmp = directory / f".{video_id}.mp4.tmp" sidecar = directory / f"{video_id}.json" sidecar_tmp = directory / f".{video_id}.json.tmp" - # Stage BOTH files, then publish. The sidecar is the pair's commit marker (list_videos scans - # mp4s but skips any without a readable sidecar), so writing the MP4 straight to its final name - # before the sidecar meant a sidecar write/replace failure left an invisible, undeletable orphan - # MP4 (gallery delete resolves by id, but the record never appears to be deleted). Rename the MP4 - # in first, then the sidecar; on ANY failure remove every artifact so nothing is stranded. + # Stage both files, rename the MP4 in, then the sidecar (the pair's commit marker: list_videos + # skips an mp4 without a readable sidecar). On any failure remove every artifact, else a sidecar + # failure would leave an invisible, undeletable orphan MP4. try: mp4_tmp.write_bytes(mp4_bytes) sidecar_tmp.write_text(json.dumps(meta), encoding = "utf-8") @@ -204,11 +202,10 @@ def list_videos( Ordered by MP4 mtime (a cheap stat ~= generation order); only the window's sidecars are read. limit=None returns everything from ``offset`` on. A file without its pair is skipped. - ``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and the - caller's has_more all count over the same accepted-record domain. Passing the route's schema - validator here keeps a sidecar that parses as JSON but fails the response schema from being - counted in this window yet dropped by the route after slicing -- which would otherwise return a - short or empty page with more remaining and stall infinite scroll.""" + ``valid`` (optional) filters records BEFORE pagination, so ``offset`` / ``limit`` and has_more + all count over the accepted-record domain. Pass the route's schema validator: a sidecar that + parses as JSON but fails the response schema would otherwise be counted here yet dropped after + slicing, stalling infinite scroll.""" try: paths = list(gallery_dir().glob("*.mp4")) except OSError: diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index b4fdf81d1e..dc5f5f4fc6 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -705,27 +705,22 @@ def discover_image_caption_pairs( caption: Optional[str] = None sidecar_present = False # 1. per-image sidecar caption file (the user's explicit edit; wins over metadata). - # An EMPTY sidecar is a deliberate tombstone (the labeling grid writes one when a user - # clears a metadata caption): it suppresses the metadata caption but leaves the image - # UNCAPTIONED so the dreambooth instance_prompt fallback below still applies. Treating - # "" as a present caption here would skip the instance prompt AND drop the image, so - # clearing every metadata caption in the grid would make a dreambooth run fail with - # "No captioned images found". + # An EMPTY sidecar is a deliberate tombstone (written when a user clears a caption): it + # suppresses the metadata caption but leaves the image uncaptioned so the instance_prompt + # fallback below still applies, rather than dropping the image. for ext in _CAPTION_EXTS: sidecar = img.with_suffix(ext) if sidecar.is_file(): sidecar_present = True caption = sidecar.read_text(encoding = "utf-8").strip() break - # 2. metadata row keyed by file name (basename or the relative path; as_posix so a - # Windows backslash path still matches the jsonl's forward-slash keys). Skipped when a - # sidecar tombstone is present (the sidecar, even empty, is authoritative). + # 2. metadata row keyed by file name (basename or relative path; as_posix so a Windows + # backslash path matches the jsonl's forward-slash keys). A sidecar, even empty, wins. if not sidecar_present: caption = meta_caption.get(img.name) or meta_caption.get( img.relative_to(root).as_posix() ) - # 3. dreambooth instance prompt for any image still without a caption (no sidecar text and - # no metadata row, or an empty tombstone). + # 3. dreambooth instance prompt for any image still without a caption. if not caption and instance_prompt: caption = instance_prompt if caption: diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 30dda0f7d0..526551aab4 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -1141,12 +1141,10 @@ class TrainingBackend: # training invisibly behind a frozen UI. Cheap enough for per-second polls. self._ensure_pump_alive() with self._lock: - # A start reserved via the compare-and-set guard in start_training but not yet - # spawned (before_spawn frees residents, then GPU auto-selection, then proc.start()) - # is already "active": the load/start guards read this to refuse a concurrent - # /images/load, /video/load, or /diffusion/start, so treating the pre-spawn window - # as idle would let another pipeline race the reserved run for VRAM. Mirrors the - # diffusion training service's reserve()/is_active(). + # A run reserved in start_training but not yet spawned (before_spawn frees residents, + # then GPU auto-selection, then proc.start()) is already active: the load/start guards + # read this to refuse a concurrent /images/load, /video/load, or /diffusion/start, so an + # idle reading here would let another pipeline race the reserved run for VRAM. if self._start_in_progress: return True diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 862942e883..23218b7f35 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2071,12 +2071,10 @@ class DiffusionGenerateRequest(BaseModel): @model_validator(mode = "after") def _batch_seeds_json_safe(self) -> "DiffusionGenerateRequest": - # A batch derives per-image seeds as seed, seed+1, ... seed+batch_size-1 (sd.cpp and - # diffusers both advance the seed per image). The base seed is bounded to 2**53-1 so it - # round-trips through the JSON gallery recipe, but the derived top-of-batch seed is not: - # with an explicit seed near the cap it can exceed Number.MAX_SAFE_INTEGER, where the - # frontend rounds it and a restored recipe replays a different image. Reject at the - # boundary so an API client can't persist an unreplayable seed. + # A batch derives per-image seeds as seed .. seed+batch_size-1. The base seed is capped at + # 2**53-1 to round-trip through the JSON recipe, but a derived top-of-batch seed near the cap + # can exceed it, where the frontend rounds it and a restored recipe replays a different + # image. Reject at the boundary so an API client can't persist an unreplayable seed. if self.seed is not None and self.seed + self.batch_size - 1 > 2**53 - 1: raise ValueError( "seed + batch_size - 1 must not exceed 2**53 - 1 so every per-image seed " diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index c452a4775a..fab3652b7f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -14512,11 +14512,10 @@ async def list_gallery_images( limit = max(1, min(limit, 200)) offset = max(0, offset) - # Validate against the response schema INSIDE the pager so offset / limit / has_more all count - # over the same accepted-record domain. A PNG whose recipe chunk has all keys but a wrong value - # type passes the presence-only read yet fails GalleryImage(**r); dropping such records only - # after pagination made a leading bad record return an empty page with has_more=True, stalling - # infinite scroll at offset 0. Filtering here keeps the window and has_more consistent. + # Validate inside the pager so offset / limit / has_more all count over the accepted domain. A + # recipe with all keys but a wrong value type passes the presence-only read yet fails + # GalleryImage(**r); dropping it only after slicing let a leading bad record return an empty + # page with has_more=True, stalling infinite scroll at offset 0. def _valid_gallery_image(record: dict) -> bool: try: GalleryImage(**record) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 71f4c2f41d..fe11446561 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -3320,13 +3320,11 @@ def _local_is_diffusers(model: "LocalModelInfo") -> bool: return True except Exception: pass - # A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json) - # ships no pipeline index and resolves to no IMAGE family either, so the checks above miss - # it. The video load route reinterprets a bare single-file local dir as a single_file load - # (routes/video.py), so it IS loadable and must be surfaced; without tagging it here - # _local_model_task returns task=null and the Video On-Device picker hides it. Match the same - # clean id / name needles (not the raw path) so a parent-dir family token can't spuriously - # match; _local_model_task then routes the video family to the text-to-video task. + # A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json) has no + # pipeline index and no image family, so the checks above miss it. The video load route loads it + # as a single_file (routes/video.py), so it must be surfaced or _local_model_task returns + # task=null and the picker hides it. Match clean id / name needles (not the raw path) so a + # parent-dir token can't spuriously match; _local_model_task then routes it to text-to-video. try: from core.inference.video_families import detect_video_family for needle in (model.model_id, model.display_name, Path(model.id).name): diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 3ce107a323..9d28bec7dc 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1616,11 +1616,8 @@ async def upload_diffusion_dataset( from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label cleaned = _clean_diffusion_dataset_name(name) - # Reject a symlinked dataset directory BEFORE any write. A bare mkdir(exist_ok=True) succeeds - # through an existing name -> external-directory symlink, and the staged upload would then - # write/replace files outside the Studio datasets root through that link. The read/caption/ - # delete endpoints already enforce this containment via _resolve_dataset_folder; the upload - # path must run the same symlink + root-containment check first so writes can never escape. + # Run the same symlink + root-containment check as the read/caption/delete endpoints before any + # write, so a name -> external-directory symlink can't make the staged upload write outside root. folder = _resolve_dataset_folder(name, must_exist = False) folder.mkdir(parents = True, exist_ok = True) @@ -1736,13 +1733,10 @@ async def upload_diffusion_dataset( ) out.write(chunk) uploaded += 1 - # Commit every staged file as one transaction. A plain replace loop is NOT atomic across - # files: if the second-or-later tmp.replace(dest) fails (disk/quota error, a Windows file - # lock, antivirus, a destination that became a directory), an earlier destination has - # already been overwritten while the request returns an error -- the user's original file - # is gone. Back up each pre-existing destination before overwriting it, then on ANY failure - # remove the versions this request installed and restore every displaced original, so the - # dataset is left exactly as it was before the upload. + # Commit every staged file as one transaction. A plain replace loop is not atomic across + # files: a mid-loop tmp.replace(dest) failure leaves earlier destinations already overwritten + # while the request errors. Back up each pre-existing destination first, then on any failure + # drop the versions this request installed and restore every displaced original. backups: list[tuple[Path, Optional[Path]]] = [] # (dest, backup path or None) installed: list[Path] = [] try: @@ -1809,11 +1803,9 @@ def _resolve_dataset_folder(name: str, *, must_exist: bool = True) -> Path: cleaned = _clean_diffusion_dataset_name(name) root = datasets_root().resolve() folder = root / cleaned - # Reject a symlinked dataset directory. _safe_dataset_image_path only proves each image path - # stays under folder.resolve(); it never proves the folder itself stays under the datasets - # root. A dataset dir that is a symlink to an external directory would therefore let image - # read / caption / delete operate on files outside Studio (a reproduced delete removed an - # external file through such a link). Prove the resolved folder is contained in the root too. + # Reject a symlinked dataset directory and prove the resolved folder stays under root: + # _safe_dataset_image_path only checks each image path, so a folder symlinked to an external + # directory would let read / caption / delete operate on files outside Studio. if folder.is_symlink(): raise HTTPException( status_code = 400, @@ -2063,9 +2055,8 @@ async def delete_diffusion_dataset_image( if thumbs_dir.is_dir(): # Thumbs are keyed on the full filename (stem + extension), so match that here too; # a stem-only glob would strand this image's thumbs or delete a same-stem sibling's. - # Escape the filename first: an uploaded name may legally contain glob metacharacters - # ('[', ']', '*', '?'), and interpolating those raw would make e.g. "[ab].png" match - # "a.png_*.jpg"/"b.png_*.jpg" -- deleting siblings' thumbs while leaving its own behind. + # Escape the name: a raw glob metacharacter (e.g. "[ab].png") would match siblings' + # thumbs while leaving its own behind. for t in thumbs_dir.glob(f"{_glob.escape(image_path.name)}_*.jpg"): t.unlink(missing_ok = True) return {"deleted": image_path.name} diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 04ea6d7d28..7e239151c5 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -85,15 +85,12 @@ async def load_video_model( backend = get_video_backend() try: - # Resolve the load kind once (gguf / single_file / pipeline) so validation and the - # load agree; a bad explicit kind raises here -> 400. + # Resolve the load kind once (gguf / single_file / pipeline) so validation and the load + # agree; a bad explicit kind raises here -> 400. kind = resolve_video_model_kind(request.gguf_filename, request.model_kind) - # A local On-Device pick can be a bare single-file .safetensors directory (no - # model_index.json): the scanner advertises it as a text-to-video model, but the - # local picker starts it as a pipeline with no filename, so a pipeline load would - # 400 on the missing model_index.json and the advertised model is unusable. If the - # directory holds exactly one checkpoint, reinterpret the pick as a single_file load - # of it, so validation and the load agree. Mirrors the image load route. + # A local On-Device pick can be a bare single-file .safetensors dir (no model_index.json) + # that the picker starts as a pipeline with no filename, which would 400 on the missing + # index. If the dir holds exactly one checkpoint, load it as a single_file. Mirrors images. if kind == "pipeline" and not request.gguf_filename: sole = await asyncio.to_thread(resolve_local_single_file, request.model_path) if sole is not None: @@ -236,11 +233,10 @@ async def list_gallery_videos( limit = max(1, min(limit, 200)) offset = max(0, offset) - # Validate against the response schema INSIDE the pager so offset / limit / has_more all count - # over the same accepted-record domain. A sidecar that parses as JSON but has a wrong value type - # passes the read yet fails GalleryVideo(**r); dropping such records only after pagination made a - # leading bad record return an empty page with has_more=True, stalling infinite scroll at offset - # 0. Filtering here keeps the window and has_more consistent. + # Validate inside the pager so offset / limit / has_more all count over the accepted domain. A + # sidecar that parses as JSON but has a wrong value type passes the read yet fails + # GalleryVideo(**r); dropping it only after slicing let a leading bad record return an empty + # page with has_more=True, stalling infinite scroll at offset 0. def _valid_gallery_video(record: dict) -> bool: try: GalleryVideo(**record) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 477577ba35..a3c7fb4bb1 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -483,11 +483,8 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): def test_generate_progress_active_during_setup(fake_runtime, tmp_path, monkeypatch): - # A generation must report active from the moment it holds the lock, BEFORE the slow - # pre-denoise setup (deferred compile / LoRA resolution / ControlNet build) runs. - # Otherwise a reload's mount probe sees idle while the lock is held and lets a second - # generate queue behind the first. _apply_loras runs inside that setup window, so probing - # generate_progress() from there exercises the gap the reviewer flagged. + # A generation must report active from the moment it holds the lock, before the slow pre-denoise + # setup. _apply_loras runs inside that window, so probe generate_progress() from there. (tmp_path / "model.gguf").write_bytes(b"weights") backend = DiffusionBackend() backend.load_pipeline( @@ -3009,10 +3006,8 @@ def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path): def test_unload_waits_for_in_flight_denoise_before_teardown(): - # Regression for the unload/denoise teardown race: unload() must WAIT for a running denoise to - # exit (acquire _generate_lock) BEFORE _unload_locked() tears down PROCESS-WIDE state (eager / - # arch attention patches, gguf compile hooks, backend flags, compile cache) that the denoise - # still depends on. Mirror the load path, which already waits on _generate_lock before teardown. + # Regression: unload() must wait for a running denoise to exit (acquire _generate_lock) before + # _unload_locked() tears down process-wide state the denoise still depends on. import threading backend = DiffusionBackend() diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index aa59adef10..60170ee273 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -496,12 +496,8 @@ def test_import_promotion_leaves_no_partial_dataset_on_failure(ds_root, monkeypa def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch): - # Re-uploading a.txt and b.txt (an allowed overwrite of two files already on disk) stages both, - # then commits them. A plain replace loop is NOT atomic: if the SECOND commit fails after the - # first destination was already overwritten, the live dataset is left partially updated with the - # request returning an error -- the user's original a.txt is gone. The transactional commit must - # back up each displaced original and, on any failure, restore every one so the dataset is left - # exactly as it was, with no stray temp/backup files. + # Re-uploading a.txt and b.txt where the SECOND commit fails must roll back the first overwrite, + # so both originals survive and no stray temp/backup files remain. from pathlib import Path app = FastAPI() @@ -518,8 +514,7 @@ def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch): state = {"failed": False} def flaky_replace(self, target, *a, **k): - # Fail exactly once, on the tmp -> b.txt promotion (source is the staged .upload-* part, NOT - # a .upload-backup-* part), so the subsequent backup -> b.txt restore still succeeds. + # Fail once on the tmp -> b.txt promotion only (not the backup restore), so rollback works. if ( not state["failed"] and str(target).endswith("b.txt") @@ -548,9 +543,8 @@ def test_upload_rolls_back_when_a_later_promotion_fails(ds_root, monkeypatch): def test_resolve_dataset_folder_rejects_symlink(ds_root, tmp_path): - # A dataset directory that is a symlink pointing OUTSIDE the datasets root must be rejected: the - # per-image containment check only proves paths stay under folder.resolve(), so without this a - # delete/caption/read could operate on external files through the link. + # A dataset dir that is a symlink outside the datasets root must be rejected, else delete / + # caption / read could operate on external files through the link. from routes.training import _resolve_dataset_folder external = tmp_path / "external" @@ -564,9 +558,8 @@ def test_resolve_dataset_folder_rejects_symlink(ds_root, tmp_path): def test_upload_through_symlinked_dataset_cannot_escape_root(client, ds_root, tmp_path): - # End to end: an upload targeting a dataset name that already exists as a symlink to an - # external directory must be refused (400) BEFORE any bytes are written, so the upload can - # never create/replace files outside the datasets root through the link. + # An upload to a dataset name that is a symlink to an external directory must be refused (400) + # before any bytes are written. external = tmp_path / "external" external.mkdir() (ds_root / "linked").symlink_to(external, target_is_directory = True) @@ -580,8 +573,7 @@ def test_upload_through_symlinked_dataset_cannot_escape_root(client, ds_root, tm def test_delete_through_symlinked_dataset_cannot_escape_root(client, ds_root, tmp_path): - # End to end: a DELETE against an image inside a symlinked dataset dir is refused (400) and the - # external file it points at is NOT removed. + # A DELETE inside a symlinked dataset dir is refused (400) and the external file survives. external = tmp_path / "external" external.mkdir() victim = external / "victim.png" @@ -594,10 +586,8 @@ def test_delete_through_symlinked_dataset_cannot_escape_root(client, ds_root, tm def test_delete_image_with_glob_chars_only_removes_own_thumbs(client, ds_root): - # A legal uploaded filename may contain glob metacharacters ('[', ']', '*', '?'). Deleting it - # must remove ONLY its own thumbnails; interpolating the raw name into Path.glob would make - # "[ab].png" match "a.png_*.jpg"/"b.png_*.jpg" -- deleting a sibling's thumbs while leaving its - # own (literally "[ab].png_32.jpg") behind. + # Deleting a filename with glob metacharacters (e.g. "[ab].png") must remove only its own + # thumbnails, not a sibling's that the raw glob would spuriously match. from urllib.parse import quote folder = ds_root / "d" diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py index 4ff9bc23c7..18b9f51b62 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -241,11 +241,8 @@ def test_no_switch_keeps_engine_and_refreshes_reason(monkeypatch): def test_activate_serializes_switch_and_concurrent_query(monkeypatch): - # Regression for the check->unload->publish race: _activate releases _lock during the slow - # unload(), so without the transition lock a second _activate could observe the not-yet-updated - # active engine, take the "no change" branch, and return the engine the first call is - # concurrently unloading. Drive both paths on threads and assert the concurrent query is blocked - # until the switch completes (i.e. the whole transition is serialized). + # Regression: without the transition lock a second _activate during the slow unload() reads the + # not-yet-updated active engine and returns it. Assert the query is blocked until the switch ends. import threading r._active_engine_name = ENGINE_DIFFUSERS @@ -279,8 +276,7 @@ def test_activate_serializes_switch_and_concurrent_query(monkeypatch): q = threading.Thread(target = _query) q.start() - # Serialized: while the switch holds the transition lock the query cannot complete. Pre-fix it - # would return immediately (active is still diffusers), setting query_done at once. + # Serialized: the query cannot complete while the switch holds the transition lock. assert not query_done.wait(0.4) release_unload.set() diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index c06d8fd45a..6ae96d97e5 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -58,11 +58,8 @@ def test_discover_sidecar_overrides_metadata_row(tmp_path): def test_discover_empty_sidecar_suppresses_metadata_but_uses_instance_prompt(tmp_path): - # Clearing a metadata caption in the labeling grid writes an EMPTY sidecar tombstone. It must - # suppress the metadata caption (so the stale text is not resurrected) yet leave the image - # UNCAPTIONED so the dreambooth instance_prompt still applies. Treating "" as a present caption - # would skip the instance prompt AND drop the image, so a dataset whose every caption was - # cleared would fail with "No captioned images found". + # An empty sidecar tombstone must suppress the metadata caption yet leave the image uncaptioned + # so the dreambooth instance_prompt still applies (not drop the image). _touch(tmp_path / "cat.png") (tmp_path / "metadata.jsonl").write_text( json.dumps({"file_name": "cat.png", "text": "old metadata caption"}) + "\n", @@ -74,9 +71,8 @@ def test_discover_empty_sidecar_suppresses_metadata_but_uses_instance_prompt(tmp def test_discover_empty_sidecar_without_instance_prompt_skips_image(tmp_path): - # With no instance prompt an empty tombstone leaves the image uncaptioned, so it is skipped and - # the suppressed metadata caption is NOT resurrected -- while a sibling with a real caption is - # still discovered. + # With no instance prompt the tombstoned image is skipped (metadata not resurrected), while a + # sibling with a real caption is still discovered. _touch(tmp_path / "cat.png") _touch(tmp_path / "cap.png") (tmp_path / "metadata.jsonl").write_text( diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index a5a825ee00..c00d9e71d5 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -344,9 +344,8 @@ def test_generate_rejects_batch_seed_past_json_safe_range(client): client.post( "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} ) - # An explicit seed at the JS-safe cap with a batch derives per-image seeds - # (seed+1 ...) that exceed Number.MAX_SAFE_INTEGER and no longer round-trip - # through the gallery JSON recipe, so the request is rejected at the boundary. + # A seed at the cap with a batch derives per-image seeds (seed+1 ...) past the JSON-safe range, + # so the request is rejected. over = client.post( "/api/inference/images/generate", json = {"prompt": "p", "seed": 2**53 - 1, "batch_size": 2}, diff --git a/studio/backend/tests/test_image_gallery.py b/studio/backend/tests/test_image_gallery.py index b9b3c0ed49..51f107ffd2 100644 --- a/studio/backend/tests/test_image_gallery.py +++ b/studio/backend/tests/test_image_gallery.py @@ -148,10 +148,8 @@ def test_list_skips_recipe_missing_required_fields(tmp_path): def test_valid_callback_paginates_over_accepted_records(): - # A record that passes _read_meta (every required key present) but fails the caller's stricter - # schema check must be filtered BEFORE pagination, so offset/limit/has_more all count over the - # accepted domain. Otherwise a leading bad record returns an empty/short page with more still - # remaining, and the frontend (which advances by valid records) stalls at offset 0. + # ``valid`` must filter before pagination, so offset/limit/has_more count over the accepted + # domain; else a leading bad record returns a short page with more remaining and stalls scroll. _save_with_mtime("BAD", 300.0) # newest, sorts first _save_with_mtime("g1", 200.0) _save_with_mtime("g2", 100.0) @@ -159,8 +157,7 @@ def test_valid_callback_paginates_over_accepted_records(): def _valid(rec): return rec.get("prompt") != "BAD" - # First page of 2 over VALID records returns both good ones -- not [g1] (bad eating a slot) - # and not [] (bad filling the whole window). + # First page of 2 returns both good records, not [g1] or []. page = gallery.list_images(limit = 2, offset = 0, valid = _valid) assert [r["prompt"] for r in page] == ["g1", "g2"] # The has_more probe (limit + 1) sees no extra VALID record beyond the two returned. @@ -168,8 +165,7 @@ def test_valid_callback_paginates_over_accepted_records(): def test_valid_callback_leading_bad_record_does_not_stall_at_offset_zero(): - # Reproduces the exact stall: every record in the first window is schema-invalid. Without - # in-pager filtering the route returned images=[] with has_more=True at offset 0 forever. + # Every record in the first window is invalid; without in-pager filtering the route stalled. for i in range(3): _save_with_mtime(f"BAD{i}", 300.0 - i) # newest three are all invalid _save_with_mtime("good", 10.0) @@ -177,15 +173,13 @@ def test_valid_callback_leading_bad_record_does_not_stall_at_offset_zero(): def _valid(rec): return not str(rec.get("prompt", "")).startswith("BAD") - # limit+1 = 3: the pager must look PAST the invalid leaders and return the one good record, - # so has_more (len > limit) is False and the client advances off offset 0. + # The pager must look past the invalid leaders and return the one good record. records = gallery.list_images(limit = 2, offset = 0, valid = _valid) assert [r["prompt"] for r in records] == ["good"] def test_save_is_atomic_no_partial_png_on_publish_failure(monkeypatch): - # A crash between writing the bytes and publishing the file must leave neither a truncated - # {id}.png nor a leftover temp: the listing only ever sees fully-written records. + # A crash before publishing must leave neither a truncated {id}.png nor a leftover temp. def _boom(*a, **k): raise OSError("simulated rename failure") diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index 7163db424b..202b9a1b72 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -181,11 +181,9 @@ def test_local_task_tags_video_pipeline_dir(tmp_path): def test_local_task_tags_video_single_file_checkpoint(tmp_path): - # A dir whose name matches a video family holding a bare single-file .safetensors (no - # model_index.json) IS loadable: the video load route reinterprets a sole single-file local - # pick as a single_file load (routes/video.py), validating BEFORE it touches the GPU. So it - # must be tagged text-to-video and surfaced in the Video On-Device picker -- not left task=null - # and hidden, which would make the advertised-and-loadable checkpoint unusable. + # A video-family dir holding a bare single-file .safetensors (no model_index.json) is loadable + # (the route loads it as a single_file), so it must be tagged text-to-video and surfaced, not + # left task=null and hidden. d = tmp_path / "ltx-loose" _touch(d / "ltx-2.safetensors") # loose weights, no model_index.json assert ( diff --git a/studio/backend/tests/test_sd_cpp_install.py b/studio/backend/tests/test_sd_cpp_install.py index 3a061c9fcf..89404fa3ae 100644 --- a/studio/backend/tests/test_sd_cpp_install.py +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -274,11 +274,8 @@ def test_install_into_empty_dir_claims_ownership(tmp_path, monkeypatch): def test_install_into_nonempty_unowned_dir_is_refused(tmp_path, monkeypatch): - # A pre-existing, non-empty directory that Studio did not create (e.g. a user's own - # stable-diffusion.cpp checkout) must NOT be extracted into. Merging the release into it would - # overwrite or mix our binaries into the user's working tree, and leaving it unowned only stops - # the uninstaller from deleting it later. install() refuses up front and leaves the dir untouched - # so the user can point us at a fresh/empty location. + # A pre-existing, non-empty directory Studio did not create (e.g. a user's own checkout) must + # not be extracted into; install() refuses up front and leaves it untouched. zb = _zip_with_sd_cli() _stub_release(monkeypatch, zip_bytes = zb, digest = "sha256:" + hashlib.sha256(zb).hexdigest()) target = tmp_path / "stable-diffusion.cpp" diff --git a/studio/backend/tests/test_training_start_offload.py b/studio/backend/tests/test_training_start_offload.py index 1d10355dda..1b53868427 100644 --- a/studio/backend/tests/test_training_start_offload.py +++ b/studio/backend/tests/test_training_start_offload.py @@ -116,11 +116,8 @@ def test_backend_start_guard_blocks_overlapping_starts(): def test_is_training_active_true_during_start_reservation(): - # While start_training holds the compare-and-set reservation but has not yet spawned - # (before_spawn frees residents, then GPU auto-selection, then proc.start()), the LLM - # training run must already read as active: /images/load, /video/load, and - # /diffusion/start all gate on is_training_active(), so an idle reading in this window - # would let another pipeline race the reserved run for the just-freed VRAM. + # A run reserved in start_training but not yet spawned must already read as active, else the + # load/start guards would let another pipeline race it for the just-freed VRAM. from core.training.training import TrainingBackend backend = TrainingBackend() diff --git a/studio/backend/tests/test_video_gallery.py b/studio/backend/tests/test_video_gallery.py index f86dbe3bcd..6ca8ea2e35 100644 --- a/studio/backend/tests/test_video_gallery.py +++ b/studio/backend/tests/test_video_gallery.py @@ -200,10 +200,8 @@ def test_list_skips_corrupt_sidecar(): def test_valid_callback_paginates_over_accepted_records(): - # A sidecar that parses as JSON (so the read accepts it) but fails the caller's stricter schema - # check must be filtered BEFORE pagination, so offset/limit/has_more count over accepted records - # only. Otherwise a leading bad record returns a short/empty page with more remaining and stalls - # infinite scroll at offset 0. + # ``valid`` must filter before pagination, so offset/limit/has_more count over accepted records; + # else a leading bad record returns a short page with more remaining and stalls scroll. _save_with_mtime("BAD", 300.0) # newest, sorts first _save_with_mtime("g1", 200.0) _save_with_mtime("g2", 100.0) @@ -231,10 +229,8 @@ def test_valid_callback_leading_bad_records_do_not_stall_at_offset_zero(): def test_save_leaves_no_orphan_mp4_when_sidecar_publish_fails(monkeypatch): - # The sidecar is the pair's commit marker. If it fails to publish, the MP4 must NOT be left - # behind as an invisible orphan (list_videos would skip it and gallery delete could never reach - # it). Fail the SECOND os.replace (the sidecar) after the mp4 is renamed in, and assert nothing - # is stranded. + # If the sidecar (the pair's commit marker) fails to publish, the MP4 must not be left as an + # invisible orphan. Fail the second os.replace and assert nothing is stranded. real_replace = gallery.os.replace calls = {"n": 0} diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index 275ca2fb6d..0793d13788 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -111,8 +111,7 @@ class _FakeBackend(video_module.VideoBackend): kind = (model_kind or ("gguf" if gguf_filename else "pipeline")).lower() if kind in ("gguf", "single_file") and not gguf_filename: raise ValueError("A gguf/single_file load needs the checkpoint filename.") - # Non-GGUF loads are gated to unsloth/* repos, the official bases, and local paths - # (the real backend trusts an existing local path via _is_trusted_video_repo). + # Non-GGUF loads are gated to unsloth/* repos, the official bases, and existing local paths. trusted = model_path.lower().startswith(("unsloth/", "lightricks/")) or ( Path(model_path).expanduser().exists() ) @@ -386,10 +385,8 @@ def test_load_progress_route(client): def test_load_local_single_file_dir_routes_through_single_file(client, tmp_path): - # An On-Device pick of a local directory named for a video family that holds exactly one - # .safetensors and no model_index.json arrives as a pipeline with no filename. The route - # reinterprets it as a single_file load of the sole checkpoint (mirrors the image route), - # so it is loadable instead of 400ing on the missing model_index.json. + # A local video-family dir with one .safetensors and no model_index.json arrives as a pipeline + # with no filename; the route reinterprets it as a single_file load of the sole checkpoint. d = tmp_path / "ltx-2.3-local" d.mkdir() (d / "ltx-dit.safetensors").write_bytes(b"0") @@ -404,8 +401,7 @@ def test_load_local_single_file_dir_routes_through_single_file(client, tmp_path) def test_load_local_pipeline_dir_stays_pipeline(client, tmp_path): - # A real diffusers directory (has model_index.json) is left as a pipeline load: - # resolve_local_single_file returns None for it, so the pick is not rewritten. + # A real diffusers directory (has model_index.json) is left as a pipeline load. d = tmp_path / "ltx-2.3-pipeline" d.mkdir() (d / "model_index.json").write_text("{}") diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 7338f23b48..2f6d70936d 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -1450,13 +1450,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) { pollTimer.current = setTimeout(() => void pollLoadProgress(), 1000); }, [dismissLoadToast, refreshStatus]); - // Re-enter the per-step generation poll for a run already in flight on the backend -- - // one this page did not start (another client called /images/generate, or this browser - // reloaded mid-generate). The backend runs generations under a serialising lock and keeps - // reporting progress, so track it here instead of showing a stale idle view. The images - // generate-progress carries only per-step progress (no terminal record), so on completion - // refresh the gallery to merge any image saved after the mount fetch. Kept separate from - // handleGenerate's own loop, which prepends its synchronous results directly. + // Re-enter the per-step poll for a generation already in flight on the backend that this page + // did not start (another client, or a reload mid-generate), instead of showing a stale idle + // view. generate-progress carries no terminal record, so refresh the gallery on completion to + // merge any image saved after the mount fetch. Separate from handleGenerate's own loop. const resumeGeneratePoll = useCallback(() => { if (genPollTimer.current) clearInterval(genPollTimer.current); if (genVisibilityListener.current) @@ -1477,8 +1474,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) { if (!isMounted.current) return; setBusy(null); setGenStep(null); - // The finished run saved its image(s) to the backend gallery; re-fetch the first - // page to merge them (a full replace, so deduped) and resync status. + // Re-fetch the first page to merge images the finished run saved, and resync status. void loadGallery(); void refreshStatus(); return; @@ -1518,11 +1514,8 @@ export function ImagesPage({ active = true }: { active?: boolean }) { } catch { // Resume is best-effort; a failed probe just leaves the idle view. } - // A generation started elsewhere -- another client, or this browser before a reload -- - // keeps running on the backend under a serialising lock. Resume tracking it so the page - // reflects the in-flight run (progress bar + disabled Generate) instead of a stale idle - // view, then refreshes the gallery when it finishes to pick up an image saved after the - // mount fetch. Mirrors the video page's mount resume. + // Resume tracking a generation started elsewhere (another client, or before a reload) so the + // page shows the in-flight run instead of a stale idle view. Mirrors the video page. try { const g = await getGenerateProgress(); if (g.active) { @@ -1561,11 +1554,9 @@ export function ImagesPage({ active = true }: { active?: boolean }) { if (lastLoad.current) return; if (seededResident.current === repoId) return; seededResident.current = repoId; - // Seed from the resolved base_repo, not repo_id: a GGUF/single_file resident (or one - // loaded from a bare local path like /models/checkpoint.safetensors) carries a repo_id - // with no family substring, so defaultsFor(repoId) would fall back to the distilled - // few-step/no-CFG recipe and the first resident generation would run with wrong defaults. - // base_repo is the resolved diffusers base (it holds the family), so prefer it when set. + // Seed from base_repo (the resolved diffusers base, holding the family), not repo_id: a + // GGUF/single_file/local-path resident has a repo_id with no family substring, so + // defaultsFor(repoId) would fall back to the wrong distilled few-step/no-CFG recipe. const d = defaultsFor(status?.base_repo ?? repoId); setSteps(d.steps); setGuidance(d.guidance); diff --git a/studio/install_sd_cpp_prebuilt.py b/studio/install_sd_cpp_prebuilt.py index 1923cd3f8e..67ff5f6635 100644 --- a/studio/install_sd_cpp_prebuilt.py +++ b/studio/install_sd_cpp_prebuilt.py @@ -401,12 +401,9 @@ def install( has no ``sd-cli``. """ target = install_dir or default_install_dir() - # Decide up front whether this install may claim ownership of ``target``. We only mark a - # directory as Studio-owned (and therefore eligible for the uninstaller's recursive delete) - # when this install actually created it or it was empty -- NEVER when it already held a user's - # own stable-diffusion.cpp checkout or unrelated files. Adopting a pre-existing, unowned, - # non-empty directory would let a later uninstall wipe the user's own work. A directory that - # already carries our marker (a prior Studio install / upgrade) stays owned. + # Only claim ownership of ``target`` (marking it for the uninstaller's recursive delete) when + # this install created it or it was empty, or it already carries our marker -- never adopt a + # pre-existing non-empty dir, else a later uninstall could wipe a user's own checkout. marker = target / ".unsloth-studio-owned" _may_own = True if target.exists(): @@ -418,11 +415,9 @@ def install( _pre_existing_entries = True # Empty dir, or one we already own, may be (re)claimed; a non-empty unowned dir may not. _may_own = (not _pre_existing_entries) or marker.is_file() - # Refuse to extract into a pre-existing, non-empty directory we do not own (a user's own - # stable-diffusion.cpp checkout, or unrelated files beside a custom Studio root). Not writing - # the ownership marker only protects the uninstaller; extracting the release here would still - # merge our binaries into the user's working tree and can overwrite same-named files. Fail with - # a clear message so the user points us at a fresh/empty location instead of corrupting theirs. + # Refuse to extract into a pre-existing, non-empty directory we do not own: merging the release + # in would overwrite or mix our binaries into the user's own files. Fail so they point us at a + # fresh/empty location. if not _may_own: raise RuntimeError( f"sd.cpp install target already exists and is not a Studio-managed directory: {target}. " @@ -470,11 +465,8 @@ def install( _make_executable(sd_server) if sd_server is not None: print(f"installed sd-server -> {sd_server}", flush = True) - # Ownership marker (the same one setup.sh/_is_studio_root use, and setup.ps1 writes into the - # Node sibling dir) so the uninstaller can tell a Studio-installed sd.cpp from a user's own - # stable-diffusion.cpp checkout beside a custom Studio root, and delete only ours. Written only - # when this install created the directory or it was empty (see _may_own above): a pre-existing, - # unowned, non-empty directory keeps its unowned status so the uninstaller leaves it alone. + # Ownership marker (the same one setup.sh/_is_studio_root use) so the uninstaller deletes only + # Studio-installed sd.cpp, not a user's own checkout. Written only when _may_own (see above). if _may_own: try: marker.touch()